diff --git a/Makefile b/Makefile index c806507..6351162 100644 --- a/Makefile +++ b/Makefile @@ -19,6 +19,13 @@ lint: # exit-zero treats all errors as warnings. flake8 ${LINT_PATHS} --count --exit-zero --statistics +ruff: + # stop the build if there are Python syntax errors or undefined names + # see https://lintlyci.github.io/Flake8Rules/ + ruff ${LINT_PATHS} --select=E9,F63,F7,F82 --show-source + # exit-zero treats all errors as warnings. + ruff ${LINT_PATHS} --exit-zero --line-length 127 + format: # Sort imports isort ${LINT_PATHS} diff --git a/docs/guide/custom_policy.rst b/docs/guide/custom_policy.rst index c9e598e..dae6048 100644 --- a/docs/guide/custom_policy.rst +++ b/docs/guide/custom_policy.rst @@ -117,11 +117,6 @@ that derives from ``BaseFeaturesExtractor`` and then pass it to the model when t ``policy_kwargs`` (both for on-policy and off-policy algorithms). -.. warning:: - If the features extractor is **non-shared**, it is **not** possible to have shared layers in the ``mlp_extractor``. - Please note that this option is **deprecated**, therefore in a future release the layers in the ``mlp_extractor`` will have to be non-shared. - - .. code-block:: python import torch as th @@ -242,41 +237,31 @@ On-Policy Algorithms Custom Networks --------------- -.. warning:: - Shared layers in the the ``mlp_extractor`` are **deprecated**. - In a future release all layers will have to be non-shared. - If needed, you can implement a custom policy network (see `advanced example below <#advanced-example>`_). - -.. warning:: - In the next Stable-Baselines3 release, the behavior of ``net_arch=[128, 128]`` will change - to match the one of off-policy algorithms: it will create **separate** networks (instead of shared currently) - for the actor and the critic, with the same architecture. - - If you need a network architecture that is different for the actor and the critic when using ``PPO``, ``A2C`` or ``TRPO``, you can pass a dictionary of the following structure: ``dict(pi=[], vf=[])``. For example, if you want a different architecture for the actor (aka ``pi``) and the critic ( value-function aka ``vf``) networks, then you can specify ``net_arch=dict(pi=[32, 32], vf=[64, 64])``. -.. Otherwise, to have actor and critic that share the same network architecture, -.. you only need to specify ``net_arch=[128, 128]`` (here, two hidden layers of 128 units each). +Otherwise, to have actor and critic that share the same network architecture, +you only need to specify ``net_arch=[128, 128]`` (here, two hidden layers of 128 units each, this is equivalent to ``net_arch=dict(pi=[128, 128], vf=[128, 128])``). + +If shared layers are needed, you need to implement a custom policy network (see `advanced example below <#advanced-example>`_). Examples ~~~~~~~~ -.. TODO(antonin): uncomment when shared network is removed -.. Same architecture for actor and critic with two layers of size 128: ``net_arch=[128, 128]`` -.. -.. .. code-block:: none -.. -.. obs -.. / \ -.. <128> <128> -.. | | -.. <128> <128> -.. | | -.. action value +Same architecture for actor and critic with two layers of size 128: ``net_arch=[128, 128]`` + +.. code-block:: none + + obs + / \ + <128> <128> + | | + <128> <128> + | | + action value Different architectures for actor and critic: ``net_arch=dict(pi=[32, 32], vf=[64, 64])`` diff --git a/docs/guide/rl_zoo.rst b/docs/guide/rl_zoo.rst index ea15832..8a611d8 100644 --- a/docs/guide/rl_zoo.rst +++ b/docs/guide/rl_zoo.rst @@ -20,6 +20,10 @@ Goals of this repository: Installation ------------ +Option 1: install the python package ``pip install rl_zoo3`` + +or: + 1. Clone the repository: :: @@ -42,7 +46,10 @@ Installation :: apt-get install swig cmake ffmpeg + # full dependencies pip install -r requirements.txt + # minimal dependencies + pip install -e . Train an Agent @@ -56,13 +63,13 @@ using: :: - python train.py --algo algo_name --env env_id + python -m rl_zoo3.train --algo algo_name --env env_id For example (with evaluation and checkpoints): :: - python train.py --algo ppo --env CartPole-v1 --eval-freq 10000 --save-freq 50000 + python -m rl_zoo3.train --algo ppo --env CartPole-v1 --eval-freq 10000 --save-freq 50000 Continue training (here, load pretrained agent for Breakout and continue @@ -70,7 +77,7 @@ training for 5000 steps): :: - python train.py --algo a2c --env BreakoutNoFrameskip-v4 -i trained_agents/a2c/BreakoutNoFrameskip-v4_1/BreakoutNoFrameskip-v4.zip -n 5000 + python -m rl_zoo3.train --algo a2c --env BreakoutNoFrameskip-v4 -i trained_agents/a2c/BreakoutNoFrameskip-v4_1/BreakoutNoFrameskip-v4.zip -n 5000 Enjoy a Trained Agent @@ -80,13 +87,13 @@ If the trained agent exists, then you can see it in action using: :: - python enjoy.py --algo algo_name --env env_id + python -m rl_zoo3.enjoy --algo algo_name --env env_id For example, enjoy A2C on Breakout during 5000 timesteps: :: - python enjoy.py --algo a2c --env BreakoutNoFrameskip-v4 --folder rl-trained-agents/ -n 5000 + python -m rl_zoo3.enjoy --algo a2c --env BreakoutNoFrameskip-v4 --folder rl-trained-agents/ -n 5000 Hyperparameter Optimization @@ -100,7 +107,7 @@ with a budget of 1000 trials and a maximum of 50000 steps: :: - python train.py --algo ppo --env MountainCar-v0 -n 50000 -optimize --n-trials 1000 --n-jobs 2 \ + python -m rl_zoo3.train --algo ppo --env MountainCar-v0 -n 50000 -optimize --n-trials 1000 --n-jobs 2 \ --sampler random --pruner median diff --git a/docs/guide/tensorboard.rst b/docs/guide/tensorboard.rst index 610afd6..2699d4a 100644 --- a/docs/guide/tensorboard.rst +++ b/docs/guide/tensorboard.rst @@ -268,11 +268,9 @@ Here is an example of how to save hyperparameters in TensorBoard: class HParamCallback(BaseCallback): - def __init__(self): - """ - Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard. - """ - super().__init__() + """ + Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard. + """ def _on_training_start(self) -> None: hparam_dict = { @@ -284,7 +282,7 @@ Here is an example of how to save hyperparameters in TensorBoard: # Tensorbaord will find & display metrics from the `SCALARS` tab metric_dict = { "rollout/ep_len_mean": 0, - "train/value_loss": 0, + "train/value_loss": 0.0, } self.logger.record( "hparams", diff --git a/docs/guide/vec_envs.rst b/docs/guide/vec_envs.rst index b074dad..d847811 100644 --- a/docs/guide/vec_envs.rst +++ b/docs/guide/vec_envs.rst @@ -122,12 +122,6 @@ StackedObservations .. autoclass:: stable_baselines3.common.vec_env.stacked_observations.StackedObservations :members: -StackedDictObservations -~~~~~~~~~~~~~~~~~~~~~~~ - -.. autoclass:: stable_baselines3.common.vec_env.stacked_observations.StackedDictObservations - :members: - VecNormalize ~~~~~~~~~~~~ diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index e12442d..767fcba 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,15 +4,19 @@ Changelog ========== -Release 1.8.0a0 (WIP) +Release 1.8.0a4 (WIP) -------------------------- Breaking Changes: ^^^^^^^^^^^^^^^^^ +- Removed shared layers in ``mlp_extractor`` (@AlexPasqua) +- Refactored ``StackedObservations`` (it now handles dict obs, ``StackedDictObservations`` was removed) New Features: ^^^^^^^^^^^^^ +- Added ``repeat_action_probability`` argument in ``AtariWrapper``. +- Only use ``NoopResetEnv`` and ``MaxAndSkipEnv`` when needed in ``AtariWrapper`` `SB3-Contrib`_ ^^^^^^^^^^^^^^ @@ -22,17 +26,26 @@ New Features: Bug Fixes: ^^^^^^^^^^ +- Fixed Atari wrapper that missed the reset condition (@luizapozzobon) +- Added the argument ``dtype`` (default to ``float32``) to the noise for consistency with gym action (@sidney-tio) +- Fixed PPO train/n_updates metric not accounting for early stopping (@adamfrly) Deprecations: ^^^^^^^^^^^^^ Others: ^^^^^^^ +- Fixed ``tests/test_tensorboard.py`` type hint +- Fixed ``tests/test_vec_normalize.py`` type hint +- Fixed ``stable_baselines3/common/monitor.py`` type hint +- Added tests for StackedObservations Documentation: ^^^^^^^^^^^^^^ - Renamed ``load_parameters`` to ``set_parameters`` (@DavyMorgan) - +- Clarified documentation about subproc multiprocessing for A2C (@Bonifatius94) +- Fixed typo in ``A2C`` docstring (@AlexPasqua) +- Renamed timesteps to episodes for ``log_interval`` description (@theSquaredError) Release 1.7.0 (2023-01-10) -------------------------- @@ -1216,4 +1229,4 @@ And all the contributors: @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede @Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875 @yuanmingqi @anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong -@DavyMorgan +@DavyMorgan @luizapozzobon @Bonifatius94 @theSquaredError diff --git a/docs/modules/a2c.rst b/docs/modules/a2c.rst index e871424..670da61 100644 --- a/docs/modules/a2c.rst +++ b/docs/modules/a2c.rst @@ -76,6 +76,24 @@ Train a A2C agent on ``CartPole-v1`` using 4 environments. env.render() +.. note:: + + A2C is meant to be run primarily on the CPU, especially when you are not using a CNN. To improve CPU utilization, try turning off the GPU and using ``SubprocVecEnv`` instead of the default ``DummyVecEnv``: + + .. code-block:: + + from stable_baselines3 import A2C + from stable_baselines3.common.env_util import make_vec_env + from stable_baselines3.common.vec_env import SubprocVecEnv + + if __name__=="__main__": + env = make_vec_env("CartPole-v1", n_envs=8, vec_env_cls=SubprocVecEnv) + model = A2C("MlpPolicy", env, device="cpu") + model.learn(total_timesteps=25_000) + + For more information, see :ref:`Vectorized Environments `, `Issue #1245 `_ or the `Multiprocessing notebook `_. + + Results ------- diff --git a/setup.cfg b/setup.cfg index ca40fa6..3698ded 100644 --- a/setup.cfg +++ b/setup.cfg @@ -39,7 +39,6 @@ exclude = (?x)( | stable_baselines3/common/envs/identity_env.py$ | stable_baselines3/common/envs/multi_input_envs.py$ | stable_baselines3/common/logger.py$ - | stable_baselines3/common/monitor.py$ | stable_baselines3/common/off_policy_algorithm.py$ | stable_baselines3/common/on_policy_algorithm.py$ | stable_baselines3/common/policies.py$ @@ -49,7 +48,6 @@ exclude = (?x)( | stable_baselines3/common/vec_env/__init__.py$ | stable_baselines3/common/vec_env/base_vec_env.py$ | stable_baselines3/common/vec_env/dummy_vec_env.py$ - | stable_baselines3/common/vec_env/stacked_observations.py$ | stable_baselines3/common/vec_env/subproc_vec_env.py$ | stable_baselines3/common/vec_env/util.py$ | stable_baselines3/common/vec_env/vec_extract_dict_obs.py$ @@ -67,9 +65,7 @@ exclude = (?x)( | stable_baselines3/td3/policies.py$ | stable_baselines3/td3/td3.py$ | tests/test_logger.py$ - | tests/test_tensorboard.py$ | tests/test_train_eval_mode.py$ - | tests/test_vec_normalize.py$ ) [flake8] diff --git a/setup.py b/setup.py index df97633..612b2be 100644 --- a/setup.py +++ b/setup.py @@ -117,7 +117,7 @@ setup( # For spelling "sphinxcontrib.spelling", # Type hints support - "sphinx-autodoc-typehints", + "sphinx-autodoc-typehints==1.21.1", # TODO: remove version constraint, see #1290 # Copy button for code snippets "sphinx_copybutton", ], diff --git a/stable_baselines3/a2c/a2c.py b/stable_baselines3/a2c/a2c.py index 972e700..9e8b40c 100644 --- a/stable_baselines3/a2c/a2c.py +++ b/stable_baselines3/a2c/a2c.py @@ -29,7 +29,7 @@ class A2C(OnPolicyAlgorithm): :param n_steps: The number of steps to run for each environment per update (i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel) :param gamma: Discount factor - :param gae_lambda: Factor for trade-off of bias vs variance for Generalized Advantage Estimator + :param gae_lambda: Factor for trade-off of bias vs variance for Generalized Advantage Estimator. Equivalent to classic advantage when set to 1. :param ent_coef: Entropy coefficient for the loss calculation :param vf_coef: Value function coefficient for the loss calculation @@ -81,7 +81,6 @@ class A2C(OnPolicyAlgorithm): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy, env, @@ -132,7 +131,6 @@ class A2C(OnPolicyAlgorithm): # This will only loop once (get all data in one go) for rollout_data in self.rollout_buffer.get(batch_size=None): - actions = rollout_data.actions if isinstance(self.action_space, spaces.Discrete): # Convert discrete action from float to long @@ -189,7 +187,6 @@ class A2C(OnPolicyAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfA2C: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/common/atari_wrappers.py b/stable_baselines3/common/atari_wrappers.py index 785d911..1e06a7f 100644 --- a/stable_baselines3/common/atari_wrappers.py +++ b/stable_baselines3/common/atari_wrappers.py @@ -12,13 +12,39 @@ except ImportError: from stable_baselines3.common.type_aliases import GymObs, GymStepReturn +class StickyActionEnv(gym.Wrapper): + """ + Sticky action. + + Paper: https://arxiv.org/abs/1709.06009 + Official implementation: https://github.com/mgbellemare/Arcade-Learning-Environment + + :param env: Environment to wrap + :param action_repeat_probability: Probability of repeating the last action + """ + + def __init__(self, env: gym.Env, action_repeat_probability: float) -> None: + super().__init__(env) + self.action_repeat_probability = action_repeat_probability + assert env.unwrapped.get_action_meanings()[0] == "NOOP" + + def reset(self, **kwargs) -> GymObs: + self._sticky_action = 0 # NOOP + return self.env.reset(**kwargs) + + def step(self, action: int) -> GymStepReturn: + if self.np_random.random() >= self.action_repeat_probability: + self._sticky_action = action + return self.env.step(self._sticky_action) + + class NoopResetEnv(gym.Wrapper): """ Sample initial states by taking random number of no-ops on reset. No-op is assumed to be action 0. - :param env: the environment to wrap - :param noop_max: the maximum value of no-ops to run + :param env: Environment to wrap + :param noop_max: Maximum value of no-ops to run """ def __init__(self, env: gym.Env, noop_max: int = 30) -> None: @@ -47,7 +73,7 @@ class FireResetEnv(gym.Wrapper): """ Take action on reset for environments that are fixed until firing. - :param env: the environment to wrap + :param env: Environment to wrap """ def __init__(self, env: gym.Env) -> None: @@ -71,7 +97,7 @@ class EpisodicLifeEnv(gym.Wrapper): Make end-of-life == end-of-episode, but only reset on true game over. Done by DeepMind for the DQN and co. since it helps value estimation. - :param env: the environment to wrap + :param env: Environment to wrap """ def __init__(self, env: gym.Env) -> None: @@ -106,7 +132,13 @@ class EpisodicLifeEnv(gym.Wrapper): obs = self.env.reset(**kwargs) else: # no-op step to advance from terminal/lost life state - obs, _, _, _ = self.env.step(0) + obs, _, done, _ = self.env.step(0) + + # The no-op step can lead to a game over, so we need to check it again + # to see if we should reset the environment and avoid the + # monitor.py `RuntimeError: Tried to step environment that needs reset` + if done: + obs = self.env.reset(**kwargs) self.lives = self.env.unwrapped.ale.lives() return obs @@ -114,9 +146,11 @@ class EpisodicLifeEnv(gym.Wrapper): class MaxAndSkipEnv(gym.Wrapper): """ Return only every ``skip``-th frame (frameskipping) + and return the max between the two last frames. - :param env: the environment - :param skip: number of ``skip``-th frame + :param env: Environment to wrap + :param skip: Number of ``skip``-th frame + The same action will be taken ``skip`` times. """ def __init__(self, env: gym.Env, skip: int = 4) -> None: @@ -150,15 +184,12 @@ class MaxAndSkipEnv(gym.Wrapper): return max_frame, total_reward, done, info - def reset(self, **kwargs) -> GymObs: - return self.env.reset(**kwargs) - class ClipRewardEnv(gym.RewardWrapper): """ - Clips the reward to {+1, 0, -1} by its sign. + Clip the reward to {+1, 0, -1} by its sign. - :param env: the environment + :param env: Environment to wrap """ def __init__(self, env: gym.Env) -> None: @@ -179,9 +210,9 @@ class WarpFrame(gym.ObservationWrapper): Convert to grayscale and warp frames to 84x84 (default) as done in the Nature paper and later work. - :param env: the environment - :param width: - :param height: + :param env: Environment to wrap + :param width: New frame width + :param height: New frame height """ def __init__(self, env: gym.Env, width: int = 84, height: int = 84) -> None: @@ -210,20 +241,29 @@ class AtariWrapper(gym.Wrapper): Specifically: - * NoopReset: obtain initial state by taking random number of no-ops on reset. + * Noop reset: obtain initial state by taking random number of no-ops on reset. * Frame skipping: 4 by default * Max-pooling: most recent two observations * Termination signal when a life is lost. * Resize to a square image: 84x84 by default * Grayscale observation * Clip reward to {-1, 0, 1} + * Sticky actions: disabled by default - :param env: gym environment - :param noop_max: max number of no-ops - :param frame_skip: the frequency at which the agent experiences the game. - :param screen_size: resize Atari frame - :param terminal_on_life_loss: if True, then step() returns done=True whenever a life is lost. + See https://danieltakeshi.github.io/2016/11/25/frame-skipping-and-preprocessing-for-deep-q-networks-on-atari-2600-games/ + for a visual explanation. + + .. warning:: + Use this wrapper only with Atari v4 without frame skip: ``env_id = "*NoFrameskip-v4"``. + + :param env: Environment to wrap + :param noop_max: Max number of no-ops + :param frame_skip: Frequency at which the agent experiences the game. + This correspond to repeating the action ``frame_skip`` times. + :param screen_size: Resize Atari frame + :param terminal_on_life_loss: If True, then step() returns done=True whenever a life is lost. :param clip_reward: If True (default), the reward is clip to {-1, 0, 1} depending on its sign. + :param action_repeat_probability: Probability of repeating the last action """ def __init__( @@ -234,9 +274,15 @@ class AtariWrapper(gym.Wrapper): screen_size: int = 84, terminal_on_life_loss: bool = True, clip_reward: bool = True, + action_repeat_probability: float = 0.0, ) -> None: - env = NoopResetEnv(env, noop_max=noop_max) - env = MaxAndSkipEnv(env, skip=frame_skip) + if action_repeat_probability > 0.0: + env = StickyActionEnv(env, action_repeat_probability) + if noop_max > 0: + env = NoopResetEnv(env, noop_max=noop_max) + # frame_skip=1 is the same as no frame-skip (action repeat) + if frame_skip > 1: + env = MaxAndSkipEnv(env, skip=frame_skip) if terminal_on_life_loss: env = EpisodicLifeEnv(env) if "FIRE" in env.unwrapped.get_action_meanings(): diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index a71043d..17be67a 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -505,7 +505,7 @@ class BaseAlgorithm(ABC): :param total_timesteps: The total number of samples (env steps) to train on :param callback: callback(s) called at every step with state of the algorithm. - :param log_interval: The number of timesteps before logging. + :param log_interval: The number of episodes before logging. :param tb_log_name: the name of the run for TensorBoard logging :param reset_num_timesteps: whether or not to reset the current timestep number (used in logging) :param progress_bar: Display a progress bar using tqdm and rich. @@ -667,6 +667,11 @@ class BaseAlgorithm(ABC): if "policy_kwargs" in data: if "device" in data["policy_kwargs"]: del data["policy_kwargs"]["device"] + # backward compatibility, convert to new format + if "net_arch" in data["policy_kwargs"] and len(data["policy_kwargs"]["net_arch"]) > 0: + saved_net_arch = data["policy_kwargs"]["net_arch"] + if isinstance(saved_net_arch, list) and isinstance(saved_net_arch[0], dict): + data["policy_kwargs"]["net_arch"] = saved_net_arch[0] if "policy_kwargs" in kwargs and kwargs["policy_kwargs"] != data["policy_kwargs"]: raise ValueError( @@ -726,7 +731,6 @@ class BaseAlgorithm(ABC): ) else: raise e - # put other pytorch variables back in place if pytorch_variables is not None: for name in pytorch_variables: diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 2dafd41..273dba9 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -240,7 +240,6 @@ class ReplayBuffer(BaseBuffer): done: np.ndarray, infos: List[Dict[str, Any]], ) -> None: - # Reshape needed when using multiple envs with discrete observations # as numpy cannot broadcast (n_discrete,) to (n_discrete, 1) if isinstance(self.observation_space, spaces.Discrete): @@ -346,7 +345,6 @@ class RolloutBuffer(BaseBuffer): gamma: float = 0.99, n_envs: int = 1, ): - super().__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs) self.gae_lambda = gae_lambda self.gamma = gamma @@ -356,7 +354,6 @@ class RolloutBuffer(BaseBuffer): self.reset() def reset(self) -> None: - self.observations = np.zeros((self.buffer_size, self.n_envs) + self.obs_shape, dtype=np.float32) self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=np.float32) self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) @@ -451,7 +448,6 @@ class RolloutBuffer(BaseBuffer): indices = np.random.permutation(self.buffer_size * self.n_envs) # Prepare the data if not self.generator_ready: - _tensor_names = [ "observations", "actions", @@ -474,7 +470,11 @@ class RolloutBuffer(BaseBuffer): yield self._get_samples(indices[start_idx : start_idx + batch_size]) start_idx += batch_size - def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> RolloutBufferSamples: + def _get_samples( + self, + batch_inds: np.ndarray, + env: Optional[VecNormalize] = None, + ) -> RolloutBufferSamples: # type: ignore[signature-mismatch] #FIXME data = ( self.observations[batch_inds], self.actions[batch_inds], @@ -603,7 +603,11 @@ class DictReplayBuffer(ReplayBuffer): self.full = True self.pos = 0 - def sample(self, batch_size: int, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: + def sample( + self, + batch_size: int, + env: Optional[VecNormalize] = None, + ) -> DictReplayBufferSamples: # type: ignore[signature-mismatch] #FIXME: """ Sample elements from the replay buffer. @@ -614,7 +618,11 @@ class DictReplayBuffer(ReplayBuffer): """ return super(ReplayBuffer, self).sample(batch_size=batch_size, env=env) - def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: + def _get_samples( + self, + batch_inds: np.ndarray, + env: Optional[VecNormalize] = None, + ) -> DictReplayBufferSamples: # type: ignore[signature-mismatch] #FIXME: # Sample randomly the env idx env_indices = np.random.randint(0, high=self.n_envs, size=(len(batch_inds),)) @@ -676,7 +684,6 @@ class DictRolloutBuffer(RolloutBuffer): gamma: float = 0.99, n_envs: int = 1, ): - super(RolloutBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs) assert isinstance(self.obs_shape, dict), "DictRolloutBuffer must be used with Dict obs space only" @@ -743,12 +750,14 @@ class DictRolloutBuffer(RolloutBuffer): if self.pos == self.buffer_size: self.full = True - def get(self, batch_size: Optional[int] = None) -> Generator[DictRolloutBufferSamples, None, None]: + def get( + self, + batch_size: Optional[int] = None, + ) -> Generator[DictRolloutBufferSamples, None, None]: # type: ignore[signature-mismatch] #FIXME assert self.full, "" indices = np.random.permutation(self.buffer_size * self.n_envs) # Prepare the data if not self.generator_ready: - for key, obs in self.observations.items(): self.observations[key] = self.swap_and_flatten(obs) @@ -767,8 +776,11 @@ class DictRolloutBuffer(RolloutBuffer): yield self._get_samples(indices[start_idx : start_idx + batch_size]) start_idx += batch_size - def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> DictRolloutBufferSamples: - + def _get_samples( + self, + batch_inds: np.ndarray, + env: Optional[VecNormalize] = None, + ) -> DictRolloutBufferSamples: # type: ignore[signature-mismatch] #FIXME return DictRolloutBufferSamples( observations={key: self.to_torch(obs[batch_inds]) for (key, obs) in self.observations.items()}, actions=self.to_torch(self.actions[batch_inds]), diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index 0ecd1ad..69a21ab 100644 --- a/stable_baselines3/common/callbacks.py +++ b/stable_baselines3/common/callbacks.py @@ -6,6 +6,8 @@ from typing import Any, Callable, Dict, List, Optional, Union import gym import numpy as np +from stable_baselines3.common.logger import Logger + try: from tqdm import TqdmExperimentalWarning @@ -29,10 +31,13 @@ class BaseCallback(ABC): :param verbose: Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages """ + # The RL model + # Type hint as string to avoid circular import + model: "base_class.BaseAlgorithm" + logger: Logger + def __init__(self, verbose: int = 0): super().__init__() - # The RL model - self.model = None # type: Optional[base_class.BaseAlgorithm] # An alias for self.model.get_env(), the environment used for training self.training_env = None # type: Union[gym.Env, VecEnv, None] # Number of time the callback was called @@ -42,7 +47,6 @@ class BaseCallback(ABC): self.verbose = verbose self.locals: Dict[str, Any] = {} self.globals: Dict[str, Any] = {} - self.logger = None # Sometimes, for event callback, it is useful # to have access to the parent object self.parent = None # type: Optional[BaseCallback] @@ -425,11 +429,9 @@ class EvalCallback(EventCallback): self._is_success_buffer.append(maybe_is_success) def _on_step(self) -> bool: - continue_training = True if self.eval_freq > 0 and self.n_calls % self.eval_freq == 0: - # Sync training and eval env if there is VecNormalize if self.model.get_vec_normalize_env() is not None: try: diff --git a/stable_baselines3/common/envs/identity_env.py b/stable_baselines3/common/envs/identity_env.py index 9635e53..a8bed17 100644 --- a/stable_baselines3/common/envs/identity_env.py +++ b/stable_baselines3/common/envs/identity_env.py @@ -71,7 +71,7 @@ class IdentityEnvBox(IdentityEnv[np.ndarray]): super().__init__(ep_length=ep_length, space=space) self.eps = eps - def step(self, action: np.ndarray) -> GymStepReturn: + def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, Dict[str, Any]]: reward = self._get_reward(action) self._choose_next_state() self.current_step += 1 @@ -83,7 +83,7 @@ class IdentityEnvBox(IdentityEnv[np.ndarray]): class IdentityEnvMultiDiscrete(IdentityEnv[np.ndarray]): - def __init__(self, dim: int = 1, ep_length: int = 100): + def __init__(self, dim: int = 1, ep_length: int = 100) -> None: """ Identity environment for testing purposes @@ -95,7 +95,7 @@ class IdentityEnvMultiDiscrete(IdentityEnv[np.ndarray]): class IdentityEnvMultiBinary(IdentityEnv[np.ndarray]): - def __init__(self, dim: int = 1, ep_length: int = 100): + def __init__(self, dim: int = 1, ep_length: int = 100) -> None: """ Identity environment for testing purposes @@ -126,7 +126,7 @@ class FakeImageEnv(gym.Env): n_channels: int = 1, discrete: bool = True, channel_first: bool = False, - ): + ) -> None: self.observation_shape = (screen_height, screen_width, n_channels) if channel_first: self.observation_shape = (n_channels, screen_height, screen_width) diff --git a/stable_baselines3/common/envs/multi_input_envs.py b/stable_baselines3/common/envs/multi_input_envs.py index 433591d..166c699 100644 --- a/stable_baselines3/common/envs/multi_input_envs.py +++ b/stable_baselines3/common/envs/multi_input_envs.py @@ -121,7 +121,7 @@ class SimpleMultiObsEnv(gym.Env): self.right_possible = [0, 1, 2, 12, 13, 14] self.up_possible = [4, 8, 12, 7, 11, 15] - def step(self, action: Union[int, float, np.ndarray]) -> GymStepReturn: + def step(self, action: Union[float, np.ndarray]) -> GymStepReturn: """ Run one timestep of the environment's dynamics. When end of episode is reached, you are responsible for calling `reset()` diff --git a/stable_baselines3/common/evaluation.py b/stable_baselines3/common/evaluation.py index ff18137..b65edf8 100644 --- a/stable_baselines3/common/evaluation.py +++ b/stable_baselines3/common/evaluation.py @@ -91,7 +91,6 @@ def evaluate_policy( current_lengths += 1 for i in range(n_envs): if episode_counts[i] < episode_count_targets[i]: - # unpack values so that the callback can access the local variables reward = rewards[i] done = dones[i] diff --git a/stable_baselines3/common/logger.py b/stable_baselines3/common/logger.py index e065992..a8aa766 100644 --- a/stable_baselines3/common/logger.py +++ b/stable_baselines3/common/logger.py @@ -5,7 +5,7 @@ import sys import tempfile import warnings from collections import defaultdict -from typing import Any, Dict, List, Optional, Sequence, TextIO, Tuple, Union +from typing import Any, Dict, List, Mapping, Optional, Sequence, TextIO, Tuple, Union import numpy as np import pandas @@ -16,7 +16,7 @@ try: from torch.utils.tensorboard import SummaryWriter from torch.utils.tensorboard.summary import hparams except ImportError: - SummaryWriter = None + SummaryWriter = None # type: ignore[misc, assignment] try: from tqdm import tqdm @@ -38,7 +38,7 @@ class Video: :param fps: frames per second """ - def __init__(self, frames: th.Tensor, fps: Union[float, int]): + def __init__(self, frames: th.Tensor, fps: float): self.frames = frames self.fps = fps @@ -80,7 +80,7 @@ class HParam: A non-empty metrics dict is required to display hyperparameters in the corresponding Tensorboard section. """ - def __init__(self, hparam_dict: Dict[str, Union[bool, str, float, int, None]], metric_dict: Dict[str, Union[float, int]]): + def __init__(self, hparam_dict: Mapping[str, Union[bool, str, float, None]], metric_dict: Mapping[str, float]): self.hparam_dict = hparam_dict if not metric_dict: raise Exception("`metric_dict` must not be empty to display hyperparameters to the HPARAMS tensorboard tab.") @@ -173,7 +173,6 @@ class HumanOutputFormat(KVWriter, SeqWriter): key2str = {} tag = None for (key, value), (_, excluded) in zip(sorted(key_values.items()), sorted(key_excluded.items())): - if excluded is not None and ("stdout" in excluded or "log" in excluded): continue @@ -329,7 +328,7 @@ class CSVOutputFormat(KVWriter): def __init__(self, filename: str): self.file = open(filename, "w+t") - self.keys = [] + self.keys: List[str] = [] self.separator = "," self.quotechar = '"' @@ -342,7 +341,7 @@ class CSVOutputFormat(KVWriter): self.file.seek(0) lines = self.file.readlines() self.file.seek(0) - for (i, key) in enumerate(self.keys): + for i, key in enumerate(self.keys): if i > 0: self.file.write(",") self.file.write(key) @@ -399,9 +398,7 @@ class TensorBoardOutputFormat(KVWriter): self.writer = SummaryWriter(log_dir=folder) def write(self, key_values: Dict[str, Any], key_excluded: Dict[str, Union[str, Tuple[str, ...]]], step: int = 0) -> None: - for (key, value), (_, excluded) in zip(sorted(key_values.items()), sorted(key_excluded.items())): - if excluded is not None and "tensorboard" in excluded: continue diff --git a/stable_baselines3/common/monitor.py b/stable_baselines3/common/monitor.py index 1e56fdb..ffc318d 100644 --- a/stable_baselines3/common/monitor.py +++ b/stable_baselines3/common/monitor.py @@ -5,7 +5,7 @@ import json import os import time from glob import glob -from typing import Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union import gym import numpy as np @@ -41,6 +41,7 @@ class Monitor(gym.Wrapper): ): super().__init__(env=env) self.t_start = time.time() + self.results_writer = None if filename is not None: self.results_writer = ResultsWriter( filename, @@ -48,18 +49,18 @@ class Monitor(gym.Wrapper): extra_keys=reset_keywords + info_keywords, override_existing=override_existing, ) - else: - self.results_writer = None + self.reset_keywords = reset_keywords self.info_keywords = info_keywords self.allow_early_resets = allow_early_resets - self.rewards = None + self.rewards: List[float] = [] self.needs_reset = True - self.episode_returns = [] - self.episode_lengths = [] - self.episode_times = [] + self.episode_returns: List[float] = [] + self.episode_lengths: List[int] = [] + self.episode_times: List[float] = [] self.total_steps = 0 - self.current_reset_info = {} # extra info about the current episode, that was passed in during reset() + # extra info about the current episode, that was passed in during reset() + self.current_reset_info: Dict[str, Any] = {} def reset(self, **kwargs) -> GymObs: """ @@ -200,7 +201,7 @@ class ResultsWriter: self.file_handler.flush() - def write_row(self, epinfo: Dict[str, Union[float, int]]) -> None: + def write_row(self, epinfo: Dict[str, float]) -> None: """ Close the file handler diff --git a/stable_baselines3/common/noise.py b/stable_baselines3/common/noise.py index 5e8632d..944408f 100644 --- a/stable_baselines3/common/noise.py +++ b/stable_baselines3/common/noise.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from typing import Iterable, List, Optional import numpy as np +from numpy.typing import DTypeLike class ActionNoise(ABC): @@ -15,7 +16,7 @@ class ActionNoise(ABC): def reset(self) -> None: """ - call end of episode reset for the noise + Call end of episode reset for the noise """ pass @@ -26,19 +27,21 @@ class ActionNoise(ABC): class NormalActionNoise(ActionNoise): """ - A Gaussian action noise + A Gaussian action noise. - :param mean: the mean value of the noise - :param sigma: the scale of the noise (std here) + :param mean: Mean value of the noise + :param sigma: Scale of the noise (std here) + :param dtype: Type of the output noise """ - def __init__(self, mean: np.ndarray, sigma: np.ndarray): + def __init__(self, mean: np.ndarray, sigma: np.ndarray, dtype: DTypeLike = np.float32) -> None: self._mu = mean self._sigma = sigma + self._dtype = dtype super().__init__() def __call__(self) -> np.ndarray: - return np.random.normal(self._mu, self._sigma) + return np.random.normal(self._mu, self._sigma).astype(self._dtype) def __repr__(self) -> str: return f"NormalActionNoise(mu={self._mu}, sigma={self._sigma})" @@ -50,11 +53,12 @@ class OrnsteinUhlenbeckActionNoise(ActionNoise): Based on http://math.stackexchange.com/questions/1287634/implementing-ornstein-uhlenbeck-in-matlab - :param mean: the mean of the noise - :param sigma: the scale of the noise - :param theta: the rate of mean reversion - :param dt: the timestep for the noise - :param initial_noise: the initial value for the noise output, (if None: 0) + :param mean: Mean of the noise + :param sigma: Scale of the noise + :param theta: Rate of mean reversion + :param dt: Timestep for the noise + :param initial_noise: Initial value for the noise output, (if None: 0) + :param dtype: Type of the output noise """ def __init__( @@ -64,11 +68,13 @@ class OrnsteinUhlenbeckActionNoise(ActionNoise): theta: float = 0.15, dt: float = 1e-2, initial_noise: Optional[np.ndarray] = None, - ): + dtype: DTypeLike = np.float32, + ) -> None: self._theta = theta self._mu = mean self._sigma = sigma self._dt = dt + self._dtype = dtype self.initial_noise = initial_noise self.noise_prev = np.zeros_like(self._mu) self.reset() @@ -81,7 +87,7 @@ class OrnsteinUhlenbeckActionNoise(ActionNoise): + self._sigma * np.sqrt(self._dt) * np.random.normal(size=self._mu.shape) ) self.noise_prev = noise - return noise + return noise.astype(self._dtype) def reset(self) -> None: """ @@ -97,11 +103,11 @@ class VectorizedActionNoise(ActionNoise): """ A Vectorized action noise for parallel environments. - :param base_noise: ActionNoise The noise generator to use - :param n_envs: The number of parallel environments + :param base_noise: Noise generator to use + :param n_envs: Number of parallel environments """ - def __init__(self, base_noise: ActionNoise, n_envs: int): + def __init__(self, base_noise: ActionNoise, n_envs: int) -> None: try: self.n_envs = int(n_envs) assert self.n_envs > 0 @@ -113,9 +119,9 @@ class VectorizedActionNoise(ActionNoise): def reset(self, indices: Optional[Iterable[int]] = None) -> None: """ - Reset all the noise processes, or those listed in indices + Reset all the noise processes, or those listed in indices. - :param indices: Optional[Iterable[int]] The indices to reset. Default: None. + :param indices: The indices to reset. Default: None. If the parameter is None, then all processes are reset to their initial position. """ if indices is None: @@ -129,7 +135,7 @@ class VectorizedActionNoise(ActionNoise): def __call__(self) -> np.ndarray: """ - Generate and stack the action noise from each noise object + Generate and stack the action noise from each noise object. """ noise = np.stack([noise() for noise in self.noises]) return noise diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 48779ef..c1ab215 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -102,7 +102,6 @@ class OffPolicyAlgorithm(BaseAlgorithm): sde_support: bool = True, supported_action_spaces: Optional[Tuple[spaces.Space, ...]] = None, ): - super().__init__( policy=policy, env=env, @@ -319,7 +318,6 @@ class OffPolicyAlgorithm(BaseAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfOffPolicyAlgorithm: - total_timesteps, callback = self._setup_learn( total_timesteps, callback, diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index bc0dda4..44d8b26 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -72,7 +72,6 @@ class OnPolicyAlgorithm(BaseAlgorithm): _init_setup_model: bool = True, supported_action_spaces: Optional[Tuple[spaces.Space, ...]] = None, ): - super().__init__( policy=policy, env=env, @@ -244,7 +243,6 @@ class OnPolicyAlgorithm(BaseAlgorithm): callback.on_training_start(locals(), globals()) 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: diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 0cf4917..457274a 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -418,8 +418,7 @@ class ActorCriticPolicy(BasePolicy): observation_space: spaces.Space, action_space: spaces.Space, lr_schedule: Schedule, - # TODO(antonin): update type annotation when we remove shared network support - net_arch: Union[List[int], Dict[str, List[int]], List[Dict[str, List[int]]], None] = None, + net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, use_sde: bool = False, @@ -434,7 +433,6 @@ class ActorCriticPolicy(BasePolicy): optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam, optimizer_kwargs: Optional[Dict[str, Any]] = None, ): - if optimizer_kwargs is None: optimizer_kwargs = {} # Small values to avoid NaN in Adam optimizer @@ -452,21 +450,15 @@ class ActorCriticPolicy(BasePolicy): normalize_images=normalize_images, ) - # Convert [dict()] to dict() as shared network are deprecated - if isinstance(net_arch, list) and len(net_arch) > 0: - if isinstance(net_arch[0], dict): - warnings.warn( - ( - "As shared layers in the mlp_extractor are deprecated and will be removed in SB3 v1.8.0, " - "you should now pass directly a dictionary and not a list " - "(net_arch=dict(pi=..., vf=...) instead of net_arch=[dict(pi=..., vf=...)])" - ), - ) - net_arch = net_arch[0] - else: - # Note: deprecation warning will be emitted - # by the MlpExtractor constructor - pass + if isinstance(net_arch, list) and len(net_arch) > 0 and isinstance(net_arch[0], dict): + warnings.warn( + ( + "As shared layers in the mlp_extractor are removed since SB3 v1.8.0, " + "you should now pass directly a dictionary and not a list " + "(net_arch=dict(pi=..., vf=...) instead of net_arch=[dict(pi=..., vf=...)])" + ), + ) + net_arch = net_arch[0] # Default network architecture, from stable-baselines if net_arch is None: @@ -488,12 +480,6 @@ class ActorCriticPolicy(BasePolicy): else: self.pi_features_extractor = self.features_extractor self.vf_features_extractor = self.make_features_extractor() - # if the features extractor is not shared, there cannot be shared layers in the mlp_extractor - # TODO(antonin): update the check once we change net_arch behavior - if isinstance(net_arch, list) and len(net_arch) > 0: - raise ValueError( - "Error: if the features extractor is not shared, there cannot be shared layers in the mlp_extractor" - ) self.log_std_init = log_std_init dist_kwargs = None @@ -770,7 +756,7 @@ class ActorCriticCnnPolicy(ActorCriticPolicy): observation_space: spaces.Space, action_space: spaces.Space, lr_schedule: Schedule, - net_arch: Union[List[int], Dict[str, List[int]], List[Dict[str, List[int]]], None] = None, + net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, use_sde: bool = False, @@ -843,7 +829,7 @@ class MultiInputActorCriticPolicy(ActorCriticPolicy): observation_space: spaces.Dict, action_space: spaces.Space, lr_schedule: Schedule, - net_arch: Union[List[int], Dict[str, List[int]], List[Dict[str, List[int]]], None] = None, + net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, use_sde: bool = False, diff --git a/stable_baselines3/common/results_plotter.py b/stable_baselines3/common/results_plotter.py index 92f67ac..dac2b6c 100644 --- a/stable_baselines3/common/results_plotter.py +++ b/stable_baselines3/common/results_plotter.py @@ -84,7 +84,7 @@ def plot_curves( plt.figure(title, figsize=figsize) max_x = max(xy[0][-1] for xy in xy_list) min_x = 0 - for (_, (x, y)) in enumerate(xy_list): + for _, (x, y) in enumerate(xy_list): plt.scatter(x, y, s=2) # Do not plot the smoothed curve at all if the timeseries is shorter than window size. if x.shape[0] >= EPISODES_WINDOW: diff --git a/stable_baselines3/common/running_mean_std.py b/stable_baselines3/common/running_mean_std.py index b48f922..9dfa4b8 100644 --- a/stable_baselines3/common/running_mean_std.py +++ b/stable_baselines3/common/running_mean_std.py @@ -1,4 +1,4 @@ -from typing import Tuple, Union +from typing import Tuple import numpy as np @@ -40,7 +40,7 @@ class RunningMeanStd: batch_count = arr.shape[0] self.update_from_moments(batch_mean, batch_var, batch_count) - def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: Union[int, float]) -> None: + def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: float) -> None: delta = batch_mean - self.mean tot_count = self.count + batch_count diff --git a/stable_baselines3/common/save_util.py b/stable_baselines3/common/save_util.py index facc55a..7ae1e22 100644 --- a/stable_baselines3/common/save_util.py +++ b/stable_baselines3/common/save_util.py @@ -367,7 +367,7 @@ def load_from_zip_file( device: Union[th.device, str] = "auto", verbose: int = 0, print_system_info: bool = False, -) -> (Tuple[Optional[Dict[str, Any]], Optional[TensorDict], Optional[TensorDict]]): +) -> Tuple[Optional[Dict[str, Any]], Optional[TensorDict], Optional[TensorDict]]: """ Load model data from a .zip archive diff --git a/stable_baselines3/common/torch_layers.py b/stable_baselines3/common/torch_layers.py index 302d9b1..44714d6 100644 --- a/stable_baselines3/common/torch_layers.py +++ b/stable_baselines3/common/torch_layers.py @@ -1,5 +1,3 @@ -import warnings -from itertools import zip_longest from typing import Dict, List, Tuple, Type, Union import gym @@ -151,98 +149,57 @@ class MlpExtractor(nn.Module): Constructs an MLP that receives the output from a previous features extractor (i.e. a CNN) or directly the observations (if no features extractor is applied) as an input and outputs a latent representation for the policy and a value network. - The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many - of them are shared between the policy network and the value network. It is assumed to be a list with the following - structure: - 1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer. - If the number of ints is zero, there will be no shared layers. - 2. An optional dict, to specify the following non-shared layers for the value network and the policy network. - It is formatted like ``dict(vf=[], pi=[])``. - If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed. + The ``net_arch`` parameter allows to specify the amount and size of the hidden layers. + It can be in either of the following forms: + 1. ``dict(vf=[], pi=[])``: to specify the amount and size of the layers in the + policy and value nets individually. If it is missing any of the keys (pi or vf), + zero layers will be considered for that key. + 2. ``[]``: "shortcut" in case the amount and size of the layers + in the policy and value nets are the same. Same as ``dict(vf=int_list, pi=int_list)`` + where int_list is the same for the actor and critic. - Deprecation note: shared layers in ``net_arch`` are deprecated, please use separate - pi and vf networks (e.g. net_arch=dict(pi=[...], vf=[...])) - - For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value - network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec - would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128 - would be specified as [128, 128]. - - Adapted from Stable Baselines. + .. note:: + If a key is not specified or an empty list is passed ``[]``, a linear network will be used. :param feature_dim: Dimension of the feature vector (can be the output of a CNN) :param net_arch: The specification of the policy and value networks. See above for details on its formatting. :param activation_fn: The activation function to use for the networks. - :param device: + :param device: PyTorch device. """ def __init__( self, feature_dim: int, - net_arch: Union[Dict[str, List[int]], List[Union[int, Dict[str, List[int]]]]], + net_arch: Union[List[int], Dict[str, List[int]]], activation_fn: Type[nn.Module], device: Union[th.device, str] = "auto", ) -> None: super().__init__() device = get_device(device) - shared_net: List[nn.Module] = [] policy_net: List[nn.Module] = [] value_net: List[nn.Module] = [] - policy_only_layers: List[int] = [] # Layer sizes of the network that only belongs to the policy network - value_only_layers: List[int] = [] # Layer sizes of the network that only belongs to the value network - last_layer_dim_shared = feature_dim + last_layer_dim_pi = feature_dim + last_layer_dim_vf = feature_dim - if isinstance(net_arch, list) and len(net_arch) > 0 and isinstance(net_arch[0], int): - warnings.warn( - ( - "Shared layers in the mlp_extractor are deprecated and will be removed in SB3 v1.8.0, " - "please use separate pi and vf networks " - "(e.g. net_arch=dict(pi=[...], vf=[...]))" - ), - DeprecationWarning, - ) - - # TODO(antonin): update behavior for net_arch=[64, 64] - # once shared networks are removed + # save dimensions of layers in policy and value nets if isinstance(net_arch, dict): - policy_only_layers = net_arch["pi"] - value_only_layers = net_arch["vf"] + # Note: if key is not specificed, assume linear network + pi_layers_dims = net_arch.get("pi", []) # Layer sizes of the policy network + vf_layers_dims = net_arch.get("vf", []) # Layer sizes of the value network else: - # Iterate through the shared layers and build the shared parts of the network - for layer in net_arch: - if isinstance(layer, int): # Check that this is a shared layer - shared_net.append(nn.Linear(last_layer_dim_shared, layer)) # add linear of size layer - shared_net.append(activation_fn()) - last_layer_dim_shared = layer - else: - assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts" - if "pi" in layer: - assert isinstance(layer["pi"], list), "Error: net_arch[-1]['pi'] must contain a list of integers." - policy_only_layers = layer["pi"] - - if "vf" in layer: - assert isinstance(layer["vf"], list), "Error: net_arch[-1]['vf'] must contain a list of integers." - value_only_layers = layer["vf"] - break # From here on the network splits up in policy and value network - - last_layer_dim_pi = last_layer_dim_shared - last_layer_dim_vf = last_layer_dim_shared - - # Build the non-shared part of the network - for pi_layer_size, vf_layer_size in zip_longest(policy_only_layers, value_only_layers): - if pi_layer_size is not None: - assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers." - policy_net.append(nn.Linear(last_layer_dim_pi, pi_layer_size)) - policy_net.append(activation_fn()) - last_layer_dim_pi = pi_layer_size - - if vf_layer_size is not None: - assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers." - value_net.append(nn.Linear(last_layer_dim_vf, vf_layer_size)) - value_net.append(activation_fn()) - last_layer_dim_vf = vf_layer_size + pi_layers_dims = vf_layers_dims = net_arch + # Iterate through the policy layers and build the policy net + for curr_layer_dim in pi_layers_dims: + policy_net.append(nn.Linear(last_layer_dim_pi, curr_layer_dim)) + policy_net.append(activation_fn()) + last_layer_dim_pi = curr_layer_dim + # Iterate through the value layers and build the value net + for curr_layer_dim in vf_layers_dims: + value_net.append(nn.Linear(last_layer_dim_vf, curr_layer_dim)) + value_net.append(activation_fn()) + last_layer_dim_vf = curr_layer_dim # Save dim, used to create the distributions self.latent_dim_pi = last_layer_dim_pi @@ -250,7 +207,6 @@ class MlpExtractor(nn.Module): # Create networks # If the list of layers is empty, the network will just act as an Identity module - self.shared_net = nn.Sequential(*shared_net).to(device) self.policy_net = nn.Sequential(*policy_net).to(device) self.value_net = nn.Sequential(*value_net).to(device) @@ -259,14 +215,13 @@ class MlpExtractor(nn.Module): :return: latent_policy, latent_value of the specified network. If all layers are shared, then ``latent_policy == latent_value`` """ - shared_latent = self.shared_net(features) - return self.policy_net(shared_latent), self.value_net(shared_latent) + return self.forward_actor(features), self.forward_critic(features) def forward_actor(self, features: th.Tensor) -> th.Tensor: - return self.policy_net(self.shared_net(features)) + return self.policy_net(features) def forward_critic(self, features: th.Tensor) -> th.Tensor: - return self.value_net(self.shared_net(features)) + return self.value_net(features) class CombinedExtractor(BaseFeaturesExtractor): diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 6dc247a..5c8f215 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -77,7 +77,7 @@ def update_learning_rate(optimizer: th.optim.Optimizer, learning_rate: float) -> param_group["lr"] = learning_rate -def get_schedule_fn(value_schedule: Union[Schedule, float, int]) -> Schedule: +def get_schedule_fn(value_schedule: Union[Schedule, float]) -> Schedule: """ Transform (if needed) learning rate and clip range (for PPO) to callable. diff --git a/stable_baselines3/common/vec_env/__init__.py b/stable_baselines3/common/vec_env/__init__.py index 33a103a..2c03637 100644 --- a/stable_baselines3/common/vec_env/__init__.py +++ b/stable_baselines3/common/vec_env/__init__.py @@ -4,7 +4,7 @@ from typing import Optional, Type, Union from stable_baselines3.common.vec_env.base_vec_env import CloudpickleWrapper, VecEnv, VecEnvWrapper from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv -from stable_baselines3.common.vec_env.stacked_observations import StackedDictObservations, StackedObservations +from stable_baselines3.common.vec_env.stacked_observations import StackedObservations from stable_baselines3.common.vec_env.subproc_vec_env import SubprocVecEnv from stable_baselines3.common.vec_env.vec_check_nan import VecCheckNan from stable_baselines3.common.vec_env.vec_extract_dict_obs import VecExtractDictObs @@ -78,7 +78,6 @@ __all__ = [ "VecEnv", "VecEnvWrapper", "DummyVecEnv", - "StackedDictObservations", "StackedObservations", "SubprocVecEnv", "VecCheckNan", diff --git a/stable_baselines3/common/vec_env/stacked_observations.py b/stable_baselines3/common/vec_env/stacked_observations.py index 8583518..a26812c 100644 --- a/stable_baselines3/common/vec_env/stacked_observations.py +++ b/stable_baselines3/common/vec_env/stacked_observations.py @@ -1,62 +1,80 @@ import warnings -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, Generic, List, Mapping, Optional, Tuple, TypeVar, Union import numpy as np from gym import spaces from stable_baselines3.common.preprocessing import is_image_space, is_image_space_channels_first +TObs = TypeVar("TObs", np.ndarray, Dict[str, np.ndarray]) -class StackedObservations: + +# Disable errors for pytype which doesn't play well with Generic[TypeVar] +# mypy check passes though +# pytype: disable=attribute-error +class StackedObservations(Generic[TObs]): """ Frame stacking wrapper for data. - Dimension to stack over is either first (channels-first) or - last (channels-last), which is detected automatically using - ``common.preprocessing.is_image_space_channels_first`` if - observation is an image space. + Dimension to stack over is either first (channels-first) or last (channels-last), which is detected automatically using + ``common.preprocessing.is_image_space_channels_first`` if observation is an image space. - :param num_envs: number of environments + :param num_envs: Number of environments :param n_stack: Number of frames to stack - :param observation_space: Environment observation space. + :param observation_space: Environment observation space :param channels_order: If "first", stack on first image dimension. If "last", stack on last dimension. - If None, automatically detect channel to stack over in case of image observation or default to "last" (default). + If None, automatically detect channel to stack over in case of image observation or default to "last". + For Dict space, channels_order can also be a dictionary. """ def __init__( self, num_envs: int, n_stack: int, - observation_space: spaces.Space, - channels_order: Optional[str] = None, - ): - + observation_space: Union[spaces.Box, spaces.Dict], # Replace by Space[TObs] in gym>=0.26 + channels_order: Optional[Union[str, Mapping[str, Optional[str]]]] = None, + ) -> None: self.n_stack = n_stack - ( - self.channels_first, - self.stack_dimension, - self.stackedobs, - self.repeat_axis, - ) = self.compute_stacking(num_envs, n_stack, observation_space, channels_order) - super().__init__() + self.observation_space = observation_space + if isinstance(observation_space, spaces.Dict): + if not isinstance(channels_order, Mapping): + channels_order = {key: channels_order for key in observation_space.spaces.keys()} + self.sub_stacked_observations = { + key: StackedObservations(num_envs, n_stack, subspace, channels_order[key]) + for key, subspace in observation_space.spaces.items() + } + self.stacked_observation_space = spaces.Dict( + {key: substack_obs.stacked_observation_space for key, substack_obs in self.sub_stacked_observations.items()} + ) # type: spaces.Dict # make mypy happy + elif isinstance(observation_space, spaces.Box): + if isinstance(channels_order, Mapping): + raise TypeError("When the observation space is Box, channels_order can't be a dict.") + + self.channels_first, self.stack_dimension, self.stacked_shape, self.repeat_axis = self.compute_stacking( + n_stack, observation_space, channels_order + ) + low = np.repeat(observation_space.low, n_stack, axis=self.repeat_axis) + high = np.repeat(observation_space.high, n_stack, axis=self.repeat_axis) + self.stacked_observation_space = spaces.Box(low=low, high=high, dtype=observation_space.dtype) + self.stacked_obs = np.zeros((num_envs,) + self.stacked_shape, dtype=observation_space.dtype) + else: + raise TypeError( + f"StackedObservations only supports Box and Dict as observation spaces. {observation_space} was provided." + ) @staticmethod def compute_stacking( - num_envs: int, - n_stack: int, - observation_space: spaces.Box, - channels_order: Optional[str] = None, - ) -> Tuple[bool, int, np.ndarray, int]: + n_stack: int, observation_space: spaces.Box, channels_order: Optional[str] = None + ) -> Tuple[bool, int, Tuple[int, ...], int]: """ Calculates the parameters in order to stack observations - :param num_envs: Number of environments in the stack - :param n_stack: The number of observations to stack - :param observation_space: The observation space - :param channels_order: The order of the channels - :return: tuple of channels_first, stack_dimension, stackedobs, repeat_axis + :param n_stack: Number of observations to stack + :param observation_space: Observation space + :param channels_order: Order of the channels + :return: Tuple of channels_first, stack_dimension, stackedobs, repeat_axis """ - channels_first = False + if channels_order is None: # Detect channel location automatically for images if is_image_space(observation_space): @@ -75,192 +93,113 @@ class StackedObservations: # This includes the vec-env dimension (first) stack_dimension = 1 if channels_first else -1 repeat_axis = 0 if channels_first else -1 - low = np.repeat(observation_space.low, n_stack, axis=repeat_axis) - stackedobs = np.zeros((num_envs,) + low.shape, low.dtype) - return channels_first, stack_dimension, stackedobs, repeat_axis + stacked_shape = list(observation_space.shape) + stacked_shape[repeat_axis] *= n_stack + return channels_first, stack_dimension, tuple(stacked_shape), repeat_axis - def stack_observation_space(self, observation_space: spaces.Box) -> spaces.Box: + def stack_observation_space(self, observation_space: Union[spaces.Box, spaces.Dict]) -> Union[spaces.Box, spaces.Dict]: """ - Given an observation space, returns a new observation space with stacked observations + This function is deprecated. + + As an alternative, use + + .. code-block:: python + + low = np.repeat(observation_space.low, stacked_observation.n_stack, axis=stacked_observation.repeat_axis) + high = np.repeat(observation_space.high, stacked_observation.n_stack, axis=stacked_observation.repeat_axis) + stacked_observation_space = spaces.Box(low=low, high=high, dtype=observation_space.dtype) :return: New observation space with stacked dimensions """ + warnings.warn( + "stack_observation_space is deprecated and will be removed in the next SB3 release. " + "Please refer to the docstring for a workaround.", + DeprecationWarning, + ) + if isinstance(observation_space, spaces.Dict): + return spaces.Dict( + { + key: sub_stacked_observation.stack_observation_space(sub_stacked_observation.observation_space) + for key, sub_stacked_observation in self.sub_stacked_observations.items() + } + ) low = np.repeat(observation_space.low, self.n_stack, axis=self.repeat_axis) high = np.repeat(observation_space.high, self.n_stack, axis=self.repeat_axis) return spaces.Box(low=low, high=high, dtype=observation_space.dtype) - def reset(self, observation: np.ndarray) -> np.ndarray: + def reset(self, observation: TObs) -> TObs: """ - Resets the stackedobs, adds the reset observation to the stack, and returns the stack + Reset the stacked_obs, add the reset observation to the stack, and return the stack. :param observation: Reset observation :return: The stacked reset observation """ - self.stackedobs[...] = 0 + if isinstance(observation, dict): + return {key: self.sub_stacked_observations[key].reset(obs) for key, obs in observation.items()} + + self.stacked_obs[...] = 0 if self.channels_first: - self.stackedobs[:, -observation.shape[self.stack_dimension] :, ...] = observation + self.stacked_obs[:, -observation.shape[self.stack_dimension] :, ...] = observation else: - self.stackedobs[..., -observation.shape[self.stack_dimension] :] = observation - return self.stackedobs + self.stacked_obs[..., -observation.shape[self.stack_dimension] :] = observation + return self.stacked_obs def update( self, - observations: np.ndarray, + observations: TObs, dones: np.ndarray, infos: List[Dict[str, Any]], - ) -> Tuple[np.ndarray, List[Dict[str, Any]]]: + ) -> Tuple[TObs, List[Dict[str, Any]]]: """ - Adds the observations to the stack and uses the dones to update the infos. + Add the observations to the stack and use the dones to update the infos. - :param observations: numpy array of observations - :param dones: numpy array of done info - :param infos: numpy array of info dicts - :return: tuple of the stacked observations and the updated infos + :param observations: Observations + :param dones: Dones + :param infos: Infos + :return: Tuple of the stacked observations and the updated infos """ - stack_ax_size = observations.shape[self.stack_dimension] - self.stackedobs = np.roll(self.stackedobs, shift=-stack_ax_size, axis=self.stack_dimension) - for i, done in enumerate(dones): + if isinstance(observations, dict): + # From [{}, {terminal_obs: {key1: ..., key2: ...}}] + # to {key1: [{}, {terminal_obs: ...}], key2: [{}, {terminal_obs: ...}]} + sub_infos = { + key: [ + {"terminal_observation": info["terminal_observation"][key]} if "terminal_observation" in info else {} + for info in infos + ] + for key in observations.keys() + } + + stacked_obs = {} + stacked_infos = {} + for key, obs in observations.items(): + stacked_obs[key], stacked_infos[key] = self.sub_stacked_observations[key].update(obs, dones, sub_infos[key]) + + # From {key1: [{}, {terminal_obs: ...}], key2: [{}, {terminal_obs: ...}]} + # to [{}, {terminal_obs: {key1: ..., key2: ...}}] + for key in stacked_infos.keys(): + for env_idx in range(len(infos)): + if "terminal_observation" in infos[env_idx]: + infos[env_idx]["terminal_observation"][key] = stacked_infos[key][env_idx]["terminal_observation"] + return stacked_obs, infos + + shift = -observations.shape[self.stack_dimension] + self.stacked_obs = np.roll(self.stacked_obs, shift, axis=self.stack_dimension) + for env_idx, done in enumerate(dones): if done: - if "terminal_observation" in infos[i]: - old_terminal = infos[i]["terminal_observation"] + if "terminal_observation" in infos[env_idx]: + old_terminal = infos[env_idx]["terminal_observation"] if self.channels_first: - new_terminal = np.concatenate( - (self.stackedobs[i, :-stack_ax_size, ...], old_terminal), - axis=0, # self.stack_dimension - 1, as there is not batch dim - ) + previous_stack = self.stacked_obs[env_idx, :shift, ...] else: - new_terminal = np.concatenate( - (self.stackedobs[i, ..., :-stack_ax_size], old_terminal), - axis=self.stack_dimension, - ) - infos[i]["terminal_observation"] = new_terminal + previous_stack = self.stacked_obs[env_idx, ..., :shift] + + new_terminal = np.concatenate((previous_stack, old_terminal), axis=self.repeat_axis) + infos[env_idx]["terminal_observation"] = new_terminal else: warnings.warn("VecFrameStack wrapping a VecEnv without terminal_observation info") - self.stackedobs[i] = 0 + self.stacked_obs[env_idx] = 0 if self.channels_first: - self.stackedobs[:, -observations.shape[self.stack_dimension] :, ...] = observations + self.stacked_obs[:, shift:, ...] = observations else: - self.stackedobs[..., -observations.shape[self.stack_dimension] :] = observations - return self.stackedobs, infos - - -class StackedDictObservations(StackedObservations): - """ - Frame stacking wrapper for dictionary data. - - Dimension to stack over is either first (channels-first) or - last (channels-last), which is detected automatically using - ``common.preprocessing.is_image_space_channels_first`` if - observation is an image space. - - :param num_envs: number of environments - :param n_stack: Number of frames to stack - :param channels_order: If "first", stack on first image dimension. If "last", stack on last dimension. - If None, automatically detect channel to stack over in case of image observation or default to "last" (default). - """ - - def __init__( - self, - num_envs: int, - n_stack: int, - observation_space: spaces.Dict, - channels_order: Optional[Union[str, Dict[str, str]]] = None, - ): - self.n_stack = n_stack - self.channels_first = {} - self.stack_dimension = {} - self.stackedobs = {} - self.repeat_axis = {} - - for key, subspace in observation_space.spaces.items(): - assert isinstance(subspace, spaces.Box), "StackedDictObservations only works with nested gym.spaces.Box" - if isinstance(channels_order, str) or channels_order is None: - subspace_channel_order = channels_order - else: - subspace_channel_order = channels_order[key] - ( - self.channels_first[key], - self.stack_dimension[key], - self.stackedobs[key], - self.repeat_axis[key], - ) = self.compute_stacking(num_envs, n_stack, subspace, subspace_channel_order) - - def stack_observation_space(self, observation_space: spaces.Dict) -> spaces.Dict: - """ - Returns the stacked version of a Dict observation space - - :param observation_space: Dict observation space to stack - :return: stacked observation space - """ - spaces_dict = {} - for key, subspace in observation_space.spaces.items(): - low = np.repeat(subspace.low, self.n_stack, axis=self.repeat_axis[key]) - high = np.repeat(subspace.high, self.n_stack, axis=self.repeat_axis[key]) - spaces_dict[key] = spaces.Box(low=low, high=high, dtype=subspace.dtype) - return spaces.Dict(spaces=spaces_dict) - - def reset(self, observation: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: # pytype: disable=signature-mismatch - """ - Resets the stacked observations, adds the reset observation to the stack, and returns the stack - - :param observation: Reset observation - :return: Stacked reset observations - """ - for key, obs in observation.items(): - self.stackedobs[key][...] = 0 - if self.channels_first[key]: - self.stackedobs[key][:, -obs.shape[self.stack_dimension[key]] :, ...] = obs - else: - self.stackedobs[key][..., -obs.shape[self.stack_dimension[key]] :] = obs - return self.stackedobs - - def update( - self, - observations: Dict[str, np.ndarray], - dones: np.ndarray, - infos: List[Dict[str, Any]], - ) -> Tuple[Dict[str, np.ndarray], List[Dict[str, Any]]]: # pytype: disable=signature-mismatch - """ - Adds the observations to the stack and uses the dones to update the infos. - - :param observations: Dict of numpy arrays of observations - :param dones: numpy array of dones - :param infos: dict of infos - :return: tuple of the stacked observations and the updated infos - """ - for key in self.stackedobs.keys(): - stack_ax_size = observations[key].shape[self.stack_dimension[key]] - self.stackedobs[key] = np.roll( - self.stackedobs[key], - shift=-stack_ax_size, - axis=self.stack_dimension[key], - ) - - for i, done in enumerate(dones): - if done: - if "terminal_observation" in infos[i]: - old_terminal = infos[i]["terminal_observation"][key] - if self.channels_first[key]: - new_terminal = np.vstack( - ( - self.stackedobs[key][i, :-stack_ax_size, ...], - old_terminal, - ) - ) - else: - new_terminal = np.concatenate( - ( - self.stackedobs[key][i, ..., :-stack_ax_size], - old_terminal, - ), - axis=self.stack_dimension[key], - ) - infos[i]["terminal_observation"][key] = new_terminal - else: - warnings.warn("VecFrameStack wrapping a VecEnv without terminal_observation info") - self.stackedobs[key][i] = 0 - if self.channels_first[key]: - self.stackedobs[key][:, -stack_ax_size:, ...] = observations[key] - else: - self.stackedobs[key][..., -stack_ax_size:] = observations[key] - return self.stackedobs, infos + self.stacked_obs[..., shift:] = observations + return self.stacked_obs, infos diff --git a/stable_baselines3/common/vec_env/vec_frame_stack.py b/stable_baselines3/common/vec_env/vec_frame_stack.py index e06d512..8a020dd 100644 --- a/stable_baselines3/common/vec_env/vec_frame_stack.py +++ b/stable_baselines3/common/vec_env/vec_frame_stack.py @@ -1,64 +1,40 @@ -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Mapping, Optional, Tuple, Union import numpy as np from gym import spaces from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper -from stable_baselines3.common.vec_env.stacked_observations import StackedDictObservations, StackedObservations +from stable_baselines3.common.vec_env.stacked_observations import StackedObservations class VecFrameStack(VecEnvWrapper): """ Frame stacking wrapper for vectorized environment. Designed for image observations. - Uses the StackedObservations class, or StackedDictObservations depending on the observations space - - :param venv: the vectorized environment to wrap + :param venv: Vectorized environment to wrap :param n_stack: Number of frames to stack :param channels_order: If "first", stack on first image dimension. If "last", stack on last dimension. If None, automatically detect channel to stack over in case of image observation or default to "last" (default). Alternatively channels_order can be a dictionary which can be used with environments with Dict observation spaces """ - def __init__(self, venv: VecEnv, n_stack: int, channels_order: Optional[Union[str, Dict[str, str]]] = None): - self.venv = venv - self.n_stack = n_stack + def __init__(self, venv: VecEnv, n_stack: int, channels_order: Optional[Union[str, Mapping[str, str]]] = None) -> None: + assert isinstance( + venv.observation_space, (spaces.Box, spaces.Dict) + ), "VecFrameStack only works with gym.spaces.Box and gym.spaces.Dict observation spaces" - wrapped_obs_space = venv.observation_space - - if isinstance(wrapped_obs_space, spaces.Box): - assert not isinstance( - channels_order, dict - ), f"Expected None or string for channels_order but received {channels_order}" - self.stackedobs = StackedObservations(venv.num_envs, n_stack, wrapped_obs_space, channels_order) - - elif isinstance(wrapped_obs_space, spaces.Dict): - self.stackedobs = StackedDictObservations(venv.num_envs, n_stack, wrapped_obs_space, channels_order) - - else: - raise Exception("VecFrameStack only works with gym.spaces.Box and gym.spaces.Dict observation spaces") - - observation_space = self.stackedobs.stack_observation_space(wrapped_obs_space) - VecEnvWrapper.__init__(self, venv, observation_space=observation_space) + self.stacked_obs = StackedObservations(venv.num_envs, n_stack, venv.observation_space, channels_order) + observation_space = self.stacked_obs.stacked_observation_space + super().__init__(venv, observation_space=observation_space) def step_wait( self, ) -> Tuple[Union[np.ndarray, Dict[str, np.ndarray]], np.ndarray, np.ndarray, List[Dict[str, Any]],]: - observations, rewards, dones, infos = self.venv.step_wait() - - observations, infos = self.stackedobs.update(observations, dones, infos) - + observations, infos = self.stacked_obs.update(observations, dones, infos) return observations, rewards, dones, infos def reset(self) -> Union[np.ndarray, Dict[str, np.ndarray]]: - """ - Reset all environments - """ observation = self.venv.reset() # pytype:disable=annotation-type-mismatch - - observation = self.stackedobs.reset(observation) + observation = self.stacked_obs.reset(observation) return observation - - def close(self) -> None: - self.venv.close() diff --git a/stable_baselines3/common/vec_env/vec_video_recorder.py b/stable_baselines3/common/vec_env/vec_video_recorder.py index 70d74eb..83d058a 100644 --- a/stable_baselines3/common/vec_env/vec_video_recorder.py +++ b/stable_baselines3/common/vec_env/vec_video_recorder.py @@ -30,7 +30,6 @@ class VecVideoRecorder(VecEnvWrapper): video_length: int = 200, name_prefix: str = "rl-video", ): - VecEnvWrapper.__init__(self, venv) self.env = venv diff --git a/stable_baselines3/ddpg/ddpg.py b/stable_baselines3/ddpg/ddpg.py index 40d67b5..c311b23 100644 --- a/stable_baselines3/ddpg/ddpg.py +++ b/stable_baselines3/ddpg/ddpg.py @@ -76,7 +76,6 @@ class DDPG(TD3): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy=policy, env=env, @@ -121,7 +120,6 @@ class DDPG(TD3): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfDDPG: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index dd8794e..ea1946a 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -94,7 +94,6 @@ class DQN(OffPolicyAlgorithm): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy, env, @@ -261,7 +260,6 @@ class DQN(OffPolicyAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfDQN: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/her/her_replay_buffer.py b/stable_baselines3/her/her_replay_buffer.py index 1518436..0c3da25 100644 --- a/stable_baselines3/her/her_replay_buffer.py +++ b/stable_baselines3/her/her_replay_buffer.py @@ -81,7 +81,6 @@ class HerReplayBuffer(DictReplayBuffer): online_sampling: bool = True, handle_timeout_termination: bool = True, ): - super().__init__(buffer_size, env.observation_space, env.action_space, device, env.num_envs) # convert goal_selection_strategy into GoalSelectionStrategy if string @@ -389,7 +388,6 @@ class HerReplayBuffer(DictReplayBuffer): done: np.ndarray, infos: List[Dict[str, Any]], ) -> None: - if self.current_idx == 0 and self.full: # Clear info buffer self.info_buffer[self.pos] = deque(maxlen=self.max_episode_length) diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index bd80736..3ea6756 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/stable_baselines3/ppo/ppo.py @@ -98,7 +98,6 @@ class PPO(OnPolicyAlgorithm): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy, env, @@ -190,7 +189,6 @@ class PPO(OnPolicyAlgorithm): clip_fractions = [] continue_training = True - # train for n_epochs epochs for epoch in range(self.n_epochs): approx_kl_divs = [] @@ -272,10 +270,10 @@ class PPO(OnPolicyAlgorithm): th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) self.policy.optimizer.step() + self._n_updates += 1 if not continue_training: break - self._n_updates += self.n_epochs explained_var = explained_variance(self.rollout_buffer.values.flatten(), self.rollout_buffer.returns.flatten()) # Logs @@ -303,7 +301,6 @@ class PPO(OnPolicyAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfPPO: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 74285b6..d1a6610 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -109,7 +109,6 @@ class SAC(OffPolicyAlgorithm): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy, env, @@ -295,7 +294,6 @@ class SAC(OffPolicyAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfSAC: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index ae442e1..c844a99 100644 --- a/stable_baselines3/td3/td3.py +++ b/stable_baselines3/td3/td3.py @@ -94,7 +94,6 @@ class TD3(OffPolicyAlgorithm): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy, env, @@ -151,7 +150,6 @@ class TD3(OffPolicyAlgorithm): actor_losses, critic_losses = [], [] for _ in range(gradient_steps): - self._n_updates += 1 # Sample replay buffer replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env) @@ -210,7 +208,6 @@ class TD3(OffPolicyAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfTD3: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 52d893b..e8175d3 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.8.0a0 +1.8.0a4 diff --git a/tests/test_custom_policy.py b/tests/test_custom_policy.py index 85c3d37..1f89b23 100644 --- a/tests/test_custom_policy.py +++ b/tests/test_custom_policy.py @@ -9,21 +9,21 @@ from stable_baselines3.common.sb2_compat.rmsprop_tf_like import RMSpropTFLike "net_arch", [ [], - dict(vf=[16], pi=[8]), - # [] behavior will change [4], [4, 4], - # All values below are deprecated - [12, dict(vf=[16], pi=[8])], - [12, dict(vf=[8, 4], pi=[8])], - [12, dict(vf=[8], pi=[8, 4])], - [12, dict(pi=[8])], + dict(vf=[16], pi=[8]), + dict(vf=[8, 4], pi=[8]), + dict(vf=[8], pi=[8, 4]), + dict(pi=[8]), + # Old format, emits a warning + [dict(vf=[8])], + [dict(vf=[8], pi=[4])], ], ) @pytest.mark.parametrize("model_class", [A2C, PPO]) def test_flexible_mlp(model_class, net_arch): - if isinstance(net_arch, list) and len(net_arch) > 0 and isinstance(net_arch[0], int): - with pytest.warns(DeprecationWarning): + if isinstance(net_arch, list) and len(net_arch) > 0 and isinstance(net_arch[0], dict): + with pytest.warns(UserWarning): _ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=64).learn(300) else: _ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=64).learn(300) @@ -62,10 +62,3 @@ def test_tf_like_rmsprop_optimizer(): def test_dqn_custom_policy(): policy_kwargs = dict(optimizer_class=RMSpropTFLike, net_arch=[32]) _ = DQN("MlpPolicy", "CartPole-v1", policy_kwargs=policy_kwargs, learning_starts=100).learn(300) - - -@pytest.mark.parametrize("model_class", [A2C, PPO]) -def test_not_shared_features_extractor(model_class): - policy_kwargs = dict(net_arch=[12, dict(vf=[16], pi=[8])], share_features_extractor=False) - with pytest.raises(ValueError): - model_class("MlpPolicy", "Pendulum-v1", policy_kwargs=policy_kwargs) diff --git a/tests/test_deterministic.py b/tests/test_deterministic.py index 4c92d26..c165e48 100644 --- a/tests/test_deterministic.py +++ b/tests/test_deterministic.py @@ -1,3 +1,4 @@ +import numpy as np import pytest from stable_baselines3 import A2C, DQN, PPO, SAC, TD3 @@ -15,7 +16,9 @@ def test_deterministic_training_common(algo): kwargs = {"policy_kwargs": dict(net_arch=[64])} env_id = "Pendulum-v1" if algo in [TD3, SAC]: - kwargs.update({"action_noise": NormalActionNoise(0.0, 0.1), "learning_starts": 100, "train_freq": 4}) + kwargs.update( + {"action_noise": NormalActionNoise(np.zeros(1), 0.1 * np.ones(1)), "learning_starts": 100, "train_freq": 4} + ) else: if algo == DQN: env_id = "CartPole-v1" diff --git a/tests/test_identity.py b/tests/test_identity.py index f5bbc49..cc7746b 100644 --- a/tests/test_identity.py +++ b/tests/test_identity.py @@ -45,6 +45,8 @@ def test_continuous(model_class): n_actions = 1 action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions)) kwargs["action_noise"] = action_noise + elif model_class in [A2C]: + kwargs["policy_kwargs"]["log_std_init"] = -0.5 model = model_class("MlpPolicy", env, **kwargs).learn(n_steps) diff --git a/tests/test_run.py b/tests/test_run.py index 71236a3..ca7548f 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -115,7 +115,6 @@ def test_dqn(): @pytest.mark.parametrize("train_freq", [4, (4, "step"), (1, "episode")]) def test_train_freq(tmp_path, train_freq): - model = SAC( "MlpPolicy", "Pendulum-v1", diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 2c35e43..9d3d537 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -648,7 +648,6 @@ def test_open_file_str_pathlib(tmp_path, pathtype): def test_open_file(tmp_path): - # path must much the type with pytest.raises(TypeError): open_path(123, None, None, None) diff --git a/tests/test_tensorboard.py b/tests/test_tensorboard.py index 8aa864d..eee0ec0 100644 --- a/tests/test_tensorboard.py +++ b/tests/test_tensorboard.py @@ -1,4 +1,5 @@ import os +from typing import Dict, Union import pytest @@ -18,21 +19,21 @@ N_STEPS = 100 class HParamCallback(BaseCallback): - def __init__(self): - """ - Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard. - """ - super().__init__() + """ + Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard. + """ def _on_training_start(self) -> None: - hparam_dict = { + hparam_dict: Dict[str, Union[str, float]] = { "algorithm": self.model.__class__.__name__, - "learning rate": self.model.learning_rate, - "gamma": self.model.gamma, + # Ignore type checking for gamma, see https://github.com/DLR-RM/stable-baselines3/pull/1194/files#r1035006458 + "gamma": self.model.gamma, # type: ignore[attr-defined] } + if isinstance(self.model.learning_rate, float): # Can also be Schedule, in that case, we don't report + hparam_dict["learning rate"] = self.model.learning_rate # define the metrics that will appear in the `HPARAMS` Tensorboard tab by referencing their tag # Tensorbaord will find & display metrics from the `SCALARS` tab - metric_dict = { + metric_dict: Dict[str, float] = { "rollout/ep_len_mean": 0, } self.logger.record( diff --git a/tests/test_utils.py b/tests/test_utils.py index eecf035..32f44c0 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -9,7 +9,7 @@ from gym import spaces import stable_baselines3 as sb3 from stable_baselines3 import A2C -from stable_baselines3.common.atari_wrappers import ClipRewardEnv, MaxAndSkipEnv +from stable_baselines3.common.atari_wrappers import MaxAndSkipEnv from stable_baselines3.common.env_util import is_wrapped, make_atari_env, make_vec_env, unwrap_wrapper from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.common.monitor import Monitor @@ -55,30 +55,54 @@ def test_make_vec_env_func_checker(): env.close() -@pytest.mark.parametrize("env_id", ["BreakoutNoFrameskip-v4"]) -@pytest.mark.parametrize("n_envs", [1, 2]) -@pytest.mark.parametrize("wrapper_kwargs", [None, dict(clip_reward=False, screen_size=60)]) -def test_make_atari_env(env_id, n_envs, wrapper_kwargs): - env = make_atari_env(env_id, n_envs, wrapper_kwargs=wrapper_kwargs, monitor_dir=None, seed=0) +# Use Asterix as it does not requires fire reset +@pytest.mark.parametrize("env_id", ["BreakoutNoFrameskip-v4", "AsterixNoFrameskip-v4"]) +@pytest.mark.parametrize("noop_max", [0, 10]) +@pytest.mark.parametrize("action_repeat_probability", [0.0, 0.25]) +@pytest.mark.parametrize("frame_skip", [1, 4]) +@pytest.mark.parametrize("screen_size", [60]) +@pytest.mark.parametrize("terminal_on_life_loss", [True, False]) +@pytest.mark.parametrize("clip_reward", [True]) +def test_make_atari_env( + env_id, noop_max, action_repeat_probability, frame_skip, screen_size, terminal_on_life_loss, clip_reward +): + n_envs = 2 + wrapper_kwargs = { + "noop_max": noop_max, + "action_repeat_probability": action_repeat_probability, + "frame_skip": frame_skip, + "screen_size": screen_size, + "terminal_on_life_loss": terminal_on_life_loss, + "clip_reward": clip_reward, + } + venv = make_atari_env( + env_id, + n_envs=2, + wrapper_kwargs=wrapper_kwargs, + monitor_dir=None, + seed=0, + ) - assert env.num_envs == n_envs + assert venv.num_envs == n_envs - obs = env.reset() + needs_fire_reset = env_id == "BreakoutNoFrameskip-v4" + expected_frame_number_low = frame_skip * 2 if needs_fire_reset else 0 # FIRE - UP on reset + expected_frame_number_high = expected_frame_number_low + noop_max + expected_shape = (n_envs, screen_size, screen_size, 1) - new_obs, reward, _, _ = env.step([env.action_space.sample() for _ in range(n_envs)]) + obs = venv.reset() + frame_numbers = [env.unwrapped.ale.getEpisodeFrameNumber() for env in venv.envs] + for frame_number in frame_numbers: + assert expected_frame_number_low <= frame_number <= expected_frame_number_high + assert obs.shape == expected_shape - assert obs.shape == new_obs.shape + new_obs, reward, _, _ = venv.step([venv.action_space.sample() for _ in range(n_envs)]) - # Wrapped into DummyVecEnv - wrapped_atari_env = env.envs[0] - if wrapper_kwargs is not None: - assert obs.shape == (n_envs, 60, 60, 1) - assert wrapped_atari_env.observation_space.shape == (60, 60, 1) - assert not isinstance(wrapped_atari_env.env, ClipRewardEnv) - else: - assert obs.shape == (n_envs, 84, 84, 1) - assert wrapped_atari_env.observation_space.shape == (84, 84, 1) - assert isinstance(wrapped_atari_env.env, ClipRewardEnv) + new_frame_numbers = [env.unwrapped.ale.getEpisodeFrameNumber() for env in venv.envs] + for frame_number, new_frame_number in zip(frame_numbers, new_frame_numbers): + assert new_frame_number - frame_number == frame_skip + assert new_obs.shape == expected_shape + if clip_reward: assert np.max(np.abs(reward)) < 1.0 diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index 7b443c2..fb37fd3 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -1,4 +1,5 @@ import operator +from typing import Any, Dict import gym import numpy as np @@ -20,7 +21,7 @@ ENV_ID = "Pendulum-v1" class DummyRewardEnv(gym.Env): - metadata = {} + metadata: Dict[str, Any] = {} def __init__(self, return_reward_idx=0): self.action_space = spaces.Discrete(2) @@ -177,7 +178,7 @@ def _make_warmstart_dict_env(**kwargs): def test_runningmeanstd(): """Test RunningMeanStd object""" - for (x_1, x_2, x_3) in [ + for x_1, x_2, x_3 in [ (np.random.randn(3), np.random.randn(4), np.random.randn(5)), (np.random.randn(3, 2), np.random.randn(4, 2), np.random.randn(5, 2)), ]: @@ -335,7 +336,6 @@ def test_normalize_dict_selected_keys(): @pytest.mark.parametrize("model_class", [SAC, TD3, HerReplayBuffer]) @pytest.mark.parametrize("online_sampling", [False, True]) def test_offpolicy_normalization(model_class, online_sampling): - if online_sampling and model_class != HerReplayBuffer: pytest.skip() diff --git a/tests/test_vec_stacked_obs.py b/tests/test_vec_stacked_obs.py new file mode 100644 index 0000000..0a7aa39 --- /dev/null +++ b/tests/test_vec_stacked_obs.py @@ -0,0 +1,314 @@ +import numpy as np +from gym import spaces + +from stable_baselines3.common.vec_env.stacked_observations import StackedObservations + +compute_stacking = StackedObservations.compute_stacking +NUM_ENVS = 2 +N_STACK = 4 +H, W, C = 16, 24, 3 + + +def test_compute_stacking_box(): + space = spaces.Box(-1, 1, (4,)) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking(N_STACK, observation_space=space) + assert not channels_first # default is channel last + assert stack_dimension == -1 + assert stacked_shape == (N_STACK * 4,) + assert repeat_axis == -1 + + +def test_compute_stacking_multidim_box(): + space = spaces.Box(-1, 1, (4, 5)) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking(N_STACK, observation_space=space) + assert not channels_first # default is channel last + assert stack_dimension == -1 + assert stacked_shape == (4, N_STACK * 5) + assert repeat_axis == -1 + + +def test_compute_stacking_multidim_box_channel_first(): + space = spaces.Box(-1, 1, (4, 5)) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking( + N_STACK, observation_space=space, channels_order="first" + ) + assert channels_first # default is channel last + assert stack_dimension == 1 + assert stacked_shape == (N_STACK * 4, 5) + assert repeat_axis == 0 + + +def test_compute_stacking_image_channel_first(): + """Detect that image is channel first and stack in that dimension.""" + space = spaces.Box(0, 255, (C, H, W), dtype=np.uint8) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking(N_STACK, observation_space=space) + assert channels_first # default is channel last + assert stack_dimension == 1 + assert stacked_shape == (N_STACK * C, H, W) + assert repeat_axis == 0 + + +def test_compute_stacking_image_channel_last(): + """Detect that image is channel last and stack in that dimension.""" + space = spaces.Box(0, 255, (H, W, C), dtype=np.uint8) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking(N_STACK, observation_space=space) + assert not channels_first # default is channel last + assert stack_dimension == -1 + assert stacked_shape == (H, W, N_STACK * C) + assert repeat_axis == -1 + + +def test_compute_stacking_image_channel_first_stack_last(): + """Detect that image is channel first and stack in that dimension.""" + space = spaces.Box(0, 255, (C, H, W), dtype=np.uint8) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking( + N_STACK, observation_space=space, channels_order="last" + ) + assert not channels_first # default is channel last + assert stack_dimension == -1 + assert stacked_shape == (C, H, N_STACK * W) + assert repeat_axis == -1 + + +def test_compute_stacking_image_channel_last_stack_first(): + """Detect that image is channel last and stack in that dimension.""" + space = spaces.Box(0, 255, (H, W, C), dtype=np.uint8) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking( + N_STACK, observation_space=space, channels_order="first" + ) + assert channels_first # default is channel last + assert stack_dimension == 1 + assert stacked_shape == (N_STACK * H, W, C) + assert repeat_axis == 0 + + +def test_reset_update_box(): + space = spaces.Box(-1, 1, (4,)) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space) + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * 4) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * 4) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate( + (np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=-1 + ), + ) + + +def test_reset_update_multidim_box(): + space = spaces.Box(-1, 1, (4, 5)) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space) + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, 4, N_STACK * 5) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, 4, N_STACK * 5) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate( + (np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=-1 + ), + ) + + +def test_reset_update_multidim_box_channel_first(): + space = spaces.Box(-1, 1, (4, 5)) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space, channels_order="first") + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * 4, 5) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * 4, 5) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate((np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=1), + ) + + +def test_reset_update_image_channel_first(): + space = spaces.Box(0, 255, (C, H, W), dtype=np.uint8) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space) + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * C, H, W) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * C, H, W) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate((np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=1), + ) + + +def test_reset_update_image_channel_last(): + space = spaces.Box(0, 255, (H, W, C), dtype=np.uint8) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space) + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, H, W, N_STACK * C) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, H, W, N_STACK * C) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate( + (np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=-1 + ), + ) + + +def test_reset_update_image_channel_first_stack_last(): + space = spaces.Box(0, 255, (C, H, W), dtype=np.uint8) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space, channels_order="last") + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, C, H, N_STACK * W) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, C, H, N_STACK * W) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate( + (np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=-1 + ), + ) + + +def test_reset_update_image_channel_last_stack_first(): + space = spaces.Box(0, 255, (H, W, C), dtype=np.uint8) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space, channels_order="first") + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * H, W, C) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * H, W, C) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate((np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=1), + ) + + +def test_reset_update_dict(): + space = spaces.Dict({"key1": spaces.Box(0, 255, (H, W, C), dtype=np.uint8), "key2": spaces.Box(-1, 1, (4, 5))}) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space, channels_order={"key1": "first", "key2": "last"}) + observations_1 = {key: np.stack([subspace.sample() for _ in range(NUM_ENVS)]) for key, subspace in space.spaces.items()} + stacked_obs = stacked_observations.reset(observations_1) + assert isinstance(stacked_obs, dict) + assert stacked_obs["key1"].shape == (NUM_ENVS, N_STACK * H, W, C) + assert stacked_obs["key2"].shape == (NUM_ENVS, 4, N_STACK * 5) + assert stacked_obs["key1"].dtype == space["key1"].dtype + assert stacked_obs["key2"].dtype == space["key2"].dtype + observations_2 = {key: np.stack([subspace.sample() for _ in range(NUM_ENVS)]) for key, subspace in space.spaces.items()} + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs["key1"].shape == (NUM_ENVS, N_STACK * H, W, C) + assert stacked_obs["key2"].shape == (NUM_ENVS, 4, N_STACK * 5) + assert stacked_obs["key1"].dtype == space["key1"].dtype + assert stacked_obs["key2"].dtype == space["key2"].dtype + + assert np.array_equal( + stacked_obs["key1"], + np.concatenate( + ( + np.zeros_like(observations_1["key1"]), + np.zeros_like(observations_1["key1"]), + observations_1["key1"], + observations_2["key1"], + ), + axis=1, + ), + ) + assert np.array_equal( + stacked_obs["key2"], + np.concatenate( + ( + np.zeros_like(observations_1["key2"]), + np.zeros_like(observations_1["key2"]), + observations_1["key2"], + observations_2["key2"], + ), + axis=-1, + ), + ) + + +def test_episode_termination_box(): + space = spaces.Box(-1, 1, (4,)) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space) + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_observations.reset(observations_1) + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_observations.update(observations_2, dones, infos) + terminal_observation = space.sample() + infos[1]["terminal_observation"] = terminal_observation # episode termination in env1 + dones[1] = True + observations_3 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs, infos = stacked_observations.update(observations_3, dones, infos) + zeros = np.zeros_like(observations_1[0]) + true_stacked_obs_env1 = np.concatenate((zeros, observations_1[0], observations_2[0], observations_3[0]), axis=-1) + true_stacked_obs_env2 = np.concatenate((zeros, zeros, zeros, observations_3[1]), axis=-1) + true_stacked_obs = np.stack((true_stacked_obs_env1, true_stacked_obs_env2)) + assert np.array_equal(true_stacked_obs, stacked_obs) + + +def test_episode_termination_dict(): + space = spaces.Dict({"key1": spaces.Box(0, 255, (H, W, 3), dtype=np.uint8), "key2": spaces.Box(-1, 1, (4, 5))}) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space, channels_order={"key1": "first", "key2": "last"}) + observations_1 = {key: np.stack([subspace.sample() for _ in range(NUM_ENVS)]) for key, subspace in space.spaces.items()} + stacked_observations.reset(observations_1) + observations_2 = {key: np.stack([subspace.sample() for _ in range(NUM_ENVS)]) for key, subspace in space.spaces.items()} + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_observations.update(observations_2, dones, infos) + terminal_observation = space.sample() + infos[1]["terminal_observation"] = terminal_observation # episode termination in env1 + dones[1] = True + observations_3 = {key: np.stack([subspace.sample() for _ in range(NUM_ENVS)]) for key, subspace in space.spaces.items()} + stacked_obs, infos = stacked_observations.update(observations_3, dones, infos) + + for key, axis in zip(observations_1.keys(), [0, -1]): + zeros = np.zeros_like(observations_1[key][0]) + true_stacked_obs_env1 = np.concatenate( + (zeros, observations_1[key][0], observations_2[key][0], observations_3[key][0]), axis + ) + true_stacked_obs_env2 = np.concatenate((zeros, zeros, zeros, observations_3[key][1]), axis) + true_stacked_obs = np.stack((true_stacked_obs_env1, true_stacked_obs_env2)) + assert np.array_equal(true_stacked_obs, stacked_obs[key])