Use pusht as example so it's fast

This commit is contained in:
Cadene 2024-04-11 13:36:50 +00:00
parent 2f4af32d3f
commit a605eec7e9
2 changed files with 5 additions and 6 deletions

View File

@ -131,7 +131,7 @@ print(lerobot.available_datasets)
# >>> ['aloha_sim_insertion_human', 'aloha_sim_insertion_scripted', 'aloha_sim_transfer_cube_human', 'aloha_sim_transfer_cube_scripted', 'pusht', 'xarm_lift_medium']
# TODO(rcadene): remove DATA_DIR
dataset = AlohaDataset("aloha_sim_transfer_cube_human", root=Path(os.environ.get("DATA_DIR")))
dataset = AlohaDataset("pusht", root=Path(os.environ.get("DATA_DIR")))
video_paths = render_dataset(
dataset,
@ -139,14 +139,13 @@ video_paths = render_dataset(
max_num_episodes=1,
)
print(video_paths)
# ['outputs/visualize_dataset/example/episode_0_top.mp4']
# ['outputs/visualize_dataset/example/episode_0.mp4']
```
Or you can achieve the same result by executing our script from the command line:
```bash
python lerobot/scripts/visualize_dataset.py \
env=aloha \
task=sim_sim_transfer_cube_human \
env=pusht \
hydra.run.dir=outputs/visualize_dataset/example
# >>> ['outputs/visualize_dataset/example/episode_0.mp4']
```

View File

@ -9,7 +9,7 @@ print(lerobot.available_datasets)
# >>> ['aloha_sim_insertion_human', 'aloha_sim_insertion_scripted', 'aloha_sim_transfer_cube_human', 'aloha_sim_transfer_cube_scripted', 'pusht', 'xarm_lift_medium']
# TODO(rcadene): remove DATA_DIR
dataset = AlohaDataset("aloha_sim_transfer_cube_human", root=Path(os.environ.get("DATA_DIR")))
dataset = AlohaDataset("pusht", root=Path(os.environ.get("DATA_DIR")))
video_paths = render_dataset(
dataset,
@ -17,4 +17,4 @@ video_paths = render_dataset(
max_num_episodes=1,
)
print(video_paths)
# ['outputs/visualize_dataset/example/episode_0_top.mp4']
# ['outputs/visualize_dataset/example/episode_0.mp4']