From 5d6649d92b2d9918f29fd8761c8fdf028c5c78e5 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 25 Nov 2019 14:54:13 +0100
Subject: [PATCH 01/47] Enable separate feature extraction for SDE

---
 tests/test_sde.py                        |  5 ++-
 torchy_baselines/a2c/a2c.py              |  1 +
 torchy_baselines/common/distributions.py | 44 +++++++++++++---------
 torchy_baselines/ppo/policies.py         | 48 +++++++++++++++++-------
 4 files changed, 66 insertions(+), 32 deletions(-)

diff --git a/tests/test_sde.py b/tests/test_sde.py
index 09b48e8..7143e9c 100644
--- a/tests/test_sde.py
+++ b/tests/test_sde.py
@@ -49,14 +49,15 @@ def test_state_dependent_exploration():
 
 
 @pytest.mark.parametrize("model_class", [A2C])
-def test_state_dependent_noise(model_class):
+@pytest.mark.parametrize("sde_net_arch", [None, [64, 64]])
+def test_state_dependent_noise(model_class, sde_net_arch):
     env_id = 'MountainCarContinuous-v0'
 
     env = VecNormalize(DummyVecEnv([lambda: Monitor(gym.make(env_id))]), norm_reward=True)
     eval_env = VecNormalize(DummyVecEnv([lambda: Monitor(gym.make(env_id))]), training=False, norm_reward=False)
 
     model = model_class('MlpPolicy', env, n_steps=200, use_sde=True, ent_coef=0.00, verbose=1, learning_rate=3e-4,
-                        policy_kwargs=dict(log_std_init=0.0, ortho_init=False), seed=None)
+                        policy_kwargs=dict(log_std_init=0.0, ortho_init=False, sde_net_arch=sde_net_arch), seed=None)
     model.learn(total_timesteps=int(1000), log_interval=5, eval_freq=500, eval_env=eval_env)
 
 
diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py
index 68a1c44..4aa2dbe 100644
--- a/torchy_baselines/a2c/a2c.py
+++ b/torchy_baselines/a2c/a2c.py
@@ -113,6 +113,7 @@ class A2C(PPO):
             # Optimization step
             self.policy.optimizer.zero_grad()
             loss.backward()
+
             # Clip grad norm
             th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
             self.policy.optimizer.step()
diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py
index f90883a..b0752e2 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -241,14 +241,17 @@ class StateDependentNoiseDistribution(Distribution):
         above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
     :param squash_output: (bool) Whether to squash the output using a tanh function,
         this allows to ensure boundaries.
+    :param learn_features: (bool) Whether to learn features for SDE or not.
+        This will enable gradients to be backpropagated through the features
+        `latent_sde` in the code.
     :param epsilon: (float) small value to avoid NaN due to numerical imprecision.
     """
     def __init__(self, action_dim, full_std=True, use_expln=False,
-                 squash_output=False, epsilon=1e-6):
+                 squash_output=False, learn_features=False, epsilon=1e-6):
         super(StateDependentNoiseDistribution, self).__init__()
         self.distribution = None
         self.action_dim = action_dim
-        self.latent_dim = None
+        self.latent_sde_dim = None
         self.mean_actions = None
         self.log_std = None
         self.weights_dist = None
@@ -256,6 +259,7 @@ class StateDependentNoiseDistribution(Distribution):
         self.use_expln = use_expln
         self.full_std = full_std
         self.epsilon = epsilon
+        self.learn_features = learn_features
         if squash_output:
             print("== Using TanhBijector ===")
             self.bijector = TanhBijector(epsilon)
@@ -284,7 +288,7 @@ class StateDependentNoiseDistribution(Distribution):
         if self.full_std:
             return std
         # Reduce the number of parameters:
-        return th.ones((self.latent_dim, self.action_dim)).to(log_std.device) * std
+        return th.ones((self.latent_sde_dim, self.action_dim)).to(log_std.device) * std
 
     def sample_weights(self, log_std):
         """
@@ -297,29 +301,32 @@ class StateDependentNoiseDistribution(Distribution):
         self.weights_dist = Normal(th.zeros_like(std), std)
         self.exploration_mat = self.weights_dist.rsample()
 
-    def proba_distribution_net(self, latent_dim, log_std_init=-2.0):
+    def proba_distribution_net(self, latent_dim, log_std_init=-2.0, latent_sde_dim=None):
         """
         Create the layers and parameter that represent the distribution:
         one output will be the deterministic action, the other parameter will be the
         standard deviation of the distribution that control the weights of the noise matrix.
 
-        :param latent_dim: (int) Dimension og the last layer of the policy (before the action layer)
+        :param latent_dim: (int) Dimension of the last layer of the policy (before the action layer)
         :param log_std_init: (float) Initial value for the log standard deviation
+        :param latent_sde_dim: (int) Dimension of the last layer of the feature extractor
+            for SDE. By default, it is shared with the policy network.
         :return: (nn.Linear, nn.Parameter)
         """
         # Network for the deterministic action, it represents the mean of the distribution
         mean_actions_net = nn.Linear(latent_dim, self.action_dim)
-
-        self.latent_dim = latent_dim
+        # When we learn features for the noise, the feature dimension
+        # can be different between the policy and the noise network
+        self.latent_sde_dim = latent_dim if latent_sde_dim is None else latent_sde_dim
         # Reduce the number of parameters if needed
-        log_std = th.ones(latent_dim, self.action_dim) if self.full_std else th.ones(latent_dim, 1)
+        log_std = th.ones(self.latent_sde_dim, self.action_dim) if self.full_std else th.ones(self.latent_sde_dim, 1)
         # Transform it to a parameter so it can be optimized
         log_std = nn.Parameter(log_std * log_std_init)
         # Sample an exploration matrix
         self.sample_weights(log_std)
         return mean_actions_net, log_std
 
-    def proba_distribution(self, mean_actions, log_std, latent_pi, deterministic=False):
+    def proba_distribution(self, mean_actions, log_std, latent_sde, deterministic=False):
         """
         Create and sample for the distribution given its parameters (mean, std)
 
@@ -328,13 +335,15 @@ class StateDependentNoiseDistribution(Distribution):
         :param deterministic: (bool)
         :return: (th.Tensor)
         """
-        variance = th.mm(latent_pi.detach() ** 2, self.get_std(log_std) ** 2)
+        # Stop gradient if we don't want to influence the features
+        latent_sde = latent_sde if self.learn_features else latent_sde.detach()
+        variance = th.mm(latent_sde ** 2, self.get_std(log_std) ** 2)
         self.distribution = Normal(mean_actions, th.sqrt(variance + self.epsilon))
 
         if deterministic:
             action = self.mode()
         else:
-            action = self.sample(latent_pi)
+            action = self.sample(latent_sde)
         return action, self
 
     def mode(self):
@@ -343,11 +352,12 @@ class StateDependentNoiseDistribution(Distribution):
             return self.bijector.forward(action)
         return action
 
-    def get_noise(self, latent_pi):
-        return th.mm(latent_pi.detach(), self.exploration_mat)
+    def get_noise(self, latent_sde):
+        latent_sde = latent_sde if self.learn_features else latent_sde.detach()
+        return th.mm(latent_sde, self.exploration_mat)
 
-    def sample(self, latent_pi):
-        noise = self.get_noise(latent_pi)
+    def sample(self, latent_sde):
+        noise = self.get_noise(latent_sde)
         action = self.distribution.mean + noise
         if self.bijector is not None:
             return self.bijector.forward(action)
@@ -357,8 +367,8 @@ class StateDependentNoiseDistribution(Distribution):
         # TODO: account for the squashing?
         return self.distribution.entropy()
 
-    def log_prob_from_params(self, mean_actions, log_std, latent_pi):
-        action, _ = self.proba_distribution(mean_actions, log_std, latent_pi)
+    def log_prob_from_params(self, mean_actions, log_std, latent_sde):
+        action, _ = self.proba_distribution(mean_actions, log_std, latent_sde)
         log_prob = self.log_prob(action)
         return action, log_prob
 
diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index 46fe75c..ed5c493 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -4,7 +4,7 @@ import torch as th
 import torch.nn as nn
 import numpy as np
 
-from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor
+from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor, create_mlp
 from torchy_baselines.common.distributions import make_proba_distribution,\
     DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution
 
@@ -30,7 +30,7 @@ class PPOPolicy(BasePolicy):
                  learning_rate, net_arch=None, device='cpu',
                  activation_fn=nn.Tanh, adam_epsilon=1e-5,
                  ortho_init=True, use_sde=False,
-                 log_std_init=0.0, full_std=True):
+                 log_std_init=0.0, full_std=True, sde_net_arch=None):
         super(PPOPolicy, self).__init__(observation_space, action_space, device)
         self.obs_dim = self.observation_space.shape[0]
 
@@ -61,9 +61,13 @@ class PPOPolicy(BasePolicy):
             dist_kwargs = {
                 'full_std': full_std,
                 'squash_output': False,
-                'use_expln': False
+                'use_expln': False,
+                'learn_features': sde_net_arch is not None
             }
 
+        self.sde_feature_extractor = None
+        self.sde_net_arch = sde_net_arch
+
         # Action distribution
         self.action_dist = make_proba_distribution(action_space, use_sde=use_sde, dist_kwargs=dist_kwargs)
 
@@ -79,9 +83,20 @@ class PPOPolicy(BasePolicy):
         self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch,
                                           activation_fn=self.activation_fn, device=self.device)
 
+        # Separate feature extractor for SDE
+        if self.sde_net_arch is not None:
+            latent_sde = create_mlp(self.features_dim, -1, self.sde_net_arch,
+                                    activation_fn=self.activation_fn, squash_out=False)
+            self.sde_feature_extractor = nn.Sequential(*latent_sde)
+
         if isinstance(self.action_dist, (DiagGaussianDistribution, StateDependentNoiseDistribution)):
             self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
                                                                                     log_std_init=self.log_std_init)
+        elif isinstance(self.action_dist, StateDependentNoiseDistribution):
+            latent_sde_dim = self.mlp_extractor.latent_dim_pi if self.sde_net_arch is None else self.sde_net_arch[-1]
+            self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
+                                                                                    latent_sde_dim=latent_sde_dim,
+                                                                                    log_std_init=self.log_std_init)
         elif isinstance(self.action_dist, CategoricalDistribution):
             self.action_net = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi)
 
@@ -102,16 +117,23 @@ class PPOPolicy(BasePolicy):
     def forward(self, obs, deterministic=False):
         if not isinstance(obs, th.Tensor):
             obs = th.FloatTensor(obs).to(self.device)
-        latent_pi, latent_vf = self._get_latent(obs)
+        latent_pi, latent_vf, latent_sde = self._get_latent(obs)
         value = self.value_net(latent_vf)
-        action, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
+        action, action_distribution = self._get_action_dist_from_latent(latent_pi, latent_sde=latent_sde,
+                                                                        deterministic=deterministic)
         log_prob = action_distribution.log_prob(action)
         return action, value, log_prob
 
     def _get_latent(self, obs):
-        return self.mlp_extractor(self.features_extractor(obs))
+        features = self.features_extractor(obs)
+        latent_pi, latent_vf = self.mlp_extractor(features)
+        # Features for sde
+        latent_sde = latent_pi
+        if self.sde_feature_extractor is not None:
+            latent_sde = self.sde_feature_extractor(features)
+        return latent_pi, latent_vf, latent_sde
 
-    def _get_action_dist_from_latent(self, latent_pi, deterministic=False):
+    def _get_action_dist_from_latent(self, latent_pi, latent_sde=None, deterministic=False):
         mean_actions = self.action_net(latent_pi)
 
         if isinstance(self.action_dist, DiagGaussianDistribution):
@@ -121,11 +143,11 @@ class PPOPolicy(BasePolicy):
             return self.action_dist.proba_distribution(mean_actions, deterministic=deterministic)
 
         elif isinstance(self.action_dist, StateDependentNoiseDistribution):
-            return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_pi, deterministic=deterministic)
+            return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde, deterministic=deterministic)
 
     def actor_forward(self, obs, deterministic=False):
-        latent_pi, _ = self._get_latent(obs)
-        action, _ = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
+        latent_pi, _, latent_sde = self._get_latent(obs)
+        action, _ = self._get_action_dist_from_latent(latent_pi, latent_sde, deterministic=deterministic)
         return action.detach().cpu().numpy()
 
     def evaluate_actions(self, obs, action, deterministic=False):
@@ -139,14 +161,14 @@ class PPOPolicy(BasePolicy):
         :return: (th.Tensor, th.Tensor, th.Tensor) estimated value, log likelihood of taking those actions
             and entropy of the action distribution.
         """
-        latent_pi, latent_vf = self._get_latent(obs)
-        _, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
+        latent_pi, latent_vf, latent_sde = self._get_latent(obs)
+        _, action_distribution = self._get_action_dist_from_latent(latent_pi, latent_sde, deterministic=deterministic)
         log_prob = action_distribution.log_prob(action)
         value = self.value_net(latent_vf)
         return value, log_prob, action_distribution.entropy()
 
     def value_forward(self, obs):
-        _, latent_vf = self._get_latent(obs)
+        _, latent_vf, _ = self._get_latent(obs)
         return self.value_net(latent_vf)
 
 

From 0885dbe74bfdf863fd69ff0af8d38e78e5d20c7c Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 25 Nov 2019 15:02:10 +0100
Subject: [PATCH 02/47] Bug fix in choosing the distribution

---
 tests/test_sde.py                | 2 +-
 torchy_baselines/ppo/policies.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_sde.py b/tests/test_sde.py
index 7143e9c..c565269 100644
--- a/tests/test_sde.py
+++ b/tests/test_sde.py
@@ -49,7 +49,7 @@ def test_state_dependent_exploration():
 
 
 @pytest.mark.parametrize("model_class", [A2C])
-@pytest.mark.parametrize("sde_net_arch", [None, [64, 64]])
+@pytest.mark.parametrize("sde_net_arch", [None, [32, 16]])
 def test_state_dependent_noise(model_class, sde_net_arch):
     env_id = 'MountainCarContinuous-v0'
 
diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index ed5c493..a0abfb4 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -89,7 +89,7 @@ class PPOPolicy(BasePolicy):
                                     activation_fn=self.activation_fn, squash_out=False)
             self.sde_feature_extractor = nn.Sequential(*latent_sde)
 
-        if isinstance(self.action_dist, (DiagGaussianDistribution, StateDependentNoiseDistribution)):
+        if isinstance(self.action_dist, DiagGaussianDistribution):
             self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
                                                                                     log_std_init=self.log_std_init)
         elif isinstance(self.action_dist, StateDependentNoiseDistribution):

From d26fcf45668eead686608e5a5df18694f4ea541f Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Tue, 26 Nov 2019 11:57:48 +0100
Subject: [PATCH 03/47] Fix grad computation for sde test

---
 tests/test_sde.py | 36 ++++++++++++++++++++++++------------
 1 file changed, 24 insertions(+), 12 deletions(-)

diff --git a/tests/test_sde.py b/tests/test_sde.py
index c565269..c39c5c5 100644
--- a/tests/test_sde.py
+++ b/tests/test_sde.py
@@ -15,37 +15,49 @@ def test_state_dependent_exploration():
     """
     n_states = 2
     state_dim = 3
-    # TODO: fix for action_dim > 1
-    action_dim = 1
-    sigma = th.ones(state_dim, 1, requires_grad=True)
+    action_dim = 10
+    sigma_hat = th.ones(state_dim, action_dim, requires_grad=True)
     # Reduce the number of parameters
     # sigma_ = th.ones(state_dim, action_dim) * sigma_
 
     # weights_dist = Normal(th.zeros_like(log_sigma), th.exp(log_sigma))
     th.manual_seed(2)
-    weights_dist = Normal(th.zeros_like(sigma), sigma)
-
+    weights_dist = Normal(th.zeros_like(sigma_hat), sigma_hat)
     weights = weights_dist.rsample()
+
     state = th.rand(n_states, state_dim)
     mu = th.ones(action_dim)
     # print(weights.shape, state.shape)
     noise = th.mm(state, weights)
 
-    variance = th.mm(state ** 2, sigma ** 2)
+    action = mu + noise
+
+    variance = th.mm(state ** 2, sigma_hat ** 2)
     action_dist = Normal(mu, th.sqrt(variance))
 
-    loss = action_dist.log_prob((mu + noise).detach()).mean()
+    # Sum over the action dimension because we assume they are independent
+    loss = action_dist.log_prob((action).detach()).sum(dim=-1).mean()
     loss.backward()
 
-    # From Rueckstiess paper
-    grad = th.zeros_like(sigma)
+    # From Rueckstiess paper: check that the computed gradient
+    # correspond to the analytical form
+    grad = th.zeros_like(sigma_hat)
     for j in range(action_dim):
+        # sigma_hat is the std of the gaussian distribution of the noise matrix weights
+        # sigma_j = sum_j(state_i **2 * sigma_hat_ij ** 2)
+        # sigma_j is the standard deviation of the policy gaussian distribution
+        sigma_j = th.sqrt(variance[:, j])
         for i in range(state_dim):
-            a = ((noise[:, j] ** 2 - variance[:, j]) / (variance[:, j] ** 2)) * (state[:, i] ** 2 * sigma[i, j])
-            grad[i, j] = a.mean()
+            # Derivative of the log probability of the jth component of the action
+            # w.r.t. the standard deviation sigma_j
+            d_log_policy_j = (noise[:, j] ** 2 - sigma_j ** 2) / sigma_j ** 3
+            # Derivative of sigma_j w.r.t. sigma_hat_ij
+            d_log_sigma_j = (state[:, i] ** 2 * sigma_hat[i, j]) / sigma_j
+            # Chain rule, average over the minibatch
+            grad[i, j] = (d_log_policy_j * d_log_sigma_j).mean()
 
     # sigma.grad should be equal to grad
-    assert sigma.grad.allclose(grad)
+    assert sigma_hat.grad.allclose(grad)
 
 
 @pytest.mark.parametrize("model_class", [A2C])

From 5483e02d1a7cf76c090e4d7e295e2e3605af4293 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Tue, 26 Nov 2019 15:26:12 +0100
Subject: [PATCH 04/47] Add SDE support for SAC

---
 tests/test_sde.py                     |  4 +-
 torchy_baselines/common/base_class.py | 10 ++-
 torchy_baselines/sac/policies.py      | 87 +++++++++++++++++++++------
 torchy_baselines/sac/sac.py           | 15 +++--
 torchy_baselines/td3/policies.py      |  6 +-
 torchy_baselines/td3/td3.py           |  4 +-
 6 files changed, 95 insertions(+), 31 deletions(-)

diff --git a/tests/test_sde.py b/tests/test_sde.py
index c39c5c5..668bc9e 100644
--- a/tests/test_sde.py
+++ b/tests/test_sde.py
@@ -4,7 +4,7 @@ import gym
 import torch as th
 from torch.distributions import Normal
 
-from torchy_baselines import A2C, TD3
+from torchy_baselines import A2C, TD3, SAC
 from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize
 from torchy_baselines.common.monitor import Monitor
 
@@ -73,7 +73,7 @@ def test_state_dependent_noise(model_class, sde_net_arch):
     model.learn(total_timesteps=int(1000), log_interval=5, eval_freq=500, eval_env=eval_env)
 
 
-@pytest.mark.parametrize("model_class", [TD3])
+@pytest.mark.parametrize("model_class", [TD3, SAC])
 def test_state_dependent_offpolicy_noise(model_class):
     model = model_class('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True,
                         verbose=1, policy_kwargs=dict(log_std_init=-2))
diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index 40dcedb..552f3a2 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -33,12 +33,14 @@ class BaseRLModel(object):
     :param monitor_wrapper: (bool) When creating an environment, whether to wrap it
         or not in a Monitor wrapper.
     :param seed: (int) Seed for the pseudo random generators
