Rename for consistency

+ add _predict to actors
+ improve sac actor code

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)

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
 

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)
+        actions = distribution.get_actions(deterministic=deterministic)
-        log_prob = distribution.log_prob(action)
+        log_prob = distribution.log_prob(actions)
-        return action, value, log_prob
+        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]:

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
+            latent_sde = latent_pi
-        else:
+            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, latent_sde
+        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)
+        mean_actions, log_std, kwargs = self.get_action_dist_params(obs)
-        kwargs = dict(latent_sde=latent_sde) if self.use_sde else {}
         # Note: the action is squashed
-        return self.action_dist.action_from_params(mean_actions, log_std,
+        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)
+        mean_actions, log_std, kwargs = self.get_action_dist_params(obs)
-        kwargs = dict(latent_sde=latent_sde) if self.use_sde else {}
         # 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):
     """

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