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

Unify A2C and TD3 SDE implementation

Browse source
This commit is contained in:
Antonin Raffin 2019-11-25 13:19:33 +01:00
parent c56865e10d
commit 5bbb14188d
4 changed files with 70 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
torchy_baselines/common/base_class.py
View file

@ -488,7 +488,7 @@ class BaseRLModel(object):
logger.logkv('time_elapsed', int(time.time() - self.start_time))
logger.logkv("total timesteps", num_timesteps)
if self.use_sde:
logger.logkv("std", th.exp(self.actor.log_std).mean().item())
logger.logkv("std", (self.actor.get_std()).mean().item())
logger.dumpkvs()
mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0

42
torchy_baselines/common/distributions.py
View file

@ -234,6 +234,8 @@ class StateDependentNoiseDistribution(Distribution):
compute the log probabilty of an action with that noise.
:param action_dim: (int) Number of continuous actions
:param full_std: (bool) Whether to use (n_features x n_actions) parameters
for the std instead of only (n_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.
@ -241,16 +243,19 @@ class StateDependentNoiseDistribution(Distribution):
this allows to ensure boundaries.
:param epsilon: (float) small value to avoid NaN due to numerical imprecision.
"""
def __init__(self, action_dim, use_expln=False,
def __init__(self, action_dim, full_std=True, use_expln=False,
squash_output=False, epsilon=1e-6):
super(StateDependentNoiseDistribution, self).__init__()
self.distribution = None
self.action_dim = action_dim
self.latent_dim = None
self.mean_actions = None
self.log_std = None
self.weights_dist = None
self.exploration_mat = None
self.use_expln = use_expln
self.full_std = full_std
self.epsilon = epsilon
if squash_output:
print("== Using TanhBijector ===")
self.bijector = TanhBijector(epsilon)
@ -269,12 +274,17 @@ class StateDependentNoiseDistribution(Distribution):
# From SDE paper, it allows to keep variance
# above zero and prevent it from growing too fast
if log_std <= 0:
return th.exp(log_std)
std = th.exp(log_std)
else:
return th.log(log_std + 1.0) + 1.0
std = th.log(log_std + 1.0) + 1.0
else:
# Use normal exponential
return th.exp(log_std)
std = th.exp(log_std)
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
def sample_weights(self, log_std):
"""
@ -283,8 +293,8 @@ class StateDependentNoiseDistribution(Distribution):
:param log_std: (th.Tensor)
"""
# TODO: reduce the number of learned dimensions (cf TD3)
self.weights_dist = Normal(th.zeros_like(log_std), self.get_std(log_std))
std = self.get_std(log_std)
self.weights_dist = Normal(th.zeros_like(std), std)
self.exploration_mat = self.weights_dist.rsample()
def proba_distribution_net(self, latent_dim, log_std_init=0.0):
@ -297,10 +307,17 @@ class StateDependentNoiseDistribution(Distribution):
:param log_std_init: (float) Initial value for the log standard deviation
:return: (nn.Linear, nn.Parameter)
"""
mean_actions = nn.Linear(latent_dim, self.action_dim)
log_std = nn.Parameter(th.ones(latent_dim, self.action_dim) * log_std_init)
# 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
# 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)
# 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, log_std
return mean_actions_net, log_std
def proba_distribution(self, mean_actions, log_std, latent_pi, deterministic=False):
"""
@ -312,7 +329,7 @@ class StateDependentNoiseDistribution(Distribution):
:return: (th.Tensor)
"""
variance = th.mm(latent_pi.detach() ** 2, self.get_std(log_std) ** 2)
self.distribution = Normal(mean_actions, th.sqrt(variance))
self.distribution = Normal(mean_actions, th.sqrt(variance + self.epsilon))
if deterministic:
action = self.mode()
@ -326,8 +343,11 @@ 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 sample(self, latent_pi):
noise = th.mm(latent_pi.detach(), self.exploration_mat)
noise = self.get_noise(latent_pi)
action = self.distribution.mean + noise
if self.bijector is not None:
return self.bijector.forward(action)

5
torchy_baselines/ppo/policies.py
View file

@ -58,6 +58,9 @@ class PPOPolicy(BasePolicy):
self._build(learning_rate)
def reset_noise_net(self):
"""
Sample new weights for the exploration matrix.
"""
self.action_dist.sample_weights(self.log_std)
def _build(self, learning_rate):
@ -75,7 +78,7 @@ class PPOPolicy(BasePolicy):
# with small initial weight for the output
if self.ortho_init:
for module in [self.mlp_extractor, self.action_net, self.value_net]:
# Values from stable-baselines check why
# Values from stable-baselines, TODO: check why
gain = {
self.mlp_extractor: np.sqrt(2),
self.action_net: 0.01,

61
torchy_baselines/td3/policies.py
View file

@ -1,8 +1,8 @@
import torch as th
import torch.nn as nn
from torch.distributions import Normal
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork
from torchy_baselines.common.distributions import StateDependentNoiseDistribution
class Actor(BaseNetwork):
@ -32,17 +32,15 @@ class Actor(BaseNetwork):
self.full_std = full_std
if use_sde:
latent_dim = net_arch[-1]
latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False)
self.latent_pi = nn.Sequential(*latent_pi)
if full_std:
self.log_std = nn.Parameter(th.ones(latent_dim, action_dim) * log_std_init, requires_grad=True)
else:
# Reduce the number of parameters:
self.log_std = nn.Parameter(th.ones(latent_dim, 1) * log_std_init, requires_grad=True)
self.latent_dim = latent_dim
self.actor_net = nn.Sequential(nn.Linear(net_arch[-1], action_dim), nn.Tanh())
# Create state dependent noise matrix (SDE)
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
squash_output=False)
action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
log_std_init=log_std_init)
# Squash output
self.actor_net = 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()
@ -50,11 +48,21 @@ class Actor(BaseNetwork):
actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True)
self.actor_net = nn.Sequential(*actor_net)
def get_log_std(self):
if self.full_std:
return self.log_std
# Reduce the number of parameters:
return th.ones((self.latent_dim, self.action_dim)).to(self.log_std.device) * self.log_std
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 _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 evaluate_actions(self, obs, action):
"""
@ -63,38 +71,37 @@ 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)
mean_actions = self.actor_net(latent_pi)
variance = th.mm(latent_pi ** 2, th.exp(self.get_log_std()) ** 2)
distribution = Normal(mean_actions, th.sqrt(variance + 1e-5))
_, distribution = self._get_action_dist_from_latent(latent_pi)
log_prob = distribution.log_prob(action)
if len(log_prob.shape) > 1:
log_prob = log_prob.sum(axis=1)
else:
log_prob = log_prob.sum()
# value = self.value_net(latent_vf)
return log_prob, distribution.entropy()
def reset_noise(self):
self.weights_dist = Normal(th.zeros_like(self.get_log_std()), th.exp(self.get_log_std()))
self.exploration_mat = self.weights_dist.rsample()
"""
Sample new weights for the exploration matrix.
"""
self.action_dist.sample_weights(self.log_std)
def forward(self, obs, deterministic=True):
if self.use_sde:
latent_pi = self.latent_pi(obs)
if deterministic:
return self.actor_net(latent_pi)
noise = th.mm(latent_pi.detach(), self.exploration_mat)
noise = self.action_dist.get_noise(latent_pi)
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
# return th.clamp(self.actor_net(latent_pi) + noise, -1, 1)
# -> 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
# action, _ = self._get_action_dist_from_latent(latent_pi)
# return action
else:
return self.actor_net(obs)