saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-15 22:10:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add first draft of SDE

Browse source
This commit is contained in:
Antonin Raffin 2019-10-28 18:24:13 +01:00
parent df1e7aa000
commit c15b4bda1e
5 changed files with 156 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
tests/test_sde.py Normal file
View file

@ -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)

9
torchy_baselines/a2c/a2c.py
View file

@ -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)

67
torchy_baselines/common/distributions.py
View file

@ -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)

37
torchy_baselines/ppo/policies.py
View file

@ -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()

12
torchy_baselines/ppo/ppo.py
View file

@ -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():