diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index bf3468c..40dcedb 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -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 diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index 9d67b39..ad02734 100644 --- a/torchy_baselines/common/distributions.py +++ b/torchy_baselines/common/distributions.py @@ -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) diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index bd85797..adcf239 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -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, diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index b883db3..97d5862 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -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)