Small logging fixes

clean_pufferl.py

@@ -284,6 +284,7 @@ def train(data):
                 losses.old_approx_kl += old_approx_kl.item() / total_minibatches
                 losses.approx_kl += approx_kl.item() / total_minibatches
                 losses.clipfrac += clipfrac.item() / total_minibatches
+                losses.discriminator += disc_loss.item() / total_minibatches
 
         if config.target_kl is not None:
             if approx_kl > config.target_kl:
@@ -423,6 +424,7 @@ def make_losses():
         approx_kl=0,
         clipfrac=0,
         explained_variance=0,
+        discriminator=0,
     )
 
 class Experience:

pufferlib/environments/morph/environment.py

@@ -17,7 +17,7 @@ def make(name, **kwargs):
 
 class PHCPufferEnv(pufferlib.PufferEnv):
     def __init__(self, motion_file, has_self_collision, num_envs=32, device_type="cuda",
-            device_id=0, headless=True, log_interval=128):
+            device_id=0, headless=True, log_interval=32):
         cfg = {
             'env': {
                 'num_envs': num_envs,
@@ -59,6 +59,7 @@ def reset(self, seed=None):
         self.env.reset()
         self.demo = self.env.demo
         self.state = self.env.state
+        self.tick = 0
         return self.observations, []
 
     def step(self, actions_np):
@@ -69,12 +70,9 @@ def step(self, actions_np):
         self.demo = self.env.demo
         self.state = self.env.state
 
-        # NOTE: rl-games reset done envs in the training script. Keeping this here for now.
-        # TODO: Move this into the env
         done_indices = torch.nonzero(self.terminals).squeeze(-1)
         if len(done_indices) > 0:
             self.observations[done_indices] = self.env.reset(done_indices)[done_indices]
-
             self._infos["episode_return"] += self.episode_returns[done_indices].tolist()
             self._infos["episode_length"] += self.episode_lengths[done_indices].tolist()
             self.episode_returns[done_indices] = 0
@@ -84,7 +82,10 @@ def step(self, actions_np):
         self.episode_lengths += 1
 
         # TODO: self.env.extras has infos. Extract useful info?
-        info = self.mean_and_log()
+        info = []
+        self.tick += 1
+        if self.tick % self.log_interval == 0:
+            info = self.mean_and_log()
 
         return self.observations, self.rewards, self.terminals, self.truncations, info