Research map / Reinforcement learning / Imitation learning
Learning from Demonstrations via Capability-Aware Goal Sampling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Cago, a Dreamer-based learning-from-demonstrations method that samples capability-aware intermediate demonstration goals (tracked via visitation counts) and explores from them with a BC explorer, outperforming IL and model-based baselines on sparse-reward MetaWorld, Adroit, and ManiSkill tasks.
Problem
Imitation learning and demonstration-based RL struggle in long-horizon sparse-reward tasks, where small replication errors accumulate and existing methods either do flat distribution matching or require unrealistic resets to arbitrary demonstration states.
Approach
Track, per demonstration step, how often the agent visits each demonstration observation; use the last well-visited index as a capability proxy and sample goals in a window around it, then run Go-Explore-style rollouts (goal-conditioned Go phase followed by a BC explorer), training a goal-conditioned Dreamer-based policy with a learned temporal distance reward and a goal predictor for test time.
Key results
- On MetaWorld very hard tasks Cago consistently outperforms all baselines in final performance and learning efficiency, averaged over 8 seeds.
- In Adroit, Cago surpasses MoDem in final performance after 1e6 environment steps, while MoDem learns faster early on.
- On unseen seeds (500), Cago reports 0.982 success on StickPush, 0.806 on Disassemble, and 0.824 on Adroit-Pen.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use capability-aware goal sampling instead of always selecting the final demonstration state (ablation) | Cago with BC Explorer · Disassemble, StickPush, and Pen | success rate | Helped |
| Use capability-aware goal sampling instead of sampling goals in proportion to training steps (ablation) | Cago with BC Explorer · Disassemble, StickPush, and Pen | success rate | Helped |
| Add a BC Explorer after the sampled-goal phase (ablation) | Cago with capability-aware goal sampling · Disassemble, StickPush, and Pen | success rate | Helped |
| Replace a uniformly random exploration policy with a BC Explorer (ablation) | Cago with capability-aware goal sampling · Disassemble, StickPush, and Pen | success rate | Helped |
| Use capability-aware goals instead of Dreamer's goal-selection strategy | Dreamer · StickPush | success rate: 0.68 -> 0.99 | Helped |
| Use capability-aware goals instead of PEG's goal-selection strategy | PEG · StickPush | success rate: 0.76 -> 0.99 | Helped |
| Use capability-aware goals instead of Dreamer's goal-selection strategy | Dreamer · Disassemble | success rate: 0.02 -> 0.80 | Helped |
| Use capability-aware goals instead of PEG's goal-selection strategy | PEG · Disassemble | success rate: 0.04 -> 0.80 | Helped |
| Use capability-aware goals instead of Dreamer's goal-selection strategy | Dreamer · Adroit-Pen | success rate: 0.52 -> 0.82 | Helped |
| Use capability-aware goals instead of PEG's goal-selection strategy | PEG · Adroit-Pen | success rate: 0.54 -> 0.82 | Helped |
| Use Cago instead of ACED | ACED · Fetch-PickPlace | success rate: 0.93 -> 1.00 | Helped |
| Use Cago instead of Task phasing | Task phasing · Fetch-PickPlace | success rate: 0.75 -> 1.00 | Helped |
Surprising findings
- Standard imitation learning baselines (GAIL, PWIL, etc.) perform poorly on these benchmarks, per the authors' experiments.
- Dreamer, sharing the same world model and policy architecture as Cago, performs significantly worse, indicating the goal sampling strategy drives the gains.
Limitations
- Requires resetting the environment to the demonstration's initial state (partially relaxed for some tasks), which may be infeasible in real-world settings.
- Relies on simple similarity metrics (L2/MSE) and a simple visitation counting scheme that may not scale to complex environments.
Experimental setup
Datasets: MetaWorld demonstrations (10 per task) Adroit demonstrations (10 per task) ManiSkill demonstrations (20 per task)
Benchmarks: MetaWorld very hard tasks (Shelf Place, Disassemble, Stick Pull, Stick Push, Pick Place Wall) Adroit (Door, Hammer, Pen) ManiSkill (PegInsertionSide, StackCube, PullCubeTool)
Metrics: success rate sample efficiency
Baselines: Dreamer JSRL MoDem Cal-QL GAIL PWIL SQIL ValueDice RLPD PEG ACED Task phasing
Builds on: Dreamer LEXA Go-Explore JSRL MoDem
Closest papers by meaning
- Demonstration-Conditioned Reinforcement Learning for Few-Shot Imitation (ICML 2021 · similarity 0.90)
- Expert Proximity as Surrogate Rewards for Single Demonstration Imitation Learning (ICML 2024 · similarity 0.90)
- GO-DICE: Goal-Conditioned Option-Aware Offline Imitation Learning via Stationary Distribution Correction Estimation (AAAI 2024 · similarity 0.90)
- Uncertainty-Based Smooth Policy Regularisation for Reinforcement Learning with Few Demonstrations (NeurIPS 2025 · similarity 0.89)
- Interactive and Hybrid Imitation Learning: Provably Beating Behavior Cloning (NeurIPS 2025 · similarity 0.88)
- Imitation Learning from Observation with Automatic Discount Scheduling (ICLR 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.