Your Own Data: Episodes and LeRobot#
Every dataset-driven task reads one format: a folder of episodes. Video, reinforcement-learning trajectories and robot demonstrations all fit it, so one dataset can train Genie, Dreamer’s or IRIS’s world model, or World Models’ VAE and MDN-RNN.
The episode#
An Episode with T frames has T - 1
transitions, the same convention as the
WorldModel protocol: actions[t] is the
action taken at frames[t], which led to frames[t + 1] and earned
rewards[t]; terminated[t] says frames[t + 1] ended the episode and
truncated[t] that it was cut off (a time limit). Only frames is required.
|
Contents |
Required |
|---|---|---|
|
|
yes |
|
|
no |
|
|
no |
|
|
no |
|
|
no |
|
|
no |
|
a task instruction |
no |
import numpy as np
from simulusion.data import Episode, EpisodeWriter
with EpisodeWriter("data/mine") as writer: # numbers files, writes meta.json
writer.write(
Episode(
frames=np.zeros((50, 64, 64, 3), np.uint8),
actions=np.zeros((49, 7), np.float32),
observations={"state": np.zeros((50, 9), np.float32)},
language="pick up the cup",
)
)
A folder may also hold loose videos (.mp4, .webm, …) and .npy frame
arrays, read as frame-only episodes, so a directory of clips is already a
dataset for Genie. meta.json records the action space, cameras, frame rate
and source; tasks read the action count from it.
Training windows#
SequenceDataset cuts every sequence_length
window that fits inside an episode (never across two) and returns a dict of
tensors: frames (L, C, H, W) uint8, and when present actions,
rewards, terminated, truncated (L - 1 each), extra cameras and
observations. Tasks configure it through their data section:
data:
path: data/pusht
camera: null # which view to train on, for multi-camera data
sequence_length: null # the model's own context when null
image_size: null # the model's own input size when null
batch_size: 16
val_fraction: 0.05 # held-out episodes for evaluation
LeRobot#
Local LeRobot dataset folders (v2.x
and v3.0) are read directly, without the lerobot package
(pip install "simulusion[lerobot]" adds PyArrow and PyAV). Convert once:
simulusion datasets import lerobot datasets/pusht data/pusht
simulusion datasets import lerobot ./my_dataset data/mine --camera top --all-cameras
simulusion datasets info data/pusht
simulusion train --config dreamer-wm data.path=data/pusht
or train on it in place with data.source=lerobot data.path=datasets/pusht,
or read it from Python:
from simulusion.data import LeRobotEpisodes
episodes = LeRobotEpisodes("datasets/pusht", camera="observation.image")
episodes[0].actions.shape # (T - 1, 2)
The chosen camera becomes frames and the others cameras; action
becomes actions without its last row (it leads past the final frame);
next.reward and next.done become rewards and terminated; numeric
observation.* features become observations (observation.state is
observations["state"]); the episode’s task is language. Cameras stored as
videos or as images in the Parquet tables both work.
RLDS episodes#
RLDS, the format of Open X-Embodiment and many other robot datasets, stores
an episode as a list of steps. episode_from_rlds()
converts one, given as plain Python and NumPy values (reading the dataset
files is left to whatever produced the steps), and
RLDSKeys says where the main camera, extra
cameras, state, actions and instruction live in a step:
from simulusion.data import EpisodeWriter, RLDSKeys, episode_from_rlds
keys = RLDSKeys(image="image", state=("state",), language="language_instruction")
with EpisodeWriter("data/robot") as writer:
for episode in my_rlds_episodes: # each {"steps": [...]}
writer.write(episode_from_rlds(episode, keys))
Dictionary actions (the RT-1 format) are flattened in RLDSKeys.action
order. RLDS stores an action in every step, so the last one (past the final
observation) is dropped; is_terminal on the final step marks the last
transition terminated, and an episode that ends without it is truncated.