Adds flag for per-batch advantage normalization (#68)

* Adds flag for per-batch advantage normalization

* make default false for backwards

Author: Mayankm96

config/dummy_config.yaml

@@ -1,6 +1,8 @@
 algorithm:
   class_name: PPO
   # training parameters
+  # -- advantage normalization
+  normalize_advantage_per_mini_batch: false
   # -- value function
   value_loss_coef: 1.0
   clip_param: 0.2

rsl_rl/algorithms/ppo.py

@@ -38,6 +38,7 @@ def __init__(
         schedule="fixed",
         desired_kl=0.01,
         device="cpu",
+        normalize_advantage_per_mini_batch=False,
         # RND parameters
         rnd_cfg: dict | None = None,
         # Symmetry parameters
@@ -48,6 +49,7 @@ def __init__(
         self.desired_kl = desired_kl
         self.schedule = schedule
         self.learning_rate = learning_rate
+        self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch
 
         # RND components
         if rnd_cfg is not None:
@@ -84,8 +86,10 @@ def __init__(
         # PPO components
         self.actor_critic = actor_critic
         self.actor_critic.to(self.device)
-        self.storage = None  # initialized later
+        # Create optimizer
         self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)
+        # Create rollout storage
+        self.storage: RolloutStorage = None  # type: ignore
         self.transition = RolloutStorage.Transition()
 
         # PPO parameters
@@ -168,7 +172,9 @@ def process_env_step(self, rewards, dones, infos):
     def compute_returns(self, last_critic_obs):
         # compute value for the last step
         last_values = self.actor_critic.evaluate(last_critic_obs).detach()
-        self.storage.compute_returns(last_values, self.gamma, self.lam)
+        self.storage.compute_returns(
+            last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch
+        )
 
     def update(self):  # noqa: C901
         mean_value_loss = 0
@@ -213,6 +219,11 @@ def update(self):  # noqa: C901
             # original batch size
             original_batch_size = obs_batch.shape[0]
 
+            # check if we should normalize advantages per mini batch
+            if self.normalize_advantage_per_mini_batch:
+                with torch.no_grad():
+                    advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8)
+
             # Perform symmetric augmentation
             if self.symmetry and self.symmetry["use_data_augmentation"]:
                 # augmentation using symmetry

rsl_rl/storage/rollout_storage.py

@@ -129,7 +129,7 @@ def _save_hidden_states(self, hidden_states):
     def clear(self):
         self.step = 0
 
-    def compute_returns(self, last_values, gamma, lam):
+    def compute_returns(self, last_values, gamma, lam, normalize_advantage: bool = True):
         advantage = 0
         for step in reversed(range(self.num_transitions_per_env)):
             # if we are at the last step, bootstrap the return value
@@ -146,9 +146,12 @@ def compute_returns(self, last_values, gamma, lam):
             # Return: R_t = A(s_t, a_t) + V(s_t)
             self.returns[step] = advantage + self.values[step]
 
-        # Compute and normalize the advantages
+        # Compute the advantages
         self.advantages = self.returns - self.values
-        self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8)
+        # Normalize the advantages if flag is set
+        # This is to prevent double normalization (i.e. if per minibatch normalization is used)
+        if normalize_advantage:
+            self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8)
 
     def get_statistics(self):
         done = self.dones