Add custom visualize_dataset.py

2024-06-01 15:45:26 +00:00
parent 042e193995
commit 5b74205e16
6 changed files with 730 additions and 163 deletions
--- a/tests/test_visualize_dataset.py
+++ b/tests/test_visualize_dataset.py
@@ -25,9 +25,8 @@ from lerobot.scripts.visualize_dataset import visualize_dataset
 def test_visualize_dataset(tmpdir, repo_id):
    rrd_path = visualize_dataset(
        repo_id,
-        episode_index=0,
-        batch_size=32,
-        save=True,
+        episode_indices=[0],
        output_dir=tmpdir,
+        serve=False,
    )
    assert rrd_path.exists()