srb agent collect — Collect Demonstrations
The srb agent collect command records episodes from the environment as a demonstration dataset. Data can come from human teleoperation, a trained policy, or built-in random/zero action sources. Collected datasets are consumed by srb agent learn for offline Imitation Learning (IL).
Usage
srb agent collect --env ENV_ID [options]
Options | Shared Agent Options
| Option | Description | Default |
|---|---|---|
| Collection | ||
--collector SOURCE | Data source: human, policy, random, or zero. If omitted, SRB uses policy when --algo or --model is present and human otherwise. | inferred |
--num_episodes N | Number of episodes to collect | 100 |
--horizon N | Maximum number of steps per episode | 400 |
--success_only | Only save episodes that end in task success | False |
| Teleop (default mode) | ||
--teleop_device [DEV …] | Input device(s) (keyboard, spacemouse, ros, gamepad, haptic) | env default, else [keyboard] |
--pos_sensitivity VAL | Translation sensitivity | 1.0 |
--rot_sensitivity VAL | Rotation sensitivity | π (≈3.14) |
--invert_controls | Invert control schema | True |
| Policy rollout | ||
--algo ALGO | RL algorithm of the trained policy | |
--model MODEL | Path to a raw policy checkpoint or SRB model artifact | |
--obs {state,visual} | Observation modality of the rolled-out policy: selects the agent-config variant and the log tree searched for its latest checkpoint | state |
Collection Modes
1. Teleoperation (default)
With no --algo or --model, the command launches a teleoperation session. Every step recorded contains the human-generated action together with the observation. Ideal for bootstrapping a dataset when no policy exists yet.
srb agent collect --env peg_in_hole --num_episodes 50 --success_only
2. Policy Rollout
Provide --algo (and optionally --model) to replay a previously trained RL policy and record its trajectories. Supported algorithms: dreamer, tdmpc2, rsl_rl_*, sb3_*, sbx_*, skrl_*, robomimic_*.
srb agent collect --env peg_in_hole --algo sbx_ppo --num_episodes 200 --success_only
If --model is omitted, the latest checkpoint from the standard training logdir is loaded automatically. For policy collection, SRB reuses the latest run directory for that algorithm so framework-native checkpoint layouts such as model_*.pt (RSL-RL) and models/*.pt (TD-MPC2) resolve the same way as evaluation.
--model may point either to a raw checkpoint or to an SRB model artifact directory. Portable artifacts can carry framework restore config (config.yaml) beside the checkpoint; this is used by stateful or config-heavy adapters such as Dreamer, TD-MPC2, and RSL-RL when collection is replayed outside the original training run.
When collection is driven by an SRB model artifact, dataset metadata records both the resolved checkpoint path (model) and the source artifact directory (model_artifact).
3. Built-in Sources
Use --collector random or --collector zero to collect baseline datasets without a human operator or trained checkpoint.
srb agent collect --env peg_in_hole --collector random --num_episodes 20
srb agent collect --env excavation --collector zero --num_episodes 5 --horizon 200
These datasets keep the collection source separate from the dataset split, so real-world collection can still use split=real while preserving controller=random, controller=zero, controller=human, or controller=policy in artifact metadata.
Output
Datasets are written as HDF5 files compatible with robomimic. Human, random, and zero collection use ${SRB_LOGS_DIR}/<env_id>/collect/<timestamp>/datasets/; policy collection writes under the resolved policy run directory, for example ${SRB_LOGS_DIR}/<env_id>/<algo>/<timestamp>/datasets/. Run metadata records the collection source, episode controls, teleop devices, and a reproduction command. The dataset path is printed at the end of the run and can be passed directly to srb agent learn --dataset.
Tips
- Use
--success_onlywhen training IL — a few clean successes outperform many failed attempts. - Collection requires
env.num_envs=1in every mode. Human, policy, random, and zero collection all raiserobomimic collection requires num_envs == 1before any dataset file is created, so set the Hydra override explicitly. - For an immersive viewpoint while teleoperating, combine with
--xr(see Extended Reality).
See Also
srb agent learn— train a policy on the collected dataset- Imitation Learning Workflow
- Imitation Learning Integration