Compute aggregated stats during save_episode

huggingface · Feb 14, 2025 · 5ce8756 · 5ce8756
1 parent 7c144fe
commit 5ce8756
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/lerobot/common/datasets/lerobot_dataset.py b/lerobot/common/datasets/lerobot_dataset.py
@@ -260,6 +260,7 @@ def save_episode(
         write_episode(episode_dict, self.root)
 
         self.episodes_stats[episode_index] = episode_stats
+        self.stats = aggregate_stats([self.stats, episode_stats]) if self.stats else episode_stats
         write_episode_stats(episode_index, episode_stats, self.root)
 
     def write_video_info(self) -> None: