diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 608e882..e6716ee 100644 --- a/docs/misc/changelog.rst +++ b/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 diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 1550e23..a07a006 100644 --- a/stable_baselines3/common/policies.py +++ b/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): """ diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 26aaba0..348e216 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.2.0 +1.2.1a0 diff --git a/tests/test_distributions.py b/tests/test_distributions.py index db8e76c..88cf864 100644 --- a/tests/test_distributions.py +++ b/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