Add support for categorical distribution

README.md

@@ -22,7 +22,6 @@ TODO:
 - flexible mlp
 - logger
 - better monitor wrapper?
-- automatic choice for action distribution
 - A2C
 
 Later:

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")

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())

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.")

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)

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