+    :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
+        instead of action noise exploration (default: False)
     """
     __metaclass__ = ABCMeta
 
     def __init__(self, policy, env, policy_base, policy_kwargs=None,
                  verbose=0, device='auto', support_multi_env=False,
-                 create_eval_env=False, monitor_wrapper=True, seed=None):
+                 create_eval_env=False, monitor_wrapper=True, seed=None, use_sde=False):
         if isinstance(policy, str) and policy_base is not None:
             self.policy = get_policy_from_name(policy_base, policy)
         else:
@@ -67,7 +69,8 @@ class BaseRLModel(object):
         self.action_noise = None
         # Used for SDE only
         self.rollout_data = None
-        self.use_sde = False
+        self.on_policy_exploration = False
+        self.use_sde = use_sde
         # Track the training progress (from 1 to 0)
         # this is used to update the learning rate
         self._current_progress = 1
@@ -395,7 +398,8 @@ class BaseRLModel(object):
         if self.use_sde:
             self.actor.reset_noise()
             # Reset rollout data
-            self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']}
+            if self.on_policy_exploration:
+                self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']}
 
         while total_steps < n_steps or total_episodes < n_episodes:
             done = False
diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index b61ae50..dd644bd 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -2,7 +2,7 @@ import torch as th
 import torch.nn as nn
 
 from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork
-from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution
+from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
 
 # CAP the standard deviation of the actor
 LOG_STD_MAX = 2
@@ -10,33 +10,83 @@ LOG_STD_MIN = -20
 
 
 class Actor(BaseNetwork):
-    def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU):
+    """
+    Actor network (policy) for SAC.
+
+    :param obs_dim: (int) Dimension of the observation
+    :param action_dim: (int) Dimension of the action space
+    :param net_arch: ([int]) Network architecture
+    :param activation_fn: (nn.Module) Activation function
+    :param use_sde: (bool) Whether to use State Dependent Exploration or not
+    :param log_std_init: (float) Initial value for the log standard deviation
+    :param full_std: (bool) Whether to use (n_features x n_actions) parameters
+        for the std instead of only (n_features,) when using SDE.
+    """
+    def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
+                 use_sde=False, log_std_init=-3, full_std=False):
         super(Actor, self).__init__()
 
-        # TODO: orthogonal initialization?
         actor_net = create_mlp(obs_dim, -1, net_arch, activation_fn)
         self.actor_net = nn.Sequential(*actor_net)
+        self.use_sde = use_sde
 
-        self.action_dist = SquashedDiagGaussianDistribution(action_dim)
-        self.mu = nn.Linear(net_arch[-1], action_dim)
-        self.log_std = nn.Linear(net_arch[-1], action_dim)
+        if self.use_sde:
+            # TODO: check for the learn_features
+            self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
+                                                               learn_features=False, squash_output=True)
+            self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
+                                                                            log_std_init=log_std_init)
+        else:
+            self.action_dist = SquashedDiagGaussianDistribution(action_dim)
+            self.mu = nn.Linear(net_arch[-1], action_dim)
+            self.log_std = nn.Linear(net_arch[-1], action_dim)
+
+    def get_std(self):
+        """
+        Retrieve the standard deviation of the action distribution.
+        Only useful when using SDE.
+        It corresponds to `th.exp(log_std)` in the normal case,
+        but is slightly different when using `expln` function
+        (cf StateDependentNoiseDistribution doc).
+
+        :return: (th.Tensor)
+        """
+        return self.action_dist.get_std(self.log_std)
+
+    def reset_noise(self):
+        """
+        Sample new weights for the exploration matrix, when using SDE.
+        """
+        self.action_dist.sample_weights(self.log_std)
 
     def get_action_dist_params(self, obs):
         latent = self.actor_net(obs)
-        mean_actions, log_std = self.mu(latent), self.log_std(latent)
-        # Original Implementation to cap the standard deviation
-        log_std = th.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX)
-        return mean_actions, log_std
+
+        if self.use_sde:
+            mean_actions, log_std = self.mu(latent), self.log_std
+        else:
+            mean_actions, log_std = self.mu(latent), self.log_std(latent)
+            # Original Implementation to cap the standard deviation
+            log_std = th.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX)
+        return mean_actions, log_std, latent
 
     def forward(self, obs, deterministic=False):
-        mean_actions, log_std = self.get_action_dist_params(obs)
-        # Note the action is squashed
-        action, _ = self.action_dist.proba_distribution(mean_actions, log_std, deterministic=deterministic)
+        mean_actions, log_std, latent = self.get_action_dist_params(obs)
+        if self.use_sde:
+            # Note the action is squashed
+            action, _ = self.action_dist.proba_distribution(mean_actions, log_std, latent, deterministic=deterministic)
+        else:
+            # Note the action is squashed
+            action, _ = self.action_dist.proba_distribution(mean_actions, log_std, deterministic=deterministic)
         return action
 
     def action_log_prob(self, obs):
-        mean_actions, log_std = self.get_action_dist_params(obs)
-        action, log_prob = self.action_dist.log_prob_from_params(mean_actions, log_std)
+        mean_actions, log_std, latent = self.get_action_dist_params(obs)
+
+        if self.use_sde:
+            action, log_prob = self.action_dist.log_prob_from_params(mean_actions, self.log_std, latent)
+        else:
+            action, log_prob = self.action_dist.log_prob_from_params(mean_actions, log_std)
         return action, log_prob
 
 
@@ -64,7 +114,7 @@ class Critic(BaseNetwork):
 class SACPolicy(BasePolicy):
     def __init__(self, observation_space, action_space,
                  learning_rate, net_arch=None, device='cpu',
-                 activation_fn=nn.ReLU):
+                 activation_fn=nn.ReLU, use_sde=False, log_std_init=-3):
         super(SACPolicy, self).__init__(observation_space, action_space, device)
 
         if net_arch is None:
@@ -80,6 +130,9 @@ class SACPolicy(BasePolicy):
             'net_arch': self.net_arch,
             'activation_fn': self.activation_fn
         }
+        self.actor_kwargs = self.net_args.copy()
+        self.actor_kwargs['use_sde'] = use_sde
+        self.actor_kwargs['log_std_init'] = log_std_init
         self.actor, self.actor_target = None, None
         self.critic, self.critic_target = None, None
 
@@ -95,7 +148,7 @@ class SACPolicy(BasePolicy):
         self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))
 
     def make_actor(self):
-        return Actor(**self.net_args).to(self.device)
+        return Actor(**self.actor_kwargs).to(self.device)
 
     def make_critic(self):
         return Critic(**self.net_args).to(self.device)
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index c14623c..aae0443 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -44,6 +44,8 @@ class SAC(BaseRLModel):
     :param action_noise: (ActionNoise) the action noise type (None by default), this can help
         for hard exploration problem. Cf common.noise for the different action noise type.
     :param gamma: (float) the discount factor
+    :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
+        instead of action noise exploration (default: False)
     :param create_eval_env: (bool) Whether to create a second environment that will be
         used for evaluating the agent periodically. (Only available when passing string for the environment)
     :param policy_kwargs: (dict) additional arguments to be passed to the policy on creation
@@ -58,12 +60,12 @@ class SAC(BaseRLModel):
                  tau=0.005, ent_coef='auto', target_update_interval=1,
                  train_freq=1, gradient_steps=1, n_episodes_rollout=-1,
                  target_entropy='auto', action_noise=None,
-                 gamma=0.99, tensorboard_log=None, create_eval_env=False,
+                 gamma=0.99, use_sde=False, tensorboard_log=None, create_eval_env=False,
                  policy_kwargs=None, verbose=0, seed=0, device='auto',
                  _init_setup_model=True):
 
         super(SAC, self).__init__(policy, env, SACPolicy, policy_kwargs, verbose, device,
-                                  create_eval_env=create_eval_env, seed=seed)
+                                  create_eval_env=create_eval_env, seed=seed, use_sde=use_sde)
 
         self.learning_rate = learning_rate
         self.target_entropy = target_entropy
@@ -124,8 +126,8 @@ class SAC(BaseRLModel):
             self.ent_coef = float(self.ent_coef)
 
         self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
-        self.policy = self.policy(self.observation_space, self.action_space,
-                                  self.learning_rate, device=self.device, **self.policy_kwargs)
+        self.policy = self.policy(self.observation_space, self.action_space, learning_rate=self.learning_rate,
+                                  use_sde=self.use_sde, device=self.device, **self.policy_kwargs)
         self.policy = self.policy.to(self.device)
         self._create_aliases()
 
@@ -167,6 +169,11 @@ class SAC(BaseRLModel):
 
             obs, action_batch, next_obs, done, reward = replay_data
 
+            # TODO: check if there is another way to fix pytorch complain
+            # if we don't sample the weights again
+            # (Trying to backward through the graph a second time)
+            if self.use_sde:
+                self.actor.reset_noise()
             # Action by the current actor for the sampled state
             action_pi, log_prob = self.actor.action_log_prob(obs)
             log_prob = log_prob.reshape(-1, 1)
diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py
index cd72950..6872213 100644
--- a/torchy_baselines/td3/policies.py
+++ b/torchy_baselines/td3/policies.py
@@ -21,7 +21,7 @@ class Actor(BaseNetwork):
         for the std instead of only (n_features,) when using SDE.
     """
     def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
-                 use_sde=False, log_std_init=-2, clip_noise=None,
+                 use_sde=False, log_std_init=-3, clip_noise=None,
                  lr_sde=3e-4, full_std=False):
         super(Actor, self).__init__()
 
@@ -84,7 +84,7 @@ class Actor(BaseNetwork):
 
     def reset_noise(self):
         """
-        Sample new weights for the exploration matrix.
+        Sample new weights for the exploration matrix, when using SDE.
         """
         self.action_dist.sample_weights(self.log_std)
 
@@ -151,7 +151,7 @@ class TD3Policy(BasePolicy):
     """
     def __init__(self, observation_space, action_space,
                  learning_rate, net_arch=None, device='cpu',
-                 activation_fn=nn.ReLU, use_sde=False, log_std_init=-2,
+                 activation_fn=nn.ReLU, use_sde=False, log_std_init=-3,
                  clip_noise=None, lr_sde=3e-4):
         super(TD3Policy, self).__init__(observation_space, action_space, device)
 
diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py
index 5ecb38e..2c1b412 100644
--- a/torchy_baselines/td3/td3.py
+++ b/torchy_baselines/td3/td3.py
@@ -62,7 +62,7 @@ class TD3(BaseRLModel):
                  seed=0, device='auto', _init_setup_model=True):
 
         super(TD3, self).__init__(policy, env, TD3Policy, policy_kwargs, verbose, device,
-                                  create_eval_env=create_eval_env, seed=seed)
+                                  create_eval_env=create_eval_env, seed=seed, use_sde=use_sde)
 
         self.buffer_size = buffer_size
         self.learning_rate = learning_rate
@@ -79,10 +79,10 @@ class TD3(BaseRLModel):
         self.target_policy_noise = target_policy_noise
 
         # State Dependent Exploration
-        self.use_sde = use_sde
         self.sde_max_grad_norm = sde_max_grad_norm
         self.sde_ent_coef = sde_ent_coef
         self.sde_log_std_scheduler = sde_log_std_scheduler
+        self.on_policy_exploration = True
 
         if _init_setup_model:
             self._setup_model()

From fe67a98711d00d0001d763d6546ef7257de74aff Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Tue, 26 Nov 2019 17:44:06 +0100
Subject: [PATCH 05/47] Log more values

---
 README.md                        | 4 +++-
 torchy_baselines/sac/policies.py | 2 +-
 torchy_baselines/sac/sac.py      | 6 ++++++
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 0b19832..bf14287 100644
--- a/README.md
+++ b/README.md
@@ -21,13 +21,15 @@ TODO:
 - save/load
 - better predict
 - complete logger
-- SDE: learn the feature extractor?
 - Refactor: buffer with numpy array instead of pytorch
 - Refactor: remove duplicated code for evaluation
+
 - plotting? -> zoo
 
 Later:
 - get_parameters / set_parameters
+- SDE: use [affine transform](https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/Affine)
+  to scale the noise after a tanh transform?
 - CNN policies + normalization
 - tensorboard support
 - DQN
diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index dd644bd..34bb3e4 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -23,7 +23,7 @@ class Actor(BaseNetwork):
         for the std instead of only (n_features,) when using SDE.
     """
     def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
-                 use_sde=False, log_std_init=-3, full_std=False):
+                 use_sde=False, log_std_init=-3, full_std=True):
         super(Actor, self).__init__()
 
         actor_net = create_mlp(obs_dim, -1, net_arch, activation_fn)
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index aae0443..49f562f 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -9,6 +9,7 @@ from torchy_baselines.common.buffers import ReplayBuffer
 from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.sac.policies import SACPolicy
 from torchy_baselines.common.vec_env import sync_envs_normalization
+from torchy_baselines.common import logger
 
 
 class SAC(BaseRLModel):
@@ -234,6 +235,11 @@ class SAC(BaseRLModel):
                 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)
 
+        # TODO: average
+        logger.logkv("ent_coef", ent_coef.item())
+        logger.logkv("actor_loss", actor_loss.item())
+        logger.logkv("critic_loss", critic_loss.item())
+
     def learn(self, total_timesteps, callback=None, log_interval=4,
               eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC",
               reset_num_timesteps=True):

From fbe29a7298380c755524886a642ab1b788326009 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Wed, 27 Nov 2019 17:29:47 +0100
Subject: [PATCH 06/47] Track down autograd error

"Trying to backward through the graph a second time" -> added a comment
---
 torchy_baselines/common/distributions.py |  2 +-
 torchy_baselines/sac/policies.py         |  2 +-
 torchy_baselines/sac/sac.py              | 23 +++++++++++++++--------
 3 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py
index b0752e2..ce7e699 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -288,7 +288,7 @@ class StateDependentNoiseDistribution(Distribution):
         if self.full_std:
             return std
         # Reduce the number of parameters:
-        return th.ones((self.latent_sde_dim, self.action_dim)).to(log_std.device) * std
+        return th.ones(self.latent_sde_dim, self.action_dim).to(log_std.device) * std
 
     def sample_weights(self, log_std):
         """
diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index 34bb3e4..0774ef0 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -33,7 +33,7 @@ class Actor(BaseNetwork):
         if self.use_sde:
             # TODO: check for the learn_features
             self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
-                                                               learn_features=False, squash_output=True)
+                                                               learn_features=True, squash_output=True)
             self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
                                                                             log_std_init=log_std_init)
         else:
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index 49f562f..5c763b5 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -170,11 +170,13 @@ class SAC(BaseRLModel):
 
             obs, action_batch, next_obs, done, reward = replay_data
 
-            # TODO: check if there is another way to fix pytorch complain
-            # if we don't sample the weights again
-            # (Trying to backward through the graph a second time)
-            if self.use_sde:
-                self.actor.reset_noise()
+            # Two options: retain_graph=True in the actor_loss.backward()
+            # or sample again the noise matrix
+            # otherwise the intermediate step `std = th.exp(log_std)`
+            # is lost and we cannot backpropagate through again
+            # if self.use_sde:
+            #     self.actor.reset_noise()
+
             # Action by the current actor for the sampled state
             action_pi, log_prob = self.actor.action_log_prob(obs)
             log_prob = log_prob.reshape(-1, 1)
@@ -196,10 +198,10 @@ class SAC(BaseRLModel):
                 ent_coef_loss.backward()
                 self.ent_coef_optimizer.step()
 
-            # Select action according to policy
-            next_action, next_log_prob = self.actor.action_log_prob(next_obs)
 
             with th.no_grad():
+                # Select action according to policy
+                next_action, next_log_prob = self.actor.action_log_prob(next_obs)
                 # Compute the target Q value
                 target_q1, target_q2 = self.critic_target(next_obs, next_action)
                 target_q = th.min(target_q1, target_q2)
@@ -227,7 +229,10 @@ class SAC(BaseRLModel):
 
             # Optimize the actor
             self.actor.optimizer.zero_grad()
-            actor_loss.backward()
+            # Cf comment above, otherwise pytorch raises an error
+            # ("Trying to backward through the graph a second time")
+            retain_graph = True if self.use_sde and gradient_steps > 1 else False
+            actor_loss.backward(retain_graph=retain_graph)
             self.actor.optimizer.step()
 
             # Update target networks
@@ -239,6 +244,8 @@ class SAC(BaseRLModel):
         logger.logkv("ent_coef", ent_coef.item())
         logger.logkv("actor_loss", actor_loss.item())
         logger.logkv("critic_loss", critic_loss.item())
+        if ent_coef_loss is not None:
+            logger.logkv("ent_coef_loss", ent_coef_loss.item())
 
     def learn(self, total_timesteps, callback=None, log_interval=4,
               eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC",

From 879191b26acb27ebb81a23562c6c800d0c1fe603 Mon Sep 17 00:00:00 2001
From: Antonin RAFFIN <antonin.raffin@ensta.org>
Date: Sun, 1 Dec 2019 13:11:13 +0100
Subject: [PATCH 07/47] Bug fix in SAC with constant ent coeff + try batch sde
 matrices

---
 torchy_baselines/common/distributions.py | 11 +++++++++--
 torchy_baselines/sac/policies.py         |  6 ++++--
 torchy_baselines/sac/sac.py              | 15 +++++++++------
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py
index ce7e699..75166af 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -290,7 +290,7 @@ class StateDependentNoiseDistribution(Distribution):
         # Reduce the number of parameters:
         return th.ones(self.latent_sde_dim, self.action_dim).to(log_std.device) * std
 
-    def sample_weights(self, log_std):
+    def sample_weights(self, log_std, batch_size=1):
         """
         Sample weights for the noise exploration matrix,
         using a centered gaussian distribution.
@@ -300,6 +300,7 @@ class StateDependentNoiseDistribution(Distribution):
         std = self.get_std(log_std)
         self.weights_dist = Normal(th.zeros_like(std), std)
         self.exploration_mat = self.weights_dist.rsample()
+        self.exploration_matrices = self.weights_dist.rsample((batch_size,))
 
     def proba_distribution_net(self, latent_dim, log_std_init=-2.0, latent_sde_dim=None):
         """
@@ -354,7 +355,13 @@ class StateDependentNoiseDistribution(Distribution):
 
     def get_noise(self, latent_sde):
         latent_sde = latent_sde if self.learn_features else latent_sde.detach()
-        return th.mm(latent_sde, self.exploration_mat)
+        if len(latent_sde) != len(self.exploration_matrices):
+            return th.mm(latent_sde, self.exploration_mat)
+        # (batch_size, n_features) -> (batch_size, 1, n_features)
+        latent_sde = latent_sde.unsqueeze(1)
+        # (batch_size, 1, n_actions)
+        noise = th.bmm(latent_sde, self.exploration_matrices)
+        return noise.squeeze(1)
 
     def sample(self, latent_sde):
         noise = self.get_noise(latent_sde)
diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index 0774ef0..0a08335 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -53,11 +53,13 @@ class Actor(BaseNetwork):
         """
         return self.action_dist.get_std(self.log_std)
 
-    def reset_noise(self):
+    def reset_noise(self, batch_size=1):
         """
         Sample new weights for the exploration matrix, when using SDE.
+
+        :param batch_size: (int)
         """
-        self.action_dist.sample_weights(self.log_std)
+        self.action_dist.sample_weights(self.log_std, batch_size=batch_size)
 
     def get_action_dist_params(self, obs):
         latent = self.actor_net(obs)
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index 5c763b5..66f3492 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -87,6 +87,7 @@ class SAC(BaseRLModel):
         self.n_episodes_rollout = n_episodes_rollout
         self.action_noise = action_noise
         self.gamma = gamma
+        self.ent_coef_optimizer = None
 
         if _init_setup_model:
             self._setup_model()
@@ -124,7 +125,7 @@ class SAC(BaseRLModel):
             # Force conversion to float
             # this will throw an error if a malformed string (different from 'auto')
             # is passed
-            self.ent_coef = float(self.ent_coef)
+            self.ent_coef = th.tensor(float(self.ent_coef)).to(self.device)
 
         self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
         self.policy = self.policy(self.observation_space, self.action_space, learning_rate=self.learning_rate,
@@ -174,15 +175,15 @@ class SAC(BaseRLModel):
             # or sample again the noise matrix
             # otherwise the intermediate step `std = th.exp(log_std)`
             # is lost and we cannot backpropagate through again
-            # if self.use_sde:
-            #     self.actor.reset_noise()
+            if self.use_sde:
+                self.actor.reset_noise(batch_size=batch_size)
 
             # Action by the current actor for the sampled state
             action_pi, log_prob = self.actor.action_log_prob(obs)
             log_prob = log_prob.reshape(-1, 1)
 
             ent_coef_loss = None
-            if not isinstance(self.ent_coef, float):
+            if self.ent_coef_optimizer is not None:
                 # Important: detach the variable from the graph
                 # so we don't change it with other losses
                 # see https://github.com/rail-berkeley/softlearning/issues/60
@@ -200,6 +201,8 @@ class SAC(BaseRLModel):
 
 
             with th.no_grad():
+                if self.use_sde:
+                    self.actor.reset_noise(batch_size=batch_size)
                 # Select action according to policy
                 next_action, next_log_prob = self.actor.action_log_prob(next_obs)
                 # Compute the target Q value
@@ -231,8 +234,8 @@ class SAC(BaseRLModel):
             self.actor.optimizer.zero_grad()
             # Cf comment above, otherwise pytorch raises an error
             # ("Trying to backward through the graph a second time")
-            retain_graph = True if self.use_sde and gradient_steps > 1 else False
-            actor_loss.backward(retain_graph=retain_graph)
+            # retain_graph = True if self.use_sde and gradient_steps > 1 else False
+            actor_loss.backward(retain_graph=False)
             self.actor.optimizer.step()
 
             # Update target networks

From 03a84f97eaea47804569fcb2f20280fea6241ab0 Mon Sep 17 00:00:00 2001
From: Antonin RAFFIN <antonin.raffin@ensta.org>
Date: Sun, 1 Dec 2019 16:46:39 +0100
Subject: [PATCH 08/47] Add monte-carlo test for SDE distribution

---
 tests/test_distributions.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/tests/test_distributions.py b/tests/test_distributions.py
index 47651e4..130a1bc 100644
--- a/tests/test_distributions.py
+++ b/tests/test_distributions.py
@@ -1,8 +1,9 @@
 import numpy as np
 import torch as th
 
+from torchy_baselines.common.utils import set_random_seed
 from torchy_baselines.common.distributions import DiagGaussianDistribution, SquashedDiagGaussianDistribution,\
-    CategoricalDistribution, TanhBijector
+    CategoricalDistribution, TanhBijector, StateDependentNoiseDistribution
 
 # TODO: more tests for the other distributions
 def test_bijector():
@@ -18,3 +19,21 @@ def test_bijector():
     assert th.max(th.abs(squashed_actions)) <= 1.0
     # Check the inverse method
     assert th.isclose(TanhBijector.inverse(squashed_actions), actions).all()
+
+
+def test_sde_distribution():
+    n_samples = int(5e6)
+    n_features = 2
+    n_actions = 1
+    deterministic_actions = th.ones(n_samples, n_actions) * 0.1
+    state = th.ones(n_samples, n_features) * 0.3
+    dist = StateDependentNoiseDistribution(n_actions, full_std=True, squash_output=False)
+
+    set_random_seed(1)
+    _, log_std = dist.proba_distribution_net(n_features)
+    dist.sample_weights(log_std, batch_size=n_samples)
+
+    actions, _ = dist.proba_distribution(deterministic_actions, log_std, state)
+
+    assert th.allclose(actions.mean(), dist.distribution.mean.mean(), rtol=1e-3)
+    assert th.allclose(actions.std(), dist.distribution.scale.mean(), rtol=1e-3)

From 21e655ecbffbc836a56d516d43d1b706872fe5d2 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 2 Dec 2019 11:47:52 +0100
Subject: [PATCH 09/47] Add test for SAC with different entropy temperature

---
 tests/test_run.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/test_run.py b/tests/test_run.py
index 32a4b30..b29a526 100644
--- a/tests/test_run.py
+++ b/tests/test_run.py
@@ -37,8 +37,9 @@ def test_onpolicy(model_class, env_id):
     # model.load("test_save")
     # os.remove("test_save.pth")
 
-def test_sac():
+@pytest.mark.parametrize("ent_coef", ['auto', 0.01])
+def test_sac(ent_coef):
     model = SAC('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]),
-                learning_starts=100, verbose=1, create_eval_env=True, ent_coef='auto',
+                learning_starts=100, verbose=1, create_eval_env=True, ent_coef=ent_coef,
                 action_noise=NormalActionNoise(np.zeros(1), np.zeros(1)))
     model.learn(total_timesteps=1000, eval_freq=500)

From 7a6a500398b00030e839f3a9754a5d968ae3c961 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 2 Dec 2019 11:48:16 +0100
Subject: [PATCH 10/47] Allow to use states directly as features for sde

---
 torchy_baselines/common/distributions.py |  4 +++-
 torchy_baselines/common/policies.py      | 20 +++++++++++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py
index 75166af..8802687 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -355,8 +355,10 @@ class StateDependentNoiseDistribution(Distribution):
 
     def get_noise(self, latent_sde):
         latent_sde = latent_sde if self.learn_features else latent_sde.detach()
-        if len(latent_sde) != len(self.exploration_matrices):
+        # Default case: only one exploration matrix
+        if len(latent_sde) == 1 or len(latent_sde) != len(self.exploration_matrices):
             return th.mm(latent_sde, self.exploration_mat)
+        # Use batch matrix multiplication for efficient computation
         # (batch_size, n_features) -> (batch_size, 1, n_features)
         latent_sde = latent_sde.unsqueeze(1)
         # (batch_size, 1, n_actions)
diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py
index 4ac7e07..b01345a 100644
--- a/torchy_baselines/common/policies.py
+++ b/torchy_baselines/common/policies.py
@@ -62,7 +62,25 @@ class BasePolicy(nn.Module):
 
 def create_mlp(input_dim, output_dim, net_arch,
                activation_fn=nn.ReLU, squash_out=False):
-    modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()]
+    """
+    Create a multi layer perceptron (MLP), which is
+    a collection of fully-connected layers each followed by an activation function.
+
+    :param input_dim: (int) Dimension of the input vector
+    :param output_dim: (int)
+    :param net_arch: ([int]) Architecture of the neural net
+        It represents the number of units per layer.
+        The length of this list is the number of layers.
+    :param activation_fn: (th.nn.Module) The activation function
+        to use after each layer.
+    :param squash_out: (bool) Whether to squash the output using a Tanh
+        activation function
+    """
+
+    if len(net_arch) > 0:
+        modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()]
+    else:
+        modules = []
 
     for idx in range(len(net_arch) - 1):
         modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1]))

