diff --git a/README.md b/README.md
index b5624ec..0b19832 100644
--- a/README.md
+++ b/README.md
@@ -21,6 +21,10 @@ 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
diff --git a/setup.py b/setup.py
index f40a7da..640112e 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.5a",
+      version="0.0.6a",
       )
 
 # python setup.py sdist
diff --git a/tests/test_sde.py b/tests/test_sde.py
index 03c8f62..09b48e8 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
+from torchy_baselines import A2C, TD3
 from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize
 from torchy_baselines.common.monitor import Monitor
 
@@ -58,3 +58,20 @@ def test_state_dependent_noise(model_class):
     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)
     model.learn(total_timesteps=int(1000), log_interval=5, eval_freq=500, eval_env=eval_env)
+
+
+@pytest.mark.parametrize("model_class", [TD3])
+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))
+    model.learn(total_timesteps=int(1000), eval_freq=500)
+
+
+def test_scheduler():
+    def scheduler(progress):
+        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)
+    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_normalize.py b/tests/test_vec_normalize.py
index 78437a3..181c42d 100644
--- a/tests/test_vec_normalize.py
+++ b/tests/test_vec_normalize.py
@@ -1,9 +1,11 @@
 import gym
+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 import CEMRL, SAC, TD3
 
 ENV_ID = 'Pendulum-v0'
 
@@ -39,3 +41,15 @@ def test_vec_env():
         actions = [env.action_space.sample()]
         obs, _, done, _ = env.step(actions)
     assert np.max(obs) <= 10
+
+
+@pytest.mark.parametrize("model_class", [SAC, TD3, CEMRL])
+def test_offpolicy_normalization(model_class):
+    env = DummyVecEnv([lambda: gym.make(ENV_ID)])
+    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.)
+
+    model = model_class('MlpPolicy', env, verbose=1)
+    model.learn(total_timesteps=1000, eval_env=eval_env, eval_freq=500)
diff --git a/torchy_baselines/__init__.py b/torchy_baselines/__init__.py
index a5896e6..b8383db 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.5a"
+__version__ = "0.0.6a"
diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py
index 140a77e..a06074b 100644
--- a/torchy_baselines/a2c/a2c.py
+++ b/torchy_baselines/a2c/a2c.py
@@ -25,6 +25,7 @@ class A2C(PPO):
         (i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel)
     :param gamma: (float) Discount factor
     :param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator
+        Equivalent to classic advantage when set to 1.
     :param ent_coef: (float) Entropy coefficient for the loss calculation
     :param vf_coef: (float) Value function coefficient for the loss calculation
     :param max_grad_norm: (float) The maximum value for the gradient clipping
@@ -92,7 +93,7 @@ class A2C(PPO):
                 action = action.long().flatten()
 
             # TODO: avoid second computation of everything because of the gradient
-            values, log_prob, entropy = self.policy.get_policy_stats(obs, action)
+            values, log_prob, entropy = self.policy.evaluate_actions(obs, action)
             values = values.flatten()
 
             # Normalize advantage (not present in the original implementation)
diff --git a/torchy_baselines/cem_rl/cem.py b/torchy_baselines/cem_rl/cem.py
index 57513c5..ee4f484 100644
--- a/torchy_baselines/cem_rl/cem.py
+++ b/torchy_baselines/cem_rl/cem.py
@@ -5,33 +5,42 @@ import numpy as np
 # or https://github.com/facebookresearch/nevergrad
 
 class CEM(object):
-
-    """
-    Cross-entropy method with diagonal covariance (separable CEM)
     """
+    Cross-entropy method with diagonal covariance (separable CEM).
 
-    def __init__(self, num_params,
-                 mu_init=None,
-                 sigma_init=1e-3,
-                 pop_size=256,
-                 damp=1e-3,
-                 damp_limit=1e-5,
-                 parents=None,
-                 elitism=False,
-                 antithetic=False):
+    :param num_params: (int) Number of parameters per individual (dimension of the problem)
+    :param mu_init: (np.ndarray) Initial mean of the population distribution
+        Taken to be zero if None is passed.
+    :param sigma_init: (float) Initial standard deviation of the population distribution
+    :param pop_size: (int) Number of individuals in the population
+    :param damping_init: (float) Initial value of damping for preventing from early convergence.
+    :param damping_final: (float) Final value of damping
+    :param parents: (int) Number of parents used to compute the new distribution
+        of individuals.
+    :param elitism: (bool) Keep the best known individual in the population
+    :param antithetic: (bool) Use a finite difference like method for sampling
+        (mu + epsilon, mu - epsilon)
+    """
+    def __init__(self, num_params, mu_init=None, sigma_init=1e-3,
+                 pop_size=256, damping_init=1e-3, damping_final=1e-5,
+                 parents=None, elitism=False, antithetic=False):
         super(CEM, self).__init__()
-        # misc
+
         self.num_params = num_params
 
-        # distribution parameters
+        # Distribution parameters
         if mu_init is None:
             self.mu = np.zeros(self.num_params)
         else:
             self.mu = np.array(mu_init)
+
         self.sigma = sigma_init
-        self.damp = damp
-        self.damp_limit = damp_limit
+        # Damping parameters
+        self.damping = damping_init
+        self.damping_final = damping_final
+        # Exponential moving average decay for damping
         self.tau = 0.95
+        # Covariance matrix, here only the diagonal
         self.cov = self.sigma * np.ones(self.num_params)
 
         # elite stuff
@@ -39,16 +48,20 @@ class CEM(object):
         self.elite = np.sqrt(self.sigma) * np.random.rand(self.num_params)
         self.elite_score = None
 
-        # sampling stuff
+        # sampling parameters
         self.pop_size = pop_size
         self.antithetic = antithetic
 
         if self.antithetic:
             assert (self.pop_size % 2 == 0), "Population size must be even"
+
         if parents is None or parents <= 0:
             self.parents = pop_size // 2
         else:
             self.parents = parents
+
+        # Weighting for computing the new mean of the distributions
+        # from the parents. The better the individual, the higher the weight
         self.weights = np.array([np.log((self.parents + 1) / i)
                                  for i in range(1, self.parents + 1)])
         self.weights /= self.weights.sum()
@@ -56,43 +69,56 @@ class CEM(object):
     def ask(self, pop_size):
         """
         Returns a list of candidates parameters
+
+        :param pop_size: (int)
+        :return: ([np.ndarray])
         """
         if self.antithetic and not pop_size % 2:
             epsilon_half = np.random.randn(pop_size // 2, self.num_params)
             epsilon = np.concatenate([epsilon_half, - epsilon_half])
-
         else:
             epsilon = np.random.randn(pop_size, self.num_params)
 
-        inds = self.mu + epsilon * np.sqrt(self.cov)
-        if self.elitism:
-            inds[-1] = self.elite
+        individuals = self.mu + epsilon * np.sqrt(self.cov)
 
-        return inds
+        # Keep the best known individual in the population
+        if self.elitism:
+            individuals[-1] = self.elite
+
+        return individuals
 
     def tell(self, solutions, scores):
         """
         Updates the distribution
+
+        :param solutions: ([np.ndarray])
+        :param scores: ([float]) episode reward.
         """
+        # Convert rewards (we want to maximize) to cost (we want to minimize)
         scores = np.array(scores)
         scores *= -1
+        # Sort the individuals by fitness
         idx_sorted = np.argsort(scores)
 
         old_mu = self.mu
-        self.damp = self.damp * self.tau + (1 - self.tau) * self.damp_limit
+        # Update damping using a moving average
+        self.damping = self.damping * self.tau + (1 - self.tau) * self.damping_final
         # self.mu = self.weights @ solutions[idx_sorted[:self.parents]]
         self.mu = self.weights.dot(solutions[idx_sorted[:self.parents]])
 
+        # CMA-ES style would be to use the new mean here
         z = (solutions[idx_sorted[:self.parents]] - old_mu)
-        self.cov = 1 / self.parents * self.weights.dot(z * z) + self.damp * np.ones(self.num_params)
+        self.cov = 1 / self.parents * self.weights.dot(z * z) + self.damping * np.ones(self.num_params)
 
+        # Retrieve the best individual
         self.elite = solutions[idx_sorted[0]]
         self.elite_score = scores[idx_sorted[0]]
-        # print(self.cov)
 
     def get_distrib_params(self):
         """
