Merge branch 'master' into nstep

docs/guide/custom_policy.rst

@@ -16,7 +16,7 @@ using ``policy_kwargs`` parameter:
 
   from stable_baselines3 import PPO
 
-  # Custom MLP policy of two layers of size 32 each with tanh activation function
+  # Custom MLP policy of two layers of size 32 each with Relu activation function
   policy_kwargs = dict(activation_fn=th.nn.ReLU, net_arch=[32, 32])
   # Create the agent
   model = PPO("MlpPolicy", "CartPole-v1", policy_kwargs=policy_kwargs, verbose=1)

docs/misc/changelog.rst

@@ -3,9 +3,11 @@
 Changelog
 ==========
 
-Pre-Release 0.8.0a4 (WIP)
+Pre-Release 0.8.0 (2020-08-03)
 ------------------------------
 
+**DQN, DDPG, bug fixes and performance matching for Atari games**
+
 Breaking Changes:
 ^^^^^^^^^^^^^^^^^
 - ``AtariWrapper`` and other Atari wrappers were updated to match SB2 ones
@@ -31,6 +33,12 @@ Bug Fixes:
 ^^^^^^^^^^
 - Fixed a bug in the ``close()`` method of ``SubprocVecEnv``, causing wrappers further down in the wrapper stack to not be closed. (@NeoExtended)
 - Fix target for updating q values in SAC: the entropy term was not conditioned by terminals states
+- Use ``cloudpickle.load`` instead of ``pickle.load`` in ``CloudpickleWrapper``. (@shwang)
+- Fixed a bug with orthogonal initialization when `bias=False` in custom policy (@rk37)
+- Fixed approximate entropy calculation in PPO and A2C. (@andyshih12)
+- Fixed DQN target network sharing feature extractor with the main network.
+- Fixed storing correct ``dones`` in on-policy algorithm rollout collection. (@andyshih12)
+- Fixed number of filters in final convolutional layer in NatureCNN to match original implementation.
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -41,10 +49,12 @@ Others:
 - Split the ``collect_rollout()`` method for off-policy algorithms
 - Added ``_on_step()`` for off-policy base class
 - Optimized replay buffer size by removing the need of ``next_observations`` numpy array
+- Optimized polyak updates (1.5-1.95 speedup) through inplace operations (@PartiallyTyped)
 - Switch to ``black`` codestyle and added ``make format``, ``make check-codestyle`` and ``commit-checks``
 - Ignored errors from newer pytype version
 - Added a check when using ``gSDE``
 - Removed codacy dependency from Dockerfile
+- Added ``common.sb2_compat.RMSpropTFLike`` optimizer, which corresponds closer to the implementation of RMSprop from Tensorflow.
 
 Documentation:
 ^^^^^^^^^^^^^^
@@ -53,6 +63,7 @@ Documentation:
 - Added Unity reacher to the projects page (@koulakis)
 - Added PyBullet colab notebook
 - Fixed typo in PPO example code (@joeljosephjin)
+- Fixed typo in custom policy doc (@RaphaelWag)
 
 
 
@@ -355,4 +366,4 @@ And all the contributors:
 @Miffyli @dwiel @miguelrass @qxcv @jaberkow @eavelardev @ruifeng96150 @pedrohbtp @srivatsankrishnan @evilsocket
 @MarvineGothic @jdossgollin @SyllogismRXS @rusu24edward @jbulow @Antymon @seheevic @justinkterry @edbeeching
 @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3
-@tirafesi @blurLake @koulakis @joeljosephjin
+@tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag

docs/modules/a2c.rst

@@ -10,6 +10,14 @@ A synchronous, deterministic variant of `Asynchronous Advantage Actor Critic (A3
 It uses multiple workers to avoid the use of a replay buffer.
 
 
+.. warning::
+  
+  If you find training unstable or want to match performance of stable-baselines A2C, consider using
+  ``RMSpropTFLike`` optimizer from ``stable_baselines3.common.sb2_compat.rmsprop_tf_like``.
+  You can change optimizer with ``A2C(policy_kwargs=dict(optimizer_class=RMSpropTFLike))``.
+  Read more `here <https://github.com/DLR-RM/stable-baselines3/pull/110#issuecomment-663255241>`_.
+
+
 Notes
 -----
 

