Learning from Demonstrations via Capability-Aware Goal Sampling

AI summary

Overview

• Research area: Imitation learning (learning from demonstrations) combined with goal-conditioned reinforcement learning and model-based RL.
• Technical level: Advanced. The paper assumes familiarity with MDPs, occupancy measures, behavior cloning, Go-Explore, Dreamer-style world models, and actor-critic temporal-distance objectives.
• Scope: The paper proposes Cago, a framework that treats expert demonstration trajectories as structured roadmaps and samples intermediate goal states at the boundary of the agent's current reaching ability to create an implicit curriculum.

What This Paper Is About

Direct imitation of expert demonstrations breaks down in long-horizon, sparse-reward tasks, because the agent cannot perfectly replicate the expert and small errors compound. The authors argue that existing methods match behavior "flatly" over the whole trajectory distribution without reasoning about which parts of the task the agent has already mastered. Cago instead monitors how far along each demonstration the agent can currently reach, then samples goals slightly beyond that point to drive steady progress toward the full task.

Key Contributions

1. Cago, a capability-aware learning-from-demonstrations framework. Demonstrations are used neither for direct imitation nor reward shaping nor offline pretraining, but as structured roadmaps for goal selection.
2. A visitation-tracking mechanism (Dict_visit) that records how often the agent visits observations along each demonstration, using a similarity metric (L2 distance for state-based environments, MSE between images for visual environments) and a matching threshold. A visitation frequency threshold λ_visit identifies the furthest demonstration index the agent is competent at reaching.
3. A capability-aware goal sampling distribution G_cap(π^G, τ^(i)) that draws subgoals from a window centered on that index, of width controlled by δ (e.g., 10% of the demonstration length), producing a curriculum of progressively harder goals.
4. A complete training pipeline combining Go-Explore-style two-phase rollouts (a goal-conditioned policy followed by a behavior-cloning Explorer policy), a Dreamer-backbone world model trained on collected data, and a goal predictor P_φ that infers the final goal from the current observation so the policy can be deployed on unseen scenarios at test time.

Main Findings

• Baseline comparison: Across MetaWorld "very hard" tasks (Shelf Place, Disassemble, Stick Pull, Stick Push, Pick Place Wall), Cago consistently outperforms all baselines in both final performance and learning efficiency. The paper states it beats Dreamer, JSRL, MoDem, and Cal-QL, plus GAIL, PWIL, SQIL, ValueDice, and RLPD evaluated in the appendix.
• Adroit suite: MoDem shows rapid early learning due to its behavior cloning pretraining and oversampling strategy, but Cago surpasses it in final performance after 1e6 environment interaction steps.
• ManiSkill suite: Given the limited demonstrations (20 per task), Cago is described as the only method capable of achieving high success rates.
• Shared-architecture control: Dreamer, which shares the same world model and policy architecture as Cago, performs significantly worse, which the authors attribute to the capability-aware goal sampling strategy.
• Goal distribution shifts over training: In the StickPush visualization, each red dot is the normalized position of a sampled goal within a demonstration (0 = start, 1.0 = final state). Early in training goals cluster at lower normalized positions; as training advances they shift toward higher positions.
• Ablations: Removing capability-aware goal sampling significantly degrades performance, because the agent often enters the Explore phase from states far outside the demonstration region. The BC-Explorer is also described as crucial for generating high-quality exploratory rollouts. Ablations (Cago-FinalGoal, Cago-StepBased, Cago-NoExplorer, Cago-RandomExplorer) are run on Disassemble, StickPush, and Pen over 5 seeds.
• Visual inputs: Cago-Visual operates on RGB images of size (64, 64, 3) with both policy and goal predictor working on image representations. It retains performance similar to the original Cago and consistently outperforms MoDem-Visual across Door, Hammer, and Pen over 5 random seeds.
• Theory: Theorem 1 bounds the expected total variation between true and learned dynamics along rollouts of the learned policy by μ + 2κ + 2ν, where κ bounds the behavior-cloning closeness, μ the model learning error under the BC policy, and ν the trajectory distribution closeness.
• Not reported: The provided text does not state the numeric final success rates, absolute performance values, or wall-clock/step budgets for the main comparisons; results are presented only as learning curves in the figures.

Methodology in Plain English

