diff --git a/docs/_static/img/split_graph.png b/docs/_static/img/split_graph.png new file mode 100644 index 0000000..c966c55 Binary files /dev/null and b/docs/_static/img/split_graph.png differ diff --git a/docs/conda_env.yml b/docs/conda_env.yml index a01d37b..98a5508 100644 --- a/docs/conda_env.yml +++ b/docs/conda_env.yml @@ -6,9 +6,9 @@ dependencies: - cpuonly=1.0=0 - pip=21.1 - python=3.7 - - pytorch=1.8.1=py3.7_cpu_0 + - pytorch=1.11=py3.7_cpu_0 - pip: - - gym>=0.17.2 + - gym==0.21 - cloudpickle - opencv-python-headless - pandas @@ -16,5 +16,5 @@ dependencies: - matplotlib - sphinx_autodoc_typehints - sphinx>=4.2 - # See https://github.com/readthedocs/sphinx_rtd_theme/issues/1115 - sphinx_rtd_theme>=1.0 + - sphinx_copybutton diff --git a/docs/conf.py b/docs/conf.py index 18898d5..b44be6f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -24,6 +24,14 @@ try: except ImportError: enable_spell_check = False +# Try to enable copy button +try: + import sphinx_copybutton # noqa: F401 + + enable_copy_button = True +except ImportError: + enable_copy_button = False + # source code directory, relative to this file, for sphinx-autobuild sys.path.insert(0, os.path.abspath("..")) @@ -51,7 +59,7 @@ with open(version_file) as file_handler: # -- Project information ----------------------------------------------------- project = "Stable Baselines3" -copyright = "2020, Stable Baselines3" +copyright = "2022, Stable Baselines3" author = "Stable Baselines3 Contributors" # The short X.Y version @@ -83,6 +91,9 @@ extensions = [ if enable_spell_check: extensions.append("sphinxcontrib.spelling") +if enable_copy_button: + extensions.append("sphinx_copybutton") + # Add any paths that contain templates here, relative to this directory. templates_path = ["_templates"] diff --git a/docs/guide/examples.rst b/docs/guide/examples.rst index a5b56b2..0d7e7c0 100644 --- a/docs/guide/examples.rst +++ b/docs/guide/examples.rst @@ -729,6 +729,16 @@ to keep track of the agent progress. model.learn(10_000) +SB3 with EnvPool or Isaac Gym +----------------------------- + +Just like Procgen (see above), `EnvPool `_ and `Isaac Gym `_ accelerate the environment by +already providing a vectorized implementation. + +To use SB3 with those tools, you must wrap the env with tool's specific ``VecEnvWrapper`` that will pre-process the data for SB3, +you can find links to those wrappers in `issue #772 `_. + + Record a Video -------------- diff --git a/docs/guide/install.rst b/docs/guide/install.rst index 3b26927..0169495 100644 --- a/docs/guide/install.rst +++ b/docs/guide/install.rst @@ -54,6 +54,17 @@ Bleeding-edge version pip install git+https://github.com/DLR-RM/stable-baselines3 +.. note:: + + If you want to use latest gym version (0.24+), you have to use + + .. code-block:: bash + + pip install git+https://github.com/carlosluis/stable-baselines3@fix_tests + + See `PR #780 `_ for more information. + + Development version ------------------- diff --git a/docs/guide/migration.rst b/docs/guide/migration.rst index 879a5fb..ef26870 100644 --- a/docs/guide/migration.rst +++ b/docs/guide/migration.rst @@ -141,7 +141,7 @@ DQN ^^^ Only the vanilla DQN is implemented right now but extensions will follow. -Default hyperparameters are taken from the nature paper, except for the optimizer and learning rate that were taken from Stable Baselines defaults. +Default hyperparameters are taken from the Nature paper, except for the optimizer and learning rate that were taken from Stable Baselines defaults. DDPG ^^^^ diff --git a/docs/guide/tensorboard.rst b/docs/guide/tensorboard.rst index dc39d8b..625c1be 100644 --- a/docs/guide/tensorboard.rst +++ b/docs/guide/tensorboard.rst @@ -26,10 +26,20 @@ You can also define custom logging name when training (by default it is the algo model.learn(total_timesteps=10_000, tb_log_name="first_run") # Pass reset_num_timesteps=False to continue the training curve in tensorboard # By default, it will create a new curve + # Keep tb_log_name constant to have continuous curve (see note below) model.learn(total_timesteps=10_000, tb_log_name="second_run", reset_num_timesteps=False) model.learn(total_timesteps=10_000, tb_log_name="third_run", reset_num_timesteps=False) +.. note:: + If you specify different ``tb_log_name`` in subsequent runs, you will have split graphs, like in the figure below. + If you want them to be continuous, you must keep the same ``tb_log_name`` (see `issue #975 `_). + And, if you still managed to get your graphs split by other means, just put tensorboard log files into the same folder. + + .. image:: ../_static/img/split_graph.png + :width: 330 + :alt: split_graph + Once the learn function is called, you can monitor the RL agent during or after the training, with the following bash command: .. code-block:: bash diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 576fbc8..2b93ebc 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,10 +3,44 @@ Changelog ========== - -Release 1.5.1a10 (WIP) +Release 1.6.1a1 (WIP) --------------------------- +Breaking Changes: +^^^^^^^^^^^^^^^^^ + +New Features: +^^^^^^^^^^^^^ +- Use MacOS Metal "mps" device when available + +SB3-Contrib +^^^^^^^^^^^ + +Bug Fixes: +^^^^^^^^^^ +- Fixed the issue that ``predict`` does not always return action as ``np.ndarray`` (@qgallouedec) +- Fixed division by zero error when computing FPS when a small number of time has elapsed in operating systems with low-precision timers. +- Added multidimensional action space support (@qgallouedec) + +Deprecations: +^^^^^^^^^^^^^ + +Others: +^^^^^^^ + +Documentation: +^^^^^^^^^^^^^^ +- Fixed typo in docstring "nature" -> "Nature" (@Melanol) +- Added info on split tensorboard logs into (@Melanol) +- Fixed typo in ppo doc (@francescoluciano) +- Fixed typo in install doc(@jlp-ue) + + +Release 1.6.0 (2022-07-11) +--------------------------- + +**Recurrent PPO (PPO LSTM), better defaults for learning from pixels with SAC/TD3** + Breaking Changes: ^^^^^^^^^^^^^^^^^ - Changed the way policy "aliases" are handled ("MlpPolicy", "CnnPolicy", ...), removing the former @@ -17,7 +51,6 @@ Breaking Changes: New Features: ^^^^^^^^^^^^^ -- Use MacOS Metal "mps" device when available - Save cloudpickle version SB3-Contrib @@ -35,6 +68,8 @@ Bug Fixes: - Added a check for unbounded actions - Fixed issues due to newer version of protobuf (tensorboard) and sphinx - Fix exception causes all over the codebase (@cool-RR) +- Prohibit simultaneous use of optimize_memory_usage and handle_timeout_termination due to a bug (@MWeltevrede) +- Fixed a bug in ``kl_divergence`` check that would fail when using numpy arrays with MultiCategorical distribution Deprecations: ^^^^^^^^^^^^^ @@ -52,6 +87,8 @@ Documentation: - Added remark about breaking Markov assumption and timeout handling - Added doc about MLFlow integration via custom logger (@git-thor) - Updated Huggingface integration doc +- Added copy button for code snippets +- Added doc about EnvPool and Isaac Gym support Release 1.5.0 (2022-03-25) @@ -981,4 +1018,5 @@ And all the contributors: @wkirgsn @AechPro @CUN-bjy @batu @IljaAvadiev @timokau @kachayev @cleversonahum @eleurent @ac-93 @cove9988 @theDebugger811 @hsuehch @Demetrio92 @thomasgubler @IperGiove @ScheiklP @simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola @buoyancy99 @caburu @xy9485 -@Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR +@Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede +@Melanol @qgallouedec @francescoluciano @jlp-ue diff --git a/docs/modules/ppo.rst b/docs/modules/ppo.rst index a562aa2..d0c425f 100644 --- a/docs/modules/ppo.rst +++ b/docs/modules/ppo.rst @@ -8,7 +8,7 @@ PPO The `Proximal Policy Optimization `_ algorithm combines ideas from A2C (having multiple workers) and TRPO (it uses a trust region to improve the actor). -The main idea is that after an update, the new policy should be not too far form the old policy. +The main idea is that after an update, the new policy should be not too far from the old policy. For that, ppo uses clipping to avoid too large update. @@ -25,7 +25,7 @@ Notes - Clear explanation of PPO on Arxiv Insights channel: https://www.youtube.com/watch?v=5P7I-xPq8u8 - OpenAI blog post: https://blog.openai.com/openai-baselines-ppo/ - Spinning Up guide: https://spinningup.openai.com/en/latest/algorithms/ppo.html -- 37 implementation details blog: https://ppo-details.cleanrl.dev//2021/11/05/ppo-implementation-details/ +- 37 implementation details blog: https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/ Can I use? diff --git a/setup.py b/setup.py index 05745e9..2816316 100644 --- a/setup.py +++ b/setup.py @@ -111,6 +111,8 @@ setup( "sphinxcontrib.spelling", # Type hints support "sphinx-autodoc-typehints", + # Copy button for code snippets + "sphinx_copybutton", ], "extra": [ # For render diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 36a73b8..e8032e7 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -422,7 +422,7 @@ class BaseAlgorithm(ABC): :param tb_log_name: the name of the run for tensorboard log :return: """ - self.start_time = time.time() + self.start_time = time.time_ns() if self.ep_info_buffer is None or reset_num_timesteps: # Initialize buffers if they don't exist, or reinitialize if resetting counters diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index d7728cb..0eb2651 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -164,6 +164,7 @@ class ReplayBuffer(BaseBuffer): at a cost of more complexity. See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195 and https://github.com/DLR-RM/stable-baselines3/pull/28#issuecomment-637559274 + Cannot be used in combination with handle_timeout_termination. :param handle_timeout_termination: Handle timeout termination (due to timelimit) separately and treat the task as infinite horizon task. https://github.com/DLR-RM/stable-baselines3/issues/284 @@ -188,6 +189,13 @@ class ReplayBuffer(BaseBuffer): if psutil is not None: mem_available = psutil.virtual_memory().available + # there is a bug if both optimize_memory_usage and handle_timeout_termination are true + # see https://github.com/DLR-RM/stable-baselines3/issues/934 + if optimize_memory_usage and handle_timeout_termination: + raise ValueError( + "ReplayBuffer does not support optimize_memory_usage = True " + "and handle_timeout_termination = True simultaneously." + ) self.optimize_memory_usage = optimize_memory_usage self.observations = np.zeros((self.buffer_size, self.n_envs) + self.obs_shape, dtype=observation_space.dtype) @@ -239,8 +247,7 @@ class ReplayBuffer(BaseBuffer): next_obs = next_obs.reshape((self.n_envs,) + self.obs_shape) # Same, for actions - if isinstance(self.action_space, spaces.Discrete): - action = action.reshape((self.n_envs, self.action_dim)) + action = action.reshape((self.n_envs, self.action_dim)) # Copy to avoid modification by reference self.observations[self.pos] = np.array(obs).copy() @@ -425,6 +432,9 @@ class RolloutBuffer(BaseBuffer): if isinstance(self.observation_space, spaces.Discrete): obs = obs.reshape((self.n_envs,) + self.obs_shape) + # Same reshape, for actions + action = action.reshape((self.n_envs, self.action_dim)) + self.observations[self.pos] = np.array(obs).copy() self.actions[self.pos] = np.array(action).copy() self.rewards[self.pos] = np.array(reward).copy() @@ -578,8 +588,7 @@ class DictReplayBuffer(ReplayBuffer): self.next_observations[key][self.pos] = np.array(next_obs[key]).copy() # Same reshape, for actions - if isinstance(self.action_space, spaces.Discrete): - action = action.reshape((self.n_envs, self.action_dim)) + action = action.reshape((self.n_envs, self.action_dim)) self.actions[self.pos] = np.array(action).copy() self.rewards[self.pos] = np.array(reward).copy() diff --git a/stable_baselines3/common/distributions.py b/stable_baselines3/common/distributions.py index 3d1ff5a..fc48625 100644 --- a/stable_baselines3/common/distributions.py +++ b/stable_baselines3/common/distributions.py @@ -4,6 +4,7 @@ from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, Tuple, Union import gym +import numpy as np import torch as th from gym import spaces from torch import nn @@ -577,10 +578,10 @@ class StateDependentNoiseDistribution(Distribution): return th.mm(latent_sde, self.exploration_mat) # Use batch matrix multiplication for efficient computation # (batch_size, n_features) -> (batch_size, 1, n_features) - latent_sde = latent_sde.unsqueeze(1) + latent_sde = latent_sde.unsqueeze(dim=1) # (batch_size, 1, n_actions) noise = th.bmm(latent_sde, self.exploration_matrices) - return noise.squeeze(1) + return noise.squeeze(dim=1) def actions_from_params( self, mean_actions: th.Tensor, log_std: th.Tensor, latent_sde: th.Tensor, deterministic: bool = False @@ -657,7 +658,6 @@ def make_proba_distribution( dist_kwargs = {} if isinstance(action_space, spaces.Box): - assert len(action_space.shape) == 1, "Error: the action space must be a vector" cls = StateDependentNoiseDistribution if use_sde else DiagGaussianDistribution return cls(get_action_dim(action_space), **dist_kwargs) elif isinstance(action_space, spaces.Discrete): @@ -688,7 +688,7 @@ def kl_divergence(dist_true: Distribution, dist_pred: Distribution) -> th.Tensor # MultiCategoricalDistribution is not a PyTorch Distribution subclass # so we need to implement it ourselves! if isinstance(dist_pred, MultiCategoricalDistribution): - assert dist_pred.action_dims == dist_true.action_dims, "Error: distributions must have the same input space" + assert np.allclose(dist_pred.action_dims, dist_true.action_dims), "Error: distributions must have the same input space" return th.stack( [th.distributions.kl_divergence(p, q) for p, q in zip(dist_true.distribution, dist_pred.distribution)], dim=1, diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 99a02ff..b841eb0 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -1,5 +1,6 @@ import io import pathlib +import sys import time import warnings from copy import deepcopy @@ -427,8 +428,8 @@ class OffPolicyAlgorithm(BaseAlgorithm): """ Write log. """ - time_elapsed = time.time() - self.start_time - fps = int((self.num_timesteps - self._num_timesteps_at_start) / (time_elapsed + 1e-8)) + time_elapsed = max((time.time_ns() - self.start_time) / 1e9, sys.float_info.epsilon) + fps = int((self.num_timesteps - self._num_timesteps_at_start) / time_elapsed) self.logger.record("time/episodes", self._episode_num, exclude="tensorboard") if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0: self.logger.record("rollout/ep_rew_mean", safe_mean([ep_info["r"] for ep_info in self.ep_info_buffer])) diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 763c108..84c89d9 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -1,3 +1,4 @@ +import sys import time from typing import Any, Dict, List, Optional, Tuple, Type, Union @@ -254,13 +255,14 @@ class OnPolicyAlgorithm(BaseAlgorithm): # Display training infos if log_interval is not None and iteration % log_interval == 0: - fps = int((self.num_timesteps - self._num_timesteps_at_start) / (time.time() - self.start_time)) + time_elapsed = max((time.time_ns() - self.start_time) / 1e9, sys.float_info.epsilon) + fps = int((self.num_timesteps - self._num_timesteps_at_start) / time_elapsed) self.logger.record("time/iterations", iteration, exclude="tensorboard") if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0: self.logger.record("rollout/ep_rew_mean", safe_mean([ep_info["r"] for ep_info in self.ep_info_buffer])) self.logger.record("rollout/ep_len_mean", safe_mean([ep_info["l"] for ep_info in self.ep_info_buffer])) self.logger.record("time/fps", fps) - self.logger.record("time/time_elapsed", int(time.time() - self.start_time), exclude="tensorboard") + self.logger.record("time/time_elapsed", int(time_elapsed), exclude="tensorboard") self.logger.record("time/total_timesteps", self.num_timesteps, exclude="tensorboard") self.logger.dump(step=self.num_timesteps) diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 51a3d37..3809b61 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -336,8 +336,8 @@ class BasePolicy(BaseModel): with th.no_grad(): actions = self._predict(observation, deterministic=deterministic) - # Convert to numpy - actions = actions.cpu().numpy() + # Convert to numpy, and reshape to the original action shape + actions = actions.cpu().numpy().reshape((-1,) + self.action_space.shape) if isinstance(self.action_space, gym.spaces.Box): if self.squash_output: @@ -350,7 +350,7 @@ class BasePolicy(BaseModel): # Remove batch dimension if needed if not vectorized_env: - actions = actions[0] + actions = actions.squeeze(axis=0) return actions, state @@ -592,6 +592,7 @@ class ActorCriticPolicy(BasePolicy): distribution = self._get_action_dist_from_latent(latent_pi) actions = distribution.get_actions(deterministic=deterministic) log_prob = distribution.log_prob(actions) + actions = actions.reshape((-1,) + self.action_space.shape) return actions, values, log_prob def _get_action_dist_from_latent(self, latent_pi: th.Tensor) -> Distribution: diff --git a/stable_baselines3/common/torch_layers.py b/stable_baselines3/common/torch_layers.py index 8fd2237..f87337c 100644 --- a/stable_baselines3/common/torch_layers.py +++ b/stable_baselines3/common/torch_layers.py @@ -50,7 +50,7 @@ class FlattenExtractor(BaseFeaturesExtractor): class NatureCNN(BaseFeaturesExtractor): """ - CNN from DQN nature paper: + CNN from DQN Nature paper: Mnih, Volodymyr, et al. "Human-level control through deep reinforcement learning." Nature 518.7540 (2015): 529-533. diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index fe8f398..0cd6dfb 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -20,7 +20,7 @@ class DQN(OffPolicyAlgorithm): Deep Q-Network (DQN) Paper: https://arxiv.org/abs/1312.5602, https://www.nature.com/articles/nature14236 - Default hyperparameters are taken from the nature paper, + Default hyperparameters are taken from the Nature paper, except for the optimizer and learning rate that were taken from Stable Baselines defaults. :param policy: The policy model to use (MlpPolicy, CnnPolicy, ...) diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index c43063f..e36b727 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.5.1a10 +1.6.1a1 diff --git a/tests/test_distributions.py b/tests/test_distributions.py index 3652b18..07920db 100644 --- a/tests/test_distributions.py +++ b/tests/test_distributions.py @@ -163,7 +163,9 @@ def test_categorical(dist, CAT_ACTIONS): BernoulliDistribution(N_ACTIONS).proba_distribution(th.rand(N_ACTIONS)), CategoricalDistribution(N_ACTIONS).proba_distribution(th.rand(N_ACTIONS)), DiagGaussianDistribution(N_ACTIONS).proba_distribution(th.rand(N_ACTIONS), th.rand(N_ACTIONS)), - MultiCategoricalDistribution([N_ACTIONS, N_ACTIONS]).proba_distribution(th.rand(1, sum([N_ACTIONS, N_ACTIONS]))), + MultiCategoricalDistribution(np.array([N_ACTIONS, N_ACTIONS])).proba_distribution( + th.rand(1, sum([N_ACTIONS, N_ACTIONS])) + ), SquashedDiagGaussianDistribution(N_ACTIONS).proba_distribution(th.rand(N_ACTIONS), th.rand(N_ACTIONS)), StateDependentNoiseDistribution(N_ACTIONS).proba_distribution( th.rand(N_ACTIONS), th.rand([N_ACTIONS, N_ACTIONS]), th.rand([N_ACTIONS, N_ACTIONS]) diff --git a/tests/test_logger.py b/tests/test_logger.py index 6fe536f..a55f88a 100644 --- a/tests/test_logger.py +++ b/tests/test_logger.py @@ -1,6 +1,7 @@ import os import time from typing import Sequence +from unittest import mock import gym import numpy as np @@ -381,3 +382,16 @@ def test_fps_logger(tmp_path, algo): # third time, FPS should be the same model.learn(100, log_interval=1, reset_num_timesteps=False) assert max_fps / 10 <= logger.name_to_value["time/fps"] <= max_fps + + +@pytest.mark.parametrize("algo", [A2C, DQN]) +def test_fps_no_div_zero(algo): + """Set time to constant and train algorithm to check no division by zero error. + + Time can appear to be constant during short runs on platforms with low-precision + timers. We should avoid division by zero errors e.g. when computing FPS in + this situation.""" + with mock.patch("time.time", lambda: 42.0): + with mock.patch("time.time_ns", lambda: 42.0): + model = algo("MlpPolicy", "CartPole-v1") + model.learn(total_timesteps=100) diff --git a/tests/test_predict.py b/tests/test_predict.py index 853f4d1..89cdb09 100644 --- a/tests/test_predict.py +++ b/tests/test_predict.py @@ -73,11 +73,13 @@ def test_predict(model_class, env_id, device): obs = env.reset() action, _ = model.predict(obs) + assert isinstance(action, np.ndarray) assert action.shape == env.action_space.shape assert env.action_space.contains(action) vec_env_obs = vec_env.reset() action, _ = model.predict(vec_env_obs) + assert isinstance(action, np.ndarray) assert action.shape[0] == vec_env_obs.shape[0] # Special case for DQN to check the epsilon greedy exploration diff --git a/tests/test_run.py b/tests/test_run.py index e4e8a2e..b0a9a11 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -10,7 +10,10 @@ normal_action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1)) @pytest.mark.parametrize("model_class", [TD3, DDPG]) -@pytest.mark.parametrize("action_noise", [normal_action_noise, OrnsteinUhlenbeckActionNoise(np.zeros(1), 0.1 * np.ones(1))]) +@pytest.mark.parametrize( + "action_noise", + [normal_action_noise, OrnsteinUhlenbeckActionNoise(np.zeros(1), 0.1 * np.ones(1))], +) def test_deterministic_pg(model_class, action_noise): """ Test for DDPG and variants (TD3). diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 2fdebbe..d7a74c5 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -375,6 +375,9 @@ def test_warn_buffer(recwarn, model_class, optimize_memory_usage): select_env(model_class), buffer_size=100, optimize_memory_usage=optimize_memory_usage, + # we cannot use optimize_memory_usage and handle_timeout_termination + # at the same time + replay_buffer_kwargs={"handle_timeout_termination": not optimize_memory_usage}, policy_kwargs=dict(net_arch=[64]), learning_starts=10, ) diff --git a/tests/test_spaces.py b/tests/test_spaces.py index b754042..0696492 100644 --- a/tests/test_spaces.py +++ b/tests/test_spaces.py @@ -33,6 +33,19 @@ class DummyMultiBinary(gym.Env): return self.observation_space.sample(), 0.0, False, {} +class DummyMultidimensionalAction(gym.Env): + def __init__(self): + super().__init__() + self.observation_space = gym.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32) + self.action_space = gym.spaces.Box(low=-1, high=1, shape=(2, 2), dtype=np.float32) + + def reset(self): + return self.observation_space.sample() + + def step(self, action): + return self.observation_space.sample(), 0.0, False, {} + + @pytest.mark.parametrize("model_class", [SAC, TD3, DQN]) @pytest.mark.parametrize("env", [DummyMultiDiscreteSpace([4, 3]), DummyMultiBinary(8)]) def test_identity_spaces(model_class, env): @@ -53,22 +66,39 @@ def test_identity_spaces(model_class, env): @pytest.mark.parametrize("model_class", [A2C, DDPG, DQN, PPO, SAC, TD3]) -@pytest.mark.parametrize("env", ["Pendulum-v1", "CartPole-v1"]) +@pytest.mark.parametrize("env", ["Pendulum-v1", "CartPole-v1", DummyMultidimensionalAction()]) def test_action_spaces(model_class, env): + kwargs = {} if model_class in [SAC, DDPG, TD3]: - supported_action_space = env == "Pendulum-v1" + supported_action_space = env == "Pendulum-v1" or isinstance(env, DummyMultidimensionalAction) + kwargs["learning_starts"] = 2 + kwargs["train_freq"] = 32 elif model_class == DQN: supported_action_space = env == "CartPole-v1" elif model_class in [A2C, PPO]: supported_action_space = True + kwargs["n_steps"] = 64 if supported_action_space: - model_class("MlpPolicy", env) + model = model_class("MlpPolicy", env, **kwargs) + if isinstance(env, DummyMultidimensionalAction): + model.learn(64) else: with pytest.raises(AssertionError): model_class("MlpPolicy", env) +def test_sde_multi_dim(): + SAC( + "MlpPolicy", + DummyMultidimensionalAction(), + learning_starts=10, + use_sde=True, + sde_sample_freq=2, + use_sde_at_warmup=True, + ).learn(20) + + @pytest.mark.parametrize("model_class", [A2C, PPO, DQN]) @pytest.mark.parametrize("env", ["Taxi-v3"]) def test_discrete_obs_space(model_class, env):