From ef50bb81e8073c5af73b832f0351aaae473ac1f5 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Tue, 8 Oct 2019 13:06:38 +0200 Subject: [PATCH] Add support for categorical distribution --- README.md | 1 - tests/test_run.py | 6 ++- torchy_baselines/common/buffers.py | 3 ++ torchy_baselines/common/distributions.py | 65 +++++++++++++++++++++++- torchy_baselines/ppo/policies.py | 27 +++++++--- torchy_baselines/ppo/ppo.py | 21 +++++++- 6 files changed, 110 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 5f70f36..a73f690 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,6 @@ TODO: - flexible mlp - logger - better monitor wrapper? -- automatic choice for action distribution - A2C Later: diff --git a/tests/test_run.py b/tests/test_run.py index f409ab5..9740921 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -1,5 +1,6 @@ import os +import pytest import numpy as np from torchy_baselines import TD3, CEMRL, PPO, SAC @@ -27,8 +28,9 @@ def test_cemrl(): os.remove("test_save.pth") -def test_ppo(): - model = PPO('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True) +@pytest.mark.parametrize("env_id", ['CartPole-v1', 'Pendulum-v0']) +def test_ppo(env_id): + model = PPO('MlpPolicy', env_id, policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True) model.learn(total_timesteps=1000, eval_freq=500) # model.save("test_save") # model.load("test_save") diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py index 61421fa..565117c 100644 --- a/torchy_baselines/common/buffers.py +++ b/torchy_baselines/common/buffers.py @@ -131,6 +131,9 @@ class RolloutBuffer(BaseBuffer): self.returns = self.advantages + self.values def add(self, obs, action, reward, done, value, log_prob): + if len(log_prob.shape) == 0: + # Reshape 0-d tensor to avoid error + log_prob = log_prob.reshape(-1, 1) self.observations[self.pos] = th.FloatTensor(np.array(obs).copy()) self.actions[self.pos] = th.FloatTensor(np.array(action).copy()) self.rewards[self.pos] = th.FloatTensor(np.array(reward).copy()) diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index bca56f5..a384ce3 100644 --- a/torchy_baselines/common/distributions.py +++ b/torchy_baselines/common/distributions.py @@ -1,7 +1,8 @@ import numpy as np import torch as th import torch.nn as nn -from torch.distributions import Normal +from torch.distributions import Normal, Categorical +from gym import spaces class Distribution(object): def __init__(self): @@ -100,8 +101,9 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution): return action, self def mode(self): + self.gaussian_action = self.distribution.mean # Squash the output - return th.tanh(self.distribution.mean) + return th.tanh(self.gaussian_action) def sample(self): self.gaussian_action = self.distribution.rsample() @@ -124,3 +126,62 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution): # Squash correction (from original SAC implementation) log_prob -= th.sum(th.log(1 - action ** 2 + self.epsilon), dim=1) return log_prob + + +class CategoricalDistribution(Distribution): + def __init__(self, action_dim): + super(CategoricalDistribution, self).__init__() + self.distribution = None + self.action_dim = action_dim + + def proba_distribution_net(self, latent_dim): + action_logits = nn.Linear(latent_dim, self.action_dim) + return action_logits + + def proba_distribution(self, action_logits, deterministic=False): + self.distribution = Categorical(logits=action_logits) + if deterministic: + action = self.mode() + else: + action = self.sample() + return action, self + + def mode(self): + return th.argmax(self.distribution.probs, dim=1) + + def sample(self): + return self.distribution.sample() + + def entropy(self): + return self.distribution.entropy() + + def log_prob_from_params(self, action_logits): + action, _ = self.proba_distribution(action_logits) + log_prob = self.log_prob(action) + return action, log_prob + + def log_prob(self, action): + log_prob = self.distribution.log_prob(action) + return log_prob + + +def make_proba_distribution(action_space): + """ + Return an instance of Distribution for the correct type of action space + + :param action_space: (Gym Space) the input action space + :return: (Distribution) the approriate Distribution object + """ + if isinstance(action_space, spaces.Box): + assert len(action_space.shape) == 1, "Error: the action space must be a vector" + return DiagGaussianDistribution(action_space.shape[0]) + elif isinstance(action_space, spaces.Discrete): + return CategoricalDistribution(action_space.n) + # elif isinstance(action_space, spaces.MultiDiscrete): + # return MultiCategoricalDistribution(action_space.nvec) + # elif isinstance(action_space, spaces.MultiBinary): + # return BernoulliDistribution(action_space.n) + else: + raise NotImplementedError("Error: probability distribution, not implemented for action space of type {}." + .format(type(action_space)) + + " Must be of type Gym Spaces: Box, Discrete, MultiDiscrete or MultiBinary.") diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 6f6f4a7..2bf5fab 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -5,7 +5,7 @@ import torch.nn as nn import numpy as np from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp -from torchy_baselines.common.distributions import DiagGaussianDistribution, SquashedDiagGaussianDistribution +from torchy_baselines.common.distributions import make_proba_distribution, DiagGaussianDistribution, CategoricalDistribution class PPOPolicy(BasePolicy): @@ -14,7 +14,6 @@ class PPOPolicy(BasePolicy): activation_fn=nn.Tanh, adam_epsilon=1e-5, ortho_init=True): super(PPOPolicy, self).__init__(observation_space, action_space, device) self.obs_dim = self.observation_space.shape[0] - self.action_dim = self.action_space.shape[0] if net_arch is None: net_arch = [64, 64] self.net_arch = net_arch @@ -30,8 +29,7 @@ class PPOPolicy(BasePolicy): self.shared_net = None self.pi_net, self.vf_net = None, None # Action distribution - self.action_dist = DiagGaussianDistribution(self.action_dim) - # self.action_dist = SquashedDiagGaussianDistribution(self.action_dim) + self.action_dist = make_proba_distribution(action_space) self._build(learning_rate) def _build(self, learning_rate): @@ -46,7 +44,11 @@ class PPOPolicy(BasePolicy): # self.action_net = nn.Linear(self.net_arch[-1], self.action_dim) # self.log_std = nn.Parameter(th.zeros(self.action_dim)) - self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.net_arch[-1]) + if isinstance(self.action_dist, DiagGaussianDistribution): + self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.net_arch[-1]) + elif isinstance(self.action_dist, CategoricalDistribution): + self.action_net = self.action_dist.proba_distribution_net(latent_dim=self.net_arch[-1]) + self.value_net = nn.Linear(self.net_arch[-1], 1) # Init weights: use orthogonal initialization # with small initial weight for the output @@ -64,6 +66,14 @@ class PPOPolicy(BasePolicy): # TODO: support linear decay of the learning rate self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate, eps=self.adam_epsilon) + # def get_action_dist_params(self, obs): + # latent_pi, _ = self._get_latent(obs) + # mean_actions = self.pi_net(latent_pi) + # if isinstance(self.action_dist, DiagGaussianDistribution): + # return {'mean_actions': mean_actions, 'log_std': self.log_std} + # elif isinstance(self.action_dist, CategoricalDistribution): + # return {'action_logits': mean_actions} + def forward(self, obs, deterministic=False): if not isinstance(obs, th.Tensor): obs = th.FloatTensor(obs).to(self.device) @@ -71,6 +81,8 @@ class PPOPolicy(BasePolicy): value = self.value_net(latent_vf) action, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic) log_prob = action_distribution.log_prob(action) + # mean_actions, log_std = self.get_action_dist_params(obs) + # action, log_prob = self.action_dist.log_prob_from_params(**self.get_action_dist_params(obs)) return action, value, log_prob def _get_latent(self, obs): @@ -82,7 +94,10 @@ class PPOPolicy(BasePolicy): def _get_action_dist_from_latent(self, latent, deterministic=False): mean_actions = self.action_net(latent) - return self.action_dist.proba_distribution(mean_actions, self.log_std, deterministic=deterministic) + if isinstance(self.action_dist, DiagGaussianDistribution): + return self.action_dist.proba_distribution(mean_actions, self.log_std, deterministic=deterministic) + elif isinstance(self.action_dist, CategoricalDistribution): + return self.action_dist.proba_distribution(mean_actions, deterministic=deterministic) def actor_forward(self, obs, deterministic=False): latent_pi, _ = self._get_latent(obs) diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 1efd412..e836555 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -3,6 +3,7 @@ import time from copy import deepcopy import gym +from gym import spaces import torch as th import torch.nn.functional as F # Check if tensorboard is available for pytorch @@ -96,7 +97,15 @@ class PPO(BaseRLModel): self._setup_model() def _setup_model(self): - state_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0] + # TODO: preprocessing: one hot vector for obs discrete + state_dim = self.observation_space.shape[0] + if isinstance(self.action_space, spaces.Box): + # Action is a 1D vector + action_dim = self.action_space.shape[0] + elif isinstance(self.action_space, spaces.Discrete): + # Action is a scalar + action_dim = 1 + # TODO: different seed for each env when n_envs > 1 if self.n_envs == 1: self.set_random_seed(self.seed) @@ -124,7 +133,10 @@ class PPO(BaseRLModel): :param deterministic: (bool) Whether or not to return deterministic actions. :return: (np.ndarray, np.ndarray) the model's action and the next state (used in recurrent policies) """ - return np.clip(self.select_action(observation), self.action_space.low, self.action_space.high) + clipped_actions = self.select_action(observation) + if isinstance(self.action_space, gym.spaces.Box): + clipped_actions = np.clip(clipped_actions, self.action_space.low, self.action_space.high) + return clipped_actions def collect_rollouts(self, env, rollout_buffer, n_rollout_steps=256, callback=None, obs=None): @@ -160,6 +172,11 @@ class PPO(BaseRLModel): for replay_data in self.rollout_buffer.get(batch_size): # Unpack obs, action, old_values, old_log_prob, advantage, return_batch = replay_data + + if isinstance(self.action_space, spaces.Discrete): + # Convert discrete action for float to long + action = action.long().flatten() + values, log_prob, entropy = self.policy.get_policy_stats(obs, action) values = values.flatten() # Normalize advantage