Render callback is called max_steps-1 times

facebookresearch · Nov 7, 2023 · acc138d · acc138d
1 parent 550c30f
commit acc138d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/benchmarl/experiment/logger.py b/benchmarl/experiment/logger.py
@@ -88,7 +88,6 @@ def log_collection(
         total_frames: int,
         step: int,
     ) -> float:
-
         to_log = {}
         json_metrics = {}
         for group in self.group_map.keys():
@@ -218,7 +217,9 @@ def log_evaluation(
         self.log(to_log, step=step)
         if video_frames is not None:
             vid = torch.tensor(
-                np.transpose(video_frames[: rollouts[0].batch_size[0]], (0, 3, 1, 2)),
+                np.transpose(
+                    video_frames[: rollouts[0].batch_size[0] - 1], (0, 3, 1, 2)
+                ),
                 dtype=torch.uint8,
             ).unsqueeze(0)
             for logger in self.loggers: