Episodes and datasets

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 simulusion.data.Episode with T frames has T - 1 transitions, the same convention as the simulusion.world_model.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.

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

simulusion.data.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. simulusion.data.episode_from_rlds converts one, given as plain Python and NumPy values (reading the dataset files is left to whatever produced the steps), and simulusion.data.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.