Integration — Imitation Learning

The Space Robotics Bench supports Imitation Learning (IL) as a complementary paradigm to Reinforcement Learning. IL is particularly useful in space robotics where reward engineering is difficult, where sparse-reward tasks stall standard RL, or where expert trajectories (from teleoperation, planners, or prior policies) are available to bootstrap learning.

Motivation

Several characteristics of space robotics tasks make IL an attractive option:

• Reward sparsity - Assembly, sample collection, and docking tasks often have reward signals that only appear at the very end of an episode, making RL credit assignment slow.
• Safety-critical actions - Demonstrations encode a known-safe envelope of behaviors, reducing the risk of catastrophic exploration on expensive hardware.
• Low-latency deployment - A policy distilled from demonstrations can be deployed to flight-grade hardware without the exploration noise that RL requires.
• Hybrid workflows - IL policies can bootstrap RL indirectly — e.g. by rolling out an IL policy to collect additional demonstrations or by shaping curricula — even though RL checkpoints cannot be warm-started directly from IL checkpoints (the frameworks use different checkpoint formats).

Available Integrations

SRB integrates the following IL stack:

• robomimic [PyTorch] - Reference implementation of BC (MLP/RNN/Transformer), BCQ, CQL, IQL, TD3-BC, GL, HBC, IRIS, and Diffusion Policy with a standard HDF5 dataset format.

All other frameworks already integrated for RL (Dreamer, TD-MPC2, Stable-Baselines3, SBX, skrl, RSL-RL) can also consume SRB-collected datasets for offline or demonstration-guided training, typically by using their replay buffer loading utilities.

Data Format

Demonstrations are stored in the robomimic-compatible HDF5 layout:

• Per-episode groups containing obs/, next_obs/, actions, rewards, and dones datasets (plus a per-episode success attribute).
• Observation keys match the SRB observation group structure (STATE, PROPRIO, COMMAND, STATE_PRIVILEGED, etc.).
• Success flags are stored per-episode to support --success_only filtering at collection time.

Datasets written by srb agent collect are directly loadable by srb agent learn without any conversion step.

How to Get Started?

1. Installation
2. Basic Usage
3. Imitation Learning Workflow

See Also

• srb agent collect — record demonstrations
• srb agent learn — train an IL policy
• Reinforcement Learning Integration — the complementary paradigm