From a2a8bbdf11983ed3244c804797a14e83c06796dc Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 2 Dec 2019 11:48:34 +0100
Subject: [PATCH 11/47] Sample n matrices for A2C/PPO when using SDE

---
 torchy_baselines/ppo/policies.py | 14 ++++++++++----
 torchy_baselines/ppo/ppo.py      |  2 +-
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index a0abfb4..50dbeea 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -73,11 +73,13 @@ class PPOPolicy(BasePolicy):
 
         self._build(learning_rate)
 
-    def reset_noise_net(self):
+    def reset_noise_net(self, n_envs=1):
         """
         Sample new weights for the exploration matrix.
+
+        :param n_envs: (int)
         """
-        self.action_dist.sample_weights(self.log_std)
+        self.action_dist.sample_weights(self.log_std, batch_size=n_envs)
 
     def _build(self, learning_rate):
         self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch,
@@ -85,15 +87,19 @@ class PPOPolicy(BasePolicy):
 
         # Separate feature extractor for SDE
         if self.sde_net_arch is not None:
+            # Special case: when using states as features (i.e. sde_net_arch is an empty list)
+            # don't use any activation function
+            sde_activation = self.activation_fn if len(self.sde_net_arch) > 0 else None
             latent_sde = create_mlp(self.features_dim, -1, self.sde_net_arch,
-                                    activation_fn=self.activation_fn, squash_out=False)
+                                    activation_fn=sde_activation, squash_out=False)
             self.sde_feature_extractor = nn.Sequential(*latent_sde)
+            latent_sde_dim = self.sde_net_arch[-1] if len(self.sde_net_arch) > 0 else self.features_dim
 
         if isinstance(self.action_dist, DiagGaussianDistribution):
             self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
                                                                                     log_std_init=self.log_std_init)
         elif isinstance(self.action_dist, StateDependentNoiseDistribution):
-            latent_sde_dim = self.mlp_extractor.latent_dim_pi if self.sde_net_arch is None else self.sde_net_arch[-1]
+            latent_sde_dim = self.mlp_extractor.latent_dim_pi if self.sde_net_arch is None else latent_sde_dim
             self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
                                                                                     latent_sde_dim=latent_sde_dim,
                                                                                     log_std_init=self.log_std_init)
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index 4257c1e..99ed1b4 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -157,7 +157,7 @@ class PPO(BaseRLModel):
         # Sample new weights for the state dependent exploration
         # TODO: ensure episodic setting?
         if self.use_sde:
-            self.policy.reset_noise_net()
+            self.policy.reset_noise_net(env.num_envs)
 
         while n_steps < n_rollout_steps:
             with th.no_grad():

From 4e39a0627c6f822d90ca557b9bd54f9655022b9b Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 2 Dec 2019 14:06:17 +0100
Subject: [PATCH 12/47] Refactor: enable sde net arch for TD3 and SAC

---
 torchy_baselines/common/policies.py | 25 +++++++++-
 torchy_baselines/ppo/policies.py    | 13 ++---
 torchy_baselines/sac/policies.py    | 77 +++++++++++++++++++++++------
 torchy_baselines/td3/policies.py    | 74 ++++++++++++++++++---------
 4 files changed, 141 insertions(+), 48 deletions(-)

diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py
index b01345a..f536624 100644
--- a/torchy_baselines/common/policies.py
+++ b/torchy_baselines/common/policies.py
@@ -93,9 +93,30 @@ def create_mlp(input_dim, output_dim, net_arch,
     return modules
 
 
-class BaseNetwork(nn.Module):
-    """docstring for BaseNetwork."""
+def create_sde_feature_extractor(features_dim, sde_net_arch, activation_fn):
+    """
+    Create the neural network that will be used to extract features
+    for the SDE.
 
+    :param features_dim: (int)
+    :param sde_net_arch: ([int])
+    :param activation_fn: (nn.Module)
+    :return: (nn.Sequential, int)
+    """
+    # Special case: when using states as features (i.e. sde_net_arch is an empty list)
+    # don't use any activation function
+    sde_activation = activation_fn if len(sde_net_arch) > 0 else None
+    latent_sde_net = create_mlp(features_dim, -1, sde_net_arch, activation_fn=sde_activation, squash_out=False)
+    latent_sde_dim = sde_net_arch[-1] if len(sde_net_arch) > 0 else features_dim
+    sde_feature_extractor = nn.Sequential(*latent_sde_net)
+    return sde_feature_extractor, latent_sde_dim
+
+
+class BaseNetwork(nn.Module):
+    """
+    Abstract class for the different networks (actor/critic)
+    that implements two helpers for using CEM with their weights.
+    """
     def __init__(self):
         super(BaseNetwork, self).__init__()
 
diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index 50dbeea..9b38317 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -25,6 +25,9 @@ class PPOPolicy(BasePolicy):
     :param log_std_init: (float) Initial value for the log standard deviation
     :param full_std: (bool) Whether to use (n_features x n_actions) parameters
         for the std instead of only (n_features,) when using SDE
+    :param sde_net_arch: ([int]) Network architecture for extracting features
+        when using SDE. If None, the latent features from the policy will be used.
+        Pass an empty list to use the states as features.
     """
     def __init__(self, observation_space, action_space,
                  learning_rate, net_arch=None, device='cpu',
@@ -34,7 +37,6 @@ class PPOPolicy(BasePolicy):
         super(PPOPolicy, self).__init__(observation_space, action_space, device)
         self.obs_dim = self.observation_space.shape[0]
 
-
         # Default network architecture, from stable-baselines
         if net_arch is None:
             net_arch = [dict(pi=[64, 64], vf=[64, 64])]
@@ -87,13 +89,8 @@ class PPOPolicy(BasePolicy):
 
         # Separate feature extractor for SDE
         if self.sde_net_arch is not None:
-            # Special case: when using states as features (i.e. sde_net_arch is an empty list)
-            # don't use any activation function
-            sde_activation = self.activation_fn if len(self.sde_net_arch) > 0 else None
-            latent_sde = create_mlp(self.features_dim, -1, self.sde_net_arch,
-                                    activation_fn=sde_activation, squash_out=False)
-            self.sde_feature_extractor = nn.Sequential(*latent_sde)
-            latent_sde_dim = self.sde_net_arch[-1] if len(self.sde_net_arch) > 0 else self.features_dim
+            self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(self.features_dim, self.sde_net_arch,
+                                                                                      self.activation_fn)
 
         if isinstance(self.action_dist, DiagGaussianDistribution):
             self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index 0a08335..be32660 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -1,7 +1,7 @@
 import torch as th
 import torch.nn as nn
 
-from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork
+from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, create_sde_feature_extractor
 from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
 
 # CAP the standard deviation of the actor
@@ -21,20 +21,30 @@ class Actor(BaseNetwork):
     :param log_std_init: (float) Initial value for the log standard deviation
     :param full_std: (bool) Whether to use (n_features x n_actions) parameters
         for the std instead of only (n_features,) when using SDE.
+    :param sde_net_arch: ([int]) Network architecture for extracting features
+        when using SDE. If None, the latent features from the policy will be used.
+        Pass an empty list to use the states as features.
     """
     def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
-                 use_sde=False, log_std_init=-3, full_std=True):
+                 use_sde=False, log_std_init=-3, full_std=True, sde_net_arch=None):
         super(Actor, self).__init__()
 
-        actor_net = create_mlp(obs_dim, -1, net_arch, activation_fn)
-        self.actor_net = nn.Sequential(*actor_net)
+        latent_pi_net = create_mlp(obs_dim, -1, net_arch, activation_fn)
+        self.latent_pi = nn.Sequential(*latent_pi_net)
         self.use_sde = use_sde
+        self.sde_feature_extractor = None
 
         if self.use_sde:
+            latent_sde_dim = net_arch[-1]
+            # Separate feature extractor for SDE
+            if sde_net_arch is not None:
+                self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch, activation_fn)
+
             # TODO: check for the learn_features
             self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
                                                                learn_features=True, squash_output=True)
             self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
+                                                                            latent_sde_dim=latent_sde_dim,
                                                                             log_std_init=log_std_init)
         else:
             self.action_dist = SquashedDiagGaussianDistribution(action_dim)
@@ -61,38 +71,55 @@ class Actor(BaseNetwork):
         """
         self.action_dist.sample_weights(self.log_std, batch_size=batch_size)
 
+    def _get_latent(self, obs):
+        latent_pi = self.latent_pi(obs)
+
+        if self.sde_feature_extractor is not None:
+            latent_sde = self.sde_feature_extractor(obs)
+        else:
+            latent_sde = latent_pi
+        return latent_pi, latent_sde
+
     def get_action_dist_params(self, obs):
-        latent = self.actor_net(obs)
+        latent_pi, latent_sde = self._get_latent(obs)
 
         if self.use_sde:
-            mean_actions, log_std = self.mu(latent), self.log_std
+            mean_actions, log_std = self.mu(latent_pi), self.log_std
         else:
-            mean_actions, log_std = self.mu(latent), self.log_std(latent)
+            mean_actions, log_std = self.mu(latent_pi), self.log_std(latent_pi)
             # Original Implementation to cap the standard deviation
             log_std = th.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX)
-        return mean_actions, log_std, latent
+        return mean_actions, log_std, latent_sde
 
     def forward(self, obs, deterministic=False):
-        mean_actions, log_std, latent = self.get_action_dist_params(obs)
+        mean_actions, log_std, latent_sde = self.get_action_dist_params(obs)
         if self.use_sde:
             # Note the action is squashed
-            action, _ = self.action_dist.proba_distribution(mean_actions, log_std, latent, deterministic=deterministic)
+            action, _ = self.action_dist.proba_distribution(mean_actions, log_std, latent_sde, deterministic=deterministic)
         else:
             # Note the action is squashed
             action, _ = self.action_dist.proba_distribution(mean_actions, log_std, deterministic=deterministic)
         return action
 
     def action_log_prob(self, obs):
-        mean_actions, log_std, latent = self.get_action_dist_params(obs)
+        mean_actions, log_std, latent_sde = self.get_action_dist_params(obs)
 
         if self.use_sde:
-            action, log_prob = self.action_dist.log_prob_from_params(mean_actions, self.log_std, latent)
+            action, log_prob = self.action_dist.log_prob_from_params(mean_actions, self.log_std, latent_sde)
         else:
             action, log_prob = self.action_dist.log_prob_from_params(mean_actions, log_std)
         return action, log_prob
 
 
 class Critic(BaseNetwork):
+    """
+    Critic network (q-value function) for SAC.
+
+    :param obs_dim: (int) Dimension of the observation
+    :param action_dim: (int) Dimension of the action space
+    :param net_arch: ([int]) Network architecture
+    :param activation_fn: (nn.Module) Activation function
+    """
     def __init__(self, obs_dim, action_dim,
                  net_arch, activation_fn=nn.ReLU):
         super(Critic, self).__init__()
@@ -114,9 +141,25 @@ class Critic(BaseNetwork):
 
 
 class SACPolicy(BasePolicy):
+    """
+    Policy class (with both actor and critic) for SAC.
+
+    :param observation_space: (gym.spaces.Space) Observation space
+    :param action_dim: (gym.spaces.Space) Action space
+    :param learning_rate: (callable) Learning rate schedule (could be constant)
+    :param net_arch: ([int or dict]) The specification of the policy and value networks.
+    :param device: (str or th.device) Device on which the code should run.
+    :param activation_fn: (nn.Module) Activation function
+    :param use_sde: (bool) Whether to use State Dependent Exploration or not
+    :param log_std_init: (float) Initial value for the log standard deviation
+    :param sde_net_arch: ([int]) Network architecture for extracting features
+        when using SDE. If None, the latent features from the policy will be used.
+        Pass an empty list to use the states as features.
+    """
     def __init__(self, observation_space, action_space,
                  learning_rate, net_arch=None, device='cpu',
-                 activation_fn=nn.ReLU, use_sde=False, log_std_init=-3):
+                 activation_fn=nn.ReLU, use_sde=False,
+                 log_std_init=-3, sde_net_arch=None):
         super(SACPolicy, self).__init__(observation_space, action_space, device)
 
         if net_arch is None:
@@ -133,8 +176,12 @@ class SACPolicy(BasePolicy):
             'activation_fn': self.activation_fn
         }
         self.actor_kwargs = self.net_args.copy()
-        self.actor_kwargs['use_sde'] = use_sde
-        self.actor_kwargs['log_std_init'] = log_std_init
+        sde_kwargs = {
+            'use_sde': use_sde,
+            'log_std_init': log_std_init,
+            'sde_net_arch': sde_net_arch
+        }
+        self.actor_kwargs.update(sde_kwargs)
         self.actor, self.actor_target = None, None
         self.critic, self.critic_target = None, None
 
diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py
index 6872213..65f2171 100644
--- a/torchy_baselines/td3/policies.py
+++ b/torchy_baselines/td3/policies.py
@@ -1,7 +1,7 @@
 import torch as th
 import torch.nn as nn
 
-from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork
+from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, create_sde_feature_extractor
 from torchy_baselines.common.distributions import StateDependentNoiseDistribution
 
 
@@ -19,10 +19,13 @@ class Actor(BaseNetwork):
     :param lr_sde: (float) Learning rate for the standard deviation of the noise
     :param full_std: (bool) Whether to use (n_features x n_actions) parameters
         for the std instead of only (n_features,) when using SDE.
+    :param sde_net_arch: ([int]) Network architecture for extracting features
+        when using SDE. If None, the latent features from the policy will be used.
+        Pass an empty list to use the states as features.
     """
     def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
                  use_sde=False, log_std_init=-3, clip_noise=None,
-                 lr_sde=3e-4, full_std=False):
+                 lr_sde=3e-4, full_std=False, sde_net_arch=None):
         super(Actor, self).__init__()
 
         self.latent_pi, self.log_std = None, None
@@ -30,23 +33,32 @@ class Actor(BaseNetwork):
         self.use_sde, self.sde_optimizer = use_sde, None
         self.action_dim = action_dim
         self.full_std = full_std
+        self.sde_feature_extractor = None
 
         if use_sde:
-            latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False)
-            self.latent_pi = nn.Sequential(*latent_pi)
+            latent_pi_net = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False)
+            self.latent_pi = nn.Sequential(*latent_pi_net)
+            latent_sde_dim = net_arch[-1]
+            learn_features = sde_net_arch is not None
+
+            # Separate feature extractor for SDE
+            if sde_net_arch is not None:
+                self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch, activation_fn)
+
             # Create state dependent noise matrix (SDE)
             self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
-                                                               squash_output=False)
+                                                               squash_output=False, learn_features=learn_features)
             action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
+                                                                               latent_sde_dim=latent_sde_dim,
                                                                                log_std_init=log_std_init)
             # Squash output
-            self.actor_net = nn.Sequential(action_net, nn.Tanh())
+            self.mu = nn.Sequential(action_net, nn.Tanh())
             self.clip_noise = clip_noise
             self.sde_optimizer = th.optim.Adam([self.log_std], lr=lr_sde)
             self.reset_noise()
         else:
             actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True)
-            self.actor_net = nn.Sequential(*actor_net)
+            self.mu = nn.Sequential(*actor_net)
 
     def get_std(self):
         """
@@ -60,9 +72,18 @@ class Actor(BaseNetwork):
         """
         return self.action_dist.get_std(self.log_std)
 
-    def _get_action_dist_from_latent(self, latent_pi):
-        mean_actions = self.actor_net(latent_pi)
-        return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_pi)
+    def _get_action_dist_from_latent(self, latent_pi, latent_sde):
+        mean_actions = self.mu(latent_pi)
+        return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde)
+
+    def _get_latent(self, obs):
+        latent_pi = self.latent_pi(obs)
+
+        if self.sde_feature_extractor is not None:
+            latent_sde = self.sde_feature_extractor(obs)
+        else:
+            latent_sde = latent_pi
+        return latent_pi, latent_sde
 
     def evaluate_actions(self, obs, action):
         """
@@ -75,9 +96,8 @@ class Actor(BaseNetwork):
         :return: (th.Tensor, th.Tensor) log likelihood of taking those actions
             and entropy of the action distribution.
         """
-        with th.no_grad():
-            latent_pi = self.latent_pi(obs)
-        _, distribution = self._get_action_dist_from_latent(latent_pi)
+        latent_pi, latent_sde = self._get_latent(obs)
+        _, distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
         log_prob = distribution.log_prob(action)
         # value = self.value_net(latent_vf)
         return log_prob, distribution.entropy()
@@ -90,20 +110,21 @@ class Actor(BaseNetwork):
 
     def forward(self, obs, deterministic=True):
         if self.use_sde:
-            latent_pi = self.latent_pi(obs)
+            latent_pi, latent_sde = self._get_latent(obs)
             if deterministic:
-                return self.actor_net(latent_pi)
-            noise = self.action_dist.get_noise(latent_pi)
+                return self.mu(latent_pi)
+
+            noise = self.action_dist.get_noise(latent_sde)
             if self.clip_noise is not None:
                 noise = th.clamp(noise, -self.clip_noise, self.clip_noise)
             # TODO: Replace with squashing -> need to account for that in the sde update
             # -> set squash_out=True in the action_dist?
             # NOTE: the clipping is done in the rollout for now
-            return self.actor_net(latent_pi) + noise
+            return self.mu(latent_pi) + noise
             # action, _ = self._get_action_dist_from_latent(latent_pi)
             # return action
         else:
-            return self.actor_net(obs)
+            return self.mu(obs)
 
 
 class Critic(BaseNetwork):
@@ -148,11 +169,14 @@ class TD3Policy(BasePolicy):
     :param activation_fn: (nn.Module) Activation function
     :param use_sde: (bool) Whether to use State Dependent Exploration or not
     :param log_std_init: (float) Initial value for the log standard deviation
+    :param sde_net_arch: ([int]) Network architecture for extracting features
+        when using SDE. If None, the latent features from the policy will be used.
+        Pass an empty list to use the states as features.
     """
     def __init__(self, observation_space, action_space,
                  learning_rate, net_arch=None, device='cpu',
                  activation_fn=nn.ReLU, use_sde=False, log_std_init=-3,
-                 clip_noise=None, lr_sde=3e-4):
+                 clip_noise=None, lr_sde=3e-4, sde_net_arch=None):
         super(TD3Policy, self).__init__(observation_space, action_space, device)
 
         # Default network architecture, from the original paper
@@ -170,10 +194,14 @@ class TD3Policy(BasePolicy):
             'activation_fn': self.activation_fn
         }
         self.actor_kwargs = self.net_args.copy()
-        self.actor_kwargs['use_sde'] = use_sde
-        self.actor_kwargs['log_std_init'] = log_std_init
-        self.actor_kwargs['clip_noise'] = clip_noise
-        self.actor_kwargs['lr_sde'] = lr_sde
+        sde_kwargs = {
+            'use_sde': use_sde,
+            'log_std_init': log_std_init,
+            'clip_noise': clip_noise,
+            'lr_sde': lr_sde,
+            'sde_net_arch': sde_net_arch
+        }
+        self.actor_kwargs.update(sde_kwargs)
 
         self.actor, self.actor_target = None, None
         self.critic, self.critic_target = None, None

From 8e9802784c7f50f139382a716d37a768fb8b80c1 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 2 Dec 2019 14:06:30 +0100
Subject: [PATCH 13/47] Revert previous changes in SAC + SDE

---
 torchy_baselines/sac/sac.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index 66f3492..82538b1 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -175,8 +175,10 @@ class SAC(BaseRLModel):
             # or sample again the noise matrix
             # otherwise the intermediate step `std = th.exp(log_std)`
             # is lost and we cannot backpropagate through again
+            # anyway, we need to sample because `log_std` may have changed between two gradient steps
             if self.use_sde:
-                self.actor.reset_noise(batch_size=batch_size)
+                # self.actor.reset_noise(batch_size=batch_size)
+                self.actor.reset_noise()
 
             # Action by the current actor for the sampled state
             action_pi, log_prob = self.actor.action_log_prob(obs)
@@ -201,8 +203,8 @@ class SAC(BaseRLModel):
 
 
             with th.no_grad():
-                if self.use_sde:
-                    self.actor.reset_noise(batch_size=batch_size)
+                # if self.use_sde:
+                #     self.actor.reset_noise(batch_size=batch_size)
                 # Select action according to policy
                 next_action, next_log_prob = self.actor.action_log_prob(next_obs)
                 # Compute the target Q value
@@ -232,10 +234,7 @@ class SAC(BaseRLModel):
 
             # Optimize the actor
             self.actor.optimizer.zero_grad()
-            # Cf comment above, otherwise pytorch raises an error
-            # ("Trying to backward through the graph a second time")
-            # retain_graph = True if self.use_sde and gradient_steps > 1 else False
-            actor_loss.backward(retain_graph=False)
+            actor_loss.backward()
             self.actor.optimizer.step()
 
             # Update target networks

From 3cdd5f20afc766a89fd2fa0a06b33292a36a2d64 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 2 Dec 2019 14:14:48 +0100
Subject: [PATCH 14/47] Bug fix + add test for sde net arch

---
 tests/test_sde.py                | 23 +++++------------------
 torchy_baselines/ppo/policies.py |  3 ++-
 2 files changed, 7 insertions(+), 19 deletions(-)

diff --git a/tests/test_sde.py b/tests/test_sde.py
index 668bc9e..851ca42 100644
--- a/tests/test_sde.py
+++ b/tests/test_sde.py
@@ -9,7 +9,7 @@ from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize
 from torchy_baselines.common.monitor import Monitor
 
 
-def test_state_dependent_exploration():
+def test_state_dependent_exploration_grad():
     """
     Check that the gradient correspond to the expected one
     """
@@ -19,7 +19,6 @@ def test_state_dependent_exploration():
     sigma_hat = th.ones(state_dim, action_dim, requires_grad=True)
     # Reduce the number of parameters
     # sigma_ = th.ones(state_dim, action_dim) * sigma_
-
     # weights_dist = Normal(th.zeros_like(log_sigma), th.exp(log_sigma))
     th.manual_seed(2)
     weights_dist = Normal(th.zeros_like(sigma_hat), sigma_hat)
@@ -60,23 +59,11 @@ def test_state_dependent_exploration():
     assert sigma_hat.grad.allclose(grad)
 
 
-@pytest.mark.parametrize("model_class", [A2C])
-@pytest.mark.parametrize("sde_net_arch", [None, [32, 16]])
-def test_state_dependent_noise(model_class, sde_net_arch):
-    env_id = 'MountainCarContinuous-v0'
-
-    env = VecNormalize(DummyVecEnv([lambda: Monitor(gym.make(env_id))]), norm_reward=True)
-    eval_env = VecNormalize(DummyVecEnv([lambda: Monitor(gym.make(env_id))]), training=False, norm_reward=False)
-
-    model = model_class('MlpPolicy', env, n_steps=200, use_sde=True, ent_coef=0.00, verbose=1, learning_rate=3e-4,
-                        policy_kwargs=dict(log_std_init=0.0, ortho_init=False, sde_net_arch=sde_net_arch), seed=None)
-    model.learn(total_timesteps=int(1000), log_interval=5, eval_freq=500, eval_env=eval_env)
-
-
-@pytest.mark.parametrize("model_class", [TD3, SAC])
-def test_state_dependent_offpolicy_noise(model_class):
+@pytest.mark.parametrize("model_class", [TD3, SAC, A2C])
+@pytest.mark.parametrize("sde_net_arch", [None, [32, 16], []])
+def test_state_dependent_offpolicy_noise(model_class, sde_net_arch):
     model = model_class('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True,
-                        verbose=1, policy_kwargs=dict(log_std_init=-2))
+                        verbose=1, policy_kwargs=dict(log_std_init=-2, sde_net_arch=sde_net_arch))
     model.learn(total_timesteps=int(1000), eval_freq=500)
 
 
diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index 9b38317..886d958 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -4,7 +4,7 @@ import torch as th
 import torch.nn as nn
 import numpy as np
 
-from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor, create_mlp
+from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor, create_mlp, create_sde_feature_extractor
 from torchy_baselines.common.distributions import make_proba_distribution,\
     DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution
 
@@ -143,6 +143,7 @@ class PPOPolicy(BasePolicy):
             return self.action_dist.proba_distribution(mean_actions, self.log_std, deterministic=deterministic)
 
         elif isinstance(self.action_dist, CategoricalDistribution):
+            # Here mean_actions are the logits before the softmax
             return self.action_dist.proba_distribution(mean_actions, deterministic=deterministic)
 
         elif isinstance(self.action_dist, StateDependentNoiseDistribution):

From 1ac1a7cad5ce7aa9a2acb6c893f2aaf5994e8235 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 2 Dec 2019 14:27:38 +0100
Subject: [PATCH 15/47] Reformat

---
 torchy_baselines/a2c/a2c.py              |  5 ++-
 torchy_baselines/common/distributions.py | 41 ++++++++++++++----------
 torchy_baselines/ppo/policies.py         | 18 +++++++----
 torchy_baselines/ppo/ppo.py              |  3 +-
 torchy_baselines/sac/policies.py         | 14 +++++---
 torchy_baselines/td3/policies.py         |  9 +++---
 6 files changed, 53 insertions(+), 37 deletions(-)

diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py
index 4aa2dbe..33eae78 100644
--- a/torchy_baselines/a2c/a2c.py
+++ b/torchy_baselines/a2c/a2c.py
@@ -4,7 +4,6 @@ import torch.nn.functional as F
 
 from torchy_baselines.common.utils import explained_variance
 from torchy_baselines.ppo.ppo import PPO
-from torchy_baselines.ppo.policies import PPOPolicy
 from torchy_baselines.common import logger
 
 
@@ -132,5 +131,5 @@ class A2C(PPO):
               eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="A2C", reset_num_timesteps=True):
 
         return super(A2C, self).learn(total_timesteps=total_timesteps, callback=callback, log_interval=log_interval,
-                  eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes,
-                  tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps)
+                                      eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes,
+                                      tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps)
diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py
index 8802687..05a5588 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -1,10 +1,9 @@
-import numpy as np
 import torch as th
 import torch.nn as nn
 from torch.distributions import Normal, Categorical
-import torch.nn.functional as F
 from gym import spaces
 
+
 class Distribution(object):
     def __init__(self):
         super(Distribution, self).__init__()
@@ -13,25 +12,25 @@ class Distribution(object):
         """
         returns the log likelihood
 
