From cdcdd32c51c6b213500e436886b7cc392451e546 Mon Sep 17 00:00:00 2001 From: Juan Rocamonde Date: Fri, 14 Oct 2022 16:45:28 +0100 Subject: [PATCH] Fix return type of `evaluate_actions` (#1118) * Fix return type of ActorCriticPolicy.evaluate_actions to optional entropy tensor * Update changelog.rst --- docs/misc/changelog.rst | 1 + stable_baselines3/common/policies.py | 5 +++-- tests/test_distributions.py | 2 ++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 1d92231..ce5e3c8 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -22,6 +22,7 @@ SB3-Contrib Bug Fixes: ^^^^^^^^^^ +- Fix return type of ``evaluate_actions`` in ``ActorCritcPolicy`` to reflect that entropy is an optional tensor (@Rocamonde) Deprecations: ^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 18f90c0..876979f 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -613,7 +613,7 @@ class ActorCriticPolicy(BasePolicy): """ return self.get_distribution(observation).get_actions(deterministic=deterministic) - def evaluate_actions(self, obs: th.Tensor, actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]: + def evaluate_actions(self, obs: th.Tensor, actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor, Optional[th.Tensor]]: """ Evaluate actions according to the current policy, given the observations. @@ -629,7 +629,8 @@ class ActorCriticPolicy(BasePolicy): distribution = self._get_action_dist_from_latent(latent_pi) log_prob = distribution.log_prob(actions) values = self.value_net(latent_vf) - return values, log_prob, distribution.entropy() + entropy = distribution.entropy() + return values, log_prob, entropy def get_distribution(self, obs: th.Tensor) -> Distribution: """ diff --git a/tests/test_distributions.py b/tests/test_distributions.py index 07920db..513429b 100644 --- a/tests/test_distributions.py +++ b/tests/test_distributions.py @@ -77,6 +77,8 @@ def test_get_distribution(dummy_model_distribution_obs_and_actions): distribution = model.policy.get_distribution(observations) log_prob_2 = distribution.log_prob(actions) entropy_2 = distribution.entropy() + assert entropy_1 is not None + assert entropy_2 is not None assert th.allclose(log_prob_1, log_prob_2) assert th.allclose(entropy_1, entropy_2)