diff --git a/tests/test_sde.py b/tests/test_sde.py new file mode 100644 index 0000000..7874ae9 --- /dev/null +++ b/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) diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 6ee6f4a..3aa1589 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/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) diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index a384ce3..d420e03 100644 --- a/torchy_baselines/common/distributions.py +++ b/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) diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index e973858..08dd7dd 100644 --- a/torchy_baselines/ppo/policies.py +++ b/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) - # self.log_std = nn.Parameter(th.zeros(self.action_dim)) - if isinstance(self.action_dist, DiagGaussianDistribution): + 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) 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() diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 13a1634..7fa318d 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/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():