-        :param x: (str) the labels of each index
-        :return: ([float]) The log likelihood of the distribution
+        :param x: (object) the taken action
+        :return: (th.Tensor) The log likelihood of the distribution
         """
         raise NotImplementedError
 
-    def kl_div(self, other):
-        """
-        Calculates the Kullback-Leibler divergence from the given probabilty distribution
-
-        :param other: ([float]) the distribution to compare with
-        :return: (float) the KL divergence of the two distributions
-        """
-        raise NotImplementedError
+    # def kl_div(self, other):
+    #     """
+    #     Calculates the Kullback-Leibler divergence from the given probabilty distribution
+    #
+    #     :param other: ([float]) the distribution to compare with
+    #     :return: (float) the KL divergence of the two distributions
+    #     """
+    #     raise NotImplementedError
 
     def entropy(self):
         """
         Returns shannon's entropy of the probability
 
-        :return: (float) the entropy
+        :return: (th.Tensor) the entropy
         """
         raise NotImplementedError
 
@@ -39,7 +38,7 @@ class Distribution(object):
         """
         returns a sample from the probabilty distribution
 
-        :return: (Tensorflow Tensor) the stochastic action
+        :return: (th.Tensor) the stochastic action
         """
         raise NotImplementedError
 
@@ -51,6 +50,7 @@ class DiagGaussianDistribution(Distribution):
 
     :param action_dim: (int)  Number of continuous actions
     """
+
     def __init__(self, action_dim):
         super(DiagGaussianDistribution, self).__init__()
         self.distribution = None
@@ -137,6 +137,7 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
     :param action_dim: (int) Number of continuous actions
     :param epsilon: (float) small value to avoid NaN due to numerical imprecision.
     """
+
     def __init__(self, action_dim, epsilon=1e-6):
         super(SquashedDiagGaussianDistribution, self).__init__(action_dim)
         # Avoid NaN (prevents division by zero or log of zero)
@@ -144,7 +145,8 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
         self.gaussian_action = None
 
     def proba_distribution(self, mean_actions, log_std, deterministic=False):
-        action, _ = super(SquashedDiagGaussianDistribution, self).proba_distribution(mean_actions, log_std, deterministic)
+        action, _ = super(SquashedDiagGaussianDistribution, self).proba_distribution(mean_actions, log_std,
+                                                                                     deterministic)
         return action, self
 
     def mode(self):
@@ -183,6 +185,7 @@ class CategoricalDistribution(Distribution):
 
     :param action_dim: (int) Number of discrete actions
     """
+
     def __init__(self, action_dim):
         super(CategoricalDistribution, self).__init__()
         self.distribution = None
@@ -246,6 +249,7 @@ class StateDependentNoiseDistribution(Distribution):
         `latent_sde` in the code.
     :param epsilon: (float) small value to avoid NaN due to numerical imprecision.
     """
+
     def __init__(self, action_dim, full_std=True, use_expln=False,
                  squash_output=False, learn_features=False, epsilon=1e-6):
         super(StateDependentNoiseDistribution, self).__init__()
@@ -256,12 +260,12 @@ class StateDependentNoiseDistribution(Distribution):
         self.log_std = None
         self.weights_dist = None
         self.exploration_mat = None
+        self.exploration_matrices = None
         self.use_expln = use_expln
         self.full_std = full_std
         self.epsilon = epsilon
         self.learn_features = learn_features
         if squash_output:
-            print("== Using TanhBijector ===")
             self.bijector = TanhBijector(epsilon)
         else:
             self.bijector = None
@@ -296,6 +300,7 @@ class StateDependentNoiseDistribution(Distribution):
         using a centered gaussian distribution.
 
         :param log_std: (th.Tensor)
+        :param batch_size: (int)
         """
         std = self.get_std(log_std)
         self.weights_dist = Normal(th.zeros_like(std), std)
@@ -333,6 +338,7 @@ class StateDependentNoiseDistribution(Distribution):
 
         :param mean_actions: (th.Tensor)
         :param log_std: (th.Tensor)
+        :param latent_sde: (th.Tensor)
         :param deterministic: (bool)
         :return: (th.Tensor)
         """
@@ -408,6 +414,7 @@ class TanhBijector(object):
 
     :param epsilon: (float) small value to avoid NaN due to numerical imprecision.
     """
+
     def __init__(self, epsilon=1e-6):
         super(TanhBijector, self).__init__()
         self.epsilon = epsilon
@@ -439,7 +446,7 @@ class TanhBijector(object):
 
     def log_prob_correction(self, x):
         # Squash correction (from original SAC implementation)
-        return th.log(1 - th.tanh(x) ** 2 + self.epsilon)
+        return th.log(1.0 - th.tanh(x) ** 2 + self.epsilon)
 
 
 def make_proba_distribution(action_space, use_sde=False, dist_kwargs=None):
diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index 886d958..19e20a5 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -4,7 +4,8 @@ import torch as th
 import torch.nn as nn
 import numpy as np
 
-from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor, create_mlp, create_sde_feature_extractor
+from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor, \
+    create_sde_feature_extractor
 from torchy_baselines.common.distributions import make_proba_distribution,\
     DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution
 
@@ -14,7 +15,7 @@ class PPOPolicy(BasePolicy):
     Policy class (with both actor and critic) for A2C and derivates (PPO).
 
     :param observation_space: (gym.spaces.Space) Observation space
-    :param action_dim: (gym.spaces.Space) Action space
+    :param action_space: (gym.spaces.Space) Action space
     :param learning_rate: (callable) Learning rate schedule (could be constant)
     :param net_arch: ([int or dict]) The specification of the policy and value networks.
     :param device: (str or th.device) Device on which the code should run.
@@ -87,21 +88,24 @@ class PPOPolicy(BasePolicy):
         self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch,
                                           activation_fn=self.activation_fn, device=self.device)
 
+        latent_dim_pi = self.mlp_extractor.latent_dim_pi
+
         # Separate feature extractor for SDE
         if self.sde_net_arch is not None:
-            self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(self.features_dim, self.sde_net_arch,
+            self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(self.features_dim,
+                                                                                      self.sde_net_arch,
                                                                                       self.activation_fn)
 
         if isinstance(self.action_dist, DiagGaussianDistribution):
-            self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
+            self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi,
                                                                                     log_std_init=self.log_std_init)
         elif isinstance(self.action_dist, StateDependentNoiseDistribution):
-            latent_sde_dim = self.mlp_extractor.latent_dim_pi if self.sde_net_arch is None else latent_sde_dim
-            self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
+            latent_sde_dim = latent_dim_pi if self.sde_net_arch is None else latent_sde_dim
+            self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi,
                                                                                     latent_sde_dim=latent_sde_dim,
                                                                                     log_std_init=self.log_std_init)
         elif isinstance(self.action_dist, CategoricalDistribution):
-            self.action_net = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi)
+            self.action_net = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi)
 
         self.value_net = nn.Linear(self.mlp_extractor.latent_dim_vf, 1)
         # Init weights: use orthogonal initialization
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index 99ed1b4..2ffc6e6 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -43,7 +43,8 @@ class PPO(BaseRLModel):
     :param n_epochs: (int) Number of epoch when optimizing the surrogate loss
     :param gamma: (float) Discount factor
     :param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator
-    :param clip_range: (float or callable) Clipping parameter, it can be a function of the current progress (from 1 to 0).
+    :param clip_range: (float or callable) Clipping parameter, it can be a function of the current progress
+        (from 1 to 0).
     :param clip_range_vf: (float or callable) Clipping parameter for the value function,
         it can be a function of the current progress (from 1 to 0).
         This is a parameter specific to the OpenAI implementation. If None is passed (default),
diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index be32660..cd1b35b 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -1,7 +1,8 @@
 import torch as th
 import torch.nn as nn
 
-from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, create_sde_feature_extractor
+from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, \
+    create_sde_feature_extractor
 from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
 
 # CAP the standard deviation of the actor
@@ -38,7 +39,8 @@ class Actor(BaseNetwork):
             latent_sde_dim = net_arch[-1]
             # Separate feature extractor for SDE
             if sde_net_arch is not None:
-                self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch, activation_fn)
+                self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch,
+                                                                                          activation_fn)
 
             # TODO: check for the learn_features
             self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
@@ -95,10 +97,12 @@ class Actor(BaseNetwork):
         mean_actions, log_std, latent_sde = self.get_action_dist_params(obs)
         if self.use_sde:
             # Note the action is squashed
-            action, _ = self.action_dist.proba_distribution(mean_actions, log_std, latent_sde, deterministic=deterministic)
+            action, _ = self.action_dist.proba_distribution(mean_actions, log_std, latent_sde,
+                                                            deterministic=deterministic)
         else:
             # Note the action is squashed
-            action, _ = self.action_dist.proba_distribution(mean_actions, log_std, deterministic=deterministic)
+            action, _ = self.action_dist.proba_distribution(mean_actions, log_std,
+                                                            deterministic=deterministic)
         return action
 
     def action_log_prob(self, obs):
@@ -145,7 +149,7 @@ class SACPolicy(BasePolicy):
     Policy class (with both actor and critic) for SAC.
 
     :param observation_space: (gym.spaces.Space) Observation space
-    :param action_dim: (gym.spaces.Space) Action space
+    :param action_space: (gym.spaces.Space) Action space
     :param learning_rate: (callable) Learning rate schedule (could be constant)
     :param net_arch: ([int or dict]) The specification of the policy and value networks.
     :param device: (str or th.device) Device on which the code should run.
diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py
index 65f2171..a472b58 100644
--- a/torchy_baselines/td3/policies.py
+++ b/torchy_baselines/td3/policies.py
@@ -1,7 +1,8 @@
 import torch as th
 import torch.nn as nn
 
-from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, create_sde_feature_extractor
+from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, \
+    create_sde_feature_extractor
 from torchy_baselines.common.distributions import StateDependentNoiseDistribution
 
 
@@ -43,7 +44,8 @@ class Actor(BaseNetwork):
 
             # Separate feature extractor for SDE
             if sde_net_arch is not None:
-                self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch, activation_fn)
+                self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch,
+                                                                                          activation_fn)
 
             # Create state dependent noise matrix (SDE)
             self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
@@ -92,7 +94,6 @@ class Actor(BaseNetwork):
 
         :param obs: (th.Tensor)
         :param action: (th.Tensor)
-        :param deterministic: (bool)
         :return: (th.Tensor, th.Tensor) log likelihood of taking those actions
             and entropy of the action distribution.
         """
@@ -162,7 +163,7 @@ class TD3Policy(BasePolicy):
     Policy class (with both actor and critic) for TD3.
 
     :param observation_space: (gym.spaces.Space) Observation space
-    :param action_dim: (gym.spaces.Space) Action space
+    :param action_space: (gym.spaces.Space) Action space
     :param learning_rate: (callable) Learning rate schedule (could be constant)
     :param net_arch: ([int or dict]) The specification of the policy and value networks.
     :param device: (str or th.device) Device on which the code should run.

From 6c423add8d5de35399591067919efa208e907f0c Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Thu, 5 Dec 2019 16:44:27 +0100
Subject: [PATCH 16/47] Bump version

---
 setup.py                     | 2 +-
 torchy_baselines/__init__.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 640112e..78a72bf 100644
--- a/setup.py
+++ b/setup.py
@@ -34,7 +34,7 @@ setup(name='torchy_baselines',
       license="MIT",
       long_description="",
       long_description_content_type='text/markdown',
-      version="0.0.6a",
+      version="0.0.7a",
       )
 
 # python setup.py sdist
diff --git a/torchy_baselines/__init__.py b/torchy_baselines/__init__.py
index b8383db..56d1dfe 100644
--- a/torchy_baselines/__init__.py
+++ b/torchy_baselines/__init__.py
@@ -4,4 +4,4 @@ from torchy_baselines.ppo import PPO
 from torchy_baselines.sac import SAC
 from torchy_baselines.td3 import TD3
 
-__version__ = "0.0.6a"
+__version__ = "0.0.7a"

From 233f346d53d5bf4ef432a3ab3d193f356fb7d7b9 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Fri, 6 Dec 2019 17:46:56 +0100
Subject: [PATCH 17/47] Update todos

---
 README.md                                | 3 ++-
 torchy_baselines/common/distributions.py | 9 ---------
 torchy_baselines/ppo/ppo.py              | 4 ++--
 3 files changed, 4 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index bf14287..388570a 100644
