Imitation Learning Workflow

Imitation Learning (IL) closes the loop between expert demonstrations and deployable policies. The Space Robotics Bench provides a unified two-command workflow — collect then learn — that works identically across every SRB environment.

1. Collect Demonstrations

Reference: srb agent collect — Collect Demonstrations

Bootstrap a dataset by teleoperating the task yourself. The --success_only flag discards failed episodes so the dataset stays clean:

srb agent collect --env peg_in_hole --num_episodes 50 --success_only

By default the keyboard is used as the input device. For higher-fidelity data, use a 6-DoF SpaceMouse or an XR headset (see Extended Reality):

srb agent collect --env peg_in_hole --teleop_device spacemouse --success_only

Datasets are written to ${SRB_LOGS_DIR}/<env_id>/collect/<timestamp>/datasets/demos_<timestamp>.hdf5 (HDF5), directly consumable by the learn command.

2. Train an IL Policy

Reference: srb agent learn — Learn from Demonstrations

The srb agent learn command runs offline and does not require Isaac Sim, which makes it fast to iterate on:

srb agent learn --algo robomimic_bc --env peg_in_hole \
  --dataset ${SRB_LOGS_DIR}/peg_in_hole/collect/<timestamp>/datasets/demos_<timestamp>.hdf5

Pass the dataset path explicitly with --dataset: srb agent learn creates a fresh run directory and does not auto-discover datasets from prior collect/ runs. The RNN variant works the same way:

srb agent learn --algo robomimic_bc_rnn --env peg_in_hole \
  --dataset ${SRB_LOGS_DIR}/peg_in_hole/collect/<timestamp>/datasets/demos_<timestamp>.hdf5

3. Evaluate the Policy

Reference: srb agent eval — Evaluate Agent

Policies trained by learn share the SRB checkpoint convention, so evaluation is just another srb agent eval invocation:

srb agent eval --algo robomimic_bc --env peg_in_hole env.num_envs=1

Robomimic runs one environment. The integration wrapper and the collection path both reject num_envs != 1 (integration.robomimic.vector_env.unsupported), so vectorized evaluation, collection, and training are unavailable for robomimic_* algorithms.

4. (Optional) Iterate

IL datasets are rarely sufficient on the first try. Two useful follow-ups:

Augment with Policy Rollouts

Use a partially trained policy to generate additional demonstrations — often called DAgger-style data augmentation. Filter by success to keep the dataset clean:

srb agent collect --algo robomimic_bc --env peg_in_hole \
  --num_episodes 200 --success_only

Rerun srb agent learn on the new demonstrations. Since srb agent learn consumes a single --dataset, merge the collection runs’ HDF5 files first to train on their union.

Combine with RL

IL provides a good starting point but rarely produces a policy that generalizes perfectly to unseen conditions. RL checkpoints cannot be warm-started directly from IL checkpoints (the frameworks use incompatible checkpoint formats), but the two paradigms still compose well:

Use the IL policy to collect additional demonstrations (see above) while training an RL agent from scratch on the same task.
Resume an interrupted RL run with srb agent train --continue_training, or pass a same-framework checkpoint explicitly via --model.

5. Deploy to Hardware

Reference: Sim-to-Real Transfer

IL-trained policies go through the same sim-to-real pipeline as RL policies. See the Sim-to-Real Transfer workflow for the real_agent gen → real_agent eval bridge.

Keyboard shortcuts

Space Robotics Bench