25 lines
830 B
Python
25 lines
830 B
Python
import os
|
|
|
|
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, root=os.environ.get("DATA_DIR"))
|
|
|
|
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']
|