stable_baselines3/a2c/a2c.py

@@ -116,6 +116,7 @@ class A2C(OnPolicyAlgorithm):
         # Update optimizer learning rate
         self._update_learning_rate(self.policy.optimizer)
 
+        # This will only loop once (get all data in one go)
         for rollout_data in self.rollout_buffer.get(batch_size=None):
 
             actions = rollout_data.actions
@@ -141,7 +142,7 @@ class A2C(OnPolicyAlgorithm):
             # Entropy loss favor exploration
             if entropy is None:
                 # Approximate entropy when no analytical form
-                entropy_loss = -log_prob.mean()
+                entropy_loss = -th.mean(-log_prob)
             else:
                 entropy_loss = -th.mean(entropy)
 

stable_baselines3/common/base_class.py

@@ -123,6 +123,7 @@ class BaseAlgorithm(ABC):
         self.tensorboard_log = tensorboard_log
         self.lr_schedule = None  # type: Optional[Callable]
         self._last_obs = None  # type: Optional[np.ndarray]
+        self._last_dones = None  # type: Optional[np.ndarray]
         # When using VecNormalize:
         self._last_original_obs = None  # type: Optional[np.ndarray]
         self._episode_num = 0
@@ -474,6 +475,7 @@ class BaseAlgorithm(ABC):
         # Avoid resetting the environment when calling ``.learn()`` consecutive times
         if reset_num_timesteps or self._last_obs is None:
             self._last_obs = self.env.reset()
+            self._last_dones = np.zeros((self.env.num_envs,), dtype=np.bool)
             # Retrieve unnormalized observation for saving into the buffer
             if self._vec_normalize_env is not None:
                 self._last_original_obs = self._vec_normalize_env.get_original_obs()

stable_baselines3/common/on_policy_algorithm.py

@@ -173,8 +173,9 @@ class OnPolicyAlgorithm(BaseAlgorithm):
             if isinstance(self.action_space, gym.spaces.Discrete):
                 # Reshape in case of discrete action
                 actions = actions.reshape(-1, 1)
-            rollout_buffer.add(self._last_obs, actions, rewards, dones, values, log_probs)
+            rollout_buffer.add(self._last_obs, actions, rewards, self._last_dones, values, log_probs)
             self._last_obs = new_obs
+            self._last_dones = dones
 
         rollout_buffer.compute_returns_and_advantage(values, dones=dones)
 

stable_baselines3/common/policies.py

