diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 5c8f55c..5a69f4e 100644 --- a/docs/misc/changelog.rst +++ b/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: ^^^^^^^^^^^^^^ diff --git a/setup.py b/setup.py index 21fedaa..ee0d0d6 100644 --- a/setup.py +++ b/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 diff --git a/tests/test_distributions.py b/tests/test_distributions.py index 7d28ad7..c63900d 100644 --- a/tests/test_distributions.py +++ b/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) diff --git a/tests/test_predict.py b/tests/test_predict.py index e68954f..2ffab6b 100644 --- a/tests/test_predict.py +++ b/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] diff --git a/tests/test_save_load.py b/tests/test_save_load.py index bdec3ed..9a0dad8 100644 --- a/tests/test_save_load.py +++ b/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 diff --git a/tests/test_sde.py b/tests/test_sde.py index 96f6c8f..cf3bbdb 100644 --- a/tests/test_sde.py +++ b/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 diff --git a/torchy_baselines/__init__.py b/torchy_baselines/__init__.py index b228acb..acf5912 100644 --- a/torchy_baselines/__init__.py +++ b/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() diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index eabe92e..cd9eb77 100644 --- a/torchy_baselines/common/base_class.py +++ b/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) diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index 6a0ab2f..5d04f83 100644 --- a/torchy_baselines/common/distributions.py +++ b/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) diff --git a/torchy_baselines/common/evaluation.py b/torchy_baselines/common/evaluation.py index ff9fa13..704057c 100644 --- a/torchy_baselines/common/evaluation.py +++ b/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: diff --git a/torchy_baselines/common/vec_env/__init__.py b/torchy_baselines/common/vec_env/__init__.py index 2cbb349..10f9288 100644 --- a/torchy_baselines/common/vec_env/__init__.py +++ b/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 diff --git a/torchy_baselines/version.txt b/torchy_baselines/version.txt new file mode 100644 index 0000000..53a75d6 --- /dev/null +++ b/torchy_baselines/version.txt @@ -0,0 +1 @@ +0.2.6