Lerobot/tests/test_visualize_dataset.py
Remi 19812ca470
Add dataset visualization with rerun.io (#131)
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
2024-05-04 16:07:14 +02:00

19 lines
372 B
Python

import pytest
from lerobot.scripts.visualize_dataset import visualize_dataset
@pytest.mark.parametrize(
"repo_id",
["lerobot/pusht"],
)
def test_visualize_dataset(tmpdir, repo_id):
rrd_path = visualize_dataset(
repo_id,
episode_index=0,
batch_size=32,
save=True,
output_dir=tmpdir,
)
assert rrd_path.exists()