@@ -188,7 +188,8 @@ class BasePolicy(BaseModel):
         """
         if isinstance(module, (nn.Linear, nn.Conv2d)):
             nn.init.orthogonal_(module.weight, gain=gain)
-            module.bias.data.fill_(0.0)
+            if module.bias is not None:
+                module.bias.data.fill_(0.0)
 
     @abstractmethod
     def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:

stable_baselines3/common/sb2_compat/rmsprop_tf_like.py

@@ -0,0 +1,126 @@
+import torch
+from torch.optim import Optimizer
+
+
+class RMSpropTFLike(Optimizer):
+    r"""Implements RMSprop algorithm with closer match to Tensorflow version.
+
+    For reproducibility with original stable-baselines. Use this
+    version with e.g. A2C for stabler learning than with the PyTorch
+    RMSProp. Based on the PyTorch v1.5.0 implementation of RMSprop.
+
+    See a more throughout conversion in pytorch-image-models repository:
+        https://github.com/rwightman/pytorch-image-models/blob/master/timm/optim/rmsprop_tf.py
+
+    Changes to the original RMSprop:
+        - Move epsilon inside square root
+        - Initialize squared gradient to ones rather than zeros
+
+    Proposed by G. Hinton in his
+    `course <http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf>`_.
+
+    The centered version first appears in `Generating Sequences
+    With Recurrent Neural Networks <https://arxiv.org/pdf/1308.0850v5.pdf>`_.
+
+    The implementation here takes the square root of the gradient average before
+    adding epsilon (note that TensorFlow interchanges these two operations). The effective
+    learning rate is thus :math:`\alpha/(\sqrt{v} + \epsilon)` where :math:`\alpha`
+    is the scheduled learning rate and :math:`v` is the weighted moving average
+    of the squared gradient.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-2)
+        momentum (float, optional): momentum factor (default: 0)
+        alpha (float, optional): smoothing constant (default: 0.99)
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        centered (bool, optional) : if ``True``, compute the centered RMSProp,
+            the gradient is normalized by an estimation of its variance
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+
+    """
+
+    def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= momentum:
+            raise ValueError("Invalid momentum value: {}".format(momentum))
+        if not 0.0 <= weight_decay:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+        if not 0.0 <= alpha:
+            raise ValueError("Invalid alpha value: {}".format(alpha))
+
+        defaults = dict(lr=lr, momentum=momentum, alpha=alpha, eps=eps, centered=centered, weight_decay=weight_decay)
+        super(RMSpropTFLike, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(RMSpropTFLike, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("momentum", 0)
+            group.setdefault("centered", False)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError("RMSpropTF does not support sparse gradients")
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # PyTorch initialized to zeros here
+                    state["square_avg"] = torch.ones_like(p, memory_format=torch.preserve_format)
+                    if group["momentum"] > 0:
+                        state["momentum_buffer"] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    if group["centered"]:
+                        state["grad_avg"] = torch.zeros_like(p, memory_format=torch.preserve_format)
+
+                square_avg = state["square_avg"]
+                alpha = group["alpha"]
+
+                state["step"] += 1
+
+                if group["weight_decay"] != 0:
+                    grad = grad.add(p, alpha=group["weight_decay"])
+
+                square_avg.mul_(alpha).addcmul_(grad, grad, value=1 - alpha)
+
+                if group["centered"]:
+                    grad_avg = state["grad_avg"]
+                    grad_avg.mul_(alpha).add_(grad, alpha=1 - alpha)
+                    # PyTorch added epsilon after square root
+                    # avg = square_avg.addcmul(grad_avg, grad_avg, value=-1).sqrt_().add_(group['eps'])
+                    avg = square_avg.addcmul(grad_avg, grad_avg, value=-1).add_(group["eps"]).sqrt_()
+                else:
+                    # PyTorch added epsilon after square root
+                    # avg = square_avg.sqrt().add_(group['eps'])
+                    avg = square_avg.add(group["eps"]).sqrt_()
+
+                if group["momentum"] > 0:
+                    buf = state["momentum_buffer"]
+                    buf.mul_(group["momentum"]).addcdiv_(grad, avg)
+                    p.add_(buf, alpha=-group["lr"])
+                else:
+                    p.addcdiv_(grad, avg, value=-group["lr"])
+
+        return loss

stable_baselines3/common/torch_layers.py

@@ -74,7 +74,7 @@ class NatureCNN(BaseFeaturesExtractor):
             nn.ReLU(),
             nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=0),
             nn.ReLU(),
-            nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=0),
+            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=0),
             nn.ReLU(),
             nn.Flatten(),
         )

stable_baselines3/common/utils.py

@@ -2,7 +2,7 @@ import glob
 import os
 import random
 from collections import deque
-from typing import Callable, Optional, Union
+from typing import Callable, Iterable, Optional, Union
 
 import gym
 import numpy as np
@@ -284,3 +284,25 @@ def safe_mean(arr: Union[np.ndarray, list, deque]) -> np.ndarray:
     :return:
     """
     return np.nan if len(arr) == 0 else np.mean(arr)
+
+
+def polyak_update(params: Iterable[th.nn.Parameter], target_params: Iterable[th.nn.Parameter], tau: float) -> None:
+    """
+    Perform a Polyak average update on ``target_params`` using ``params``:
+    target parameters are slowly updated towards the main parameters.
+    ``tau``, the soft update coefficient controls the interpolation:
+    ``tau=1`` corresponds to copying the parameters to the target ones whereas nothing happens when ``tau=0``.
+    The Polyak update is done in place, with ``no_grad``, and therefore does not create intermediate tensors,
+    or a computation graph, reducing memory cost and improving performance.  We scale the target params
+    by ``1-tau`` (in-place), add the new weights, scaled by ``tau`` and store the result of the sum in the target
+    params (in place).
+    See https://github.com/DLR-RM/stable-baselines3/issues/93
+
+    :param params: (Iterable[th.nn.Parameter]) parameters to use to update the target params
+    :param target_params: (Iterable[th.nn.Parameter]) parameters to update
+    :param tau: (float) the soft update coefficient ("Polyak update", between 0 and 1)
+    """
+    with th.no_grad():
+        for param, target_param in zip(params, target_params):
+            target_param.data.mul_(1 - tau)
+            th.add(target_param.data, param.data, alpha=tau, out=target_param.data)

