Collect data & train a policy

The imitation-learning loop: demonstrate a task, record episodes, export a LeRobot dataset, train, and run the policy back on the robot

End-to-end: demonstrate a task on an arm, record episodes to a session DB, convert that DB into a LeRobot or HDF5 dataset, train an imitation-learning policy on it, and run that policy back on the robot.

demonstrate ─▶ episodes in a session .db ─▶ dimos dataprep ─▶ LeRobot dataset ─▶ train ─▶ run on the robot
 (teach / VR)     (recorded)                (LeRobot/HDF5)     (LeRobot, ACT)   (policy module)

Each arrow is one command. DimOS ships collection, dataset export, and — on the Galaxea A1Z today — running a trained checkpoint back on hardware. Training itself runs in LeRobot.

1. Record demonstrations

Two supported ways to demonstrate. Whichever you use, three data-quality rules apply:

• Write a real task description (--task / default_task_label). Language-conditioned policies train on this text; a placeholder bakes in uselessness.
• Stay out of the camera frame. The policy trains on exactly what the camera sees; at deployment there is no human in the scene.
• Record the gripper. If the task grasps, the gripper channel must be in the observation and action.

Option A: teach mode (kinesthetic) - recommended

On arms that support it (currently the Galaxea A1Z), put the arm in gravity-compensated zero-force mode and drag it through the task:

uv run dimos a1z teach --task "pick up the object"

Each saved episode contains camera frames plus measured joints and gripper position. The command prints the session .db path. Full details, safety notes, and camera options: Galaxea A1Z.

Option B: Quest VR teleop

Run a collection blueprint. Add --simulation to drive MuJoCo; omit it for real hardware (a RealSense + the arm):

dimos --simulation run learning-collect-quest-xarm7   # XArm7 in sim
dimos run learning-collect-quest-piper                # Piper on real hardware

This brings up teleop, the camera, the episode monitor, and the recorder, wired together.

Button	Action
A (right) / X (left)	Hold to engage - the arm tracks the controller only while held
B	Toggle record - press to start an episode, press again to save it
Y	Discard the in-progress episode

A take is: hold A to move the arm into place, press B to start, perform the task, press B to save (or Y to throw it away). The terminal prints one line per transition:

[collect] ▶ RECORDING episode  (state=recording  saved=0  discarded=0)
[collect] ✓ SAVED episode      (state=idle       saved=1  discarded=0)

End each good take with B before quitting - an episode still recording at shutdown is dropped.

(Keyboard teleop also exists, but do not collect training data with it. See choosing a teleop method.)

Where recordings go

~/.local/state/dimos/recordings/

A new timestamped .db file per session, nothing overwritten. Quest sessions record color_image, coordinator_joint_state, and status (episode start/save/discard markers); the exact path is printed when the recorder starts.

2. Build a dataset

DataPrep is an offline batch step that reads a session DB and writes a dataset. The observation/action stream mapping comes from a JSON config.

For the A1Z, a ready-made config ships in the repo:

uv run dimos dataprep build \
  --source ~/.local/state/dimos/recordings/a1z_teach_<timestamp>.db \
  --output ./a1z_lerobot_dataset \
  --format lerobot \
  --config dimos/learning/dataprep/galaxea_a1z_state_config.json

For other arms, start from dimos/learning/dataprep/example_config.json and edit to taste:

dimos dataprep build \
  --source ~/.local/state/dimos/recordings/session_xarm7_20260622_120000.db \
  --config dimos/learning/dataprep/example_config.json

# HDF5 instead of LeRobot
dimos dataprep build -s <session.db> -c <config.json> -f hdf5

--source / --output / --format override the config, so one config serves many sessions. Inspect the result (features, shapes, dtypes, episode/frame counts):

dimos dataprep inspect ./a1z_lerobot_dataset

Each dataset gets a dimos_meta.json sidecar recording exactly how it was built.

Config fields that matter

• source - the session .db.
• observation / action - map a dataset feature name to a recorded {stream, field}. Action defaults to the next frame’s joint state (action_shift: 1), giving a next-state behavioral-cloning target.
• sync - resample everything onto one timeline: anchor stream, rate_hz, nearest-match tolerance_ms.
• output - format (lerobot | hdf5), path, and metadata (robot, default_task_label, …).

3. Train a policy

The LeRobot dataset is the handoff point. What it unlocks, roughly in order of effort:

• ACT - the reliable baseline. On the order of 50-100 good episodes per task; trains in hours on a single consumer GPU with LeRobot’s standard training scripts.
• SmolVLA - a small language-conditioned VLA, fine-tunable on a consumer GPU.
• pi0 / pi0-FAST - fine-tunable in LeRobot, but wants datacenter-class GPUs.
• GR00T - NVIDIA’s fine-tuning pipeline consumes LeRobot datasets directly.

A completed training run does not by itself mean the policy generalizes - quality tracks demonstration count, consistency, and coverage. The ACT checkpoint you deploy lands at outputs/<job_name>_act/checkpoints/last/pretrained_model.

4. Run the policy on the robot

LeRobotPolicyModule wraps a checkpoint as a DimOS module: camera frames and joint states in, joint commands out, at the policy’s control rate. It never moves hardware at startup - its execute_learned_policy skill starts inference explicitly, and stop_learned_policy halts it with the robot holding position.

On the Galaxea A1Z it is one command:

uv run dimos a1z run-policy \
  outputs/a1z_act/checkpoints/last/pretrained_model \
  --task "pick up the object" \
  --duration 20

Loading and hardware initialization ask for confirmation, and inference starts only once live camera and joint observations are flowing. To run several trained behaviors behind a language agent, see the multi-policy agent blueprint.

For other arms, compose LeRobotPolicyModule into a blueprint next to a camera and the coordinator.

Status: collection and dataset export work on every supported arm. Running a trained checkpoint back on hardware ships today on the A1Z; for other arms, compose LeRobotPolicyModule yourself. Training runs in LeRobot outside of DimOS. The exported dataset remains the stable interface between the two halves.

Notes

• Sim vs real camera - under --simulation the MuJoCo camera supplies color_image; on real hardware a RealSense does. The blueprint picks the right one automatically.
• “action” is the measured next joint state, not a recorded command. For true commanded actions you would record joint_command and map action to it.
• Old sessions - recordings made before the coordinator_joint_state rename use the old stream name; point a matching config at them, or re-record.