-        Returns the parameters of the distrubtion:
-        the mean and sigma
+        Returns the parameters of the distribution:
+        the mean and standard deviation.
+
+        :return: (np.ndarray, np.ndarray)
         """
         return np.copy(self.mu), np.copy(self.cov)
diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py
index 3305cb5..41efdf3 100644
--- a/torchy_baselines/cem_rl/cem_rl.py
+++ b/torchy_baselines/cem_rl/cem_rl.py
@@ -5,21 +5,53 @@ import torch as th
 from torchy_baselines.cem_rl.cem import CEM
 from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.td3.td3 import TD3
+from torchy_baselines.common.vec_env import sync_envs_normalization
 
 
 class CEMRL(TD3):
     """
-    Implementation of CEM-RL
+    Implementation of CEM-RL, in fact CEM combined with TD3.
 
     Paper: https://arxiv.org/abs/1810.01222
     Code: https://github.com/apourchot/CEM-RL
-    """
 
+    :param policy: (TD3Policy or str) The policy model to use (MlpPolicy, CnnPolicy, ...)
+    :param env: (Gym environment or str) The environment to learn from (if registered in Gym, can be str)
+    :param sigma_init: (float) Initial standard deviation of the population distribution
+    :param pop_size: (int) Number of individuals in the population
+    :param damping_init: (float)  Initial value of damping for preventing from early convergence.
+    :param damping_final: (float) Final value of damping
+    :param elitism: (bool) Keep the best known individual in the population
+    :param n_grad: (int) Number of individuals that will receive a gradient update.
+        Half of the population size in the paper.
+    :param buffer_size: (int) size of the replay buffer
+    :param learning_rate: (float or callable) learning rate for adam optimizer,
+        the same learning rate will be used for all networks (Q-Values and Actor networks)
+        it can be a function of the current progress (from 1 to 0)
+    :param policy_delay: (int) Policy and target networks will only be updated once every policy_delay steps
+        per training steps. The Q values will be updated policy_delay more often (update every training step).
+    :param learning_starts: (int) how many steps of the model to collect transitions for before learning starts
+    :param gamma: (float) the discount factor
+    :param batch_size: (int) Minibatch size for each gradient update
+    :param tau: (float) the soft update coefficient ("polyak update" of the target networks, between 0 and 1)
+    :param action_noise: (ActionNoise) the action noise type. Cf common.noise for the different action noise type.
+    :param target_policy_noise: (float) Standard deviation of gaussian noise added to target policy
+        (smoothing noise)
+    :param target_noise_clip: (float) Limit for absolute value of target policy smoothing noise.
+    :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
+    :param verbose: (int) the verbosity level: 0 none, 1 training information, 2 tensorflow debug
+    :param seed: (int) Seed for the pseudo random generators
+    :param device: (str or th.device) Device (cpu, cuda, ...) on which the code should be run.
+        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, sigma_init=1e-3, pop_size=10,
-                 damp=1e-3, damp_limit=1e-5, elitism=False, n_grad=5,
-                 policy_delay=2, batch_size=100,
-                 buffer_size=int(1e6), learning_rate=1e-3,
-                 action_noise=None, learning_starts=100, tau=0.005,
+                 damping_init=1e-3, damping_final=1e-5, elitism=False, n_grad=5,
+                 buffer_size=int(1e6), learning_rate=1e-3, policy_delay=2,
+                 learning_starts=100, gamma=0.99, batch_size=100, tau=0.005,
+                 action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5,
                  n_episodes_rollout=1, update_style='original',
                  tensorboard_log=None, create_eval_env=False,
                  policy_kwargs=None, verbose=0, seed=0, device='auto',
@@ -27,18 +59,21 @@ class CEMRL(TD3):
 
         super(CEMRL, self).__init__(policy, env,
                                     buffer_size=buffer_size, learning_rate=learning_rate, seed=seed, device=device,
-                                    action_noise=action_noise, learning_starts=learning_starts,
-                                    n_episodes_rollout=n_episodes_rollout, tau=tau,
+                                    action_noise=action_noise, target_policy_noise=target_policy_noise,
+                                    target_noise_clip=target_noise_clip, learning_starts=learning_starts,
+                                    n_episodes_rollout=n_episodes_rollout, tau=tau, gamma=gamma,
                                     policy_kwargs=policy_kwargs, verbose=verbose,
                                     policy_delay=policy_delay, batch_size=batch_size,
-                                    create_eval_env=create_eval_env,
+                                    create_eval_env=create_eval_env, tensorboard_log=tensorboard_log,
                                     _init_setup_model=False)
 
+        # Evolution strategy method that follows cma-es interface (ask-tell)
+        # for now, only CEM is implemented
         self.es = None
         self.sigma_init = sigma_init
         self.pop_size = pop_size
-        self.damp = damp
-        self.damp_limit = damp_limit
+        self.damping_init = damping_init
+        self.damping_final = damping_final
         self.elitism = elitism
         self.n_grad = n_grad
         self.es_params = None
@@ -52,7 +87,7 @@ class CEMRL(TD3):
         super(CEMRL, self)._setup_model()
         params_vector = self.actor.parameters_to_vector()
         self.es = CEM(len(params_vector), mu_init=params_vector,
-                      sigma_init=self.sigma_init, damp=self.damp, damp_limit=self.damp_limit,
+                      sigma_init=self.sigma_init, damping_init=self.damping_init, damping_final=self.damping_final,
                       pop_size=self.pop_size, antithetic=not self.pop_size % 2, parents=self.pop_size // 2,
                       elitism=self.elitism)
 
@@ -103,7 +138,7 @@ class CEMRL(TD3):
                             n_training_steps = 2 * (actor_steps // self.n_grad)
                         for it in range(n_training_steps):
                             # Sample replay buffer
-                            replay_data = self.replay_buffer.sample(self.batch_size)
+                            replay_data = self.replay_buffer.sample(self.batch_size, env=self._vec_normalize_env)
                             self.train_critic(replay_data=replay_data)
 
                             # Delayed policy updates
@@ -118,6 +153,7 @@ class CEMRL(TD3):
                 timesteps_since_eval %= eval_freq
 
                 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)
                 evaluations.append(mean_reward)
@@ -150,10 +186,6 @@ class CEMRL(TD3):
                 actor_steps += episode_timesteps
                 self.fitnesses.append(episode_reward)
 
-                if self.verbose > 1:
-                    print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
-                        self.num_timesteps, episode_num, episode_timesteps, episode_reward))
-
             self._update_current_progress(self.num_timesteps, total_timesteps)
             self.es.tell(self.es_params, self.fitnesses)
             timesteps_since_eval += actor_steps
diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index ccb3f5e..dfa9ab1 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -11,7 +11,7 @@ import numpy as np
 
 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
+from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize
 from torchy_baselines.common.monitor import Monitor
 from torchy_baselines.common import logger
 from torchy_baselines.common.save_util import data_to_json, json_to_data
@@ -24,13 +24,19 @@ class BaseRLModel(object):
     :param policy: (BasePolicy) Policy object
     :param env: (Gym environment) The environment to learn from
                 (if registered in Gym, can be str. Can be None for loading trained models)
-    :param verbose: (int) the verbosity level: 0 none, 1 training information, 2 debug
     :param policy_base: (BasePolicy) the base policy used by this method
-    :param device: (str or th.device) Device on which the code should.
+    :param policy_kwargs: (dict) additional arguments to be passed to the policy on creation
+    :param verbose: (int) the verbosity level: 0 none, 1 training information, 2 debug
+    :param device: (str or th.device) Device on which the code should run.
         By default, it will try to use a Cuda compatible device and fallback to cpu
         if it is not possible.
+    :param support_multi_env: (bool) Whether the algorithm supports training
+        with multiple environments (as in A2C)
+    :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 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
     """
     __metaclass__ = ABCMeta
 
@@ -50,6 +56,8 @@ class BaseRLModel(object):
             print("Using {} device".format(self.device))
 
         self.env = env
