Learning from Demonstrations via Capability-Aware Goal Sampling
Authors: Yuanlin Duan
Yuning Wang
Wenjie Qiu
He Zhu
Venue: NeurIPS 2025 poster ()
Keywords: Reinforcement Learning
Demonstration
Goal-Conditioned Reinforcement Learning
Model-Based Reinforcement Learning
Exploration Strategies
Imitation Learning
World Models
Primary area: reinforcement_learning

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.

Learning from Demonstrations via Capability-Aaare
Goal Sampling
Yuanlin Duan Rutgers University yuanlin in. duan@rutgers.
Yuning Wang Rutgers University yw8950cs rutgers..drrrr....
Wenjie Qiu Rutgers University wq378cs rutgerrs...rr.... edu
He Zhu
Rutgers University hz3750cs rutgers. edu
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-Awar Goal Sam- pling), novel rearning-from-demeed from-tomonstratioon method that mitigates the brittle dependence on expert trajectories for direct imitation. Unlike prior methods that rely demonstrations only for policy initialization reward shaping, Cago dy- namically tracks the agent's competence along expert trajectories and uses this signal to select intermediate deps--goll that are just beyond the agent's 's current reach-to guide learning. This results 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 range sparse-reward, goal-conditioned tasks, consistently outperforming existing learning from-demonstrations basclines.
1 Introduction
Imitation Learning (IL) provides a powerful paradigm for training agents using expert demonstra- tions, effectively alleviating the exploration challenges common in Deep Reinforcement Learning (DRL) (Arulkumaran and Lillrank 2024). The simplest form II. Behavior Cloning (BC), which directly supervises policy actions based the states visited by the expert (Bain and Sammut [9955 Torabi et al.) 2018). However, BC often suffers from compounding errors when the learned policy deviates from expert trajectories. To overcome this limitation, modern approaches such as GAIL (Ho and Ermon 2016), PWIL (Dadashi et al.) 2020), and AdRIL (Eysenbach et al. 2021) seek to align the aatennation distributions of the agent and the expert through adversarial or cittribution-matchinn objectives. In parallel, Inverse Reinforcement Learning (IRL) methods (Ziebart et al., 20088 aim to infer underlying reward functions from demonstrations, which then guide reinforcement learning. More recently, advances in offline and offline-to-online RL, such as CQL (Kumar et al.) 2020) and Cal-QL (Nakamoto et al. 2023), integrate demonstrations anchors regularize policy learning. These methods penalize value estimates that diverge from demonstrated behavior, mitigating overestimation and instability caused by out-of-distribution actions.
However, existing IL methods often struggle with complex, hong-horizon tasks because they fail to reason about which parts of the task the agent has already mastered and which remain challenging. In particular, distribution-mate matching approaches perform flat matching srreching-attmmttnnn attempting to align occupancy measures over the entire trajectory distribution without considering the agent's evolving capabilities. This leads to poor exploration guidance, especially in the early stages of training when the agent
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