Cago keeps a tally of how often the agent's observations resemble each state in each demonstration. Because the tally is updated online, it reveals the furthest point along a demonstration the agent can reliably reach. The method then sets an upper limit at the last demonstration index whose visit count exceeds a frequency threshold (the paper gives 100 as an example). Around that index it defines a sampling window whose half-width is a fraction of the demonstration length (the paper gives 10% as an example), and it randomly picks a goal from inside that window. Goals can therefore be familiar ones or ones slightly beyond current ability, but not so far ahead that training collapses.

Each episode begins at the initial state of a randomly chosen demonstration. The agent first tries to reach the sampled goal with its goal-conditioned policy, then hands control to a behavior-cloning Explorer policy trained on the demonstrations, which continues forward and collects task-relevant data. This two-phase scheme anchors collected trajectories near the demonstration distribution while still allowing exploration. Collected data trains a predictive world model, and the goal-conditioned policy is then improved on imagined rollouts generated inside that model, using an actor-critic setup with a self-supervised temporal distance function that estimates how many steps separate a state from a goal. The reward is defined as the negative of that estimated temporal distance.

A practical detail: Cago resets the environment only to demonstration initial states, not to intermediate states. The authors argue this is more realistic than methods that reset to arbitrary demonstration states, because physical quantities such as joint velocities and angular momentum are hard to reproduce. To handle unseen test scenarios where no final goal is given, a goal predictor trained by minimizing mean squared error between predicted goals and true final observations supplies the goal condition at deployment.

Why This Matters

• Impact on research: The paper reframes demonstrations from something to imitate into something to navigate. It offers an alternative to reset-based curricula, distribution matching, and offline pretraining, and connects capability estimation directly to goal-conditioned model-based RL. Theorem 1 gives a formal account of how a behavior-cloned data-collection policy influences model prediction error.
• Real-world applications:
  • Robot manipulation from a handful of teleoperated demonstrations, where sparse rewards make exploration impractical.
  • Dexterous hand manipulation such as door opening, hammering, and pen manipulation, where fine-grained finger control matters.
  • Industrial assembly-style tasks such as peg insertion, cube stacking, and tool pulling, which the paper evaluates in ManiSkill.
  • Vision-based robot control, since Cago-Visual works from raw pixel observations rather than privileged state vectors.
• Industry relevance: Cago's data efficiency (10 demonstrations per MetaWorld and Adroit task, 20 per ManiSkill task) and its avoidance of intermediate-state resets lower the barrier to applying demonstration-driven training on physical hardware, where precise resetting to arbitrary expert states is usually infeasible.

Future Directions

• Demonstration quality and quantity: The paper references appendix experiments on how the number and quality of demonstrations, including suboptimal ones, affect performance, leaving the practical limits of degradation an open question.
• Hyperparameter sensitivity: The behavior of the visitation threshold λ_visit and the window width δ is examined in the appendix; how to set them automatically per task is not addressed in the main text.
• Theoretical tightness: Theorem 1 depends on assumptions about behavior cloning closeness, model learning error under the BC policy, and trajectory distribution closeness, which the paper does not empirically measure.
• Beyond the evaluated suites: Whether capability-aware goal sampling transfers to tasks where demonstrations do not form a single well-defined forward progress axis, or where goals are not directly comparable to states (S = G is assumed here), remains untested in the provided content.

Target Audience

Researchers and practitioners working on imitation learning, goal-conditioned RL, and model-based RL, especially those applying demonstrations to sparse-reward long-horizon robotics tasks. Readers interested in curriculum generation without environment resets will find the method most relevant. Readers seeking fully worked benchmark numbers or an introductory treatment of imitation learning will find the paper less suited, since the provided text reports results primarily through learning curves.

Authors’ abstract

Despite its promise, imitation learning often fails in long-horizon environments where perfect replication of demonstrations is unrealistic and small errors can accumulate catastrophically. We introduce Cago (Capability-Aware Goal Sampling), a novel learning-from-demonstrations method that mitigates the brittle dependence on expert trajectories for direct imitation. Unlike prior methods that rely on demonstrations only for policy initialization or reward shaping, Cago dynamically tracks the agent's competence along expert trajectories and uses this signal to select intermediate steps--goals that are just beyond the agent's current reach--to guide learning. This results in an adaptive curriculum that enables steady progress toward solving the full task. Empirical results demonstrate that Cago significantly improves sample efficiency and final performance across a range of sparse-reward, goal-conditioned tasks, consistently outperforming existing learning from-demonstrations baselines.