--- a/README.md
+++ b/README.md
@@ -18,11 +18,11 @@ PyTorch version of [Stable Baselines](https://github.com/hill-a/stable-baselines
 ## Roadmap
 
 TODO:
-- save/load
 - better predict
 - complete logger
 - Refactor: buffer with numpy array instead of pytorch
 - Refactor: remove duplicated code for evaluation
+- double check the shape of log prob
 
 - plotting? -> zoo
 
@@ -30,6 +30,7 @@ Later:
 - get_parameters / set_parameters
 - SDE: use [affine transform](https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/Affine)
   to scale the noise after a tanh transform?
+- Use MultivariateNormal with full covariance matrix?
 - CNN policies + normalization
 - tensorboard support
 - DQN
diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py
index 05a5588..97eaf6f 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -17,15 +17,6 @@ class Distribution(object):
         """
         raise NotImplementedError
 
-    # def kl_div(self, other):
-    #     """
-    #     Calculates the Kullback-Leibler divergence from the given probabilty distribution
-    #
-    #     :param other: ([float]) the distribution to compare with
-    #     :return: (float) the KL divergence of the two distributions
-    #     """
-    #     raise NotImplementedError
-
     def entropy(self):
         """
         Returns shannon's entropy of the probability
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index 6fa2875..113e050 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -311,8 +311,8 @@ class PPO(BaseRLModel):
     def get_opt_parameters(self):
         """
         Returns a dict of all the optimizers and their parameters
-        
-        :return: (dict) of optimizer names and their state_dict 
+
+        :return: (dict) of optimizer names and their state_dict
         """
         return {"opt": self.policy.optimizer.state_dict()}
 

From d63cef7693e630fe42541471902168e655210375 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Fri, 6 Dec 2019 18:32:57 +0100
Subject: [PATCH 18/47] Add gradient clipping for SAC

---
 README.md                   | 2 +-
 torchy_baselines/sac/sac.py | 7 ++++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 388570a..8ed3ae3 100644
--- a/README.md
+++ b/README.md
@@ -23,7 +23,7 @@ TODO:
 - Refactor: buffer with numpy array instead of pytorch
 - Refactor: remove duplicated code for evaluation
 - double check the shape of log prob
-
+- try squashing both mean and output when using SAC + SDE
 - plotting? -> zoo
 
 Later:
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index b5d2a17..e170cd0 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -44,6 +44,7 @@ class SAC(BaseRLModel):
     :param target_entropy: (str or float) target entropy when learning ent_coef (ent_coef = 'auto')
     :param action_noise: (ActionNoise) the action noise type (None by default), this can help
         for hard exploration problem. Cf common.noise for the different action noise type.
+    :param max_grad_norm: (float) The maximum value for the gradient clipping (None by default)
     :param gamma: (float) the discount factor
     :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
         instead of action noise exploration (default: False)
@@ -61,7 +62,7 @@ class SAC(BaseRLModel):
                  learning_starts=100, batch_size=256,
                  tau=0.005, ent_coef='auto', target_update_interval=1,
                  train_freq=1, gradient_steps=1, n_episodes_rollout=-1,
-                 target_entropy='auto', action_noise=None,
+                 target_entropy='auto', action_noise=None, max_grad_norm=None,
                  gamma=0.99, use_sde=False, tensorboard_log=None, create_eval_env=False,
                  policy_kwargs=None, verbose=0, seed=0, device='auto',
                  _init_setup_model=True):
@@ -89,6 +90,7 @@ class SAC(BaseRLModel):
         self.action_noise = action_noise
         self.gamma = gamma
         self.ent_coef_optimizer = None
+        self.max_grad_norm = max_grad_norm
 
         if _init_setup_model:
             self._setup_model()
@@ -236,6 +238,9 @@ class SAC(BaseRLModel):
             # Optimize the actor
             self.actor.optimizer.zero_grad()
             actor_loss.backward()
+            # Clip grad norm
+            if self.max_grad_norm is not None:
+                th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
             self.actor.optimizer.step()
 
             # Update target networks

From 919dfee45239ddd2c23291baf722f0ac870c5ad4 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Tue, 17 Dec 2019 11:14:44 +0100
Subject: [PATCH 19/47] Try to clip grad norm

---
 torchy_baselines/sac/sac.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index e170cd0..ee5ee4c 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -202,6 +202,9 @@ class SAC(BaseRLModel):
             if ent_coef_loss is not None:
                 self.ent_coef_optimizer.zero_grad()
                 ent_coef_loss.backward()
+                # Clip grad norm
+                if self.max_grad_norm is not None:
+                    th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
                 self.ent_coef_optimizer.step()
 
 
@@ -227,6 +230,9 @@ class SAC(BaseRLModel):
             # Optimize the critic
             self.critic.optimizer.zero_grad()
             critic_loss.backward()
+            # Clip grad norm
+            if self.max_grad_norm is not None:
+                th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
             self.critic.optimizer.step()
 
             # Compute actor loss

From 1d6f9bf1006e6e122637b0ae11ad91bc8dfe0025 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Tue, 17 Dec 2019 11:47:21 +0100
Subject: [PATCH 20/47] Add sample freq for SDE

---
 torchy_baselines/a2c/a2c.py           |  8 +++++---
 torchy_baselines/common/base_class.py | 10 +++++++++-
 torchy_baselines/ppo/policies.py      |  2 +-
 torchy_baselines/ppo/ppo.py           | 15 ++++++++++-----
 torchy_baselines/sac/sac.py           |  7 +++++--
 torchy_baselines/td3/td3.py           |  6 ++++--
 6 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py
index 33eae78..6fa667c 100644
--- a/torchy_baselines/a2c/a2c.py
+++ b/torchy_baselines/a2c/a2c.py
@@ -33,6 +33,8 @@ class A2C(PPO):
     :param use_rms_prop: (bool) Whether to use RMSprop (default) or Adam as optimizer
     :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
         instead of action noise exploration (default: False)
+    :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
+        Default: -1 (only sample at the beginning of the rollout)
     :param normalize_advantage: (bool) Whether to normalize or not the advantage
     :param tensorboard_log: (str) the log location for tensorboard (if None, no logging)
     :param create_eval_env: (bool) Whether to create a second environment that will be
@@ -44,11 +46,10 @@ class A2C(PPO):
         Setting it to auto, the code will be run on the GPU if possible.
     :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance
     """
-
     def __init__(self, policy, env, learning_rate=7e-4,
                  n_steps=5, gamma=0.99, gae_lambda=1.0,
                  ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5,
-                 rms_prop_eps=1e-5, use_rms_prop=True, use_sde=False,
+                 rms_prop_eps=1e-5, use_rms_prop=True, use_sde=False, sde_sample_freq=-1,
                  normalize_advantage=False, tensorboard_log=None, create_eval_env=False,
                  policy_kwargs=None, verbose=0, seed=0, device='auto',
                  _init_setup_model=True):
@@ -56,7 +57,8 @@ class A2C(PPO):
         super(A2C, self).__init__(policy, env, learning_rate=learning_rate,
                                   n_steps=n_steps, batch_size=None, n_epochs=1,
                                   gamma=gamma, gae_lambda=gae_lambda, ent_coef=ent_coef,
-                                  vf_coef=vf_coef, max_grad_norm=max_grad_norm, use_sde=use_sde,
+                                  vf_coef=vf_coef, max_grad_norm=max_grad_norm,
+                                  use_sde=use_sde, sde_sample_freq=sde_sample_freq,
                                   tensorboard_log=tensorboard_log, policy_kwargs=policy_kwargs,
                                   verbose=verbose, device=device, create_eval_env=create_eval_env,
                                   seed=seed, _init_setup_model=False)
diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index 576bea2..9073296 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -39,12 +39,15 @@ class BaseRLModel(object):
     :param seed: (int) Seed for the pseudo random generators
     :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
         instead of action noise exploration (default: False)
+    :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
+        Default: -1 (only sample at the beginning of the rollout)
     """
     __metaclass__ = ABCMeta
 
     def __init__(self, policy, env, policy_base, policy_kwargs=None,
                  verbose=0, device='auto', support_multi_env=False,
-                 create_eval_env=False, monitor_wrapper=True, seed=None, use_sde=False):
+                 create_eval_env=False, monitor_wrapper=True, seed=None,
+                 use_sde=False, sde_sample_freq=-1):
         if isinstance(policy, str) and policy_base is not None:
             self.policy_class = get_policy_from_name(policy_base, policy)
         else:
@@ -74,6 +77,7 @@ class BaseRLModel(object):
         self.rollout_data = None
         self.on_policy_exploration = False
         self.use_sde = use_sde
+        self.sde_sample_freq = sde_sample_freq
         # Track the training progress (from 1 to 0)
         # this is used to update the learning rate
         self._current_progress = 1
@@ -516,6 +520,10 @@ class BaseRLModel(object):
             episode_reward, episode_timesteps = 0.0, 0
 
             while not done:
+                if self.use_sde and self.sde_sample_freq > 0 and  n_steps % self.sde_sample_freq == 0:
+                    # Sample a new noise matrix
+                    self.actor.reset_noise()
+
                 # Select action randomly or according to policy
                 if num_timesteps < learning_starts:
                     # Warmup phase
diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index 19e20a5..8c4df3f 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -76,7 +76,7 @@ class PPOPolicy(BasePolicy):
 
         self._build(learning_rate)
 
-    def reset_noise_net(self, n_envs=1):
+    def reset_noise(self, n_envs=1):
         """
         Sample new weights for the exploration matrix.
 
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index 113e050..bc2c7d8 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -55,6 +55,8 @@ class PPO(BaseRLModel):
     :param max_grad_norm: (float) The maximum value for the gradient clipping
     :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
         instead of action noise exploration (default: False)
+    :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
+        Default: -1 (only sample at the beginning of the rollout)    
     :param target_kl: (float) Limit the KL divergence between updates,
         because the clipping is not enough to prevent large update
         see issue #213 (cf https://github.com/hill-a/stable-baselines/issues/213)
@@ -69,17 +71,17 @@ class PPO(BaseRLModel):
         Setting it to auto, the code will be run on the GPU if possible.
     :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance
     """
-
     def __init__(self, policy, env, learning_rate=3e-4,
                  n_steps=2048, batch_size=64, n_epochs=10,
                  gamma=0.99, gae_lambda=0.95, clip_range=0.2, clip_range_vf=None,
-                 ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5, use_sde=False,
+                 ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5,
+                 use_sde=False, sde_sample_freq=-1,
                  target_kl=None, tensorboard_log=None, create_eval_env=False,
                  policy_kwargs=None, verbose=0, seed=0, device='auto',
                  _init_setup_model=True):
 
         super(PPO, self).__init__(policy, env, PPOPolicy, policy_kwargs=policy_kwargs,
-                                  verbose=verbose, device=device,
+                                  verbose=verbose, device=device, use_sde=use_sde, sde_sample_freq=sde_sample_freq,
                                   create_eval_env=create_eval_env, support_multi_env=True, seed=seed)
 
         self.learning_rate = learning_rate
@@ -97,7 +99,6 @@ class PPO(BaseRLModel):
         self.target_kl = target_kl
         self.tensorboard_log = tensorboard_log
         self.tb_writer = None
-        self.use_sde = use_sde
 
         if _init_setup_model:
             self._setup_model()
@@ -158,9 +159,13 @@ class PPO(BaseRLModel):
         # Sample new weights for the state dependent exploration
         # TODO: ensure episodic setting?
         if self.use_sde:
-            self.policy.reset_noise_net(env.num_envs)
+            self.policy.reset_noise(env.num_envs)
 
         while n_steps < n_rollout_steps:
+            if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
+                # Sample a new noise matrix
+                self.policy.reset_noise(env.num_envs)
+
             with th.no_grad():
                 actions, values, log_probs = self.policy.forward(obs)
             actions = actions.cpu().numpy()
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index ee5ee4c..bc24aa3 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -48,6 +48,8 @@ class SAC(BaseRLModel):
     :param gamma: (float) the discount factor
     :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
         instead of action noise exploration (default: False)
+    :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
+        Default: -1 (only sample at the beginning of the rollout)
     :param create_eval_env: (bool) Whether to create a second environment that will be
         used for evaluating the agent periodically. (Only available when passing string for the environment)
     :param policy_kwargs: (dict) additional arguments to be passed to the policy on creation
@@ -63,12 +65,13 @@ class SAC(BaseRLModel):
                  tau=0.005, ent_coef='auto', target_update_interval=1,
                  train_freq=1, gradient_steps=1, n_episodes_rollout=-1,
                  target_entropy='auto', action_noise=None, max_grad_norm=None,
-                 gamma=0.99, use_sde=False, tensorboard_log=None, create_eval_env=False,
+                 gamma=0.99, use_sde=False, sde_sample_freq=-1,
+                 tensorboard_log=None, create_eval_env=False,
                  policy_kwargs=None, verbose=0, seed=0, device='auto',
                  _init_setup_model=True):
 
         super(SAC, self).__init__(policy, env, SACPolicy, policy_kwargs, verbose, device,
-                                  create_eval_env=create_eval_env, seed=seed, use_sde=use_sde)
+                                  create_eval_env=create_eval_env, seed=seed, use_sde=use_sde, sde_sample_freq=sde_sample_freq)
 
         self.learning_rate = learning_rate
         self.target_entropy = target_entropy
diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py
index 06eaff5..3e1dfb7 100644
--- a/torchy_baselines/td3/td3.py
+++ b/torchy_baselines/td3/td3.py
@@ -40,6 +40,8 @@ class TD3(BaseRLModel):
     :param target_noise_clip: (float) Limit for absolute value of target policy smoothing noise.
     :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
         instead of action noise exploration (default: False)
+    :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
+        Default: -1 (only sample at the beginning of the rollout)
     :param sde_max_grad_norm: (float)
     :param sde_ent_coef: (float)
     :param sde_log_std_scheduler: (callable)
@@ -57,12 +59,12 @@ class TD3(BaseRLModel):
                  policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100,
                  train_freq=-1, gradient_steps=-1, n_episodes_rollout=1,
                  tau=0.005, action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5,
-                 use_sde=False, sde_max_grad_norm=1, sde_ent_coef=0.0, sde_log_std_scheduler=None,
+                 use_sde=False, sde_sample_freq=-1, sde_max_grad_norm=1, sde_ent_coef=0.0, sde_log_std_scheduler=None,
                  tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0,
                  seed=0, device='auto', _init_setup_model=True):
 
         super(TD3, self).__init__(policy, env, TD3Policy, policy_kwargs, verbose, device,
-                                  create_eval_env=create_eval_env, seed=seed, use_sde=use_sde)
+                                  create_eval_env=create_eval_env, seed=seed, use_sde=use_sde, sde_sample_freq=sde_sample_freq)
 
         self.buffer_size = buffer_size
         self.learning_rate = learning_rate

From 57708a628c779e8b0f8b683b2bbe73665713da47 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Tue, 17 Dec 2019 15:01:08 +0100
Subject: [PATCH 21/47] Add value function for SDE + TD3

---
 torchy_baselines/common/base_class.py | 14 ++++++++++----
 torchy_baselines/td3/policies.py      | 27 +++++++++++++++++++++++++++
 torchy_baselines/td3/td3.py           | 26 +++++++++++++++-----------
 3 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index 9073296..6200bc1 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -511,7 +511,7 @@ class BaseRLModel(object):
             self.actor.reset_noise()
             # Reset rollout data
             if self.on_policy_exploration:
-                self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']}
+                self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones', 'values']}
 
         while total_steps < n_steps or total_episodes < n_episodes:
             done = False
@@ -574,6 +574,7 @@ class BaseRLModel(object):
                     self.rollout_data['actions'].append(scaled_action[0].copy())
                     self.rollout_data['rewards'].append(reward[0].copy())
                     self.rollout_data['dones'].append(np.array(done_bool[0]).copy())
+                    self.rollout_data['values'].append(self.vf_net(th.FloatTensor(obs).to(self.device))[0].cpu().detach().numpy())
 
                 obs = new_obs
                 # Save the true unnormalized observation
@@ -615,19 +616,24 @@ class BaseRLModel(object):
 
         # Post processing
         if self.rollout_data is not None:
-            for key in ['observations', 'actions', 'rewards', 'dones']:
+            for key in ['observations', 'actions', 'rewards', 'dones', 'values']:
                 self.rollout_data[key] = th.FloatTensor(np.array(self.rollout_data[key])).to(self.device)
 
             self.rollout_data['returns'] = self.rollout_data['rewards'].clone()
-            # Compute return
+            self.rollout_data['advantage'] = self.rollout_data['rewards'].clone()
+
+            # Compute return and advantage
             last_return = 0.0
             for step in reversed(range(len(self.rollout_data['rewards']))):
                 if step == len(self.rollout_data['rewards']) - 1:
-                    last_return = self.rollout_data['rewards'][step]
+                    next_non_terminal = 1.0 - done[0]
+                    next_value = self.vf_net(th.FloatTensor(obs).to(self.device))[0].detach()
+                    last_return = self.rollout_data['rewards'][step] + next_non_terminal * next_value
                 else:
                     next_non_terminal = 1.0 - self.rollout_data['dones'][step + 1]
                     last_return = self.rollout_data['rewards'][step] + self.gamma * last_return * next_non_terminal
                 self.rollout_data['returns'][step] = last_return
+            self.rollout_data['advantage'] = self.rollout_data['returns'] - self.rollout_data['values']
 
         return mean_reward, total_steps, total_episodes, obs
 
diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py
index a472b58..1fd6a20 100644
--- a/torchy_baselines/td3/policies.py
+++ b/torchy_baselines/td3/policies.py
@@ -158,6 +158,27 @@ class Critic(BaseNetwork):
         return self.q_networks[0](th.cat([obs, action], dim=1))
 
 
+class ValueFunction(BaseNetwork):
+    """
+    Value function for TD3 when doing on-policy exploration with SDE.
+
+    :param obs_dim: (int) Dimension of the observation
+    :param net_arch: ([int]) Network architecture
+    :param activation_fn: (nn.Module) Activation function
+    """
+    def __init__(self, obs_dim, net_arch=None, activation_fn=nn.Tanh):
+        super(ValueFunction, self).__init__()
+
+        if net_arch is None:
+            net_arch = [64, 64]
+
+        vf_net = create_mlp(obs_dim, 1, net_arch, activation_fn)
+        self.vf_net = nn.Sequential(*vf_net)
+
+    def forward(self, obs):
+        return self.vf_net(obs)
+
+
 class TD3Policy(BasePolicy):
     """
     Policy class (with both actor and critic) for TD3.
@@ -206,7 +227,9 @@ class TD3Policy(BasePolicy):
 
         self.actor, self.actor_target = None, None
         self.critic, self.critic_target = None, None
+        # For SDE only
         self.use_sde = use_sde
+        self.vf_net = None
         self.log_std_init = log_std_init
         self._build(learning_rate)
 
@@ -221,6 +244,10 @@ class TD3Policy(BasePolicy):
         self.critic_target.load_state_dict(self.critic.state_dict())
         self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))
 
+        if self.use_sde:
+            self.vf_net = ValueFunction(self.obs_dim)
+            self.actor.sde_optimizer.add_param_group({'params': self.vf_net.parameters()})
+
     def reset_noise(self):
         return self.actor.reset_noise()
 
diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py
index 3e1dfb7..9b5cafd 100644
--- a/torchy_baselines/td3/td3.py
+++ b/torchy_baselines/td3/td3.py
@@ -9,6 +9,7 @@ from torchy_baselines.common.buffers import ReplayBuffer
 from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.td3.policies import TD3Policy
 from torchy_baselines.common.vec_env import sync_envs_normalization
+from torchy_baselines.ppo.policies import MlpPolicy
 
 
 class TD3(BaseRLModel):
@@ -85,6 +86,7 @@ class TD3(BaseRLModel):
         self.sde_ent_coef = sde_ent_coef
         self.sde_log_std_scheduler = sde_log_std_scheduler
         self.on_policy_exploration = True
+        self.sde_vf = None
 
         if _init_setup_model:
             self._setup_model()
@@ -104,6 +106,7 @@ class TD3(BaseRLModel):
         self.actor_target = self.policy.actor_target
         self.critic = self.policy.critic
         self.critic_target = self.policy.critic_target
+        self.vf_net = self.policy.vf_net
 
     def select_action(self, observation, deterministic=True):
         # Normally not needed
@@ -209,25 +212,26 @@ class TD3(BaseRLModel):
         # self._update_learning_rate(self.policy.optimizer)
 
         # Unpack
-        obs, action, returns = [self.rollout_data[key] for key in ['observations', 'actions', 'returns']]
+        obs, action, advantage, returns = [self.rollout_data[key] for key in ['observations', 'actions', 'advantage', 'returns']]
 
-        # TODO: avoid second computation of everything because of the gradient
         log_prob, entropy = self.actor.evaluate_actions(obs, action)
+        values = self.vf_net(obs).flatten()
 
-        # Normalize returns
-        # returns = (returns - returns.mean()) / (returns.std() + 1e-8)
-        # returns = (returns - returns.mean())
-        with th.no_grad():
-            current_q1, current_q2 = self.critic(obs, action)
-        # Alternatively use the q value
-        returns = (returns - th.min(current_q1, current_q2))
+        # Normalize advantage
+        # if self.normalize_advantage:
+        #     advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8)
 
-        policy_loss = -(returns * log_prob).mean()
+        # Value loss using the TD(gae_lambda) target
+        value_loss = F.mse_loss(returns, values)
+
+        # A2C loss
+        policy_loss = -(advantage * log_prob).mean()
 
         # Entropy loss favor exploration
         entropy_loss = -th.mean(entropy)
 
-        loss = policy_loss + self.sde_ent_coef * entropy_loss
+        vf_coef = 0.5
+        loss = policy_loss + self.sde_ent_coef * entropy_loss + vf_coef * value_loss
 
         # Optimization step
         self.actor.sde_optimizer.zero_grad()

From e49d97bf98a0340f4173d5d8d51b42a6170dc453 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Wed, 18 Dec 2019 13:45:33 +0100
Subject: [PATCH 22/47] Fix infs in SAC by bounding the mean

---
 torchy_baselines/sac/policies.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index cd1b35b..610b59e 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -48,6 +48,8 @@ class Actor(BaseNetwork):
             self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
                                                                             latent_sde_dim=latent_sde_dim,
                                                                             log_std_init=log_std_init)
+            # Avoid saturation by limiting the mean of the gaussian to be in [-1, 1]
+            self.mu = nn.Sequential(self.mu, nn.Tanh())
         else:
             self.action_dist = SquashedDiagGaussianDistribution(action_dim)
             self.mu = nn.Linear(net_arch[-1], action_dim)

From 07345e5e2752de58ff694257641e955246935180 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Wed, 18 Dec 2019 13:45:56 +0100
Subject: [PATCH 23/47] Test for differential entropy

---
 tests/test_distributions.py | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/tests/test_distributions.py b/tests/test_distributions.py
index 130a1bc..3d896ab 100644
--- a/tests/test_distributions.py
+++ b/tests/test_distributions.py
@@ -1,3 +1,4 @@
+import pytest
 import numpy as np
 import torch as th
 
@@ -5,6 +6,7 @@ from torchy_baselines.common.utils import set_random_seed
 from torchy_baselines.common.distributions import DiagGaussianDistribution, SquashedDiagGaussianDistribution,\
     CategoricalDistribution, TanhBijector, StateDependentNoiseDistribution
 
+
 # TODO: more tests for the other distributions
 def test_bijector():
     """
@@ -37,3 +39,32 @@ def test_sde_distribution():
 
     assert th.allclose(actions.mean(), dist.distribution.mean.mean(), rtol=1e-3)
     assert th.allclose(actions.std(), dist.distribution.scale.mean(), rtol=1e-3)
+
+
+N_ACTIONS = 1
+
+# TODO: fix for num action > 1
+# TODO: analytical form for squashed Gaussian?
+@pytest.mark.parametrize("dist", [
+    DiagGaussianDistribution(N_ACTIONS),
+    StateDependentNoiseDistribution(N_ACTIONS, squash_output=False),
+])
+def test_entropy(dist):
+    # The entropy can be approximated by averaging the negative log likelihood
+    # mean negative log likelihood == differential entropy
+    n_samples = int(5e6)
+    n_features = 3
+    set_random_seed(1)
+    state = th.rand(n_samples, n_features)
+    deterministic_actions = th.rand(n_samples, N_ACTIONS)
+    _, log_std = dist.proba_distribution_net(n_features, log_std_init=th.log(th.tensor(0.2)))
+
+    if isinstance(dist, DiagGaussianDistribution):
+        actions, dist = dist.proba_distribution(deterministic_actions, log_std)
+    else:
+        dist.sample_weights(log_std, batch_size=n_samples)
+        actions, dist = dist.proba_distribution(deterministic_actions, log_std, state)
+
+    entropy = dist.entropy()
+    log_prob = dist.log_prob(actions)
+    assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=5e-3)

From c05c9902857e86c869aed53851f8858e5bd16641 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Wed, 18 Dec 2019 16:56:51 +0100
Subject: [PATCH 24/47] Remove norm clipping

---
 torchy_baselines/sac/sac.py | 13 +------------
 1 file changed, 1 insertion(+), 12 deletions(-)

diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index bc24aa3..6169943 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -44,7 +44,6 @@ class SAC(BaseRLModel):
     :param target_entropy: (str or float) target entropy when learning ent_coef (ent_coef = 'auto')
     :param action_noise: (ActionNoise) the action noise type (None by default), this can help
         for hard exploration problem. Cf common.noise for the different action noise type.
-    :param max_grad_norm: (float) The maximum value for the gradient clipping (None by default)
     :param gamma: (float) the discount factor
     :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
         instead of action noise exploration (default: False)
@@ -64,7 +63,7 @@ class SAC(BaseRLModel):
                  learning_starts=100, batch_size=256,
                  tau=0.005, ent_coef='auto', target_update_interval=1,
                  train_freq=1, gradient_steps=1, n_episodes_rollout=-1,
-                 target_entropy='auto', action_noise=None, max_grad_norm=None,
+                 target_entropy='auto', action_noise=None,
                  gamma=0.99, use_sde=False, sde_sample_freq=-1,
                  tensorboard_log=None, create_eval_env=False,
                  policy_kwargs=None, verbose=0, seed=0, device='auto',
@@ -93,7 +92,6 @@ class SAC(BaseRLModel):
         self.action_noise = action_noise
         self.gamma = gamma
         self.ent_coef_optimizer = None
-        self.max_grad_norm = max_grad_norm
 
         if _init_setup_model:
             self._setup_model()
@@ -205,9 +203,6 @@ class SAC(BaseRLModel):
             if ent_coef_loss is not None:
                 self.ent_coef_optimizer.zero_grad()
                 ent_coef_loss.backward()
-                # Clip grad norm
-                if self.max_grad_norm is not None:
-                    th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
                 self.ent_coef_optimizer.step()
 
 
@@ -233,9 +228,6 @@ class SAC(BaseRLModel):
             # Optimize the critic
             self.critic.optimizer.zero_grad()
             critic_loss.backward()
-            # Clip grad norm
-            if self.max_grad_norm is not None:
-                th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
             self.critic.optimizer.step()
 
             # Compute actor loss
@@ -247,9 +239,6 @@ class SAC(BaseRLModel):
             # Optimize the actor
             self.actor.optimizer.zero_grad()
             actor_loss.backward()
-            # Clip grad norm
-            if self.max_grad_norm is not None:
-                th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
             self.actor.optimizer.step()
 
             # Update target networks

From bff0ca0ea86ec911fdee681d95964d964fb051fd Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Thu, 19 Dec 2019 11:59:00 +0100
Subject: [PATCH 25/47] Use HardTanh to relax the constrain

---
 torchy_baselines/sac/policies.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index 610b59e..61080c0 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -49,7 +49,9 @@ class Actor(BaseNetwork):
                                                                             latent_sde_dim=latent_sde_dim,
                                                                             log_std_init=log_std_init)
             # Avoid saturation by limiting the mean of the gaussian to be in [-1, 1]
-            self.mu = nn.Sequential(self.mu, nn.Tanh())
+            # self.mu = nn.Sequential(self.mu, nn.Tanh())
+            # TODO: test with small positive slope to have non zero gradient
+            self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-2.0, max_val=2.0))
         else:
             self.action_dist = SquashedDiagGaussianDistribution(action_dim)
             self.mu = nn.Linear(net_arch[-1], action_dim)

From a5c3418765e371cdfd831711110339dad58c8fa6 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Thu, 19 Dec 2019 15:28:36 +0100
Subject: [PATCH 26/47] Update README (roadmap moved to github)

---
 README.md | 24 +++---------------------
 1 file changed, 3 insertions(+), 21 deletions(-)

diff --git a/README.md b/README.md
index 8ed3ae3..d8f5fa6 100644
--- a/README.md
+++ b/README.md
@@ -14,27 +14,9 @@ PyTorch version of [Stable Baselines](https://github.com/hill-a/stable-baselines
 - SAC
 - TD3
 
+- SDE support for A2C, PPO, SAC and TD3.
+
 
 ## Roadmap
 
-TODO:
-- better predict
-- complete logger
-- Refactor: buffer with numpy array instead of pytorch
-- Refactor: remove duplicated code for evaluation
-- double check the shape of log prob
-- try squashing both mean and output when using SAC + SDE
-- plotting? -> zoo
-
-Later:
-- get_parameters / set_parameters
-- SDE: use [affine transform](https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/Affine)
-  to scale the noise after a tanh transform?
-- Use MultivariateNormal with full covariance matrix?
-- CNN policies + normalization
-- tensorboard support
-- DQN
-- TRPO
-- ACER
-- DDPG
-- HER -> use stable-baselines because does not depends on tf?
+- cf github Roadmap

From 84ebc3d7da04ca1d8a3731ee28f4755260f9fc55 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Thu, 19 Dec 2019 15:28:51 +0100
Subject: [PATCH 27/47] Relax the HardTanh limit

---
 torchy_baselines/sac/policies.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index 61080c0..1cc6fdf 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -51,7 +51,7 @@ class Actor(BaseNetwork):
             # Avoid saturation by limiting the mean of the gaussian to be in [-1, 1]
             # self.mu = nn.Sequential(self.mu, nn.Tanh())
             # TODO: test with small positive slope to have non zero gradient
-            self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-2.0, max_val=2.0))
+            self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-3.0, max_val=3.0))
         else:
             self.action_dist = SquashedDiagGaussianDistribution(action_dim)
             self.mu = nn.Linear(net_arch[-1], action_dim)

From 69428346bdb2d271ca1610cbb601092b52026ebd Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Thu, 19 Dec 2019 15:32:03 +0100
Subject: [PATCH 28/47] Bump version

---
 setup.py                     | 2 +-
 torchy_baselines/__init__.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 78a72bf..71242ee 100644
--- a/setup.py
+++ b/setup.py
@@ -34,7 +34,7 @@ setup(name='torchy_baselines',
       license="MIT",
       long_description="",
       long_description_content_type='text/markdown',
-      version="0.0.7a",
+      version="0.0.8a0",
       )
 
 # python setup.py sdist
diff --git a/torchy_baselines/__init__.py b/torchy_baselines/__init__.py
index 56d1dfe..3f3ec0d 100644
--- a/torchy_baselines/__init__.py
+++ b/torchy_baselines/__init__.py
@@ -4,4 +4,4 @@ from torchy_baselines.ppo import PPO
 from torchy_baselines.sac import SAC
 from torchy_baselines.td3 import TD3
 
-__version__ = "0.0.7a"
+__version__ = "0.0.8a0"

From e894f1f11b3b7770722c7684747de3fbfd68c6b0 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Thu, 19 Dec 2019 18:20:02 +0100
Subject: [PATCH 29/47] Add leakyclip

---
 torchy_baselines/sac/policies.py | 27 +++++++++++++++++++++++++--
 1 file changed, 25 insertions(+), 2 deletions(-)

diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index 1cc6fdf..7ede952 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -10,6 +10,28 @@ LOG_STD_MAX = 2
 LOG_STD_MIN = -20
 
 
+class LeakyClip(nn.Module):
+    """
+    Cip values outside a certain range
+    (it is not a hard clip, there is a small slope to have non-zero gradient)
+
+    :param min_val: (float)
+    :param max_val: (float)
+    :param slope: (float)
+    """
+    def __init__(self, min_val=-2.0, max_val=2.0, slope=0.01):
+        super(LeakyClip, self).__init__()
+        self.min_val = min_val
+        self.max_val = max_val
+        self.slope = slope
+
+    def forward(self, x):
+        linear_part = x * (x >= self.min_val) * (x <= self.max_val)
+        above_max_val = self.slope * (x - self.max_val) * (x > self.max_val)
+        below_min_val = self.slope * (x - self.min_val) * (x < self.min_val)
+        return linear_part + below_min_val + above_max_val
+
+
 class Actor(BaseNetwork):
     """
     Actor network (policy) for SAC.
