diff --git a/.readthedocs.yml b/.readthedocs.yml index 6753d87..dbb2fad 100644 --- a/.readthedocs.yml +++ b/.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" diff --git a/docs/conda_env.yml b/docs/conda_env.yml index 8c71fe9..53fecf2 100644 --- a/docs/conda_env.yml +++ b/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 diff --git a/docs/conf.py b/docs/conf.py index 8aeae8b..bd63657 100644 --- a/docs/conf.py +++ b/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") diff --git a/docs/guide/vec_envs.rst b/docs/guide/vec_envs.rst index ea99444..10bba85 100644 --- a/docs/guide/vec_envs.rst +++ b/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)``. diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 1ef7155..cd6bc01 100644 --- a/docs/misc/changelog.rst +++ b/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 diff --git a/setup.py b/setup.py index cdaf263..0a1ab30 100644 --- a/setup.py +++ b/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", ], diff --git a/stable_baselines3/a2c/a2c.py b/stable_baselines3/a2c/a2c.py index fda20c9..718571f 100644 --- a/stable_baselines3/a2c/a2c.py +++ b/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, diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index a230a31..306b435 100644 --- a/stable_baselines3/common/buffers.py +++ b/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 diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index 54f1b97..2898df8 100644 --- a/stable_baselines3/common/callbacks.py +++ b/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( diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 2caaf8e..e8dcac4 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/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 diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 1e0f9e6..ddd0f8d 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/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 diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index e975349..5f57672 100644 --- a/stable_baselines3/common/policies.py +++ b/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 diff --git a/stable_baselines3/common/type_aliases.py b/stable_baselines3/common/type_aliases.py index 2f98ee1..4a0a878 100644 --- a/stable_baselines3/common/type_aliases.py +++ b/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, ...) diff --git a/stable_baselines3/common/vec_env/base_vec_env.py b/stable_baselines3/common/vec_env/base_vec_env.py index 16518a1..5e06a5c 100644 --- a/stable_baselines3/common/vec_env/base_vec_env.py +++ b/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() diff --git a/stable_baselines3/common/vec_env/dummy_vec_env.py b/stable_baselines3/common/vec_env/dummy_vec_env.py index 2908981..15ecfb6 100644 --- a/stable_baselines3/common/vec_env/dummy_vec_env.py +++ b/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: diff --git a/stable_baselines3/common/vec_env/subproc_vec_env.py b/stable_baselines3/common/vec_env/subproc_vec_env.py index dbc7002..8375884 100644 --- a/stable_baselines3/common/vec_env/subproc_vec_env.py +++ b/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: diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index 9bd83b0..ea7cf5e 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/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, diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 47f323d..b7120ad 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -2.2.0a6 +2.2.0a9 diff --git a/tests/test_buffers.py b/tests/test_buffers.py index e7d4a1c..2ea366a 100644 --- a/tests/test_buffers.py +++ b/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) diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index f8b0e54..d159c43 100644 --- a/tests/test_callbacks.py +++ b/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 diff --git a/tests/test_vec_envs.py b/tests/test_vec_envs.py index 61740c4..a9516ae 100644 --- a/tests/test_vec_envs.py +++ b/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)