Failure-Aware Imitation Learning with Residual Action Correction
Behavior Cloning + Recovery & Correction (BCRAC) — a closed-loop system that learns from its own failures.
A complete robotics ML pipeline that goes beyond vanilla behavior cloning. We train a base imitation policy, identify where and why it fails, then surgically fix those failures through residual correction — all in an automated loop.
The core idea: don't just clone expert behavior. Clone it, break it under stress, cluster the failures, and fine-tune a correction layer that activates only when the base policy is uncertain.
Expert Demos ──→ BC Base Policy ──→ Stress Rollouts ──→ Failure Mining
│
┌─────────────────────────────────────────┘
▼
Cluster Failures ──→ Auto-Label ──→ Target Hardest Clusters
│
┌─────────────────────────────────────────┘
▼
Recovery Fine-tune ──→ Eval (Pre/Post) ──→ Report + Dashboard
│
└──→ Repeat (Flywheel)
The policy itself is a BCRAC (Behavior Cloning + Residual Action Correction) network:
┌─────────────────────┐
observation ──→ │ Base MLP (BC) │──→ a_base
+ │ │
stage_id ──→ │ Residual MLPs (3) │──→ a_residual
│ + Manager Head │
│ + Stage Gate │
│ + Uncertainty Gate │──→ g(u) · g(sw) · α · a_residual
└─────────────────────┘
│
a_final = a_base + gated_residual
Key components:
- MC-Dropout uncertainty — run the base policy N times with dropout on, measure variance. High variance = the policy doesn't know what to do here.
- Uncertainty gate
g(u) = σ((u - threshold) · temperature)— only apply corrections when the base policy is confused. - Switch-point learning — a learned binary classifier predicting "should the residual be active?" trained on base-policy error signal.
- Hierarchical residual — 3 sub-controllers with a softmax manager + stage-conditioned gating.
- Safety envelope — joint/velocity limits, obstacle checks, action clipping with full diagnostics.
Evaluated under stress conditions (randomized initial states, mid-episode disturbances, sensor noise, object offsets):
| Metric | Before | After | Delta |
|---|---|---|---|
| Overall Success | 80.0% | 91.7% | +11.7% |
| Push Success | 80.0% | 100.0% | +20.0% |
| Drawer Success | 100.0% | 100.0% | +0.0% |
| Sequential | 60.0% | 75.0% | +15.0% |
| Robustness | 0.800 | 0.917 | +0.117 |
The flywheel automatically identified the 2 hardest failure clusters and fine-tuned on them:
| Cluster | Label | Hardness | Before | After |
|---|---|---|---|---|
| cluster_0 | [push] off-axis stable-gated |
0.828 | 100% | 100% |
| cluster_6 | [push] off-axis stable-gated |
0.997 | 50% | 100% |
Cluster 6 (the hardest cluster, high uncertainty push failures) went from 50% to 100% success after targeted fine-tuning.
Failure cluster structure in PCA space — 8 clusters across push and drawer tasks
Uncertainty-gate correlation (r=0.996) — the gate activates precisely when the policy is uncertain
Uncertainty vs error proxy — high uncertainty correlates with high-error states
Residual action magnitudes by cluster — different clusters trigger different correction intensities
Success rate over flywheel iterations
Stress profile: state_randomization=medium, object_offset=0.2, action_noise=0.08, sensor_noise=0.02, base_degrade=0.08
Cluster Selection: hardest (k=2)
Selected: [cluster_0, cluster_6] — hardness threshold = 0.783 (67th percentile)
Cluster Labels:
cluster_0 (73 samples): [push] off-axis stable-gated | hardness=0.828 | FAILURE
cluster_1 (2361 samples): [push] off-axis high-uncertainty | hardness=0.733 | nominal
cluster_2 (49 samples): [push] off-axis high-uncertainty | hardness=0.806 | FAILURE
cluster_3 (1004 samples): [drawer] drawer-unstable | hardness=0.304 | nominal
cluster_4 (926 samples): [drawer] drawer-unstable | hardness=0.002 | nominal
cluster_5 (985 samples): [drawer] drawer-unstable | hardness=0.010 | nominal
cluster_6 (57 samples): [push] off-axis stable-gated | hardness=0.997 | FAILURE
cluster_7 (1242 samples): [drawer] drawer-unstable | hardness=0.308 | nominal
Recovery Guardrails:
Nominal validation samples: 2,149
Best nominal val loss: 0.4917
LR backoff factor: 0.5 on nominal degradation > 10%
| Component | What it does |
|---|---|
lerobot-collect-bcrac-mujoco-demos |
Collect expert demos from MetaWorld tasks |
lerobot-train-bcrac |
Train the BCRAC policy on demo data |
lerobot-benchmark-bcrac-mujoco |
Evaluate with perturbations, mid-episode disturbances, recovery tracking |
lerobot-failure-miner |
Cluster failure states, auto-label, compute hardness scores |
lerobot-train-bcrac-recovery |
Fine-tune on targeted failure clusters with nominal validation guardrails |
lerobot-eval-flywheel |
Before/after evaluation with cluster-conditioned metrics |
lerobot-behavior-flywheel |
End-to-end orchestrator (rollout → mine → finetune → eval → report) |
lerobot-ablation-study |
Train & eval five ablated variants (BC → Full model) |
lerobot-data-efficiency |
Sweep over 50/100/200 demos to show data efficiency |
lerobot-dashboard |
Live web dashboard with charts, search, job runner |
lerobot-behavior-search |
CLI for searching episodes by task, cluster, uncertainty |
lerobot-behavior-plots |
Generate diagnostic plots from behavior logs |
# Install
pip install -e .
# Collect 100 demos per task
lerobot-collect-bcrac-mujoco-demos --episodes-per-task 100 --out demos.npz
# Train base policy
lerobot-train-bcrac --data demos.npz --output policy/ --steps 6000 --device cpu
# Run full flywheel (stress + cluster targeting + recovery + eval + report)
lerobot-behavior-flywheel \
--base-policy policy/ \
--base-data demos.npz \
--workdir flywheel_run/ \
--episodes 20 \
--state-randomization-level medium \
--mid-episode-disturbance-prob 0.15 \
--cluster-selection hardest \
--device cpu
# Launch dashboard
lerobot-dashboard --workdir flywheel_run/ --port 8501
# Run ablation study
lerobot-ablation-study --data demos.npz --steps 4000 --eval-episodes 10 --device cpuFull technical writeup: A Data Flywheel Kernel for Robot Learning
Built on top of LeRobot by Hugging Face.