@@ -50,8 +72,9 @@ class Actor(BaseNetwork):
                                                                             log_std_init=log_std_init)
             # Avoid saturation by limiting the mean of the gaussian to be in [-1, 1]
             # self.mu = nn.Sequential(self.mu, nn.Tanh())
-            # TODO: test with small positive slope to have non zero gradient
-            self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-3.0, max_val=3.0))
+            self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-2.0, max_val=2.0))
+            # Small positive slope to have non-zero gradient
+            self.mu = nn.Sequential(self.mu, LeakyClip())
         else:
             self.action_dist = SquashedDiagGaussianDistribution(action_dim)
             self.mu = nn.Linear(net_arch[-1], action_dim)

From 161c608f9cebaebb1d2e7118a62b032347bb1fc7 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Fri, 20 Dec 2019 11:28:20 +0100
Subject: [PATCH 30/47] Re-sample noise matrix for PPO

---
 torchy_baselines/ppo/ppo.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index bc2c7d8..66f459d 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -56,7 +56,7 @@ class PPO(BaseRLModel):
     :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
         instead of action noise exploration (default: False)
     :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
-        Default: -1 (only sample at the beginning of the rollout)    
+        Default: -1 (only sample at the beginning of the rollout)
     :param target_kl: (float) Limit the KL divergence between updates,
         because the clipping is not enough to prevent large update
         see issue #213 (cf https://github.com/hill-a/stable-baselines/issues/213)
@@ -210,6 +210,12 @@ class PPO(BaseRLModel):
                     # Convert discrete action for float to long
                     action = action.long().flatten()
 
+                # Re-sample the noise matrix because the log_std has changed
+                # TODO: investigate why there is no issue with the gradient
+                # if that line is commented (as in SAC)
+                if self.use_sde:
+                    self.policy.reset_noise(batch_size)
+
                 values, log_prob, entropy = self.policy.evaluate_actions(obs, action)
                 values = values.flatten()
                 # Normalize advantage

From 9b3b34c9c4b4b6fafc52ec3b4e4fcb4573a5df45 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Fri, 20 Dec 2019 11:28:44 +0100
Subject: [PATCH 31/47] Sample `batch_size` noise matrices for SAC

---
 torchy_baselines/sac/sac.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index 6169943..cd54497 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -181,8 +181,8 @@ class SAC(BaseRLModel):
             # is lost and we cannot backpropagate through again
             # anyway, we need to sample because `log_std` may have changed between two gradient steps
             if self.use_sde:
-                # self.actor.reset_noise(batch_size=batch_size)
-                self.actor.reset_noise()
+                self.actor.reset_noise(batch_size=batch_size)
+                # self.actor.reset_noise()
 
             # Action by the current actor for the sampled state
             action_pi, log_prob = self.actor.action_log_prob(obs)

From 7f34108ed6fddc63f037831b07db830f6ae8197e Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Fri, 20 Dec 2019 18:02:01 +0100
Subject: [PATCH 32/47] Fix `exp_ln` computation

---
 torchy_baselines/common/distributions.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py
index 97eaf6f..ba37cd3 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -272,10 +272,11 @@ class StateDependentNoiseDistribution(Distribution):
         if self.use_expln:
             # From SDE paper, it allows to keep variance
             # above zero and prevent it from growing too fast
-            if log_std <= 0:
-                std = th.exp(log_std)
-            else:
-                std = th.log(log_std + 1.0) + 1.0
+            below_threshold = th.exp(log_std) * (log_std <= 0)
+            # Avoid NaN: zeros values that are below zero
+            safe_log_std = log_std * (log_std > 0) + self.epsilon
+            above_threshold = (th.log1p(safe_log_std) + 1.0) * (log_std > 0)
+            std = below_threshold + above_threshold
         else:
             # Use normal exponential
             std = th.exp(log_std)

From f6c475a44b10500d63cfc77558c21de2aeb053ab Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Fri, 20 Dec 2019 18:10:24 +0100
Subject: [PATCH 33/47] Add `use_expln` as a policy argument

---
 torchy_baselines/sac/policies.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index 7ede952..58d5a99 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -47,9 +47,13 @@ class Actor(BaseNetwork):
     :param sde_net_arch: ([int]) Network architecture for extracting features
         when using SDE. If None, the latent features from the policy will be used.
         Pass an empty list to use the states as features.
+    :param use_expln: (bool) Use `expln()` function instead of `exp()` when using SDE to ensure
+        a positive standard deviation (cf paper). It allows to keep variance
+        above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
     """
     def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
-                 use_sde=False, log_std_init=-3, full_std=True, sde_net_arch=None):
+                 use_sde=False, log_std_init=-3, full_std=True,
+                 sde_net_arch=None, use_expln=False):
         super(Actor, self).__init__()
 
         latent_pi_net = create_mlp(obs_dim, -1, net_arch, activation_fn)
@@ -65,7 +69,7 @@ class Actor(BaseNetwork):
                                                                                           activation_fn)
 
             # TODO: check for the learn_features
-            self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
+            self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=use_expln,
                                                                learn_features=True, squash_output=True)
             self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
                                                                             latent_sde_dim=latent_sde_dim,
@@ -186,11 +190,14 @@ class SACPolicy(BasePolicy):
     :param sde_net_arch: ([int]) Network architecture for extracting features
         when using SDE. If None, the latent features from the policy will be used.
         Pass an empty list to use the states as features.
+    :param use_expln: (bool) Use `expln()` function instead of `exp()` when using SDE to ensure
+        a positive standard deviation (cf paper). It allows to keep variance
+        above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
     """
     def __init__(self, observation_space, action_space,
                  learning_rate, net_arch=None, device='cpu',
                  activation_fn=nn.ReLU, use_sde=False,
-                 log_std_init=-3, sde_net_arch=None):
+                 log_std_init=-3, sde_net_arch=None, use_expln=False):
         super(SACPolicy, self).__init__(observation_space, action_space, device)
 
         if net_arch is None:
@@ -210,7 +217,8 @@ class SACPolicy(BasePolicy):
         sde_kwargs = {
             'use_sde': use_sde,
             'log_std_init': log_std_init,
-            'sde_net_arch': sde_net_arch
+            'sde_net_arch': sde_net_arch,
+            'use_expln': use_expln
         }
         self.actor_kwargs.update(sde_kwargs)
         self.actor, self.actor_target = None, None

From 3a7508ac16794b63c28794096553b7fdb277ad24 Mon Sep 17 00:00:00 2001
From: Antonin RAFFIN <antonin.raffin@ensta.org>
Date: Sun, 22 Dec 2019 13:56:30 +0100
Subject: [PATCH 34/47] Fix double clip

---
 torchy_baselines/sac/policies.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index 58d5a99..9de90e6 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -76,7 +76,7 @@ class Actor(BaseNetwork):
                                                                             log_std_init=log_std_init)
             # Avoid saturation by limiting the mean of the gaussian to be in [-1, 1]
             # self.mu = nn.Sequential(self.mu, nn.Tanh())
-            self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-2.0, max_val=2.0))
+            # self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-2.0, max_val=2.0))
             # Small positive slope to have non-zero gradient
             self.mu = nn.Sequential(self.mu, LeakyClip())
         else:

From 57c890f3e9f1b16f2d0e83d813ebaff79d58c018 Mon Sep 17 00:00:00 2001
From: Antonin RAFFIN <antonin.raffin@ensta.org>
Date: Sun, 22 Dec 2019 14:38:30 +0100
Subject: [PATCH 35/47] LeakyClip not working yet

---
 torchy_baselines/sac/policies.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index 9de90e6..bc5fbcf 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -76,9 +76,9 @@ class Actor(BaseNetwork):
                                                                             log_std_init=log_std_init)
             # Avoid saturation by limiting the mean of the gaussian to be in [-1, 1]
             # self.mu = nn.Sequential(self.mu, nn.Tanh())
-            # self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-2.0, max_val=2.0))
+            self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-2.0, max_val=2.0))
             # Small positive slope to have non-zero gradient
-            self.mu = nn.Sequential(self.mu, LeakyClip())
+            # self.mu = nn.Sequential(self.mu, LeakyClip())
         else:
             self.action_dist = SquashedDiagGaussianDistribution(action_dim)
             self.mu = nn.Linear(net_arch[-1], action_dim)

From 4a79f7e5a76bcbdef7fd43dc768882742cb1fb4c Mon Sep 17 00:00:00 2001
From: Antonin RAFFIN <antonin.raffin@ensta.org>
Date: Tue, 24 Dec 2019 13:12:04 +0100
Subject: [PATCH 36/47] Print std reward for evaluation

---
 torchy_baselines/common/evaluation.py | 2 +-
 torchy_baselines/sac/sac.py           | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/torchy_baselines/common/evaluation.py b/torchy_baselines/common/evaluation.py
index 3d76d9f..bc8fc70 100644
--- a/torchy_baselines/common/evaluation.py
+++ b/torchy_baselines/common/evaluation.py
@@ -19,4 +19,4 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, render=F
                 env.render()
         episode_rewards.append(episode_reward)
 
-    return np.mean(episode_rewards), n_steps
+    return np.mean(episode_rewards), np.std(episode_rewards)
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index cd54497..e4f24d2 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -291,10 +291,10 @@ class SAC(BaseRLModel):
             if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
                 timesteps_since_eval %= eval_freq
                 sync_envs_normalization(self.env, eval_env)
-                mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
+                mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes)
                 evaluations.append(mean_reward)
                 if self.verbose > 0:
-                    print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
+                    print("Eval num_timesteps={}, mean_reward={:.2f}, std_reward={:.2f}".format(self.num_timesteps, mean_reward, std_reward))
                     print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
 
         return self

From aa7b91333e1970fa84346fd1c690427a30eab03b Mon Sep 17 00:00:00 2001
From: Antonin RAFFIN <antonin.raffin@ensta.org>
Date: Mon, 30 Dec 2019 12:01:37 +0100
Subject: [PATCH 37/47] Add seeding for subproc vecenv

---
 torchy_baselines/common/vec_env/subproc_vec_env.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/torchy_baselines/common/vec_env/subproc_vec_env.py b/torchy_baselines/common/vec_env/subproc_vec_env.py
index b00e465..f89cbc6 100644
--- a/torchy_baselines/common/vec_env/subproc_vec_env.py
+++ b/torchy_baselines/common/vec_env/subproc_vec_env.py
@@ -188,6 +188,17 @@ class SubprocVecEnv(VecEnv):
         for remote in target_remotes:
             remote.recv()
 
+    def seed(self, seed, indices=None):
+        """
+        :param seed: (int or [int])
+        :param indices: ([int])
+        """
+        indices = self._get_indices(indices)
+        if not hasattr(seed, 'len'):
+            seed = [seed] * len(indices)
+        assert len(seed) == len(indices)
+        return [self.env_method('seed', seed[i], indices=i) for i in indices]
+
     def env_method(self, method_name, *method_args, **method_kwargs):
         """Call instance methods of vectorized environments."""
         indices = method_kwargs.get('indices')

From 8831eff163bc061727e16b7f728162776dc5ba27 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Tue, 7 Jan 2020 14:00:03 +0100
Subject: [PATCH 38/47] Unify evaluation

---
 torchy_baselines/cem_rl/cem_rl.py     |  5 ++--
 torchy_baselines/common/base_class.py | 34 ++++++++++++++++++----
 torchy_baselines/common/evaluation.py | 42 +++++++++++++++++++++++----
 torchy_baselines/ppo/ppo.py           | 21 ++++----------
 torchy_baselines/sac/sac.py           | 13 ++-------
 torchy_baselines/td3/td3.py           | 15 ++--------
 6 files changed, 80 insertions(+), 50 deletions(-)

diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py
index 41efdf3..e3a322b 100644
--- a/torchy_baselines/cem_rl/cem_rl.py
+++ b/torchy_baselines/cem_rl/cem_rl.py
@@ -155,11 +155,12 @@ class CEMRL(TD3):
                 self.actor.load_from_vector(self.es.mu)
                 sync_envs_normalization(self.env, eval_env)
 
-                mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
+                mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes)
                 evaluations.append(mean_reward)
 
                 if self.verbose > 0:
-                    print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
+                    print("Eval num_timesteps={}, "
+                          "episode_reward={:.2f} +/- {:.2f}".format(self.num_timesteps, mean_reward, std_reward))
                     print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
 
             actor_steps = 0
diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index 6200bc1..041ce83 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -9,11 +9,12 @@ import gym
 import torch as th
 import numpy as np
 
+from torchy_baselines.common import logger
 from torchy_baselines.common.policies import get_policy_from_name
 from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate
-from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize
+from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, sync_envs_normalization
 from torchy_baselines.common.monitor import Monitor
-from torchy_baselines.common import logger
+from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.common.save_util import data_to_json, json_to_data
 
 
@@ -287,9 +288,9 @@ class BaseRLModel(object):
         :param log_interval: (int) The number of timesteps before logging.
         :param tb_log_name: (str) the name of the run for tensorboard log
         :param reset_num_timesteps: (bool) whether or not to reset the current timestep number (used in logging)
-        :param eval_env: (gym.Env)
-        :param eval_freq: (int)
-        :param n_eval_episodes: (int)
+        :param eval_env: (gym.Env) Environment that will be used to evaluate the agent
+        :param eval_freq: (int) Evaluate the agent every `eval_freq` timesteps (this may vary a little)
+        :param n_eval_episodes: (int) Number of episode to evaluate the agent
         :return: (BaseRLModel) the trained model
         """
         pass
@@ -712,3 +713,26 @@ class BaseRLModel(object):
         params_to_save = self.get_policy_parameters()
         opt_params_to_save = self.get_opt_parameters()
         self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save)
+
+    def _eval_policy(self, eval_freq, eval_env, n_eval_episodes,
+                     timesteps_since_eval, deterministic=True):
+        """
+        Evaluate the current policy on a test environment.
+
+        :param eval_env: (gym.Env) Environment that will be used to evaluate the agent
+        :param eval_freq: (int) Evaluate the agent every `eval_freq` timesteps (this may vary a little)
+        :param n_eval_episodes: (int) Number of episode to evaluate the agent
+        :parma timesteps_since_eval: (int) Number of timesteps since last evaluation
+        :param deterministic: (bool) Whether to use deterministic or stochastic actions
+        :return: (int) Number of timesteps since last evaluation
+        """
+        if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
+            timesteps_since_eval %= eval_freq
+            # Synchronise the normalization stats if needed
+            sync_envs_normalization(self.env, eval_env)
+            mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes, deterministic=deterministic)
+            if self.verbose > 0:
+                print("Eval num_timesteps={}, "
+                      "episode_reward={:.2f} +/- {:.2f}".format(self.num_timesteps, mean_reward, std_reward))
+                print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
+        return timesteps_since_eval
diff --git a/torchy_baselines/common/evaluation.py b/torchy_baselines/common/evaluation.py
index bc8fc70..04ae296 100644
--- a/torchy_baselines/common/evaluation.py
+++ b/torchy_baselines/common/evaluation.py
@@ -1,10 +1,34 @@
+# Copied from stable_baselines
 import numpy as np
 
+from torchy_baselines.common.vec_env import VecEnv
 
-def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, render=False):
+
+def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True,
+                    render=False, callback=None, reward_threshold=None,
+                    return_episode_rewards=False):
     """
-    Runs policy for n episodes and returns average reward
+    Runs policy for `n_eval_episodes` episodes and returns average reward.
+    This is made to work only with one env.
+
+    :param model: (BaseRLModel) The RL agent you want to evaluate.
+    :param env: (gym.Env or VecEnv) The gym environment. In the case of a `VecEnv`
+        this must contain only one environment.
+    :param n_eval_episodes: (int) Number of episode to evaluate the agent
+    :param deterministic: (bool) Whether to use deterministic or stochastic actions
+    :param render: (bool) Whether to render the environment or not
+    :param callback: (callable) callback function to do additional checks,
+        called after each step.
+    :param reward_threshold: (float) Minimum expected reward per episode,
+        this will raise an error if the performance is not met
+    :param return_episode_rewards: (bool) If True, a list of reward per episode
+        will be returned instead of the mean.
+    :return: (float, float) Mean reward per episode, std of reward per episode
+        returns ([float], int) when `return_episode_rewards` is True
     """
+    if isinstance(env, VecEnv):
+        assert env.num_envs == 1, "You must pass only one environment when using this function"
+
     episode_rewards, n_steps = [], 0
     for _ in range(n_eval_episodes):
         obs = env.reset()
@@ -12,11 +36,19 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, render=F
         episode_reward = 0.0
         while not done:
             action = model.predict(obs, deterministic=deterministic)
-            obs, reward, done, _ = env.step(action)
+            obs, reward, done, _info = env.step(action)
             episode_reward += reward
+            if callback is not None:
+                callback(locals(), globals())
             n_steps += 1
             if render:
                 env.render()
         episode_rewards.append(episode_reward)
-
-    return np.mean(episode_rewards), np.std(episode_rewards)
+    mean_reward = np.mean(episode_rewards)
+    std_reward = np.std(episode_rewards)
+    if reward_threshold is not None:
+        assert mean_reward > reward_threshold, 'Mean reward below threshold: '\
+                                         '{:.2f} < {:.2f}'.format(mean_reward, reward_threshold)
+    if return_episode_rewards:
+        return episode_rewards, n_steps
+    return mean_reward, std_reward
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index 66f459d..8641552 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -14,10 +14,8 @@ except ImportError:
 import numpy as np
 
 from torchy_baselines.common.base_class import BaseRLModel
-from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.common.buffers import RolloutBuffer
 from torchy_baselines.common.utils import explained_variance, get_schedule_fn
-from torchy_baselines.common.vec_env import sync_envs_normalization
 from torchy_baselines.common import logger
 from torchy_baselines.ppo.policies import PPOPolicy
 
@@ -303,19 +301,12 @@ class PPO(BaseRLModel):
 
             self.train(self.n_epochs, batch_size=self.batch_size)
 
-            # Evaluate agent
-            if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
-                timesteps_since_eval %= eval_freq
-                sync_envs_normalization(self.env, eval_env)
-
-                mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
-                if self.tb_writer is not None:
-                    self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps)
-
-                evaluations.append(mean_reward)
-                if self.verbose > 0:
-                    print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
-                    print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
+            # Evaluate the agent
+            timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
+                                                     timesteps_since_eval, deterministic=True)
+            # For tensorboard integration
+            # if self.tb_writer is not None:
+            #     self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps)
 
         return self
 
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index e4f24d2..f3f45fb 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -6,9 +6,7 @@ import numpy as np
 
 from torchy_baselines.common.base_class import BaseRLModel
 from torchy_baselines.common.buffers import ReplayBuffer
-from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.sac.policies import SACPolicy
-from torchy_baselines.common.vec_env import sync_envs_normalization
 from torchy_baselines.common import logger
 
 
@@ -287,15 +285,8 @@ class SAC(BaseRLModel):
 
                 self.train(gradient_steps, batch_size=self.batch_size)
 
-            # Evaluate episode
-            if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
-                timesteps_since_eval %= eval_freq
-                sync_envs_normalization(self.env, eval_env)
-                mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes)
-                evaluations.append(mean_reward)
-                if self.verbose > 0:
-                    print("Eval num_timesteps={}, mean_reward={:.2f}, std_reward={:.2f}".format(self.num_timesteps, mean_reward, std_reward))
-                    print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
+            timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
+                                                     timesteps_since_eval, deterministic=True)
 
         return self
 
diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py
index 9b5cafd..a24e0db 100644
--- a/torchy_baselines/td3/td3.py
+++ b/torchy_baselines/td3/td3.py
@@ -6,10 +6,7 @@ import numpy as np
 
 from torchy_baselines.common.base_class import BaseRLModel
 from torchy_baselines.common.buffers import ReplayBuffer
-from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.td3.policies import TD3Policy
-from torchy_baselines.common.vec_env import sync_envs_normalization
-from torchy_baselines.ppo.policies import MlpPolicy
 
 
 class TD3(BaseRLModel):
@@ -290,15 +287,9 @@ class TD3(BaseRLModel):
                 gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps
                 self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay)
 
-            # Evaluate episode
-            if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
-                timesteps_since_eval %= eval_freq
-                sync_envs_normalization(self.env, eval_env)
-                mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
-                evaluations.append(mean_reward)
-                if self.verbose > 0:
-                    print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
-                    print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
+            # Evaluate the agent
+            timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
+                                                     timesteps_since_eval, deterministic=True)
 
         return self
 

From 299ca007b55ccf3a20a7476afba61e5441bfcbab Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Tue, 7 Jan 2020 17:36:26 +0100
Subject: [PATCH 39/47] Add comment about warmup phase

---
 torchy_baselines/common/base_class.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index 041ce83..f517c58 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -526,6 +526,8 @@ class BaseRLModel(object):
                     self.actor.reset_noise()
 
                 # Select action randomly or according to policy
+                # TODO: use action from policy when using SDE during the warmup phase?
+                # if num_timesteps < learning_starts and not self.use_sde:
                 if num_timesteps < learning_starts:
                     # Warmup phase
                     unscaled_action = np.array([self.action_space.sample()])

From d3a718b94e771fa154d52d5e6e71e02d5fd1f8b3 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Wed, 8 Jan 2020 11:26:57 +0100
Subject: [PATCH 40/47] Add extra dependency

---
 setup.py                             | 3 +++
 torchy_baselines/common/save_util.py | 2 +-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 71242ee..e5e411f 100644
--- a/setup.py
+++ b/setup.py
@@ -23,6 +23,9 @@ setup(name='torchy_baselines',
             'sphinx',
             'sphinx-autobuild',
             'sphinx-rtd-theme'
+        ],
+        'render': [
+            'opencv-python'
         ]
       },
       description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.',
diff --git a/torchy_baselines/common/save_util.py b/torchy_baselines/common/save_util.py
index 30dd0b2..7b6a2be 100644
--- a/torchy_baselines/common/save_util.py
+++ b/torchy_baselines/common/save_util.py
@@ -131,4 +131,4 @@ def json_to_data(json_string, custom_objects=None):
         else:
             # Read as it is
             return_data[data_key] = data_item
-    return return_data
\ No newline at end of file
+    return return_data

From 60d5f4463d3fe27c08a12872919cf286091908a6 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Wed, 8 Jan 2020 17:04:28 +0100
Subject: [PATCH 41/47] Add `use_expln` option for td3

---
 torchy_baselines/td3/policies.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py
index 1fd6a20..f016dab 100644
--- a/torchy_baselines/td3/policies.py
+++ b/torchy_baselines/td3/policies.py
@@ -23,10 +23,13 @@ class Actor(BaseNetwork):
     :param sde_net_arch: ([int]) Network architecture for extracting features
         when using SDE. If None, the latent features from the policy will be used.
         Pass an empty list to use the states as features.
+    :param use_expln: (bool) Use `expln()` function instead of `exp()` when using SDE to ensure
+        a positive standard deviation (cf paper). It allows to keep variance
+        above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
     """
     def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
                  use_sde=False, log_std_init=-3, clip_noise=None,
