PPO

hive/agents/__init__.py

@@ -3,6 +3,7 @@
 from hive.agents.ddpg import DDPG
 from hive.agents.dqn import DQNAgent
 from hive.agents.legal_moves_rainbow import LegalMovesRainbowAgent
+from hive.agents.ppo import PPOAgent
 from hive.agents.rainbow import RainbowDQNAgent
 from hive.agents.random import RandomAgent
 from hive.agents.td3 import TD3
@@ -13,6 +14,7 @@
     {
         "DDPG": DDPG,
         "DQNAgent": DQNAgent,
+        "PPOAgent": PPOAgent,
         "LegalMovesRainbowAgent": LegalMovesRainbowAgent,
         "RainbowDQNAgent": RainbowDQNAgent,
         "RandomAgent": RandomAgent,

hive/agents/qnets/ppo_nets.py

@@ -0,0 +1,113 @@
+import gym
+import numpy as np
+import torch
+
+from hive.agents.qnets.base import FunctionApproximator
+from hive.agents.qnets.utils import calculate_output_dim
+
+
+# TODO: add Multi-Discrete
+class CategoricalHead(torch.nn.Module):
+    """A module that implements a discrete actor head. It uses the ouput from the
+    :obj:`actor_net`, and adds creates a :py:class:`~torch.distributions.categorical.Categorical`
+    object to compute the action distribution."""
+
+    def __init__(self, feature_dim, action_space: gym.spaces.Discrete) -> None:
+        """
+        Args:
+            feature dim: Expected output shape of the actor network.
+            action_shape: Expected shape of actions.
+        """
+        super().__init__()
+        self.network = torch.nn.Linear(feature_dim, action_space.n)
+        self.distribution = torch.distributions.categorical.Categorical
+
+    def forward(self, x):
+        logits = self.network(x)
+        return self.distribution(logits=logits)
+
+
+class GaussianPolicyHead(torch.nn.Module):
+    """A module that implements a continuous actor head. It uses the output from the
+    :obj:`actor_net` and state independent learnable parameter :obj:`policy_logstd` to
+    create a :py:class:`~torch.distributions.normal.Normal`  object to compute
+    the action distribution."""
+
+    def __init__(self, feature_dim, action_space: gym.spaces.Box) -> None:
+        """
+        Args:
+            feature dim: Expected output shape of the actor network.
+            action_shape: Expected shape of actions.
+        """
+        super().__init__()
+        self._action_shape = action_space.shape
+        self.policy_mean = torch.nn.Sequential(
+            torch.nn.Linear(feature_dim, np.prod(self._action_shape))
+        )
+        self.policy_logstd = torch.nn.Parameter(
+            torch.zeros(1, np.prod(action_space.shape))
+        )
+        self.distribution = torch.distributions.normal.Normal
+
+    def forward(self, x):
+        _mean = self.policy_mean(x)
+        _std = self.policy_logstd.repeat(x.shape[0], 1).exp()
+        distribution = self.distribution(
+            torch.reshape(_mean, (x.size(0), *self._action_shape)),
+            torch.reshape(_std, (x.size(0), *self._action_shape)),
+        )
+        return distribution
+
+
+class PPOActorCriticNetwork(torch.nn.Module):
+    """A module that implements the TD3 actor computation. It puts together the
+    :obj:`representation_network` and :obj:`actor_net`, and adds a final
+    :py:class:`~torch.nn.Linear` layer to compute the action."""
+
+    def __init__(
+        self,
+        representation_network,
+        actor_net,
+        critic_net,
+        network_output_dim,
+        action_space,
+        continuous_action,
+    ) -> None:
+        super().__init__()
+        self._network = representation_network
+        self._continuous_action = continuous_action
+        if actor_net is None:
+            actor_network = torch.nn.Identity()
+        else:
+            actor_network = actor_net(network_output_dim)
+        feature_dim = np.prod(calculate_output_dim(actor_network, network_output_dim))
+        actor_head = GaussianPolicyHead if self._continuous_action else CategoricalHead
+
+        self.actor = torch.nn.Sequential(
+            actor_network,
+            torch.nn.Flatten(),
+            actor_head(feature_dim, action_space),
+        )
+
+        if critic_net is None:
+            critic_network = torch.nn.Identity()
+        else:
+            critic_network = critic_net(network_output_dim)
+        feature_dim = np.prod(calculate_output_dim(critic_network, network_output_dim))
+        self.critic = torch.nn.Sequential(
+            critic_network,
+            torch.nn.Flatten(),
+            torch.nn.Linear(feature_dim, 1),
+        )
+
+    def forward(self, x, action=None):
+        hidden_state = self._network(x)
+        distribution = self.actor(hidden_state)
+        value = self.critic(hidden_state)
+        if action is None:
+            action = distribution.sample()
+
+        logprob, entropy = distribution.log_prob(action), distribution.entropy()
+        if self._continuous_action:
+            logprob, entropy = logprob.sum(dim=-1), entropy.sum(dim=-1)
+        return action, logprob, entropy, value

hive/configs/atari/ppo.yml

@@ -0,0 +1,64 @@
+run_name: &run_name 'atari-ppo'
+train_steps: 10000000
+test_frequency: 250000
+test_episodes: 10
+max_steps_per_episode: 27000
+stack_size: &stack_size 4
+save_dir: 'experiment'
+saving_schedule:
+  name: 'PeriodicSchedule'
+  kwargs:
+    off_value: False
+    on_value: True
+    period: 1000000
+environment:
+  name: 'AtariEnv'
+  kwargs:
+    env_name: 'Breakout'
+
+agent:
+  name: 'PPOAgent'
+  kwargs:
+    representation_net:
+      name: 'ConvNetwork'
+      kwargs:
+        channels: [32, 64, 64]
+        kernel_sizes: [8, 4, 3]
+        strides: [4, 2, 1]
+        paddings: [2, 2, 1]
+        mlp_layers: [512]
+    optimizer_fn:
+      name: 'Adam'
+      kwargs:
+        lr: .00025
+    init_fn: 
+      name: 'orthogonal'
+    replay_buffer:
+      name: 'PPOReplayBuffer'
+      kwargs:
+        transitions_per_update: 4096
+        stack_size: *stack_size
+        use_gae: True
+        gae_lambda: .95
+    discount_rate: .99
+    grad_clip: .5
+    clip_coef: .1
+    ent_coef: .0
+    clip_vloss: True
+    vf_coef: .5
+    num_epochs_per_update: 4
+    normalize_advantages: True
+    batch_size: 256
+    device: 'cuda'
+    id: 'agent'
+# List of logger configs used.
+loggers:
+  -
+    name: ChompLogger
+  -
+    name: WandbLogger
+    kwargs:
+      project: Hive
+      name: *run_name
+      resume: "allow"
+      start_method: "fork"

hive/configs/gym/ppo.yml

@@ -0,0 +1,64 @@
+run_name: &run_name 'ppo'
+train_steps: 500000
+test_frequency: 1000
+test_episodes: 10
+max_steps_per_episode: 500
+stack_size: &stack_size 1
+save_dir: 'experiment'
+saving_schedule:
+  name: 'PeriodicSchedule'
+  kwargs:
+    off_value: False
+    on_value: True
+    period: 10000
+
+environment:
+  name: 'GymEnv'
+  kwargs:
+    env_name: 'CartPole-v1'
+
+agent:
+  name: 'PPOAgent'
+  kwargs:
+    actor_net:
+      name: 'MLPNetwork'
+      kwargs: 
+        hidden_units: [256, 256]
+    critic_net:
+      name: 'MLPNetwork'
+      kwargs: 
+        hidden_units: [256, 256]
+    replay_buffer:
+      name: 'PPOReplayBuffer'
+      kwargs:
+        transitions_per_update: 2048
+        use_gae: True
+        gae_lambda: .95
+    optimizer_fn:
+      name: 'Adam'
+      kwargs:
+        lr: .00025
+    discount_rate: .99
+    grad_clip: .5
+    clip_coef: .2
+    ent_coef: .01
+    clip_vloss: True
+    vf_coef: .5
+    num_epochs_per_update: 4
+    normalize_advantages: True
+    batch_size: 128
+    device: 'cuda'
+    id: 'agent'
+    init_fn: 
+      name: 'orthogonal'
+# List of logger configs used.
+loggers:
+  -
+    name: ChompLogger
+  -
+    name: WandbLogger
+    kwargs:
+      project: Hive
+      name: *run_name
+      resume: "allow"
+      start_method: "fork"

hive/configs/mujoco/ppo.yml

@@ -0,0 +1,65 @@
+run_name: &run_name 'ppo'
+train_steps: 2000000
+test_frequency: 20000
+test_episodes: 5
+max_steps_per_episode: 100000
+stack_size: &stack_size 1
+save_dir: 'experiment'
+saving_schedule:
+  name: 'PeriodicSchedule'
+  kwargs:
+    off_value: False
+    on_value: True
+    period: 500000
+environment:
+  name: 'GymEnv'
+  kwargs:
+    env_name: 'Hopper-v3'
+    mujoco_wrapper: True
+
+agent:
+  name: 'PPOAgent'
+  kwargs:
+    actor_net:
+      name: 'MLPNetwork'
+      kwargs: 
+        hidden_units: [64, 64]
+    critic_net:
+      name: 'MLPNetwork'
+      kwargs: 
+        hidden_units: [64, 64]
+    replay_buffer:
+      name: 'PPOReplayBuffer'
+      kwargs:
+        transitions_per_update: 2048
+        use_gae: True
+        gae_lambda: .95
+    optimizer_fn:
+      name: 'Adam'
+      kwargs:
+        lr: .0003
+    discount_rate: .99
+    grad_clip: .5
+    clip_coef: .2
+    ent_coef: .0
+    clip_vloss: True
+    vf_coef: .5
+    num_epochs_per_update: 10
+    normalize_advantages: True
+    batch_size: 64
+    device: 'cuda'
+    id: 'agent'
+    init_fn: 
+      name: 'orthogonal'
+
+# List of logger configs used.
+loggers:
+  -
+    name: ChompLogger
+  -
+    name: WandbLogger
+    kwargs:
+      project: Hive
+      name: *run_name
+      resume: "allow"
+      start_method: "fork"

hive/envs/gym_env.py

@@ -1,5 +1,5 @@
 import gym
-
+import numpy as np
 from hive.envs.base import BaseEnv
 from hive.envs.env_spec import EnvSpec
 
@@ -29,7 +29,25 @@ def create_env(self, env_name, **kwargs):
         Args:
             env_name (str): Name of the environment
         """
-        self._env = gym.make(env_name)
+        env = gym.make(env_name)
+        if kwargs.get("mujoco_wrapper", False):
+            env = gym.wrappers.RecordEpisodeStatistics(env)
+            env = gym.wrappers.ClipAction(env)
+            env = gym.wrappers.NormalizeObservation(env)
+            env = gym.wrappers.TransformObservation(
+                env, lambda obs: np.clip(obs, -10, 10)
+            )
+            env = gym.wrappers.NormalizeReward(env)
+            env = gym.wrappers.TransformReward(
+                env, lambda reward: np.clip(reward, -10, 10)
+            )
+
+        elif kwargs.get("atari_wrapper", False):
+            env = gym.wrappers.NoopResetEnv(env, noop_max=30)
+            env = gym.wrappers.EpisodicLifeEnv(env)
+            if "FIRE" in env.unwrapped.get_action_meanings():
+                env = gym.wrappers.FireResetEnv(env)
+        self._env = env
 
     def create_env_spec(self, env_name, **kwargs):
         """Function used to create the specification. Subclasses can override this method

hive/replays/__init__.py

@@ -1,5 +1,6 @@
 from hive.replays.circular_replay import CircularReplayBuffer, SimpleReplayBuffer
 from hive.replays.legal_moves_replay import LegalMovesBuffer
+from hive.replays.ppo_replay import PPOReplayBuffer
 from hive.replays.prioritized_replay import PrioritizedReplayBuffer
 from hive.replays.replay_buffer import BaseReplayBuffer
 from hive.utils.registry import registry
@@ -11,6 +12,7 @@
         "SimpleReplayBuffer": SimpleReplayBuffer,
         "PrioritizedReplayBuffer": PrioritizedReplayBuffer,
         "LegalMovesBuffer": LegalMovesBuffer,
+        "PPOReplayBuffer": PPOReplayBuffer,
     },
 )
 

hive/replays/circular_replay.py

@@ -242,7 +242,7 @@ def sample(self, batch_size):
 
         if self._n_step == 1:
             is_terminal = terminals
-            trajectory_lengths = np.ones(batch_size)
+            trajectory_lengths = np.ones(terminals.shape[0])
         else:
             is_terminal = terminals.any(axis=1).astype(int)
             trajectory_lengths = (