seldom reaches meaningful parts of the state space. As a result, the learned reward function tends to assign uniformly low rewards, yielding uninformative gradients and hindering effective policy improvement. Some prior work proposes demonstration-guiel to solve tasks by starting near the goal or high-reward states and gradually expanding to earlier parts of the trajectories (Resnick et al., 2018, 8, Salimans and Chen 2018 Tao et al., (024). However, these approaches rely on the ability to reset the agent to arbitrary demonstration states- an assumption impractical in real-world settings due to challenges in replicating physical conditions like joint velocities and angular momentum.
curriculum learning that trains agents
We propose Cago (Capability-Awwe AAware Goal Sampling), a new learning- foaming-frrm--ie demonstrations framework that explicitly aligns the agent's 's learning process with its evolving capabilities. Unlike prior methods that use demonstrations for direct imitation. reward shaping, or offline pretraining, Cago treats demonstrations as structured roadmaps. It continuously monitors which parts of a demonstration the agent can already reach and leverages this signal to sample intermediate goal states in the demonstration, those at the boudary of the agent's current goal-reaching capabilities. At each episode, goal-conditioned agent (Liu et al. 2022 Plappert et al. 2018) first attempls to reach the sampled goal and then explores forward from it, generating informative, task-relevant data for policy optimization. This iterative process of capability-awar goal selection and curriculum-aligned exploration enables steadily progress toward solving the full task.
We evaluate Cago across several sparse-reward environments and demonstrate substantial improve- ments in both sample efficiency and final task performance over existing imitation-based baselines. Our experiments highlight that capability-awan goal sampling provides powerful signal for struc- turing learning, particularly in long-horizon tasks.
2 Background and Problem Setup
Reinforcement learning (RL) aims to enable agents to learn optimal behaviors through trial-and-error interactions with environment. An RL problem is formulated as a Markov Decision Process (MDP), represented as a tuple (S, A, T, G, 77. R. Po). The agent operates within a state space S and takes actions from an action space A. transitioning between states according to the dynamics T(s' |s, a). R(s, a) E IR is the reward function and Po is the initial state distribution. Given a policy n, consider the trajectory т (8o; (n2 81, .., sampled by æ,, i.e., 8p Po. 0t rT(1ss), and 81+1 T(- at)- The goal of RL is to learn a return- couumm-mxiiizzng policy π arg max, E, (ass((s)ss)
1 7'r(st, at)]
where y E (0, 1] is the discount factor.
Learning from Demonstrations. In imitation learning, the agent is provided a dataset of demon- strations Dammo collected from some expert policy Texpent. The objective to learn a policy π that reproduces the expert's 's behavior by generalizing from these demonstrations. The simplest approach, behavioral cloning (BC), treats this as a supervised learning problem, minimizing the discrepancy between the agent's 's predicted actions and the expert's. Another line of work, inverse reinforcement learning (IRL) (Abbeel and Ng 2004), aims to infer an underlying reward function that explains the expert's 's behavior and then optimizes policy through RL on this learned reward, thus decoupling reward inference from policy optimization. Building on ideas from IRL, Generative Adversarial Imitation Learning (GAIL) (Ho and Ermon 2016) bypasses explicit reward recovery by training policy and a discriminator in an adversarial game: the discriminator distinguishes expert from agent trajectories, while its output serves an implicit, learned reward signal guiding the policy. More broadly, many recent imitation learning algorithms can be interpreted as minimizing a divergence between the expert and agent occupancy measures.
State Reset. Several methods attempt to mitigate the exploration challenge in sparse-reward RL environments by resetting the agent to states from expert demonstrations, thereby bypassing the need to discover those states through the agent's own exploration. These strategies include initializing the
agent to states sampled uniformly from demonstration trajectories (Nair al. 2018 Peng et al. 2018 Hosu and Rebedea 2016), employing a hand-crafted curriculum (Zhu et al. 2018), or using a reverse curriculum that progressively trains the agent from goal or high-reward states backward (Resnick) jet al.] 2018 (Salimans and Chen 2018: Tao et al. 2024). These approaches assume the ability to reset the agent to arbitrary demonstration states an assumption that is unrealistic in real-world settings.
Goal-Conditioned RL (GCRL) extends the standard RI. framework by conditioning policies on specific target goals, guiding agents toward desired goals. The MDPs are augmented with
2

---

goal space G and are associated with states via mapping 1 S - G, ensuring that each state corresponds an achieved goal. In GCRL, the reward signal from the environment typ- 1)((²) g). We assume that each episode has a fixed horizon T and S = G. The agent's objective is to train a goal-conditioned policy r (-1s;.g) to achieve a given goal g E G through maximizing the expected cumulative reward
ically sparse and is defined as: R(s, n,
J (T) = L Sggpy, st.g)
== 7² R(ss,, Ct St+1, g) where Py is the goal distribution.
This Paper. We introduce Cago, novel approach that leverages demonstrations as scaffold for goal-directed reinforcement learning. Rather than direct imitation, Cago uses demonstrations to guide exploration by training a goal-conditioned policy (a | s; that Iearns to progressively reach intermediate states along demonstration trajectories, effectively inducing a curriculum that facilitates steady progress toward solving the full task. In addition, Cago learns goal predictor P(s) that infers the final goal state 9r from the current state s. The resulting task policy is defined as (s) x (s. P(s)) that enables automatic inference of goal conditions at test time for previously unseen situations.
Method
The main idea of Cago is to continuously monitor the agent's evolving capabilities to reach various stages of demonstration trajectories during training. It dynamically selects the most appropriate goal from the demonstrations, conditioned on the agent's current performance ceiling. The selected goal guides online exploration, with the agent first attempting to reach it using its current policy. From there, it continues to explore, collecting task-relevant trajectories in an Go-Explore style (Ecoffet et al., 2019). By anchoring exploration in achievable yet progressively harder goals, this process effectively constructs an implicit curriculum, where the agent is gradually exposed to more challenging states aligned with its growing competence.
3.1 Observation Visit Tracking with Demonstration Alignment
Cago assumes the existence of a limited amount of expert demonstrations Docmo : {-(0) {(so, au)( ...., (SL, aL,) ) where M is the number of the demonstrations and L; is the length of i-th demonstration 7(0) To select goals the boundary of the agent's current reaching capabilities, It is crucial to determine the stage at which the agent can accomplish its task completion. Central to our method maintaining a dictionary Dictvisa that tracks the visitation frequencies of observations 8; across demonstrations. For each demonstration r(t), we initialize an all-zero list of the same length as its steps. Each element in this list records the visitation count of the corresponding observation in the demonstration. At each environment step, the record list is updated to reflect whether the agent has visited observations from the demonstration, based on similarity metrics sim(..) such as L2 distances for state-based environments or mean squared errors (MSE) between images in visual environments. Formally, we define the visitation record dictionary as:
Dict,isi = - [0,0,.... E No. i = 1,2, 1,2,..., M
During online exploration, for each new episode, we first sample a demonstration ((() from D Dommo and reset the environment the initial state of (() Our strategy is more practical in the real-worid setting than related methods (e.g. (Tao et al., 2024 Nair et al. 2018)) that reset the environment to intermediate demonstration states, which are often infeasible to reproduce due to unobservable or difficult-to-control physical factors such as velocity and angular momentum. Given a rollout T (So, .i,...) from the environment, we update Dictrisu(>>)) as follows:
Dict.isl (11] += if sim (19) L. L, E L,
where is the agent's 's observation state at timestep i, T.r is the total length of the rollout 83 is the j-th observation state in the i-th demonstration ((i), sim(..) is the similarity metric (e.g.. L.2 distance for state-based environments or MSE for masee-based environments), and is a matching threshold. This simple record dictionary effectively tracks the agent's progress and helps identify its goal-reaching capability limits along task demonstrations.
3

---

Toe hord capey
Chalegrngt smirtlegee
Dictvisit 9
grinal
g find
Expest Toastayy
Go
Explor
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 d may not yet be capable of reaching it. The shaded region illustrates the set of states currently reachable under TC Attempting reach Pinal (i.e., executing x (-|-, (ima)) causes the agent to diverge from the demonstration trajectory. Right: Cago improves learning by leveraging visitation frequency dictionary Dict,isir built from demonstrations. Given a demonstration trajectory with subgoals !i 99..9........ In. the agent selects the furthest subgoal yi that remains within its current capabilities for Go-Explore sampling. enabling curriculum of progressively more challenging goals aligned with the demonstration.
3.2 Capability-Aware Goal Sampling
Cago leverages demonstration- bbase visitation counts to guide goal selection and trajectory collection in capability-aware manner. Figure illlstraates our method. After resetting the environment the initial state of randomly sampled demonstration ,(t), select a goal for the agent explore:
g g~ Geap (Tt (9)),
where ( ((0) denotes capability awware goal sampling distribution over subgoals whose reachability is aligned with the current goal-reaching capability of the policy πº. Cago examines the visitation frequency list to identify the last index where the frequency exceeds predefined threshold:
max { Dictoisu[ (7]]]
Avisit
(4)
where pict,ss[T<(1]] denotes the vis- itation frequency of j-th observation
Algorithm 1 Capabilityy Aware Goal Sampling (Cago)
sy of 7(0) under policy α² and Avisit
1: Input: Demonstration 7(1), Visitation recurd Dictvaa
is a frequency threshold (c.g. 100).
2: Output: capability-aware goal g
3: Identify capability-aware upper limit point using visitation thresh- old Aua (Eq. (
This index indicates the latest point in the demonstration that the agent sufficiently competent at reach- ing- efgeeffectvvely serving as proxy for the limit the agent's s current goal- reaching capability. Cago constructs goal sampling range centered around j". The sampled goal drawn from this range, which allows the agent to either revisit familiar goals or attempt slightly goal sampling
4: Define sampling region Geap(a ,(()) centered around (Eq. 51
5: Sample subgoal ~ Geng (n ,700)
6: return
more challenging ones that are just beyond its current capability. This capability- strategy introduces controlled diversity into the training process and encourages progressive learning, while also avoiding excessively difficult goals that could derail training. The corresponding goal sampling region is defined as:
G (00) - ( (00)
E (+)) 1
(5)
where Li is the length of (6) and (0. 1] controls the window size for goal sampling (e.g., 10% of the demonstration length). Goals are then sampled at random from this set. Our capability- goal sampling scheme introduces a curriculum-aligne learning signal that progressively guides the agent with steady improvement toward successful task completion. The overall goal sampling process in Cago is described in Algorithm
3.3 Learning Framework
Go-Explore. Cago trains goal-conditioned agent following the Go-Explore paradigm (Ecoffet et al. 2019), which divides each episode into two sequential phases: the Go Go-phase and the Explore-phase.

---

In the Go phase, the agent is guided towards a sampled goal state using the coul-conditioned policy (-I-, g). reaching an intermediate state sg. To improve environment exploration beyond the agent's current capabilities, the Explore-phase phase takes over from Sg, where an exploration policy x is used for the remaining time steps. Since we have access to a limited set of task demonstrations Daemo- we implement π: as a Behavior Clone (BC) Explorer trained on Demo- The BC Explorer outputs a stochastic action distribution that enables the agent to balance between exploration and imitation. This two-phase strategy ensures that collected trajectories stay anchored near the demonstration distribution in Ddemo- while encouraging exploration. We further analyze the impact the BC Explorer through ablation studies, detailed in Section 4
As Cago actively resets environments initial states drawn from the demonstration set Deemo key question is how the agent generalizes beyond Diemo Our solution is to train the conliconditioned agent T.
using a richer set of imagined rollouts generated by a world model via model-based RI..
World Model and Policy Training, Cago stores the trajectories generated under the Go-Explore paradigm with capability-aware goal sampling in a dataset Deap {(8,, ae, st+1)(=1) for world model and policy training. A pre-
Algorithm 2 The main training framework Cago
Input: GC Policy a² World Model M. Demonstrations Daems
2: Initialize replay buffer Deap and Dictyier (Eq. (I]
3: Train explorer policy = using Behavior Cloning Deeno
4: Train goal predictor P. Diemo (Fg.
5: for
to Ntrwin do Initialize empty trajectory T
dictive world model M approximates the transition dynamics ((''ls, a) in the real world M as ((s'|s, a). Our model learning algorithm is based on the Dreamer backbone (Hafner et al. 2019, 2020 (023). which updates
6:
7:
Randomly sample demonstration r()) E Demoo
8:
Initialize the environment to the initial state of (1)
9:
Sample a capability-awart goal by Algorithm for z to I.- do
10:
agent has not reached g < Tgo then x (s.9)
12:
the world model M via supervised
13:
14:
(s)
learning using Deap- Once the world
15:
Step in the real environment using and add this step to Update Dict,sa(F()) (Eq. (2)
model is updated, we train the goal- conditioned policy T using imagined trajectories generated by the world model M. Intuitively, since Deap collected by exploring around demon- stration states in Paemo- the learned model enables the agent to generate imagined trajectories that remain grounded in task-relevant regions of the state space. Each imagined trajectory from the learned world model M begins at sg, a state randomly sampled from trajectory T in Ddemo Decpp and is rolled out for H steps using the goal-conditioned policy - (aglse: g). The goal state g is selected as future state from the same trajectory т. The objective is to train x to reinforce trajectories that efficiently reach g in the imagined rollouts from Su under the learned dynamies M. To achieve this, we adopt an actor-critic critic algorithm that leverages a sulf-super temporal distance function (s g) (Mendonca et al. 2021), which estimates the number of steps required transition from state 8 to goal g. The reward function is defined as: ( = - Dr(8, s,, g). This formulation encourages the policy generate actions that minimize the estimated temporal distance to the goal. The temporal distance estimator D; is trained by extracting state pairs (st, s(+k) from simulated trajectories generated by the world model. The function learns predict the normalized temporal difference between two states: D, ( (s:), (((++)))
16:
17:
Dey Dap (r)
18: Update M with Desp
19: Update n using imagined rollouts with M
à to states (e.g.. embedding them into the world model's N where V denotes transformation applied latent space), and H is the length of the generated rollout. More details on the model-based learning algorithm and the full training procedure for Dr can be found in Appendix B.1 and Appendix B.2
Goal Predictor. At training time, the goal-conditioned policy G 1, .g) is trained using intermediate states from demon- stration trajectories
P
goal conditions (recall S 9). This assumes access to demonstrations, with the final states used as the target goal condition. However, test time, this assumption no longer holds: for unseen scenarios, the true final goal state not available. This raises the challenge of how to specify an appropriate goal condition based solely on the agent's ' current observation. We introduce goal predictor Pa, learned model
So
Figure 2: The workflow of the goal predictor Pa.
5

---

that infers a goal state و given the current observation s:
Ps 8 9, where 9 = (p(s)
(6)
The mapping learned by P is illustrated in Figure 2 It is trained using demonstration trajecto-
ries Daemo by minimizing the mean squared enror between the predicted goal and the true final observation:
min E (r(1)=0))) ( ( Po(() 81 10)2
Once trained, the goal predictor enables T to generalize to new tasks. Given a test-time state glext the predicted goal g Po serves as the planning target for the agent nt glest glest. The complete training pipeline of Cago is detailed in Algorithm 2
Rationale Behind Cago's Design. Let J [TT, .M) and J(T", M) be the expected return of the agent's
policy π and expert policy x¹ in the real-worlh MDP M.. We want to bound their return difference:
min [[(" MM - M)],
(8)
Let Raax be the maximum of the reward with unknown dynamics: Rasx (s.(s.a) R(s, a) and
(), = (1 7) Ei=o 7 p((s = = a) be the discounted state-action visitation distribution of policy ı in the real world MDP M. Suppose that the total variation of learned dynamics model
M from the true transitions .M is bounded such DTv (T a). ((s. S a V(s. a) E S A.
According to previous work (Rafailov al., 2021 DeMoss al. 2023 Kolev et al. 2024). have:
Rmax (1-7))
Rmax Drv (prn
(TTS,,) - J(T, a
+
y
model prediction error
udaptation error
where PM is the discounted visitation distribution of the expert policy and Dry denote total variation distance. The model prediction error with respect to the true environment dynamics can be reduced by collecting more real-world data. In contrast, the adaptation error depends on the total variation distance between the distribution of trajectories generated by policy и under the learned world model M and the expert distribution under the true dynamics M. Thus, the learning problem reduces to bounding the deviation between the behavior of the learned policy π under the learned model and the expert behavior under the true environment. To this end, given any H-step trajectory (#o, 81,..... 8 sampled from expert demonstrations D Cago encourages the agent to match expert behavior by rewarding it for reaching the final state = sи starting from so under the learned dynamics model M (Line 19 in Algorithm 2).
We further show that Cago effectively reduces the model prediction error by leveraging the BC explore policy _ BC for data collection. In the following, state-action distribution time induced by policy π in the environment M. We assume df dPane ss d: d,,."", where Ddcmo is the dataset of demonstrations generated by the expert x;, and is sufficiently representative to approximate the true marginal distributions at each timestep. Assuming: (1) x BC accurately approximates the expert policy in Ddemo, (2) the world model M accurately trained on state transitions induced by N BC and (3) the learned policy π generates trajectories in M that closely match the expert's behavior, we can bound the model prediction error along the imagined rollouts generated by x under M:
d:M. to denote the marginal
Theorem 1. Let M denote the true dynamics model and M the learned model. Let TBC be a behavior-cloned policy, and w a new policy. Let Dáemo be a dataset of expert demon- strations from unknown expert policy. Suppose that, for all t = O.I,...,TT (1) Close- ness of behavior cloning:
Div
M,.Bc dPamoo S K, (2) Model learning error under BC: S p. and (3) Trajectory distribution closeness:
E (aaa)aaMipe
DIV (,)
Drv P
< v. Then all t = 0.1...., T, we have:
(DIv (MG- ss,)) ((.11sa)) Sp + 2x + 2v.
E (s.a)mafiix

---

4 Experiments
We evaluate Cago across a diverse set of challenging robotic manipulation environments address the following research questions: (Q1) Does Cago outperform existing imitation learning baselines that leverage demonstrations in alternative ways? (Q2) Can Cago effectively realize capability- goal sampling that aligns with the agent's ss learning progress? (Q3) How essential are the proposed capa goal sampling and BC-Explorer components to the overall performance of Cago?
Environments. For our experiments, we evaluate and compare Cago against several baselines across three robot environment suites with sparse rewards: World (Yu et al. 2020), Adroit (Rajeswaran et al. 2017), and Maniskill (Gu et al. [023] Tao et al., 2025). We adopt the five "very hard" level environments from MetaWorld, as categorized by Seo et al. (2023): Shelf Place. Disassemble, Stick Pull, Stick Push, Pick Place Wall. These environments are considered the most challenging tasks in Metaworld, requiring precise robotic arm control with only sparse task completion rewards. We also use three dexterous hand manipulation tasks from the Adroit suite: Door, Hammer, Pen. To succeed in these three environments, the agent must perform fine-graine mne-grained and intricate finger manipulations, enabling the grasping and movement of different objects. We also selected three challenging tasks from the ManiSkill benchmark: PegInsertionSide, StackCube, and PullCubeTool. The sparse-reward ManiSkill environments are the most difficult tasks in our benchmarks due to their high-dimensional state and action spaces. During training, we used only 10 demonstration trajectories per task for the Meta World and Adroit environments, and 20 demonstration trajectories per task for the ManiSkill environments. More details about each task can be found in Appendix E
Baselines. Our approach is developed on top of the Dreamer framework (Hafner et al. 2019 2020
Hu et al. 2023 Duan et al., 2024a|b), bb. making it a key model-based RL baseline for evaluating the performance gains of Cago. Jump-Start Reinforcement Learning (JSRL) (Uchendu et al.) (203) is a curriculum-base early-stage exploration during online training. At the beginning of each training episode, the agent follows the guide-policy for number of steps determined by curriculum progression, after which control is handed over to the online policy. In our JSRL implementation. due to the limited number of demonstrations available, training a reliable guide-policy becomes challenging. Therefore, we directly use the demonstration trajectories as the guide-policy. p Specifically, reset the environment to a demonstration initial state, enabling the agent to replicate expert behavior during the initial phase of each episode before switching to the online policy. Our JSRL implementation is also built on top of the Dreamer framework. MoDem (Hansen et al. 2022) represents one of the most efficient frameworks in the model-based RL literature. It pretrains its policy using small set of demonstrations and repeatedly oversamples the demonstrations to train both the world model and the policy. We consider MoDem be the strongest baseline due to its fast convergence and low data requirements. Cal-QL. [Nakamoto et al. 2023) is a state-of-the-art algorithm following the offline- to-online RL paradigm. It uses demonstrations to pretrain the Q-function and applies calibration to mitigate performance drop when transitioning from offline to online learning phases. In addition to these representative baselines, we compare against four more imitation learning basclines: GAIL (Ho) and Ermon 2016). PWIL (Dadashi et al. 2020). SQIL (Reddy et al.) 2019), ValueDice (Kostrikov al. 2019). and RLPD (Ball al. 2023). well-tuned variant of SAC that leverages offline data when learning online, in the Appendix
approach that leverages a guide-policy pretrained from offline data to guide
During training, we uniformly sample a demonstration and reset the environment using the same seed that was used to collect it. All baselines share the same seeds and demonstration data. To evaluate generalization,
Main Results.
test on unseen environment seeds. For these, Cago uses the goal predictor (Section 3.3) to infer goal conditions for the goal-conditioned policy r Fach method is evaluated on 100 held-out seeds, and we report the average success rate over these episodes. Figure depicts the mean learning performance of Cago and all baselines in terms of the agent's 's task success rate averaged over 8 random seeds. On the Meta World very hard tasks, Cago consistently outperforms all the baselines in both final performance and learning efficiency. In the Adroit suite, although Modem exhibits rapid early learning due to its behavior cloning (BC) pretraining and oversampling strategy, Cago surpasses it in final performance after 1eб environment interaction steps. Notably, Dreamer, which shares the same world model and policy architecture as Cago, performs significantly worse, underscoring the effectiveness of the capability-awar goal sampling strategy. Our JSRL baseline, based on the same world model architecture, adopts a uniform curriculum to reduce the guide-steps from demonstrations. It lags behind Cago in both learning speed and final success
7

---

Ohee
Orys
Adl
SU2Ce50 otc
Dal-
SUccess ntte
ette
Dr-CL
96
02
02
02
30
0.
08
10
0.2
6.4
0.6
0.
23
0.4
09
13
0.2
04
00
steps
steps
staps
staps
(a) Disassemble
(b) PickPlace Wall
(c) ShelfPlace
(d) StickPull
1.0
1.0
06
Rodee
05
Foderr
06
Kodee
0.5
Fodkun
SUcCess 185
siccess
stagss
slags
stags
(c) StickPush
(1) Door
(g) Hammer
(h) Pen
Cape
5100050 obte
Or-CL
O,
Success mite
08
Pool
(i) PegInsertionSide
(k) StackCube
Figure 3: Experiment results comparing Cago with the baselines 8 random seeds. The solid line denotes the average success rate in evaluation, while the shaded region signifies the standard deviation.
rate, highlighting the effectiveness of our goal-sampling strategy in adapting to the agent's 's evolving capabilities. In the ManiSkill environments, given the limited demonstrations, Cago stands out as the only method capable of achieving high success rates.
-apability-Aware Goal Distribution. To answer Q2, we visualize the progression of capability- aware goal sampling throughout the training process in the StickPush environment in Figure 4d) Each red dot represents the normalized position of sampled goal within demonstration trajectory, with 0 indicating the start and 1.0 indicating the final demonstration state. Early in training, the agent predominantly samples goals lower normalized positions, focusing on subgoals near the beginning of the trajectory that are within its current capabilities. As training advances, goal sampling gradually shifts toward higher normalized positions, indicating the agent's increasing ability to pursue more challenging goals closer to task completion. By continuously targeting goals just at the boundary of the agent's 's current capability, Cago facilitates efficient learning in reward, long- horizon tasks.
Ablation Studies. To answer Q3, we assess the individual contributions of (a) capability goal sampling and (b) the BC-Explorer component the overall performance. The first ablation, Cago- FinalGoal, retains only (b): it uses BC-based exploration but always selects the final observation from a demonstration the goal phase of our Go-Ex Go-Explore sampling paradigm. ignoring the agent's 's current goal-reaching capability. The BC Explorer takes control from the goal- condiiiond policy halfway through each rollout. The second ablation, Steo-Ste,Based, also retains (b), but samples goals from demonstrations in proportion to training steps (current training step / predefined total training steps). However, it does not assess the agent's s actual capabilities, and may therefore sample goals that are either too easy or too difficult for the agent at its current learning stage. The third ablation,
8

---

Drin
Capp sad
Eage seroliiee
SUCCes0 otc
Eoue fie Aantiuun
SUCCeS Rite
02
30
0.
0.4
066
08
10
0.
6.4
0.6
-3
04 0.
08
Splonte Number
steps
steps
steps
(a) Disassemble
(b) StickPush
(c) Pen
(d) Sampling Progress
Figure 4: (b)) 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.
Explorr,,,,,,,r, keeps only (a): it uses capability-awar goal sampling, but does not explore beyond the sampled goal with the BC Explorer. The fourth ablation, Bago-RandomExplorer, replaces the BC Explorer with uniformly random policy during the exploration phase of our Go-Explore-style rollout strategy. We conduct the ablation study on the Disassemble and StickPush tasks from Meta World, and the Pen task from Adroit. As shown in Figure 4 removing capability aware goal sampling significantly degrades performance. Without it, the agent often enters the Explore phase from states far outside the demonstration region, making it difficult for the BC-Explorer to make meaningful progress. The BC-Expe BC-Explorer itself is also crucial, as it accelerates learning by generating high-quality exploratory rollouts. We further examine how the number and quality of demonstrations (including suboptimal ones), as well as the hyperparameter Aviest and 6, affect performance; see Appendix F for details.
Visual-input Environments. We further assess Cago's applicability in high-dimensional visual settings. Specifically, we extend our framework to raw pixel observations by replacing the vector- based states and goals with RGB image inputs of size (64, 64, 3), resulting in a variant referred to Cago-Vishe Visual. In this setting, both the policy and the goal predictor Pe operate representations. We benchmark Cago- Visual against Mooem--Visual that similarly utilizes image-bases based observations and demonstration. As shown in Figure 5 Cago- Visual not only retains performance similar to the original Cago, but also consistently outperforms Modem- Modem--Vsual, highlighting the robustness of our method in visual domains. Details of the goal predictor employed in our visual-input experiments are provided in Appendix F2
image strong model based baseline
Cage
Cage-Viuu Vouu
Cago visual
Cage Vieua
0.8
Midoo Visual
0.8
Nodom Visual
0.8
Modon Visual
SUcceSS refe
success rete
Success rete
0.6
0.2
0.2
0.2
0.000
0.60
0.00
steps
steps
steps
(a) Door
(b) Hammer
(c)
Pen
Figure 5: Visual input experiment results over 5 random seeds.
Related Work
Imitation Learning. Demonstrations are a key tool for improving the efficiency of RL, with prior work integrating them across various stages of RL pipeline (Arulkumaran and Lillrank 2024 Nair et al. 2018 Cui et al. 2025). A prominent approach uses demonstrations for direct learning Via Behavior Cloning (BC) and its variants (Bain and Sammut) 1995] Torabi al., 2018). The introduction of the Generative Adversarial Imitation Learning (GAIL) algorithm (Ho and Ermon

---

2016) has driven significant advances in scalable deep imitation learning methods (Fu et al. 2017: Ghasemipour et al., 2020 Kostrikov et al., 2018 Jena et al. 2020; Finn et al. 2 Blondé and Kalousis 2019| Orsini et al.) 2021 Eysenbach et al. 2021). Beyond adversarial approaches, sev- eral imitation learning algorithms aim to match the state action distributions of the expert and the agent through adverarrallial techniques, such as non-parametric models (Kim and Park (018), random network distillation (Wang et al. 2011)), support estimation (Brantley et al. 2020), Wasser- stein distance minimization (Dadashi et al. 2020), and moment matching (Swamy et al. 2021).
Demonstrations have also been the bridge between imitation learning and offline RL. Conservative Q-learning (CQL) (Kumar et al.) 2020) and Cal-QL (Nakamoto et al., 2023) regularize Q-values using demonstration data to better estimate out-of- of-distribution actions. Traditional imitation learning typically requires direct learning from demonstrations, either by recovering reward function from demonstrations, or matching expert and agent state-action distributions (e.g. GAIL-(Ha and Ermon 2016)), or incorporating demonstrations anchors to mitigate overestimation and instability caused by out-of-distribun novel use of demonstrations by treating them structured roadmap for building an adaptive curriculum, scaffolding the agent's learning to enable steady progress toward solving the full task.
actions (e.g. Cal-QL((Naaaman
et al. ). In contrast, Cago introduces
Curriculum for Learning from Demonstrations. Several prior works have explored curriculum design using demonstration. Yengera et al. (2021); Tao et al. (2024) introduce difficulty scores to rank demonstrations, offering theoretical framework for selecting optimal trajectories to scaffold learning. Task Phasing (Bajaj et al., (223)) automatically extracts curriculum phases from demonstrations and dynamically transitions the agent through them during training. JSRL (Uchendu et al., 2023) design a curriculum-based approach that leverages guide-policy pretrained from demonstration to guide early-stage exploration during online training. Compared to these curriculum learning methods using demonstration, Cago employs a goal-level curriculum that incrementally samples intermediate goals from demonstrations based on the agent's evolving capabilities. A number of methods have explored the use of expert demonstration states as informative priors for guiding agent learning. A common strategy involves resetting the agent from states sampled along demonstration trajectories (Nair eet al. 2018: Peng et al.) 2018 Hosu and Rebedeal (0166), enabling the agent to experience these expert regions of the state space without exhaustive exploration. Some approaches employ structured curricula to sequence these resets, either through manually designed progressions (Zhu et al.) 20118) or automated strategies such as reverse curricula, which gradually increase the exploration horizon by training the agent from goal states backward (Resnick et al.] 2018, [Sllmaans and Chen 2018
Tao et al., 2024). While effective in simulation, these methods rely on the ability to precisely reset the environment to arbitrary demonstration states requirement that poses significant challenges in physical systems, where replicating exact configurations, including latent dynamics like joint velocities, is often infeasible. Rather than forcibly placing the agent into expert states a strong form of intervention-Can Cago encourages the agent to actively reach intermediate goals that are sampled to match its current competence. This self-directed learning process allows the agent to internalize groleem-solving skills more effectively, promoting deeper understanding through its own attempts rather than directly resetting. While demonstrations provide useful guidance, true mastery requires the agent to explore and overcome challenges through trial and error.
Conclusion
We introduce Cago, a novel method that leverages demonstrations in a dynamic, goal-guided manner to tackle exploration challenges in sparse-reward environments. By continuously monitoring the agent's capabilities, Cago constructs an adaptive curriculum that incrementally samples intermediate goals from demonstrations, effectively scaffolding learning and enabling steady progress toward solving the full task. Extensive experiments show consistent improvements over baselines, A key limitation of Cago lies in its reliance on resetting environments to the initial states specified in the demonstrations. Appendix C provides a discussion on how this requirement could be relaxed to improve applicability in real-worke settings.
Reproducibility Statement
The code for Cago is available atlattps///githui http:///gggghtb.ooo/////tttttatt com/RUUAAttomated Reasoning-Group. For hyperparameter settings, please refer to Appendix H
10

---

Acknowledgements
We thank the anonymous reviewers for their comments and suggestions. This work was supported by NSF Award #CCF-2124155.
References
Abbeel, P. and Ng, A. Y. (2004). Apprenticeship learning via inverse reinforcement learning, In Brodley, C. E., editor, Machine Learning. Proceedings of the Twenty-first International Conference (ICML 2004), Banff. Alberta, Canada, July 4-8, 2004, volume 69 of ACM International Conference Proceeding Series. ACM.
Arulkumaran, K. and Lillrank, D. O. (2024). A pragmatic look at deep imitation learning. In Asian Conference on Machine Learning, pages 58-73, PMLR.
Bain, M. and Sammut, C. (1995). A framework for behavioural cloning. In Machine intelligence 15, pages 103-129.
Bajaj, V., Sharon, G., and Stone, P. (2023). Task phasing: Automated curriculum learning from demonstrations. In Proceedings of the International Conference Automated Planning and Scheduling, volume 33, pages 542-550.
Ball, P. J., Smith, L., Kostrikov, I,, and Levine, S. (2023). Efficient online reinforcement learning with offline data. In International Conference on Machine Learning, pages 1577-1594. PMLR.
Blondé, L. and Kalousis, A. (2019). Sample-efficient imitation learning via generative adversarial nets. In Chaudhuri, and Sugiyama, M., editors, The 22nd International Conference on Artificial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan, volume 89 of PMLR.
Proceedings of Machine Learning Research, pages
Brantley, K., Sun, W., and Henaff, M. (2020). Disagreemeent-regulaiieee imitation learning, In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
Cui, G., Wang, Y., Mao, W., Duan, Y., and Zhu, H. (2025). Abstraction refinement-guidee program synthesis for robot learning from demonstrations. Proceedings of the ACM on Programming Languages, (00PSLA2)):555
Dadashi, R., Hussenot, L., Geist, M., and Pietquin, O. (2020). Primal wasserstein imitation learning. arXiv preprini arXiv:2006.04677
Dai, S., Hofmann, A., and Williams, B. (2021). Automatic curricula via expert demonstrations. arXiv preprint arXiv:2106.091::
DeMoss, B., Duckworth, P., Hawes, N., and Posner, (2023). Ditto: Offline imitation learning with world models. arXiv preprint arXiv:2302.03000 03006........
Duan, Y., Cui, G., and Zhu, H. (2024a). Exploring the edges of latent state clusters for goal- conditioned reinforcement learning. arXiv preprint 1rXiv:2411.01394
Duan, Y. Mao, W., and Zhu, H. (2024b). Learning world models for unconstrained goal navigation. Advances in Neural Information Processing Systems, 37:59236-59265.
Ecoffet, A., Huizinga, J., Lehman, J., Stanley, K. O., and Clune, (2019). Go-explore: approach for hard-exploration problems. arXiv preprint 1rXivv/901.100
Eysenbach, B., Levine, S., and Salakhutdinov, R. R. (2021). Replacing rewards with examples: Example-basee policy search via recursive classification. Advances in Neural Information Process- ing Systems, 34:115 1541115111552
Finn, C., Christiano, P. F., Abbeel, P., and Levine, (2016). A connection between generative adver- sarial networks, inverse reinforcement learning, and energy- based models. CoRR, 0ts//611.03855

---

Fu, ,, Luo, K., and Levine. S. (2017). Learning robust rewards with adversarial inverse reinforcement learning, arXiv preprint arXiv:1710.11245.
Ghasemipour, S. K. S., Zemel, R., and Gu, S. (2020). A divergence minimization perspective on imitation learning methods. In Conference on robot learning, pages 1259 1277. PMLR.
Gu, J., Xiang, F., Li, X., Ling, Z., Liu, X., Mu, T., Tang, Y., Tao, S., Wei, X., Yao, Y., al. (2023). Maniskill2: A unified benchmark for generalizable manipulation skills, arXiv preprint 234.:29302.44659
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. (2019). Dream to control: Learning behaviors by latent imagination. arXiv preprint
Hafner, D., Lillicrap, T., Norouzi, M., and Ba, J. (2020). Mastering atari with discrete world models. arXiv preprint ar arXiv:2010.02193.
Hafner, D.. Pasukonis, J., Ba, J., and Lillicrap. T. (2023). Mastering diverse domains through world models. arXiv preprint arXiv:2301.044 00100444.
Hansen, N., Lin, Y., Su, H., Wang. X., Kumar, V., and Rajeswaran, A. (2022). Modem: Accelerating visual model-based reinforcement learning with demonstrations. arXiv preprint arXiv:2212.05668.
Hermann, L., Argus, M., Eitel, A., Amiranashvili, A., Burgard, W., and Brox, T. (2020). Adaptive curriculum generation from demonstrations for sim-to-real visuomotor control. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 6498-6505 IEEE.
Ho, J. and Ermon, S. (2016). Generative adversarial imitation learning. Advances neural informa- fion processing systems, 29.
Hosu, I. and Rehedea, T. (2016). Playing atari games with deep reinforcement learning and human checkpoint replay. CoRR, 0ass/167.05000
Hu, E. S., Chang, R., Rybkin, O., and Jayaraman, D. (2023). Planning goals for exploration. arXiv preprint 2rXIv:2303.13002.
Jena, R., Agrawal, S., and Sycara, K. P. (2020). Addressing reward bias in adversarial imitation learning with neutral reward functions. CoRR, abs/2009.09a. 009477.
Kim, K. and Park, H. S. (2018). Imitation learning via kernel mean embedding. In Mcllraith, A. and Weinberger, K. Q. editors, Proceedings of the Thirty-Secome AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pages 3415-3422. AAAI Press.
Kolev, V., Rafailov, R., Hatch, K., Wu, J., and Finn, C. (2024). Efficient imitation learning with conservative world models. In 6th Annual Learning for Dynamics & Control Conference, pages 1777-1790. PMLR.
Kostrikov, I., Agrawal, K. K., Dwibedi, D., Levine, S., and Tompson, (2018). Discriminator-acttnn critic: Addressing sample inefficiency and reward bias in adversarial imitation learning. arXiv preprint 11Xiv:1809.088
Kostrikov, I., Nachum, O., and Tompson, J. (2019). Imitation learning via off-policy distribution matching. arXiv preprint arXiv:19/2.0500
Kumar, A., Zhou, A., Tucker, G., and Levine, S. (2020). Conservative q-learning for offline reinforcement learning. Advances in neural information processing systems, 33:1 33:1179-1191.
Liu. M., Zhu, M., and Zhang, W. (2022). Goal-conditioned reinforcement learning: Problems and solutions. arXiv preprint ImXiv:2201.08299
Mendonca, R., Rybkin, O., Daniilidis, K., Hafner, D., and Pathak, D. (2021). Discovering and achieving goals via world models. Advances in Neural Information Processing Systems, 34:24379- 1.391.
12

---

Nair, A., McGrew, B., Andrychowicz, M., Zaremba, W., and Abbeel, P. (2018). Overcoming exploration in reinforcement learning with demonstrations. In 2018 IEEE international conference on robolics and automation (ICRA), pages 6292-6299. IEEF.
Nakamoto, M., Zhai, S., Singh, A., Sobol Mark, M., Ma. Y., Finn, C., Kumar, A., and Levine, S. (2023). Cal-ql: Calibrated offline pre-training for efficient online fine-tuning. Advances in Neural Information Processing Systems, 36:62244-62288
Orsini, M., Raichuk, A., Hussenot, L., Vincent, D., Dadashi, R., Girgin, S., Geist, M., Bachem, O., Pietquin, O., and Andrychowicz, M. (2021). What matters for adversarial imitation learning? In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W., editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, viriual, pages 14656-14668.
Peng. X. B., Abbeel, P., Levine, S., and van de Panne, M. (2018). Deepmimic: example-guide deep reinforcement learning of physics-based character skills. ACM Trans. Graph., (7):
Plappert, M., Andrychowicz, M., Ray, A., McGrew, B., Baker, B., Powell, G., Schneider, J., Tobin, J., Chociej, M., Welinder, P., et al. (2018). Multi-goal reinforcement learning: Challenging robotics environments and request for research, arXiv preprint 1802.0946
Rafailov, R., Yu, T., Rajeswaran, A., and Finn, C. (2021). Visual adversarial imitation learning using variational models. Advances in Neural Information Processing Systems, 34:3016-3028.
Rajeswaran, A., Kumar, V., Gupta, A., Vezzani, G., Schulman, J., Todorov, E., and Levine, S. (2017). Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. arXiv preprint arXiv: 1ee::1109.00990
Reddy, S., Dragan, A. D., and Levine, S. (2019). Sqil: Imitation learning via reinforcement learning with sparse rewards. arXiv preprint arXiv:7903 1080.
Resnick, C.,, Raileanu, R., Kapoor, S., Peysakhovvchh, A., Cho, K., and Bruna, J. (2018). Backplay: "man muss immer umkehren". CoRR, abs/1880 1807.06919.
Salimans, and Chen, R. (2018). Learning montezuma's revenge from a single demonstration. CoRR, abs/1812.0333
Sco, Y., Hafner, D., Liu, H., Liu, F., James, S., Lee, K., and Abbeel, P. (2023). Masked world models for visual control. In Conference on Robot Learning, pages 1332-1344. PMLR.
Swamy, G., Choudhury, S., Bagnell, J. A., and Wu, S. (2021). Of moments and matching: A game- theoretic framework for closing the imitation gap. In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 10022-10033 PMLR.
Tao, S., Shukla, A., Chan, T., and Su, H. (2024). Reverse forward curriculum learning for extreme sample and demo efficiency. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net.
Tao, S., Xiang, F., Shukla, A., Qin, Y., Hinrichsen, X., Yuan, X., Bao, C.,, Lin, X., Liu, Y., kai Chan, T., Gao, Y., Li, X., Mu, T., Xiao, N., Gurha, A., Rajesh, V. N.N Choi, Y. W., Chen, Y.-.., Huang, Z.. Calandra, R., Chen, R., Luo, S., and Su, H. (2025). Maniskill3; Gpu parallelized robotics simulation and rendering for generalizable embodied ai. Robotics: Science and Systems.
Torabi, F., Warnell, G., and Stone, P. (2018). Behavioral cloning from observation. arXiv preprint arXiv:1805.0195..
Uchendu, I., Xiao, T., Lu, Y., Zhu, B., Yan, M., Simon, J., Bennice, M., Fu, C.,, Ma, C., Jiao, J., et al. (2023). Jump-start reinforcement learning. In International Conference on Machine Learning, pages 34556-3458 PMLR.
13

---

Wang, R., Ciliberto, C., Amadori, P. V., and Demiris, Y. (2019). Random expert distillation: Imitation learning via expert policy support estimation. In Chaudhuri. K. and Salakhutdinov, R., editors. Proceedings of the 36th International Conference on Machine Learning, ICMI. 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 3544. PMLR.
Yengera, G., Devidze, R., Kamalaruban, P., and Singla, A. (2021). Curriculum design for teaching via demonstrations: theory and applications. Advances Neural Information Processing Systems, 34:10496-
Yu, T., Quillen, D., He, Z., Julian, R., Hausman, K., Finn, C., and Levine, S. (2020). Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages 1094-1 100. PMLR.
Zhu, Y., Wang. Z., Merel, J., Rusu, A. A., Erez, T., Cabi, S., Tunyasuvunakool, S., Kramár, J., Hadsell, R., de Freitas, N., and Heess, N. (2018). Reinforcement and imitation learning for diverse visuomotor skills. Kress-Gazit, H., Srinivasa, S., Howard, T., and Atanasov, N., editors, Robotics: Science and Systems XIV, Carnegie Mellon University, Pittsburgh, Pennsylvania, USA, June 36-30, 2018.
Ziebart, B. D., Maas, A. L., Bagnell, A., and Dey, A. K. (2008). Maximum entropy inverse reinforcement learning. In Fox, D. and Gomes, P., editors, Proceedings of the Twenty-Third AAAI Conference on Artificial Intelligence, AAAI 2008, Chicago, Illinois, USA, July 13-17, 2008, pages 1433-1438. AAAI Press.
14

---

Appendix
A Proof of Theorem 1
Theorem 1. Let M denote the true dynamics model and M the learned model. Let TBC be behavior-cloned policy, and ı new policy. Let Demo be dataset of expert demon- Close-
strations from an unknown expert policy. Suppose that, for all t = 0, .,
(1)
ness of behavior cloning: Drv (eMiBBC, dPam I K, (2) Model learning error under BC:
E (sa) Drv (M s, a), (M(-1s,))
M(. 8, a) I p. and (3) Trajectory distribution closeness:
Drv PM
v. Then for all = 00,,
T, we have:
E (s.a)rd am Dww (M( s.a). s.ai)] s 2x + 2v.
Proof. By triangle inequality expectations and total variation:
(saa)aafa Min DTTV (MC- s, a(. (0((-11s))) s,))
DDv (M(. a), [s.a)]] 2 Drv (df (d
$ E (s.a)ra Puam
We now bound each term separately:
For the first term, apply triangle incquality again:
DTV (MC (M) ( a)) s E (o.a) Minuc (DrV (.(.ss,), s,ai)]
E
(s.a)rdPPomm Peoero
+ 2 DTv (dP (dRamm , d.M.TBC)
$ μ + 21s
using assumptions (1) and (2).
For the second term, use that marginal total variation is bounded by trajectory total variation:
Drv (o
M. ,dpamo) Drv (pMa. .-,.F-) I v,
by assumption (3).
Combining:
(s.a)aafhi Mir DTV (M( (M(-1,,a),MMp ((.1,1,1)]] S ut 2x + 2u.
E
B Extended Background
B.1 Dreamer World Model
The RSSM consists of encoder, recurrent model, representation model, transition predictor, and a decoder, as formulated in Equation [0] And employs an end-to-end training methodology, where its parameters are jointly optimized based the loss functions of various components, including dynamic transition prediction, reward prediction, and observation recoding-decoding These components often operate in latent space rather than the original observation space, as encoded by the World Model. Therefore, during end-to-end training, the losses of all components indirectly optimize the latent space.
The encoder f encodes the input state 2, into embed state et, which is then fed with the deter- ministic state h; into the representation model qu to generate the posterior state 2t, The transition
15

---

predictor Pe predicts the prior state & based on the deterministic state h, without access to the current input state 2t. Using the concatenation of either (h,, 2x) or (he, 2x) input, the recurrent transition function f iteratively updates the deterministic state he with given action de.
Encoder:
et fE(ee(zz)
Recurrent model:
hr - ,, ii- at-1)
Representation model:
21 g lef ex)
(10)
Transition predictor:
& ~ Po(Eellte)
Decoder:
fo [D(eelles
B.2 Temporal Distance Training in LEXA
The goal-reaching reward n is defined by the self-supervised temporal distance objective (Mendonca) et al., 2021) which aims to minimize the number of action steps needed to transition from the current state to a goal state within imagined rollouts. We use be to denote the concatenate of the deterministic state h, and the posterior state 2e at time step t.
b = (h,, 2x)
(11)
The temporal distance Dr trained by sampling pairs imagined states b,, de+x from imagined rollouts and predicting the action steps number & between the embedding of them, with predicted embedding & from b, to approximate the true embedding e, of the observation Zt.-
Predicted embedding:
emb(b.)) = és R Ct,
where e = [E(())
(12)
Temporal distance: Di(ée, (x++) R k/H
where é = emb(br) êtt2 = emb(br+*)
(13)
Cb be4n) = -Di(er, (c4x)
(14)
C Discussion
In this section, we discuss several key questions and extensions regarding the design, applicability, and limitations of Cago. Our goal is to clarify how Cago fundamentally differs from existing imitation learning paradigms, what design choices drive its efffcttiveness, and how it can be adapted or extended to broader settings. We also explain and clarify its current assumptions- such as initial state resetting and model-based training. Through this discussion, we aim to provide deeper insights into the generality, scalability, and future potential of Cago as a new framework for learning from demonstrations.
How does Cago differ from traditional Imitation Learning methods?
Cago is fundamentally different from traditional Imitation Learning (IL. methods, both in methodol- ogy and motivation. IL typically requires direct learning from demonstrations, either by recovering reward function from demonstrations, or matching expert and agent state-action distributions (e.g.
GHIL(Ho and Ermon 2016)), or incorporating demonstrations anchors mitigate overestimation and instability caused by out-of-distributin actions (e.g. Cal-QL(Nakamad et al. (023)). In contrast, Cago presents new paradigm utilize demonstrations: it dynamically tracks the agent's 's competence along demonstrated trajectories and uses this signal to select intermediate states in the demonstrations that are just beyond the agent's current reach as goals to guide online trajectories collecting. To evaluate this novel perspective, we compared Cago against state-of-the the--at baselines that represent diverse strategies for learning from demonstrations, including distribution mathcing (e.g. GAIL, and Ermon 2016) and PWIL (Dadashi et al.) 2020)), curriculum learning (e.g. J.RISSRUUU. (Uchnnu et al. 2023), model exploration (e.g. MoDem(Hansse
al. (022)) and offline-to-online fine-tuning (c.g. (alQL(Nakamon et al. 2023 and RLPD Ball et al. 2023
16

---

2. Can Cago be extended with more sophisticated similarity metrics to handle complex tasks?
In our implementation of Cago, we use the observation space as the goal space and adopt simple similarity metrics such as MSE or 1.2 distance to demonstrate the applicability and robustness of Cago's curriculum strategy in both state-based and pixel-based environments. Cago can be naturally extended to support more complex settings. For example, in high-dimensional visual environments, Cago could adopt more expressive similarity metries such SSIM (Structural Similarity Index) or FSIM (Feature Similarity Index). In tasks involving semantic or language-conditioned goals, one could compute a scalar score function that estimates how well a visual state satisfies goal, and then compare states based on the difference in their scores with respect to the same goal. This score function can be instantiated using pretrained vision- -aniunne models, allowing goal representations to move beyond raw states to language embeddings or other semantic forms. Our main contribution is a general framework for goal sampling from demonstrations- based on the agent's 's capability- which can be paired with any goal space and similarity metric appropriate to the domain.
3. How does Cago perform relative curriculum learning baselines?
In our experiments, we included Jump-Start Reinforcement Learning (JSRL(Uchendu et al. 2023 ).
recent and strong curriculum learning baseline. JSRL pretrains guide policy offline data to provide a curriculum of starting states for the RL policy, which significantly simplifies the exploration problem. As the RL policy improves, the effect of the guide-policy is gradually diminished, leading to an RL-only policy that is capable of further autonomous improvement. We discuss four additional
curriculum learning baselines: Yengera et al. all](222)) (2021): Bajaj al. 2023); y Dai al. (202T) Hermann et al. 22020) Yengera et al. (2021) constructs a personalized demonstration curriculum by computing task task-specific difficulty scores based the teacher's s and learner's policies. In contrast, Cago does not require hand-crafted difficulty metrics-instead it extracts intermediate goals directly from the demonstration, and guides the agent to reach them progressively based on its own learning progress. This avoids the challenge of designing difficulty metrics for complex manipulation tasks, which is non-trivial without expert knowledge. [Bjjaj et al. (2023) assumes access demonstrator policy, either via imitation learning or inverse reinforcement learning. However, our experiments show that standard imitation learning performs poorly on our benchmarks. limiting the effectiveness of such approaches in our domain. Dai et al. (202T): Hermann et al. (2020) segment demonstrations and construct a curriculum by resetting episodes to intermediate states along the demonstration trajectory, starting from the end and moving backward as learning progresses. Cago explicitly avoids such state resets, which are often unrealistic in real-world robotic settings without simulator support. we conducted additional experiments comparing Cago to ACED et al. (202I) and Task phasing(Baja) et al. 2023) on the Fetch-PickPlace and Fetch-Slide environments that they were evaluated on. The tables below report the final success rates after IM training steps, averaged over 5 random seeds (We were unable to run ACED on Fetch-Slide).
Table I: Comparison of Cago and ACED, Task phasing across Petch-PickPlace and Fetch-Slide tasks.
Task
Cago (Ours) ACED (Dai et al. 2021
Task phasing (Bajaj et al. 2023
Fetch-PickPlacc
1.00
0.93
0.75
Fetch-Slide
0.48
0.20
4. Is the need for environment resets in Cago a fundamental limitation?
Indeed, resetting to specific initial states is limitation of Cago, but as we mentioned in the Section 5 it is already a substantial improvement over arbitrary state resets, which are often infcasible in real- world robotics. Importantly, Cago does not require exact resets to the full demonstration initial states. For instance, in the Peg Insertion task, the peg's initial position can differ from the demonstration, and Cago can learn to reach states within the early portion of the demonstration. However, some elements liik the hole and the box-are are mon-movable and randomly initialized by the environment. Since the agent cannot manipulate these objects, we have to reset them to match the demonstration configuration to ensure goal reachability. The constraint over other controllable components, such as the robot arm and peg, can be relaxed. This partial reset strategy reduces the burden and makes Cago more practical for real-world settings.
17

---

5. What motivates the choice of model-based RL formulation for Cago, rather than model-free alternative?
In early experiments, we attempted to train Cago in a model-free manner by learning the goal- conditioned policy directly from real environment rollouts. But further find that instead of directly using Peup to train T in Model-free manner, the world model can offer richer set of imagined roll- outs to train the T², which improves the learning efficiency. In our setting, the simulated trajectories have starting states and goal states randomly sampled from the same trajectory in Deap- As result, the simulated trajectories still resemble segments from Deap. while significantly increasing trajectory diversity and data richness. The LLBRL-Ieveragee
simulated trajectories allows us to generate more diverse and abundant training data. This promotes generalization across the environment. Cago builds upon the Dreamer framework (Hafner et al. 2019 2020 Huu al. 2023 Duan al. 2024a b). which serves as the model-based cornerstone of our approach. Dreamer learns latent world model and trains a goal-conditioned policy using an actor-critic algorithm. samples trajectories by set- ting the final observation of a demonstration as the goal, while resetting the environment using the same seed as the corresponding demonstration. PEG (Hu et al., 2023 introduces an alternative goal-selection strategy that enhances exploration within the Dreamer framework. To demonstrate the improvement brought by Cago over its cornerstone, we conduct comparisons against both Dreamer and PEG. Table 2 presents the performance comparison of Cago against Dreamer and PEG across three environments (results are averaged over 5 random sceds). Cago clearly surpasses its model-basee comnerstone, Dreamer, well as the competitive goal-pies goal-picking strategy, PEG, across all environments.
Table 2: Performance comparison of Cago against Dreamer and PEG.
Environment
Cago (Ours)
Dreamer
PEG
StickPush
0.99
0.68
0.76
Disassemble
0.80
0.02
0.04
Adroit-Pen 0.82
0.52
0.54
6. Can Cago be applied to broader scenarios-suce as real-world settings with diverse initial states and uncertain dynamics?
Cago can be extended to support diverse initial states and uncertain dynamics. For diverse initial states in real-world settings, Cago can gradually prioritize those where the agent has the most learning potential-starting from demonstration initial states and expanding to nearby ones where the agent can still benefit from the demonstration trajectories for dynamic capability tracking and progressive goal selection. As training progresses, we incorporate successful rollouts from these nearby initial states into the demonstration set, broadening its coverage. This enables natural curriculum that eventually transitions to diverse, weal-world initial states. For uncertain, stochastic dynamics, Cago remains robust by not requiring exact reproduction of actions in demonstration trajectories. Instead, it adapts goals based on observed agent successes, making the curriculum self-correcting in the face of noise or variability. If the agent consistently fails reach specific demonstration state, Cago's capability metric naturally shifts goal sampling toward more achievable targets. Additionally, goal selection can be extended with probabilistic estimates such as goal-rea likelihood future state predication uncertainty from an ensemble of world which we view as promising direction for high-stochastieity domains.
Does Cago risk overfitting to a limited number of demonstrations?
Cago mitigates overfitting to a small number of demonstrations by training its goal-conditioned policy T within a world model. Rather than training TC to reach states solely from demonstrations, Cago instructs G to learn to reach diverse states sampled from real rollouts in Deup. which records the trajectories generated under our Go-Explore paradigm with capability-aware aware goal sampling. in its world model. This promotes generalization across the environment. To evaluate the generalization capability, we tested Cago on 500 unseen initial states generated from random seeds, each differing from those in the demonstrations. We report both the average 1.2 deviation from demonstration initial states and the average success rate across 500 unseen seeds across three environments Table The results are consistent with Figure]
18

---

Table 3: Initial demo observation deviation and average success rate across 500 unseen seeds.
Environment
Average Initial Observation Deviation
Average Success Rate(500 seeds)
StickPush
0.0547
0.982
Disassemble
0,1405
0,806
Adroit-Pen
1.2749
0.824
D Limitations and Future Work
While Cago shows promising results in leveraging demonstrations to improve exploration and learning efficiency in sparse-reward, long-horizon tasks, several limitations remain. In our experiments, we employ image MSE similarity(or 1.2 distance for state) and simple counting scheme, which generalizable and applicable across diverse environments, as illustrative measures of evaluating agent capability in Cago. However, in much more complex environments, such simple approaches may not constitute the optimal strategy. In the future, incorporating uncertainty-awail models for both agent capability estimation and subgoal selection could improve robustness and adaptability of Cago framework. Extending Cago to multi-agent or real-world robotics scenarios is another promising direction, where the complexity coordination and physical constraints introduces new challenges for efficient demonstration utilization and goal guidance.
E Environment Details
We evaluate and compare Cago against several baselines across three robot environment suites using sparse rewards: MetaWorld (Yu et al.] 2020), Adroit (Rajeswaran et al. 2017), and Maniskill (Gu al. 2023 Tao et al. 2025). In this section, we provide more details about each benchmark and the specific experimental setup.
E.1 MetaWorld
Disassemble
PickPlaceWall
ShelfPlace
StickPull
StickPush
Figure 6: very hard level environments from MetaWorld World
MetaWorld (Yu et al. 2020) is widely used benchmark suite designed to evaluate the generalization and manipulation capabilities of reinforcement learning algorithms in robotic control tasks. It consists of a diverse collection of continuous control environments simulated using the MuJoCo physics interact with objects the scene achieve goal-directed
engine. Each task requires robotic behavior under sparse or dense reward settings. MetaWorld includes 50 distinct tasks of varying difficulty, ranging from simple reaching to complex object manipulation. our experiments, we focus on the "very hard" subset of tasks identified in prior studies (Sco et al. 2023 Hansen et al. 2022). which are characterized by sparse rewards, delayed feedback, and the need for precise lowelevel control, making them particularly suitable for benchmarking sample efficiency and generalization in learning frameworks. The five very hard tasks we choose are: Shelf Place, Disassemble, StickPull, Stick Push, Pick Place Wall. Shelf Place: The agent must grasp an object and place accurately onto shelf, requiring precise vertical and lateral arm control. Disassemble: The task involves picking nut out of the peg, demanding strong grasp and directional pulling motion. Stick Pull: The robot needs to grasp a stick and pull a bottle, requiring fine foree control and coordinated motion. Stick Push: The goal is grasp stick and push bottle, emphasizing
demonstration-augnen
19

---

three envi- controlled contact and alignment with the target location. Pick Place Wall: The agent must pick up an object and place it over wall barrier onto a specified target location, combining lifting, positioning. and obstacle avoidance. We use the 1.2-distance to calculate the similarity between observations to judge if the agent has reached the demonstration observation. The threshold for similarity is 0.05. We use 10 demonstrations for training.
E.2 Adroit
Hammer
Pen
Door
Figure environments from Adroit Suite
The Adroit Suite (Rajeswaran al. 2017 is a set of high-dimensional dexterous hand manipulation tasks that emphasize fine motor control, contact-rich interactions, and sparse-reward learning. is built upon 24-DoF ShadowHand robotic hand, presenting significant challenges in both control and generalization. Each task requires the agent to coordinate multiple fingers and joints to manipulate objects with high precision under partial observability and complex dynamics. We ronments from this suit. Hammer: The agent must grasp a hammer and use it to drive a nail into box. This task demands stable object manipulation, precise tool orientation, and effective force transmission. Pen: The objective is to reorient and position pen in an assigned direction. requires careful control of finger articulation and rotational dexterity. Door: The task involves unlatching and opening a door by manipulating the handle and applying a pulling motion. It tests the agent's ability perform multi-stage interactions and coordinate wrist and finger movement to exert torque in the correct direction. We use the Mean square eπror between images rendered to calculate the similarity judge if the agent has reached the demonstration observation. Each rendered image will be reshaped to size 1000,, We use 10 demonstrations for training.
E.3 Maniskill
StackCube
PegInsertionSide
PullCubeTTool
Figure environments from Maniskill
ManiSkill (Gu et al. 2023 Tao et al. 2025) is a comprehensive benchmark suite designed to evaluate generalizable robotic manipulation skills simulation, emphasizing real-world task diversity, object variety, and generalization across instances. It provides high-quality 3D environments with continuous control, supporting both visual and proprioceptive observations. The benchmark is particularly challenging under sparse reward settings, tasks often require multi-step reasoning, long- Hong-horizon
20

---

planning, and precise control to accomplish. We pick three environments from this benchmark. PullCubeTool: Given an L-shaped tool that is within the reach of the robot, the agent needs to leverage the tool to pull a cube that is out of it's reach. PegInsertionSide: The robot must align and insert a peg into a side-entry slot. The task demands precise pose estimation, spatial reasoning, and careful control to avoid misalignment or jamming. StackCube: This task involves picking up a cube and accurately stacking it on top of another. We use the Mean square error between images rendered to calculate the similarity judge if the agent has reached the demonstration observation. Each rendered image will also be reshaped to a size (100,100,3). We use 20 demonstrations for scale up the position(x.y.2) 10 times and normalize
training. On StackCube and PegInsertionSide, the degree of griper opening for more stable learning. We set the clearance of the hole to 0.01 in PegInsertionSide so that the peg could be inserted more easily.
F
More Experiments
F1 More Baselines
Generative Adversarial Imitation Learning (GAIL) (Ho and Ermon 2016 Arulkumaran and Lillrank 2024) adopts an adversarial learning framework where discriminator trained to differentiate between expert and agent trajectories; the discriminator's 's output is then used as the reward signal for the agent. Primal Wasserstein Imitation Learning (PWIL) (Dadashi al. 2020; Arulkumaran and Lillran 2024) formulates imitation as a primal optimization problem that minimizes the Wasser- stein distance between expert and agent trajectory distributions. It constructs a shaped reward function directly from this distance, encouraging the agent to produce expert-like behaviors. Soft Q Imitation Learning (SQIL) (Reddy al. 2019) simplifies imitation learning by assigning fixed rewards to expert and agent transitions, effectively transforming imitation into a standard reinforce- ment learning problem with sparse binary rewards. Value-based Distribution Correction Estimation (ValueDice) (Kostrikov et al.) 2019) takes a issribution-matchin perspective by minimizing a divergence between expert and agent state-action occupancy measures, providing principled con- nection between imitation and value-based reinforcement learning. Reinforcement Learning with Prior Data (RLPD) (Ball et al. 2023) is a state-of-the-an baseline improving the efficiency of online reinforcement learning by leveraging offline data. We run these basclines over 5 random experimental seeds and report the average success rate.
As the results shown in Figure Cago generally outperforms all additional baseline approaches across diverse set of manipulation tasks. In tasks such Disassemble, StickPull, Hammer, and Pen, Cago demonstrates significantly faster convergence and higher final success rates, indicating its superior learning efficiency and robustness. Particularly in Maniskill environments, Cago is the only method that achieves meaningful learning progress, while all baselines fail to get any success, highlighting the goal sampling in challenging. sparse-reward environments. Although
importance of capability- GAIL achieves the best performance in the ShelfPlace environment, this success is not representative of its overall effectiveness. In all other tasks, GAIL performs poorly, exhibiting highly unstable learning process and low final success rates.
F.2 Visual goal predictor
To evaluate the generalization ability of our visual goal predictor Pt, we visualize its predicted goal images given initial observations at test time in Figure I0 For each task, we compare the predicted goal (middle column) to the -roun final observation from demonstration trajectory (right column). Importantly, these demonstrations are collected from unseen seeds that were never used during training. These results show that Ps is capable of accurately inferring the final goal state purely from a single initial image observation, even in unseen evaluation seeds. This predictive capability allows the goal-conditioned policy nt to finish tasks effecttvvly for any environment seeds.
21

---

Daee
Cloce
Cran
Call
Goll
512Ce50 otC
FOE
SUccess Rite
ette
02
02
02
30
08
20
0.2
20
0.4
00
13
0.2
04
stops
steps
(a) Disassemble
(b) PickPlace Wall
(c) ShelfPlace
(d) StickPull
1.0
1.0
1.0
06
PHLL
05
Fark
FALL
0.8
SU/CKS 155
sUccess
stapss
s0ags
(c) StickPush
(1) Door
(g) Hammer
(h) Pen
RUPO
5100e50 Dbt
50L
success mie
08
Pool
(i) PegInsertionSide
(k) StackCube
Figure 9: Experiment results comparing Cago with the additional baselines over 5 random seeds.
Door
Hammer
P
Pen
Initial Observation
Predicted Goal
True Goal
Figure 10: Visual goal prediction results from our learned goal predictor Pa. Each row corresponds to different task (Door, Hammer, Pen). From left to right: the agent's initial observation at test time, the goal image predicted by Pg. and the ground-truth final observation from demonstration trajectory. Notably, these demonstrations are drawn from unseen seeds not used during training. The predicted goals closely match the actual final states, illustrating strong generalization of Ps novel environment seeds.
22

---

F33 Ablation on Number of Demonstration
This section investigates how the number of demonstration trajectories used for training influences the performance of Cago across various tasks. We present success rate curves under different numbers of demonstrations for four representative environments from Metaworld, Adroit, and ManiSkill. The results show that Cago maintains strong performance even with limited amount of expert data (as few as 5 demonstrations), particularly in tasks such as Disassemble, StickPush, and Pen, demonstrating its robustness and stability under data-scarce conditions. Increasing the number of demonstrations in general yield higher final success rates. This effect is especially evident in more challenging tasks, such as the PegInsertionSide task in Maniskill. These findings suggest that Cago's performance can be further improved with sufficient expert demonstrations.
Cape-
damcc Cagle Mepo--deeenna
Capec Lo---dd-eena
50-demcc
Capen Modeeraa
Oagen
Gagio 20- do--erica
0.
Cape Miveemes
0.
Caral Studemes
05
Cagoo Studemes
0.
SU0CKs ne8
SUcces reft
SUccess
00
Stops
stops
Stops
staps
(a) Disassemble
(b) StickPush
(c) Pen
(d) PegInsertionSide
Figure 11: Success rates under different numbers of demonstrations. Results are averaged over 5
random seeds.
F4 Ablation on Demonstration Quality
common concern in demonstrationn-aased reinforcement learning the reliance on high-quality trajectories that adequately cover the task space. While our theoretical analysis (Theorem 1) assumes such demonstrations order to establish guarantees, real-world data are often noisy, incomplete. or even contain failed attempts. To evaluate the robustness of Cago under such imperfections, we conducted ablation studies on both suboptimal and failed demonstrations.
Suboptimal Demonstrations. We first constructed three types of perturbed demonstrations: (i) Missing Observations: We randomly removed 20% of the observations from each original demon- stration trajectory to simulate incomplete data. (ii) Noisy Actions: We added Gaussian noise (scaled by 0.1) to each action output by the expert policy during expert trajectory collection. (iii) Random Actions: We replaced the expert's action with a randomly sampled action with a probability of 20% each timestep during expert trajectory collection.
These settings simulate three kinds of suboptimality; missing state information, noisy expert behavior, and corrupted decision decision-making. Using these demonstrations, we re-evaluated Cago three bench- mark environments: AdroitPen, Stick-Push, and Disassemble. Because we do not have the expert policy for Adroit, only conducted the "Missing Obs" experiment for Adroit-Pen. Table reports the final success rates after 1M training steps, averaged over 5 random seeds (each evaluated over 100 test episodes). The results suggests that Cago is robust to the suboptimality in demonstrations. This is because: First, Cago only uses demonstration trajectories to extract sequences of observations that serve as curriculum goals. It does not attempt imitate the demonstration directly, nor does it rely on demonstrations to infer a reward function. This significantly reduces Cago's dependency on high-quality demonstrations, Second, even when the demonstration trajectories are suboptimal-as long as they represent a successful sequence of observations-Caan can still learn well. We observe that in the Disassemble environment, Cago achieves significantly higher success rates when using the "Random Actions" type of suboptimal demonstrations compared to using the original demonstrations. Those imperfect demonstrations withhhhhhh--hhhhh pauses or detours- cover broader region of the task space, thereby improving the final policy's generalization. More importantly, Cago is guided by a reward function shaped using a Temporal Distance Network, which estimates how many steps are needed to reach goal state. This enables Cago to discover more efficient, often shorter, alternative paths for goal reaching. Therefore, even when the demonstration are suboptimal, Cago may still learn the optimal paths.
23

---

Table Success rates under suboptimal demonstrations (after 1M steps, averaged over 5 seeds).
Environment
Original Demo
Missing Obs
Noisy Actions
Random Actions
StickPush
0,99
0,99
0.97
0,96 0.96
Disassemble 0.80
0.77
0.88
Adroit-Pen 0.82
0.88
Failed Demonstrations. We next considered demonstrations that include failed trajectories. For MetaWorld, injected uniform noise (scale 0.3) into expert actions, producing datasets with 30% and 50% failed demonstrations. We then compared Cago against several strong baselines (GAIL, PWIL, JSRL, and Modem) Adroit-Pen, Stick-Push, and Disassemble. We assume the environments used to collect demonstrations provide sparse (binary) reward for goal reaching, enabling us to train goal predictors exclusively on successful demonstrations. The results are shown Figure 12
13
1.0
Modemm 30%6Far
Moden 3016Fa1
Modem 30%6fal
Success rate
SUcCESS rate
SUcCesS rate
0.4
64
0.4
0.0
6.0
0.2
D.4
0.0
0.5
0.0
0.3
03.5
0.5
0.0
6.0
02
0.6
0.8
0.5
Lo
steps
steps
steps
(a) Disassemble
(b) StickPush
(c) Pen
Figure 12: Success rates under 30% failed demonstrations (after 1M steps, averaged over 5 seeds).
1.0
1.0
0.8
F Modemm 10%6/al
S05iFal Modemr-50%
ISRL-R Moder-
SucCess rate
Success rate
success rate
0..
C4
0.4
0.0
0.2
0.4
0.0
0.5
0.
0.2
0.4
0.5
0.0
0.0
02
0.4
0.0
0.8
Lo
Steps
steps
steps
(a) Disassemble
(b) StickPush
(c) Pen
Figure 13: Success rates under 50% failed demonstrations (after 1M steps, averaged over 5 seeds).
Cago maintains strong performance even with many failed demonstrations, while the baselines degrade significantly. This robustness to suboptimal or noisy data stems from two key factors: (I) Cago prioritizes goals that are both challenging and feasible, regardless of whether they originate from successful or failed demonstrations. While failed trajectories may not directly lead to task success, they can expose the agent diverse regions of the environment, potentially aiding exploration. This allows Cago to extract useful learning signals from a broader range of experiences.
(2) Cago neither imitates demonstrations nor infers reward functions from them. Instead, it uses demonstrations solely to extract sequences of observations that serve candidate goals. These goals are selected based on the agent's 's own learning progress, forming a curriculum that guides policy improvement. This design reduces sensitivity to demonstration quality.
24

---

F,5
Ablation on Parameters
Auisit controls the threshold for how frequently state must be visited before it is considered "mastered" by the agent, and thus eligible for sampling more difficult goals. it is set too small, Cago may prematurely progress to harder goals before acquiring sufficient competence in easier ones. Conversely, if it is set too large, Cago may become overly conservative, spending excessive time on seready-masteree states. We have conducted a sensitivity analysis by setting Avisit 50, 100, 200, and evaluated Cago on three environments: Stick-Push, Disassemble, and AdroitPen. The final success rates after 1M training steps, averaged over 5 random seeds are shown in Table 5
Table Performance with varying Aeisit-
Environment
Avisit = 50
Auisit = 100
Apisit = 200
StickPush 0,94
0,96
0,99
Disassemble 0.79
0.82
0,80
Adroit-Pen 0.77
0.82
0.83
ó defines the goal sampling window size, i.e., the proportion of expert states from which new goals are sampled during training, A large ó value may lead to sampling overly distant or inappropriate goals, affecting stability and learning efficiency. In contrast, very small ó may overly restrict exploration and hinder curriculum progression. We added a sensitivity experiment comparing the default 8 10% with a smaller value 5% and larger value 20%, across the same three environments. Table shows that while larger values can sometimes slightly reduce performance, Cago remains capable of learning effective policies in all settings.
Table Performance with varying 6.
Environment 6 = 5%
6 10%
20%
StickPush
1.00
0.99
0.99
Disassemble 0.81
0.80
0.75
Adroit-Pen
0.76
0.82
0.81
G Runtime
G.1 Experiment total runtimes
Table 7: Runtimes per experiment.
Environment
Runtime (h)
Benchmark
Steps
ShelfPlace
72
MetaWorke
le6
Disassemble
MetaWorld
1c6
StickPull
72
MorlD
le6
StickPush
Meta World
le6
PickPlaceWall
72
Meta World
le6
Adroit-Door
Adroit
le6
Adroit-Hammer
85
Adroit
le6
Adroit-Pen
Adroit
le6
PullCube Tool
78
ManiSkill
le6
PegInsertionSide
143
ManiSkill
5e6
StackCube
155
ManiSkill
5e6
25

---

G.2 Computation Time for Updating the Cago Visitation Record Dictionary
In this section, we analyze the computational cost associated with updating the visitation record dictionary. Let the length of sampled trajectory be denoted Ly, and let - represent the demonstration trajectory associated with the same environment reset seed, having length L5. The visitation record dictionary Dictysa updated according to Equation
Dict,isi [T]]] +=1 if sim(st, 8 Wel,
W......., ......,,,,,V.,... W.....,,..V...........
This update rule implies that for each step in the sampled trajectory, a similarity check is performed against all steps in the corresponding demonstration trajectory. Thus, the time required to perform an update of Dictyisit can approximated by:
Time(Update) R Total Steps x L, x Time(Similarity Calculation)
The computational cost is therefore influenced by three main factors: the total number of interac- tion steps, the length of each demonstration trajectory, and the cost of computing the similarity metric. Importantly, the similarity function sim(; :) differs by environment, which directly affects computation time. For MetaWorld environments, we utilize L2-distance in the state vector space (i.e.. Llw--dimension numerical vectors). This calculation is computationally efficient, typically requiring only simple element- operations over vector entries. In contrast, for the Adroit and Maniskill environments, the similarity is computed based on MSE in the image space. This involves pixel-wise comparison over image observations, which increases the computational load due to the large input dimensionality (e.g., 100 100 x 3) result, while the update rule remains structurally the same, the actual runtime overhead for image-based similarity can substantially higher than that for state-bas baste similarity. The table below summarizes the total runtime, update time, and similarity function used for each environment:
Table 8: Computation time and similarity function for updating the visitation dictionary Dictyisi-
Environment
Steps
Runtime (h)
Update Time (h)
Similarity
Disassemble
le6
0.05
State I.2
WickPlaceWall
le6
72
0.05
State L2
ShelfPlace
le6
0.05
State L22
StickPull
le6
72
0.05
State L2
StickPush
le6
0.05
State L.2
Adroit-Door
le6
80
5.3
Image MSE
Adroit-Hammn Hammer
le6
85
5.3
Image MSE
Adroit-Pen
le6
83
5.3
Image MSE
PullCube' Tool
le6
78
2.8
Image MSE
PeglnsertionSide
5e6
143
13.8
Image MSE
StackCube
5e6
155
14.2
Image MSE
H Hyperparameters
We adopt the default hyperparameters from the LEXA backbone model-based RL (MBRL) agent-suce such as the learning rate, optimizer, and network architecture- and maintain them con- sistently across environments. The primary hyperparameter tuning for Cago focuses on the following aspects: (1) the episode length L,; (2) the proportion of L, allocated to the goal-directed phase g (3) the number of demonstrations Na used both dictionary construction and en- vironment resetting; (4) the visit frequency threshold Auisit used in Algorithm for filtering goal candidates; and (5) the similarity calculate metries in Equation 2 (6) the similarity threshold in
Equation] 2
26

---

Table 9: Hyperparameters of Cago.
Environment
L
Tgo rate
Ndemo
Apisit
sim (:.
e
Disassemble
300
0.6
10
200
10%
State L2
0.1
PickPlaceWall
150
0.6
10
200
10%
State L2
0.08
ShelfPlace
150
0.6
10
50
10%
State L2
0.05
StickPull
200
0.6
10
200
10%
State L2
0.08
StickPush
150
0.7
10
200
10%
State L2
0.05
Adroit-Door
200
0.7
10
100
10%
Image MSE
200.0
Adroit-Hammer
100
0.6
10
200
10%
Image MSE
50.0
Adroit-Pen
100
0.7
200
10%
Image MSE
200.0
PullCube
100
0.7
10
100
10%
Image MSE 100.0
Peg SiglnsertionSide
100
0.7
100
10%
Image MSE 100.0
StackCube
100
0.7
20
100
10%
Image MSE
50.0
27

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], , [No] or [NA]. .
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] "provided proper justification is given (c.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). ). general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our abstract and introduction accurately represent the primary contributions of the paper, which include the development of the Cago algorithm designed to solve sparse- reward tasks through a capability key challenges in GCRL, specifically with sparse rewards, and how Cago addresses these by improving the quality of online sampled trajectories using demonstration. These claims are well-supported by the theoretical underpinnings and experimental results presented in the paper, reflecting the scope and impact of the proposed method.
goal picking strategy. The introduction outlines the
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals not attained by the paper.
28

---

2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Our paper thoroughly discusses the limitations of the Cago framework in Appendix.
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Yes, our paper provides the full set of assumptions and a complete (and correct) proof in the Appendix foror Theorem
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
29

---

Answer: [Yes]
Justification: In the Experiment section and appendix of our paper, elaborate the procedure and configuration of our experiments. This includes the sources and modifications of all testing environments, and implementation methods for all basclines, the devices and memory utilized, as well as specific values of hyperparameters employed. Concurrently, we have open-sourced our code; please refer the Reproducibility Statement section for further details.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: As we mentioned the previous justification, we have not only open-sourced our code but also provided detailed steps and settings for reproducing our main experimental results. In the Experiment section and Appendix, we elaborate on the sources and modi- fications of the environments, baseline implementation details, and Cago implementation specifics.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.ce.
for more details.
30

---

While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https //nips wttpccccppppppupcueeee guides/ wwt////odeseiiiiioooooo www.iiiissiooiiioooocooooo.omiii for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, The authors should provide scripts to reproduce all experimental results the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide comprchensive details regarding the hyperparameters essential for understanding the experiments, including those specific to our Cago framework. The table presented (Table. (9] outlines these hyperparameters for each task, facilitating reproducibility and comparison.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We conducted each experiment minimum of five times using different random seeds, and upon plotting the results, as demonstrated in the Experiment section, we incorporated the experimental error. The solid line denotes the average success rate, while the shaded region signifies the standard deviation among the repeated experimental outcomes.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
31

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We clearly specifies the computer resources (8 Nvidia A100 GPU) and the amount of GPU memory required (approximately 2.4GB).
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:://www.aas.cccppccccceeeeea
Answer: [Yes]
The research conducted in our paper aligns with the NeurIPS Code of Ethics. We have thoroughly reviewed the guidelines and ensured that research adheres to ethical standards. Additionally, we have implemented measures to safeguard anonymity and comply with pertinent laws and regulations.
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our research aims to address the learning efficiency problem in Reinforcement Learning (RL) within the GCRL environment. is currently in the theoretical research stage and has minimal societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
32

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Our paper properly credits the creators original owners of assets used, including code, data, and models.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
33

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/dataseta
datasets.papersstt
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
Justification: We have documented our code and provided detailed instructions its usage, licenses, and permissible scope of use. Additionally, we have included the documentation alongside the assets to ensure accessibility and clarity for users.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooell as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification:
Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
34

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: Our method does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https://neurd) hbttep///eerrrppccc///ooeeeeeeee /c//////ooeeeeeee//22////.. for what should or should not be described.
35


Reviews:
[{"review_id": "oWPifB73tk", "reviewer": "Reviewer_xD15", "summary": "This 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.", "questions": "1. 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?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "# 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.", "was_revised": true, "final_justification": "The authors' responses have addressed most of my concerns. They provided additional explanations of their method and experiments to improve the clarity and the quality of validation.", "rebuttal": "We appreciate your insightful feedback and constructive comments!\n\n**1. Cago should be compared with baselines using demonstrations to generate a curriculum [1-4].**\n\nIn our experiments, we included Jump-Start Reinforcement Learning (JSRL, ICML’23), a recent and strong curriculum learning baseline. JSRL pretrains a guide policy on offline data to provide a curriculum of starting states for the RL policy, which significantly simplifies the exploration problem. As the RL policy improves, the effect of the guide-policy is gradually diminished, leading to an RL-only policy that is capable of further autonomous improvement. \n\nRegarding the specific curriculum-based approaches in [1–4]:\n\n[1] constructs a personalized demonstration curriculum by computing task-specific difficulty scores based on the teacher’s and learner’s policies. In contrast, Cago does not require hand-crafted difficulty metrics—instead, it extracts intermediate goals directly from the demonstration, and guides the agent to reach them progressively based on its own learning progress. This avoids the challenge of designing difficulty metrics for complex manipulation tasks, which is non-trivial without expert knowledge. [2] assumes access to a demonstrator policy, either via imitation learning or inverse reinforcement learning. However, our experiments show that standard imitation learning performs poorly on our benchmarks, limiting the effectiveness of such approaches in our domain. [3,4] segment demonstrations and construct a curriculum by resetting episodes to intermediate states along the demonstration trajectory, starting from the end and moving backward as learning progresses. Cago explicitly avoids such state resets, which are often unrealistic in real-world robotic settings without simulator support (Line 83).\n\nTo further address this concern, we conducted additional experiments comparing Cago to [2] and [3] on the Fetch-PickPlace and Fetch-Slide environments that they were evaluated on. The tables below report the final success rates after 1M training steps, averaged over 5 random seeds (each evaluated over 100 test episodes):\n\n| Task                     |   Cago(Ours)  |    ACED[3]          | Task phasing[2]   |\n|-------------------------|--------------------|-----------------------|------------------------|\n| Fetch-PickPlace   |        **1.00**         |          0.93           |         0.75          |\n|  Fetch-Slide          |        **0.48**         |             -              |        0.24          |\n(We were unable to run ACED on Fetch-Slide.)\n\nWe will incorporate these comparisons in the revised paper.\n\n**2. BC Explorer appears to be a crucial component for Cago to work, but its training details are unclear.**\n\nThe BC-Explorer is trained using behavior cloning on a limited set of demonstrations $\\mathcal{D}_{\\text{demo}}$ (Line 3, Algorithm 2). It is a stochastic policy trained to mimic expert actions by maximizing their log-likelihood. This stochasticity—combined with limited demonstrations and controlled training steps—prevents overfitting and ensures a good balance between exploration and imitation. \n\n**3. Cago includes multiple trainable modules. How well does each module need to perform to make Cago work?**\n\nCago consists of five main modules: (1) a capability-aware goal sampler, (2) a BC explorer, (3) a goal-conditioned policy, (4) a world model, and (5) a goal predictor. Cago trains the world model and goal-conditioned policy in an end-to-end manner. The world model is updated using real trajectories collected by the goal-conditioned policy, whose goal is set by the capability-aware goal sampler (Algorithm 2, Line 18). Upon goal reaching, a pretrained BC explorer takes over for the remaining time steps in a rollout. Simultaneously, the goal-conditioned policy is trained using imagined rollouts generated by the world model (Line 19). These two components are tightly coupled and co-adapt during training. The goal predictor is trained separately on demonstrations to infer the final goal state from an initial observation. It is used only during evaluation and enables the goal-conditioned policy to generalize to test environments where the goal state is not explicitly provided.\n\nWe assess the importance of each component as follows. \n\n(1) To ablate the capability-aware goal sampler, we designed two additional goal sampling strategies: \n\n* Step-based curriculum: This baseline samples goals from demonstrations in proportion to training steps (i.e., current training step / predefined total training steps). However, it does not assess the agent's actual capabilities, and may therefore sample goals that are either too easy or too difficult for the agent at its current learning stage.\n* Final-goal sampling (Line 329): This strategy always selects the final observation from a demonstration in the goal phase of our Go-Explore sampling paradigm, ignoring the agent's current goal-reaching capability. The BC Explorer takes control from the goal-conditioned policy halfway through each rollout. The table below reports the results averaged over 5 random seeds:\n\n| Environment     | Cago (Ours) | Step-based Curriculum | Final-goal Sampling |\n|----------------|-------------|------------------------|---------------------|\n| Adroit-Pen     |    **0.82**         |           0.48             |          0.52           |\n| Stick-Push     |      **0.99**       |            0.92            |          0.68           |\n| Disassemble    |      **0.8**       |           0.04             |           0.02          |\n\nCago's goal sampling strategy significantly outperforms the alternatives, highlighting its critical role in the overall effectiveness of our approach.\n\n(2) As for BC-Explorer, we provide detailed ablation studies in Section 4 (Line 328). Specifically, Cago-NoExplorer removes the BC Explorer. We also evaluate Cago-RandomExplorer, which replaces the BC Explorer with a uniformly random policy during the exploration phase of our Go-Explore-style rollout strategy. The results are in Fig. 4. The table below reports the final success rates averaged over 5 random seeds (each evaluated over 100 test episodes):\n\n| Environment     | Cago-BC-Explorer (Ours) | Cago-NoExplorer | Cago-RandomExplorer |\n|----------------|-------------|------------------------|---------------------|\n| Adroit-Pen     |   **0.82**         |           0.60             |          0.69           |\n| Stick-Push     |      **0.99**       |            0.72            |          0.79           |\n| Disassemble    |      **0.8**       |           0.61             |           0.72          |\n\nThe results demonstrate that both ablations significantly reduce performance, highlighting the importance of BC Explorer for generating relevant exploratory rollouts to accelerate learning.\n\n(3) For the goal-conditioned policy, we observed that its critic loss decreases over time in correlation with the success rates reported in Fig. 3, underscoring its crucial role in overall performance.\n\n(4)  In early development, we attempted a model-free variant of Cago by training a goal-conditioned policy directly from environment rollouts starting at demonstration initial states using actor-critic methods. However, due to the limited number of demonstrations, this approach failed to generalize. We therefore adopted the model-based strategy that leverages a world model to generate diverse imagined rollouts (Line 173), enabling more effective and sample-efficient policy learning.\n\n(5) Empirically, we find that the Goal Predictor reliably infers final goals from held-out demonstrations, enabling effective zero-shot execution in unseen test tasks.\n\n**4. Given a small number of demonstrations with complex model structures, it is very likely that the model would overfit to the demonstrations.**\n\nCago mitigates overfitting to a small number of demonstrations by training its goal-conditioned policy $\\pi^G$ within a world model. Rather than training $\\pi^G$ to reach states solely from demonstrations, Cago instructs $\\pi^G$ to learn to reach diverse states sampled from real rollouts in $\\mathcal{D}_{\\text{cap}}$, which records the trajectories generated under our Go-Explore paradigm with capability-aware goal sampling, in its world model. This promotes generalization across the environment. To evaluate the generalization capability, we tested Cago on 500 unseen initial states generated from random seeds, each differing from those in the demonstrations. We report both the average L2 deviation from demonstration initial states and the success rate across these trials.\n\n| Environment     | Average initial observation deviation | Average success rate(500 seeds) |\n|----------------|-------------|------------------------|\n| Adroit-Pen     |      1.2749     |          0.824              |\n| Stick-Push     |     0.0547        |           0.982             |\n| Disassemble    |      0.1405       |            0.806           |\n\nThe results are consistent with Fig 3.\n\n**5. Can reaching a state with a trajectory that is very different from the demonstrations indicate that the agent has learned the ideal skills?**\n\nCago adopts a curriculum-style sampling strategy, where goals are selected progressively from easier to harder states along the demonstration trajectories. This naturally encourages the agent to incrementally reproduce and internalize the demonstrated behaviors, as illustrated in Fig.1 (Right) and Fig.4(d).\n\n**6. Cago does not consistently or significantly outperform other methods.**\n\nShelfPlace is the only environment where Cago is outperformed by GAIL. In all other environments, Cago either outperforms the baselines or achieves near 100% success. On ShelfPlace, we observe that GAIL learns a shortcut to achieve the goal, whereas Cago follows the demonstration more faithfully, which results in slower task completion."}, {"review_id": "YLIIXNrb2h", "reviewer": "Reviewer_x9ni", "summary": "The 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.", "questions": "* 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.", "limitations": "Limitations around environment resets and similarity metric could be further discussed.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The 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.", "was_revised": true, "final_justification": "The authors have addressed my concerns. I am increasing my score and I am curious to see the CSIL baseline results.", "rebuttal": "We appreciate your insightful feedback and constructive comments!\n\n**Q1. Provide a clearer framing, I would suggest committing to the IL narrative.**\n\nWe would like to clarify that Cago is fundamentally different from traditional IL methods, both in methodology and motivation. IL typically requires direct learning from demonstrations, either by recovering a reward function from demonstrations, or matching expert and agent state-action distributions (e.g. GAIL), or incorporating demonstrations as anchors to mitigate overestimation and instability caused by out-of-distribution actions (e.g. Cal-QL). In contrast, Cago presents a new paradigm to utilize demonstrations: it dynamically tracks the agent's competence along demonstrated trajectories and uses this signal to select intermediate states in the demonstrations that are just beyond the agent's current reach as goals to guide\nonline trajectories collecting. To evaluate this novel perspective, we compared Cago against state-of-the-art baselines that represent diverse strategies for learning from demonstrations, including distribution mathcing (GAIL and PWIL), curriculum learning (JSRL), model-based exploration (MoDem), and offline-to-online fine-tuning (CalQL and RLPD). These comparisons are essential to highlight the effectiveness of Cago as a new paradigm for learning from demonstrations. \n\n***We clarify that our baselines can access an absolute sparse reward where +1 is rewarded only on success states and 0 is given otherwise. This reward is, however, not accessible to Cago during training.***\n\n**Q2. Extend the related work to better cover imitation learning**\n\nIn the revision, we will expand the Related Work section to include a broader and more up-to-date survey of IL literature and clarify more explicitly how Cago differs from standard IL methods.\n\n**Q3. Provide more IL state-of-the-art baseline**\n\nWe appreciate the suggestion to include more imitation learning baselines for clarity and completeness. To address this, we conducted additional experiments comparing Cago with two suggested state-of-the-art imitation learning methods—SQIL and ValueDice—in addition to our existing IL baselines, GAIL and PWIL, across three environments: Stick-Push, Disassemble, and Adroit-Pen. The table below reports the final success rates averaged over 5 random seeds (each evaluated over 100 test episodes):\n\n| Task            | Cago (Ours) | SQIL [1] | ValueDice [2] | GAIL | PWIL |\n|-----------------|-------------|----------|----------------|------|------|\n| Stick-Push      |     **0.99**       |    0.78     |      0.48       |  0.02   |  0.94   |\n| Disassemble     |     **0.80**       |    0.34     |       0.18        |  0.05   |  0.02   |\n| Adroit-Pen      |     **0.82**       |   0.44   |       0.24        |  0.15   |  0.21   |\n\nCago consistently outperforms these baselines. We will update the paper to include the additional baselines suggested by the reviewer across all our benchmark tasks.\n\n**Q4. Include results on D4RL for which there are results on established baselines.**\n\nWe would like to clarify that our evaluation already includes tasks from the D4RL benchmark. Specifically, the Adroit environments used in our experiments—Door, Hammer, and Pen—are directly taken from D4RL. \n\n**Q5. Provide further discussion on initial state resets, how this impacts real world robot experiments, and what alternatives there are, if any?**\n\nIndeed, resetting to specific initial states is a limitation of Cago, but as the reviewer notes, it is already a substantial improvement over arbitrary state resets, which are often infeasible in real-world robotics. Importantly, Cago does not require exact resets to the full demonstration initial states. For instance, in the Peg Insertion task, the peg's initial position can differ from the demonstration, and Cago can learn to reach states within the early portion of the demonstration. However, some elements—like the hole and the box—are non-movable and randomly initialized by the environment. Since the agent cannot manipulate these objects, we have to reset them to match the demonstration configuration to ensure goal reachability. The constraint over other controllable components, such as the robot arm and peg, can be relaxed. This partial reset strategy reduces the burden and makes Cago more practical for real-world settings.\n\n**Q6. Further details on the implications of theorem 1 would be helpful. Also, it feels like assumption (2) would mean BC is good enough.**\n\nTheorem 1 is intended to formalize that the imagined rollouts generated by the world model are relevant and close to trajectories that could have been sampled in the real environment. This property is important because it ensures that policies trained over imagined rollouts can generalize to the real world. Assumption (2) states that the world model can be accurately learned via supervised learning from trajectories collected by the BC Explore policy. This is a reasonable assumption, and it does not require the BC Explore policy itself to be of high quality. \n\n**Q7. $D_{cap}$ is not used to train $\\pi_{G}$ in Algorithm 2 and it feels a bit wasteful.**\n\n$D_{cap}$ plays a crucial role in training $\\pi_{G}$. We learn $\\pi_{G}$ based on imagined rollouts from the learned world model $\\widehat{\\mathcal{M}}$. Each rollout begins at $s_0$, a state randomly sampled from a trajectory $\\tau$ in $\\mathcal{D}_{cap}$, and is rolled out for $H$ steps using the goal-conditioned policy $\\pi^G(a_t | s_t, g)$. The goal state $g$ is selected as a future state $s_H$ from the same trajectory $\\tau$. The objective is to train $\\pi^G$ to reinforce trajectories that efficiently reach $g$ in the imagined rollouts from $s_0$ under the learned dynamics $\\widehat{\\mathcal{M}}$.\n\n**Q8. Cago applied to Model-free RL.**\n\nIn early experiments, we attempted to train Cago in a model-free manner by learning the goal-conditioned policy directly from real environment rollouts. But we further find that instead of directly using $D_{cap}$ to train $\\pi_G$ in Model-free manner, the world model can offer a richer set of imagined rollouts to train the $\\pi^G$, which improves the learning efficiency. In our setting, the simulated trajectories have starting states and goal states randomly sampled from the same trajectory in $D_{cap}$. As a result, the simulated trajectories still resemble segments from $D_{cap}$​, while significantly increasing trajectory diversity and data richness. The MBRL—leveraging simulated trajectories allows us to generate more diverse and abundant training data. This promotes generalization across the environment.\n\n**Q9. The use of the similarity metric to count visitation could be a limitation for more complex tasks.**\n\nWe have evaluated Cago in visual environments (Appendix F.2) and found it performs comparably to the state-based setting. While that experiment used MSE to compute differences between image-based observations, more advanced metrics—such as SSIM (Structural Similarity Index) or FSIM(Feature Similarity Index)—can be incorporated. Importantly, the similarity metric used in our paper serves as a simple instantiation, rather than a core component of the paradigm. Cago can be configured with alternative metrics, such as similarity between state embeddings in a world model's latent space, to reduce input dimensionality, or use other signals like future state predication uncertainty from an ensemble of world models to assess whether the surrounding space has been sufficiently explored."}, {"review_id": "chJ5yodlKT", "reviewer": "Reviewer_1nz5", "summary": "This 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.", "questions": "1. 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.", "limitations": "The 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.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "### 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.", "was_revised": true, "final_justification": "The additional experiments and clarifications have improved the submission in several aspects, including new results under missing and noisy demonstrations, analyses of hyperparameters, and inclusion of PEG as a GCRL baseline. However, some core limitations remain only partially addressed—most notably, the reliance on successful demonstrations and the limited scalability to high-dimensional or semantic goal spaces.", "rebuttal": "We appreciate your insightful feedback and constructive comments!\n\n**1. This methodology critically relies on \"high-quality demonstration trajectories that adequately cover the task space.\" The methodology lacks experimental analysis or solutions to noisy, incomplete, or suboptimal expert data.**\n\nWe clarify that Cago only assumes access to \"high-quality demonstration trajectories that adequately cover the task space\" (Line 246) in order to derive the theoretical guarantee stated in Theorem 1. In practice, the demonstrations used in the Adroit environments (door, hammer, and pen) are in fact noisy and suboptimal, and yet Cago remains effective. This highlights the robustness of our method to imperfect demonstrations.\n\nTo further address this concern, we have conducted additional experiments specifically designed to evaluate how Cago performs when provided with noisy, incomplete demonstrations. We collected three different types of suboptimal demonstrations:\n\n* (1) Missing Observations: We randomly removed 20% of the observations from each original demonstration trajectory to simulate incomplete data.\n\n* (2) Noisy Expert Actions: We added Gaussian noise (scaled by 0.1) to each action output by the expert policy during expert trajectory collection.\n\n* (3) Random Actions: We replaced the expert’s action with a randomly sampled action with a probability of 20% at each timestep during expert trajectory collection.\n\nThese settings simulate three kinds of suboptimality: missing state information, noisy expert behavior, and corrupted decision-making. Using these demonstrations, we re-evaluated Cago on three benchmark environments: AdroitPen, Stick-Push, and Disassemble. Because we do not have the expert policy for Adroit, we only conducted the \"Missing Obs\" experiment for Adroit-Pen. The table below reports the final success rates after 1M training steps, averaged over 5 random seeds (each evaluated over 100 test episodes):\n\n| Environment     | Original Demo | Missing Obs | Noisy Actions | Random Actions |\n|----------------|-------------------|--------------------|------------------------|----------------------|\n| Adroit-Pen     | 0.82                |  **0.88**               |  -                    |   -                 |\n| Stick-Push     | **0.99**                 | **0.99**                  | 0.97                      | 0.96                    |\n| Disassemble    | 0.80                 | 0.77                  | 0.88                      | **0.96**                    |\n\n The results suggests that Cago is robust to the suboptimality in demonstrations. This is because:\n\n* First, Cago only uses demonstration trajectories to extract sequences of observations that serve as curriculum goals. It does not attempt to imitate the demonstration directly, nor does it rely on demonstrations to infer a reward function. This significantly reduces Cago's dependency on high-quality demonstrations. \n\n* Second, even when the demonstration trajectories are suboptimal—as long as they represent a successful sequence of observations—Cago can still learn well. We observe that in the Disassemble environment, Cago achieves significantly higher success rates when using the \"Random Actions\" type of suboptimal demonstrations compared to using the original demonstrations. Those imperfect demonstrations—with pauses or detours—cover a broader region of the task space, thereby improving the final policy's generalization. More importantly, Cago is guided by a reward function shaped using a Temporal Distance Network (see Main Paper Appendix B.2), which estimates how many steps are needed to reach a goal state. This enables Cago to discover more efficient, often shorter, alternative paths for goal reaching. Therefore, even when the demonstration are suboptimal, Cago may still learn the optimal paths. We will revise the paper to highlight Cago's robustness to suboptimal demonstrations by including these results.\n\n**2. Here is a clear lack of discussion or evidence regarding its scalability to situations where goals are difficult to define directly from demonstration trajectories.**\n\nOur work specifically targets robotic manipulation tasks, where it is both common and effective to define the goal space directly in terms of the observation space. This approach has been widely adopted in prior works on goal-conditioned reinforcement learning [1,2,3]. In Appendix F.2, we demonstrate that Cago performs comparably in high-dimensional visual environments, even when using a simple similarity metric such as mean squared error (MSE) between image-based observations. Since Cago guides agents toward low-level demonstration states in a progressive manner, it does not rely on or infer high-level abstract goals—doing so is outside the intended scope of our method.\n\n[1] Planning goals for exploration, ICLR 2023\n\n[2] Learning World Models for Unconstrained Goal Navigation, NeurIPS 2024\n\n[3] Discovering and Achieving Goals via World Models, NeurIPS 2021\n\n**3. More sensitivity analysis on how lambda_visit and delta affect performance.**\n\nRegarding $\\lambda_{visit}$: This parameter controls the threshold for how frequently a state must be visited before it is considered “mastered” by the agent, and thus eligible for sampling more difficult goals. If it is set too small, Cago may prematurely progress to harder goals before acquiring sufficient competence in easier ones. Conversely, if it is set too large, Cago may become overly conservative, spending excessive time on already-mastered states. We have conducted a sensitivity analysis by setting $\\lambda_{visit} \\in {50, 100, 200}$, and evaluated Cago on three environments: Stick-Push, Disassemble, and AdroitPen. The final success rates after 1M training steps, averaged over 5 random seeds are shown below.\n\n| Environment     | $\\lambda_{visit}$ = 50 | $\\lambda_{visit}$ = 100 | $\\lambda_{visit}$ = 200 |\n|----------------|--------------|---------------|---------------|\n| Stick-Push     | 0.94         | 0.96          | **0.99**          |\n| Disassemble    | 0.79      | **0.82**          | 0.80          |\n| Adroit-Pen     | 0.77         | 0.82          | **0.83**          |\n\nResults show that while very small $\\lambda_{visit}$ may slightly influence learning, Cago consistently reaches similar final performance, indicating robustness to this parameter across a broad range.\n\nRegarding $\\delta$: This parameter defines the goal sampling window size, i.e., the proportion of expert states from which new goals are sampled during training. A large $\\delta$ value may lead to sampling overly distant or inappropriate goals, affecting stability and learning efficiency. In contrast, a very small $\\delta$ may overly restrict exploration and hinder curriculum progression. We added a sensitivity experiment comparing the default $\\delta$ = 10% with a smaller value $\\delta$ = 5% and a larger value $\\delta$ = 20%, across the same three environments. The table below shows that while larger values can sometimes slightly reduce performance, Cago remains capable of learning effective policies in all settings. We will include these results to the paper to provide a more thorough discussion of both parameters. \n\n| Environment     | $\\delta$ = 5% | $\\delta$ = 10% | $\\delta$ = 20% |\n|----------------|---------|---------|---------|\n| Stick-Push     | **1.00**    | 0.99    | 0.99    |\n| Disassemble    | **0.81**    | 0.80    | 0.75    |\n| Adroit-Pen     | 0.76    | **0.82**    | 0.81    |\n\n**4. Although Cago is compared to demonstration-based baselines, it overlooks comparisons with broader goal-conditioned reinforcement learning methods beyond demonstrations.**\n\nCago introduces a novel paradigm for leveraging demonstrations, with goal-conditioned RL (GCRL) serving only as the underlying framework. Our focus is on demonstration utilization, so we adapt a broad set of strong baselines that differ in how they use demonstrations. To further address the reviewer’s concern, we include PEG [1], a recent competitive model-based GCRL method, alongside the existing GCRL-Dreamer baseline. As shown in the table below, Cago consistently outperforms these GCRL baselines.\n\nTable4: Results of Cago and GCRL baselines.\n| Environment     | Cago(Ours)| Dreamer | PEG |\n|----------------|--------------|---------------|---------------|\n| Stick-Push     |    **0.99**     |     0.68     |    0.76      |\n| Disassemble    |    **0.80**     |     0.02     |     0.04      |\n| Adroit-Pen     |    **0.82**     |      0.52    |      0.54    |\n\n[1] Hu. et al., Planning goals for exploration, ICLR 2023\n\n**5. How Cago's learned policy behaves during the exploration phase, and how Cago succeeds in specific scenarios where baselines fail.**\n\nWe visualized the exploration phase of Cago for the StickPush environment in Figure 4(d). Each red dot represents the normalized position of a sampled goal within a demonstration trajectory, with 0 indicating the start and 1.0 indicating the final demonstration state. Early in training, the agent predominantly samples goals at lower normalized positions, focusing on subgoals near the beginning of the trajectory that are within its current capabilities. As training advances, goal sampling gradually shifts toward higher normalized positions, indicating the agent’s increasing ability to pursue more challenging goals closer to task completion. By continuously targeting goals just at the boundary of the agent’s current capability, Cago facilitates efficient learning in this long-horizon task.  In challenging tasks like Peg-insertion, Cago's curriculum-like progression—guided by internal capability estimation—significantly improves learning efficiency. In contrast, we found the baseline policies, while attempting to imitate the demonstrations, move the manipulator to the final pose even the it does not yet grasp the peg. Similar exploration trends are observed across other benchmarks, and we will include these figures in the revised paper."}, {"review_id": "Z2N9sWT7Se", "reviewer": "Reviewer_NtvZ", "summary": "The 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.", "questions": "In 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.", "limitations": "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.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "After carefully considering the authors' rebuttal and the discussion among reviewers, I have updated my assessment as follows:\n\n**Resolved Issues:**  \n- The authors provided comprehensive ablation studies and analyses clarifying the role and robustness of each module in their framework, which addressed my concerns about the method's complexity and the importance of individual components.\n- They gave detailed responses on the method’s applicability to broader, real-world scenarios, including diverse initial states and stochastic dynamics.\n- The rationale for not using behavioral cloning (BC) pre-training was explained clearly, and the approach was justified in the context of sparse demonstrations.\n- Additional experiments and ablations on the number of demonstrations demonstrated that the method is effective even with limited expert data, which is important for practical deployment.\n\n**Remaining Issues:**  \n- While the authors addressed most technical questions, some open questions remain regarding the method's scalability to highly abstract or non-robotic tasks and the potential overhead in real-world deployment.\n- The paper could still benefit from deeper qualitative analysis and a more explicit discussion of limitations in the main text.\n\n**Overall Assessment:**  \nThe rebuttal has largely addressed my main concerns, especially regarding robustness, practical applicability, and experimental validation. The remaining issues are relatively minor and do not outweigh the strengths demonstrated. Therefore, I have raised my score to reflect the improved confidence in the paper's contribution.", "rebuttal": "We appreciate your insightful feedback and constructive comments!\n\n**W1. The approach involves several components. It would be helpful to understand how robust the overall method is to each individual component.**\n\nCago consists of five main modules: (1) a capability-aware goal sampler, (2) a BC explorer, (3) a goal-conditioned policy, (4) a world model, and (5) a goal predictor. At a high level, Cago trains the world model and goal-conditioned policy in an end-to-end manner. The world model is updated using real trajectories collected by the goal-conditioned policy, whose goal is set by the capability-aware goal sampler (Algorithm 2, Line 18). Upon goal reaching, a pretrained BC explorer takes over for the remaining time steps in a rollout. Simultaneously, the goal-conditioned policy is trained using imagined rollouts generated by the world model (Line 19). These two components are tightly coupled and co-adapt during training, making their joint performance critical to overall success. The goal predictor is trained separately on demonstrations to infer the final goal state from an initial observation. It is used only during evaluation and enables the goal-conditioned policy to generalize to test environments where the goal state is not explicitly provided.\n\nWe assess the importance of each component as follows. \n\n(1) To ablate the capability-aware goal sampler, we designed two additional goal sampling strategies: \n\n* Step-based curriculum: This baseline samples goals from demonstrations in proportion to training steps (i.e., current training step / predefined total training steps). However, it does not assess the agent's actual capabilities, and may therefore sample goals that are either too easy or too difficult for the agent at its current learning stage.\n* Final-goal sampling (Line 329): This strategy always selects the final observation from a demonstration in the goal phase of our Go-Explore sampling paradigm, ignoring the agent's current goal-reaching capability. The BC Explorer takes control from the goal-conditioned policy halfway through each rollout. The table below reports the final success rates averaged over 5 random seeds (each evaluated over 100 test episodes):\n\n| Environment     | Cago (Ours) | Step-based Curriculum | Final-goal Sampling |\n|----------------|-------------|------------------------|---------------------|\n| Adroit-Pen     |    **0.82**         |           0.48             |          0.52           |\n| Stick-Push     |      **0.99**       |            0.92            |          0.68           |\n| Disassemble    |      **0.80**       |           0.04             |           0.02          |\n\nCago's goal sampling strategy significantly outperforms the alternatives, highlighting its critical role in the overall effectiveness of our approach.\n\n(2) As for BC-Explorer, we provide detailed ablation studies in Section 4 (Line 328). Specifically, Cago-NoExplorer removes the BC Explorer. It does not further explore after reaching a sampled goal. We also evaluate Cago-RandomExplorer, which replaces the BC Explorer with a uniformly random policy during the exploration phase of our Go-Explore-style rollout strategy. The results are in Fig. 4. The table below reports the final success rates averaged over 5 random seeds:\n\n| Environment     | Cago-BC-Explorer (Ours) | Cago-NoExplorer | Cago-RandomExplorer |\n|----------------|-------------|------------------------|---------------------|\n| Adroit-Pen     |    **0.82**         |           0.60             |          0.69           |\n| Stick-Push     |      **0.99**       |            0.72            |          0.79           |\n| Disassemble    |      **0.80**       |           0.61             |           0.72          |\n\nThe results highlight the importance of BC Explorer for generating relevant exploratory rollouts to accelerate learning.\n\n(3) For the goal-conditioned policy, we observed that its critic loss decreases over time in correlation with the success rates reported in Fig. 3, underscoring its crucial role in overall performance.\n\n(4) In early development, we attempted a model-free variant of Cago by training a goal-conditioned policy directly from environment rollouts starting at demonstration initial states using actor-critic methods. However, due to the limited number of demonstrations, this approach failed to generalize. We therefore adopted the model-based strategy that leverages a world model to generate diverse imagined rollouts (Line 173), enabling more effective and sample-efficient policy learning.\n\n(5) Empirically, we find that the Goal Predictor reliably infers final goals from held-out demonstrations, enabling effective zero-shot execution in unseen test tasks.\n\n**W2. 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.**\n\nCago can be extended to support diverse initial states and uncertain dynamics.\n* For diverse initial states in real-world settings, Cago can gradually prioritize those where the agent has the most learning potential—starting from demonstration initial states and expanding to nearby ones where the agent can still benefit from the demonstration trajectories for dynamic capability tracking and progressive goal selection. As training progresses, we incorporate successful rollouts from these nearby initial states into the demonstration set, broadening its coverage. This enables a natural curriculum that eventually transitions to diverse, real-world initial states.\n* For uncertain, stochastic dynamics, Cago remains robust by not requiring exact reproduction of actions in demonstration trajectories. Instead, it adapts goals based on observed agent successes, making the curriculum self-correcting in the face of noise or variability. If the agent consistently fails to reach a specific demonstration state, Cago's capability metric naturally shifts goal sampling toward more achievable targets. Additionally, goal selection can be extended with probabilistic estimates—such as goal-reaching likelihood or future state predication uncertainty from an ensemble of world models—which we view as a promising direction for high-stochasticity domains.\n\n**Q1. Why not initialize the goal-conditioned policy with behavioral cloning (BC) pre-training?**\n\nWe considered the possibility of initializing $\\pi^G$ with behavioral cloning (BC) on demonstration data, and conducted preliminary experiments with BC pretraining in Cago’s early development stage. However, our findings led us to deliberately decide against using BC initialization in our final framework, for the following reasons:\n\n * (1) Risk of overfitting under sparse demonstrations: Cago is specifically designed for the low-data regime, where only a small number of demonstrations (typically 5 or 10) are available. In such cases, BC pretraining tends to overfit to the limited demonstration data, resulting in poor generalization when the agent starts to explore beyond the demonstrated trajectories. \n\n* (2) Misalignment with the training framework of Cago: Cago’s goal-conditioned policy $\\pi^G$ is trained jointly with a Temporal Distance Network (TDN), which estimates how many steps are needed to reach a goal. The TDN serves as a dense feedback signal and is itself trained on imagined trajectories from the world model. Crucially, all three components—the policy, world model, and TDN—are updated jointly in an end-to-end manner. In the early phase of training, the TDN’s estimates are often inaccurate. Therefore, even a well-behaved BC-initialized policy would receive noisy or misleading training signals, causing it to quickly deviate from the pretraining distribution. This undermines the value of BC initialization and can sometimes even slow down learning. These observations informed our design decision to exclude BC pretraining. \n\n**Q2. How does the method’s performance scale with the number of demonstrations? ... It would be helpful to justify the choice of using only 10–20 demonstrations.**\n\nOur method is explicitly designed for the low-demonstration regime, which is typical in real-world robotic settings where expert trajectories are expensive, time-consuming, or risky to collect. We intentionally evaluate Cago with only 10–20 demonstrations to highlight its practical relevance, emphasizing sample efficiency and generalization from limited supervision, enabled by the use of a world model to generate imagined rollouts that enrich training. A more accurate model around the demonstration distribution can benefit policy generalization. We conducted ablation study in the Supplemental Material (Figure 1)  using 5, 10, and 20 demonstrations. We further conducted an experiment using 50 demonstrations to evaluate how Cago performs under a large set of expert data. The table below reports the final success rates averaged over 5 random seeds (each evaluated over 100 test episodes):\n\n| Environment | Cago-5-demons | Cago-10-demons | Cago-20-demons | Cago-50-demons |\n|----------------|-------------|------------------------|---------------------|---------------------|\n| Adroit-Pen     |    0.71         |           0.82             |          0.92           |          **0.93**           |\n| Stick-Push     |      0.81       |            0.99            |          1.00           |          **1.00**           |\n| Disassemble    |      0.69       |           0.80             |           0.81          |          **0.98**           |\n\nWe found that with 50 demonstrations, Cago performs exceptionally well, exhibiting significantly better generalization and achieving higher final success rates compared to training with fewer demonstrations. However, we also found that larger demonstration sets may delay learning onset due to longer warm-up times to fit world models. Studying this trade-off is an interesting direction."}]

Author rebuttal:

Meta-review: This paper presents an empirical method for inverse reinforcement learning in goal-condition RL, whch combines behavioral cloning with a generative moel, and evaluates it in sparse reward environments. The setting is not pure IRL, as the agent uses the generative model to obtain more data using a policy of its choice. 

+ The performance is generally strong
- The method is somewhat adhoc, but it makes intuitive sense to some extent.
- Since this is not pure IRL, it would make sense to compare with other methods combining IRL with RL

All reviewers are for acceptance.
Decision: Accept (poster)