Add get_distribution for on-policy algorithms (#566)

* feat: get_distribution method for ActorCriticPolicy

New method get_distribution for class ActorCriticPolicy returning current action distribution given observations

* doc: updating changelog.rst

- adding block for Release 1.2.1a0
- adding cyprienc to contributors

* style: make format

* fix: updating version.txt

Changing version from 1.2.0 to 1.2.1a0

* Update changelog

* Add test for get distribution

Co-authored-by: Cyprien <courtot.c@gmail.com>

docs/misc/changelog.rst

@@ -4,6 +4,31 @@ Changelog
 ==========
 
 
+Release 1.2.1a0 (WIP)
+---------------------------
+
+
+Breaking Changes:
+^^^^^^^^^^^^^^^^^
+
+New Features:
+^^^^^^^^^^^^^
+- Added method ``get_distribution`` for ``ActorCriticPolicy`` for A2C/PPO/TRPO (@cyprienc)
+
+Bug Fixes:
+^^^^^^^^^^
+
+Deprecations:
+^^^^^^^^^^^^^
+
+Others:
+^^^^^^^
+
+Documentation:
+^^^^^^^^^^^^^^
+
+
+
 Release 1.2.0 (2021-09-03)
 ---------------------------
 
@@ -746,4 +771,4 @@ And all the contributors:
 @diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @megan-klaiber @thisray
 @tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio @JadenTravnik @decodyng @ardabbour @lorenz-h @mschweizer @lorepieri8 @vwxyzjn
 @ShangqunYu @PierreExeter @JacopoPan @ltbd78 @tom-doerr @Atlis @liusida @09tangriro @amy12xx @juancroldan @benblack769 @bstee615
-@c-rizz @skandermoalla @MihaiAnca13 @davidblom603 @ayeright
+@c-rizz @skandermoalla @MihaiAnca13 @davidblom603 @ayeright @cyprienc

stable_baselines3/common/policies.py

@@ -664,6 +664,16 @@ class ActorCriticPolicy(BasePolicy):
         values = self.value_net(latent_vf)
         return values, log_prob, distribution.entropy()
 
+    def get_distribution(self, obs: th.Tensor) -> Distribution:
+        """
+        Get the current policy distribution given the observations.
+
+        :param obs:
+        :return: the action distribution.
+        """
+        latent_pi, _, latent_sde = self._get_latent(obs)
+        return self._get_action_dist_from_latent(latent_pi, latent_sde)
+
 
 class ActorCriticCnnPolicy(ActorCriticPolicy):
     """

stable_baselines3/version.txt

@@ -1 +1 @@
-1.2.0
+1.2.1a0

tests/test_distributions.py

@@ -1,5 +1,7 @@
 from copy import deepcopy
 
+import gym
+import numpy as np
 import pytest
 import torch as th
 
@@ -51,6 +53,23 @@ def test_squashed_gaussian(model_class):
     assert th.max(th.abs(actions)) <= 1.0
 
 
+def test_get_distribution():
+    env = gym.make("Pendulum-v0")
+    model = A2C("MlpPolicy", env, seed=23)
+    random_obs = np.array([env.observation_space.sample() for _ in range(10)])
+    random_actions = np.array([env.action_space.sample() for _ in range(10)])
+    # Check that evaluate actions return the same thing as get_distribution
+    with th.no_grad():
+        observations, _ = model.policy.obs_to_tensor(random_obs)
+        actions = th.tensor(random_actions, device=observations.device).float()
+        _, log_prob_1, entropy_1 = model.policy.evaluate_actions(observations, actions)
+        distribution = model.policy.get_distribution(observations)
+        log_prob_2 = distribution.log_prob(actions)
+        entropy_2 = distribution.entropy()
+        assert th.allclose(log_prob_1, log_prob_2)
+        assert th.allclose(entropy_1, entropy_2)
+
+
 def test_sde_distribution():
     n_actions = 1
     deterministic_actions = th.ones(N_SAMPLES, n_actions) * 0.1