Merge branch 'master' into feat/mps-support

.readthedocs.yml

@@ -14,3 +14,8 @@ formats: all
 # Set requirements using conda env
 conda:
   environment: docs/conda_env.yml
+
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "mambaforge-22.9"

docs/conda_env.yml

@@ -14,7 +14,6 @@ dependencies:
     - pandas
     - numpy
     - matplotlib
-    - sphinx_autodoc_typehints
-    - sphinx>=5.3,<7.0
-    - sphinx_rtd_theme>=1.0
+    - sphinx>=5,<8
+    - sphinx_rtd_theme>=1.3.0
     - sphinx_copybutton

docs/conf.py

@@ -64,7 +64,6 @@ release = __version__
 # ones.
 extensions = [
     "sphinx.ext.autodoc",
-    "sphinx_autodoc_typehints",
     "sphinx.ext.autosummary",
     "sphinx.ext.mathjax",
     "sphinx.ext.ifconfig",
@@ -73,6 +72,8 @@ extensions = [
     # 'sphinx.ext.doctest'
 ]
 
+autodoc_typehints = "description"
+
 if enable_spell_check:
     extensions.append("sphinxcontrib.spelling")
 

docs/guide/vec_envs.rst

@@ -89,8 +89,8 @@ SB3 VecEnv API is actually close to Gym 0.21 API but differs to Gym 0.26+ API:
   if no mode is passed or ``mode="rgb_array"`` is passed when calling ``vec_env.render`` then we use the default mode, otherwise, we use the OpenCV display.
   Note that if ``render_mode != "rgb_array"``, you can only call ``vec_env.render()`` (without argument or with ``mode=env.render_mode``).
 
-- the ``reset()`` method doesn't take any parameter. If you want to seed the pseudo-random generator,
-  you should call ``vec_env.seed(seed=seed)`` and ``obs = vec_env.reset()`` afterward.
+- the ``reset()`` method doesn't take any parameter. If you want to seed the pseudo-random generator or pass options,
+  you should call ``vec_env.seed(seed=seed)``/``vec_env.set_options(options)`` and ``obs = vec_env.reset()`` afterward (seed and options are discarded after each call to ``reset()``).
 
 - methods and attributes of the underlying Gym envs can be accessed, called and set using ``vec_env.get_attr("attribute_name")``,
   ``vec_env.env_method("method_name", args1, args2, kwargs1=kwargs1)`` and ``vec_env.set_attr("attribute_name", new_value)``.

docs/misc/changelog.rst

@@ -3,17 +3,22 @@
 Changelog
 ==========
 
-Release 2.2.0a6 (WIP)
+Release 2.2.0a9 (WIP)
 --------------------------
+**Support for options at reset, bug fixes and better error messages**
 
 Breaking Changes:
 ^^^^^^^^^^^^^^^^^
 - Switched to ``ruff`` for sorting imports (isort is no longer needed), black and ruff version now require a minimum version
+- Dropped ``x is False`` in favor of ``not x``, which means that callbacks that wrongly returned None (instead of a boolean) will cause the training to stop (@iwishiwasaneagle)
 
 New Features:
 ^^^^^^^^^^^^^
 - Improved error message of the ``env_checker`` for env wrongly detected as GoalEnv (``compute_reward()`` is defined)
 - Improved error message when mixing Gym API with VecEnv API (see GH#1694)
+- Add support for setting ``options`` at reset with VecEnv via the ``set_options()`` method. Same as seeds logic, options are reset at the end of an episode (@ReHoss)
+- Added ``rollout_buffer_class`` and ``rollout_buffer_kwargs`` arguments to on-policy algorithms (A2C and PPO)
+
 
 Bug Fixes:
 ^^^^^^^^^^
@@ -34,9 +39,9 @@ Bug Fixes:
 `RL Zoo`_
 ^^^^^^^^^
 
-`SBX`_
-^^^^^^^^^
-- Added ``DDPG`` and ``TD3``
+`SBX`_ (SB3 + Jax)
+^^^^^^^^^^^^^^^^^^
+- Added ``DDPG`` and ``TD3`` algorithms
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -52,6 +57,8 @@ Others:
 - Fixed ``stable_baselines3/common/buffers.py`` type hints
 - Fixed ``stable_baselines3/her/her_replay_buffer.py`` type hints
 - Buffers do no call an additional ``.copy()`` when storing new transitions
+- Fixed ``ActorCriticPolicy.extract_features()`` signature by adding an optional ``features_extractor`` argument
+- Update dependencies (accept newer Shimmy/Sphinx version and remove ``sphinx_autodoc_typehints``)
 
 Documentation:
 ^^^^^^^^^^^^^^
@@ -1463,8 +1470,8 @@ And all the contributors:
 @eleurent @ac-93 @cove9988 @theDebugger811 @hsuehch @Demetrio92 @thomasgubler @IperGiove @ScheiklP
 @simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola @buoyancy99 @caburu @xy9485
 @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede
-@carlosluis @arjun-kg @tlpss @JonathanKuelz @Gabo-Tor
+@carlosluis @arjun-kg @tlpss @JonathanKuelz @Gabo-Tor @iwishiwasaneagle
 @Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875
-@anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong
+@anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong @ReHoss
 @DavyMorgan @luizapozzobon @Bonifatius94 @theSquaredError @harveybellini @DavyMorgan @FieteO @jonasreiher @npit @WeberSamuel @troiganto
 @lutogniew @lbergmann1 @lukashass @BertrandDecoster @pseudo-rnd-thoughts @stefanbschneider @kyle-he @PatrickHelm @corentinlger

setup.py

@@ -85,7 +85,7 @@ extra_no_roms = [
     "tqdm",
     "rich",
     # For atari games,
-    "shimmy[atari]~=1.1.0",
+    "shimmy[atari]~=1.3.0",
     "pillow",
 ]
 
@@ -126,13 +126,11 @@ setup(
             "black>=23.9.1,<24",
         ],
         "docs": [
-            "sphinx>=5.3,<7.0",
+            "sphinx>=5,<8",
             "sphinx-autobuild",
-            "sphinx-rtd-theme",
+            "sphinx-rtd-theme>=1.3.0",
             # For spelling
             "sphinxcontrib.spelling",
-            # Type hints support
-            "sphinx-autodoc-typehints",
             # Copy button for code snippets
             "sphinx_copybutton",
         ],

stable_baselines3/a2c/a2c.py

@@ -4,6 +4,7 @@ import torch as th
 from gymnasium import spaces
 from torch.nn import functional as F
 
+from stable_baselines3.common.buffers import RolloutBuffer
 from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
 from stable_baselines3.common.policies import ActorCriticCnnPolicy, ActorCriticPolicy, BasePolicy, MultiInputActorCriticPolicy
 from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule
@@ -41,6 +42,8 @@ class A2C(OnPolicyAlgorithm):
         instead of action noise exploration (default: False)
     :param sde_sample_freq: Sample a new noise matrix every n steps when using gSDE
         Default: -1 (only sample at the beginning of the rollout)
+    :param rollout_buffer_class: Rollout buffer class to use. If ``None``, it will be automatically selected.
+    :param rollout_buffer_kwargs: Keyword arguments to pass to the rollout buffer on creation.
     :param normalize_advantage: Whether to normalize or not the advantage
     :param stats_window_size: Window size for the rollout logging, specifying the number of episodes to average
         the reported success rate, mean episode length, and mean reward over
@@ -75,6 +78,8 @@ class A2C(OnPolicyAlgorithm):
         use_rms_prop: bool = True,
         use_sde: bool = False,
         sde_sample_freq: int = -1,
+        rollout_buffer_class: Optional[Type[RolloutBuffer]] = None,
+        rollout_buffer_kwargs: Optional[Dict[str, Any]] = None,
         normalize_advantage: bool = False,
         stats_window_size: int = 100,
         tensorboard_log: Optional[str] = None,
@@ -96,6 +101,8 @@ class A2C(OnPolicyAlgorithm):
             max_grad_norm=max_grad_norm,
             use_sde=use_sde,
             sde_sample_freq=sde_sample_freq,
+            rollout_buffer_class=rollout_buffer_class,
+            rollout_buffer_kwargs=rollout_buffer_kwargs,
             stats_window_size=stats_window_size,
             tensorboard_log=tensorboard_log,
             policy_kwargs=policy_kwargs,

stable_baselines3/common/buffers.py

@@ -561,7 +561,7 @@ class DictReplayBuffer(ReplayBuffer):
         if psutil is not None:
             mem_available = psutil.virtual_memory().available
 
-        assert optimize_memory_usage is False, "DictReplayBuffer does not support optimize_memory_usage"
+        assert not optimize_memory_usage, "DictReplayBuffer does not support optimize_memory_usage"
         # disabling as this adds quite a bit of complexity
         # https://github.com/DLR-RM/stable-baselines3/pull/243#discussion_r531535702
         self.optimize_memory_usage = optimize_memory_usage

stable_baselines3/common/callbacks.py

@@ -1,7 +1,7 @@
 import os
 import warnings
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 import gymnasium as gym
 import numpy as np
@@ -19,10 +19,13 @@ except ImportError:
     # if the progress bar is used
     tqdm = None
 
-from stable_baselines3.common import base_class
+
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv, sync_envs_normalization
 
+if TYPE_CHECKING:
+    from stable_baselines3.common import base_class
+
 
 class BaseCallback(ABC):
     """
@@ -554,7 +557,6 @@ class StopTrainingOnRewardThreshold(BaseCallback):
 
     def _on_step(self) -> bool:
         assert self.parent is not None, "``StopTrainingOnMinimumReward`` callback must be used with an ``EvalCallback``"
-        # Convert np.bool_ to bool, otherwise callback() is False won't work
         continue_training = bool(self.parent.best_mean_reward < self.reward_threshold)
         if self.verbose >= 1 and not continue_training:
             print(

stable_baselines3/common/off_policy_algorithm.py

@@ -330,7 +330,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
                 log_interval=log_interval,
             )
 
-            if rollout.continue_training is False:
+            if not rollout.continue_training:
                 break
 
             if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
@@ -556,7 +556,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
             # Give access to local variables
             callback.update_locals(locals())
             # Only stop training if return value is False, not when it is None.
-            if callback.on_step() is False:
+            if not callback.on_step():
                 return RolloutReturn(num_collected_steps * env.num_envs, num_collected_episodes, continue_training=False)
 
             # Retrieve reward and episode length if using Monitor wrapper

stable_baselines3/common/on_policy_algorithm.py

@@ -37,6 +37,8 @@ class OnPolicyAlgorithm(BaseAlgorithm):
         instead of action noise exploration (default: False)
     :param sde_sample_freq: Sample a new noise matrix every n steps when using gSDE
         Default: -1 (only sample at the beginning of the rollout)
+    :param rollout_buffer_class: Rollout buffer class to use. If ``None``, it will be automatically selected.
+    :param rollout_buffer_kwargs: Keyword arguments to pass to the rollout buffer on creation.
     :param stats_window_size: Window size for the rollout logging, specifying the number of episodes to average
         the reported success rate, mean episode length, and mean reward over
     :param tensorboard_log: the log location for tensorboard (if None, no logging)
@@ -68,6 +70,8 @@ class OnPolicyAlgorithm(BaseAlgorithm):
         max_grad_norm: float,
         use_sde: bool,
         sde_sample_freq: int,
+        rollout_buffer_class: Optional[Type[RolloutBuffer]] = None,
+        rollout_buffer_kwargs: Optional[Dict[str, Any]] = None,
         stats_window_size: int = 100,
         tensorboard_log: Optional[str] = None,
         monitor_wrapper: bool = True,
@@ -100,6 +104,8 @@ class OnPolicyAlgorithm(BaseAlgorithm):
         self.ent_coef = ent_coef
         self.vf_coef = vf_coef
         self.max_grad_norm = max_grad_norm
+        self.rollout_buffer_class = rollout_buffer_class
+        self.rollout_buffer_kwargs = rollout_buffer_kwargs or {}
 
         if _init_setup_model:
             self._setup_model()
@@ -108,9 +114,13 @@ class OnPolicyAlgorithm(BaseAlgorithm):
         self._setup_lr_schedule()
         self.set_random_seed(self.seed)
 
-        buffer_cls = DictRolloutBuffer if isinstance(self.observation_space, spaces.Dict) else RolloutBuffer
+        if self.rollout_buffer_class is None:
+            if isinstance(self.observation_space, spaces.Dict):
+                self.rollout_buffer_class = DictRolloutBuffer
+            else:
+                self.rollout_buffer_class = RolloutBuffer
 
-        self.rollout_buffer = buffer_cls(
+        self.rollout_buffer = self.rollout_buffer_class(
             self.n_steps,
             self.observation_space,  # type: ignore[arg-type]
             self.action_space,
@@ -118,6 +128,7 @@ class OnPolicyAlgorithm(BaseAlgorithm):
             gamma=self.gamma,
             gae_lambda=self.gae_lambda,
             n_envs=self.n_envs,
+            **self.rollout_buffer_kwargs,
         )
         self.policy = self.policy_class(  # type: ignore[assignment]
             self.observation_space, self.action_space, self.lr_schedule, use_sde=self.use_sde, **self.policy_kwargs
@@ -186,7 +197,7 @@ class OnPolicyAlgorithm(BaseAlgorithm):
 
             # Give access to local variables
             callback.update_locals(locals())
-            if callback.on_step() is False:
+            if not callback.on_step():
                 return False
 
             self._update_info_buffer(infos)
@@ -265,7 +276,7 @@ class OnPolicyAlgorithm(BaseAlgorithm):
         while self.num_timesteps < total_timesteps:
             continue_training = self.collect_rollouts(self.env, callback, self.rollout_buffer, n_rollout_steps=self.n_steps)
 
-            if continue_training is False:
+            if not continue_training:
                 break
 
             iteration += 1

stable_baselines3/common/policies.py

@@ -90,7 +90,7 @@ class BaseModel(nn.Module):
         self.features_extractor_class = features_extractor_class
         self.features_extractor_kwargs = features_extractor_kwargs
         # Automatically deactivate dtype and bounds checks
-        if normalize_images is False and issubclass(features_extractor_class, (NatureCNN, CombinedExtractor)):
+        if not normalize_images and issubclass(features_extractor_class, (NatureCNN, CombinedExtractor)):
             self.features_extractor_kwargs.update(dict(normalized_image=True))
 
     def _update_features_extractor(
@@ -123,9 +123,9 @@ class BaseModel(nn.Module):
         """
         Preprocess the observation if needed and extract features.
 
-         :param obs: The observation
-         :param features_extractor: The features extractor to use.
-         :return: The extracted features
+        :param obs: Observation
+        :param features_extractor: The features extractor to use.
+        :return: The extracted features
         """
         preprocessed_obs = preprocess_obs(obs, self.observation_space, normalize_images=self.normalize_images)
         return features_extractor(preprocessed_obs)
@@ -642,16 +642,26 @@ class ActorCriticPolicy(BasePolicy):
         actions = actions.reshape((-1, *self.action_space.shape))
         return actions, values, log_prob
 
-    def extract_features(self, obs: th.Tensor) -> Union[th.Tensor, Tuple[th.Tensor, th.Tensor]]:
+    def extract_features(
+        self, obs: th.Tensor, features_extractor: Optional[BaseFeaturesExtractor] = None
+    ) -> Union[th.Tensor, Tuple[th.Tensor, th.Tensor]]:
         """
         Preprocess the observation if needed and extract features.
 
         :param obs: Observation
-        :return: the output of the features extractor(s)
+        :param features_extractor: The features extractor to use. If None, then ``self.features_extractor`` is used.
+        :return: The extracted features. If features extractor is not shared, returns a tuple with the
+            features for the actor and the features for the critic.
         """
         if self.share_features_extractor:
-            return super().extract_features(obs, self.features_extractor)
+            return super().extract_features(obs, self.features_extractor if features_extractor is None else features_extractor)
         else:
+            if features_extractor is not None:
+                warnings.warn(
+                    "Provided features_extractor will be ignored because the features extractor is not shared.",
+                    UserWarning,
+                )
+
             pi_features = super().extract_features(obs, self.pi_features_extractor)
             vf_features = super().extract_features(obs, self.vf_features_extractor)
             return pi_features, vf_features

stable_baselines3/common/type_aliases.py

@@ -1,15 +1,17 @@
 """Common aliases for type hints"""
-
 from enum import Enum
-from typing import Any, Callable, Dict, List, NamedTuple, Optional, Protocol, SupportsFloat, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple, Optional, Protocol, SupportsFloat, Tuple, Union
 
 import gymnasium as gym
 import numpy as np
 import torch as th
 
-from stable_baselines3.common import callbacks, vec_env
+# Avoid circular imports, we use type hint as string to avoid it too
+if TYPE_CHECKING:
+    from stable_baselines3.common.callbacks import BaseCallback
+    from stable_baselines3.common.vec_env import VecEnv
 
-GymEnv = Union[gym.Env, vec_env.VecEnv]
+GymEnv = Union[gym.Env, "VecEnv"]
 GymObs = Union[Tuple, Dict[str, Any], np.ndarray, int]
 GymResetReturn = Tuple[GymObs, Dict]
 AtariResetReturn = Tuple[np.ndarray, Dict[str, Any]]
@@ -17,7 +19,7 @@ GymStepReturn = Tuple[GymObs, float, bool, bool, Dict]
 AtariStepReturn = Tuple[np.ndarray, SupportsFloat, bool, bool, Dict[str, Any]]
 TensorDict = Dict[str, th.Tensor]
 OptimizerStateDict = Dict[str, Any]
-MaybeCallback = Union[None, Callable, List[callbacks.BaseCallback], callbacks.BaseCallback]
+MaybeCallback = Union[None, Callable, List["BaseCallback"], "BaseCallback"]
 
 # A schedule takes the remaining progress as input
 # and ouputs a scalar (e.g. learning rate, clip range, ...)

stable_baselines3/common/vec_env/base_vec_env.py

@@ -1,6 +1,7 @@
 import inspect
 import warnings
 from abc import ABC, abstractmethod
+from copy import deepcopy
 from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, Union
 
 import cloudpickle
@@ -67,6 +68,8 @@ class VecEnv(ABC):
         self.reset_infos: List[Dict[str, Any]] = [{} for _ in range(num_envs)]
         # seeds to be used in the next call to env.reset()
         self._seeds: List[Optional[int]] = [None for _ in range(num_envs)]
+        # options to be used in the next call to env.reset()
+        self._options: List[Dict[str, Any]] = [{} for _ in range(num_envs)]
 
         try:
             render_modes = self.get_attr("render_mode")
@@ -95,6 +98,12 @@ class VecEnv(ABC):
         """
         self._seeds = [None for _ in range(self.num_envs)]
 
+    def _reset_options(self) -> None:
+        """
+        Reset the options that are going to be used at the next reset.
+        """
+        self._options = [{} for _ in range(self.num_envs)]
+
     @abstractmethod
     def reset(self) -> VecEnvObs:
         """
@@ -283,6 +292,22 @@ class VecEnv(ABC):
         self._seeds = [seed + idx for idx in range(self.num_envs)]
         return self._seeds
 
+    def set_options(self, options: Optional[Union[List[Dict], Dict]] = None) -> None:
+        """
+        Set environment options for all environments.
+        If a dict is passed instead of a list, the same options will be used for all environments.
+        WARNING: Those options will only be passed to the environment at the next reset.
+
+        :param options: A dictionary of environment options to pass to each environment at the next reset.
+        """
+        if options is None:
+            options = {}
+        # Use deepcopy to avoid side effects
+        if isinstance(options, dict):
+            self._options = deepcopy([options] * self.num_envs)
+        else:
+            self._options = deepcopy(options)
+
     @property
     def unwrapped(self) -> "VecEnv":
         if isinstance(self, VecEnvWrapper):
@@ -354,6 +379,9 @@ class VecEnvWrapper(VecEnv):
     def seed(self, seed: Optional[int] = None) -> Sequence[Union[None, int]]:
         return self.venv.seed(seed)
 
+    def set_options(self, options: Optional[Union[List[Dict], Dict]] = None) -> None:
+        return self.venv.set_options(options)
+
     def close(self) -> None:
         return self.venv.close()
 

stable_baselines3/common/vec_env/dummy_vec_env.py

@@ -73,10 +73,12 @@ class DummyVecEnv(VecEnv):
 
     def reset(self) -> VecEnvObs:
         for env_idx in range(self.num_envs):
-            obs, self.reset_infos[env_idx] = self.envs[env_idx].reset(seed=self._seeds[env_idx])
+            maybe_options = {"options": self._options[env_idx]} if self._options[env_idx] else {}
+            obs, self.reset_infos[env_idx] = self.envs[env_idx].reset(seed=self._seeds[env_idx], **maybe_options)
             self._save_obs(env_idx, obs)
-        # Seeds are only used once
+        # Seeds and options are only used once
         self._reset_seeds()
+        self._reset_options()
         return self._obs_from_buf()
 
     def close(self) -> None:

stable_baselines3/common/vec_env/subproc_vec_env.py

@@ -42,7 +42,8 @@ def _worker(
                     observation, reset_info = env.reset()
                 remote.send((observation, reward, done, info, reset_info))
             elif cmd == "reset":
-                observation, reset_info = env.reset(seed=data)
+                maybe_options = {"options": data[1]} if data[1] else {}
+                observation, reset_info = env.reset(seed=data[0], **maybe_options)
                 remote.send((observation, reset_info))
             elif cmd == "render":
                 remote.send(env.render())
@@ -132,11 +133,12 @@ class SubprocVecEnv(VecEnv):
 
     def reset(self) -> VecEnvObs:
         for env_idx, remote in enumerate(self.remotes):
-            remote.send(("reset", self._seeds[env_idx]))
+            remote.send(("reset", (self._seeds[env_idx], self._options[env_idx])))
         results = [remote.recv() for remote in self.remotes]
         obs, self.reset_infos = zip(*results)
-        # Seeds are only used once
+        # Seeds and options are only used once
         self._reset_seeds()
+        self._reset_options()
         return _flatten_obs(obs, self.observation_space)
 
     def close(self) -> None:

stable_baselines3/ppo/ppo.py

@@ -6,6 +6,7 @@ import torch as th
 from gymnasium import spaces
 from torch.nn import functional as F
 
+from stable_baselines3.common.buffers import RolloutBuffer
 from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
 from stable_baselines3.common.policies import ActorCriticCnnPolicy, ActorCriticPolicy, BasePolicy, MultiInputActorCriticPolicy
 from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule
@@ -52,6 +53,8 @@ class PPO(OnPolicyAlgorithm):
         instead of action noise exploration (default: False)
     :param sde_sample_freq: Sample a new noise matrix every n steps when using gSDE
         Default: -1 (only sample at the beginning of the rollout)
+    :param rollout_buffer_class: Rollout buffer class to use. If ``None``, it will be automatically selected.
+    :param rollout_buffer_kwargs: Keyword arguments to pass to the rollout buffer on creation
     :param target_kl: Limit the KL divergence between updates,
         because the clipping is not enough to prevent large update
         see issue #213 (cf https://github.com/hill-a/stable-baselines/issues/213)
@@ -92,6 +95,8 @@ class PPO(OnPolicyAlgorithm):
         max_grad_norm: float = 0.5,
         use_sde: bool = False,
         sde_sample_freq: int = -1,
+        rollout_buffer_class: Optional[Type[RolloutBuffer]] = None,
+        rollout_buffer_kwargs: Optional[Dict[str, Any]] = None,
         target_kl: Optional[float] = None,
         stats_window_size: int = 100,
         tensorboard_log: Optional[str] = None,
@@ -113,6 +118,8 @@ class PPO(OnPolicyAlgorithm):
             max_grad_norm=max_grad_norm,
             use_sde=use_sde,
             sde_sample_freq=sde_sample_freq,
+            rollout_buffer_class=rollout_buffer_class,
+            rollout_buffer_kwargs=rollout_buffer_kwargs,
             stats_window_size=stats_window_size,
             tensorboard_log=tensorboard_log,
             policy_kwargs=policy_kwargs,

stable_baselines3/version.txt

@@ -1 +1 @@
-2.2.0a6
+2.2.0a9

tests/test_buffers.py

@@ -4,6 +4,7 @@ import pytest
 import torch as th
 from gymnasium import spaces
 
+from stable_baselines3 import A2C
 from stable_baselines3.common.buffers import DictReplayBuffer, DictRolloutBuffer, ReplayBuffer, RolloutBuffer
 from stable_baselines3.common.env_checker import check_env
 from stable_baselines3.common.env_util import make_vec_env
@@ -150,3 +151,16 @@ def test_device_buffer(replay_buffer_cls, device):
                 assert value[key].device.type == desired_device
         elif isinstance(value, th.Tensor):
             assert value.device.type == desired_device
+
+
+def test_custom_rollout_buffer():
+    A2C("MlpPolicy", "Pendulum-v1", rollout_buffer_class=RolloutBuffer, rollout_buffer_kwargs=dict())
+
+    with pytest.raises(TypeError, match="unexpected keyword argument 'wrong_keyword'"):
+        A2C("MlpPolicy", "Pendulum-v1", rollout_buffer_class=RolloutBuffer, rollout_buffer_kwargs=dict(wrong_keyword=1))
+
+    with pytest.raises(TypeError, match="got multiple values for keyword argument 'gamma'"):
+        A2C("MlpPolicy", "Pendulum-v1", rollout_buffer_class=RolloutBuffer, rollout_buffer_kwargs=dict(gamma=1))
+
+    with pytest.raises(AssertionError, match="DictRolloutBuffer must be used with Dict obs space only"):
+        A2C("MlpPolicy", "Pendulum-v1", rollout_buffer_class=DictRolloutBuffer)

tests/test_callbacks.py

@@ -4,9 +4,11 @@ import shutil
 import gymnasium as gym
 import numpy as np
 import pytest
+import torch as th
 
 from stable_baselines3 import A2C, DDPG, DQN, PPO, SAC, TD3, HerReplayBuffer
 from stable_baselines3.common.callbacks import (
+    BaseCallback,
     CallbackList,
     CheckpointCallback,
     EvalCallback,
@@ -123,6 +125,40 @@ def test_eval_callback_vec_env():
     assert eval_callback.last_mean_reward == 100.0
 
 
+class AlwaysFailCallback(BaseCallback):
+    def __init__(self, *args, callback_false_value, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.callback_false_value = callback_false_value
+
+    def _on_step(self) -> bool:
+        return self.callback_false_value
+
+
+@pytest.mark.parametrize(
+    "model_class,model_kwargs",
+    [
+        (A2C, dict(n_steps=1, stats_window_size=1)),
+        (
+            SAC,
+            dict(
+                learning_starts=1,
+                buffer_size=1,
+                batch_size=1,
+            ),
+        ),
+    ],
+)
+@pytest.mark.parametrize("callback_false_value", [False, np.bool_(0), th.tensor(0, dtype=th.bool)])
+def test_callbacks_can_cancel_runs(model_class, model_kwargs, callback_false_value):
+    assert not callback_false_value  # Sanity check to ensure parametrized values are valid
+    env_id = select_env(model_class)
+    model = model_class("MlpPolicy", env_id, **model_kwargs, policy_kwargs=dict(net_arch=[2]))
+    alwaysfailcallback = AlwaysFailCallback(callback_false_value=callback_false_value)
+    model.learn(10, callback=alwaysfailcallback)
+
+    assert alwaysfailcallback.n_calls == 1
+
+
 def test_eval_success_logging(tmp_path):
     n_bits = 2
     n_envs = 2

tests/test_vec_envs.py

@@ -30,11 +30,13 @@ class CustomGymEnv(gym.Env):
         self.current_step = 0
         self.ep_length = 4
         self.render_mode = render_mode
+        self.current_options: Optional[Dict] = None
 
     def reset(self, *, seed: Optional[int] = None, options: Optional[Dict] = None):
         if seed is not None:
             self.seed(seed)
         self.current_step = 0
+        self.current_options = options
         self._choose_next_state()
         return self.state, {}
 
@@ -160,6 +162,25 @@ def test_vecenv_custom_calls(vec_env_class, vec_env_wrapper):
     assert getattr_result == [12] + [0 for _ in range(N_ENVS - 2)] + [12]
     assert vec_env.get_attr("current_step", indices=[-1]) == [12]
 
+    # Checks that options are correctly passed
+    assert vec_env.get_attr("current_options")[0] is None
+    # Same options for all envs
+    options = {"hello": 1}
+    vec_env.set_options(options)
+    assert vec_env.get_attr("current_options")[0] is None
+    # Only effective at reset
+    vec_env.reset()
+    assert vec_env.get_attr("current_options") == [options] * N_ENVS
+    vec_env.reset()
+    # Options are reset
+    assert vec_env.get_attr("current_options")[0] is None
+    # Use a list of options, different for the first env
+    options = [{"hello": 1}] * N_ENVS
+    options[0] = {"other_option": 2}
+    vec_env.set_options(options)
+    vec_env.reset()
+    assert vec_env.get_attr("current_options") == options
+
     vec_env.close()
 
 
@@ -487,7 +508,14 @@ def test_vec_deterministic(vec_env_class):
     vec_env.seed(3)
     new_obs = vec_env.reset()
     assert np.allclose(new_obs, obs)
-    vec_env.close()
+    # Test with VecNormalize (VecEnvWrapper should call self.venv.seed())
+    vec_normalize = VecNormalize(vec_env)
+    vec_normalize.seed(3)
+    obs = vec_env.reset()
+    vec_normalize.seed(3)
+    new_obs = vec_env.reset()
+    assert np.allclose(new_obs, obs)
+    vec_normalize.close()
     # Similar test but with make_vec_env
     vec_env_1 = make_vec_env("Pendulum-v1", n_envs=N_ENVS, vec_env_cls=vec_env_class, seed=0)
     vec_env_2 = make_vec_env("Pendulum-v1", n_envs=N_ENVS, vec_env_cls=vec_env_class, seed=0)