Skip to content

Record episodes and export a dataset

Record what a running workcell does into episodes, then turn finished episodes into a LeRobot dataset to train a policy on. This runs in the sim as it is; on a bench, the runtime and the cameras are the ones your script already has.

"""Records two episodes of an arm in the sim, with its wrist camera, and exports them as one LeRobot dataset."""

import tempfile
from pathlib import Path

import anvil.camera
import anvil.control
import anvil.hw
import anvil.recorder
import anvil.sim
import anvil.workcell

# The single-arm yam with the camera on its wrist, as a bench that has one would save it.
workcell = anvil.workcell.WorkcellConfig.OPENYAM_V1_SINGLE
workcell.cameras.append(
    anvil.workcell.CameraConfig(
        name="wrist",
        device="",
        mode=anvil.camera.Mode(width_px=640, height_px=480, fps=30),
        link="wrist_cam",
    )
)
robot = workcell.robot
# Where the episodes go: a scratch directory here, wherever your data lives on a bench.
data_root = Path(tempfile.mkdtemp())

with (
    anvil.sim.World(robot=robot, scene="table") as world,
    anvil.camera.Cameras(workcell=workcell, world=world) as cameras,
):
    hw = anvil.hw.SimHw(world=world)
    hw.enable()
    trajectory_controller = anvil.control.JointTrajectoryController(robot=robot)
    with anvil.control.Runtime(hw=hw, controllers=[trajectory_controller]) as runtime:
        runtime.start()
        handle = runtime.controller(kind=anvil.control.JointTrajectoryController)
        runtime.activate(controller=handle)

        session = anvil.recorder.Session(root=data_root, task="home the arm and fold it back")
        recorder = anvil.recorder.Recorder(workcell=workcell, runtime=runtime, cameras=cameras)
        for _ in range(2):
            recorder.start(session=session)
            handle.play(trajectory=robot.home_trajectory(from_q=runtime.latest().q), timeout_s=30.0)
            handle.play(trajectory=robot.dehome_trajectory(from_q=runtime.latest().q), timeout_s=30.0)
            stats = recorder.stop(outcome=anvil.recorder.Outcome.SUCCESS)
            print(f"recorded an episode: {stats.frames_received['wrist']} frames, {stats.ticks_missed} ticks missed")

# Encoding the video is the slow part; H.264 is several times faster than LeRobot's default AV1.
episodes = sorted(session.path.glob("*.mcap"))
export = anvil.recorder.DatasetExport(
    episodes=episodes, out_dir=data_root / "dataset", codec=anvil.recorder.VideoCodec.H264
)
dataset = export.run()
print(f"exported {dataset.num_episodes} episodes, {dataset.num_frames} frames at {dataset.fps} fps to {dataset.path}")

Record

  1. Open a session. A Session is a directory under a data root you choose, under a name you give or one generated for it. Every episode recorded into it is its next file, 0001.mcap, 0002.mcap, each with the task written into it.
  2. Build a recorder over the workcell, the running runtime and the cameras to record, as Cameras opens them. On real arms, pass the Report the bench came up on as readiness, so each episode says which physical drives it was recorded on.
  3. Start, do the task, stop. start returns at once and records on threads of its own; the loop never waits on it. stop finishes the file and returns what it counted. Say how it went as you stop it, or label it later; a failed episode is kept, labeled, not deleted. Session.discard_last_episode deletes one you do not want.

By default each source is recorded at its own rate, the loop at 100 Hz and each camera at its frame rate; RecorderConfig says otherwise. An episode opens in Foxglove as it is, with the layout written beside it.

Export

DatasetExport writes the episodes, in the order given, into a new dataset directory that LeRobot opens with LeRobotDataset(root=...). The episodes must be of one robot, the same arms with the same cameras at the same rates; a retune or a homing run taught between them is no matter. The dataset runs at the anchor camera's rate, the first camera by name unless you name one: each frame is the anchor's, carrying each actuator's measured position as observation.state and its commanded position as action. Frames where the loop commanded no position, as while no controller is active, are left out.

Encoding the video is the slow part, several core-seconds per second of video per camera, and needs ffmpeg on PATH. It runs at the lowest priority; on the robot, call cancel from another thread the moment the robot needs the CPU, or export on another machine.