Learning from Demonstrations via Capability-Aware Goal Sampling

TL;DR

Cago introduces a capability-aware goal sampling framework for learning from demonstrations in sparse-reward, long-horizon tasks. By continuously tracking the agent’s capability along demonstration trajectories and sampling intermediate goals just beyond current reach, it creates an adaptive curriculum that guides exploration via a Go-Explore style framework and a Dreamer-inspired world model. A goal predictor enables test-time planning without access to final demonstration states. Empirical results across MetaWorld, Adroit, and ManiSkill show improved sample efficiency and final performance over diverse imitation and curriculum baselines, with ablations confirming the importance of capability-aware sampling and the BC Explorer component.

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.

Figures (15)

• Figure 1: Illustration of the Cago. Left: Directly setting the final goal as the agent’s target often leads to failure, as the current policy $\pi^G$ may not yet be capable of reaching it. The shaded region illustrates the set of states currently reachable under $\pi^G$. Attempting to reach $g_{\text{final}}$ (i.e., executing $\pi^G(\cdot|\cdot, g_{\text{final}})$) causes the agent to diverge from the demonstration trajectory. Right: Cago improves learning by leveraging a visitation frequency dictionary $\text{Dict}_{\text{visit}}$ built from demonstrations. Given a demonstration trajectory with subgoals $g_1, g_2, \dots, g_n$, the agent selects the furthest subgoal $g_i$ that remains within its current capabilities for Go-Explore sampling, enabling a curriculum of progressively more challenging goals aligned with the demonstration.
• Figure 2: The workflow of the goal predictor $\mathcal{P}_\phi$.
• Figure 3: Experiment results comparing Cago with the baselines over 8 random seeds. The solid line denotes the average success rate in evaluation, while the shaded region signifies the standard deviation.
• Figure 4: Figure(a),(b),(c) are the results of ablation study on the importance of each component of Cago over 5 seeds. Figure(d) shows the progress of capability-aware goal sampling in Stickpush.
• Figure 5: Visual input experiment results over 5 random seeds.

Theorems & Definitions (2)

• Theorem 1
• proof