Add method predict_values for ActorCriticPolicy (#569)

* feat: add method predict_values for ActorCriticPolicy

* Fixes for new gym version

* Reformat

Co-authored-by: Antonin Raffin <antonin.raffin@ensta.org>

docs/misc/changelog.rst

@@ -4,7 +4,7 @@ Changelog
 ==========
 
 
-Release 1.2.1a0 (WIP)
+Release 1.2.1a1 (WIP)
 ---------------------------
 
 
@@ -13,16 +13,18 @@ Breaking Changes:
 
 New Features:
 ^^^^^^^^^^^^^
-- Added method ``get_distribution`` for ``ActorCriticPolicy`` for A2C/PPO/TRPO (@cyprienc)
+- Added methods ``get_distribution`` and ``predict_values`` for ``ActorCriticPolicy`` for A2C/PPO/TRPO (@cyprienc)
 
 Bug Fixes:
 ^^^^^^^^^^
+- Fixed ``dtype`` of observations for ``SimpleMultiObsEnv``
 
 Deprecations:
 ^^^^^^^^^^^^^
 
 Others:
 ^^^^^^^
+- Cap gym max version to 0.19 to avoid issues with atari-py and other breaking changes
 
 Documentation:
 ^^^^^^^^^^^^^^

setup.py

@@ -73,7 +73,7 @@ setup(
     packages=[package for package in find_packages() if package.startswith("stable_baselines3")],
     package_data={"stable_baselines3": ["py.typed", "version.txt"]},
     install_requires=[
-        "gym>=0.17",
+        "gym>=0.17,<0.20",  # gym 0.20 breaks atari-py behavior
         "numpy",
         "torch>=1.8.1",
         # For saving models

stable_baselines3/common/envs/multi_input_envs.py

@@ -59,7 +59,7 @@ class SimpleMultiObsEnv(gym.Env):
 
         self.observation_space = gym.spaces.Dict(
             spaces={
-                "vec": gym.spaces.Box(0, 1, (self.vector_size,)),
+                "vec": gym.spaces.Box(0, 1, (self.vector_size,), dtype=np.float64),
                 "img": gym.spaces.Box(0, 255, self.img_size, dtype=np.uint8),
             }
         )
@@ -87,7 +87,7 @@ class SimpleMultiObsEnv(gym.Env):
         # Each column is represented by a random vector
         col_vecs = np.random.random((num_col, self.vector_size))
         # Each row is represented by a random image
-        row_imgs = np.random.randint(0, 255, (num_row, 64, 64), dtype=np.int32)
+        row_imgs = np.random.randint(0, 255, (num_row, 64, 64), dtype=np.uint8)
 
         for i in range(num_col):
             for j in range(num_row):

stable_baselines3/common/policies.py

@@ -674,6 +674,16 @@ class ActorCriticPolicy(BasePolicy):
         latent_pi, _, latent_sde = self._get_latent(obs)
         return self._get_action_dist_from_latent(latent_pi, latent_sde)
 
+    def predict_values(self, obs: th.Tensor) -> th.Tensor:
+        """
+        Get the estimated values according to the current policy given the observations.
+
+        :param obs:
+        :return: the estimated values.
+        """
+        _, latent_vf, _ = self._get_latent(obs)
+        return self.value_net(latent_vf)
+
 
 class ActorCriticCnnPolicy(ActorCriticPolicy):
     """

stable_baselines3/version.txt

@@ -1 +1 @@
-1.2.1a0
+1.2.1a1

tests/test_distributions.py

@@ -1,4 +1,5 @@
 from copy import deepcopy
+from typing import Tuple
 
 import gym
 import numpy as np
@@ -53,11 +54,21 @@ def test_squashed_gaussian(model_class):
     assert th.max(th.abs(actions)) <= 1.0
 
 
-def test_get_distribution():
+@pytest.fixture()
+def dummy_model_distribution_obs_and_actions() -> Tuple[A2C, np.array, np.array]:
+    """
+    Fixture creating a Pendulum-v0 gym env, an A2C model and sampling 10 random observations and actions from the env
+    :return: A2C model, random observations, random actions
+    """
     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)])
+    return model, random_obs, random_actions
+
+
+def test_get_distribution(dummy_model_distribution_obs_and_actions):
+    model, random_obs, random_actions = dummy_model_distribution_obs_and_actions
     # Check that evaluate actions return the same thing as get_distribution
     with th.no_grad():
         observations, _ = model.policy.obs_to_tensor(random_obs)
@@ -70,6 +81,17 @@ def test_get_distribution():
         assert th.allclose(entropy_1, entropy_2)
 
 
+def test_predict_values(dummy_model_distribution_obs_and_actions):
+    model, random_obs, random_actions = dummy_model_distribution_obs_and_actions
+    # Check that evaluate_actions return the same thing as predict_values
+    with th.no_grad():
+        observations, _ = model.policy.obs_to_tensor(random_obs)
+        actions = th.tensor(random_actions, device=observations.device).float()
+        values_1, _, _ = model.policy.evaluate_actions(observations, actions)
+        values_2 = model.policy.predict_values(observations)
+        assert th.allclose(values_1, values_2)
+
+
 def test_sde_distribution():
     n_actions = 1
     deterministic_actions = th.ones(N_SAMPLES, n_actions) * 0.1