Wandb works, One output dir

lerobot/scripts/visualize_dataset.py

@@ -15,7 +15,15 @@ from lerobot.common.datasets.factory import make_offline_buffer
 
 
 @hydra.main(version_base=None, config_name="default", config_path="../configs")
-def visualize_dataset(cfg: dict):
+def visualize_dataset_cli(cfg: dict):
+    visualize_dataset(
+        cfg, out_dir=hydra.core.hydra_config.HydraConfig.get().runtime.output_dir
+    )
+
+
+def visualize_dataset(cfg: dict, out_dir=None):
+    if out_dir is None:
+        raise NotImplementedError()
 
     sampler = SliceSamplerWithoutReplacement(
         num_slices=1,
@@ -40,10 +48,10 @@ def visualize_dataset(cfg: dict):
             dim=0,
         )
 
-        video_dir = Path(cfg.video_dir)
+        video_dir = Path(out_dir) / "visualize_dataset"
         video_dir.mkdir(parents=True, exist_ok=True)
         # TODO(rcadene): make fps configurable
-        video_path = video_dir / f"eval_episode_{ep_idx}.mp4"
+        video_path = video_dir / f"episode_{ep_idx}.mp4"
 
         assert ep_frames.min().item() >= 0
         assert ep_frames.max().item() > 1, "Not mendatory, but sanity check"
@@ -59,4 +67,4 @@ def visualize_dataset(cfg: dict):
 
 
 if __name__ == "__main__":
-    visualize_dataset()
+    visualize_dataset_cli()