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'] # >>> ['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 # 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( video_paths = render_dataset(
dataset, dataset,
@ -139,14 +139,13 @@ video_paths = render_dataset(
max_num_episodes=1, max_num_episodes=1,
) )
print(video_paths) 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: Or you can achieve the same result by executing our script from the command line:
```bash ```bash
python lerobot/scripts/visualize_dataset.py \ python lerobot/scripts/visualize_dataset.py \
env=aloha \ env=pusht \
task=sim_sim_transfer_cube_human \
hydra.run.dir=outputs/visualize_dataset/example hydra.run.dir=outputs/visualize_dataset/example
# >>> ['outputs/visualize_dataset/example/episode_0.mp4'] # >>> ['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'] # >>> ['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 # 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( video_paths = render_dataset(
dataset, dataset,
@ -17,4 +17,4 @@ video_paths = render_dataset(
max_num_episodes=1, max_num_episodes=1,
) )
print(video_paths) print(video_paths)
# ['outputs/visualize_dataset/example/episode_0_top.mp4'] # ['outputs/visualize_dataset/example/episode_0.mp4']