mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-07-30 20:18:15 +00:00
Merge branch 'master' into feat/mps-support
This commit is contained in:
Quentin Gallouédec
GitHub
commit
7f11843afd
50 changed files with 319 additions and 139 deletions
3
.github/workflows/ci.yml
vendored
3
.github/workflows/ci.yml
vendored
|
|
@ -11,6 +11,9 @@ on:
|
|||
|
||||
jobs:
|
||||
build:
|
||||
env:
|
||||
TERM: xterm-256color
|
||||
FORCE_COLOR: 1
|
||||
# Skip CI if [ci skip] in the commit message
|
||||
if: "! contains(toJSON(github.event.commits.*.message), '[ci skip]')"
|
||||
runs-on: ubuntu-latest
|
||||
|
|
|
|||
2
Makefile
2
Makefile
|
|
@ -8,7 +8,7 @@ pytype:
|
|||
pytype -j auto
|
||||
|
||||
mypy:
|
||||
MYPY_FORCE_COLOR=1 mypy ${LINT_PATHS}
|
||||
mypy ${LINT_PATHS}
|
||||
|
||||
type: pytype mypy
|
||||
|
||||
|
|
|
|||
14
README.md
14
README.md
|
|
@ -126,13 +126,15 @@ env = gym.make("CartPole-v1")
|
|||
model = PPO("MlpPolicy", env, verbose=1)
|
||||
model.learn(total_timesteps=10_000)
|
||||
|
||||
obs = env.reset()
|
||||
vec_env = model.get_env()
|
||||
obs = vec_env.reset()
|
||||
for i in range(1000):
|
||||
action, _states = model.predict(obs, deterministic=True)
|
||||
obs, reward, done, info = env.step(action)
|
||||
env.render()
|
||||
if done:
|
||||
obs = env.reset()
|
||||
obs, reward, done, info = vec_env.step(action)
|
||||
vec_env.render()
|
||||
# VecEnv resets automatically
|
||||
# if done:
|
||||
# obs = env.reset()
|
||||
|
||||
env.close()
|
||||
```
|
||||
|
|
@ -247,7 +249,7 @@ If you want to contribute, please read [**CONTRIBUTING.md**](./CONTRIBUTING.md)
|
|||
|
||||
## Acknowledgments
|
||||
|
||||
The initial work to develop Stable Baselines3 was partially funded by the project *Reduced Complexity Models* from the *Helmholtz-Gemeinschaft Deutscher Forschungszentren*.
|
||||
The initial work to develop Stable Baselines3 was partially funded by the project *Reduced Complexity Models* from the *Helmholtz-Gemeinschaft Deutscher Forschungszentren*, and by the EU's Horizon 2020 Research and Innovation Programme under grant number 951992 ([VeriDream](https://www.veridream.eu/)).
|
||||
|
||||
The original version, Stable Baselines, was created in the [robotics lab U2IS](http://u2is.ensta-paristech.fr/index.php?lang=en) ([INRIA Flowers](https://flowers.inria.fr/) team) at [ENSTA ParisTech](http://www.ensta-paristech.fr/en).
|
||||
|
||||
|
|
|
|||
|
|
@ -333,11 +333,11 @@ If your task requires even more granular control over the policy/value architect
|
|||
:return: (th.Tensor, th.Tensor) latent_policy, latent_value of the specified network.
|
||||
If all layers are shared, then ``latent_policy == latent_value``
|
||||
"""
|
||||
return self.policy_net(features), self.value_net(features)
|
||||
|
||||
return self.forward_actor(features), self.forward_critic(features)
|
||||
|
||||
def forward_actor(self, features: th.Tensor) -> th.Tensor:
|
||||
return self.policy_net(features)
|
||||
|
||||
|
||||
def forward_critic(self, features: th.Tensor) -> th.Tensor:
|
||||
return self.value_net(features)
|
||||
|
||||
|
|
|
|||
|
|
@ -94,11 +94,12 @@ In the following example, we will train, save and load a DQN model on the Lunar
|
|||
mean_reward, std_reward = evaluate_policy(model, model.get_env(), n_eval_episodes=10)
|
||||
|
||||
# Enjoy trained agent
|
||||
obs = env.reset()
|
||||
vec_env = model.get_env()
|
||||
obs = vec_env.reset()
|
||||
for i in range(1000):
|
||||
action, _states = model.predict(obs, deterministic=True)
|
||||
obs, rewards, dones, info = env.step(action)
|
||||
env.render()
|
||||
obs, rewards, dones, info = vec_env.step(action)
|
||||
vec_env.render()
|
||||
|
||||
|
||||
Multiprocessing: Unleashing the Power of Vectorized Environments
|
||||
|
|
|
|||
|
|
@ -19,13 +19,15 @@ Here is a quick example of how to train and run A2C on a CartPole environment:
|
|||
model = A2C("MlpPolicy", env, verbose=1)
|
||||
model.learn(total_timesteps=10_000)
|
||||
|
||||
obs = env.reset()
|
||||
vec_env = model.get_env()
|
||||
obs = vec_env.reset()
|
||||
for i in range(1000):
|
||||
action, _state = model.predict(obs, deterministic=True)
|
||||
obs, reward, done, info = env.step(action)
|
||||
env.render()
|
||||
if done:
|
||||
obs = env.reset()
|
||||
obs, reward, done, info = vec_env.step(action)
|
||||
vec_env.render()
|
||||
# VecEnv resets automatically
|
||||
# if done:
|
||||
# obs = vec_env.reset()
|
||||
|
||||
.. note::
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ Changelog
|
|||
==========
|
||||
|
||||
|
||||
Release 1.7.0a1 (WIP)
|
||||
Release 1.7.0a5 (WIP)
|
||||
--------------------------
|
||||
|
||||
Breaking Changes:
|
||||
|
|
@ -17,6 +17,7 @@ Breaking Changes:
|
|||
New Features:
|
||||
^^^^^^^^^^^^^
|
||||
- Introduced mypy type checking
|
||||
- Added ``with_bias`` argument to ``create_mlp``
|
||||
|
||||
SB3-Contrib
|
||||
^^^^^^^^^^^
|
||||
|
|
@ -26,7 +27,9 @@ Bug Fixes:
|
|||
- Fix return type of ``evaluate_actions`` in ``ActorCritcPolicy`` to reflect that entropy is an optional tensor (@Rocamonde)
|
||||
- Fix type annotation of ``policy`` in ``BaseAlgorithm`` and ``OffPolicyAlgorithm``
|
||||
- Allowed model trained with Python 3.7 to be loaded with Python 3.8+ without the ``custom_objects`` workaround
|
||||
- Raise an error when the same gym environment instance is passed as separate environments when creating a vectorized environment with more than one environment. (@Rocamonde)
|
||||
- Fix type annotation of ``model`` in ``evaluate_policy``
|
||||
- Fixed ``Self`` return type using ``TypeVar``
|
||||
|
||||
Deprecations:
|
||||
^^^^^^^^^^^^^
|
||||
|
|
@ -34,11 +37,20 @@ Deprecations:
|
|||
Others:
|
||||
^^^^^^^
|
||||
- Used issue forms instead of issue templates
|
||||
- Fixed flake8 config to be compatible with flake8 6+
|
||||
- Goal-conditioned environments are now characterized by the availability of the ``compute_reward`` method, rather than by their inheritance to ``gym.GoalEnv``
|
||||
- Replaced ``CartPole-v0`` by ``CartPole-v1`` is tests
|
||||
- Fixed ``tests/test_distributions.py`` type hint
|
||||
- Fixed ``stable_baselines3/common/type_aliases.py`` type hint
|
||||
- Fixed ``stable_baselines3/common/torch_layers.py`` type hint
|
||||
- Fixed ``stable_baselines3/common/env_util.py`` type hint
|
||||
- Exposed modules in ``__init__.py`` with the ``__all__`` attribute (@ZikangXiong)
|
||||
|
||||
Documentation:
|
||||
^^^^^^^^^^^^^^
|
||||
- Updated Hugging Face Integration page (@simoninithomas)
|
||||
|
||||
- Changed ``env`` to ``vec_env`` when environment is vectorized
|
||||
- Update custom policy documentation (@athatheo)
|
||||
|
||||
Release 1.6.2 (2022-10-10)
|
||||
--------------------------
|
||||
|
|
@ -74,7 +86,6 @@ Documentation:
|
|||
^^^^^^^^^^^^^^
|
||||
- Extended docstring of the ``wrapper_class`` parameter in ``make_vec_env`` (@AlexPasqua)
|
||||
|
||||
|
||||
Release 1.6.1 (2022-09-29)
|
||||
---------------------------
|
||||
|
||||
|
|
@ -1120,4 +1131,4 @@ And all the contributors:
|
|||
@simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola @buoyancy99 @caburu @xy9485
|
||||
@Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede
|
||||
@Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875
|
||||
@anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer
|
||||
@anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong
|
||||
|
|
|
|||
|
|
@ -60,7 +60,7 @@ This example is only to demonstrate the use of the library and its functions, an
|
|||
|
||||
from stable_baselines3 import DQN
|
||||
|
||||
env = gym.make("CartPole-v0")
|
||||
env = gym.make("CartPole-v1")
|
||||
|
||||
model = DQN("MlpPolicy", env, verbose=1)
|
||||
model.learn(total_timesteps=10000, log_interval=4)
|
||||
|
|
|
|||
|
|
@ -19,10 +19,12 @@ It creates "virtual" transitions by relabeling transitions (changing the desired
|
|||
but a replay buffer class ``HerReplayBuffer`` that must be passed to an off-policy algorithm
|
||||
when using ``MultiInputPolicy`` (to have Dict observation support).
|
||||
|
||||
|
||||
.. warning::
|
||||
|
||||
HER requires the environment to inherits from `gym.GoalEnv <https://github.com/openai/gym/blob/3394e245727c1ae6851b504a50ba77c73cd4c65b/gym/core.py#L160>`_
|
||||
HER requires the environment to follow the legacy `gym_robotics.GoalEnv interface <https://github.com/Farama-Foundation/Gymnasium-Robotics/blob/a35b1c1fa669428bf640a2c7101e66eb1627ac3a/gym_robotics/core.py#L8>`_
|
||||
In short, the ``gym.Env`` must have:
|
||||
- a vectorized implementation of ``compute_reward()``
|
||||
- a dictionary observation space with three keys: ``observation``, ``achieved_goal`` and ``desired_goal``
|
||||
|
||||
|
||||
.. warning::
|
||||
|
|
|
|||
21
setup.cfg
21
setup.cfg
|
|
@ -36,7 +36,6 @@ exclude = (?x)(
|
|||
| stable_baselines3/common/buffers.py$
|
||||
| stable_baselines3/common/callbacks.py$
|
||||
| stable_baselines3/common/distributions.py$
|
||||
| stable_baselines3/common/env_util.py$
|
||||
| stable_baselines3/common/envs/bit_flipping_env.py$
|
||||
| stable_baselines3/common/envs/identity_env.py$
|
||||
| stable_baselines3/common/envs/multi_input_envs.py$
|
||||
|
|
@ -48,8 +47,6 @@ exclude = (?x)(
|
|||
| stable_baselines3/common/preprocessing.py$
|
||||
| stable_baselines3/common/save_util.py$
|
||||
| stable_baselines3/common/sb2_compat/rmsprop_tf_like.py$
|
||||
| stable_baselines3/common/torch_layers.py$
|
||||
| stable_baselines3/common/type_aliases.py$
|
||||
| stable_baselines3/common/utils.py$
|
||||
| stable_baselines3/common/vec_env/__init__.py$
|
||||
| stable_baselines3/common/vec_env/base_vec_env.py$
|
||||
|
|
@ -72,7 +69,6 @@ exclude = (?x)(
|
|||
| stable_baselines3/sac/sac.py$
|
||||
| stable_baselines3/td3/policies.py$
|
||||
| stable_baselines3/td3/td3.py$
|
||||
| tests/test_distributions.py$
|
||||
| tests/test_logger.py$
|
||||
| tests/test_tensorboard.py$
|
||||
| tests/test_train_eval_mode.py$
|
||||
|
|
@ -80,20 +76,13 @@ exclude = (?x)(
|
|||
)
|
||||
|
||||
[flake8]
|
||||
ignore = W503,W504,E203,E231 # line breaks before and after binary operators
|
||||
# line breaks before and after binary operators
|
||||
ignore = W503,W504,E203,E231
|
||||
# Ignore import not used when aliases are defined
|
||||
per-file-ignores =
|
||||
./stable_baselines3/__init__.py:F401
|
||||
./stable_baselines3/common/__init__.py:F401
|
||||
./stable_baselines3/common/envs/__init__.py:F401
|
||||
./stable_baselines3/a2c/__init__.py:F401
|
||||
./stable_baselines3/ddpg/__init__.py:F401
|
||||
./stable_baselines3/dqn/__init__.py:F401
|
||||
./stable_baselines3/her/__init__.py:F401
|
||||
./stable_baselines3/ppo/__init__.py:F401
|
||||
./stable_baselines3/sac/__init__.py:F401
|
||||
./stable_baselines3/td3/__init__.py:F401
|
||||
./stable_baselines3/common/vec_env/__init__.py:F401
|
||||
# Default implementation in abstract methods
|
||||
./stable_baselines3/common/callbacks.py:B027
|
||||
./stable_baselines3/common/noise.py:B027
|
||||
exclude =
|
||||
# No need to traverse our git directory
|
||||
.git,
|
||||
|
|
|
|||
13
setup.py
13
setup.py
|
|
@ -48,13 +48,16 @@ env = gym.make("CartPole-v1")
|
|||
model = PPO("MlpPolicy", env, verbose=1)
|
||||
model.learn(total_timesteps=10_000)
|
||||
|
||||
obs = env.reset()
|
||||
vec_env = model.get_env()
|
||||
obs = vec_env.reset()
|
||||
for i in range(1000):
|
||||
action, _states = model.predict(obs, deterministic=True)
|
||||
obs, reward, done, info = env.step(action)
|
||||
env.render()
|
||||
if done:
|
||||
obs = env.reset()
|
||||
obs, reward, done, info = vec_env.step(action)
|
||||
vec_env.render()
|
||||
# VecEnv resets automatically
|
||||
# if done:
|
||||
# obs = vec_env.reset()
|
||||
|
||||
```
|
||||
|
||||
Or just train a model with a one liner if [the environment is registered in Gym](https://www.gymlibrary.ml/content/environment_creation/) and if [the policy is registered](https://stable-baselines3.readthedocs.io/en/master/guide/custom_policy.html):
|
||||
|
|
|
|||
|
|
@ -20,3 +20,15 @@ def HER(*args, **kwargs):
|
|||
"Since Stable Baselines 2.1.0, `HER` is now a replay buffer class `HerReplayBuffer`.\n "
|
||||
"Please check the documentation for more information: https://stable-baselines3.readthedocs.io/"
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"A2C",
|
||||
"DDPG",
|
||||
"DQN",
|
||||
"PPO",
|
||||
"SAC",
|
||||
"TD3",
|
||||
"HerReplayBuffer",
|
||||
"get_system_info",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
from stable_baselines3.a2c.a2c import A2C
|
||||
from stable_baselines3.a2c.policies import CnnPolicy, MlpPolicy, MultiInputPolicy
|
||||
|
||||
__all__ = ["CnnPolicy", "MlpPolicy", "MultiInputPolicy", "A2C"]
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from stable_baselines3.common.policies import ActorCriticCnnPolicy, ActorCriticP
|
|||
from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule
|
||||
from stable_baselines3.common.utils import explained_variance
|
||||
|
||||
A2CSelf = TypeVar("A2CSelf", bound="A2C")
|
||||
SelfA2C = TypeVar("SelfA2C", bound="A2C")
|
||||
|
||||
|
||||
class A2C(OnPolicyAlgorithm):
|
||||
|
|
@ -181,14 +181,14 @@ class A2C(OnPolicyAlgorithm):
|
|||
self.logger.record("train/std", th.exp(self.policy.log_std).mean().item())
|
||||
|
||||
def learn(
|
||||
self: A2CSelf,
|
||||
self: SelfA2C,
|
||||
total_timesteps: int,
|
||||
callback: MaybeCallback = None,
|
||||
log_interval: int = 100,
|
||||
tb_log_name: str = "A2C",
|
||||
reset_num_timesteps: bool = True,
|
||||
progress_bar: bool = False,
|
||||
) -> A2CSelf:
|
||||
) -> SelfA2C:
|
||||
|
||||
return super().learn(
|
||||
total_timesteps=total_timesteps,
|
||||
|
|
|
|||
|
|
@ -86,7 +86,7 @@ class EpisodicLifeEnv(gym.Wrapper):
|
|||
# then update lives to handle bonus lives
|
||||
lives = self.env.unwrapped.ale.lives()
|
||||
if 0 < lives < self.lives:
|
||||
# for Qbert sometimes we stay in lives == 0 condtion for a few frames
|
||||
# for Qbert sometimes we stay in lives == 0 condition for a few frames
|
||||
# so its important to keep lives > 0, so that we only reset once
|
||||
# the environment advertises done.
|
||||
done = True
|
||||
|
|
|
|||
|
|
@ -38,6 +38,8 @@ from stable_baselines3.common.vec_env import (
|
|||
unwrap_vec_normalize,
|
||||
)
|
||||
|
||||
SelfBaseAlgorithm = TypeVar("SelfBaseAlgorithm", bound="BaseAlgorithm")
|
||||
|
||||
|
||||
def maybe_make_env(env: Union[GymEnv, str, None], verbose: int) -> Optional[GymEnv]:
|
||||
"""If env is a string, make the environment; otherwise, return env.
|
||||
|
|
@ -53,9 +55,6 @@ def maybe_make_env(env: Union[GymEnv, str, None], verbose: int) -> Optional[GymE
|
|||
return env
|
||||
|
||||
|
||||
BaseAlgorithmSelf = TypeVar("BaseAlgorithmSelf", bound="BaseAlgorithm")
|
||||
|
||||
|
||||
class BaseAlgorithm(ABC):
|
||||
"""
|
||||
The base of RL algorithms
|
||||
|
|
@ -126,7 +125,7 @@ class BaseAlgorithm(ABC):
|
|||
# Used for computing fps, it is updated at each call of learn()
|
||||
self._num_timesteps_at_start = 0
|
||||
self.seed = seed
|
||||
self.action_noise = None # type: Optional[ActionNoise]
|
||||
self.action_noise: Optional[ActionNoise] = None
|
||||
self.start_time = None
|
||||
self.policy = None
|
||||
self.learning_rate = learning_rate
|
||||
|
|
@ -491,14 +490,14 @@ class BaseAlgorithm(ABC):
|
|||
|
||||
@abstractmethod
|
||||
def learn(
|
||||
self: BaseAlgorithmSelf,
|
||||
self: SelfBaseAlgorithm,
|
||||
total_timesteps: int,
|
||||
callback: MaybeCallback = None,
|
||||
log_interval: int = 100,
|
||||
tb_log_name: str = "run",
|
||||
reset_num_timesteps: bool = True,
|
||||
progress_bar: bool = False,
|
||||
) -> BaseAlgorithmSelf:
|
||||
) -> SelfBaseAlgorithm:
|
||||
"""
|
||||
Return a trained model.
|
||||
|
||||
|
|
@ -544,6 +543,7 @@ class BaseAlgorithm(ABC):
|
|||
return
|
||||
set_random_seed(seed, using_cuda=self.device.type == th.device("cuda").type)
|
||||
self.action_space.seed(seed)
|
||||
# self.env is always a VecEnv
|
||||
if self.env is not None:
|
||||
self.env.seed(seed)
|
||||
|
||||
|
|
@ -617,7 +617,7 @@ class BaseAlgorithm(ABC):
|
|||
|
||||
@classmethod
|
||||
def load(
|
||||
cls: Type[BaseAlgorithmSelf],
|
||||
cls: Type[SelfBaseAlgorithm],
|
||||
path: Union[str, pathlib.Path, io.BufferedIOBase],
|
||||
env: Optional[GymEnv] = None,
|
||||
device: Union[th.device, str] = "auto",
|
||||
|
|
@ -625,7 +625,7 @@ class BaseAlgorithm(ABC):
|
|||
print_system_info: bool = False,
|
||||
force_reset: bool = True,
|
||||
**kwargs,
|
||||
) -> BaseAlgorithmSelf:
|
||||
) -> SelfBaseAlgorithm:
|
||||
"""
|
||||
Load the model from a zip-file.
|
||||
Warning: ``load`` re-creates the model from scratch, it does not update it in-place!
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Probability distributions."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
|
|
@ -12,6 +12,16 @@ from torch.distributions import Bernoulli, Categorical, Normal
|
|||
|
||||
from stable_baselines3.common.preprocessing import get_action_dim
|
||||
|
||||
SelfDistribution = TypeVar("SelfDistribution", bound="Distribution")
|
||||
SelfDiagGaussianDistribution = TypeVar("SelfDiagGaussianDistribution", bound="DiagGaussianDistribution")
|
||||
SelfSquashedDiagGaussianDistribution = TypeVar(
|
||||
"SelfSquashedDiagGaussianDistribution", bound="SquashedDiagGaussianDistribution"
|
||||
)
|
||||
SelfCategoricalDistribution = TypeVar("SelfCategoricalDistribution", bound="CategoricalDistribution")
|
||||
SelfMultiCategoricalDistribution = TypeVar("SelfMultiCategoricalDistribution", bound="MultiCategoricalDistribution")
|
||||
SelfBernoulliDistribution = TypeVar("SelfBernoulliDistribution", bound="BernoulliDistribution")
|
||||
SelfStateDependentNoiseDistribution = TypeVar("SelfStateDependentNoiseDistribution", bound="StateDependentNoiseDistribution")
|
||||
|
||||
|
||||
class Distribution(ABC):
|
||||
"""Abstract base class for distributions."""
|
||||
|
|
@ -28,7 +38,7 @@ class Distribution(ABC):
|
|||
concrete classes."""
|
||||
|
||||
@abstractmethod
|
||||
def proba_distribution(self, *args, **kwargs) -> "Distribution":
|
||||
def proba_distribution(self: SelfDistribution, *args, **kwargs) -> SelfDistribution:
|
||||
"""Set parameters of the distribution.
|
||||
|
||||
:return: self
|
||||
|
|
@ -141,7 +151,9 @@ class DiagGaussianDistribution(Distribution):
|
|||
log_std = nn.Parameter(th.ones(self.action_dim) * log_std_init, requires_grad=True)
|
||||
return mean_actions, log_std
|
||||
|
||||
def proba_distribution(self, mean_actions: th.Tensor, log_std: th.Tensor) -> "DiagGaussianDistribution":
|
||||
def proba_distribution(
|
||||
self: SelfDiagGaussianDistribution, mean_actions: th.Tensor, log_std: th.Tensor
|
||||
) -> SelfDiagGaussianDistribution:
|
||||
"""
|
||||
Create the distribution given its parameters (mean, std)
|
||||
|
||||
|
|
@ -207,7 +219,9 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
|
|||
self.epsilon = epsilon
|
||||
self.gaussian_actions = None
|
||||
|
||||
def proba_distribution(self, mean_actions: th.Tensor, log_std: th.Tensor) -> "SquashedDiagGaussianDistribution":
|
||||
def proba_distribution(
|
||||
self: SelfSquashedDiagGaussianDistribution, mean_actions: th.Tensor, log_std: th.Tensor
|
||||
) -> SelfSquashedDiagGaussianDistribution:
|
||||
super().proba_distribution(mean_actions, log_std)
|
||||
return self
|
||||
|
||||
|
|
@ -271,7 +285,7 @@ class CategoricalDistribution(Distribution):
|
|||
action_logits = nn.Linear(latent_dim, self.action_dim)
|
||||
return action_logits
|
||||
|
||||
def proba_distribution(self, action_logits: th.Tensor) -> "CategoricalDistribution":
|
||||
def proba_distribution(self: SelfCategoricalDistribution, action_logits: th.Tensor) -> SelfCategoricalDistribution:
|
||||
self.distribution = Categorical(logits=action_logits)
|
||||
return self
|
||||
|
||||
|
|
@ -323,7 +337,9 @@ class MultiCategoricalDistribution(Distribution):
|
|||
action_logits = nn.Linear(latent_dim, sum(self.action_dims))
|
||||
return action_logits
|
||||
|
||||
def proba_distribution(self, action_logits: th.Tensor) -> "MultiCategoricalDistribution":
|
||||
def proba_distribution(
|
||||
self: SelfMultiCategoricalDistribution, action_logits: th.Tensor
|
||||
) -> SelfMultiCategoricalDistribution:
|
||||
self.distribution = [Categorical(logits=split) for split in th.split(action_logits, tuple(self.action_dims), dim=1)]
|
||||
return self
|
||||
|
||||
|
|
@ -376,7 +392,7 @@ class BernoulliDistribution(Distribution):
|
|||
action_logits = nn.Linear(latent_dim, self.action_dims)
|
||||
return action_logits
|
||||
|
||||
def proba_distribution(self, action_logits: th.Tensor) -> "BernoulliDistribution":
|
||||
def proba_distribution(self: SelfBernoulliDistribution, action_logits: th.Tensor) -> SelfBernoulliDistribution:
|
||||
self.distribution = Bernoulli(logits=action_logits)
|
||||
return self
|
||||
|
||||
|
|
@ -520,8 +536,8 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
return mean_actions_net, log_std
|
||||
|
||||
def proba_distribution(
|
||||
self, mean_actions: th.Tensor, log_std: th.Tensor, latent_sde: th.Tensor
|
||||
) -> "StateDependentNoiseDistribution":
|
||||
self: SelfStateDependentNoiseDistribution, mean_actions: th.Tensor, log_std: th.Tensor, latent_sde: th.Tensor
|
||||
) -> SelfStateDependentNoiseDistribution:
|
||||
"""
|
||||
Create the distribution given its parameters (mean, std)
|
||||
|
||||
|
|
@ -663,7 +679,7 @@ def make_proba_distribution(
|
|||
elif isinstance(action_space, spaces.Discrete):
|
||||
return CategoricalDistribution(action_space.n, **dist_kwargs)
|
||||
elif isinstance(action_space, spaces.MultiDiscrete):
|
||||
return MultiCategoricalDistribution(action_space.nvec, **dist_kwargs)
|
||||
return MultiCategoricalDistribution(list(action_space.nvec), **dist_kwargs)
|
||||
elif isinstance(action_space, spaces.MultiBinary):
|
||||
return BernoulliDistribution(action_space.n, **dist_kwargs)
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import warnings
|
||||
from typing import Union
|
||||
from typing import Any, Dict, Union
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
|
|
@ -93,6 +93,62 @@ def _check_nan(env: gym.Env) -> None:
|
|||
_, _, _, _ = vec_env.step(action)
|
||||
|
||||
|
||||
def _is_goal_env(env: gym.Env) -> bool:
|
||||
"""
|
||||
Check if the env uses the convention for goal-conditioned envs (previously, the gym.GoalEnv interface)
|
||||
"""
|
||||
if isinstance(env, gym.Wrapper): # We need to unwrap the env since gym.Wrapper has the compute_reward method
|
||||
return _is_goal_env(env.unwrapped)
|
||||
return hasattr(env, "compute_reward")
|
||||
|
||||
|
||||
def _check_goal_env_obs(obs: dict, observation_space: spaces.Dict, method_name: str) -> None:
|
||||
"""
|
||||
Check that an environment implementing the `compute_rewards()` method
|
||||
(previously known as GoalEnv in gym) contains three elements,
|
||||
namely `observation`, `desired_goal`, and `achieved_goal`.
|
||||
"""
|
||||
assert len(observation_space.spaces) == 3, (
|
||||
"A goal conditioned env must contain 3 observation keys: `observation`, `desired_goal`, and `achieved_goal`."
|
||||
f"The current observation contains {len(observation_space.spaces)} keys: {list(observation_space.spaces.keys())}"
|
||||
)
|
||||
|
||||
for key in ["observation", "achieved_goal", "desired_goal"]:
|
||||
if key not in observation_space.spaces:
|
||||
raise AssertionError(
|
||||
f"The observation returned by the `{method_name}()` method of a goal-conditioned env requires the '{key}' "
|
||||
"key to be part of the observation dictionary. "
|
||||
f"Current keys are {list(observation_space.spaces.keys())}"
|
||||
)
|
||||
|
||||
|
||||
def _check_goal_env_compute_reward(
|
||||
obs: Dict[str, Union[np.ndarray, int]],
|
||||
env: gym.Env,
|
||||
reward: float,
|
||||
info: Dict[str, Any],
|
||||
):
|
||||
"""
|
||||
Check that reward is computed with `compute_reward`
|
||||
and that the implementation is vectorized.
|
||||
"""
|
||||
achieved_goal, desired_goal = obs["achieved_goal"], obs["desired_goal"]
|
||||
assert reward == env.compute_reward( # type: ignore[attr-defined]
|
||||
achieved_goal, desired_goal, info
|
||||
), "The reward was not computed with `compute_reward()`"
|
||||
|
||||
achieved_goal, desired_goal = np.array(achieved_goal), np.array(desired_goal)
|
||||
batch_achieved_goals = np.array([achieved_goal, achieved_goal])
|
||||
batch_desired_goals = np.array([desired_goal, desired_goal])
|
||||
if isinstance(achieved_goal, int) or len(achieved_goal.shape) == 0:
|
||||
batch_achieved_goals = batch_achieved_goals.reshape(2, 1)
|
||||
batch_desired_goals = batch_desired_goals.reshape(2, 1)
|
||||
batch_infos = np.array([info, info])
|
||||
rewards = env.compute_reward(batch_achieved_goals, batch_desired_goals, batch_infos) # type: ignore[attr-defined]
|
||||
assert rewards.shape == (2,), f"Unexpected shape for vectorized computation of reward: {rewards.shape} != (2,)"
|
||||
assert rewards[0] == reward, f"Vectorized computation of reward differs from single computation: {rewards[0]} != {reward}"
|
||||
|
||||
|
||||
def _check_obs(obs: Union[tuple, dict, np.ndarray, int], observation_space: spaces.Space, method_name: str) -> None:
|
||||
"""
|
||||
Check that the observation returned by the environment
|
||||
|
|
@ -141,7 +197,11 @@ def _check_returned_values(env: gym.Env, observation_space: spaces.Space, action
|
|||
# because env inherits from gym.Env, we assume that `reset()` and `step()` methods exists
|
||||
obs = env.reset()
|
||||
|
||||
if isinstance(observation_space, spaces.Dict):
|
||||
if _is_goal_env(env):
|
||||
# Make mypy happy, already checked
|
||||
assert isinstance(observation_space, spaces.Dict)
|
||||
_check_goal_env_obs(obs, observation_space, "reset")
|
||||
elif isinstance(observation_space, spaces.Dict):
|
||||
assert isinstance(obs, dict), "The observation returned by `reset()` must be a dictionary"
|
||||
|
||||
if not obs.keys() == observation_space.spaces.keys():
|
||||
|
|
@ -167,7 +227,12 @@ def _check_returned_values(env: gym.Env, observation_space: spaces.Space, action
|
|||
# Unpack
|
||||
obs, reward, done, info = data
|
||||
|
||||
if isinstance(observation_space, spaces.Dict):
|
||||
if _is_goal_env(env):
|
||||
# Make mypy happy, already checked
|
||||
assert isinstance(observation_space, spaces.Dict)
|
||||
_check_goal_env_obs(obs, observation_space, "step")
|
||||
_check_goal_env_compute_reward(obs, env, reward, info)
|
||||
elif isinstance(observation_space, spaces.Dict):
|
||||
assert isinstance(obs, dict), "The observation returned by `step()` must be a dictionary"
|
||||
|
||||
if not obs.keys() == observation_space.spaces.keys():
|
||||
|
|
@ -190,15 +255,16 @@ def _check_returned_values(env: gym.Env, observation_space: spaces.Space, action
|
|||
assert isinstance(done, bool), "The `done` signal must be a boolean"
|
||||
assert isinstance(info, dict), "The `info` returned by `step()` must be a python dictionary"
|
||||
|
||||
if isinstance(env, gym.GoalEnv):
|
||||
# For a GoalEnv, the keys are checked at reset
|
||||
# Goal conditioned env
|
||||
if _is_goal_env(env):
|
||||
assert reward == env.compute_reward(obs["achieved_goal"], obs["desired_goal"], info)
|
||||
|
||||
|
||||
def _check_spaces(env: gym.Env) -> None:
|
||||
"""
|
||||
Check that the observation and action spaces are defined
|
||||
and inherit from gym.spaces.Space.
|
||||
Check that the observation and action spaces are defined and inherit from gym.spaces.Space. For
|
||||
envs that follow the goal-conditioned standard (previously, the gym.GoalEnv interface) we check
|
||||
the observation space is gym.spaces.Dict
|
||||
"""
|
||||
# Helper to link to the code, because gym has no proper documentation
|
||||
gym_spaces = " cf https://github.com/openai/gym/blob/master/gym/spaces/"
|
||||
|
|
@ -209,6 +275,11 @@ def _check_spaces(env: gym.Env) -> None:
|
|||
assert isinstance(env.observation_space, spaces.Space), "The observation space must inherit from gym.spaces" + gym_spaces
|
||||
assert isinstance(env.action_space, spaces.Space), "The action space must inherit from gym.spaces" + gym_spaces
|
||||
|
||||
if _is_goal_env(env):
|
||||
assert isinstance(
|
||||
env.observation_space, spaces.Dict
|
||||
), "Goal conditioned envs (previously gym.GoalEnv) require the observation space to be gym.spaces.Dict"
|
||||
|
||||
|
||||
# Check render cannot be covered by CI
|
||||
def _check_render(env: gym.Env, warn: bool = True, headless: bool = False) -> None: # pragma: no cover
|
||||
|
|
|
|||
|
|
@ -116,7 +116,7 @@ def make_atari_env(
|
|||
monitor_dir: Optional[str] = None,
|
||||
wrapper_kwargs: Optional[Dict[str, Any]] = None,
|
||||
env_kwargs: Optional[Dict[str, Any]] = None,
|
||||
vec_env_cls: Optional[Union[DummyVecEnv, SubprocVecEnv]] = None,
|
||||
vec_env_cls: Optional[Union[Type[DummyVecEnv], Type[SubprocVecEnv]]] = None,
|
||||
vec_env_kwargs: Optional[Dict[str, Any]] = None,
|
||||
monitor_kwargs: Optional[Dict[str, Any]] = None,
|
||||
) -> VecEnv:
|
||||
|
|
@ -138,22 +138,16 @@ def make_atari_env(
|
|||
:param monitor_kwargs: Keyword arguments to pass to the ``Monitor`` class constructor.
|
||||
:return: The wrapped environment
|
||||
"""
|
||||
if wrapper_kwargs is None:
|
||||
wrapper_kwargs = {}
|
||||
|
||||
def atari_wrapper(env: gym.Env) -> gym.Env:
|
||||
env = AtariWrapper(env, **wrapper_kwargs)
|
||||
return env
|
||||
|
||||
return make_vec_env(
|
||||
env_id,
|
||||
n_envs=n_envs,
|
||||
seed=seed,
|
||||
start_index=start_index,
|
||||
monitor_dir=monitor_dir,
|
||||
wrapper_class=atari_wrapper,
|
||||
wrapper_class=AtariWrapper,
|
||||
env_kwargs=env_kwargs,
|
||||
vec_env_cls=vec_env_cls,
|
||||
vec_env_kwargs=vec_env_kwargs,
|
||||
monitor_kwargs=monitor_kwargs,
|
||||
wrapper_kwargs=wrapper_kwargs,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,3 +7,14 @@ from stable_baselines3.common.envs.identity_env import (
|
|||
IdentityEnvMultiDiscrete,
|
||||
)
|
||||
from stable_baselines3.common.envs.multi_input_envs import SimpleMultiObsEnv
|
||||
|
||||
__all__ = [
|
||||
"BitFlippingEnv",
|
||||
"FakeImageEnv",
|
||||
"IdentityEnv",
|
||||
"IdentityEnvBox",
|
||||
"IdentityEnvMultiBinary",
|
||||
"IdentityEnvMultiDiscrete",
|
||||
"SimpleMultiObsEnv",
|
||||
"SimpleMultiObsEnv",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -2,13 +2,13 @@ from collections import OrderedDict
|
|||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
from gym import GoalEnv, spaces
|
||||
from gym import Env, spaces
|
||||
from gym.envs.registration import EnvSpec
|
||||
|
||||
from stable_baselines3.common.type_aliases import GymStepReturn
|
||||
|
||||
|
||||
class BitFlippingEnv(GoalEnv):
|
||||
class BitFlippingEnv(Env):
|
||||
"""
|
||||
Simple bit flipping env, useful to test HER.
|
||||
The goal is to flip all the bits to get a vector of ones.
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ class ActionNoise(ABC):
|
|||
The action noise base class
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def reset(self) -> None:
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ from stable_baselines3.common.utils import safe_mean, should_collect_more_steps
|
|||
from stable_baselines3.common.vec_env import VecEnv
|
||||
from stable_baselines3.her.her_replay_buffer import HerReplayBuffer
|
||||
|
||||
OffPolicyAlgorithmSelf = TypeVar("OffPolicyAlgorithmSelf", bound="OffPolicyAlgorithm")
|
||||
SelfOffPolicyAlgorithm = TypeVar("SelfOffPolicyAlgorithm", bound="OffPolicyAlgorithm")
|
||||
|
||||
|
||||
class OffPolicyAlgorithm(BaseAlgorithm):
|
||||
|
|
@ -311,14 +311,14 @@ class OffPolicyAlgorithm(BaseAlgorithm):
|
|||
)
|
||||
|
||||
def learn(
|
||||
self: OffPolicyAlgorithmSelf,
|
||||
self: SelfOffPolicyAlgorithm,
|
||||
total_timesteps: int,
|
||||
callback: MaybeCallback = None,
|
||||
log_interval: int = 4,
|
||||
tb_log_name: str = "run",
|
||||
reset_num_timesteps: bool = True,
|
||||
progress_bar: bool = False,
|
||||
) -> OffPolicyAlgorithmSelf:
|
||||
) -> SelfOffPolicyAlgorithm:
|
||||
|
||||
total_timesteps, callback = self._setup_learn(
|
||||
total_timesteps,
|
||||
|
|
@ -602,7 +602,6 @@ class OffPolicyAlgorithm(BaseAlgorithm):
|
|||
# Log training infos
|
||||
if log_interval is not None and self._episode_num % log_interval == 0:
|
||||
self._dump_logs()
|
||||
|
||||
callback.on_rollout_end()
|
||||
|
||||
return RolloutReturn(num_collected_steps * env.num_envs, num_collected_episodes, continue_training)
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedul
|
|||
from stable_baselines3.common.utils import obs_as_tensor, safe_mean
|
||||
from stable_baselines3.common.vec_env import VecEnv
|
||||
|
||||
OnPolicyAlgorithmSelf = TypeVar("OnPolicyAlgorithmSelf", bound="OnPolicyAlgorithm")
|
||||
SelfOnPolicyAlgorithm = TypeVar("SelfOnPolicyAlgorithm", bound="OnPolicyAlgorithm")
|
||||
|
||||
|
||||
class OnPolicyAlgorithm(BaseAlgorithm):
|
||||
|
|
@ -223,14 +223,14 @@ class OnPolicyAlgorithm(BaseAlgorithm):
|
|||
raise NotImplementedError
|
||||
|
||||
def learn(
|
||||
self: OnPolicyAlgorithmSelf,
|
||||
self: SelfOnPolicyAlgorithm,
|
||||
total_timesteps: int,
|
||||
callback: MaybeCallback = None,
|
||||
log_interval: int = 1,
|
||||
tb_log_name: str = "OnPolicyAlgorithm",
|
||||
reset_num_timesteps: bool = True,
|
||||
progress_bar: bool = False,
|
||||
) -> OnPolicyAlgorithmSelf:
|
||||
) -> SelfOnPolicyAlgorithm:
|
||||
iteration = 0
|
||||
|
||||
total_timesteps, callback = self._setup_learn(
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ from stable_baselines3.common.torch_layers import (
|
|||
from stable_baselines3.common.type_aliases import Schedule
|
||||
from stable_baselines3.common.utils import get_device, is_vectorized_observation, obs_as_tensor
|
||||
|
||||
BaseModelSelf = TypeVar("BaseModelSelf", bound="BaseModel")
|
||||
SelfBaseModel = TypeVar("SelfBaseModel", bound="BaseModel")
|
||||
|
||||
|
||||
class BaseModel(nn.Module):
|
||||
|
|
@ -159,7 +159,7 @@ class BaseModel(nn.Module):
|
|||
th.save({"state_dict": self.state_dict(), "data": self._get_constructor_parameters()}, path)
|
||||
|
||||
@classmethod
|
||||
def load(cls: Type[BaseModelSelf], path: str, device: Union[th.device, str] = "auto") -> BaseModelSelf:
|
||||
def load(cls: Type[SelfBaseModel], path: str, device: Union[th.device, str] = "auto") -> SelfBaseModel:
|
||||
"""
|
||||
Load model from path.
|
||||
|
||||
|
|
|
|||
|
|
@ -28,9 +28,6 @@ class BaseFeaturesExtractor(nn.Module):
|
|||
def features_dim(self) -> int:
|
||||
return self._features_dim
|
||||
|
||||
def forward(self, observations: th.Tensor) -> th.Tensor:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class FlattenExtractor(BaseFeaturesExtractor):
|
||||
"""
|
||||
|
|
@ -99,6 +96,7 @@ def create_mlp(
|
|||
net_arch: List[int],
|
||||
activation_fn: Type[nn.Module] = nn.ReLU,
|
||||
squash_output: bool = False,
|
||||
with_bias: bool = True,
|
||||
) -> List[nn.Module]:
|
||||
"""
|
||||
Create a multi layer perceptron (MLP), which is
|
||||
|
|
@ -113,21 +111,22 @@ def create_mlp(
|
|||
to use after each layer.
|
||||
:param squash_output: Whether to squash the output using a Tanh
|
||||
activation function
|
||||
:param with_bias: If set to False, the layers will not learn an additive bias
|
||||
:return:
|
||||
"""
|
||||
|
||||
if len(net_arch) > 0:
|
||||
modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()]
|
||||
modules = [nn.Linear(input_dim, net_arch[0], bias=with_bias), activation_fn()]
|
||||
else:
|
||||
modules = []
|
||||
|
||||
for idx in range(len(net_arch) - 1):
|
||||
modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1]))
|
||||
modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1], bias=with_bias))
|
||||
modules.append(activation_fn())
|
||||
|
||||
if output_dim > 0:
|
||||
last_layer_dim = net_arch[-1] if len(net_arch) > 0 else input_dim
|
||||
modules.append(nn.Linear(last_layer_dim, output_dim))
|
||||
modules.append(nn.Linear(last_layer_dim, output_dim, bias=with_bias))
|
||||
if squash_output:
|
||||
modules.append(nn.Tanh())
|
||||
return modules
|
||||
|
|
@ -171,9 +170,11 @@ class MlpExtractor(nn.Module):
|
|||
):
|
||||
super().__init__()
|
||||
device = get_device(device)
|
||||
shared_net, policy_net, value_net = [], [], []
|
||||
policy_only_layers = [] # Layer sizes of the network that only belongs to the policy network
|
||||
value_only_layers = [] # Layer sizes of the network that only belongs to the value network
|
||||
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
|
||||
|
||||
# Iterate through the shared layers and build the shared parts of the network
|
||||
|
|
@ -252,7 +253,7 @@ class CombinedExtractor(BaseFeaturesExtractor):
|
|||
# TODO we do not know features-dim here before going over all the items, so put something there. This is dirty!
|
||||
super().__init__(observation_space, features_dim=1)
|
||||
|
||||
extractors = {}
|
||||
extractors: Dict[str, nn.Module] = {}
|
||||
|
||||
total_concat_size = 0
|
||||
for key, subspace in observation_space.spaces.items():
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ class RolloutBufferSamples(NamedTuple):
|
|||
returns: th.Tensor
|
||||
|
||||
|
||||
class DictRolloutBufferSamples(RolloutBufferSamples):
|
||||
class DictRolloutBufferSamples(NamedTuple):
|
||||
observations: TensorDict
|
||||
actions: th.Tensor
|
||||
old_values: th.Tensor
|
||||
|
|
@ -53,7 +53,7 @@ class ReplayBufferSamples(NamedTuple):
|
|||
rewards: th.Tensor
|
||||
|
||||
|
||||
class DictReplayBufferSamples(ReplayBufferSamples):
|
||||
class DictReplayBufferSamples(NamedTuple):
|
||||
observations: TensorDict
|
||||
actions: th.Tensor
|
||||
next_observations: TensorDict
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
# flake8: noqa F401
|
||||
import typing
|
||||
from copy import deepcopy
|
||||
from typing import Optional, Type, Union
|
||||
|
|
@ -72,3 +71,25 @@ def sync_envs_normalization(env: "GymEnv", eval_env: "GymEnv") -> None:
|
|||
eval_env_tmp.ret_rms = deepcopy(env_tmp.ret_rms)
|
||||
env_tmp = env_tmp.venv
|
||||
eval_env_tmp = eval_env_tmp.venv
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CloudpickleWrapper",
|
||||
"VecEnv",
|
||||
"VecEnvWrapper",
|
||||
"DummyVecEnv",
|
||||
"StackedDictObservations",
|
||||
"StackedObservations",
|
||||
"SubprocVecEnv",
|
||||
"VecCheckNan",
|
||||
"VecExtractDictObs",
|
||||
"VecFrameStack",
|
||||
"VecMonitor",
|
||||
"VecNormalize",
|
||||
"VecTransposeImage",
|
||||
"VecVideoRecorder",
|
||||
"unwrap_vec_wrapper",
|
||||
"unwrap_vec_normalize",
|
||||
"is_vecenv_wrapped",
|
||||
"sync_envs_normalization",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -19,10 +19,21 @@ class DummyVecEnv(VecEnv):
|
|||
|
||||
:param env_fns: a list of functions
|
||||
that return environments to vectorize
|
||||
:raises ValueError: If the same environment instance is passed as the output of two or more different env_fn.
|
||||
"""
|
||||
|
||||
def __init__(self, env_fns: List[Callable[[], gym.Env]]):
|
||||
self.envs = [fn() for fn in env_fns]
|
||||
if len(set([id(env.unwrapped) for env in self.envs])) != len(self.envs):
|
||||
raise ValueError(
|
||||
"You tried to create multiple environments, but the function to create them returned the same instance "
|
||||
"instead of creating different objects. "
|
||||
"You are probably using `make_vec_env(lambda: env)` or `DummyVecEnv([lambda: env] * n_envs)`. "
|
||||
"You should replace `lambda: env` by a `make_env` function that "
|
||||
"creates a new instance of the environment at every call "
|
||||
"(using `gym.make()` for instance). You can take a look at the documentation for an example. "
|
||||
"Please read https://github.com/DLR-RM/stable-baselines3/issues/1151 for more information."
|
||||
)
|
||||
env = self.envs[0]
|
||||
VecEnv.__init__(self, len(env_fns), env.observation_space, env.action_space)
|
||||
obs_space = env.observation_space
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
from stable_baselines3.ddpg.ddpg import DDPG
|
||||
from stable_baselines3.ddpg.policies import CnnPolicy, MlpPolicy, MultiInputPolicy
|
||||
|
||||
__all__ = ["CnnPolicy", "MlpPolicy", "MultiInputPolicy", "DDPG"]
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedul
|
|||
from stable_baselines3.td3.policies import TD3Policy
|
||||
from stable_baselines3.td3.td3 import TD3
|
||||
|
||||
DDPGSelf = TypeVar("DDPGSelf", bound="DDPG")
|
||||
SelfDDPG = TypeVar("SelfDDPG", bound="DDPG")
|
||||
|
||||
|
||||
class DDPG(TD3):
|
||||
|
|
@ -113,14 +113,14 @@ class DDPG(TD3):
|
|||
self._setup_model()
|
||||
|
||||
def learn(
|
||||
self: DDPGSelf,
|
||||
self: SelfDDPG,
|
||||
total_timesteps: int,
|
||||
callback: MaybeCallback = None,
|
||||
log_interval: int = 4,
|
||||
tb_log_name: str = "DDPG",
|
||||
reset_num_timesteps: bool = True,
|
||||
progress_bar: bool = False,
|
||||
) -> DDPGSelf:
|
||||
) -> SelfDDPG:
|
||||
|
||||
return super().learn(
|
||||
total_timesteps=total_timesteps,
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
from stable_baselines3.dqn.dqn import DQN
|
||||
from stable_baselines3.dqn.policies import CnnPolicy, MlpPolicy, MultiInputPolicy
|
||||
|
||||
__all__ = ["CnnPolicy", "MlpPolicy", "MultiInputPolicy", "DQN"]
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedul
|
|||
from stable_baselines3.common.utils import get_linear_fn, get_parameters_by_name, is_vectorized_observation, polyak_update
|
||||
from stable_baselines3.dqn.policies import CnnPolicy, DQNPolicy, MlpPolicy, MultiInputPolicy
|
||||
|
||||
DQNSelf = TypeVar("DQNSelf", bound="DQN")
|
||||
SelfDQN = TypeVar("SelfDQN", bound="DQN")
|
||||
|
||||
|
||||
class DQN(OffPolicyAlgorithm):
|
||||
|
|
@ -253,14 +253,14 @@ class DQN(OffPolicyAlgorithm):
|
|||
return action, state
|
||||
|
||||
def learn(
|
||||
self: DQNSelf,
|
||||
self: SelfDQN,
|
||||
total_timesteps: int,
|
||||
callback: MaybeCallback = None,
|
||||
log_interval: int = 4,
|
||||
tb_log_name: str = "DQN",
|
||||
reset_num_timesteps: bool = True,
|
||||
progress_bar: bool = False,
|
||||
) -> DQNSelf:
|
||||
) -> SelfDQN:
|
||||
|
||||
return super().learn(
|
||||
total_timesteps=total_timesteps,
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
from stable_baselines3.her.goal_selection_strategy import GoalSelectionStrategy
|
||||
from stable_baselines3.her.her_replay_buffer import HerReplayBuffer
|
||||
|
||||
__all__ = ["GoalSelectionStrategy", "HerReplayBuffer"]
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
from stable_baselines3.ppo.policies import CnnPolicy, MlpPolicy, MultiInputPolicy
|
||||
from stable_baselines3.ppo.ppo import PPO
|
||||
|
||||
__all__ = ["CnnPolicy", "MlpPolicy", "MultiInputPolicy", "PPO"]
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ from stable_baselines3.common.policies import ActorCriticCnnPolicy, ActorCriticP
|
|||
from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule
|
||||
from stable_baselines3.common.utils import explained_variance, get_schedule_fn
|
||||
|
||||
PPOSelf = TypeVar("PPOSelf", bound="PPO")
|
||||
SelfPPO = TypeVar("SelfPPO", bound="PPO")
|
||||
|
||||
|
||||
class PPO(OnPolicyAlgorithm):
|
||||
|
|
@ -295,14 +295,14 @@ class PPO(OnPolicyAlgorithm):
|
|||
self.logger.record("train/clip_range_vf", clip_range_vf)
|
||||
|
||||
def learn(
|
||||
self: PPOSelf,
|
||||
self: SelfPPO,
|
||||
total_timesteps: int,
|
||||
callback: MaybeCallback = None,
|
||||
log_interval: int = 1,
|
||||
tb_log_name: str = "PPO",
|
||||
reset_num_timesteps: bool = True,
|
||||
progress_bar: bool = False,
|
||||
) -> PPOSelf:
|
||||
) -> SelfPPO:
|
||||
|
||||
return super().learn(
|
||||
total_timesteps=total_timesteps,
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
from stable_baselines3.sac.policies import CnnPolicy, MlpPolicy, MultiInputPolicy
|
||||
from stable_baselines3.sac.sac import SAC
|
||||
|
||||
__all__ = ["CnnPolicy", "MlpPolicy", "MultiInputPolicy", "SAC"]
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedul
|
|||
from stable_baselines3.common.utils import get_parameters_by_name, polyak_update
|
||||
from stable_baselines3.sac.policies import CnnPolicy, MlpPolicy, MultiInputPolicy, SACPolicy
|
||||
|
||||
SACSelf = TypeVar("SACSelf", bound="SAC")
|
||||
SelfSAC = TypeVar("SelfSAC", bound="SAC")
|
||||
|
||||
|
||||
class SAC(OffPolicyAlgorithm):
|
||||
|
|
@ -287,14 +287,14 @@ class SAC(OffPolicyAlgorithm):
|
|||
self.logger.record("train/ent_coef_loss", np.mean(ent_coef_losses))
|
||||
|
||||
def learn(
|
||||
self: SACSelf,
|
||||
self: SelfSAC,
|
||||
total_timesteps: int,
|
||||
callback: MaybeCallback = None,
|
||||
log_interval: int = 4,
|
||||
tb_log_name: str = "SAC",
|
||||
reset_num_timesteps: bool = True,
|
||||
progress_bar: bool = False,
|
||||
) -> SACSelf:
|
||||
) -> SelfSAC:
|
||||
|
||||
return super().learn(
|
||||
total_timesteps=total_timesteps,
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
from stable_baselines3.td3.policies import CnnPolicy, MlpPolicy, MultiInputPolicy
|
||||
from stable_baselines3.td3.td3 import TD3
|
||||
|
||||
__all__ = ["CnnPolicy", "MlpPolicy", "MultiInputPolicy", "TD3"]
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedul
|
|||
from stable_baselines3.common.utils import get_parameters_by_name, polyak_update
|
||||
from stable_baselines3.td3.policies import CnnPolicy, MlpPolicy, MultiInputPolicy, TD3Policy
|
||||
|
||||
TD3Self = TypeVar("TD3Self", bound="TD3")
|
||||
SelfTD3 = TypeVar("SelfTD3", bound="TD3")
|
||||
|
||||
|
||||
class TD3(OffPolicyAlgorithm):
|
||||
|
|
@ -150,7 +150,6 @@ class TD3(OffPolicyAlgorithm):
|
|||
self._update_learning_rate([self.actor.optimizer, self.critic.optimizer])
|
||||
|
||||
actor_losses, critic_losses = [], []
|
||||
|
||||
for _ in range(gradient_steps):
|
||||
|
||||
self._n_updates += 1
|
||||
|
|
@ -203,14 +202,14 @@ class TD3(OffPolicyAlgorithm):
|
|||
self.logger.record("train/critic_loss", np.mean(critic_losses))
|
||||
|
||||
def learn(
|
||||
self: TD3Self,
|
||||
self: SelfTD3,
|
||||
total_timesteps: int,
|
||||
callback: MaybeCallback = None,
|
||||
log_interval: int = 4,
|
||||
tb_log_name: str = "TD3",
|
||||
reset_num_timesteps: bool = True,
|
||||
progress_bar: bool = False,
|
||||
) -> TD3Self:
|
||||
) -> SelfTD3:
|
||||
|
||||
return super().learn(
|
||||
total_timesteps=total_timesteps,
|
||||
|
|
|
|||
|
|
@ -1 +1 @@
|
|||
1.7.0a1
|
||||
1.7.0a5
|
||||
|
|
|
|||
|
|
@ -102,7 +102,7 @@ def test_callbacks(tmp_path, model_class):
|
|||
|
||||
def select_env(model_class) -> str:
|
||||
if model_class is DQN:
|
||||
return "CartPole-v0"
|
||||
return "CartPole-v1"
|
||||
else:
|
||||
return "Pendulum-v1"
|
||||
|
||||
|
|
|
|||
|
|
@ -55,7 +55,7 @@ def test_squashed_gaussian(model_class):
|
|||
|
||||
|
||||
@pytest.fixture()
|
||||
def dummy_model_distribution_obs_and_actions() -> Tuple[A2C, np.array, np.array]:
|
||||
def dummy_model_distribution_obs_and_actions() -> Tuple[A2C, np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Fixture creating a Pendulum-v1 gym env, an A2C model and sampling 10 random observations and actions from the env
|
||||
:return: A2C model, random observations, random actions
|
||||
|
|
@ -165,9 +165,7 @@ 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(np.array([N_ACTIONS, N_ACTIONS])).proba_distribution(
|
||||
th.rand(1, sum([N_ACTIONS, N_ACTIONS]))
|
||||
),
|
||||
MultiCategoricalDistribution([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])
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ ENV_CLASSES = [
|
|||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("env_id", ["CartPole-v0", "Pendulum-v1"])
|
||||
@pytest.mark.parametrize("env_id", ["CartPole-v1", "Pendulum-v1"])
|
||||
def test_env(env_id):
|
||||
"""
|
||||
Check that environmnent integrated in Gym pass the test.
|
||||
|
|
|
|||
|
|
@ -40,7 +40,7 @@ def test_auto_wrap(model_class):
|
|||
"""Test auto wrapping of env into a VecEnv."""
|
||||
# Use different environment for DQN
|
||||
if model_class is DQN:
|
||||
env_name = "CartPole-v0"
|
||||
env_name = "CartPole-v1"
|
||||
else:
|
||||
env_name = "Pendulum-v1"
|
||||
env = gym.make(env_name)
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ def select_env(model_class: BaseAlgorithm) -> gym.Env:
|
|||
if model_class == DQN:
|
||||
return IdentityEnv(10)
|
||||
else:
|
||||
return IdentityEnvBox(10)
|
||||
return IdentityEnvBox(-10, 10)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_class", MODEL_LIST)
|
||||
|
|
|
|||
|
|
@ -45,6 +45,16 @@ def test_make_vec_env(env_id, n_envs, vec_env_cls, wrapper_class):
|
|||
env.close()
|
||||
|
||||
|
||||
def test_make_vec_env_func_checker():
|
||||
"""The functions in ``env_fns'' must return distinct instances since we need distinct environments."""
|
||||
env = gym.make("CartPole-v1")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
make_vec_env(lambda: env, n_envs=2)
|
||||
|
||||
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)])
|
||||
|
|
@ -238,7 +248,7 @@ def test_evaluate_policy_monitors(vec_env_class):
|
|||
# Also test VecEnvs
|
||||
n_eval_episodes = 3
|
||||
n_envs = 2
|
||||
env_id = "CartPole-v0"
|
||||
env_id = "CartPole-v1"
|
||||
model = A2C("MlpPolicy", env_id, seed=0)
|
||||
|
||||
def make_eval_env(with_monitor, wrapper_class=gym.Wrapper):
|
||||
|
|
@ -288,13 +298,13 @@ def test_evaluate_policy_monitors(vec_env_class):
|
|||
episode_rewards, episode_lengths = evaluate_policy(
|
||||
model, eval_env, n_eval_episodes, return_episode_rewards=True, warn=False
|
||||
)
|
||||
assert all(map(lambda l: l == 1, episode_lengths)), "AlwaysDoneWrapper did not fix episode lengths to one"
|
||||
assert all(map(lambda length: length == 1, episode_lengths)), "AlwaysDoneWrapper did not fix episode lengths to one"
|
||||
eval_env.close()
|
||||
|
||||
# Should get longer episodes with with Monitor (true episodes)
|
||||
eval_env = make_eval_env(with_monitor=True, wrapper_class=AlwaysDoneWrapper)
|
||||
episode_rewards, episode_lengths = evaluate_policy(model, eval_env, n_eval_episodes, return_episode_rewards=True)
|
||||
assert all(map(lambda l: l > 1, episode_lengths)), "evaluate_policy did not get episode lengths from Monitor"
|
||||
assert all(map(lambda length: length > 1, episode_lengths)), "evaluate_policy did not get episode lengths from Monitor"
|
||||
eval_env.close()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -62,6 +62,16 @@ class CustomGymEnv(gym.Env):
|
|||
return np.ones((dim_0, dim_1))
|
||||
|
||||
|
||||
def test_vecenv_func_checker():
|
||||
"""The functions in ``env_fns'' must return distinct instances since we need distinct environments."""
|
||||
env = CustomGymEnv(gym.spaces.Box(low=np.zeros(2), high=np.ones(2)))
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
DummyVecEnv([lambda: env for _ in range(N_ENVS)])
|
||||
|
||||
env.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("vec_env_class", VEC_ENV_CLASSES)
|
||||
@pytest.mark.parametrize("vec_env_wrapper", VEC_ENV_WRAPPERS)
|
||||
def test_vecenv_custom_calls(vec_env_class, vec_env_wrapper):
|
||||
|
|
|
|||
|
|
@ -140,7 +140,7 @@ def test_vec_monitor_ppo(recwarn):
|
|||
|
||||
# No warnings because using `VecMonitor`
|
||||
evaluate_policy(model, monitor_env)
|
||||
assert len(recwarn) == 0
|
||||
assert len(recwarn) == 0, f"{[str(warning) for warning in recwarn]}"
|
||||
|
||||
|
||||
def test_vec_monitor_warn():
|
||||
|
|
|
|||
|
|
@ -40,7 +40,7 @@ class DummyRewardEnv(gym.Env):
|
|||
return np.array([self.returned_rewards[self.return_reward_idx]])
|
||||
|
||||
|
||||
class DummyDictEnv(gym.GoalEnv):
|
||||
class DummyDictEnv(gym.Env):
|
||||
"""
|
||||
Dummy gym goal env for testing purposes
|
||||
"""
|
||||
|
|
|
|||
Loading…
Reference in a new issue