diff --git a/README.md b/README.md
index 0b19832..d8f5fa6 100644
--- a/README.md
+++ b/README.md
@@ -14,24 +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:
-- 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
-- CNN policies + normalization
-- tensorboard support
-- DQN
-- TRPO
-- ACER
-- DDPG
-- HER -> use stable-baselines because does not depends on tf?
+- cf github Roadmap
diff --git a/setup.py b/setup.py
index 640112e..037e4bd 100644
--- a/setup.py
+++ b/setup.py
@@ -23,6 +23,12 @@ setup(name='torchy_baselines',
             'sphinx',
             'sphinx-autobuild',
             'sphinx-rtd-theme'
+        ],
+        'extra': [
+            # For render
+            'opencv-python',
+            # For reading logs
+            'pandas'
         ]
       },
       description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.',
@@ -34,7 +40,7 @@ setup(name='torchy_baselines',
       license="MIT",
       long_description="",
       long_description_content_type='text/markdown',
-      version="0.0.6a",
+      version="0.0.8a0",
       )
 
 # python setup.py sdist
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 47651e4..7e24cda 100644
--- a/tests/test_distributions.py
+++ b/tests/test_distributions.py
@@ -1,8 +1,10 @@
-import numpy as np
+import pytest
 import torch as th
 
-from torchy_baselines.common.distributions import DiagGaussianDistribution, SquashedDiagGaussianDistribution,\
-    CategoricalDistribution, TanhBijector
+from torchy_baselines.common.distributions import DiagGaussianDistribution, TanhBijector, \
+    StateDependentNoiseDistribution
+from torchy_baselines.common.utils import set_random_seed
+
 
 # TODO: more tests for the other distributions
 def test_bijector():
@@ -18,3 +20,51 @@ 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)
+
+
+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)
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/tests/test_run.py b/tests/test_run.py
index a5569e5..e551a8d 100644
--- a/tests/test_run.py
+++ b/tests/test_run.py
@@ -37,9 +37,10 @@ def test_onpolicy(model_class, env_id):
     os.remove("test_save.zip")
 
 
-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)
     model.save("test_save")
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 09b48e8..497f398 100644
--- a/tests/test_sde.py
+++ b/tests/test_sde.py
@@ -1,69 +1,65 @@
 import pytest
-
-import gym
 import torch as th
 from torch.distributions import Normal
 
-from torchy_baselines import A2C, TD3
-from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize
-from torchy_baselines.common.monitor import Monitor
+from torchy_baselines import A2C, TD3, SAC
 
 
-def test_state_dependent_exploration():
+def test_state_dependent_exploration_grad():
     """
     Check that the gradient correspond to the expected one
     """
     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])
-def test_state_dependent_noise(model_class):
-    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)
-    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):
+@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)
 
 
@@ -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/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/__init__.py b/torchy_baselines/__init__.py
index b8383db..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.6a"
+__version__ = "0.0.8a0"
diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py
index a06074b..6fa667c 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
 
 
@@ -34,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
@@ -45,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):
@@ -57,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)
@@ -113,6 +114,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/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 dfa9ab1..7a29847 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
 
 
@@ -37,12 +38,17 @@ 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)
+    :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):
+                 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:
@@ -70,7 +76,9 @@ 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
+        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
@@ -217,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):
@@ -250,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):
@@ -280,12 +271,12 @@ 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
+        raise NotImplementedError()
 
     @abstractmethod
     def predict(self, observation, state=None, mask=None, deterministic=False):
@@ -298,13 +289,14 @@ 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=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()
@@ -342,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
@@ -357,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):
@@ -395,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:
@@ -503,7 +497,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', 'values']}
 
         while total_steps < n_steps or total_episodes < n_episodes:
             done = False
@@ -512,7 +507,13 @@ 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
+                # 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()])
@@ -562,6 +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())
+                    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
@@ -580,6 +583,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()
 
@@ -603,19 +607,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
 
