Title: Learning from Demonstrations via Capability-Aware Goal Sampling
Venue: NeurIPS 2025
Accepted: yes
Keywords: Reinforcement Learning

Demonstration

Goal-Conditioned Reinforcement Learning

Model-Based Reinforcement Learning

Exploration Strategies

Imitation Learning

World Models
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents Capability-Aware Goal Sampling (Cago), a new method for imitation learning in tasks that require reaching goals over many steps. Cago works by first picking goals that are just beyond what the agent can currently do, helping it learn in manageable steps. In the next phase, an exploration policy trained from demonstrations helps the agent explore the environment. The approach also uses model-based reinforcement learning, where a world model and the goal-conditioned policy are trained together. Experiments show that Cago leads to much better sample efficiency and final performance than current leading methods.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n+ The paper clearly defines the problem: existing imitation learning methods struggle with long-horizon tasks due to compounding errors and a lack of awareness of the agent’s own capabilities. The descriptions of the proposed method and the experimental setup are also clear and well-organized.\n+ The experiments demonstrate significant improvements over baseline methods.\n\n**Weaknesses:**\n+ The approach involves several components—goal sampling, goal-conditioned policy, behavioral cloning (BC) policy, world model, and goal predictor. It would be helpful to understand how robust the overall method is to the performance of each individual component.\n+ The discussion on how the method could be applied to broader scenarios—such as real-world settings with diverse initial states and uncertain dynamics—could be expanded. This would help clarify the method’s generalizability and practical impact.\n\nQUESTIONS:\nIn addition to the questions raised in the weaknesses section, addressing the following points would help alleviate my concerns:\n+ Why not initialize the goal-conditioned policy with behavioral cloning (BC) pre-training, for example by sampling future states as goals?\n+ How does the method’s performance scale with the number of demonstrations? Improving the BC policy and the $L_\\tau$ component in Algorithm 2 could potentially make the online phase much more efficient. Given the high training cost of your method, it would be helpful to justify the choice of using only 10–20 demonstrations.\n\nLIMITATIONS:\nThe discussion on how the method could be applied to broader scenarios—such as real-world settings with diverse initial states and uncertain dynamics—could be expanded.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes Cago (Capability-Aware Goal Sampling), a novel method to enhance the efficiency of Imitation Learning (IL) in long-horizon, sparse-reward environments. Cago dynamically tracks the agent's current capabilities and, based on this, intelligently samples intermediate goals from expert demonstrations that are just beyond the agent's current reach, thereby adaptively adjusting the learning curriculum. This approach aims to mitigate the limitations of fixed curricula or direct imitation of demonstrations, such as error accumulation and inefficient exploration of the task space. The proposed framework integrates a Go-Explore paradigm, which leverages goal-conditioned policies and behavior cloning explorers, along with a predictive world model, enabling efficient data collection and policy optimization. The paper presents experimental results demonstrating Cago's superior performance compared to existing IL baselines across various robotic manipulation tasks.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n- The core idea of \"Capability-Aware Goal Sampling\" is itself interesting and represents a novel attempt to address the limitations of imitation learning. The adaptive curriculum approach, which dynamically sets goals based on the agent's capabilities, is theoretically appealing.\n- The construction of a framework that integrates multiple components, such as observation visit tracking, the Go-Explore paradigm, and a world model, demonstrates the complex design of the methodology.\n- The attempt to evaluate Cago's performance across various robotic manipulation environments showcases an effort to demonstrate the methodology's applicability.\n\n### Weaknesses\n\n- This methodology critically relies on \"high-quality demonstration trajectories that adequately cover the task space.\" While mentioned as a limitation in the paper, this core assumption is neither sufficiently justified nor mitigated, considering the significant difficulty in reliably obtaining high-quality demonstrations in the real world. The methodology severely lacks robustness to noisy, incomplete, or suboptimal expert demonstrations, with no experimental analysis or proposed solutions provided. This fundamentally limits the practical applicability and generality of the proposed method.\n- The `Dict_visit`based capability tracking and visit-count-based goal sampling appear tailored for environments where goals can be clearly defined as states directly from demonstrations. However, there is a clear lack of discussion or evidence regarding its scalability to highly-dimensional or abstract continuous goal spaces, or situations where goals are difficult to define directly from demonstration trajectories (e.g., abstract goals like \"achieve emotional state\"). Without in-depth analysis in this area, the general utility of the methodology remains low.\n- Key hyperparameters in Cago, such as `lambda_visit` (visitation frequency threshold) and `delta` (window size for goal sampling), are crucial in determining the behavior of the adaptive curriculum. However, there is an absence of an in-depth sensitivity analysis on how these parameters affect performance. The paper lacks a clear explanation of how misconfigurations of these parameters might disrupt the learning process (e.g., repeatedly selecting only overly easy goals, or consistently pursuing unattainable goals), leading to significant uncertainty in applying the method to new tasks.\n- Although Cago is compared to demonstration-based baselines, it overlooks comparisons with broader goal-conditioned reinforcement learning methods beyond demonstrations. This weakens the clarity of Cago's distinct contributions and its unique advantages over existing approaches.\n- Although quantitative results are presented, a deeper qualitative analysis of *how* Cago's learned policy behaves during the exploration phase, and how Cago succeeds in specific scenarios where baselines fail, is missing. This hinders a comprehensive understanding of the methodology's internal workings and the fundamental reasons for its success.\n\nQUESTIONS:\n1. Real-world demonstration data inevitably contains noise or can be incomplete. This paper assumes high-quality demonstrations. Have you conducted experimental analyses to see how Cago's performance changes when the demonstration data is intentionally noised (e.g., deleting some observations, using demonstrations with unclear target states), when the expertise of the demonstrator is reduced, or when the number of demonstrations is severely limited? Without such analysis, it is difficult to be convinced of Cago's utility in real-world environments.\n2. The `Dict_visit`based capability tracking and goal sampling mechanism might work well in clear state spaces like robot arm positions. However, in situations where 'goal states' are difficult to define clearly within demonstration trajectories, such as tasks involving linguistic instructions or needing to set goals within high-dimensional, abstract visual feature spaces, how can Cago sample goals and track capabilities?\n3. The values of key hyperparameters like `lambda_visit` and `delta` are critically important for Cago's adaptive curriculum. A detailed sensitivity analysis on how changes in these values affect learning efficiency and final performance should be included in the paper. Currently, there is a lack of clear explanation on how these parameters were determined, and no guidelines are provided on how to tune them for new tasks.\n\nLIMITATIONS:\nThe authors do not adequately address the limitations of their work within the main body of the paper. While the \"Checklist\" in the appendix indicates that \"Limitations\" are discussed, a dedicated section or clear discussion within the core text is absent. This makes it difficult for readers to quickly identify and understand the boundaries and potential shortcomings of the proposed method.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a novel method to guide learning when demonstrations are available. It suggest to track the progress along the demonstrations by keeping track of visitation counts. Goals are then sampled around the capability limits of the agent to support steady progress. This is implemented as part of a goal-conditioned model-based RL agent and tested across a variety of tasks.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper has a good structure and reads well. \n\nThe core idea of guiding learning by leveraging demonstrations is very interesting and addresses a highly relevant problem in imitation learning. While certain design decision may be debatable (e.g. the choice of a visitation dictionary and the use of a simple similarity metrics), the paper shows a successful application of the idea.\n\nThe proposed set of ablations is very interesting and confirms the algorithmic design choices.\n\nHaving said that, the framing of the method could be improved. While the method is likely applicable in different setups where demonstrations are available, the motivation and introduction mainly suggest application in imitation learning (using references to GAIL or PWIL). The paper then slightly deviates towards goal-conditioned methods, which still maintain a link to IL and is a sensible choice with suggested goal sampling. However, during evaluation many of the suggested baselines (JSRL, MoDem, CalQL, RLPD) are RL algorithms which leverage demonstrations in conjunction with a reward. As such they cannot directly serve as a imitation learning baseline and I wonder whether the environment reward is used? If there are used as RL backbone in conjunction with a Cago driven imitation reward they would rather represent ablations. On the other hand, there are not many online state-of-the-art imitation learning algorithm present among the baselines except for PWIL. I would suggest sticking with the IL narrative in which case further baselines would be required, e.g. SQIL [Reddy2020], CSIL [Watson2023], ValueDice [Kostrikov2020], or ReCOIL [Sikchi2024].\n\nWhile there is already a wide set of tasks in the paper I would suggest further evaluating on D4RL [Fu2020] in line with the above baselines.\n\nAccording to Algorithm 2 $D_{cap}$ is not used to train $\\pi^G$. Is this correct? So the distance function and policy are only trained on simulated data? It feels a bit wasteful to use $D_{cap}$ for the world model only. According to the rational behind eq (9) the left term could be attenuated with real world transitions, no? Overall, I would very much like to see the method applied to model-free RL.\n\nThe use of the similarity metric to count visitation could be a limitation for more complex tasks, in particular when moving images are involved. But also state spaces with highly varying ranges could probably be challenging. Could you maybe further comment?\n\nWhile not as bad as arbitrary resets, having to reset to specific initial states remains an extra burden, in particular for real world robots which usually contain non-controllable objects. It would thus be good to better understand the challenge with sampling parts of the initial states away from the demonstrations and whether this could maybe be addressed differently?\n\nI am not sure I fully understand the importance of Theorem 1. The model prediction error in eq (9) is over the full state space, but the theorem bounds it under the visitation of the learned policy and model. Also, it feels like assumption (2) would mean BC is good enough.\n\n[Reddy2020] Reddy et al., SQIL: Imitation Learning via Reinforcement Learning with Sparse Rewards, ICLR 2020.\n[Kostrikov2020] Kostrikov et al., Imitation Learning via Off-Policy Distribution Matching, ICLR 2020.\n[CSIL2023] Watson et al., Coherent Soft Imitation Learning, NeurIPS 2023.\n[Sikchi2024] Sikchi et al., Dual RL: Unification and New Methods for Reinforcement and Imitation Learning, ICLR 2024.\n[Fu2020] Fu et al., D4rl: Datasets for deep data-driven reinforcement learning.\n\nQUESTIONS:\n* Provide a clearer framing, I would suggest committing to the IL narative.\n* Extend the related work to better cover imitation learning (e.g. [Reddy2020], [Kostrikov2020] [CSIL2023] [Sikchi2024] from above)\n* Provide more IL state-of-the-art baseline.\n* Include results on D4RL for which there are results on established baselines\n* Provide further discussion on initial state resets, how this impacts real world robot experiments, and what alternatives there are, if any?\n* Further details on the implications of theorem 1 would be helpful.\n\nLIMITATIONS:\nLimitations around environment resets and similarity metric could be further discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to enable agents to learn from demonstrations using goal-conditioned RL. The proposed method, Cago (Capability-Aware Goal Sampling), maintains the visitation frequencies of observations across demonstrations to identify an agent’s currently reachable observations and generates states around those observations as the goal for policy optimization. Cago also learns an exploration policy and a final goal predictor from demonstration data, as well as a model for transition dynamics from data collected online. All these modules are used in the Go-Explore paradigm to learn the final goal-conditioned policies. \n\nCago is evaluated on multiple robot environments and compared with model-based RL (Dreamer and MoDem), Jump-Start RL, offline data for online learning (Cal-QL and RLPD), and imitation learning algorithms (GAIL and PWIL). The results demonstrate that Cago outperforms most of the selected baselines. The contributions of each component in Cago are also analyzed.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n1. The paper is written with a clear flow, and one can easily catch the high-level idea of the proposed method.\n2. Cago is a comprehensive framework that encompasses several key components, extending beyond the core idea of capability-aware goal sampling. \n3. The experimental results demonstrate the effectiveness of Cago.\n\n# Weaknesses\n1. The core idea of Cago is to use demonstrations to generate a curriculum and incorporate this curriculum into goal-conditioned RL. Prior work has proposed similar ideas, e.g., [1-4], but the experiments do not compare with them. Instead, the baselines evaluated, although falling into the category of learning from demonstrations, are not the most closely related to the proposed approach. \n2. I do not fully get the purpose of Section 2. It includes relevant work, problem formulation, and background knowledge, but all these parts are mixed together. I will recommend restructuring this section. \n3. BC Explorer appears to be a crucial component for Cago to work, but its training details are unclear. How is it trained? \n4. Cago includes multiple trainable modules, including a world model and a goal predictor. How well does each module need to perform to make Cago work? Or are they all end-to-end trained, so their individual performance does not matter? I would like to see more discussions on this. \n5. Following the previous point, although the paper mentions that Cago could generalize beyond the demonstration data, given a small number of demonstrations with complex model structures, it is very likely that the model would overfit to the demonstrations. I would like to see some experimental verification or discussions regarding the generalizability of Cago to unseen tasks. \n6. All experiments are run with five random seeds, which may not be sufficient to demonstrate statistical significance. \n\n# References\n[1] Yengera, Gaurav, et al. \"Curriculum design for teaching via demonstrations: theory and applications.\" Advances in Neural Information Processing Systems 34 (2021): 10496-10509.\n\n[2] Bajaj, Vaibhav, Guni Sharon, and Peter Stone. \"Task phasing: Automated curriculum learning from demonstrations.\" Proceedings of the International Conference on Automated Planning and Scheduling. Vol. 33. 2023.\n\n[3] Dai, Siyu, Andreas Hofmann, and Brian Williams. \"Automatic curricula via expert demonstrations.\" arXiv preprint arXiv:2106.09159 (2021).\n\n[4] Hermann, Lukas, et al. \"Adaptive curriculum generation from demonstrations for sim-to-real visuomotor control.\" 2020 IEEE International Conference on Robotics and Automation (ICRA). IEEE, 2020.\n\nQUESTIONS:\n1. In Figures 8 and 9, Cago does not consistently or significantly outperform other methods. What are the possible reasons? \n2. Cago mainly uses state similarity instead of trajectory similarity to identify an agent’s current capability. Can reaching a state with a trajectory that is very different from the demonstrations indicate that the agent has learned the ideal skills?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}