+        # get VecNormalize object if needed
+        self._vec_normalize_env = unwrap_vec_normalize(env)
         self.verbose = verbose
         self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs
         self.observation_space = None
@@ -60,10 +68,14 @@ class BaseRLModel(object):
         self.replay_buffer = None
         self.seed = seed
         self.action_noise = None
+        # Used for SDE only
+        self.rollout_data = None
+        self.use_sde = False
         # Track the training progress (from 1 to 0)
         # this is used to update the learning rate
         self._current_progress = 1
 
+        # Create and wrap the env if needed
         if env is not None:
             if isinstance(env, str):
                 if create_eval_env:
@@ -93,6 +105,12 @@ class BaseRLModel(object):
                                  " environment.")
 
     def _get_eval_env(self, eval_env):
+        """
+        Return the environment that will be used for evaluation.
+
+        :param eval_env: (gym.Env or VecEnv)
+        :return: (VecEnv)
+        """
         if eval_env is None:
             eval_env = self.eval_env
 
@@ -106,6 +124,9 @@ class BaseRLModel(object):
         """
         Rescale the action from [low, high] to [-1, 1]
         (no need for symmetric action space)
+
+        :param action: (np.ndarray)
+        :return: (np.ndarray)
         """
         low, high = self.action_space.low, self.action_space.high
         return 2.0 * ((action - low) / (high - low)) - 1.0
@@ -114,6 +135,9 @@ class BaseRLModel(object):
         """
         Rescale the action from [-1, 1] to [low, high]
         (no need for symmetric action space)
+
+        :param scaled_action: (np.ndarray)
+        :return: (np.ndarray)
         """
         low, high = self.action_space.low, self.action_space.high
         return low + (0.5 * (scaled_action + 1.0) * (high - low))
@@ -126,7 +150,7 @@ class BaseRLModel(object):
         """
         Compute current progress (from 1 to 0)
 
-        :param num_timesteps: (int)
+        :param num_timesteps: (int) current number of timesteps
         :param total_timesteps: (int)
         """
         self._current_progress = 1.0 - float(num_timesteps) / float(total_timesteps)
@@ -162,7 +186,7 @@ class BaseRLModel(object):
         """
         returns the current environment (can be None if not defined)
 
-        :return: (Gym Environment) The current environment
+        :return: (gym.Env) The current environment
         """
         return self.env
 
@@ -190,7 +214,7 @@ class BaseRLModel(object):
          - observation_space
          - action_space
 
-        :param env: (Gym Environment) The environment for learning a policy
+        :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")
@@ -251,12 +275,14 @@ class BaseRLModel(object):
         Return a trained model.
 
         :param total_timesteps: (int) The total number of samples to train on
-        :param seed: (int) The initial seed for training, if None: keep current seed
         :param callback: (function (dict, dict)) -> boolean function called at every steps with state of the algorithm.
             It takes the local and global variables. If it returns False, training is aborted.
         :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)
         :return: (BaseRLModel) the trained model
         """
         pass
@@ -407,21 +433,33 @@ class BaseRLModel(object):
             self.eval_env.seed(seed)
 
     def _setup_learn(self, eval_env):
+        """
+        Initialize different variables needed for training.
+
+        :param eval_env: (gym.Env or VecEnv)
+        :return: (int, int, [float], np.ndarray, VecEnv)
+        """
         self.start_time = time.time()
         self.ep_info_buffer = deque(maxlen=100)
+
         if self.action_noise is not None:
             self.action_noise.reset()
+
         timesteps_since_eval, episode_num = 0, 0
         evaluations = []
+
         if eval_env is not None and self.seed is not None:
             eval_env.seed(self.seed)
+
         eval_env = self._get_eval_env(eval_env)
         obs = self.env.reset()
         return timesteps_since_eval, episode_num, evaluations, obs, eval_env
 
     def _update_info_buffer(self, infos):
         """
-        Retrieve reward and episode length if using Monitor wrapper.
+        Retrieve reward and episode length and update the buffer
+        if using Monitor wrapper.
+
         :param infos: ([dict])
         """
         for info in infos:
@@ -434,13 +472,39 @@ class BaseRLModel(object):
                          learning_starts=0, num_timesteps=0,
                          replay_buffer=None, obs=None,
                          episode_num=0, log_interval=None):
+        """
+        Collect rollout using the current policy (and possibly fill the replay buffer)
+        TODO: move this method to off-policy base class.
 
+        :param env: (VecEnv)
+        :param n_episodes: (int)
+        :param n_steps: (int)
+        :param action_noise: (ActionNoise)
+        :param deterministic: (bool)
+        :param callback: (callable)
+        :param learning_starts: (int)
+        :param num_timesteps: (int)
+        :param replay_buffer: (ReplayBuffer)
+        :param obs: (np.ndarray)
+        :param episode_num: (int)
+        :param log_interval: (int)
+        """
         episode_rewards = []
         total_timesteps = []
         total_steps, total_episodes = 0, 0
         assert isinstance(env, VecEnv)
         assert env.num_envs == 1
 
+        # Retrieve unnormalized observation for saving into the buffer
+        if self._vec_normalize_env is not None:
+            obs_ = self._vec_normalize_env.get_original_obs()
+
+        self.rollout_data = None
+        if self.use_sde:
+            self.actor.reset_noise()
+            # Reset rollout data
+            self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']}
+
         while total_steps < n_steps or total_episodes < n_episodes:
             done = False
             # Reset environment: not needed for VecEnv
@@ -450,20 +514,29 @@ class BaseRLModel(object):
             while not done:
                 # Select action randomly or according to policy
                 if num_timesteps < learning_starts:
-                    action = np.array([self.action_space.sample()])
+                    # Warmup phase
+                    unscaled_action = np.array([self.action_space.sample()])
                 else:
-                    action = self.predict(obs, deterministic=deterministic)
+                    unscaled_action = self.predict(obs, deterministic=not self.use_sde)
 
                 # Rescale the action from [low, high] to [-1, 1]
-                action = self.scale_action(action)
+                scaled_action = self.scale_action(unscaled_action)
+
+                if self.use_sde:
+                    # When using SDE, the action can be out of bounds
+                    # TODO: fix with squashing and account for that in the proba distribution
+                    clipped_action = np.clip(scaled_action, -1, 1)
+                else:
+                    clipped_action = scaled_action
 
                 # Add noise to the action (improve exploration)
                 if action_noise is not None:
                     # NOTE: in the original implementation of TD3, the noise was applied to the unscaled action
-                    action = np.clip(action + action_noise(), -1, 1)
+                    # Update(October 2019): Not anymore
+                    clipped_action = np.clip(clipped_action + action_noise(), -1, 1)
 
                 # Rescale and perform action
-                new_obs, reward, done, infos = env.step(self.unscale_action(action))
+                new_obs, reward, done, infos = env.step(self.unscale_action(clipped_action))
 
                 done_bool = [float(done[0])]
                 episode_reward += reward
@@ -473,14 +546,34 @@ class BaseRLModel(object):
 
                 # Store data in replay buffer
                 if replay_buffer is not None:
-                    replay_buffer.add(obs, new_obs, action, reward, done_bool)
+                    # Store only the unnormalized version
+                    if self._vec_normalize_env is not None:
+                        new_obs_ = self._vec_normalize_env.get_original_obs()
+                        reward_ = self._vec_normalize_env.get_original_reward()
+                    else:
+                        # Avoid changing the original ones
+                        obs_, new_obs_, reward_ = obs, new_obs, reward
+
+                    replay_buffer.add(obs_, new_obs_, clipped_action, reward_, done_bool)
+
+                if self.rollout_data is not None:
+                    # Assume only one env
+                    self.rollout_data['observations'].append(obs[0].copy())
+                    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())
 
                 obs = new_obs
+                # Save the true unnormalized observation
+                # otherwise obs_ = self._vec_normalize_env.unnormalize_obs(obs)
+                # is a good approximation
+                if self._vec_normalize_env is not None:
+                    obs_ = new_obs_
 
                 num_timesteps += 1
                 episode_timesteps += 1
                 total_steps += 1
-                if n_steps > 0 and total_steps >= n_steps:
+                if 0 < n_steps <= total_steps:
                     break
 
             if done:
@@ -495,29 +588,46 @@ class BaseRLModel(object):
                             episode_num + total_episodes) % log_interval == 0:
                     fps = int(num_timesteps / (time.time() - self.start_time))
                     logger.logkv("episodes", episode_num + total_episodes)
-                    # logger.logkv("mean 100 episode reward", mean_reward)
                     if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0:
                         logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer]))
                         logger.logkv('ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer]))
                     # logger.logkv("n_updates", n_updates)
-                    # logger.logkv("current_lr", current_lr)
                     logger.logkv("fps", fps)
                     logger.logkv('time_elapsed', int(time.time() - self.start_time))
                     logger.logkv("total timesteps", num_timesteps)
+                    if self.use_sde:
+                        logger.logkv("std", (self.actor.get_std()).mean().item())
                     logger.dumpkvs()
 
         mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0
+
+        # Post processing
+        if self.rollout_data is not None:
+            for key in ['observations', 'actions', 'rewards', 'dones']:
+                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
+            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]
+                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
+
         return mean_reward, total_steps, total_episodes, obs
 
     @staticmethod
     def _save_to_file_zip(save_path, data=None, params=None, opt_params=None):
         """Save model to a zip archive