@@ -652,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():
@@ -664,7 +673,7 @@ class BaseRLModel(object):
 
         :return: ([str]) List of parameters that should be excluded from save
         """
-        return ["env", "eval_env", "replay_buffer", "rollout_buffer"]
+        return ["env", "eval_env", "replay_buffer", "rollout_buffer", "_vec_normalize_env"]
 
     def save(self, path, exclude=None, include=None):
         """
@@ -694,3 +703,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/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 8478871..9d6189c 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -1,8 +1,6 @@
-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
 
 
@@ -14,17 +12,8 @@ class Distribution(object):
         """
         returns the log likelihood
 
-        :param x: (str) the labels of each index
-        :return: ([float]) 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
+        :param x: (object) the taken action
+        :return: (th.Tensor) The log likelihood of the distribution
         """
         raise NotImplementedError
 
@@ -32,7 +21,7 @@ class Distribution(object):
         """
         Returns shannon's entropy of the probability
 
-        :return: (float) the entropy
+        :return: (th.Tensor) the entropy
         """
         raise NotImplementedError
 
@@ -40,7 +29,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
 
@@ -52,6 +41,7 @@ class DiagGaussianDistribution(Distribution):
 
     :param action_dim: (int)  Number of continuous actions
     """
+
     def __init__(self, action_dim):
         super(DiagGaussianDistribution, self).__init__()
         self.distribution = None
@@ -138,6 +128,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)
@@ -185,6 +176,7 @@ class CategoricalDistribution(Distribution):
 
     :param action_dim: (int) Number of discrete actions
     """
+
     def __init__(self, action_dim):
         super(CategoricalDistribution, self).__init__()
         self.distribution = None
@@ -243,23 +235,28 @@ 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
         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
