Add first draft of SDE

tests/test_sde.py

@@ -0,0 +1,47 @@
+import pytest
+
+import torch as th
+from torch.distributions import Normal
+
+from torchy_baselines import A2C
+
+
+def test_state_dependent_exploration():
+    state_dim = 3
+    # TODO: fix for action_dim > 1
+    action_dim = 1
+    sigma = th.ones(state_dim, action_dim, requires_grad=True)
+
+    # log_sigma = th.ones(2, 1, requires_grad=True)
+
+    # 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 = weights_dist.rsample()
+    state = th.rand(1, state_dim)
+    # state = (th.ones(state_dim,) * 2).view(1, -1)
+    mu = th.ones(action_dim)
+    # print(weights.shape, state.shape)
+    noise = th.mm(state, weights)
+    # variance = th.mm(state ** 2, th.exp(log_sigma) ** 2)
+    variance = th.mm(state ** 2, sigma ** 2)
+    action_dist = Normal(mu, th.sqrt(variance))
+
+    loss = action_dist.log_prob((mu + noise).detach()).mean()
+    loss.backward()
+
+    # From Rueckstiess paper
+    grad = th.zeros_like(sigma)
+    for j in range(action_dim):
+        for i in range(state_dim):
+            grad[i, j] = ((noise[:, j] ** 2 - variance[:, j]) / (variance[:, j] ** 2)) * (state[:, i] ** 2 * sigma[i, j])
+
+    # sigma.grad should be equal to grad
+    assert sigma.grad.allclose(grad)
+
+
+@pytest.mark.parametrize("model_class", [A2C])
+def test_state_dependent_noise(model_class):
+    model = model_class('MlpPolicy', 'Pendulum-v0', n_steps=200, use_sde=True, verbose=1, create_eval_env=True)
+    model.learn(total_timesteps=int(1e6), log_interval=10, eval_freq=10000)

torchy_baselines/a2c/a2c.py

@@ -5,6 +5,7 @@ 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
 
 
 class A2C(PPO):
@@ -30,6 +31,8 @@ class A2C(PPO):
     :param rms_prop_eps: (float) RMSProp epsilon. It stabilizes square root computation in denominator
         of RMSProp update
     :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 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,7 +48,7 @@ class A2C(PPO):
     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,
+                 rms_prop_eps=1e-5, use_rms_prop=True, use_sde=False,
                  normalize_advantage=False, tensorboard_log=None, create_eval_env=False,
                  policy_kwargs=None, verbose=0, seed=0, device='auto',
                  _init_setup_model=True):
@@ -53,7 +56,7 @@ 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,
+                                  vf_coef=vf_coef, max_grad_norm=max_grad_norm, use_sde=use_sde,
                                   tensorboard_log=tensorboard_log, policy_kwargs=policy_kwargs,
                                   verbose=verbose, device=device, create_eval_env=create_eval_env,
                                   seed=seed, _init_setup_model=False)
@@ -73,6 +76,8 @@ class A2C(PPO):
                                                      eps=self.rms_prop_eps, weight_decay=0)
 
     def train(self, gradient_steps, batch_size=None):
+        if self.use_sde:
+            logger.logkv("noise net std", th.exp(self.policy.log_std).mean().item())
 
         # Update optimizer learning rate
         self._update_learning_rate(self.policy.optimizer)

torchy_baselines/common/distributions.py

@@ -165,15 +165,80 @@ class CategoricalDistribution(Distribution):
         return log_prob
 
 
-def make_proba_distribution(action_space):
+class StateDependentNoiseDistribution(Distribution):
+    def __init__(self, features_dim, action_dim):
+        super(StateDependentNoiseDistribution, self).__init__()
+        self.distribution = None
+        self.action_dim = action_dim
+        self.features_dim = features_dim
+        self.mean_actions = None
+        self.log_std = None
+        self.weights_dist = None
+        self.noise_weights = None
+
+    @staticmethod
+    def get_std(log_std):
+        # TODO: use expln instead of exp only to avoid sigma growing too fast
+        return th.exp(log_std)
+
+    def sample_weights(self, log_std):
+        self.weights_dist = Normal(th.zeros_like(log_std), self.get_std(log_std))
+        self.noise_weights = self.weights_dist.rsample()
+
+    def proba_distribution_net(self, latent_dim, log_std_init=0.0):
+        mean_actions = nn.Linear(latent_dim, self.action_dim)
+        log_std = nn.Parameter(th.zeros(self.features_dim, self.action_dim))
+        self.sample_weights(log_std)
+        return mean_actions, log_std
+
+    def proba_distribution(self, mean_actions, log_std, observations, deterministic=False):
+        variance = th.mm(observations ** 2, self.get_std(log_std) ** 2)
+        self.distribution = Normal(mean_actions, th.sqrt(variance))
+
+        if deterministic:
+            action = self.mode()
+        else:
+            action = self.sample(observations)
+        return action, self
+
+    def mode(self):
+        return self.distribution.mean
+
+    def sample(self, observations):
+        noise = th.mm(observations, self.noise_weights)
+        return self.distribution.mean + noise
+
+    def entropy(self):
+        return self.distribution.entropy()
+
+    def log_prob_from_params(self, mean_actions, log_std, observations):
+        action, _ = self.proba_distribution(mean_actions, log_std, observations)
+        log_prob = self.log_prob(action)
+        return action, log_prob
+
+    def log_prob(self, action):
+        log_prob = self.distribution.log_prob(action)
+        if len(log_prob.shape) > 1:
+            log_prob = log_prob.sum(axis=1)
+        else:
+            log_prob = log_prob.sum()
+        return log_prob
+
+
+def make_proba_distribution(action_space, features_dim=None, use_sde=False):
     """
     Return an instance of Distribution for the correct type of action space
 
     :param action_space: (Gym Space) the input action space
+    :param feature_dim: (int) Dimension of the feature vector
+    :param use_sde: (bool) Force the use of StateDependentNoiseDistribution
+        instead of DiagGaussianDistribution
     :return: (Distribution) the approriate Distribution object
     """
     if isinstance(action_space, spaces.Box):
         assert len(action_space.shape) == 1, "Error: the action space must be a vector"