stable_baselines3/common/vec_env/base_vec_env.py

@@ -1,5 +1,4 @@
 import inspect
-import pickle
 from abc import ABC, abstractmethod
 from typing import List, Optional, Sequence, Union
 
@@ -349,7 +348,7 @@ class VecEnvWrapper(VecEnv):
         return shadowed_wrapper_class
 
 
-class CloudpickleWrapper(object):
+class CloudpickleWrapper:
     def __init__(self, var):
         """
         Uses cloudpickle to serialize contents (otherwise multiprocessing tries to use pickle)
@@ -362,4 +361,4 @@ class CloudpickleWrapper(object):
         return cloudpickle.dumps(self.var)
 
     def __setstate__(self, obs):
-        self.var = pickle.loads(obs)
+        self.var = cloudpickle.loads(obs)

stable_baselines3/dqn/dqn.py

@@ -8,7 +8,7 @@ from stable_baselines3.common import logger
 from stable_baselines3.common.buffers import ReplayBuffer
 from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
 from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
-from stable_baselines3.common.utils import get_linear_fn
+from stable_baselines3.common.utils import get_linear_fn, polyak_update
 from stable_baselines3.dqn.policies import DQNPolicy
 
 
@@ -141,8 +141,7 @@ class DQN(OffPolicyAlgorithm):
         This method is called in ``collect_rollout()`` after each step in the environment.
         """
         if self.num_timesteps % self.target_update_interval == 0:
-            for param, target_param in zip(self.q_net.parameters(), self.q_net_target.parameters()):
-                target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+            polyak_update(self.q_net.parameters(), self.q_net_target.parameters(), self.tau)
 
         self.exploration_rate = self.exploration_schedule(self._current_progress_remaining)
         logger.record("rollout/exploration rate", self.exploration_rate)

stable_baselines3/dqn/policies.py

@@ -133,8 +133,6 @@ class DQNPolicy(BasePolicy):
             else:
                 net_arch = []
 
-        self.features_extractor = features_extractor_class(self.observation_space, **self.features_extractor_kwargs)
-        self.features_dim = self.features_extractor.features_dim
         self.net_arch = net_arch
         self.activation_fn = activation_fn
         self.normalize_images = normalize_images