@@ -275,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)
@@ -286,57 +284,65 @@ 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):
+    def sample_weights(self, log_std, batch_size=1):
         """
         Sample weights for the noise exploration matrix,
         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)
         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):
+    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)
 
         :param mean_actions: (th.Tensor)
         :param log_std: (th.Tensor)
+        :param latent_sde: (th.Tensor)
         :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):
@@ -345,11 +351,20 @@ 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()
+        # 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)
+        noise = th.bmm(latent_sde, self.exploration_matrices)
+        return noise.squeeze(1)
 
-    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)
@@ -359,8 +374,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
 
@@ -391,11 +406,13 @@ 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
 
-    def forward(self, x):
+    @staticmethod
+    def forward(x):
         return th.tanh(x)
 
     @staticmethod
@@ -422,7 +439,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/common/evaluation.py b/torchy_baselines/common/evaluation.py
index 3d76d9f..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), n_steps
+    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/common/logger.py b/torchy_baselines/common/logger.py
index 5b59b76..a909285 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
@@ -515,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='#')
diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py
index 4ac7e07..f536624 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]))
@@ -75,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/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
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/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):
         """
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')
diff --git a/torchy_baselines/common/vec_env/vec_normalize.py b/torchy_baselines/common/vec_env/vec_normalize.py
index 0b9797f..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,
+            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):
         """
diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py
index 46fe75c..5dd746f 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
+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.
@@ -25,16 +26,24 @@ 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.
+    :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):
+                 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]
 
-
         # Default network architecture, from stable-baselines
         if net_arch is None:
             net_arch = [dict(pi=[64, 64], vf=[64, 64])]
@@ -60,30 +69,49 @@ 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
             }
 
+        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)
 
         self._build(learning_rate)
 
-    def reset_noise_net(self):
+    def reset_noise(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,
                                           activation_fn=self.activation_fn, device=self.device)
 
-        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,
+        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.activation_fn)
+
+        if isinstance(self.action_dist, DiagGaussianDistribution):
+            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 = 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
@@ -102,30 +130,39 @@ 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):
             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):
-            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 +176,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)
 
 
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index 44a411a..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
 
@@ -43,7 +41,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),
@@ -54,6 +53,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)
@@ -68,17 +69,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
@@ -96,7 +97,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()
@@ -157,9 +157,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()
+            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()
@@ -204,6 +208,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
@@ -291,27 +301,20 @@ 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
 
     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()}
 
diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index b61ae50..12b23cc 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -1,46 +1,160 @@
 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.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
 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):
-    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.
+    :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_expln=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)
+        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
 
-        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:
+            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=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,
+                                                                            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))
+            # 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)
+            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, 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, 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)
-        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_pi, latent_sde = self._get_latent(obs)
+
+        if self.use_sde:
+            mean_actions, log_std = self.mu(latent_pi), self.log_std
+        else:
+            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_sde
 
     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_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)
+        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_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_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__()
@@ -62,9 +176,28 @@ 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_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.
+    :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.
+    :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):
+                 activation_fn=nn.ReLU, use_sde=False,
+                 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:
@@ -80,6 +213,14 @@ class SACPolicy(BasePolicy):
             'net_arch': self.net_arch,
             'activation_fn': self.activation_fn
         }
+        self.actor_kwargs = self.net_args.copy()
+        sde_kwargs = {
+            'use_sde': use_sde,
+            'log_std_init': log_std_init,
+            'sde_net_arch': sde_net_arch,
+            'use_expln': use_expln
+        }
+        self.actor_kwargs.update(sde_kwargs)
         self.actor, self.actor_target = None, None
         self.critic, self.critic_target = None, None
 
@@ -95,11 +236,14 @@ 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)
 
+    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 a42a8c9..f5a18d2 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -1,14 +1,11 @@
-import time
-
 import torch as th
 import torch.nn.functional as F
 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
 
 
 class SAC(BaseRLModel):
@@ -44,6 +41,10 @@ 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 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
@@ -59,12 +60,14 @@ 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, 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)
+                                  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
@@ -85,6 +88,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()
@@ -122,11 +126,12 @@ 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_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()
 
@@ -168,12 +173,21 @@ class SAC(BaseRLModel):
 
             obs, action_batch, next_obs, done, reward = replay_data
 
+            # 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
+            # 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()
+
             # 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
@@ -189,10 +203,11 @@ 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():
+                # 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
                 target_q1, target_q2 = self.critic_target(next_obs, next_action)
                 target_q = th.min(target_q1, target_q2)
@@ -228,6 +243,13 @@ 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())
+        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",
               reset_num_timesteps=True):
@@ -262,15 +284,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, _ = 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)))
+            timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
+                                                     timesteps_since_eval, deterministic=True)
 
         return self
 
@@ -278,7 +293,7 @@ class SAC(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
         """
         opt_dict = {"actor": self.actor.optimizer.state_dict(), "critic": self.critic.optimizer.state_dict()}
         if self.ent_coef_optimizer is not None:
diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py
index cd72950..f016dab 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
+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 +20,16 @@ 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.
+    :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=-2, clip_noise=None,
-                 lr_sde=3e-4, full_std=False):
+                 use_sde=False, log_std_init=-3, clip_noise=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
@@ -30,23 +37,33 @@ 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)
+            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,
                                                                                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 +77,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):
         """
@@ -71,39 +97,38 @@ 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.
         """
-        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()
 
     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)
 
     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):
@@ -136,23 +161,50 @@ 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.
 
     :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.
     :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.
+    :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=-2,
-                 clip_noise=None, lr_sde=3e-4):
+                 activation_fn=nn.ReLU, use_sde=False, log_std_init=-3,
+                 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
@@ -170,14 +222,21 @@ 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,
+            'use_expln': use_expln
+        }
+        self.actor_kwargs.update(sde_kwargs)
 
         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)
 
@@ -192,6 +251,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 c8847a5..d9ce51b 100644
--- a/torchy_baselines/td3/td3.py
+++ b/torchy_baselines/td3/td3.py
@@ -1,14 +1,10 @@
-import time
-
 import torch as th
 import torch.nn.functional as F
 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
 
 
 class TD3(BaseRLModel):
@@ -40,6 +36,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 +55,13 @@ 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)
+                                  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
@@ -79,10 +78,11 @@ 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
+        self.sde_vf = None
 
         if _init_setup_model:
             self._setup_model()
@@ -93,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()
 
@@ -102,6 +103,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
@@ -207,25 +209,27 @@ 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()
@@ -284,15 +288,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