+        if use_sde:
+            return StateDependentNoiseDistribution(features_dim, action_space.shape[0])
         return DiagGaussianDistribution(action_space.shape[0])
     elif isinstance(action_space, spaces.Discrete):
         return CategoricalDistribution(action_space.n)

torchy_baselines/ppo/policies.py

@@ -6,7 +6,8 @@ import torch.nn as nn
 import numpy as np
 
 from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp
-from torchy_baselines.common.distributions import make_proba_distribution, DiagGaussianDistribution, CategoricalDistribution
+from torchy_baselines.common.distributions import make_proba_distribution,\
+    DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution
 
 
 class MlpExtractor(nn.Module):
@@ -101,7 +102,8 @@ class MlpExtractor(nn.Module):
 class PPOPolicy(BasePolicy):
     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):
+                 activation_fn=nn.Tanh, adam_epsilon=1e-5,
+                 ortho_init=True, use_sde=False):
         super(PPOPolicy, self).__init__(observation_space, action_space, device)
         self.obs_dim = self.observation_space.shape[0]
         if net_arch is None:
@@ -118,20 +120,31 @@ class PPOPolicy(BasePolicy):
         }
         self.shared_net = None
         self.pi_net, self.vf_net = None, None
-        # Action distribution
-        self.action_dist = make_proba_distribution(action_space)
         # In the future, feature_extractor will be replaced with a CNN
         self.features_extractor = nn.Flatten()
         self.features_dim = self.obs_dim
+        # Action distribution
+        self.action_dist = make_proba_distribution(action_space, self.features_dim, use_sde=use_sde)
+
         self._build(learning_rate)
 
+    def reset_noise_net(self):
+        self.action_dist.sample_weights(self.log_std)
+        # weights_dist = Normal(th.zeros_like(self.noise_log_sigma), th.exp(self.noise_log_sigma))
+        # self.noise_net = weights_dist.rsample()
+        # noise = th.mm(state, weights)
+        # variance = th.mm(state ** 2, sigma ** 2)
+        # action_dist = Normal(mu, th.sqrt(variance))
+        # # action_dist.log_prob((mu + noise).detach())
+        # action_dist.log_prob(action)
+        # # action_dist = Normal(mu_j + noise_j, sum of s_i * sigma_ij)
+        # # log_prob = distribution.log_prob(self.noise_net)
+
     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)
 
-        # self.action_net = nn.Linear(self.net_arch[-1], self.action_dim)
+        if isinstance(self.action_dist, (DiagGaussianDistribution, StateDependentNoiseDistribution)):
-        # self.log_std = nn.Parameter(th.zeros(self.action_dim))
-        if isinstance(self.action_dist, DiagGaussianDistribution):
             self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi)
         elif isinstance(self.action_dist, CategoricalDistribution):
             self.action_net = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi)
@@ -155,28 +168,30 @@ class PPOPolicy(BasePolicy):
             obs = th.FloatTensor(obs).to(self.device)
         latent_pi, latent_vf = 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, obs, 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))
 
-    def _get_action_dist_from_latent(self, latent, deterministic=False):
+    def _get_action_dist_from_latent(self, latent, obs, deterministic=False):
         mean_actions = self.action_net(latent)
         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):
             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, obs, 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)
+        action, _ = self._get_action_dist_from_latent(latent_pi, obs, deterministic=deterministic)
         return action.detach().cpu().numpy()
 
     def get_policy_stats(self, obs, action):
         latent_pi, latent_vf = self._get_latent(obs)
-        _, action_distribution = self._get_action_dist_from_latent(latent_pi)
+        _, action_distribution = self._get_action_dist_from_latent(latent_pi, obs)
         log_prob = action_distribution.log_prob(action)
         value = self.value_net(latent_vf)
         return value, log_prob, action_distribution.entropy()

torchy_baselines/ppo/ppo.py

@@ -52,6 +52,8 @@ class PPO(BaseRLModel):
     :param ent_coef: (float) Entropy coefficient for the loss calculation
     :param vf_coef: (float) Value function coefficient for the loss calculation
     :param max_grad_norm: (float) The maximum value for the gradient clipping
+    :param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
+        instead of action noise exploration (default: False)
     :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)
@@ -70,7 +72,7 @@ class PPO(BaseRLModel):
     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,
+                 ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5, use_sde=False,
                  target_kl=None, tensorboard_log=None, create_eval_env=False,
                  policy_kwargs=None, verbose=0, seed=0, device='auto',
                  _init_setup_model=True):
@@ -94,6 +96,7 @@ 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()
@@ -116,7 +119,8 @@ class PPO(BaseRLModel):
         self.rollout_buffer = RolloutBuffer(self.n_steps, state_dim, action_dim, self.device,
                                             gamma=self.gamma, gae_lambda=self.gae_lambda, n_envs=self.n_envs)
         self.policy = self.policy(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.clip_range = get_schedule_fn(self.clip_range)
@@ -150,6 +154,10 @@ class PPO(BaseRLModel):
 
         n_steps = 0
         rollout_buffer.reset()
+        # Sample new weights for the state dependent exploration
+        # TODO: ensure episodic setting?
+        if self.use_sde:
+            self.policy.reset_noise_net()
 
         while n_steps < n_rollout_steps:
             with th.no_grad():