-    
+
         :param save_path: (str) Where to store the model
         :param data: (dict) Class parameters being stored
         :param params: (dict) Model parameters being stored expected to be state_dict
         :param opt_params: (dict) Optimizer parameters being stored expected to contain an entry for every
-                                         optimizer with its name and the state_dict            
+                                         optimizer with its name and the state_dict
         """
 
         # data/params can be None, so do not
@@ -551,7 +661,7 @@ class BaseRLModel(object):
         """
         Returns the names of the parameters that should be excluded by default
         when saving the model.
-        
+
         :return: ([str]) List of parameters that should be excluded from save
         """
         return ["env", "eval_env", "replay_buffer", "rollout_buffer"]
diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py
index 7ca61b8..ae31a1f 100644
--- a/torchy_baselines/common/buffers.py
+++ b/torchy_baselines/common/buffers.py
@@ -1,8 +1,19 @@
 import numpy as np
 import torch as th
 
+from torchy_baselines.common.vec_env import unwrap_vec_normalize
+
 
 class BaseBuffer(object):
+    """
+    Base class that represent a buffer (rollout or replay)
+
+    :param buffer_size: (int) Max number of element in the buffer
+    :param obs_dim: (int) Dimension of the observation
+    :param action_dim: (int) Dimension of the action space
+    :param device: (th.device)
+    :param n_envs: (int) Number of parallel environments
+    """
     def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1):
         super(BaseBuffer, self).__init__()
         self.buffer_size = buffer_size
@@ -29,35 +40,68 @@ class BaseBuffer(object):
         return tensor.transpose(0, 1).reshape(shape[0] * shape[1], *shape[2:])
 
     def size(self):
+        """
+        :return: (int) The current size of the buffer
+        """
         if self.full:
             return self.buffer_size
         return self.pos
 
-    def get_pos(self):
-        return self.pos
-
     def add(self, *args, **kwargs):
+        """
+        Add elements to the buffer.
+        """
         raise NotImplementedError()
 
     def reset(self):
+        """
+        Reset the buffer.
+        """
         self.pos = 0
         self.full = False
 
-    def sample(self, batch_size):
+    def sample(self, batch_size, env=None):
+        """
+        :param batch_size: (int) Number of element to sample
+        :param env: (VecNormalize) [Optional] associated gym VecEnv
+            to normalize the observations/rewards when sampling
+        """
         upper_bound = self.buffer_size if self.full else self.pos
         batch_inds = th.LongTensor(
             np.random.randint(0, upper_bound, size=batch_size))
-        return self._get_samples(batch_inds)
+        return self._get_samples(batch_inds, env=env)
 
-    def _get_samples(self, batch_inds):
+    def _get_samples(self, batch_inds, env=None):
+        """
+        :param batch_inds: (th.Tensor)
+        :param env: (gym.Env)
+        :return: ([th.Tensor])
+        """
         raise NotImplementedError()
 
+    def _normalize_obs(self, obs, env=None):
+        if env is not None:
+            # TODO: get rid of pytorch - numpy conversion
+            return th.FloatTensor(env.normalize_obs(obs.numpy()))
+        return obs
+
+    def _normalize_reward(self, reward, env=None):
+        if env is not None:
+            return th.FloatTensor(env.normalize_reward(reward.numpy()))
+        return reward
+
 
 class ReplayBuffer(BaseBuffer):
     """
-    Taken from https://github.com/apourchot/CEM-RL
-    """
+    Replay buffer used in off-policy algorithms like SAC/TD3.
+    Adapted from from https://github.com/apourchot/CEM-RL
 
+    :param buffer_size: (int) Max number of element in the buffer
+    :param obs_dim: (int) Dimension of the observation
+    :param action_dim: (int) Dimension of the action space
+    :param device: (th.device)
+    :param n_envs: (int) Number of parallel environments
+    """
     def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1):
         super(ReplayBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs)
 
@@ -81,15 +125,27 @@ class ReplayBuffer(BaseBuffer):
             self.full = True
             self.pos = 0
 
-    def _get_samples(self, batch_inds):
-        return (self.observations[batch_inds, 0, :].to(self.device),
+    def _get_samples(self, batch_inds, env=None):
+        return (self._normalize_obs(self.observations[batch_inds, 0, :], env).to(self.device),
                 self.actions[batch_inds, 0, :].to(self.device),
-                self.next_observations[batch_inds, 0, :].to(self.device),
+                self._normalize_obs(self.next_observations[batch_inds, 0, :], env).to(self.device),
                 self.dones[batch_inds].to(self.device),
-                self.rewards[batch_inds].to(self.device))
+                self._normalize_reward(self.rewards[batch_inds], env).to(self.device))
 
 
 class RolloutBuffer(BaseBuffer):
+    """
+    Rollout buffer used in on-policy algorithms like A2C/PPO.
+
+    :param buffer_size: (int) Max number of element in the buffer
+    :param obs_dim: (int) Dimension of the observation
+    :param action_dim: (int) Dimension of the action space
+    :param device: (th.device)
+    :param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator
+        Equivalent to classic advantage when set to 1.
+    :param gamma: (float) Discount factor
+    :param n_envs: (int) Number of parallel environments
+    """
     def __init__(self, buffer_size, obs_dim, action_dim, device='cpu',
                  gae_lambda=1, gamma=0.99, n_envs=1):
         super(RolloutBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs)
@@ -115,7 +171,10 @@ class RolloutBuffer(BaseBuffer):
 
     def compute_returns_and_advantage(self, last_value, dones=False, use_gae=True):
         """
-        From Stable-Baselines PPO2
+        Post-processing step: compute the returns (sum of discounted rewards)
+        and advantage (A(s) = R - V(S)).
+        Adapted from Stable-Baselines PPO2.
+
         :param last_value: (th.Tensor)
         :param dones: ([bool])
         :param use_gae: (bool) Whether to use Generalized Advantage Estimation
@@ -151,6 +210,16 @@ class RolloutBuffer(BaseBuffer):
             self.advantages = self.returns - self.values
 
     def add(self, obs, action, reward, done, value, log_prob):
+        """
+        :param obs: (np.ndarray) Observation
+        :param action: (np.ndarray) Action
+        :param reward: (np.ndarray)
+        :param done: (np.ndarray) End of episode signal.
+        :param value: (th.Tensor) estimated value of the current state
+            following the current policy.
+        :param log_prob: (th.Tensor) log probability of the action
+            following the current policy.
+        """
         if len(log_prob.shape) == 0:
             # Reshape 0-d tensor to avoid error
             log_prob = log_prob.reshape(-1, 1)
@@ -184,7 +253,7 @@ class RolloutBuffer(BaseBuffer):
             yield self._get_samples(indices[start_idx:start_idx + batch_size])
             start_idx += batch_size
 
-    def _get_samples(self, batch_inds):
+    def _get_samples(self, batch_inds, env=None):
         return (self.observations[batch_inds].to(self.device),
                 self.actions[batch_inds].to(self.device),
                 self.values[batch_inds].flatten().to(self.device),
diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py
index 07eeea4..8478871 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -236,6 +236,8 @@ class StateDependentNoiseDistribution(Distribution):
     compute the log probabilty of an action with that noise.
 
     :param action_dim: (int) Number of continuous actions
+    :param full_std: (bool) Whether to use (n_features x n_actions) parameters
+        for the std instead of only (n_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.
@@ -243,16 +245,19 @@ class StateDependentNoiseDistribution(Distribution):
         this allows to ensure boundaries.
     :param epsilon: (float) small value to avoid NaN due to numerical imprecision.
     """
