PPO

hive/agents/ppo.py

@@ -7,8 +7,11 @@
 
 from hive.agents.agent import Agent
 from hive.agents.qnets.base import FunctionApproximator
-from hive.agents.qnets.normalizer import NormalizationFn
-from hive.agents.qnets.ppo_nets import PPOActorCriticNetwork
+from hive.agents.qnets.normalizer import (
+    MovingAvgNormalizer,
+    RewardNormalizer,
+)
+from hive.agents.qnets.ac_nets import ActorCriticNetwork
 from hive.agents.qnets.utils import (
     InitializationFn,
     calculate_output_dim,
@@ -33,8 +36,8 @@ def __init__(
         init_fn: InitializationFn = None,
         optimizer_fn: OptimizerFn = None,
         critic_loss_fn: LossFn = None,
-        observation_normalization_fn: NormalizationFn = None,
-        reward_normalization_fn: NormalizationFn = None,
+        observation_normalizer: MovingAvgNormalizer = None,
+        reward_normalizer: RewardNormalizer = None,
         stack_size: int = 1,
         replay_buffer: OnPolicyReplayBuffer = None,
         discount_rate: float = 0.99,
@@ -75,9 +78,10 @@ def __init__(
                 If None, defaults to :py:class:`~torch.optim.Adam`.
             critic_loss_fn (LossFn): The loss function used to optimize the critic. If
                 None, defaults to :py:class:`~torch.nn.MSELoss`.
-            observation_normalizer (NormalizationFn): The function for normalizing
-                observations
-            reward_normalizer (NormalizationFn): The function for normalizing rewards
+            observation_normalizer (MovingAvgNormalizer): The function for
+                normalizing observations
+            reward_normalizer (RewardNormalizer): The function for normalizing
+                rewards
             stack_size (int): Number of observations stacked to create the state fed
                 to the agent.
             replay_buffer (OnPolicyReplayBuffer): The replay buffer that the agent will
@@ -123,17 +127,15 @@ def __init__(
             actor_net,
             critic_net,
         )
-        if observation_normalization_fn is not None:
-            self._observation_normalization_fn = observation_normalization_fn(
-                self._state_size
-            )
+        if observation_normalizer is not None:
+            self._observation_normalizer = observation_normalizer(self._state_size)
         else:
-            self._observation_normalization_fn = None
+            self._observation_normalizer = None
 
-        if reward_normalization_fn is not None:
-            self._reward_normalization_fn = reward_normalization_fn(discount_rate)
+        if reward_normalizer is not None:
+            self._reward_normalizer = reward_normalizer(discount_rate)
         else:
-            self._reward_normalization_fn = None
+            self._reward_normalizer = None
 
         if optimizer_fn is None:
             optimizer_fn = torch.optim.Adam
@@ -187,7 +189,7 @@ def create_networks(self, representation_net, actor_net, critic_net):
             network = representation_net(self._state_size)
 
         network_output_shape = calculate_output_dim(network, self._state_size)
-        self._actor_critic = PPOActorCriticNetwork(
+        self._actor_critic = ActorCriticNetwork(
             network,
             actor_net,
             critic_net,
@@ -215,15 +217,15 @@ def preprocess_update_info(self, update_info, agent_traj_state):
             update_info: Contains the information from the current timestep that the
                 agent should use to update itself.
         """
-        if self._observation_normalization_fn:
-            update_info["observation"] = self._observation_normalization_fn(
+        if self._observation_normalizer:
+            update_info["observation"] = self._observation_normalizer(
                 update_info["observation"]
             )
 
         done = update_info["terminated"] or update_info["truncated"]
-        if self._reward_normalization_fn:
-            self._reward_normalization_fn.update(update_info["reward"], done)
-            update_info["reward"] = self._reward_normalization_fn(update_info["reward"])
+        if self._reward_normalizer:
+            self._reward_normalizer.update(update_info["reward"], done)
+            update_info["reward"] = self._reward_normalizer(update_info["reward"])
 
         preprocessed_update_info = {
             "observation": update_info["observation"],
@@ -278,9 +280,9 @@ def act(self, observation, agent_traj_state=None):
         """
         if agent_traj_state is None:
             agent_traj_state = {}
-        if self._observation_normalization_fn:
-            self._observation_normalization_fn.update(observation)
-            observation = self._observation_normalization_fn(observation)
+        if self._observation_normalizer:
+            self._observation_normalizer.update(observation)
+            observation = self._observation_normalizer(observation)
         action, logprob, value = self.get_action_logprob_value(observation)
         agent_traj_state["logprob"] = logprob
         agent_traj_state["value"] = value
@@ -305,8 +307,8 @@ def update(self, update_info, agent_traj_state=None):
         )
 
         if self._replay_buffer.size() >= self._transitions_per_update - 1:
-            if self._observation_normalization_fn:
-                update_info["next_observation"] = self._observation_normalization_fn(
+            if self._observation_normalizer:
+                update_info["next_observation"] = self._observation_normalizer(
                     update_info["next_observation"]
                 )
             _, _, values = self.get_action_logprob_value(
@@ -376,29 +378,33 @@ def update(self, update_info, agent_traj_state=None):
                         approx_kl = ((ratios - 1) - logratios).mean()
 
                     if self._logger.should_log(self._timescale):
-                        self._logger.log_scalar(
-                            "actor_loss", actor_loss, self._timescale
-                        )
-                        self._logger.log_scalar(
-                            "critic_loss", critic_loss, self._timescale
-                        )
-                        self._logger.log_scalar(
-                            "entropy_loss", entr_loss, self._timescale
-                        )
-                        self._logger.log_scalar("approx_kl", approx_kl, self._timescale)
-
-                if self._target_kl is not None:
-                    if approx_kl > self._target_kl:
-                        break
+                        metrics = {
+                            "actor_loss": actor_loss,
+                            "critic_loss": critic_loss,
+                            "entropy_loss": entr_loss,
+                            "approx_kl": approx_kl,
+                        }
+                        self._logger.log_metrics(metrics, prefix=self._timescale)
+                if self._target_kl is not None and self._target_kl < approx_kl:
+                    break
             self._replay_buffer.reset()
         return agent_traj_state
 
     def save(self, dname):
-        torch.save(
+        state_dict = (
             {
                 "actor_critic": self._actor_critic.state_dict(),
                 "optimizer": self._optimizer.state_dict(),
             },
+        )
+        if self._observation_normalizer:
+            state_dict[
+                "observation_normalizer"
+            ] = self._observation_normalizer.state_dict()
+        if self._reward_normalizer:
+            state_dict["reward_normalizer"] = self._reward_normalizer.state_dict()
+        torch.save(
+            state_dict,
             os.path.join(dname, "agent.pt"),
         )
         replay_dir = os.path.join(dname, "replay")
@@ -410,3 +416,9 @@ def load(self, dname):
         self._actor_critic.load_state_dict(checkpoint["actor_critic"])
         self._optimizer.load_state_dict(checkpoint["optimizer"])
         self._replay_buffer.load(os.path.join(dname, "replay"))
+        if self._observation_normalizer:
+            self._observation_normalizer.load_state_dict(
+                checkpoint["observation_normalizer"]
+            )
+        if self._reward_normalizer:
+            self._reward_normalizer.load_state_dict(checkpoint["reward_normalizer"])

hive/agents/qnets/ppo_nets.py → hive/agents/qnets/ac_nets.py

@@ -64,8 +64,8 @@ def forward(self, x):
         return distribution
 
 
-class PPOActorCriticNetwork(torch.nn.Module):
-    """A module that implements the PPO actor and critic computation. It puts together
+class ActorCriticNetwork(torch.nn.Module):
+    """A module that implements the actor and critic computation. It puts together
     the :obj:`representation_network`, :obj:`actor_net` and :obj:`critic_net`, then
     adds two final :py:class:`~torch.nn.Linear` layers to compute the action and state
     value."""

hive/agents/qnets/normalizer.py

@@ -1,3 +1,4 @@
+import abc
 from typing import Tuple
 
 import numpy as np
@@ -46,24 +47,44 @@ def update_mean_var_count_from_moments(
 
         return new_mean, new_var, new_count
 
+    def state_dict(self):
+        """Returns the state as a dictionary."""
+        return {"mean": self.mean, "var": self.var, "count": self.count}
 
-class BaseNormalizationFn(object):
-    """Implements the base normalization function."""
+    def load_state_dict(self, state_dict):
+        """Loads the state from a dictionary."""
+        self.mean = state_dict["mean"]
+        self.var = state_dict["var"]
+        self.count = state_dict["count"]
 
-    def __init__(self, *args, **kwds):
-        pass
 
-    def __call__(self, *args, **kwds):
-        return NotImplementedError
+class Normalizer(Registrable):
+    """A wrapper for callables that produce normalization functions.
 
-    def update(self, *args, **kwds):
-        return NotImplementedError
+    These wrapped callables can be partially initialized through configuration
+    files or command line arguments.
+    """
 
+    @classmethod
+    def type_name(cls):
+        """
+        Returns:
+            "norm_fn"
+        """
+        return "norm_fn"
+
+    @abc.abstractmethod
+    def state_dict(self):
+        """Returns the state of the normalizer as a dictionary."""
 
-class ObservationNormalizationFn(BaseNormalizationFn):
-    """Implements a normalization function. Transforms output by
-    normalising the input data by the running :obj:`mean` and
-    :obj:`std`, and clipping the normalised data on :obj:`clip`
+    def load_state_dict(self, state_dict):
+        """Loads the normalizer state from a dictionary."""
+
+
+class MovingAvgNormalizer(Normalizer):
+    """Implements a moving average normalization and clipping function. Normalizes
+    input data with the running mean and std. The normalized data is then clipped
+    within the specified range.
     """
 
     def __init__(
@@ -76,26 +97,35 @@ def __init__(
             clip (np.float32): The clip value for the normalised data.
         """
         super().__init__()
-        self.obs_rms = MeanStd(epsilon, shape)
+        self._rms = MeanStd(epsilon, shape)
         self._shape = shape
         self._epsilon = epsilon
         self._clip = clip
 
-    def __call__(self, obs):
-        obs = np.array([obs])
-        obs = ((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self._epsilon))[0]
+    def __call__(self, input_data):
+        input_data = np.array([input_data])
+        input_data = (
+            (input_data - self._rms.mean) / np.sqrt(self._rms.var + self._epsilon)
+        )[0]
         if self._clip is not None:
-            obs = np.clip(obs, -self._clip, self._clip)
-        return obs
+            input_data = np.clip(input_data, -self._clip, self._clip)
+        return input_data
+
+    def update(self, input_data):
+        self._rms.update(input_data)
 
-    def update(self, obs):
-        self.obs_rms.update(obs)
+    def state_dict(self):
+        return self._rms.state_dict()
 
+    def load_state_dict(self, state_dict):
+        self._rms.load_state_dict(state_dict)
 
-class RewardNormalizationFn(BaseNormalizationFn):
-    """Implements a normalization function. Transforms output by
-    normalising the input data by the running :obj:`mean` and
-    :obj:`std`, and clipping the normalised data on :obj:`clip`
+
+class RewardNormalizer(Normalizer):
+    """Normalizes and clips rewards from the environment. Applies a discount-based
+    scaling scheme, where the rewards are divided by the standard deviation of a
+    rolling discounted sum of the rewards. The scaled rewards are then clipped within
+    specified range.
     """
 
     def __init__(self, gamma: float, epsilon: float = 1e-4, clip: np.float32 = np.inf):
@@ -106,48 +136,41 @@ def __init__(self, gamma: float, epsilon: float = 1e-4, clip: np.float32 = np.in
             clip (np.float32): The clip value for the normalised data.
         """
         super().__init__()
-        self.return_rms = MeanStd(epsilon, ())
+        self._return_rms = MeanStd(epsilon, ())
         self._epsilon = epsilon
         self._clip = clip
         self._gamma = gamma
         self._returns = np.zeros(1)
 
     def __call__(self, rew):
         rew = np.array([rew])
-        rew = (rew / np.sqrt(self.return_rms.var + self._epsilon))[0]
+        rew = (rew / np.sqrt(self._return_rms.var + self._epsilon))[0]
         if self._clip is not None:
             rew = np.clip(rew, -self._clip, self._clip)
         return rew
 
     def update(self, rew, done):
         self._returns = self._returns * self._gamma + rew
-        self.return_rms.update(self._returns)
+        self._return_rms.update(self._returns)
         self._returns *= 1 - done
 
+    def state_dict(self):
+        state_dict = self._return_rms.state_dict()
+        state_dict["returns"] = self._returns
+        return state_dict
 
-class NormalizationFn(Registrable):
-    """A wrapper for callables that produce normalization functions.
-
-    These wrapped callables can be partially initialized through configuration
-    files or command line arguments.
-    """
-
-    @classmethod
-    def type_name(cls):
-        """
-        Returns:
-            "norm_fn"
-        """
-        return "norm_fn"
+    def load_state_dict(self, state_dict):
+        self._returns = state_dict["returns"]
+        state_dict.pop("returns")
+        self._return_rms.load_state_dict(state_dict)
 
 
 registry.register_all(
-    NormalizationFn,
+    Normalizer,
     {
-        "BaseNormalization": BaseNormalizationFn,
-        "RewardNormalization": RewardNormalizationFn,
-        "ObservationNormalization": ObservationNormalizationFn,
+        "RewardNormalizer": RewardNormalizer,
+        "MovingAvgNormalizer": MovingAvgNormalizer,
     },
 )
 
-get_norm_fn = getattr(registry, f"get_{NormalizationFn.type_name()}")
+get_norm_fn = getattr(registry, f"get_{Normalizer.type_name()}")

hive/configs/mujoco/ppo.yml

@@ -38,12 +38,12 @@ kwargs:
         name: 'MLPNetwork'
         kwargs: 
           hidden_units: [64, 64]
-      observation_normalization_fn:
-        name: 'ObservationNormalization'
+      observation_normalizer:
+        name: 'MovingAvgNormalizer'
         kwargs:
           clip: 10
-      reward_normalization_fn:
-        name: 'RewardNormalization'
+      reward_normalizer:
+        name: 'RewardNormalizer'
         kwargs:
           clip: 10
       replay_buffer:

hive/envs/gym_env.py

@@ -46,7 +46,7 @@ def create_env(self, env_name, env_wrappers, **kwargs):
         Args:
             env_name (str): Name of the environment
         """
-        self._env = gym.make(env_name)
+        self._env = gym.make(env_name, **kwargs)
 
         if env_wrappers is not None:
             self._env = apply_wrappers(self._env, env_wrappers)