From c264403816836aef2f0a51cb93600c48231e3b48 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Tue, 31 Mar 2020 17:48:23 +0200 Subject: [PATCH] Rename for consistency + add _predict to actors + improve sac actor code --- tests/test_distributions.py | 8 ++--- torchy_baselines/common/distributions.py | 26 +++++++-------- torchy_baselines/ppo/policies.py | 10 +++--- torchy_baselines/sac/policies.py | 41 +++++++++++++++--------- torchy_baselines/td3/policies.py | 15 +++++---- 5 files changed, 56 insertions(+), 44 deletions(-) diff --git a/tests/test_distributions.py b/tests/test_distributions.py index 4a1294c..7d8f228 100644 --- a/tests/test_distributions.py +++ b/tests/test_distributions.py @@ -39,7 +39,7 @@ def test_squashed_gaussian(model_class): dist = SquashedDiagGaussianDistribution(N_ACTIONS) _, log_std = dist.proba_distribution_net(N_FEATURES) dist = dist.proba_distribution(gaussian_mean, log_std) - actions = dist.get_action() + actions = dist.get_actions() assert th.max(th.abs(actions)) <= 1.0 def test_sde_distribution(): @@ -53,7 +53,7 @@ def test_sde_distribution(): dist.sample_weights(log_std, batch_size=N_SAMPLES) dist = dist.proba_distribution(deterministic_actions, log_std, state) - actions = dist.get_action() + actions = dist.get_actions() assert th.allclose(actions.mean(), dist.distribution.mean.mean(), rtol=1e-3) assert th.allclose(actions.std(), dist.distribution.scale.mean(), rtol=1e-3) @@ -78,7 +78,7 @@ def test_entropy(dist): dist.sample_weights(log_std, batch_size=N_SAMPLES) dist = dist.proba_distribution(deterministic_actions, log_std, state) - actions = dist.get_action() + actions = dist.get_actions() entropy = dist.entropy() log_prob = dist.log_prob(actions) assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=5e-3) @@ -93,7 +93,7 @@ def test_categorical(): action_logits = th.rand(N_SAMPLES, N_ACTIONS) dist = dist.proba_distribution(action_logits) - actions = dist.get_action() + actions = dist.get_actions() entropy = dist.entropy() log_prob = dist.log_prob(actions) assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=1e-4) diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index 62d14b0..dd1090b 100644 --- a/torchy_baselines/common/distributions.py +++ b/torchy_baselines/common/distributions.py @@ -48,7 +48,7 @@ class Distribution(object): """ raise NotImplementedError - def get_action(self, deterministic: bool = False) -> th.Tensor: + def get_actions(self, deterministic: bool = False) -> th.Tensor: """ Return an action according to the probabilty distribution. @@ -60,7 +60,7 @@ class Distribution(object): else: return self.sample() - def action_from_params(self, *args, **kwargs) -> th.Tensor: + def actions_from_params(self, *args, **kwargs) -> th.Tensor: """ Returns a sample from the probabilty distribution given its parameters. @@ -149,12 +149,12 @@ class DiagGaussianDistribution(Distribution): def entropy(self) -> th.Tensor: return sum_independent_dims(self.distribution.entropy()) - def action_from_params(self, mean_actions: th.Tensor, + def actions_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor, deterministic: bool = False) -> th.Tensor: # Update the proba distribution self.proba_distribution(mean_actions, log_std) - return self.get_action(deterministic=deterministic) + return self.get_actions(deterministic=deterministic) def log_prob_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: @@ -166,7 +166,7 @@ class DiagGaussianDistribution(Distribution): :param log_std: (th.Tensor) :return: (Tuple[th.Tensor, th.Tensor]) """ - action = self.action_from_params(mean_actions, log_std) + action = self.actions_from_params(mean_actions, log_std) log_prob = self.log_prob(action) return action, log_prob @@ -219,7 +219,7 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution): def log_prob_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: - action = self.action_from_params(mean_actions, log_std) + action = self.actions_from_params(mean_actions, log_std) log_prob = self.log_prob(action, self.gaussian_action) return action, log_prob @@ -277,14 +277,14 @@ class CategoricalDistribution(Distribution): def entropy(self) -> th.Tensor: return self.distribution.entropy() - def action_from_params(self, action_logits: th.Tensor, + def actions_from_params(self, action_logits: th.Tensor, deterministic: bool = False) -> th.Tensor: # Update the proba distribution self.proba_distribution(action_logits) - return self.get_action(deterministic=deterministic) + return self.get_actions(deterministic=deterministic) def log_prob_from_params(self, action_logits: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: - action = self.action_from_params(action_logits) + action = self.actions_from_params(action_logits) log_prob = self.log_prob(action) return action, log_prob @@ -419,7 +419,7 @@ class StateDependentNoiseDistribution(Distribution): self.distribution = Normal(mean_actions, th.sqrt(variance + self.epsilon)) return self - def get_action(self, deterministic: bool = False) -> th.Tensor: + def get_actions(self, deterministic: bool = False) -> th.Tensor: if deterministic: return self.mode() else: @@ -457,18 +457,18 @@ class StateDependentNoiseDistribution(Distribution): return None return sum_independent_dims(self.distribution.entropy()) - def action_from_params(self, mean_actions: th.Tensor, + def actions_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor, latent_sde: th.Tensor, deterministic: bool = False) -> th.Tensor: # Update the proba distribution self.proba_distribution(mean_actions, log_std, latent_sde) - return self.get_action(deterministic=deterministic) + return self.get_actions(deterministic=deterministic) def log_prob_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor, latent_sde: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: - action = self.action_from_params(mean_actions, log_std, latent_sde) + action = self.actions_from_params(mean_actions, log_std, latent_sde) log_prob = self.log_prob(action) return action, log_prob diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 1589247..5ea14a7 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -155,11 +155,11 @@ class PPOPolicy(BasePolicy): """ latent_pi, latent_vf, latent_sde = self._get_latent(obs) # Evaluate the values for the given observations - value = self.value_net(latent_vf) + values = self.value_net(latent_vf) distribution = self._get_action_dist_from_latent(latent_pi, latent_sde=latent_sde) - action = distribution.get_action(deterministic=deterministic) - log_prob = distribution.log_prob(action) - return action, value, log_prob + actions = distribution.get_actions(deterministic=deterministic) + log_prob = distribution.log_prob(actions) + return actions, values, log_prob def _get_latent(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]: """ @@ -212,7 +212,7 @@ class PPOPolicy(BasePolicy): """ latent_pi, _, latent_sde = self._get_latent(observation) distribution = self._get_action_dist_from_latent(latent_pi, latent_sde) - return distribution.get_action(deterministic=deterministic) + return distribution.get_actions(deterministic=deterministic) def evaluate_actions(self, obs: th.Tensor, actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]: diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index f561c6e..c7417fc 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -1,4 +1,4 @@ -from typing import Optional, List, Tuple, Callable, Union, Type +from typing import Optional, List, Tuple, Callable, Union, Type, Dict import gym import torch as th @@ -108,34 +108,43 @@ class Actor(BasePolicy): 'reset_noise() is only available when using SDE' self.action_dist.sample_weights(self.log_std, batch_size=batch_size) - def get_action_dist_params(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]: + def get_action_dist_params(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, Dict[str, th.Tensor]]: + """ + Get the parameters for the action distribution. + + :param obs: (th.Tensor) + :return: (Tuple[th.Tensor, th.Tensor, Dict[str, th.Tensor]]) + Mean, standard deviation and optional keyword arguments. + """ features = self.extract_features(obs) latent_pi = self.latent_pi(features) - latent_sde = self.sde_features_extractor(features) if self.sde_features_extractor is not None else latent_pi - mean_actions = self.mu(latent_pi) if self.use_sde: - log_std = self.log_std - else: - log_std = self.log_std(latent_pi) - # 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_sde + latent_sde = latent_pi + if self.sde_features_extractor is not None: + latent_sde = self.sde_features_extractor(features) + return mean_actions, self.log_std, dict(latent_sde=latent_sde) + # Unstructured exploration (Original implementation) + log_std = self.log_std(latent_pi) + # Original Implementation to cap the standard deviation + log_std = th.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX) + return mean_actions, log_std, {} def forward(self, obs: th.Tensor, deterministic: bool = False) -> th.Tensor: - mean_actions, log_std, latent_sde = self.get_action_dist_params(obs) - kwargs = dict(latent_sde=latent_sde) if self.use_sde else {} + mean_actions, log_std, kwargs = self.get_action_dist_params(obs) # Note: the action is squashed - return self.action_dist.action_from_params(mean_actions, log_std, - deterministic=deterministic, **kwargs) + return self.action_dist.actions_from_params(mean_actions, log_std, + deterministic=deterministic, **kwargs) def action_log_prob(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: - mean_actions, log_std, latent_sde = self.get_action_dist_params(obs) - kwargs = dict(latent_sde=latent_sde) if self.use_sde else {} + mean_actions, log_std, kwargs = self.get_action_dist_params(obs) # return action and associated log prob return self.action_dist.log_prob_from_params(mean_actions, log_std, **kwargs) + def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor: + return self.forward(observation, deterministic) + class Critic(BasePolicy): """ diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index c10e4f2..30705b3 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -110,7 +110,7 @@ class Actor(BasePolicy): latent_sde = self.sde_features_extractor(features) if self.sde_features_extractor is not None else latent_pi return latent_pi, latent_sde - def evaluate_actions(self, obs: th.Tensor, action: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: + def evaluate_actions(self, obs: th.Tensor, actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: """ Evaluate actions according to the current policy, given the observations. Only useful when using SDE. @@ -123,7 +123,7 @@ class Actor(BasePolicy): latent_pi, latent_sde = self._get_latent(obs) mean_actions = self.mu(latent_pi) distribution = self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde) - log_prob = distribution.log_prob(action) + log_prob = distribution.log_prob(actions) return log_prob, distribution.entropy() def reset_noise(self) -> None: @@ -149,6 +149,9 @@ class Actor(BasePolicy): features = self.extract_features(obs) return self.mu(features) + def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor: + return self.forward(observation, deterministic=deterministic) + class Critic(BasePolicy): """ @@ -184,14 +187,14 @@ class Critic(BasePolicy): q2_net = create_mlp(features_dim + action_dim, 1, net_arch, activation_fn) self.q2_net = nn.Sequential(*q2_net) - def forward(self, obs: th.Tensor, action: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: + def forward(self, obs: th.Tensor, actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: features = self.extract_features(obs) - qvalue_input = th.cat([features, action], dim=1) + qvalue_input = th.cat([features, actions], dim=1) return self.q1_net(qvalue_input), self.q2_net(qvalue_input) - def q1_forward(self, obs: th.Tensor, action: th.Tensor) -> th.Tensor: + def q1_forward(self, obs: th.Tensor, actions: th.Tensor) -> th.Tensor: features = self.extract_features(obs) - return self.q1_net(th.cat([features, action], dim=1)) + return self.q1_net(th.cat([features, actions], dim=1)) class ValueFunction(BasePolicy):