-    def __init__(self, action_dim, use_expln=False,
+    def __init__(self, action_dim, full_std=True, use_expln=False,
                  squash_output=False, epsilon=1e-6):
         super(StateDependentNoiseDistribution, self).__init__()
         self.distribution = None
         self.action_dim = action_dim
+        self.latent_dim = None
         self.mean_actions = None
         self.log_std = None
         self.weights_dist = None
         self.exploration_mat = None
         self.use_expln = use_expln
+        self.full_std = full_std
+        self.epsilon = epsilon
         if squash_output:
             print("== Using TanhBijector ===")
             self.bijector = TanhBijector(epsilon)
@@ -271,12 +276,17 @@ class StateDependentNoiseDistribution(Distribution):
             # From SDE paper, it allows to keep variance
             # above zero and prevent it from growing too fast
             if log_std <= 0:
-                return th.exp(log_std)
+                std = th.exp(log_std)
             else:
-                return th.log(log_std + 1.0) + 1.0
+                std = th.log(log_std + 1.0) + 1.0
         else:
             # Use normal exponential
-            return th.exp(log_std)
+            std = th.exp(log_std)
+
+        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
 
     def sample_weights(self, log_std):
         """
@@ -285,11 +295,11 @@ class StateDependentNoiseDistribution(Distribution):
 
         :param log_std: (th.Tensor)
         """
-        # TODO: reduce the number of learned dimensions (cf TD3)
-        self.weights_dist = Normal(th.zeros_like(log_std), self.get_std(log_std))
+        std = self.get_std(log_std)
+        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=0.0):
+    def proba_distribution_net(self, latent_dim, log_std_init=-2.0):
         """
         Create the layers and parameter that represent the distribution:
         one output will be the deterministic action, the other parameter will be the
@@ -299,10 +309,17 @@ class StateDependentNoiseDistribution(Distribution):
         :param log_std_init: (float) Initial value for the log standard deviation
         :return: (nn.Linear, nn.Parameter)
         """
-        mean_actions = nn.Linear(latent_dim, self.action_dim)
-        log_std = nn.Parameter(th.ones(latent_dim, self.action_dim) * log_std_init)
+        # 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
+        # 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)
+        # 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, log_std
+        return mean_actions_net, log_std
 
     def proba_distribution(self, mean_actions, log_std, latent_pi, deterministic=False):
         """
@@ -314,7 +331,7 @@ class StateDependentNoiseDistribution(Distribution):
         :return: (th.Tensor)
         """
         variance = th.mm(latent_pi.detach() ** 2, self.get_std(log_std) ** 2)
-        self.distribution = Normal(mean_actions, th.sqrt(variance))
+        self.distribution = Normal(mean_actions, th.sqrt(variance + self.epsilon))
 
         if deterministic:
             action = self.mode()
@@ -328,8 +345,11 @@ 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 sample(self, latent_pi):
-        noise = th.mm(latent_pi.detach(), self.exploration_mat)
+        noise = self.get_noise(latent_pi)
         action = self.distribution.mean + noise
         if self.bijector is not None:
             return self.bijector.forward(action)
@@ -405,26 +425,30 @@ class TanhBijector(object):
         return th.log(1 - th.tanh(x) ** 2 + self.epsilon)
 
 
-def make_proba_distribution(action_space, use_sde=False):
+def make_proba_distribution(action_space, use_sde=False, dist_kwargs=None):
     """
     Return an instance of Distribution for the correct type of action space
 
     :param action_space: (Gym Space) the input action space
     :param use_sde: (bool) Force the use of StateDependentNoiseDistribution
         instead of DiagGaussianDistribution
+    :param dist_kwargs: (dict) Keyword arguments to pass to the probabilty distribution
     :return: (Distribution) the approriate Distribution object
     """
+    if dist_kwargs is None:
+        dist_kwargs = {}
+
     if isinstance(action_space, spaces.Box):
         assert len(action_space.shape) == 1, "Error: the action space must be a vector"
         if use_sde:
-            return StateDependentNoiseDistribution(action_space.shape[0])
-        return DiagGaussianDistribution(action_space.shape[0])
+            return StateDependentNoiseDistribution(action_space.shape[0], **dist_kwargs)
+        return DiagGaussianDistribution(action_space.shape[0], **dist_kwargs)
     elif isinstance(action_space, spaces.Discrete):
-        return CategoricalDistribution(action_space.n)
+        return CategoricalDistribution(action_space.n, **dist_kwargs)
     # elif isinstance(action_space, spaces.MultiDiscrete):
-    #     return MultiCategoricalDistribution(action_space.nvec)
+    #     return MultiCategoricalDistribution(action_space.nvec, **dist_kwargs)
     # elif isinstance(action_space, spaces.MultiBinary):
