saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-29 20:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add SDE support for SAC

Browse source
This commit is contained in:
Antonin Raffin 2019-11-26 15:26:12 +01:00
parent d26fcf4566
commit 5483e02d1a
6 changed files with 95 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
tests/test_sde.py
View file

@ -4,7 +4,7 @@ import gym
import torch as th
from torch.distributions import Normal
from torchy_baselines import A2C, TD3
from torchy_baselines import A2C, TD3, SAC
from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize
from torchy_baselines.common.monitor import Monitor
@ -73,7 +73,7 @@ def test_state_dependent_noise(model_class, sde_net_arch):
model.learn(total_timesteps=int(1000), log_interval=5, eval_freq=500, eval_env=eval_env)
@pytest.mark.parametrize("model_class", [TD3])
@pytest.mark.parametrize("model_class", [TD3, SAC])
def test_state_dependent_offpolicy_noise(model_class):
model = model_class('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True,
verbose=1, policy_kwargs=dict(log_std_init=-2))

10
torchy_baselines/common/base_class.py
View file

@ -33,12 +33,14 @@ 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)
"""
__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):
if isinstance(policy, str) and policy_base is not None:
self.policy = get_policy_from_name(policy_base, policy)
else:
@ -67,7 +69,8 @@ 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
# Track the training progress (from 1 to 0)
# this is used to update the learning rate
self._current_progress = 1
@ -395,7 +398,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']}
while total_steps < n_steps or total_episodes < n_episodes:
done = False

87
torchy_baselines/sac/policies.py
View file

@ -2,7 +2,7 @@ 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.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
# CAP the standard deviation of the actor
LOG_STD_MAX = 2
@ -10,33 +10,83 @@ LOG_STD_MIN = -20
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.
"""
def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
use_sde=False, log_std_init=-3, full_std=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)
self.use_sde = use_sde
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:
# TODO: check for the learn_features
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
learn_features=False, squash_output=True)
self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
log_std_init=log_std_init)
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):
"""
Sample new weights for the exploration matrix, when using SDE.
"""
self.action_dist.sample_weights(self.log_std)
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
if self.use_sde:
mean_actions, log_std = self.mu(latent), self.log_std
else:
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
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 = 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, 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 = 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)
else:
action, log_prob = self.action_dist.log_prob_from_params(mean_actions, log_std)
return action, log_prob
@ -64,7 +114,7 @@ class Critic(BaseNetwork):
class SACPolicy(BasePolicy):
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):
super(SACPolicy, self).__init__(observation_space, action_space, device)
if net_arch is None:
@ -80,6 +130,9 @@ class SACPolicy(BasePolicy):
'net_arch': self.net_arch,
'activation_fn': self.activation_fn
}
self.actor_kwargs = self.net_args.copy()
self.actor_kwargs['use_sde'] = use_sde
self.actor_kwargs['log_std_init'] = log_std_init
self.actor, self.actor_target = None, None
self.critic, self.critic_target = None, None
@ -95,7 +148,7 @@ 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)

15
torchy_baselines/sac/sac.py
View file

@ -44,6 +44,8 @@ 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 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
@ -58,12 +60,12 @@ 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, 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)
self.learning_rate = learning_rate
self.target_entropy = target_entropy
@ -124,8 +126,8 @@ class SAC(BaseRLModel):
self.ent_coef = float(self.ent_coef)
self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
self.policy = self.policy(self.observation_space, self.action_space,
self.learning_rate, device=self.device, **self.policy_kwargs)
self.policy = self.policy(self.observation_space, self.action_space, learning_rate=self.learning_rate,
use_sde=self.use_sde, device=self.device, **self.policy_kwargs)
self.policy = self.policy.to(self.device)
self._create_aliases()
@ -167,6 +169,11 @@ class SAC(BaseRLModel):
obs, action_batch, next_obs, done, reward = replay_data
# TODO: check if there is another way to fix pytorch complain
# if we don't sample the weights again
# (Trying to backward through the graph a second time)
if self.use_sde:
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)

6
torchy_baselines/td3/policies.py
View file

@ -21,7 +21,7 @@ class Actor(BaseNetwork):
for the std instead of only (n_features,) when using SDE.
"""
def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
use_sde=False, log_std_init=-2, clip_noise=None,
use_sde=False, log_std_init=-3, clip_noise=None,
lr_sde=3e-4, full_std=False):
super(Actor, self).__init__()
@ -84,7 +84,7 @@ class Actor(BaseNetwork):
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)
@ -151,7 +151,7 @@ class TD3Policy(BasePolicy):
"""
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,
activation_fn=nn.ReLU, use_sde=False, log_std_init=-3,
clip_noise=None, lr_sde=3e-4):
super(TD3Policy, self).__init__(observation_space, action_space, device)

4
torchy_baselines/td3/td3.py
View file

@ -62,7 +62,7 @@ class TD3(BaseRLModel):
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)
self.buffer_size = buffer_size
self.learning_rate = learning_rate
@ -79,10 +79,10 @@ 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
if _init_setup_model:
self._setup_model()