diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index baa291f..07374d2 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -19,10 +19,8 @@ from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.buffers import ReplayBuffer from stable_baselines3.common.save_util import save_to_pkl, load_from_pkl -_BufferType = TypeVar("_BufferType", bound=ReplayBuffer) - -class OffPolicyAlgorithm(BaseAlgorithm, Generic[_BufferType]): +class OffPolicyAlgorithm(BaseAlgorithm): """ The base for Off-Policy algorithms (ex: SAC/TD3) @@ -75,7 +73,7 @@ class OffPolicyAlgorithm(BaseAlgorithm, Generic[_BufferType]): env: Union[GymEnv, str], policy_base: Type[BasePolicy], learning_rate: Union[float, Callable], - replay_buffer_cls: Type[_BufferType] = None, + replay_buffer_cls: Type[ReplayBuffer] = None, replay_buffer_kwargs: Optional[Mapping[str, Any]] = None, buffer_size: int = int(1e6), learning_starts: int = 100, diff --git a/tests/test_buffers.py b/tests/test_buffers.py index a441678..1bb7f9c 100644 --- a/tests/test_buffers.py +++ b/tests/test_buffers.py @@ -101,4 +101,4 @@ def test_with_algo(algo): agent.replay_buffer_cls = NstepReplayBuffer agent.replay_buffer_kwargs.update({"gamma":agent.gamma, "n_step": 10}) agent._setup_model() - agent.learn(500) \ No newline at end of file + agent.learn(500)