-    #     return BernoulliDistribution(action_space.n)
+    #     return BernoulliDistribution(action_space.n, **dist_kwargs)
     else:
         raise NotImplementedError("Error: probability distribution, not implemented for action space of type {}."
                                   .format(type(action_space)) +
diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py
index 55e356c..4ac7e07 100644
--- a/torchy_baselines/common/policies.py
+++ b/torchy_baselines/common/policies.py
@@ -1,3 +1,5 @@
+from itertools import zip_longest
+
 import torch as th
 import torch.nn as nn
 
@@ -141,3 +143,92 @@ def register_policy(name, policy):
         raise ValueError("Error: the name {} is alreay registered for a different policy, will not override."
                          .format(name))
     _policy_registry[sub_class][name] = policy
+
+
+class MlpExtractor(nn.Module):
+    """
+    Constructs an MLP that receives observations as an input and outputs a latent representation for the policy and
+    a value network. The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many
+    of them are shared between the policy network and the value network. It is assumed to be a list with the following
+    structure:
+
+    1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer.
+       If the number of ints is zero, there will be no shared layers.
+    2. An optional dict, to specify the following non-shared layers for the value network and the policy network.
+       It is formatted like ``dict(vf=[<value layer sizes>], pi=[<policy layer sizes>])``.
+       If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed.
+
+    For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value
+    network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec
+    would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128
+    would be specified as [128, 128].
+
+    Adapted from Stable Baselines.
+
+    :param feature_dim: (int) Dimension of the feature vector (can be the output of a CNN)
+    :param net_arch: ([int or dict]) The specification of the policy and value networks.
+        See above for details on its formatting.
+    :param activation_fn: (nn.Module) The activation function to use for the networks.
+    :param device: (th.device)
+    """
+    def __init__(self, feature_dim, net_arch, activation_fn, device='cpu'):
+        super(MlpExtractor, self).__init__()
+
+        shared_net, policy_net, value_net = [], [], []
+        policy_only_layers = []  # Layer sizes of the network that only belongs to the policy network
+        value_only_layers = []  # Layer sizes of the network that only belongs to the value network
+        last_layer_dim_shared = feature_dim
+
+        # Iterate through the shared layers and build the shared parts of the network
+        for idx, layer in enumerate(net_arch):
+            if isinstance(layer, int):  # Check that this is a shared layer
+                layer_size = layer
+                # TODO: give layer a meaningful name
+                shared_net.append(nn.Linear(last_layer_dim_shared, layer_size))
+                shared_net.append(activation_fn())
+                last_layer_dim_shared = layer_size
+            else:
+                assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts"
+                if 'pi' in layer:
+                    assert isinstance(layer['pi'], list), "Error: net_arch[-1]['pi'] must contain a list of integers."
+                    policy_only_layers = layer['pi']
+
+                if 'vf' in layer:
+                    assert isinstance(layer['vf'], list), "Error: net_arch[-1]['vf'] must contain a list of integers."
+                    value_only_layers = layer['vf']
+                break  # From here on the network splits up in policy and value network
+
+        last_layer_dim_pi = last_layer_dim_shared
+        last_layer_dim_vf = last_layer_dim_shared
+
+        # Build the non-shared part of the network
+        for idx, (pi_layer_size, vf_layer_size) in enumerate(zip_longest(policy_only_layers, value_only_layers)):
+            if pi_layer_size is not None:
+                assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers."
+                policy_net.append(nn.Linear(last_layer_dim_pi, pi_layer_size))
+                policy_net.append(activation_fn())
+                last_layer_dim_pi = pi_layer_size
+
+            if vf_layer_size is not None:
+                assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers."
+                value_net.append(nn.Linear(last_layer_dim_vf, vf_layer_size))
+                value_net.append(activation_fn())
+                last_layer_dim_vf = vf_layer_size
+
+        # Save dim, used to create the distributions
+        self.latent_dim_pi = last_layer_dim_pi
+        self.latent_dim_vf = last_layer_dim_vf
+
+        # Create networks
+        # If the list of layers is empty, the network will just act as an Identity module
+        self.shared_net = nn.Sequential(*shared_net).to(device)
+        self.policy_net = nn.Sequential(*policy_net).to(device)
+        self.value_net = nn.Sequential(*value_net).to(device)
+
+    def forward(self, features):
+        """
+        :return: (th.Tensor, th.Tensor) latent_policy, latent_value of the specified network.
+            If all layers are shared, then ``latent_policy == latent_value``
+        """
+        shared_latent = self.shared_net(features)
+        return self.policy_net(shared_latent), self.value_net(shared_latent)
diff --git a/torchy_baselines/common/vec_env/__init__.py b/torchy_baselines/common/vec_env/__init__.py
index 97f6022..2c542e5 100644
--- a/torchy_baselines/common/vec_env/__init__.py
+++ b/torchy_baselines/common/vec_env/__init__.py
@@ -1,7 +1,38 @@
 # flake8: noqa F401
+from copy import deepcopy
+
 from torchy_baselines.common.vec_env.base_vec_env import AlreadySteppingError, NotSteppingError,\
     VecEnv, VecEnvWrapper, CloudpickleWrapper
 from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv
 from torchy_baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
 from torchy_baselines.common.vec_env.vec_frame_stack import VecFrameStack
 from torchy_baselines.common.vec_env.vec_normalize import VecNormalize
+
+
+def unwrap_vec_normalize(env):
+    """
+    :param env: (gym.Env)
+    :return: (VecNormalize)
+    """
+    env_tmp = env
+    while isinstance(env_tmp, VecEnvWrapper):
+        if isinstance(env_tmp, VecNormalize):
+            return env_tmp
+        env_tmp = env_tmp.venv
+    return None
+
+
+# Define here to avoid circular import
+def sync_envs_normalization(env, eval_env):
+    """
+    Sync eval env and train env when using VecNormalize
+
+    :param env: (gym.Env)
+    :param eval_env: (gym.Env)
+    """
+    env_tmp, eval_env_tmp = env, eval_env
+    while isinstance(env_tmp, VecEnvWrapper):
+        if isinstance(env_tmp, VecNormalize):
+            eval_env_tmp.obs_rms = deepcopy(env_tmp.obs_rms)
+        env_tmp = env_tmp.venv
+        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 6f2b2b3..0b9797f 100644
--- a/torchy_baselines/common/vec_env/vec_normalize.py
+++ b/torchy_baselines/common/vec_env/vec_normalize.py
@@ -36,6 +36,7 @@ class VecNormalize(VecEnvWrapper):
         self.norm_obs = norm_obs
         self.norm_reward = norm_reward
         self.old_obs = np.array([])
+        self.old_reward = np.array([])
 
     def step_wait(self):
         """
@@ -46,12 +47,13 @@ class VecNormalize(VecEnvWrapper):
         """
         obs, rews, news, infos = self.venv.step_wait()
         self.ret = self.ret * self.gamma + rews
-        self.old_obs = obs
+        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 = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward)
+            rews = self.normalize_reward(rews)
         self.ret[news] = 0
         return obs, rews, news, infos
 
@@ -62,12 +64,31 @@ class VecNormalize(VecEnvWrapper):
         if self.norm_obs:
             if self.training:
                 self.obs_rms.update(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
+            return self.normalize_obs(obs)
         else:
             return obs
 
+    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)
+        return obs
+
+    def normalize_reward(self, reward):
+        if self.norm_reward:
+            return np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward)
+        return reward
+
+    def unnormalize_obs(self, obs):
+        if self.norm_obs:
+            return (obs * np.sqrt(self.obs_rms.var + self.epsilon)) + self.obs_rms.mean
+        return obs
+
+    def unnormalize_reward(self, reward):
+        if self.norm_reward:
+            return reward * np.sqrt(self.ret_rms.var + self.epsilon)
+        return reward
+
     def get_original_obs(self):
         """
         returns the unnormalized observation
@@ -76,6 +97,14 @@ class VecNormalize(VecEnvWrapper):
         """
         return self.old_obs
 
