saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-04 20:23:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge branch 'master' into feat/mps-support

Browse source
This commit is contained in:
Quentin Gallouédec 2023-02-14 10:11:19 +01:00 committed by GitHub
commit b235c8eea0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
48 changed files with 805 additions and 552 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
Makefile
View file

@ -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}

45
docs/guide/custom_policy.rst
View file

@ -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=[<actor network architecture>], vf=[<critic network architecture>])``.
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])``

19
docs/guide/rl_zoo.rst
View file

@ -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

10
docs/guide/tensorboard.rst
View file

@ -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",

6
docs/guide/vec_envs.rst
View file

@ -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
~~~~~~~~~~~~

19
docs/misc/changelog.rst
View file

@ -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

18
docs/modules/a2c.rst
View file

@ -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 <vec_env>`, `Issue #1245 <https://github.com/DLR-RM/stable-baselines3/issues/1245>`_ or the `Multiprocessing notebook <https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/multiprocessing_rl.ipynb>`_.
Results
-------

4
setup.cfg
View file

@ -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]

2
setup.py
View file

@ -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",
],

5
stable_baselines3/a2c/a2c.py
View file

@ -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,

92
stable_baselines3/common/atari_wrappers.py
View file

@ -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():

8
stable_baselines3/common/base_class.py
View file

@ -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:

36
stable_baselines3/common/buffers.py
View file

@ -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]),

12
stable_baselines3/common/callbacks.py
View file

@ -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:

8
stable_baselines3/common/envs/identity_env.py
View file

@ -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)

2
stable_baselines3/common/envs/multi_input_envs.py
View file

@ -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()`

1
stable_baselines3/common/evaluation.py
View file

@ -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]

15
stable_baselines3/common/logger.py
View file

@ -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

19
stable_baselines3/common/monitor.py
View file

@ -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

44
stable_baselines3/common/noise.py
View file

@ -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

2
stable_baselines3/common/off_policy_algorithm.py
View file

@ -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,

2
stable_baselines3/common/on_policy_algorithm.py
View file

@ -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:

38
stable_baselines3/common/policies.py
View file

@ -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,

2
stable_baselines3/common/results_plotter.py
View file

@ -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:

4
stable_baselines3/common/running_mean_std.py
View file

@ -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

2
stable_baselines3/common/save_util.py
View file

@ -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

109
stable_baselines3/common/torch_layers.py
View file

@ -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=[<value layer sizes>], pi=[<policy layer sizes>])``.
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=[<list of layer sizes>], pi=[<list of layer sizes>])``: 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. ``[<list of layer sizes>]``: "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):

2
stable_baselines3/common/utils.py
View file

@ -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.

3
stable_baselines3/common/vec_env/__init__.py
View file

@ -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",

317
stable_baselines3/common/vec_env/stacked_observations.py
View file

@ -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

48
stable_baselines3/common/vec_env/vec_frame_stack.py
View file

@ -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()

1
stable_baselines3/common/vec_env/vec_video_recorder.py
View file

@ -30,7 +30,6 @@ class VecVideoRecorder(VecEnvWrapper):
video_length: int = 200,
name_prefix: str = "rl-video",
):
VecEnvWrapper.__init__(self, venv)
self.env = venv

2
stable_baselines3/ddpg/ddpg.py
View file

@ -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,

2
stable_baselines3/dqn/dqn.py
View file

@ -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,

2
stable_baselines3/her/her_replay_buffer.py
View file

@ -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)

5
stable_baselines3/ppo/ppo.py
View file

@ -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,

2
stable_baselines3/sac/sac.py
View file

@ -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,

3
stable_baselines3/td3/td3.py
View file

@ -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,

2
stable_baselines3/version.txt
View file

@ -1 +1 @@
1.8.0a0
1.8.0a4

25
tests/test_custom_policy.py
View file

@ -9,21 +9,21 @@ from stable_baselines3.common.sb2_compat.rmsprop_tf_like import RMSpropTFLike
"net_arch",
[
[],
dict(vf=[16], pi=[8]),
# [<layer_sizes>] 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)

5
tests/test_deterministic.py
View file

@ -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"

2
tests/test_identity.py
View file

@ -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)

1
tests/test_run.py
View file

@ -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",

1
tests/test_save_load.py
View file

@ -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)

19
tests/test_tensorboard.py
View file

@ -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(

64
tests/test_utils.py
View file

@ -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

6
tests/test_vec_normalize.py
View file

@ -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()

314
tests/test_vec_stacked_obs.py Normal file
View file

@ -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])