forked from tangger/lerobot
1a1308d62f
add fixes for reproducibility only try to start env if it is closed revision fix normalization and data type Improve README Improve README Tests are passing, Eval pretrained model works, Add gif Update gif Update gif Update gif Update gif Update README Update README update minor Update README.md Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com> Update README.md Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com> Address suggestions Update thumbnail + stats Update thumbnail + stats Update README.md Co-authored-by: Alexander Soare <alexander.soare159@gmail.com> Add more comments Add test_examples.py
23 lines
786 B
Python
23 lines
786 B
Python
from torchrl.data.replay_buffers import SamplerWithoutReplacement
|
|
|
|
import lerobot
|
|
from lerobot.common.datasets.aloha import AlohaDataset
|
|
from lerobot.scripts.visualize_dataset import render_dataset
|
|
|
|
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']
|
|
|
|
# we use this sampler to sample 1 frame after the other
|
|
sampler = SamplerWithoutReplacement(shuffle=False)
|
|
|
|
dataset = AlohaDataset("aloha_sim_transfer_cube_human", sampler=sampler)
|
|
|
|
video_paths = render_dataset(
|
|
dataset,
|
|
out_dir="outputs/visualize_dataset/example",
|
|
max_num_samples=300,
|
|
fps=50,
|
|
)
|
|
print(video_paths)
|
|
# ['outputs/visualize_dataset/example/episode_0.mp4']
|