+    def get_original_reward(self):
+        """
+        returns the unnormalized observation
+
+        :return: (numpy float)
+        """
+        return self.old_reward
+
     def reset(self):
         """
         Reset all environments
diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index 1e3b25c..46fe75c 100644
--- a/torchy_baselines/ppo/policies.py
+++ b/torchy_baselines/ppo/policies.py
@@ -1,113 +1,44 @@
 from functools import partial
-from itertools import zip_longest
 
 import torch as th
 import torch.nn as nn
 import numpy as np
 
-from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp
+from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor
 from torchy_baselines.common.distributions import make_proba_distribution,\
     DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution
 
 
-class MlpExtractor(nn.Module):
-    """
-    Constructs an MLP that receives observations as an input and outputs a latent representation for the policy and
-    a value network. The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many
-    of them are shared between the policy network and the value network. It is assumed to be a list with the following
-    structure:
-
-    1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer.
-       If the number of ints is zero, there will be no shared layers.
-    2. An optional dict, to specify the following non-shared layers for the value network and the policy network.
-       It is formatted like ``dict(vf=[<value layer sizes>], pi=[<policy layer sizes>])``.
-       If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed.
-
-    For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value
-    network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec
-    would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128
-    would be specified as [128, 128].
-
-    Adapted from Stable Baselines.
-
-    :param flat_observations: (th.Tensor) The observations to base policy and value function on.
-    :param net_arch: ([int or dict]) The specification of the policy and value networks.
-        See above for details on its formatting.
-    :param activation_fn: (nn.Module) The activation function to use for the networks.
-    :param device: (th.device)
-    """
-    def __init__(self, feature_dim, net_arch, activation_fn, device='cpu'):
-        super(MlpExtractor, self).__init__()
-
-        shared_net, policy_net, value_net = [], [], []
-        policy_only_layers = []  # Layer sizes of the network that only belongs to the policy network
-        value_only_layers = []  # Layer sizes of the network that only belongs to the value network
-        last_layer_dim_shared = feature_dim
-
-        # Iterate through the shared layers and build the shared parts of the network
-        for idx, layer in enumerate(net_arch):
-            if isinstance(layer, int):  # Check that this is a shared layer
-                layer_size = layer
-                # TODO: give layer a meaningful name
-                shared_net.append(nn.Linear(last_layer_dim_shared, layer_size))
-                shared_net.append(activation_fn())
-                last_layer_dim_shared = layer_size
-            else:
-                assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts"
-                if 'pi' in layer:
-                    assert isinstance(layer['pi'], list), "Error: net_arch[-1]['pi'] must contain a list of integers."
-                    policy_only_layers = layer['pi']
-
-                if 'vf' in layer:
-                    assert isinstance(layer['vf'], list), "Error: net_arch[-1]['vf'] must contain a list of integers."
-                    value_only_layers = layer['vf']
-                break  # From here on the network splits up in policy and value network
-
-        last_layer_dim_pi = last_layer_dim_shared
-        last_layer_dim_vf = last_layer_dim_shared
-
-        # Build the non-shared part of the network
-        for idx, (pi_layer_size, vf_layer_size) in enumerate(zip_longest(policy_only_layers, value_only_layers)):
-            if pi_layer_size is not None:
-                assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers."
-                policy_net.append(nn.Linear(last_layer_dim_pi, pi_layer_size))
-                policy_net.append(activation_fn())
-                last_layer_dim_pi = pi_layer_size
-
-            if vf_layer_size is not None:
-                assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers."
-                value_net.append(nn.Linear(last_layer_dim_vf, vf_layer_size))
-                value_net.append(activation_fn())
-                last_layer_dim_vf = vf_layer_size
-
-        # Save dim, used to create the distributions
-        self.latent_dim_pi = last_layer_dim_pi
-        self.latent_dim_vf = last_layer_dim_vf
-
-        # Create networks
-        # If the list of layers is empty, the network will just act as an Identity module
-        self.shared_net = nn.Sequential(*shared_net).to(device)
-        self.policy_net = nn.Sequential(*policy_net).to(device)
-        self.value_net = nn.Sequential(*value_net).to(device)
-
-    def forward(self, features):
-        """
-        :return: (th.Tensor, th.Tensor) latent_policy, latent_value of the specified network.
-            If all layers are shared, then ``latent_policy == latent_value``
-        """
-        shared_latent = self.shared_net(features)
-        return self.policy_net(shared_latent), self.value_net(shared_latent)
-
-
 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 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 adam_epsilon: (float) Small values to avoid NaN in ADAM optimizer
+    :param ortho_init: (bool) Whether to use or not orthogonal initialization
+    :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, 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):
+                 ortho_init=True, use_sde=False,
+                 log_std_init=0.0, full_std=True):
         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], vf=[64])]
+            net_arch = [dict(pi=[64, 64], vf=[64, 64])]
+
         self.net_arch = net_arch
         self.activation_fn = activation_fn
         self.adam_epsilon = adam_epsilon
@@ -124,12 +55,24 @@ class PPOPolicy(BasePolicy):
         self.features_extractor = nn.Flatten()
         self.features_dim = self.obs_dim
         self.log_std_init = log_std_init
+        dist_kwargs = None
+        # Keyword arguments for SDE distribution
+        if use_sde:
+            dist_kwargs = {
+                'full_std': full_std,
+                'squash_output': False,
+                'use_expln': False
+            }
+
         # Action distribution
-        self.action_dist = make_proba_distribution(action_space, use_sde=use_sde)
+        self.action_dist = make_proba_distribution(action_space, use_sde=use_sde, dist_kwargs=dist_kwargs)
 
         self._build(learning_rate)
 
     def reset_noise_net(self):
+        """
+        Sample new weights for the exploration matrix.
+        """
         self.action_dist.sample_weights(self.log_std)
 
     def _build(self, learning_rate):
@@ -147,7 +90,7 @@ class PPOPolicy(BasePolicy):
         # with small initial weight for the output
         if self.ortho_init:
             for module in [self.mlp_extractor, self.action_net, self.value_net]:
-                # Values from stable-baselines check why
+                # Values from stable-baselines, TODO: check why
                 gain = {
                     self.mlp_extractor: np.sqrt(2),
                     self.action_net: 0.01,
@@ -185,15 +128,26 @@ class PPOPolicy(BasePolicy):
         action, _ = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
         return action.detach().cpu().numpy()
 
-    def get_policy_stats(self, obs, action, deterministic=False):
+    def evaluate_actions(self, obs, action, deterministic=False):
+        """
+        Evaluate actions according to the current policy,
+        given the observations.
+
+        :param obs: (th.Tensor)
+        :param action: (th.Tensor)
+        :param deterministic: (bool)
+        :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)
         log_prob = action_distribution.log_prob(action)
         value = self.value_net(latent_vf)
         return value, log_prob, action_distribution.entropy()
 
-    def value_forward(self):
-        pass
+    def value_forward(self, obs):
+        _, latent_vf = self._get_latent(obs)
+        return self.value_net(latent_vf)
 
 
 MlpPolicy = PPOPolicy
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index b0e9de2..44a411a 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -1,6 +1,5 @@
 import os
 import time
-from copy import deepcopy
 
 import gym
 from gym import spaces
@@ -18,7 +17,7 @@ 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 VecNormalize, VecEnvWrapper
+from torchy_baselines.common.vec_env import sync_envs_normalization
 from torchy_baselines.common import logger
 from torchy_baselines.ppo.policies import PPOPolicy
 
@@ -205,7 +204,7 @@ class PPO(BaseRLModel):
                     # Convert discrete action for float to long
                     action = action.long().flatten()
 
-                values, log_prob, entropy = self.policy.get_policy_stats(obs, action)
+                values, log_prob, entropy = self.policy.evaluate_actions(obs, action)
                 values = values.flatten()
                 # Normalize advantage
                 advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8)
@@ -241,7 +240,8 @@ class PPO(BaseRLModel):
                 approx_kl_divs.append(th.mean(old_log_prob - log_prob).detach().cpu().numpy())
 
             if self.target_kl is not None and np.mean(approx_kl_divs) > 1.5 * self.target_kl:
-                print("Early stopping at step {} due to reaching max kl: {:.2f}".format(it, np.mean(approx_kl_divs)))
+                print("Early stopping at step {} due to reaching max kl: {:.2f}".format(gradient_step,
+                                                                                        np.mean(approx_kl_divs)))
                 break
 
         explained_var = explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(),
@@ -294,14 +294,8 @@ class PPO(BaseRLModel):
             # Evaluate agent
             if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
                 timesteps_since_eval %= eval_freq
-                # TODO: move that to the base class
-                # Sync eval env and train env when using VecNormalize
-                env_tmp, eval_env_tmp = self.env, eval_env
-                while isinstance(env_tmp, VecEnvWrapper):
-                    if isinstance(env_tmp, VecNormalize):
-                        eval_env_tmp.obs_rms = deepcopy(env_tmp.obs_rms)
-                    env_tmp = env_tmp.venv
-                    eval_env_tmp.venv
+                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)
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index 5222754..a42a8c9 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -8,6 +8,7 @@ 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
 
 
 class SAC(BaseRLModel):
@@ -54,7 +55,7 @@ class SAC(BaseRLModel):
     """
 
     def __init__(self, policy, env, learning_rate=3e-4, buffer_size=int(1e6),
-                 learning_starts=100, batch_size=64,
+                 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,
@@ -163,7 +164,7 @@ class SAC(BaseRLModel):
 
         for gradient_step in range(gradient_steps):
             # Sample replay buffer
-            replay_data = self.replay_buffer.sample(batch_size)
+            replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
 
             obs, action_batch, next_obs, done, reward = replay_data
 
@@ -257,9 +258,6 @@ class SAC(BaseRLModel):
             self._update_current_progress(self.num_timesteps, total_timesteps)
 
             if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
-                if self.verbose > 1:
-                    print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
-                        self.num_timesteps, episode_num, episode_timesteps, episode_reward))
                 gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps
 
                 self.train(gradient_steps, batch_size=self.batch_size)
@@ -267,6 +265,7 @@ class SAC(BaseRLModel):
             # 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:
diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py
index 4cf32f3..cd72950 100644
--- a/torchy_baselines/td3/policies.py
+++ b/torchy_baselines/td3/policies.py
@@ -2,21 +2,120 @@ 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 StateDependentNoiseDistribution
 
 
 class Actor(BaseNetwork):
-    def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU):
+    """
+    Actor network (policy) for TD3.
+
+    :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 clip_noise: (float) Clip the magnitude of the noise
+    :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.
+    """
+    def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
+                 use_sde=False, log_std_init=-2, clip_noise=None,
+                 lr_sde=3e-4, full_std=False):
         super(Actor, self).__init__()
 
-        # TODO: orthogonal initialization?
-        actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True)
-        self.actor_net = nn.Sequential(*actor_net)
+        self.latent_pi, self.log_std = None, None
+        self.weights_dist, self.exploration_mat = None, None
+        self.use_sde, self.sde_optimizer = use_sde, None
+        self.action_dim = action_dim
+        self.full_std = full_std
 
-    def forward(self, obs):
-        return self.actor_net(obs)
+        if use_sde:
+            latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False)
+            self.latent_pi = nn.Sequential(*latent_pi)
+            # Create state dependent noise matrix (SDE)
+            self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
+                                                               squash_output=False)
+            action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
+                                                                               log_std_init=log_std_init)
+            # Squash output
+            self.actor_net = 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)
+
+    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 _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 evaluate_actions(self, obs, action):
+        """
+        Evaluate actions according to the current policy,
+        given the observations. Only useful when using SDE.
+
+        :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.
+        """
+        with th.no_grad():
+            latent_pi = self.latent_pi(obs)
+        _, distribution = self._get_action_dist_from_latent(latent_pi)
+        log_prob = distribution.log_prob(action)
+        # value = self.value_net(latent_vf)
+        return log_prob, distribution.entropy()
+
+    def reset_noise(self):
+        """
+        Sample new weights for the exploration matrix.
+        """
+        self.action_dist.sample_weights(self.log_std)
+
+    def forward(self, obs, deterministic=True):
+        if self.use_sde:
+            latent_pi = self.latent_pi(obs)
+            if deterministic:
+                return self.actor_net(latent_pi)
+            noise = self.action_dist.get_noise(latent_pi)
+            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
+            # action, _ = self._get_action_dist_from_latent(latent_pi)
+            # return action
+        else:
+            return self.actor_net(obs)
 
 
 class Critic(BaseNetwork):
+    """
+    Critic network for TD3,
+    in fact it represents the action-state value function (Q-value function)
+
+    :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__()
@@ -38,11 +137,25 @@ class Critic(BaseNetwork):
 
 
 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 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
