diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 47ca919..02e832a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -16,7 +16,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [3.6, 3.7] # 3.8 not supported yet by pytype + python-version: [3.6, 3.7, 3.8] steps: - uses: actions/checkout@v2 diff --git a/Dockerfile b/Dockerfile index 2ed1d23..a946a36 100644 --- a/Dockerfile +++ b/Dockerfile @@ -14,7 +14,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ libglib2.0-0 && \ rm -rf /var/lib/apt/lists/* -# Install anaconda abd dependencies +# Install Anaconda and dependencies RUN curl -o ~/miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \ chmod +x ~/miniconda.sh && \ ~/miniconda.sh -b -p /opt/conda && \ diff --git a/docs/guide/migration.rst b/docs/guide/migration.rst index 2bc8e65..fa23584 100644 --- a/docs/guide/migration.rst +++ b/docs/guide/migration.rst @@ -33,6 +33,13 @@ You can also take a look at the `rl-zoo3 `_ of SB2 to have a concrete example of successful migration. +.. note:: + + If you experience massive slow-down switching to PyTorch, you may need to play with the number of threads used, + using ``torch.set_num_threads(1)`` or ``OMP_NUM_THREADS=1``, see `issue #122 `_ + and `issue #90 `_. + + Breaking Changes ================ diff --git a/docs/guide/rl_tips.rst b/docs/guide/rl_tips.rst index c207c21..d75d347 100644 --- a/docs/guide/rl_tips.rst +++ b/docs/guide/rl_tips.rst @@ -119,14 +119,14 @@ Discrete Actions Discrete Actions - Single Process ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -DQN with extensions (double DQN, prioritized replay, ...) are the recommended algorithms. -We notably provide QR-DQN in our :ref:`contrib repo `. -DQN is usually slower to train (regarding wall clock time) but is the most sample efficient (because of its replay buffer). +``DQN`` with extensions (double DQN, prioritized replay, ...) are the recommended algorithms. +We notably provide ``QR-DQN`` in our :ref:`contrib repo `. +``DQN`` is usually slower to train (regarding wall clock time) but is the most sample efficient (because of its replay buffer). Discrete Actions - Multiprocessed ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -You should give a try to PPO or A2C. +You should give a try to ``PPO`` or ``A2C``. Continuous Actions @@ -142,7 +142,7 @@ Please use the hyperparameters in the `RL zoo `_ +Take a look at ``PPO`` or ``A2C``. Again, don't forget to take the hyperparameters from the `RL zoo `_ for continuous actions problems (cf *Bullet* envs). .. note:: @@ -155,12 +155,12 @@ Goal Environment ----------------- If your environment follows the ``GoalEnv`` interface (cf :ref:`HER `), then you should use -HER + (SAC/TD3/DDPG/DQN/TQC) depending on the action space. +HER + (SAC/TD3/DDPG/DQN/QR-DQN/TQC) depending on the action space. .. note:: - The number of workers is an important hyperparameters for experiments with HER + The ``batch_size`` is an important hyperparameter for experiments with :ref:`HER ` diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index fcde57e..c3e3253 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,17 +3,26 @@ Changelog ========== -Release 1.0rc0 (2021-02-28) +Release 1.0rc1 (WIP) ------------------------------- Breaking Changes: ^^^^^^^^^^^^^^^^^ - Removed ``stable_baselines3.common.cmd_util`` (already deprecated), please use ``env_util`` instead +New Features: +^^^^^^^^^^^^^ +- Added support for ``custom_objects`` when loading models + +Bug Fixes: +^^^^^^^^^^ +- Fixed a bug with ``DQN`` predict method when using ``deterministic=False`` with image space + Documentation: ^^^^^^^^^^^^^^ - Fixed examples - Added new project using SB3: rl_reach (@PierreExeter) +- Added note about slow-down when switching to PyTorch - Add a note on continual learning and resetting environment diff --git a/docs/modules/a2c.rst b/docs/modules/a2c.rst index 4c25040..54c8a42 100644 --- a/docs/modules/a2c.rst +++ b/docs/modules/a2c.rst @@ -53,13 +53,12 @@ Train a A2C agent on ``CartPole-v1`` using 4 environments. import gym from stable_baselines3 import A2C - from stable_baselines3.a2c import MlpPolicy from stable_baselines3.common.env_util import make_vec_env # Parallel environments - env = make_vec_env('CartPole-v1', n_envs=4) + env = make_vec_env("CartPole-v1", n_envs=4) - model = A2C(MlpPolicy, env, verbose=1) + model = A2C("MlpPolicy", env, verbose=1) model.learn(total_timesteps=25000) model.save("a2c_cartpole") diff --git a/docs/modules/ddpg.rst b/docs/modules/ddpg.rst index c14f5da..08f5bf4 100644 --- a/docs/modules/ddpg.rst +++ b/docs/modules/ddpg.rst @@ -63,13 +63,13 @@ Example from stable_baselines3 import DDPG from stable_baselines3.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise - env = gym.make('Pendulum-v0') + env = gym.make("Pendulum-v0") # The noise objects for DDPG n_actions = env.action_space.shape[-1] action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions)) - model = DDPG('MlpPolicy', env, action_noise=action_noise, verbose=1) + model = DDPG("MlpPolicy", env, action_noise=action_noise, verbose=1) model.learn(total_timesteps=10000, log_interval=10) model.save("ddpg_pendulum") env = model.get_env() diff --git a/docs/modules/dqn.rst b/docs/modules/dqn.rst index 8b97afa..76490dd 100644 --- a/docs/modules/dqn.rst +++ b/docs/modules/dqn.rst @@ -56,11 +56,10 @@ Example import numpy as np from stable_baselines3 import DQN - from stable_baselines3.dqn import MlpPolicy - env = gym.make('CartPole-v0') + env = gym.make("CartPole-v0") - model = DQN(MlpPolicy, env, verbose=1) + model = DQN("MlpPolicy", env, verbose=1) model.learn(total_timesteps=10000, log_interval=4) model.save("dqn_pendulum") diff --git a/docs/modules/her.rst b/docs/modules/her.rst index 6894c01..fc0c696 100644 --- a/docs/modules/her.rst +++ b/docs/modules/her.rst @@ -52,7 +52,7 @@ Notes Can I use? ---------- -Please refer to the used model (DQN, SAC, TD3 or DDPG) for that section. +Please refer to the used model (DQN, QR-DQN, SAC, TQC, TD3, or DDPG) for that section. Example ------- diff --git a/docs/modules/ppo.rst b/docs/modules/ppo.rst index 091323a..09e2e63 100644 --- a/docs/modules/ppo.rst +++ b/docs/modules/ppo.rst @@ -54,13 +54,12 @@ Train a PPO agent on ``Pendulum-v0`` using 4 environments. import gym from stable_baselines3 import PPO - from stable_baselines3.ppo import MlpPolicy from stable_baselines3.common.env_util import make_vec_env # Parallel environments - env = make_vec_env('CartPole-v1', n_envs=4) + env = make_vec_env("CartPole-v1", n_envs=4) - model = PPO(MlpPolicy, env, verbose=1) + model = PPO("MlpPolicy", env, verbose=1) model.learn(total_timesteps=25000) model.save("ppo_cartpole") diff --git a/docs/modules/sac.rst b/docs/modules/sac.rst index bbe6bfc..63d4245 100644 --- a/docs/modules/sac.rst +++ b/docs/modules/sac.rst @@ -68,11 +68,10 @@ Example import numpy as np from stable_baselines3 import SAC - from stable_baselines3.sac import MlpPolicy - env = gym.make('Pendulum-v0') + env = gym.make("Pendulum-v0") - model = SAC(MlpPolicy, env, verbose=1) + model = SAC("MlpPolicy", env, verbose=1) model.learn(total_timesteps=10000, log_interval=4) model.save("sac_pendulum") diff --git a/docs/modules/td3.rst b/docs/modules/td3.rst index 2ecf8c9..7a1a0f3 100644 --- a/docs/modules/td3.rst +++ b/docs/modules/td3.rst @@ -61,16 +61,15 @@ Example import numpy as np from stable_baselines3 import TD3 - from stable_baselines3.td3.policies import MlpPolicy from stable_baselines3.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise - env = gym.make('Pendulum-v0') + env = gym.make("Pendulum-v0") # The noise objects for TD3 n_actions = env.action_space.shape[-1] action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions)) - model = TD3(MlpPolicy, env, action_noise=action_noise, verbose=1) + model = TD3("MlpPolicy", env, action_noise=action_noise, verbose=1) model.learn(total_timesteps=10000, log_interval=10) model.save("td3_pendulum") env = model.get_env() diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 42669bf..43edc2f 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -119,3 +119,8 @@ cuda Polyak gSDE rollouts +Pyro +softmax +stdout +Contrib +Quantile diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 6be7ca3..06bef17 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -586,6 +586,7 @@ class BaseAlgorithm(ABC): path: Union[str, pathlib.Path, io.BufferedIOBase], env: Optional[GymEnv] = None, device: Union[th.device, str] = "auto", + custom_objects: Optional[Dict[str, Any]] = None, **kwargs, ) -> "BaseAlgorithm": """ @@ -596,9 +597,15 @@ class BaseAlgorithm(ABC): :param env: the new environment to run the loaded model on (can be None if you only need prediction from a trained model) has priority over any saved environment :param device: Device on which the code should run. + :param custom_objects: Dictionary of objects to replace + upon loading. If a variable is present in this dictionary as a + key, it will not be deserialized and the corresponding item + will be used instead. Similar to custom_objects in + ``keras.models.load_model``. Useful when you have an object in + file that can not be deserialized. :param kwargs: extra arguments to change the model when loading """ - data, params, pytorch_variables = load_from_zip_file(path, device=device) + data, params, pytorch_variables = load_from_zip_file(path, device=device, custom_objects=custom_objects) # Remove stored device information and replace with ours if "policy_kwargs" in data: @@ -625,7 +632,7 @@ class BaseAlgorithm(ABC): env = data["env"] # noinspection PyArgumentList - model = cls( + model = cls( # pytype: disable=not-instantiable,wrong-keyword-args policy=data["policy_class"], env=env, device=device, diff --git a/stable_baselines3/common/distributions.py b/stable_baselines3/common/distributions.py index 7a3506c..ebe06fa 100644 --- a/stable_baselines3/common/distributions.py +++ b/stable_baselines3/common/distributions.py @@ -623,7 +623,7 @@ class TanhBijector(object): """ Inverse of Tanh - Taken from pyro: https://github.com/pyro-ppl/pyro + Taken from Pyro: https://github.com/pyro-ppl/pyro 0.5 * torch.log((1 + x ) / (1 - x)) """ return 0.5 * (x.log1p() - (-x).log1p()) diff --git a/stable_baselines3/common/evaluation.py b/stable_baselines3/common/evaluation.py index cb773bb..a352000 100644 --- a/stable_baselines3/common/evaluation.py +++ b/stable_baselines3/common/evaluation.py @@ -41,7 +41,7 @@ def evaluate_policy( called after each step. Gets locals() and globals() passed as parameters. :param reward_threshold: Minimum expected reward per episode, this will raise an error if the performance is not met - :param return_episode_rewards: If True, a list of rewards and episde lengths + :param return_episode_rewards: If True, a list of rewards and episode lengths per episode will be returned instead of the mean. :param warn: If True (default), warns user about lack of a Monitor wrapper in the evaluation environment. diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 164174b..7c7412f 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -174,7 +174,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): self.device, optimize_memory_usage=self.optimize_memory_usage, ) - self.policy = self.policy_class( + self.policy = self.policy_class( # pytype:disable=not-instantiable self.observation_space, self.action_space, self.lr_schedule, diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 0a930d9..c55acb8 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -114,7 +114,7 @@ class OnPolicyAlgorithm(BaseAlgorithm): gae_lambda=self.gae_lambda, n_envs=self.n_envs, ) - self.policy = self.policy_class( + self.policy = self.policy_class( # pytype:disable=not-instantiable self.observation_space, self.action_space, self.lr_schedule, diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 543e6c7..ea6fa97 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -19,11 +19,10 @@ from stable_baselines3.common.distributions import ( StateDependentNoiseDistribution, make_proba_distribution, ) -from stable_baselines3.common.preprocessing import get_action_dim, is_image_space, preprocess_obs +from stable_baselines3.common.preprocessing import get_action_dim, maybe_transpose, preprocess_obs from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, MlpExtractor, NatureCNN, create_mlp from stable_baselines3.common.type_aliases import Schedule from stable_baselines3.common.utils import get_device, is_vectorized_observation -from stable_baselines3.common.vec_env import VecTransposeImage from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper @@ -266,17 +265,7 @@ class BasePolicy(BaseModel): # Handle the different cases for images # as PyTorch use channel first format - if is_image_space(self.observation_space): - if not ( - observation.shape == self.observation_space.shape or observation.shape[1:] == self.observation_space.shape - ): - # Try to re-order the channels - transpose_obs = VecTransposeImage.transpose_image(observation) - if ( - transpose_obs.shape == self.observation_space.shape - or transpose_obs.shape[1:] == self.observation_space.shape - ): - observation = transpose_obs + observation = maybe_transpose(observation, self.observation_space) vectorized_env = is_vectorized_observation(observation, self.observation_space) diff --git a/stable_baselines3/common/preprocessing.py b/stable_baselines3/common/preprocessing.py index 881970a..63052f3 100644 --- a/stable_baselines3/common/preprocessing.py +++ b/stable_baselines3/common/preprocessing.py @@ -61,6 +61,26 @@ def is_image_space(observation_space: spaces.Space, channels_last: bool = True, return False +def maybe_transpose(observation: np.ndarray, observation_space: spaces.Space) -> np.ndarray: + """ + Handle the different cases for images as PyTorch use channel first format. + + :param observation: + :param observation_space: + :return: channel first observation if observation is an image + """ + # Avoid circular import + from stable_baselines3.common.vec_env import VecTransposeImage + + if is_image_space(observation_space): + if not (observation.shape == observation_space.shape or observation.shape[1:] == observation_space.shape): + # Try to re-order the channels + transpose_obs = VecTransposeImage.transpose_image(observation) + if transpose_obs.shape == observation_space.shape or transpose_obs.shape[1:] == observation_space.shape: + observation = transpose_obs + return observation + + def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space, normalize_images: bool = True) -> th.Tensor: """ Preprocess observation to be to a neural network. diff --git a/stable_baselines3/common/save_util.py b/stable_baselines3/common/save_util.py index 4b2f1d6..d41c79c 100644 --- a/stable_baselines3/common/save_util.py +++ b/stable_baselines3/common/save_util.py @@ -137,7 +137,7 @@ def json_to_data(json_string: str, custom_objects: Optional[Dict[str, Any]] = No upon loading. If a variable is present in this dictionary as a key, it will not be deserialized and the corresponding item will be used instead. Similar to custom_objects in - `keras.models.load_model`. Useful when you have an object in + ``keras.models.load_model``. Useful when you have an object in file that can not be deserialized. :return: Loaded class parameters. """ @@ -162,7 +162,7 @@ def json_to_data(json_string: str, custom_objects: Optional[Dict[str, Any]] = No try: base64_object = base64.b64decode(serialization.encode()) deserialized_object = cloudpickle.loads(base64_object) - except RuntimeError: + except (RuntimeError, TypeError): warnings.warn( f"Could not deserialize object {data_key}. " + "Consider using `custom_objects` argument to replace " @@ -359,6 +359,7 @@ def load_from_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], verbose: in def load_from_zip_file( load_path: Union[str, pathlib.Path, io.BufferedIOBase], load_data: bool = True, + custom_objects: Optional[Dict[str, Any]] = None, device: Union[th.device, str] = "auto", verbose: int = 0, ) -> (Tuple[Optional[Dict[str, Any]], Optional[TensorDict], Optional[TensorDict]]): @@ -368,6 +369,12 @@ def load_from_zip_file( :param load_path: Where to load the model from :param load_data: Whether we should load and return data (class parameters). Mainly used by 'load_parameters' to only load model parameters (weights) + :param custom_objects: Dictionary of objects to replace + upon loading. If a variable is present in this dictionary as a + key, it will not be deserialized and the corresponding item + will be used instead. Similar to custom_objects in + ``keras.models.load_model``. Useful when you have an object in + file that can not be deserialized. :param device: Device on which the code should run. :return: Class parameters, model state_dicts (aka "params", dict of state_dict) and dict of pytorch variables @@ -392,7 +399,7 @@ def load_from_zip_file( # Load class parameters that are stored # with either JSON or pickle (not PyTorch variables). json_data = archive.read("data").decode() - data = json_to_data(json_data) + data = json_to_data(json_data, custom_objects=custom_objects) # Check for all .pth files and load them using th.load. # "pytorch_variables.pth" stores PyTorch variables, and any other .pth diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 0972993..87ef92a 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -21,7 +21,8 @@ from stable_baselines3.common.type_aliases import GymEnv, Schedule, TrainFreq, T def set_random_seed(seed: int, using_cuda: bool = False) -> None: """ - Seed the different random generators + Seed the different random generators. + :param seed: :param using_cuda: """ diff --git a/stable_baselines3/common/vec_env/obs_dict_wrapper.py b/stable_baselines3/common/vec_env/obs_dict_wrapper.py index 5b1dd1a..d07ad24 100644 --- a/stable_baselines3/common/vec_env/obs_dict_wrapper.py +++ b/stable_baselines3/common/vec_env/obs_dict_wrapper.py @@ -61,8 +61,8 @@ class ObsDictWrapper(VecEnvWrapper): Concatenate observation and (desired) goal of observation dict. :param observation_dict: Dictionary with observation. - :param observation_key: Key of observation in dicitonary. - :param goal_key: Key of (desired) goal in dicitonary. + :param observation_key: Key of observation in dictionary. + :param goal_key: Key of (desired) goal in dictionary. :return: Concatenated observation. """ return np.concatenate([observation_dict[observation_key], observation_dict[goal_key]], axis=-1) diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 77e9d30..241f9ae 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -7,6 +7,7 @@ from torch.nn import functional as F from stable_baselines3.common import logger from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm +from stable_baselines3.common.preprocessing import maybe_transpose from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule from stable_baselines3.common.utils import get_linear_fn, is_vectorized_observation, polyak_update from stable_baselines3.dqn.policies import DQNPolicy @@ -201,7 +202,7 @@ class DQN(OffPolicyAlgorithm): (used in recurrent policies) """ if not deterministic and np.random.rand() < self.exploration_rate: - if is_vectorized_observation(observation, self.observation_space): + if is_vectorized_observation(maybe_transpose(observation, self.observation_space), self.observation_space): n_batch = observation.shape[0] action = np.array([self.action_space.sample() for _ in range(n_batch)]) else: diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index f4fb83c..0f82de4 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.0rc0 +1.0rc1 diff --git a/tests/test_cnn.py b/tests/test_cnn.py index 5c5f8cf..c7d5cc4 100644 --- a/tests/test_cnn.py +++ b/tests/test_cnn.py @@ -25,7 +25,11 @@ def test_cnn(tmp_path, model_class): else: # Avoid memory error when using replay buffer # Reduce the size of the features - kwargs = dict(buffer_size=250, policy_kwargs=dict(features_extractor_kwargs=dict(features_dim=32))) + kwargs = dict( + buffer_size=250, + policy_kwargs=dict(features_extractor_kwargs=dict(features_dim=32)), + seed=1, + ) model = model_class("CnnPolicy", env, **kwargs).learn(250) # FakeImageEnv is channel last by default and should be wrapped @@ -33,6 +37,13 @@ def test_cnn(tmp_path, model_class): obs = env.reset() + # Test stochastic predict with channel last input + if model_class == DQN: + model.exploration_rate = 0.9 + + for _ in range(10): + model.predict(obs, deterministic=False) + action, _ = model.predict(obs, deterministic=True) model.save(tmp_path / SAVE_NAME) diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 369b82a..1815877 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -183,17 +183,17 @@ def test_set_env(model_class): # create model model = model_class("MlpPolicy", env, policy_kwargs=dict(net_arch=[16]), **kwargs) # learn - model.learn(total_timesteps=300) + model.learn(total_timesteps=128) # change env model.set_env(env2) # learn again - model.learn(total_timesteps=300) + model.learn(total_timesteps=128) # change env test wrapping model.set_env(env3) # learn again - model.learn(total_timesteps=300) + model.learn(total_timesteps=128) @pytest.mark.parametrize("model_class", MODEL_LIST) @@ -220,8 +220,14 @@ def test_exclude_include_saved_params(tmp_path, model_class): # Check if include works model.save(tmp_path / "test_save", exclude=["verbose"], include=["verbose"]) del model - model = model_class.load(str(tmp_path / "test_save.zip")) + # Load with custom objects + custom_objects = dict(learning_rate=2e-5, dummy=1.0) + model = model_class.load(str(tmp_path / "test_save.zip"), custom_objects=custom_objects) assert model.verbose == 2 + # Check that the custom object was taken into account + assert model.learning_rate == custom_objects["learning_rate"] + # Check that only parameters that are here already are replaced + assert not hasattr(model, "dummy") # clear file from os os.remove(tmp_path / "test_save.zip")