diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 35af80c..36ad34b 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -2,7 +2,7 @@ import io import pathlib import time import warnings -from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Type, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union import gym import numpy as np @@ -76,8 +76,8 @@ class OffPolicyAlgorithm(BaseAlgorithm): env: Union[GymEnv, str], policy_base: Type[BasePolicy], learning_rate: Union[float, Callable], - replay_buffer_cls: Type[ReplayBuffer] = None, - replay_buffer_kwargs: Optional[Mapping[str, Any]] = None, + replay_buffer_class: Optional[Type[ReplayBuffer]] = None, + replay_buffer_kwargs: Optional[Dict[str, Any]] = None, buffer_size: int = int(1e6), learning_starts: int = 100, batch_size: int = 256, @@ -148,13 +148,13 @@ class OffPolicyAlgorithm(BaseAlgorithm): self.use_sde_at_warmup = use_sde_at_warmup self.replay_buffer = None - self.replay_buffer_cls = replay_buffer_cls or ReplayBuffer + self.replay_buffer_class = replay_buffer_class or ReplayBuffer self.replay_buffer_kwargs = dict(replay_buffer_kwargs or {}) def _setup_model(self) -> None: self._setup_lr_schedule() self.set_random_seed(self.seed) - self.replay_buffer = self.replay_buffer_cls( + self.replay_buffer = self.replay_buffer_class( self.buffer_size, self.observation_space, self.action_space, diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index ed0f655..b030b43 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -70,6 +70,8 @@ class DQN(OffPolicyAlgorithm): gradient_steps: int = 1, n_episodes_rollout: int = -1, optimize_memory_usage: bool = False, + replay_buffer_class: Optional[Type[ReplayBuffer]] = None, + replay_buffer_kwargs: Optional[Dict[str, Any]] = None, target_update_interval: int = 10000, exploration_fraction: float = 0.1, exploration_initial_eps: float = 1.0, @@ -89,8 +91,8 @@ class DQN(OffPolicyAlgorithm): env, DQNPolicy, learning_rate, - ReplayBuffer, - None, + replay_buffer_class, + replay_buffer_kwargs, buffer_size, learning_starts, batch_size, diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 396087b..f631ec6 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -86,6 +86,8 @@ class SAC(OffPolicyAlgorithm): n_episodes_rollout: int = -1, action_noise: Optional[ActionNoise] = None, optimize_memory_usage: bool = False, + replay_buffer_class: Optional[Type[ReplayBuffer]] = None, + replay_buffer_kwargs: Optional[Dict[str, Any]] = None, ent_coef: Union[str, float] = "auto", target_update_interval: int = 1, target_entropy: Union[str, float] = "auto", @@ -106,8 +108,8 @@ class SAC(OffPolicyAlgorithm): env, SACPolicy, learning_rate, - ReplayBuffer, - None, + replay_buffer_class, + replay_buffer_kwargs, buffer_size, learning_starts, batch_size,