+    """
     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=-2,
+                 clip_noise=None, lr_sde=3e-4):
         super(TD3Policy, self).__init__(observation_space, action_space, device)
 
+        # Default network architecture, from the original paper
         if net_arch is None:
             net_arch = [400, 300]
 
@@ -56,8 +169,16 @@ class TD3Policy(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_kwargs['clip_noise'] = clip_noise
+        self.actor_kwargs['lr_sde'] = lr_sde
+
         self.actor, self.actor_target = None, None
         self.critic, self.critic_target = None, None
+        self.use_sde = use_sde
+        self.log_std_init = log_std_init
         self._build(learning_rate)
 
     def _build(self, learning_rate):
@@ -71,14 +192,17 @@ 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))
 
+    def reset_noise(self):
+        return self.actor.reset_noise()
+
     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)
 
-    def forward(self, obs):
-        return self.actor(obs)
+    def forward(self, obs, deterministic=True):
+        return self.actor(obs, deterministic=deterministic)
 
 
 MlpPolicy = TD3Policy
diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py
index 645ecf2..c8847a5 100644
--- a/torchy_baselines/td3/td3.py
+++ b/torchy_baselines/td3/td3.py
@@ -8,6 +8,7 @@ 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
 
 
 class TD3(BaseRLModel):
@@ -37,6 +38,11 @@ class TD3(BaseRLModel):
     :param target_policy_noise: (float) Standard deviation of gaussian noise added to target policy
         (smoothing noise)
     :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_max_grad_norm: (float)
+    :param sde_ent_coef: (float)
+    :param sde_log_std_scheduler: (callable)
     :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
@@ -51,6 +57,7 @@ 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,
                  tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0,
                  seed=0, device='auto', _init_setup_model=True):
 
@@ -58,7 +65,6 @@ class TD3(BaseRLModel):
                                   create_eval_env=create_eval_env, seed=seed)
 
         self.buffer_size = buffer_size
-        # TODO: accept callables
         self.learning_rate = learning_rate
         self.learning_starts = learning_starts
         self.train_freq = train_freq
@@ -72,6 +78,12 @@ class TD3(BaseRLModel):
         self.target_noise_clip = target_noise_clip
         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
+
         if _init_setup_model:
             self._setup_model()
 
@@ -81,7 +93,7 @@ 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, 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()
 
@@ -91,12 +103,12 @@ class TD3(BaseRLModel):
         self.critic = self.policy.critic
         self.critic_target = self.policy.critic_target
 
-    def select_action(self, observation):
+    def select_action(self, observation, deterministic=True):
         # Normally not needed
         observation = np.array(observation)
         with th.no_grad():
             observation = th.FloatTensor(observation.reshape(1, -1)).to(self.device)
-            return self.actor(observation).cpu().numpy()
+            return self.actor(observation, deterministic=deterministic).cpu().numpy()
 
     def predict(self, observation, state=None, mask=None, deterministic=True):
         """
@@ -108,7 +120,7 @@ class TD3(BaseRLModel):
         :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)
         """
-        return self.unscale_action(self.select_action(observation))
+        return self.unscale_action(self.select_action(observation, deterministic=deterministic))
 
     def train_critic(self, gradient_steps=1, batch_size=100, replay_data=None, tau=0.0):
         # Update optimizer learning rate
@@ -117,7 +129,7 @@ class TD3(BaseRLModel):
         for gradient_step in range(gradient_steps):
             # Sample replay buffer
             if replay_data is None:
-                obs, action, next_obs, done, reward = self.replay_buffer.sample(batch_size)
+                obs, action, next_obs, done, reward = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
             else:
                 obs, action, next_obs, done, reward = replay_data
 
@@ -158,7 +170,7 @@ class TD3(BaseRLModel):
         for gradient_step in range(gradient_steps):
             # Sample replay buffer
             if replay_data is None:
-                obs, _, next_obs, done, reward = self.replay_buffer.sample(batch_size)
+                obs, _, next_obs, done, reward = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
             else:
                 obs, _, next_obs, done, reward = replay_data
 
@@ -183,13 +195,53 @@ class TD3(BaseRLModel):
         for gradient_step in range(gradient_steps):
 
             # Sample replay buffer
-            replay_data = self.replay_buffer.sample(batch_size)
+            replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
             self.train_critic(replay_data=replay_data)
 
             # Delayed policy updates
             if gradient_step % policy_delay == 0:
                 self.train_actor(replay_data=replay_data, tau_actor=self.tau, tau_critic=self.tau)
 
+    def train_sde(self):
+        # Update optimizer learning rate
+        # self._update_learning_rate(self.policy.optimizer)
+
+        # Unpack
+        obs, action, returns = [self.rollout_data[key] for key in ['observations', 'actions', 'returns']]
+
+        # TODO: avoid second computation of everything because of the gradient
+        log_prob, entropy = self.actor.evaluate_actions(obs, action)
+
+        # 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))
+
+        policy_loss = -(returns * log_prob).mean()
+
+        # Entropy loss favor exploration
+        entropy_loss = -th.mean(entropy)
+
+        loss = policy_loss + self.sde_ent_coef * entropy_loss
+
+        # Optimization step
+        self.actor.sde_optimizer.zero_grad()
+        loss.backward()
+
+        assert not th.isnan(log_prob).any(), log_prob
+        assert not th.isnan(entropy).any()
+        assert not th.isnan(self.actor.log_std.grad).any()
+        assert not th.isnan(self.actor.log_std).any()
+
+        # Clip grad norm
+        th.nn.utils.clip_grad_norm_([self.actor.log_std], self.sde_max_grad_norm)
+        self.actor.sde_optimizer.step()
+
+        del self.rollout_data
+
     def learn(self, total_timesteps, callback=None, log_interval=4,
               eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="TD3", reset_num_timesteps=True):
 
@@ -219,9 +271,15 @@ class TD3(BaseRLModel):
             self._update_current_progress(self.num_timesteps, total_timesteps)
 
             if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
-                if self.verbose > 1:
-                    print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
-                        self.num_timesteps, episode_num, episode_timesteps, episode_reward))
+
+                if self.use_sde:
+                    if self.sde_log_std_scheduler is not None:
+                        # Call the scheduler
+                        value = self.sde_log_std_scheduler(self._current_progress)
+                        self.actor.log_std.data = th.ones_like(self.actor.log_std) * value
+                    else:
+                        # On-policy gradient
+                        self.train_sde()
 
                 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)
@@ -229,6 +287,7 @@ class TD3(BaseRLModel):
             # 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:
@@ -241,7 +300,7 @@ class TD3(BaseRLModel):
         """
         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 {"actor": self.actor.optimizer.state_dict(), "critic": self.critic.optimizer.state_dict()}