Merge pull request #61 from Antonin-Raffin/feat/more-tests

More tests and sync SB

docs/misc/changelog.rst

@@ -11,6 +11,7 @@ Breaking Changes:
 ^^^^^^^^^^^^^^^^^
 - Removed default seed
 - Bump dependencies (PyTorch and Gym)
+- ``predict()`` now returns a tuple to match Stable-Baselines behavior
 
 New Features:
 ^^^^^^^^^^^^^
@@ -20,6 +21,7 @@ Bug Fixes:
 ^^^^^^^^^^
 - Synced callbacks with Stable-Baselines
 - Fixed colors in `results_plotter`
+- Fix entropy computation (now summed over action dim)
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -32,6 +34,8 @@ Others:
 - More typing
 - Add test for ``expln``
 - Renamed ``learning_rate`` to ``lr_schedule``
+- Add ``version.txt``
+- Add more tests for distribution
 
 Documentation:
 ^^^^^^^^^^^^^^

setup.py

@@ -1,7 +1,11 @@
+import os
 import sys
 import subprocess
 from setuptools import setup, find_packages
 
+with open(os.path.join('torchy_baselines', 'version.txt'), 'r') as file_handler:
+    __version__ = file_handler.read()
+
 
 setup(name='torchy_baselines',
       packages=[package for package in find_packages()
@@ -48,7 +52,7 @@ setup(name='torchy_baselines',
       license="MIT",
       long_description="",
       long_description_content_type='text/markdown',
-      version="0.2.4",
+      version=__version__,
       )
 
 # python setup.py sdist

tests/test_distributions.py

@@ -2,18 +2,22 @@ import pytest
 import torch as th
 
 from torchy_baselines import A2C, PPO
-from torchy_baselines.common.distributions import DiagGaussianDistribution, TanhBijector, \
-    StateDependentNoiseDistribution
+from torchy_baselines.common.distributions import (DiagGaussianDistribution, TanhBijector,
+                                                   StateDependentNoiseDistribution,
+                                                   CategoricalDistribution)
 from torchy_baselines.common.utils import set_random_seed
 
 
-# TODO: more tests for the other distributions
+N_ACTIONS = 2
+N_FEATURES = 3
+N_SAMPLES = int(5e6)
+
+
 def test_bijector():
     """
     Test TanhBijector
     """
     actions = th.ones(5) * 2.0
-
     bijector = TanhBijector()
 
     squashed_actions = bijector.forward(actions)
@@ -33,16 +37,14 @@ def test_squashed_gaussian(model_class):
 
 
 def test_sde_distribution():
-    n_samples = int(5e6)
-    n_features = 2
     n_actions = 1
-    deterministic_actions = th.ones(n_samples, n_actions) * 0.1
-    state = th.ones(n_samples, n_features) * 0.3
+    deterministic_actions = th.ones(N_SAMPLES, n_actions) * 0.1
+    state = th.ones(N_SAMPLES, N_FEATURES) * 0.3
     dist = StateDependentNoiseDistribution(n_actions, full_std=True, squash_output=False)
 
     set_random_seed(1)
-    _, log_std = dist.proba_distribution_net(n_features)
-    dist.sample_weights(log_std, batch_size=n_samples)
+    _, log_std = dist.proba_distribution_net(N_FEATURES)
+    dist.sample_weights(log_std, batch_size=N_SAMPLES)
 
     actions, _ = dist.proba_distribution(deterministic_actions, log_std, state)
 
@@ -50,10 +52,6 @@ def test_sde_distribution():
     assert th.allclose(actions.std(), dist.distribution.scale.mean(), rtol=1e-3)
 
 
-N_ACTIONS = 1
-
-
-# TODO: fix for num action > 1
 # TODO: analytical form for squashed Gaussian?
 @pytest.mark.parametrize("dist", [
     DiagGaussianDistribution(N_ACTIONS),
@@ -62,19 +60,31 @@ N_ACTIONS = 1
 def test_entropy(dist):
     # The entropy can be approximated by averaging the negative log likelihood
     # mean negative log likelihood == differential entropy
-    n_samples = int(5e6)
-    n_features = 3
     set_random_seed(1)
-    state = th.rand(n_samples, n_features)
-    deterministic_actions = th.rand(n_samples, N_ACTIONS)
-    _, log_std = dist.proba_distribution_net(n_features, log_std_init=th.log(th.tensor(0.2)))
+    state = th.rand(N_SAMPLES, N_FEATURES)
+    deterministic_actions = th.rand(N_SAMPLES, N_ACTIONS)
+    _, log_std = dist.proba_distribution_net(N_FEATURES, log_std_init=th.log(th.tensor(0.2)))
 
     if isinstance(dist, DiagGaussianDistribution):
         actions, dist = dist.proba_distribution(deterministic_actions, log_std)
     else:
-        dist.sample_weights(log_std, batch_size=n_samples)
+        dist.sample_weights(log_std, batch_size=N_SAMPLES)
         actions, dist = dist.proba_distribution(deterministic_actions, log_std, state)
 
     entropy = dist.entropy()
     log_prob = dist.log_prob(actions)
     assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=5e-3)
+
+
+def test_categorical():
+    # The entropy can be approximated by averaging the negative log likelihood
+    # mean negative log likelihood == entropy
+    dist = CategoricalDistribution(N_ACTIONS)
+    set_random_seed(1)
+    state = th.rand(N_SAMPLES, N_FEATURES)
+    action_logits = th.rand(N_SAMPLES, N_ACTIONS)
+    actions, dist = dist.proba_distribution(action_logits)
+
+    entropy = dist.entropy()
+    log_prob = dist.log_prob(actions)
+    assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=1e-4)

tests/test_predict.py

@@ -34,10 +34,10 @@ def test_predict(model_class, env_id):
     vec_env = DummyVecEnv([lambda: gym.make(env_id), lambda: gym.make(env_id)])
 
     obs = env.reset()
-    action = model.predict(obs)
+    action, _ = model.predict(obs)
     assert action.shape == env.action_space.shape
     assert env.action_space.contains(action)
 
     vec_env_obs = vec_env.reset()
-    action = model.predict(vec_env_obs)
+    action, _ = model.predict(vec_env_obs)
     assert action.shape[0] == vec_env_obs.shape[0]

tests/test_save_load.py

@@ -54,7 +54,7 @@ def test_save_load(model_class):
     params = new_params
 
     # get selected actions
-    selected_actions = model.predict(observations, deterministic=True)
+    selected_actions, _ = model.predict(observations, deterministic=True)
 
     # Check
     model.save("test_save.zip")
@@ -69,7 +69,7 @@ def test_save_load(model_class):
         assert th.allclose(params[key], new_params[key]), "Model parameters not the same after save and load."
 
     # check if model still selects the same actions
-    new_selected_actions = model.predict(observations, deterministic=True)
+    new_selected_actions, _ = model.predict(observations, deterministic=True)
     assert np.allclose(selected_actions, new_selected_actions, 1e-4)
 
     # check if learn still works

tests/test_sde.py

@@ -22,7 +22,6 @@ def test_state_dependent_exploration_grad():
 
     state = th.rand(n_states, state_dim)
     mu = th.ones(action_dim)
-    # print(weights.shape, state.shape)
     noise = th.mm(state, weights)
 
     action = mu + noise

torchy_baselines/__init__.py

@@ -1,7 +1,12 @@
+import os
+
 from torchy_baselines.a2c import A2C
 from torchy_baselines.cem_rl import CEMRL
 from torchy_baselines.ppo import PPO
 from torchy_baselines.sac import SAC
 from torchy_baselines.td3 import TD3
 
-__version__ = "0.2.4"
+# Read version from file
+version_file = os.path.join(os.path.dirname(__file__), 'version.txt')
+with open(version_file, 'r') as file_handler:
+    __version__ = file_handler.read()

torchy_baselines/common/base_class.py

@@ -369,7 +369,7 @@ class BaseRLModel(ABC):
     def predict(self, observation: np.ndarray,
                 state: Optional[np.ndarray] = None,
                 mask: Optional[np.ndarray] = None,
-                deterministic: bool = False) -> np.ndarray:
+                deterministic: bool = False) -> Tuple[np.ndarray, Optional[np.ndarray]]:
         """
         Get the model's action(s) from an observation
 
@@ -377,7 +377,7 @@ class BaseRLModel(ABC):
         :param state: (Optional[np.ndarray]) The last states (can be None, used in recurrent policies)
         :param mask: (Optional[np.ndarray]) The last masks (can be None, used in recurrent policies)
         :param deterministic: (bool) Whether or not to return deterministic actions.
-        :return: (np.ndarray) the model's action and the next state (used in recurrent policies)
+        :return: (Tuple[np.ndarray, Optional[np.ndarray]]) the model's action and the next state (used in recurrent policies)
         """
         # if state is None:
         #     state = self.initial_state
@@ -409,9 +409,7 @@ class BaseRLModel(ABC):
                 raise ValueError("Error: The environment must be vectorized when using recurrent policies.")
             clipped_actions = clipped_actions[0]
 
-        # TODO: switch to stable baselines API
-        # return clipped_actions, state
-        return clipped_actions
+        return clipped_actions, state
 
     @classmethod
     def load(cls, load_path: str, env: Optional[GymEnv] = None, **kwargs):
@@ -896,7 +894,7 @@ class OffPolicyRLModel(BaseRLModel):
                 else:
                     # Note: we assume that the policy uses tanh to scale the action
                     # We use non-deterministic action in the case of SAC, for TD3, it does not matter
-                    unscaled_action = self.predict(obs, deterministic=False)
+                    unscaled_action, _ = self.predict(obs, deterministic=False)
 
                 # Rescale the action from [low, high] to [-1, 1]
                 scaled_action = self.scale_action(unscaled_action)

torchy_baselines/common/distributions.py

@@ -38,6 +38,22 @@ class Distribution(object):
         raise NotImplementedError
 
 
+def sum_independent_dims(tensor: th.Tensor) -> th.Tensor:
+    """
+    Continuous actions are usually considered to be independent,
+    so we can sum the components for the ``log_prob``
+    or the entropy.
+
+    :param tensor: (th.Tensor) shape: (n_batch, n_actions) or (n_batch,)
+    :return: (th.Tensor) shape: (n_batch,)
+    """
+    if len(tensor.shape) > 1:
+        tensor = tensor.sum(axis=1)
+    else:
+        tensor = tensor.sum()
+    return tensor
+
+
 class DiagGaussianDistribution(Distribution):
     """
     Gaussian distribution with diagonal covariance matrix,
@@ -95,7 +111,7 @@ class DiagGaussianDistribution(Distribution):
         return self.distribution.rsample()
 
     def entropy(self) -> th.Tensor:
-        return self.distribution.entropy()
+        return sum_independent_dims(self.distribution.entropy())
 
     def log_prob_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
         """
@@ -113,18 +129,14 @@ class DiagGaussianDistribution(Distribution):
     def log_prob(self, action: th.Tensor) -> th.Tensor:
         """
         Get the log probabilty of an action given a distribution.
-        Note that you must call `proba_distribution()` method
+        Note that you must call ``proba_distribution()`` method
         before.
 
         :param action: (th.Tensor)
         :return: (th.Tensor)
         """
         log_prob = self.distribution.log_prob(action)
-        if len(log_prob.shape) > 1:
-            log_prob = log_prob.sum(axis=1)
-        else:
-            log_prob = log_prob.sum()
-        return log_prob
+        return sum_independent_dims(log_prob)
 
 
 class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
@@ -243,14 +255,14 @@ class StateDependentNoiseDistribution(Distribution):
     :param action_dim: (int) Number of continuous actions
     :param full_std: (bool) Whether to use (n_features x n_actions) parameters
         for the std instead of only (n_features,)
-    :param use_expln: (bool) Use `expln()` function instead of `exp()` to ensure
+    :param use_expln: (bool) Use ``expln()`` function instead of ``exp()`` to ensure
         a positive standard deviation (cf paper). It allows to keep variance
-        above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
+        above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
     :param squash_output: (bool) Whether to squash the output using a tanh function,
         this allows to ensure boundaries.
     :param learn_features: (bool) Whether to learn features for SDE or not.
         This will enable gradients to be backpropagated through the features
-        `latent_sde` in the code.
+        ``latent_sde`` in the code.
     :param epsilon: (float) small value to avoid NaN due to numerical imprecision.
     """
 
@@ -396,7 +408,7 @@ class StateDependentNoiseDistribution(Distribution):
         # entropy needs to be estimated using -log_prob.mean()
         if self.bijector is not None:
             return None
-        return self.distribution.entropy()
+        return sum_independent_dims(self.distribution.entropy())
 
     def log_prob_from_params(self, mean_actions: th.Tensor,
                              log_std: th.Tensor,
@@ -412,11 +424,8 @@ class StateDependentNoiseDistribution(Distribution):
             gaussian_action = action
         # log likelihood for a gaussian
         log_prob = self.distribution.log_prob(gaussian_action)
-
-        if len(log_prob.shape) > 1:
-            log_prob = log_prob.sum(axis=1)
-        else:
-            log_prob = log_prob.sum()
+        # Sum along action dim
+        log_prob = sum_independent_dims(log_prob)
 
         if self.bijector is not None:
             # Squash correction (from original SAC implementation)

torchy_baselines/common/evaluation.py

@@ -8,11 +8,11 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True,
                     render=False, callback=None, reward_threshold=None,
                     return_episode_rewards=False):
     """
-    Runs policy for `n_eval_episodes` episodes and returns average reward.
+    Runs policy for ``n_eval_episodes`` episodes and returns average reward.
     This is made to work only with one env.
 
     :param model: (BaseRLModel) The RL agent you want to evaluate.
-    :param env: (gym.Env or VecEnv) The gym environment. In the case of a `VecEnv`
+    :param env: (gym.Env or VecEnv) The gym environment. In the case of a ``VecEnv``
         this must contain only one environment.
     :param n_eval_episodes: (int) Number of episode to evaluate the agent
     :param deterministic: (bool) Whether to use deterministic or stochastic actions
@@ -24,7 +24,7 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True,
     :param return_episode_rewards: (bool) If True, a list of reward per episode
         will be returned instead of the mean.
     :return: (float, float) Mean reward per episode, std of reward per episode
-        returns ([float], [int]) when `return_episode_rewards` is True
+        returns ([float], [int]) when ``return_episode_rewards`` is True
     """
     if isinstance(env, VecEnv):
         assert env.num_envs == 1, "You must pass only one environment when using this function"
@@ -32,11 +32,11 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True,
     episode_rewards, episode_lengths = [], []
     for _ in range(n_eval_episodes):
         obs = env.reset()
-        done = False
+        done, state = False, None
         episode_reward = 0.0
         episode_length = 0
         while not done:
-            action = model.predict(obs, deterministic=deterministic)
+            action, state = model.predict(obs, state=state, deterministic=deterministic)
             obs, reward, done, _info = env.step(action)
             episode_reward += reward
             if callback is not None:

torchy_baselines/common/vec_env/__init__.py

@@ -3,8 +3,8 @@ import typing
 from typing import Optional
 from copy import deepcopy
 
-from torchy_baselines.common.vec_env.base_vec_env import AlreadySteppingError, NotSteppingError,\
-    VecEnv, VecEnvWrapper, CloudpickleWrapper
+from torchy_baselines.common.vec_env.base_vec_env import (AlreadySteppingError, NotSteppingError,
+                                                          VecEnv, VecEnvWrapper, CloudpickleWrapper)
 from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv
 from torchy_baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
 from torchy_baselines.common.vec_env.vec_frame_stack import VecFrameStack

torchy_baselines/version.txt

@@ -0,0 +1 @@
+0.2.6