-                 lr_sde=3e-4, full_std=False, sde_net_arch=None):
+                 lr_sde=3e-4, full_std=False, sde_net_arch=None, use_expln=False):
         super(Actor, self).__init__()
 
         self.latent_pi, self.log_std = None, None
@@ -48,7 +51,7 @@ class Actor(BaseNetwork):
                                                                                           activation_fn)
 
             # Create state dependent noise matrix (SDE)
-            self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
+            self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=use_expln,
                                                                squash_output=False, learn_features=learn_features)
             action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
                                                                                latent_sde_dim=latent_sde_dim,
@@ -194,11 +197,14 @@ class TD3Policy(BasePolicy):
     :param sde_net_arch: ([int]) Network architecture for extracting features
         when using SDE. If None, the latent features from the policy will be used.
         Pass an empty list to use the states as features.
+    :param use_expln: (bool) Use `expln()` function instead of `exp()` when using SDE to ensure
+        a positive standard deviation (cf paper). It allows to keep variance
+        above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
     """
     def __init__(self, observation_space, action_space,
                  learning_rate, net_arch=None, device='cpu',
                  activation_fn=nn.ReLU, use_sde=False, log_std_init=-3,
-                 clip_noise=None, lr_sde=3e-4, sde_net_arch=None):
+                 clip_noise=None, lr_sde=3e-4, sde_net_arch=None, use_expln=False):
         super(TD3Policy, self).__init__(observation_space, action_space, device)
 
         # Default network architecture, from the original paper
@@ -221,7 +227,8 @@ class TD3Policy(BasePolicy):
             'log_std_init': log_std_init,
             'clip_noise': clip_noise,
             'lr_sde': lr_sde,
-            'sde_net_arch': sde_net_arch
+            'sde_net_arch': sde_net_arch,
+            'use_expln': use_expln
         }
         self.actor_kwargs.update(sde_kwargs)
 

From 03e853997a5c78cab38ca69989a5ee73c3e2ed58 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Wed, 15 Jan 2020 13:21:20 +0100
Subject: [PATCH 42/47] Add squash_output and expln as policy param for ppo and
 a2c

---
 torchy_baselines/ppo/policies.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index 8c4df3f..d3b897e 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -29,12 +29,18 @@ class PPOPolicy(BasePolicy):
     :param sde_net_arch: ([int]) Network architecture for extracting features
         when using SDE. If None, the latent features from the policy will be used.
         Pass an empty list to use the states as features.
+    :param use_expln: (bool) Use `expln()` function instead of `exp()` to ensure
+        a positive standard deviation (cf paper). It allows to keep variance
+        above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
+    :param squash_output: (bool) Whether to squash the output using a tanh function,
+        this allows to ensure boundaries when using SDE.
     """
     def __init__(self, observation_space, action_space,
                  learning_rate, net_arch=None, device='cpu',
                  activation_fn=nn.Tanh, adam_epsilon=1e-5,
                  ortho_init=True, use_sde=False,
-                 log_std_init=0.0, full_std=True, sde_net_arch=None):
+                 log_std_init=0.0, full_std=True,
+                 sde_net_arch=None, use_expln=False, squash_output=False):
         super(PPOPolicy, self).__init__(observation_space, action_space, device)
         self.obs_dim = self.observation_space.shape[0]
 
@@ -63,8 +69,8 @@ class PPOPolicy(BasePolicy):
         if use_sde:
             dist_kwargs = {
                 'full_std': full_std,
-                'squash_output': False,
-                'use_expln': False,
+                'squash_output': squash_output,
+                'use_expln': use_expln,
                 'learn_features': sde_net_arch is not None
             }
 

From ea20721632477a7f9da5e9ca65827391ddc06522 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Wed, 15 Jan 2020 15:58:45 +0100
Subject: [PATCH 43/47] Add TODO

---
 torchy_baselines/common/base_class.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index f517c58..4c7a764 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -596,6 +596,7 @@ class BaseRLModel(object):
                 total_episodes += 1
                 episode_rewards.append(episode_reward)
                 total_timesteps.append(episode_timesteps)
+                # TODO: reset SDE matrix at the end of the episode?
                 if action_noise is not None:
                     action_noise.reset()
 

From c5420096414dc2c1d63f0bf58c052e6f99f5317b Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 20 Jan 2020 11:17:55 +0100
Subject: [PATCH 44/47] Clean up code + bug fixes

---
 tests/test_custom_policy.py                   |  5 +--
 tests/test_distributions.py                   |  6 +--
 tests/test_save_load.py                       | 10 ++---
 tests/test_sde.py                             |  8 +---
 tests/test_vec_envs.py                        |  7 ++--
 torchy_baselines/common/base_class.py         | 41 +++++++------------
 torchy_baselines/common/buffers.py            |  5 +--
 torchy_baselines/common/distributions.py      |  3 +-
 torchy_baselines/common/logger.py             | 14 +------
 torchy_baselines/common/vec_env/__init__.py   |  2 +-
 .../common/vec_env/vec_normalize.py           |  2 +-
 torchy_baselines/ppo/policies.py              |  3 +-
 torchy_baselines/sac/policies.py              |  3 ++
 torchy_baselines/sac/sac.py                   |  9 ++--
 torchy_baselines/td3/td3.py                   | 11 ++---
 15 files changed, 51 insertions(+), 78 deletions(-)

diff --git a/tests/test_custom_policy.py b/tests/test_custom_policy.py
index d45be88..2555016 100644
--- a/tests/test_custom_policy.py
+++ b/tests/test_custom_policy.py
@@ -1,6 +1,3 @@
-import os
-
-import gym
 import pytest
 
 from torchy_baselines import PPO
@@ -15,4 +12,4 @@ from torchy_baselines import PPO
     [12, dict(pi=[8])],
 ])
 def test_flexible_mlp(net_arch):
-    model = PPO('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000)
+    _ = PPO('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000)
diff --git a/tests/test_distributions.py b/tests/test_distributions.py
index 3d896ab..7e24cda 100644
--- a/tests/test_distributions.py
+++ b/tests/test_distributions.py
@@ -1,10 +1,9 @@
 import pytest
-import numpy as np
 import torch as th
 
+from torchy_baselines.common.distributions import DiagGaussianDistribution, TanhBijector, \
+    StateDependentNoiseDistribution
 from torchy_baselines.common.utils import set_random_seed
-from torchy_baselines.common.distributions import DiagGaussianDistribution, SquashedDiagGaussianDistribution,\
-    CategoricalDistribution, TanhBijector, StateDependentNoiseDistribution
 
 
 # TODO: more tests for the other distributions
@@ -43,6 +42,7 @@ def test_sde_distribution():
 
 N_ACTIONS = 1
 
+
 # TODO: fix for num action > 1
 # TODO: analytical form for squashed Gaussian?
 @pytest.mark.parametrize("dist", [
diff --git a/tests/test_save_load.py b/tests/test_save_load.py
index 3a2cb42..f4adc0c 100644
--- a/tests/test_save_load.py
+++ b/tests/test_save_load.py
@@ -1,13 +1,12 @@
+import numpy as np
 import os
 import pytest
-from copy import deepcopy
-import numpy as np
-
 import torch as th
+from copy import deepcopy
 
 from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
+from torchy_baselines.common.identity_env import IdentityEnvBox
 from torchy_baselines.common.vec_env import DummyVecEnv
-from torchy_baselines.common.identity_env import IdentityEnvBox, IdentityEnv
 
 MODEL_LIST = [
     CEMRL,
@@ -81,7 +80,8 @@ def test_save_load(model_class):
     for optimizer, opt_state in opt_params.items():
         for param_group_idx, param_group in enumerate(opt_state['param_groups']):
             for param_key, param_value in param_group.items():
-                if param_key == 'params':  # don't know how to handle params correctly, therefore only check if we have the same amount
+                # don't know how to handle params correctly, therefore only check if we have the same amount
+                if param_key == 'params':
                     assert len(param_value) == len(
                         new_opt_params[optimizer]['param_groups'][param_group_idx][param_key])
                 else:
diff --git a/tests/test_sde.py b/tests/test_sde.py
index 851ca42..497f398 100644
--- a/tests/test_sde.py
+++ b/tests/test_sde.py
@@ -1,12 +1,8 @@
 import pytest
-
-import gym
 import torch as th
 from torch.distributions import Normal
 
 from torchy_baselines import A2C, TD3, SAC
-from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize
-from torchy_baselines.common.monitor import Monitor
 
 
 def test_state_dependent_exploration_grad():
@@ -35,7 +31,7 @@ def test_state_dependent_exploration_grad():
     action_dist = Normal(mu, th.sqrt(variance))
 
     # Sum over the action dimension because we assume they are independent
-    loss = action_dist.log_prob((action).detach()).sum(dim=-1).mean()
+    loss = action_dist.log_prob(action.detach()).sum(dim=-1).mean()
     loss.backward()
 
     # From Rueckstiess paper: check that the computed gradient
@@ -72,6 +68,6 @@ def test_scheduler():
         return -2.0 * progress + 1
 
     model = TD3('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True,
-                        verbose=1, sde_log_std_scheduler=scheduler)
+                verbose=1, sde_log_std_scheduler=scheduler)
     model.learn(total_timesteps=int(1000), eval_freq=500)
     assert th.isclose(model.actor.log_std, th.ones_like(model.actor.log_std)).all()
diff --git a/tests/test_vec_envs.py b/tests/test_vec_envs.py
index efa5119..f2dd1c2 100644
--- a/tests/test_vec_envs.py
+++ b/tests/test_vec_envs.py
@@ -325,9 +325,8 @@ def test_vecenv_wrapper_getattr():
     assert wrapped.name_test() == CustomWrapperBB
 
     double_wrapped = CustomWrapperA(CustomWrapperB(wrapped))
-    dummy = double_wrapped.var_a  # should not raise as it is directly defined here
+    _ = double_wrapped.var_a  # should not raise as it is directly defined here
     with pytest.raises(AttributeError):  # should raise due to ambiguity
-        dummy = double_wrapped.var_b
+        _ = double_wrapped.var_b
     with pytest.raises(AttributeError):  # should raise as does not exist
-        dummy = double_wrapped.nonexistent_attribute
-    del dummy  # keep linter happy
+        _ = double_wrapped.nonexistent_attribute
diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index 4c7a764..0f20211 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -225,7 +225,8 @@ class BaseRLModel(object):
         :param env: (gym.Env) The environment for learning a policy
         """
         if self.check_env(env, self.observation_space, self.action_space) is False:
-            raise ValueError("The given environment is not compatible with model: observation and action spaces do not match")
+            raise ValueError("The given environment is not compatible with model: "
+                             "observation and action spaces do not match")
         # it must be coherent now
         # if it is not a VecEnv, make it a VecEnv
         if not isinstance(env, VecEnv):
@@ -258,24 +259,6 @@ class BaseRLModel(object):
         """
         raise NotImplementedError()
 
-    def pretrain(self, dataset, n_epochs=10, learning_rate=1e-4,
-                 adam_epsilon=1e-8, val_interval=None):
-        """
-        Pretrain a model using behavior cloning:
-        supervised learning given an expert dataset.
-
-        NOTE: only Box and Discrete spaces are supported for now.
-
-        :param dataset: (ExpertDataset) Dataset manager
-        :param n_epochs: (int) Number of iterations on the training set
-        :param learning_rate: (float) Learning rate
-        :param adam_epsilon: (float) the epsilon value for the adam optimizer
-        :param val_interval: (int) Report training and validation losses every n epochs.
-            By default, every 10th of the maximum number of epochs.
-        :return: (BaseRLModel) the pretrained model
-        """
-        raise NotImplementedError()
-
     @abstractmethod
     def learn(self, total_timesteps, callback=None, log_interval=100, tb_log_name="run",
               eval_env=None, eval_freq=-1, n_eval_episodes=5, reset_num_timesteps=True):
@@ -308,11 +291,12 @@ class BaseRLModel(object):
         """
         pass
 
-    def load_parameters(self, load_dict, opt_params=None):
+    def load_parameters(self, load_dict, opt_params):
         """
         Load model parameters from a dictionary
         load_dict should contain all keys from torch.model.state_dict()
-        If opt_params are given this does also load agent's optimizer-parameters, but can only be handled in child classes.
+        If opt_params are given this does also load agent's optimizer-parameters,
+        but can only be handled in child classes.
 
 
         :param load_dict: (dict) dict of parameters from model.state_dict()
@@ -350,6 +334,7 @@ class BaseRLModel(object):
             env = data["env"]
 
         # first create model, but only setup if a env was given
+        # noinspection PyArgumentList
         model = cls(policy=data["policy_class"], env=env, _init_setup_model=env is not None)
 
         # load parameters
@@ -365,7 +350,8 @@ class BaseRLModel(object):
         :param load_path: (str) Where to load the model from
         :param load_data: (bool) Whether we should load and return data
             (class parameters). Mainly used by 'load_parameters' to only load model parameters (weights)
-        :return: (dict),(dict),(dict) Class parameters, model parameters (state_dict) and dict of optimizer parameters (dict of state_dict)
+        :return: (dict),(dict),(dict) Class parameters, model parameters (state_dict)
+            and dict of optimizer parameters (dict of state_dict)
         """
         # Check if file exists if load_path is a string
         if isinstance(load_path, str):
@@ -403,7 +389,7 @@ class BaseRLModel(object):
 
                 # check for all other .pth files
                 other_files = [file_name for file_name in namelist if
-                              os.path.splitext(file_name)[1] == ".pth" and file_name != "params.pth"]
+                               os.path.splitext(file_name)[1] == ".pth" and file_name != "params.pth"]
                 # if there are any other files which end with .pth and aren't "params.pth"
                 # assume that they each are optimizer parameters
                 if len(other_files) > 0:
@@ -521,7 +507,7 @@ class BaseRLModel(object):
             episode_reward, episode_timesteps = 0.0, 0
 
             while not done:
-                if self.use_sde and self.sde_sample_freq > 0 and  n_steps % self.sde_sample_freq == 0:
+                if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
                     # Sample a new noise matrix
                     self.actor.reset_noise()
 
@@ -577,7 +563,8 @@ class BaseRLModel(object):
                     self.rollout_data['actions'].append(scaled_action[0].copy())
                     self.rollout_data['rewards'].append(reward[0].copy())
                     self.rollout_data['dones'].append(np.array(done_bool[0]).copy())
-                    self.rollout_data['values'].append(self.vf_net(th.FloatTensor(obs).to(self.device))[0].cpu().detach().numpy())
+                    obs_tensor = th.FloatTensor(obs).to(self.device)
+                    self.rollout_data['values'].append(self.vf_net(obs_tensor)[0].cpu().detach().numpy())
 
                 obs = new_obs
                 # Save the true unnormalized observation
@@ -674,9 +661,9 @@ class BaseRLModel(object):
                 with archive.open('params.pth', mode="w") as param_file:
                     th.save(params, param_file)
             if opt_params is not None:
-                for file_name, dict in opt_params.items():
+                for file_name, dict_ in opt_params.items():
                     with archive.open(file_name + '.pth', mode="w") as opt_param_file:
-                        th.save(dict, opt_param_file)
+                        th.save(dict_, opt_param_file)
 
     @staticmethod
     def excluded_save_params():
diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py
index ae31a1f..369841f 100644
--- a/torchy_baselines/common/buffers.py
+++ b/torchy_baselines/common/buffers.py
@@ -1,8 +1,6 @@
 import numpy as np
 import torch as th
 
-from torchy_baselines.common.vec_env import unwrap_vec_normalize
-
 
 class BaseBuffer(object):
     """
@@ -79,7 +77,8 @@ class BaseBuffer(object):
         """
         raise NotImplementedError()
 
-    def _normalize_obs(self, obs, env=None):
+    @staticmethod
+    def _normalize_obs(obs, env=None):
         if env is not None:
             # TODO: get rid of pytorch - numpy conversion
             return th.FloatTensor(env.normalize_obs(obs.numpy()))
diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py
index ba37cd3..9d6189c 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -411,7 +411,8 @@ class TanhBijector(object):
         super(TanhBijector, self).__init__()
         self.epsilon = epsilon
 
-    def forward(self, x):
+    @staticmethod
+    def forward(x):
         return th.tanh(x)
 
     @staticmethod
diff --git a/torchy_baselines/common/logger.py b/torchy_baselines/common/logger.py
index 5b59b76..15528eb 100644
--- a/torchy_baselines/common/logger.py
+++ b/torchy_baselines/common/logger.py
@@ -1,11 +1,10 @@
 """
 Taken from stable-baselines
 """
-import os
 import sys
-import json
-import time
 import datetime
+import json
+import os
 import tempfile
 import warnings
 from collections import defaultdict
@@ -185,15 +184,6 @@ class CSVOutputFormat(KVWriter):
         self.file.close()
 
 
-def summary_val(key, value):
-    """
-    :param key: (str)
-    :param value: (float)
-    """
-    kwargs = {'tag': key, 'simple_value': float(value)}
-    return tf.Summary.Value(**kwargs)
-
-
 def valid_float_value(value):
     """
     Returns True if the value can be successfully cast into a float
diff --git a/torchy_baselines/common/vec_env/__init__.py b/torchy_baselines/common/vec_env/__init__.py
index 2c542e5..38099af 100644
--- a/torchy_baselines/common/vec_env/__init__.py
+++ b/torchy_baselines/common/vec_env/__init__.py
@@ -35,4 +35,4 @@ def sync_envs_normalization(env, eval_env):
         if isinstance(env_tmp, VecNormalize):
             eval_env_tmp.obs_rms = deepcopy(env_tmp.obs_rms)
         env_tmp = env_tmp.venv
-        eval_env_tmp.venv
+        eval_env_tmp = eval_env_tmp.venv
diff --git a/torchy_baselines/common/vec_env/vec_normalize.py b/torchy_baselines/common/vec_env/vec_normalize.py
index 0b9797f..3cb3c67 100644
--- a/torchy_baselines/common/vec_env/vec_normalize.py
+++ b/torchy_baselines/common/vec_env/vec_normalize.py
@@ -71,7 +71,7 @@ class VecNormalize(VecEnvWrapper):
     def normalize_obs(self, obs):
         if self.norm_obs:
             return np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs,
-                          self.clip_obs)
+                           self.clip_obs)
         return obs
 
     def normalize_reward(self, reward):
diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index d3b897e..5dd746f 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -157,7 +157,8 @@ class PPOPolicy(BasePolicy):
             return self.action_dist.proba_distribution(mean_actions, deterministic=deterministic)
 
         elif isinstance(self.action_dist, StateDependentNoiseDistribution):
-            return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde, deterministic=deterministic)
+            return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde,
+                                                       deterministic=deterministic)
 
     def actor_forward(self, obs, deterministic=False):
         latent_pi, _, latent_sde = self._get_latent(obs)
diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index bc5fbcf..12b23cc 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -241,6 +241,9 @@ class SACPolicy(BasePolicy):
     def make_critic(self):
         return Critic(**self.net_args).to(self.device)
 
+    def forward(self, obs):
+        return self.actor(obs)
+
 
 MlpPolicy = SACPolicy
 
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index f3f45fb..f5a18d2 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -1,5 +1,3 @@
-import time
-
 import torch as th
 import torch.nn.functional as F
 import numpy as np
@@ -68,7 +66,8 @@ class SAC(BaseRLModel):
                  _init_setup_model=True):
 
         super(SAC, self).__init__(policy, env, SACPolicy, policy_kwargs, verbose, device,
-                                  create_eval_env=create_eval_env, seed=seed, use_sde=use_sde, sde_sample_freq=sde_sample_freq)
+                                  create_eval_env=create_eval_env, seed=seed,
+                                  use_sde=use_sde, sde_sample_freq=sde_sample_freq)
 
         self.learning_rate = learning_rate
         self.target_entropy = target_entropy