@@ -142,8 +140,6 @@ class DQNPolicy(BasePolicy):
         self.net_args = {
             "observation_space": self.observation_space,
             "action_space": self.action_space,
-            "features_extractor": self.features_extractor,
-            "features_dim": self.features_dim,
             "net_arch": self.net_arch,
             "activation_fn": self.activation_fn,
             "normalize_images": normalize_images,
@@ -169,7 +165,10 @@ class DQNPolicy(BasePolicy):
         self.optimizer = self.optimizer_class(self.parameters(), lr=lr_schedule(1), **self.optimizer_kwargs)
 
     def make_q_net(self) -> QNetwork:
-        return QNetwork(**self.net_args).to(self.device)
+        # Make sure we always have separate networks for feature extractors etc
+        features_extractor = self.features_extractor_class(self.observation_space, **self.features_extractor_kwargs)
+        features_dim = features_extractor.features_dim
+        return QNetwork(features_extractor=features_extractor, features_dim=features_dim, **self.net_args).to(self.device)
 
     def forward(self, obs: th.Tensor, deterministic: bool = True) -> th.Tensor:
         return self._predict(obs, deterministic=deterministic)

stable_baselines3/ppo/ppo.py

@@ -198,7 +198,7 @@ class PPO(OnPolicyAlgorithm):
                 # Entropy loss favor exploration
                 if entropy is None:
                     # Approximate entropy when no analytical form
-                    entropy_loss = -log_prob.mean()
+                    entropy_loss = -th.mean(-log_prob)
                 else:
                     entropy_loss = -th.mean(entropy)
 

stable_baselines3/sac/sac.py

@@ -9,6 +9,7 @@ from stable_baselines3.common.buffers import ReplayBuffer
 from stable_baselines3.common.noise import ActionNoise
 from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
 from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
+from stable_baselines3.common.utils import polyak_update
 from stable_baselines3.sac.policies import SACPolicy
 
 
@@ -259,8 +260,7 @@ class SAC(OffPolicyAlgorithm):
 
             # Update target networks
             if gradient_step % self.target_update_interval == 0:
-                for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
-                    target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+                polyak_update(self.critic.parameters(), self.critic_target.parameters(), self.tau)
 
         self._n_updates += gradient_steps
 

stable_baselines3/td3/td3.py

@@ -8,6 +8,7 @@ from stable_baselines3.common.buffers import ReplayBuffer
 from stable_baselines3.common.noise import ActionNoise
 from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
 from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
+from stable_baselines3.common.utils import polyak_update
 from stable_baselines3.td3.policies import TD3Policy
 
 
@@ -169,12 +170,8 @@ class TD3(OffPolicyAlgorithm):
                 actor_loss.backward()
                 self.actor.optimizer.step()
 
-                # Update the frozen target networks
-                for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
-                    target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
-
-                for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):
-                    target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
+                polyak_update(self.critic.parameters(), self.critic_target.parameters(), self.tau)
+                polyak_update(self.actor.parameters(), self.actor_target.parameters(), self.tau)
 
         self._n_updates += gradient_steps
         logger.record("train/n_updates", self._n_updates, exclude="tensorboard")

stable_baselines3/version.txt

@@ -1 +1 @@
-0.8.0a4
+0.8.0

tests/test_custom_policy.py

@@ -2,6 +2,7 @@ import pytest
 import torch as th
 
 from stable_baselines3 import A2C, PPO, SAC, TD3
+from stable_baselines3.common.sb2_compat.rmsprop_tf_like import RMSpropTFLike
 
 
 @pytest.mark.parametrize(
@@ -32,3 +33,8 @@ def test_custom_offpolicy(model_class, net_arch):
 def test_custom_optimizer(model_class, optimizer_kwargs):
     policy_kwargs = dict(optimizer_class=th.optim.AdamW, optimizer_kwargs=optimizer_kwargs, net_arch=[32])
     _ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs).learn(1000)
+
+
+def test_tf_like_rmsprop_optimizer():
+    policy_kwargs = dict(optimizer_class=RMSpropTFLike, net_arch=[32])
+    _ = A2C("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs).learn(1000)

tests/test_utils.py

@@ -4,6 +4,7 @@ import shutil
 import gym
 import numpy as np
 import pytest
+import torch as th
 
 from stable_baselines3 import A2C
 from stable_baselines3.common.atari_wrappers import ClipRewardEnv
@@ -11,6 +12,7 @@ from stable_baselines3.common.cmd_util import make_atari_env, make_vec_env
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.monitor import Monitor
 from stable_baselines3.common.noise import ActionNoise, OrnsteinUhlenbeckActionNoise, VectorizedActionNoise
+from stable_baselines3.common.utils import polyak_update
 from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
 
 
@@ -152,3 +154,16 @@ def test_vec_noise():
         vec.noises = [base] * (num_envs - 1)
     assert all(isinstance(noise, type(base)) for noise in vec.noises)
     assert len(vec.noises) == num_envs
+
+
+def test_polyak():
+    param1, param2 = th.nn.Parameter(th.ones((5, 5))), th.nn.Parameter(th.zeros((5, 5)))
+    target1, target2 = th.nn.Parameter(th.ones((5, 5))), th.nn.Parameter(th.zeros((5, 5)))
+    tau = 0.1
+    polyak_update([param1], [param2], tau)
+    with th.no_grad():
+        for param, target_param in zip([target1], [target2]):
+            target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)
+
+    assert th.allclose(param1, target1)
+    assert th.allclose(param2, target2)