@@ -131,7 +130,8 @@ class SAC(BaseRLModel):
 
         self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
         self.policy = self.policy_class(self.observation_space, self.action_space,
-                                        self.learning_rate, use_sde=self.use_sde, device=self.device, **self.policy_kwargs)
+                                        self.learning_rate, use_sde=self.use_sde,
+                                        device=self.device, **self.policy_kwargs)
         self.policy = self.policy.to(self.device)
         self._create_aliases()
 
@@ -203,7 +203,6 @@ class SAC(BaseRLModel):
                 ent_coef_loss.backward()
                 self.ent_coef_optimizer.step()
 
-
             with th.no_grad():
                 # if self.use_sde:
                 #     self.actor.reset_noise(batch_size=batch_size)
diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py
index a24e0db..d9ce51b 100644
--- a/torchy_baselines/td3/td3.py
+++ b/torchy_baselines/td3/td3.py
@@ -1,5 +1,3 @@
-import time
-
 import torch as th
 import torch.nn.functional as F
 import numpy as np
@@ -62,7 +60,8 @@ class TD3(BaseRLModel):
                  seed=0, device='auto', _init_setup_model=True):
 
         super(TD3, self).__init__(policy, env, TD3Policy, policy_kwargs, verbose, device,
-                                  create_eval_env=create_eval_env, seed=seed, use_sde=use_sde, sde_sample_freq=sde_sample_freq)
+                                  create_eval_env=create_eval_env, seed=seed,
+                                  use_sde=use_sde, sde_sample_freq=sde_sample_freq)
 
         self.buffer_size = buffer_size
         self.learning_rate = learning_rate
@@ -94,7 +93,8 @@ class TD3(BaseRLModel):
         self.set_random_seed(self.seed)
         self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
         self.policy = self.policy_class(self.observation_space, self.action_space,
-                                        self.learning_rate, use_sde=self.use_sde, device=self.device, **self.policy_kwargs)
+                                        self.learning_rate, use_sde=self.use_sde,
+                                        device=self.device, **self.policy_kwargs)
         self.policy = self.policy.to(self.device)
         self._create_aliases()
 
@@ -209,7 +209,8 @@ class TD3(BaseRLModel):
         # self._update_learning_rate(self.policy.optimizer)
 
         # Unpack
-        obs, action, advantage, returns = [self.rollout_data[key] for key in ['observations', 'actions', 'advantage', 'returns']]
+        obs, action, advantage, returns = [self.rollout_data[key] for key in
+                                           ['observations', 'actions', 'advantage', 'returns']]
 
         log_prob, entropy = self.actor.evaluate_actions(obs, action)
         values = self.vf_net(obs).flatten()

From 89db65b1fb57eee5bd9a6394d5ba43d09b26aa84 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 20 Jan 2020 11:58:00 +0100
Subject: [PATCH 45/47] Improve logger testing + add readers

---
 setup.py                          |  7 ++-
 tests/test_logger.py              | 82 +++++++++++++++++++++++++++++++
 torchy_baselines/common/logger.py | 65 ++++++++++++++++++++++++
 3 files changed, 152 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_logger.py

diff --git a/setup.py b/setup.py
index e5e411f..037e4bd 100644
--- a/setup.py
+++ b/setup.py
@@ -24,8 +24,11 @@ setup(name='torchy_baselines',
             'sphinx-autobuild',
             'sphinx-rtd-theme'
         ],
-        'render': [
-            'opencv-python'
+        'extra': [
+            # For render
+            'opencv-python',
+            # For reading logs
+            'pandas'
         ]
       },
       description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.',
diff --git a/tests/test_logger.py b/tests/test_logger.py
new file mode 100644
index 0000000..b55a616
--- /dev/null
+++ b/tests/test_logger.py
@@ -0,0 +1,82 @@
+import os
+import shutil
+
+import pytest
+import numpy as np
+
+from torchy_baselines.common.logger import (make_output_format, read_csv, read_json, DEBUG, ScopedConfigure,
+    info, debug, set_level, configure, logkv, logkvs, dumpkvs, logkv_mean, warn, error, reset)
+
+KEY_VALUES = {
+    "test": 1,
+    "b": -3.14,
+    "8": 9.9,
+    "l": [1, 2],
+    "a": np.array([1, 2, 3]),
+    "f": np.array(1),
+    "g": np.array([[[1]]]),
+}
+
+LOG_DIR = '/tmp/torchy_baselines/'
+
+
+def test_main():
+    """
+    tests for the logger module
+    """
+    info("hi")
+    debug("shouldn't appear")
+    set_level(DEBUG)
+    debug("should appear")
+    folder = "/tmp/testlogging"
+    if os.path.exists(folder):
+        shutil.rmtree(folder)
+    configure(folder=folder)
+    logkv("a", 3)
+    logkv("b", 2.5)
+    dumpkvs()
+    logkv("b", -2.5)
+    logkv("a", 5.5)
+    dumpkvs()
+    info("^^^ should see a = 5.5")
+    logkv_mean("b", -22.5)
+    logkv_mean("b", -44.4)
+    logkv("a", 5.5)
+    dumpkvs()
+    with ScopedConfigure(None, None):
+        info("^^^ should see b = 33.3")
+
+    with ScopedConfigure("/tmp/test-logger/", ["json"]):
+        logkv("b", -2.5)
+        dumpkvs()
+
+    reset()
+    logkv("a", "longasslongasslongasslongasslongasslongassvalue")
+    dumpkvs()
+    warn("hey")
+    error("oh")
+    logkvs({"test": 1})
+
+
+@pytest.mark.parametrize('_format', ['stdout', 'log', 'json', 'csv'])
+def test_make_output(_format):
+    """
+    test make output
+
+    :param _format: (str) output format
+    """
+    writer = make_output_format(_format, LOG_DIR)
+    writer.writekvs(KEY_VALUES)
+    if _format == "csv":
+        read_csv(LOG_DIR + 'progress.csv')
+    elif _format == 'json':
+        read_json(LOG_DIR + 'progress.json')
+    writer.close()
+
+
+def test_make_output_fail():
+    """
+    test value error on logger
+    """
+    with pytest.raises(ValueError):
+        make_output_format('dummy_format', LOG_DIR)
diff --git a/torchy_baselines/common/logger.py b/torchy_baselines/common/logger.py
index 15528eb..a909285 100644
--- a/torchy_baselines/common/logger.py
+++ b/torchy_baselines/common/logger.py
@@ -505,3 +505,68 @@ def configure(folder=None, format_strs=None):
 
     Logger.CURRENT = Logger(folder=folder, output_formats=output_formats)
     log('Logging to %s' % folder)
+
+
+def reset():
+    """
+    reset the current logger
+    """
+    if Logger.CURRENT is not Logger.DEFAULT:
+        Logger.CURRENT.close()
+        Logger.CURRENT = Logger.DEFAULT
+        log('Reset logger')
+
+
+class ScopedConfigure(object):
+    def __init__(self, folder=None, format_strs=None):
+        """
+        Class for using context manager while logging
+
+        usage:
+        with ScopedConfigure(folder=None, format_strs=None):
+            {code}
+
+        :param folder: (str) the logging folder
+        :param format_strs: ([str]) the list of output logging format
+        """
+        self.dir = folder
+        self.format_strs = format_strs
+        self.prevlogger = None
+
+    def __enter__(self):
+        self.prevlogger = Logger.CURRENT
+        configure(folder=self.dir, format_strs=self.format_strs)
+
+    def __exit__(self, *args):
+        Logger.CURRENT.close()
+        Logger.CURRENT = self.prevlogger
+
+
+# ================================================================
+# Readers
+# ================================================================
+
+def read_json(fname):
+    """
+    read a json file using pandas
+
+    :param fname: (str) the file path to read
+    :return: (pandas DataFrame) the data in the json
+    """
+    import pandas
+    data = []
+    with open(fname, 'rt') as file_handler:
+        for line in file_handler:
+            data.append(json.loads(line))
+    return pandas.DataFrame(data)
+
+
+def read_csv(fname):
+    """
+    read a csv file using pandas
+
+    :param fname: (str) the file path to read
+    :return: (pandas DataFrame) the data in the csv
+    """
+    import pandas
+    return pandas.read_csv(fname, index_col=None, comment='#')

From e5c6601726a79fc874c5458b8c41b82cb1f39bd8 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 20 Jan 2020 11:58:16 +0100
Subject: [PATCH 46/47] Update VecNormalize (pickling) and improve tests

---
 tests/test_vec_normalize.py                   | 127 ++++++++++++++---
 .../common/vec_env/vec_normalize.py           | 128 +++++++++++++-----
 2 files changed, 205 insertions(+), 50 deletions(-)

diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py
index 181c42d..ab0f7f1 100644
--- a/tests/test_vec_normalize.py
+++ b/tests/test_vec_normalize.py
@@ -3,12 +3,49 @@ import pytest
 import numpy as np
 
 from torchy_baselines.common.running_mean_std import RunningMeanStd
-from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv
-from torchy_baselines.common.vec_env.vec_normalize import VecNormalize
+from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize, VecFrameStack, sync_envs_normalization
 from torchy_baselines import CEMRL, SAC, TD3
 
 ENV_ID = 'Pendulum-v0'
 
+def make_env():
+    return gym.make(ENV_ID)
+
+def check_rms_equal(rmsa, rmsb):
+    assert np.all(rmsa.mean == rmsb.mean)
+    assert np.all(rmsa.var == rmsb.var)
+    assert np.all(rmsa.count == rmsb.count)
+
+
+def check_vec_norm_equal(norma, normb):
+    assert norma.observation_space == normb.observation_space
+    assert norma.action_space == normb.action_space
+    assert norma.num_envs == normb.num_envs
+
+    check_rms_equal(norma.obs_rms, normb.obs_rms)
+    check_rms_equal(norma.ret_rms, normb.ret_rms)
+    assert norma.clip_obs == normb.clip_obs
+    assert norma.clip_reward == normb.clip_reward
+    assert norma.norm_obs == normb.norm_obs
+    assert norma.norm_reward == normb.norm_reward
+
+    assert np.all(norma.ret == normb.ret)
+    assert norma.gamma == normb.gamma
+    assert norma.epsilon == normb.epsilon
+    assert norma.training == normb.training
+
+def _make_warmstart_cartpole():
+    """Warm-start VecNormalize by stepping through CartPole"""
+    venv = DummyVecEnv([lambda: gym.make("CartPole-v1")])
+    venv = VecNormalize(venv)
+    venv.reset()
+    venv.get_original_obs()
+
+    for _ in range(100):
+        actions = [venv.action_space.sample()]
+        venv.step(actions)
+    return venv
+
 
 def test_runningmeanstd():
     """Test RunningMeanStd object"""
@@ -27,29 +64,87 @@ def test_runningmeanstd():
         assert np.allclose(moments_1, moments_2)
 
 
-def test_vec_env():
+def test_vec_env(tmpdir):
     """Test VecNormalize Object"""
+    clip_obs = 0.5
+    clip_reward = 5.0
 
-    def make_env():
-        return gym.make(ENV_ID)
-
-    env = DummyVecEnv([make_env])
-    env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10., clip_reward=10.)
-    _, done = env.reset(), [False]
-    obs = None
+    orig_venv = DummyVecEnv([make_env])
+    norm_venv = VecNormalize(orig_venv, norm_obs=True, norm_reward=True, clip_obs=clip_obs, clip_reward=clip_reward)
+    _, done = norm_venv.reset(), [False]
     while not done[0]:
-        actions = [env.action_space.sample()]
-        obs, _, done, _ = env.step(actions)
-    assert np.max(obs) <= 10
+        actions = [norm_venv.action_space.sample()]
+        obs, rew, done, _ = norm_venv.step(actions)
+        assert np.max(np.abs(obs)) <= clip_obs
+        assert np.max(np.abs(rew)) <= clip_reward
+
+    path = str(tmpdir.join("vec_normalize"))
+    norm_venv.save(path)
+    deserialized = VecNormalize.load(path, venv=orig_venv)
+    check_vec_norm_equal(norm_venv, deserialized)
+
+
+def test_get_original():
+    venv = _make_warmstart_cartpole()
+    for _ in range(3):
+        actions = [venv.action_space.sample()]
+        obs, rewards, _, _ = venv.step(actions)
+        obs = obs[0]
+        orig_obs = venv.get_original_obs()[0]
+        rewards = rewards[0]
+        orig_rewards = venv.get_original_reward()[0]
+
+        assert np.all(orig_rewards == 1)
+        assert orig_obs.shape == obs.shape
+        assert orig_rewards.dtype == rewards.dtype
+        assert not np.array_equal(orig_obs, obs)
+        assert not np.array_equal(orig_rewards, rewards)
+        np.testing.assert_allclose(venv.normalize_obs(orig_obs), obs)
+        np.testing.assert_allclose(venv.normalize_reward(orig_rewards), rewards)
+
+
+def test_normalize_external():
+    venv = _make_warmstart_cartpole()
+
+    rewards = np.array([1, 1])
+    norm_rewards = venv.normalize_reward(rewards)
+    assert norm_rewards.shape == rewards.shape
+    # Episode return is almost always >= 1 in CartPole. So reward should shrink.
+    assert np.all(norm_rewards < 1)
 
 
 @pytest.mark.parametrize("model_class", [SAC, TD3, CEMRL])
 def test_offpolicy_normalization(model_class):
-    env = DummyVecEnv([lambda: gym.make(ENV_ID)])
+    env = DummyVecEnv([make_env])
     env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10., clip_reward=10.)
 
-    eval_env = DummyVecEnv([lambda: gym.make(ENV_ID)])
-    eval_env = VecNormalize(eval_env, norm_obs=True, norm_reward=False, clip_obs=10., clip_reward=10.)
+    eval_env = DummyVecEnv([make_env])
+    eval_env = VecNormalize(eval_env, training=False, norm_obs=True, norm_reward=False, clip_obs=10., clip_reward=10.)
 
     model = model_class('MlpPolicy', env, verbose=1)
     model.learn(total_timesteps=1000, eval_env=eval_env, eval_freq=500)
+
+
+def test_sync_vec_normalize():
+    env = DummyVecEnv([make_env])
+    env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10., clip_reward=10.)
+    env = VecFrameStack(env, 1)
+
+    eval_env = DummyVecEnv([make_env])
+    eval_env = VecNormalize(eval_env, training=False, norm_obs=True, norm_reward=True, clip_obs=10., clip_reward=10.)
+    eval_env = VecFrameStack(eval_env, 1)
+
+    env.reset()
+    # Initialize running mean
+    for _ in range(100):
+        env.step([env.action_space.sample()])
+
+    obs = env.reset()
+    original_obs = env.get_original_obs()
+    # Normalization must be different
+    assert not np.allclose(obs, eval_env.normalize_obs(original_obs))
+
+    sync_envs_normalization(env, eval_env)
+
+    # Now they must be synced
+    assert np.allclose(obs, eval_env.normalize_obs(original_obs))
diff --git a/torchy_baselines/common/vec_env/vec_normalize.py b/torchy_baselines/common/vec_env/vec_normalize.py
index 3cb3c67..0824f6a 100644
--- a/torchy_baselines/common/vec_env/vec_normalize.py
+++ b/torchy_baselines/common/vec_env/vec_normalize.py
@@ -38,6 +38,45 @@ class VecNormalize(VecEnvWrapper):
         self.old_obs = np.array([])
         self.old_reward = np.array([])
 
+    def __getstate__(self):
+        """
+        Gets state for pickling.
+
+        Excludes self.venv, as in general VecEnv's may not be pickleable."""
+        state = self.__dict__.copy()
+        # these attributes are not pickleable
+        del state['venv']
+        del state['class_attributes']
+        # these attributes depend on the above and so we would prefer not to pickle
+        del state['ret']
+        return state
+
+    def __setstate__(self, state):
+        """
+        Restores pickled state.
+
+        User must call set_venv() after unpickling before using.
+
+        :param state: (dict)"""
+        self.__dict__.update(state)
+        assert 'venv' not in state
+        self.venv = None
+
+    def set_venv(self, venv):
+        """
+        Sets the vector environment to wrap to venv.
+
+        Also sets attributes derived from this such as `num_env`.
+
+        :param venv: (VecEnv)
+        """
+        if self.venv is not None:
+            raise ValueError("Trying to set venv of already initialized VecNormalize wrapper.")
+        VecEnvWrapper.__init__(self, venv)
+        if self.obs_rms.mean.shape != self.observation_space.shape:
+            raise ValueError("venv is incompatible with current statistics.")
+        self.ret = np.zeros(self.num_envs)
+
     def step_wait(self):
         """
         Apply sequence of actions to sequence of environments
@@ -46,37 +85,44 @@ class VecNormalize(VecEnvWrapper):
         where 'news' is a boolean vector indicating whether each element is new.
         """
         obs, rews, news, infos = self.venv.step_wait()
-        self.ret = self.ret * self.gamma + rews
-        self.old_obs = obs.copy()
-        self.old_reward = rews.copy()
-        obs = self._normalize_observation(obs)
-        if self.norm_reward:
-            if self.training:
-                self.ret_rms.update(self.ret)
-            rews = self.normalize_reward(rews)
+        self.old_obs = obs
+        self.old_rews = rews
+
+        if self.training:
+            self.obs_rms.update(obs)
+        obs = self.normalize_obs(obs)
+
+        if self.training:
+            self._update_reward(rews)
+        rews = self.normalize_reward(rews)
+
         self.ret[news] = 0
         return obs, rews, news, infos
 
-    def _normalize_observation(self, obs):
-        """
-        :param obs: (numpy tensor)
-        """
-        if self.norm_obs:
-            if self.training:
-                self.obs_rms.update(obs)
-            return self.normalize_obs(obs)
-        else:
-            return obs
+    def _update_reward(self, reward):
+        """Update reward normalization statistics."""
+        self.ret = self.ret * self.gamma + reward
+        self.ret_rms.update(self.ret)
 
     def normalize_obs(self, obs):
+        """
+        Normalize observations using this VecNormalize's observations statistics.
+        Calling this method does not update statistics.
+        """
         if self.norm_obs:
-            return np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs,
-                           self.clip_obs)
+            obs = np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon),
+                          -self.clip_obs,
+                          self.clip_obs)
         return obs
 
     def normalize_reward(self, reward):
+        """
+        Normalize rewards using this VecNormalize's rewards statistics.
+        Calling this method does not update statistics.
+        """
         if self.norm_reward:
-            return np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward)
+            reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon),
+                           -self.clip_reward, self.clip_reward)
         return reward
 
     def unnormalize_obs(self, obs):
@@ -91,31 +137,45 @@ class VecNormalize(VecEnvWrapper):
 
     def get_original_obs(self):
         """
-        returns the unnormalized observation
-
-        :return: (numpy float)
+        Returns an unnormalized version of the observations from the most recent
+        step or reset.
         """
-        return self.old_obs
+        return self.old_obs.copy()
 
     def get_original_reward(self):
         """
-        returns the unnormalized observation
-
-        :return: (numpy float)
+        Returns an unnormalized version of the rewards from the most recent step.
         """
-        return self.old_reward
+        return self.old_rews.copy()
 
     def reset(self):
         """
         Reset all environments
         """
         obs = self.venv.reset()
-        if len(np.array(obs).shape) == 1:  # for when num_cpu is 1
-            self.old_obs = [obs]
-        else:
-            self.old_obs = obs
+        self.old_obs = obs
         self.ret = np.zeros(self.num_envs)
-        return self._normalize_observation(obs)
+        if self.training:
+            self._update_reward(self.ret)
+        return self.normalize_obs(obs)
+
+    @staticmethod
+    def load(load_path, venv):
+        """
+        Loads a saved VecNormalize object.
+
+        :param load_path: the path to load from.
+        :param venv: the VecEnv to wrap.
+        :return: (VecNormalize)
+        """
+        with open(load_path, "rb") as file_handler:
+            vec_normalize = pickle.load(file_handler)
+        vec_normalize.set_venv(venv)
+        return vec_normalize
+
+    def save(self, save_path):
+        with open(save_path, "wb") as file_handler:
+            pickle.dump(self, file_handler)
 
     def save_running_average(self, path):
         """

From 0bed698ec5dd28bf098344735bd7d6237b4d37a2 Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@dlr.de>
Date: Mon, 20 Jan 2020 12:57:40 +0100
Subject: [PATCH 47/47] Raise error for abstract methods

---
 torchy_baselines/common/base_class.py           |  4 ++--
 torchy_baselines/common/vec_env/base_vec_env.py | 14 +++++++-------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index 0f20211..7a29847 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -276,7 +276,7 @@ class BaseRLModel(object):
         :param n_eval_episodes: (int) Number of episode to evaluate the agent
         :return: (BaseRLModel) the trained model
         """
-        pass
+        raise NotImplementedError()
 
     @abstractmethod
     def predict(self, observation, state=None, mask=None, deterministic=False):
@@ -289,7 +289,7 @@ class BaseRLModel(object):
         :param deterministic: (bool) Whether or not to return deterministic actions.
         :return: (np.ndarray, np.ndarray) the model's action and the next state (used in recurrent policies)
         """
-        pass
+        raise NotImplementedError()
 
     def load_parameters(self, load_dict, opt_params):
         """
diff --git a/torchy_baselines/common/vec_env/base_vec_env.py b/torchy_baselines/common/vec_env/base_vec_env.py
index 15a27e1..abc36a8 100644
--- a/torchy_baselines/common/vec_env/base_vec_env.py
+++ b/torchy_baselines/common/vec_env/base_vec_env.py
@@ -58,7 +58,7 @@ class VecEnv(object):
 
         :return: ([int] or [float]) observation
         """
-        pass
+        raise NotImplementedError()
 
     @abstractmethod
     def step_async(self, actions):
@@ -70,7 +70,7 @@ class VecEnv(object):
         You should not call this if a step_async run is
         already pending.
         """
-        pass
+        raise NotImplementedError()
 
     @abstractmethod
     def step_wait(self):
@@ -79,14 +79,14 @@ class VecEnv(object):
 
         :return: ([int] or [float], [float], [bool], dict) observation, reward, done, information
         """
-        pass
+        raise NotImplementedError()
 
     @abstractmethod
     def close(self):
         """
         Clean up the environment's resources.
         """
-        pass
+        raise NotImplementedError()
 
     @abstractmethod
     def get_attr(self, attr_name, indices=None):
@@ -97,7 +97,7 @@ class VecEnv(object):
         :param indices: (list,int) Indices of envs to get attribute from
         :return: (list) List of values of 'attr_name' in all environments
         """
-        pass
+        raise NotImplementedError()
 
     @abstractmethod
     def set_attr(self, attr_name, value, indices=None):
@@ -109,7 +109,7 @@ class VecEnv(object):
         :param indices: (list,int) Indices of envs to assign value
         :return: (NoneType)
         """
-        pass
+        raise NotImplementedError()
 
     @abstractmethod
     def env_method(self, method_name, *method_args, **method_kwargs):
@@ -122,7 +122,7 @@ class VecEnv(object):
         :param method_kwargs: (dict) Any keyword arguments to provide in the call
         :return: (list) List of items returned by the environment's method call
         """
-        pass
+        raise NotImplementedError()
 
     def step(self, actions):
         """