Multi-Stage Manipulation with Demonstration-Augmented Reward, Policy, and World Model Learning
Authors: Adrià López Escoriza
Nicklas Hansen
Stone Tao
Tongzhou Mu
Hao Su
Venue: ICML 2025 poster (2025)
Keywords: Reinforcement Learning
Learning from Demonstrations
Robotics
Manipulation
Primary area: reinforcement_learning->deep_rl
TL;DR: We propose a framework multi-stage manipulation tasks with sparse rewards and visual inputs; our framework combines learned dense rewards, model-based RL, a bi-phasic training scheme, and a small number of demonstrations.

Abstract: Long-horizon tasks in robotic manipulation present significant challenges in reinforcement learning (RL) due to the difficulty of designing dense reward functions and effectively exploring the expansive state-action space. However, despite a lack of dense rewards, these tasks often have a multi-stage structure, which can be leveraged to decompose the overall objective into manageable sub-goals. In this work, we propose DEMO³, a framework that exploits this structure for efficient learning from visual inputs. Specifically, our approach incorporates multi-stage dense reward learning, a bi-phasic training scheme, and world model learning into a carefully designed demonstration-augmented RL framework that strongly mitigates the challenge of exploration in long-horizon tasks. Our evaluations demonstrate that our method improves data-efficiency by an average of 40% and by 70% on particularly difficult tasks
compared to state-of-the-art approaches. We validate this across 16 sparse-reward tasks spanning four domains, including challenging humanoid visual control tasks using as few as five demonstrations.

Multi-Stage Manipulation with Demonstration-AAumeente Reward, Policy, and World Model Learning
Adrià López Escoriza 12
Nicklas Hansen
Stone Tao 13 Tongzhou Mu Hao Su 13
Abstract
Zhang et al., 2018; Kidambi al., 2020; Hafner et al., 2020; Yu et al., 2020; Hansen et al., 2022; 2024; Sferrazza et al., 2024). However, designing accurate reward functions is challenging. Poorly designed rewards can lead agents to become trapped in local min- ima or exploit unintended shortcuts, resulting in undesirable behaviors (Clark & Amodei, 2016). More critically, scal- ing reward design to complex tasks is highly impractical: the larger the state space and the longer the horizon, the more intricate the reward must be. While recent approaches leveraging Large Language Models (Kwon et al., 2023; Ma et al., 2023; Xie et al., 2024) and Vision-Language Models (Rocamonde al., 2024; Baumli al., 2024) for reward generation show promise, they still struggle with high-precision requirements, particularly in manipulation problems. In contrast, sparse rewards, such as binary signals indicating task or subtask completion, are much easier to obtain. However, traditional RL methods still struggle to learn effectively from sparse rewards.
Schmidhuber, 2018;
Long-horizon tasks in robotic manipulation present significant challenges in reinforcement learning (RL) due to the difficulty of designing dense reward functions and effectively explor- ing the expansive state-action space. However, despite a lack of dense rewards, these tasks of- ten have a multi-stage structure, which can be leveraged to decompose the overall objective into manageable subgoals. In this work, we propose Demonstration-AAugeens
Reward, Policy, and World Model Learning (DEMO"), framework that exploits this structure for efficient learning from visual inputs. Specifically, our approach incorporates multi-stage dense reward learning. a bi-phasic training scheme, and world model learning into a carefully designed demonstration- augmented RL framework that strongly mitigates the challenge of exploration in long-horizon horizon tasks. Our evaluations demonstrate that our method im- proves data-efficiency by an average of 40% and by 70% on particularly difficult tasks compared to state-of-the-arr approaches. We validate this across 16 rewrre tasks spanning four do- mains, including challenging humanoid visual control tasks using as few as five demonstrations. Website with code and visualizations can be found at https://adrriloopzzssoor wwtprriwwwagttbb..//ommmoo.
Fortunately, long-horizon tasks do offer opportunities to simplify the problem. Typically, such tasks exhibit natural multi-stage structure. For example, pick- pick-and-pp aand-lacce task can be broken down into subtasks such as grasping, lifting, and placing. Each of these can be associated with stage indicators or rewards that can easily be queried from the environment. This multi-stage structure allows these tasks to be decomposed into more manageable subgoals, enabling the agent to collect rewards more frequently (Smith et al., 2020; Di Palo & Johns, 2021). However, subgoal sparse rewards can still be insufficient if the distance between sub- goals is too great, leading back to the exploration problem.
1. Introduction
Reinforcement learning (RL) with dense rewards has en- abled significant progress in high-dimensional control tasks. Many such tasks are now solvable when the reward func- tion is carefully designed for the specific goal. In partic- ular, model-based RL (MBRL) has demonstrated strong performance in these high-dimensional problems (Ha &
Prior work shows that Learning from Demonstrations (Lff) can help mitigate exploration issues in sparse reward set- tings. Algorithms such as CoDER (Zhan al., 2022) and MoDem (Hansen et al., 2023; Lancaster al., 2024) lever- age demonstrations to populate the replay buffer of an off- policy RL algorithm (Sutton & Barto, 2018), but they often scale poorly with task complexity as they need demonstra- tions that sufficiently cover the behavior space. Inverse RL methods (Trott et al., 2019; Wu et al., 2021; Memarian et al., 2021; Escontrela et al., 2022) train RL on a reward function that is learned from demonstrations, but inverse RL alone
Universitty of California San Diego "ETH Zürich, Switzer- land 'Hillbot. Correspondence to: Adrià López Escoriza <alopez@ethz.ch>>
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Demonstration-AAggeeiis Reward, Policy, and World Model Learning
ManiSkill Manipulation
Meta-World
ManiSkill Humanoids
Robosuite
100%
90.5
76.4
78.0
71.4
50%
41.7
33.8
31.0
24.0
25.3
16.0
18.0
16.7
0.0
4.4
7.7
0.0
0%
Ours
TD-MPC2
MoDem
LaNE
BC
Figure 1. Summary of results. Final success rate (%) achieved by method and set of strong baselines, averaged across all tasks within each domains. Average 5 seeds. Given handful demonstrations, our method achieves high success rates in challenging visual manipulation tasks with sparse rewards, far exceeding previous state-of-the methods. See Appendix A for per-task results,
often struggles as the learned reward function can have poor predictions on unseen states. Lastly, these methods typically require vast amount of samples to learn a reward func- tion (Kumar et al., 2022; Mu et al., 2024) before any policy learning can begin.
In this paper, build on demoostrationnaaggeeeet RL to tackle multi-stage manipulation tasks with sparse stage- wise rewards. We introduce DEMO model-based RL algorithm that leverages a limited number of demonstrations for three key purposes: learning policy, world model, and a dense reward. DEMO exploits the multi-stage structure of long-horizon tasks to transform sparse stage indicators into stage-wise dense reward. This enables dense feedback in structured way, prioritizing achieving subgoals over following demonstrations. Unlike prior work, our dense reward is learned online, alongside policy and world model learning.
Figure 2. Task domains. We evaluate methods on 16 multi-stage image-based sparse-reward tasks spanning four domains: Meta- World (Yu et al., 2021), Robosuite (Zhu et al., 2022), as well manipulation and humanoid tasks from ManiSkilI3 (Tao et al., 2024b). See Appendix D for complete overview of tasks.
2. Preliminaries
Problem formulation. We aim to learn control policies for multi-stage, long-horizon tasks, which model as infinite- horizon Markov Decision Processes (Bellman, 1957) de- fined by the tuple (S, A, P. R, 7). Here, S and A denote the state and action spaces, respectively. P is the (unknown) state transition probability function, is a sparse reward function, and ς E [0, 1) the discount factor. Our goal is learn policy me S A parameterized by 0 that maximizes the expected cumulative reward (return) over infinite time horizon, formalized as max Exo IE ?'r:].
evaluate our method on a range of challenging manipu- lation tasks from Meta--orlk (Yu al., 2021), Robosuite (Zhu et al., 2022), well as both humanoid and tabletop manipulation tasks from ManiSkill3 (Tao et al., 2024b). Our results (see Figure 1) demonstrate that our method outper- forms state-of- oifte-of-hh-arr the-hht-t methods by an average of 40%, and for more complex, longer-horizon tasks, this performance gap increases to 70%. Our main contributions can summarized as follows:
Multi-stage sparse rewards. In this work, we focus on multi-stage tasks where the overall objective can be de- composed into a sequence N subgoals or stages. Stage indicators are often easy to obtain; for example, in manipu- lation tasks, it is straightforward to query whether the agent has grasped object. We model the (sparse) reward stage indicator function S (1,2...., N) that maps each state to its corresponding stage. We assume no access to any privileged information from the environment (e.g. object configurations), and instead consider multi-modal observations o (x, q) where denotes raw RGB images coming from the agent's cameras, and q denotes the propri- oceptive state of the robot. This constitutes realistic sensory inputs available in typical robotic platforms.
We introduce DEMO*, an MBRL algorithm for highly data-efficient robotic manipulation from visual inputs and sparse rewards. Our method integrates online dense reward learning into RL for multi-stage tasks.
2. We conduct extensive experiments in 16 tasks across 4 domains to demonstrate the data-efficiency and robust- ness of our approach compared to existing methods.
3. We analyze the relative importance of each component our framework, and are open-soureing all code and demonstrations used in this work.
TD-MPC2 (Hansen et al., 2022: 2024) is model-based RL algorithm that combines Model Predictive Control (MPC), a learned latent-space world model, and a terminal value func- tion learned via temporal difference (TD) learning. Specif-
2

---

Demonstration-Auumeenee Reward, Policy, and World Model Learning
Phase Pretraining
Phase Interaclive Learning
Ft
9
St,
3g
TBC
Environment
Planning
80%
enc
Disc.
Replay buffer
O
Latent state
Learned dense reward
Demo datasel
Sparse reward
Figure 3. Method overview. We present two-phase framework for multi-stage visual manipulation from sparse rewards that leverages a handful of demonstrations for dense reward learning and MBRL. Phase (left): policy and encoder is pre-trained on the available demonstrations using behavioral cloning, which serve initialization for the phase. Phase 2 (right): the agent iteratively collects environment data via planning and uses all available data update its world model as well a latent state discriminator; this discriminator used transform sparse environment rewards a learned dense reward for world model learning subsequent planning.
ically, TD-MPC2 learns a representation ho( that maps digeedimensional observation into compact rep- resentation 2, as well
world model, and a dense reward function simultaneously.
As our main algorithmic contribution, we introduce stage- specific reward learning. In particular, we extend the strat- egy on reward learning from demonstrations (LID) pre- sented in Mu et al. (2024) to online reward learning within world model. By learning structured, multi-stage rewards online alongside world model and policy, our method pro- vides more frequent and meaningful training signals the agent than prior work on demonstration augmeeted RL.
dynamics model in this latent space z' = de (z, a). In addition, TD-MPC2 also learns prediction heads, Rg. Qe. Tg. for (i) instantaneous reward r r= r=Rg(z [z (, a), (ii) state-action value (z, a). and (iii) a pol- icy prior To (z). The policy prior To serves to "guide" planning towards high-return trajectories and is optimized to maximize temporally weighted Q-values. The remaining components are jointly optimized to minimize TD)errors, and reward and latent state prediction enror,, minimizing
Our approach builds directly upon the strengths of prior work. In particular, we leverage MoDem's multi-phase accelcrated learning framework and use TD-MPC2 as our backhone for its robustness and generalizability.
t+ H
(TD-MPC(0)- = ())
(x)) (1)
where o',, z/ are the (latent) states at time t During envi- ronment interaction, TD-MPC2 selects actions via sample- based planner MPPI (Williams et al., 2015) and the learned world model. We adopt TD-MPC2 as choice of visual MBRL algorithm due to its simplicity and strong empiri- cal performance but emphasize that our framework can be instantiated with any MBRI. algorithm.
3.1. Model-based RL with online reward learning
Sparse rewards are a major challenge in RL, particularly for long-horizon tasks comprising multiple stages. To over- come this, we learn densify sparse rewards with a small number of demonstrations. For this, we introduce series of discriminators {6x stage l {O...NT N). The objective of each discriminator is to predict the likelihood of progressing to the next stage based on the latent state representation Z, produced by the back-bone world model.
1. Ji= each corresponding to task
3. Method
In this work, we address the challenge of solving multi-stage manipulation tasks from sparse rewards. Such long- horizon tasks are particularly difficult due to the combinatorial com- plexity of the state-action space and the lack of informative feedback across extended horizons. To overcome these is- sues, we propose DEMO3, a novel RL method that uses demonstrations for a three-fold purpose: to learn a policy, a
Therefore, each discriminator & acts stage classifier trained to distinguish states as either leading or not leading to successful stage transitions. For each stage k. we use typical Binary Cross Entropy (BCE) loss:
Lon =
E (o.r)=
[BCE(1,,24 og (h(o))))) (2)
3

---

Demonstration-AAugeente Reward, Policy, and World Model Learning
Algorithm DEMO' (Phase 2)
Zt
6x
Stanh
Require:
Demonstration dataset DD number of stages N
Discriminator
Learned dense reward
1: Initialize discriminators {6x )=
Initialize replay buffer B (0)
(t)
Rollout
Damos Replay
3: for each environment step do
4: if rand() 2 then
5: Agent step: a ~ mcc
6: else
Figure 4. Dense reward learning. At each update step, the contin- uous output of a stage discriminator is added to the environment sparse reward. The discriminator output is normalized to the [-3, 8] interval with a tanh operator.
7: Agent step: a, t (Op
8: end if
9: Env step: (7)+ + Env(or, ae)
10: Save sample: f- 7 U (Or, a 1,-1,
11: if episode done then
12: Reset environment
where h denotes the world model encoder, B is the replay buffer, and 9; represents the maximum stage that will be reached by the trajectory after the given sample:
13: Compute maximum stage labels: {sr]ö
14: Save trajectory: B + U (T U {si]8)
15: end if
16: f- min(1, on
8/-
max TV
F 2
17: end for
Update
We refer to 8 as the maximum stage label of a sub-trajectory. Thus, each trajectory, T; = {(or, ag. Tt, 0t+1, 8t) s is annotated with maximum stage labels 8, that serve as suc- cess labels for the stage discriminators. Then, presented in Algorithm I, at each update of the world model, the dis- criminators are updated as an additional part of the model. Specifically, for a given sample from the replay buffer (Ot, at, Oc+1, Tr, 8r) ~ B, the sparse reward associated with given stage, k- rr. will tell us which discriminator, 8,. will be updated by that sample. If the maximum stage la- bel st is greater than the current stage reward k rr, the sample belonging to trajectory with successful stage tran- sition will be treated as a positive example in the classifier loss. Note that in the event that no samples with a stage reward, r2 = k, would appear in a given batch, the discrimi- nator 6 for that stage would simply not get updated at that step. Therefore, while the algorithm is capable of working without any demonstrations, using small demonstration dataset can significantly accelerate the training of the world model and discriminators.
18: for each update step do
19:
Sample: (or. 0t+1, at, rt, 81) 81) 1otH }
(B D)
20:
Predict dense reward (ft)t)
21:
Compute world model losses: LR. Lq, Lis Lr.
22:
Compute discriminator loss: Ls = N Σx CJ
23:
Gradient step: <- 0 + PVLr
24:
end for
that do not, encouraging the agent to explore regions with higher probability of transition.
The total world model loss integrating these signals becomes
Lp LR + Lq + Ln + Ls:
(5)
where LR. LQ. and En represent the TD-MPC2 world model losses: (i) reconstruction, (ii) value estimation, and (iii) latent dynamics losses. Importantly, LR is computed to predict the learned dense reward produced by the discrimi- nator, i8. thus providing a richer learning signal than pure sparse rewards. Finally, L3 the average loss of all the stage discriminators. As in Hansen et al. (2024), the total loss is used to compute gradients for the world model, meaning that Ls is also used to learn the observation encoder.
While training the world model, the discriminators are used to generate dense rewards as per
p = + B (an((6,..2e))
(4)
where the output of the discriminator is mapped the [-3, 3) interval. The process is illustrated by Figure 4. We set 3 to be hyperparameter with < 1/3 to ensure that rewards never cross between different stage regions. This is to ensure that states belonging to more advanced stage always get a higher reward than lower-stage states. Effectively, our method rewards states that have a higher chance of transitioning to the next stage and penalizes those
3.2. Training scheme
In order to further boost the data-efficiency of DEMO3, we build upon previous work on accelerating MBRL with demonstrations. Specifically, we draw inspiration from Mo- Dem (Hansen et al., 2023) and propose a bi-phase training scheme in which we first use demonstrations to pre-train an initial policy through behavioral cloning (Atkeson & Schaal,
4

---

Demonstration-Augmenne Reward, Policy, and World Model Learning
ManiSkill Manipulation
Meta-World
ManiSkill Humanoids
Robosuite
100%
50%
0%
0
250K
500K 0
250K
500K
50K
100K 0
50K
100K
Interaction Steps
Ours
TD-MPC2
MoDem
LaNE
BC
Figure 5. Learning curves. Success rate as function of interaction steps for each of the four domains that we consider, averaged across all tasks and 5 random seeds. The shaded area corresponds to a 95% confidence interval. Our method consistently outperforms baselines.
1997; Pomerleau, 1988), TBC to collect informative samples during early stages of training. In phase 2, we gradually phase out the (frozen) pre-trained policy and start collecting samples by planning through the world model, which is learned via interactive learning. An overall diagram of our training strategy can be found in Figure 3.
Table 1. Experimental setup. We consider 16 challenging visual manipulation tasks in 4 different domains. Domains that empiri- cally present slower convergence are given bigger budget of interactions. The number of stages determined according to the nature of the task and by the typical horizon of demonstrations.
Domain
Tasks
Demos
Interactions
Stages
Phase 1: Pretraining. One of the main bottlenecks of RL in long-horizon sparse reward tasks is the low-informative data that is collected at the early stages of training. As early data tends to contain no rewards, learning a meaningful representation for such states becomes challenging. Tradi- tional methods tend start collecting data with randomly initialized policy that usually struggles find any rewards in the environment. For this reason, we jointly pre-train policy TBC and an encoder hac on the full demonstration dataset D using the classic behavioral cloning (BC) loss as an objective function:
ManiSkill
5
5-100
500k
3
Meta-World
5
5
500k
2
Robosuite
4
5-25
100k
1
Humanoids
2
5
100k
3
from. This approach prevents collected data from quickly outnumbering the demonstrations. While the sampling ratio is tunable hyperparameter, we empirically found that an initial 50% demonstration ratio works well for most tasks.
Therefore, as detailed in Algorithm 1, at each update step, the world model gets updated as explained in Section 3.1. The agent will then proceed to interact with the environment to collect more data that will be stored in the replay buffer.
LBC
= E (o.a)rD
[o To (a
(6)
where the policy learns to imitate the behaviors encoded in the dataset. At interaction time, the interactive policy TeL and the world model encoder, h, are initialized with their pre-trained analogs TBC, hBc.
Similar to Lancaster et al. (2024), we use annealing to con- trol the probability of a sample coming from TBC or from the planning module of the world model. In this way, the data distribution of the replay buffer B is initially biased toward the one of the dataset D. This technique aims to collect more informative data than the one purely random policy would collect during early stages of training. As the world model starts learning and TRL is able to collect more meaningful samples, we increase the annealing coefficient a to improve the diversity of collected samples and stop re- lying on the suboptimal pre-trained policy TBC- Eventually, a, will converge to 1, which point all samples will come from planning with TRL (see Algorithm 1).
Given that we focus on datasets with limited demonstrations, behavioral cloning can be prone overfitting (Peters al., 2010; Parisi al., 2022; Duan al., 2017). To mitigate this, we regularly evaluate TBC during pretraining by rolling out episodes in the environment. We use early stopping on the evaluation set (Y (Yao et al., 2007) to select the best-performing policy for the interactive learning phase.
Interactive Learning. After initial pretraining of the encoder and policy, the agent starts collecting data from the environment to learn using offline reinforcement learning (RL). In order to utilize the demonstrations, we follow Hansen et al. (2023) by sampling from the replay buffer B and the demonstration dataset at each update step. Specifically, every time we sample a batch, a fraction of the samples come from D while the remaining fraction comes
Phase 2:
4. Experiments
We consider 16 challenging visual multi-stage manipulation tasks with a long-horizon for our experimental evaluation. This includes 5 manipulation tasks from ManiSkil13 (Tao
5

---

Demonstration-Auggeent Reward, Policy, and World Model Learning
Stack Cube
Stick Push
Place Apple
Door
100%
50%
0%
Peg Insertion
Stick Pull
Transport Box
Pick Place Can
100%
50%
0%
250K
500K
250K
500K
50K
100K
50K
100K
Interaction Steps
Ours
MoDem
TD-MPC2
LaNE
Figure 6. Challenging tasks. Success rate of our method and baselines on the hardest tasks from each domain. Averaged across 5 random seeds. LaNE results in Robosuite are kindly provided by Zhao al. (2024). Shaded areas correspond to 95% confidence intervals.
et al., 2024b), 5 manipulation tasks from Meta- World (Yu al., 2021) and 4 manipulation tasks from Robosuite (Zhu et al., 2022). Additionally, we include 2 humanoid manipu- lation tasks from ManiSkill3 with a high-dimensional action space, which we refer to as ManiSkill Humanoids in our evaluations. We place a strong emphasis on long-horizon precise manipulation, which is why we select the most chal- lenging tasks from each domain. We relate the difficulty to the required precision of task, its horizon, and the level of randomization in the scene (see Appendix D.3 for further details on difficulty categorization). For each task, the agent is given a constrained budget of demonstrations and inter- action steps (see Table 1). To allow most baselines to solve the task, we set the interaction budget and demonstrations to different amount for each task. For a complete list of details on our experimental setup, please refer to Table I and Appendix D. Through our evaluations, aim answer the following questions:
with sparse rewards. Similarly to our method, MoDem em- ploys three-phase framework: policy pretraining, seeding, and interactive learning with oversampling of demonstra- tion data. The authors show state-of-the-as art performance on Meta-World and Adroit domains from visual inputs and sparse rewards.
LaNE (Zhao et al., 2024) is a data-efficient model-free RL method for sparse-reward tasks from visual inputs. LaNE utilizes pre-trained feature extractor to learn embedding space and rewards the agent for exploring regions near the demonstrations within this latent space. The authors also show sate-of-theear data-efficiency in the Robosuite environment with a limited amount of demonstrations.
TD-MPC2 (Hansen et al., 2022; 2024) is the state-of-the art MBRL algorithm for control tasks. combines temporal difference learning with model predictive control (MPC) and constitutes the backbone of our approach. Compared to TD-MPC, TD-MPC2 includes a series of algorithmic changes that improve robustness and scaling.
1. Can our proposed method effectively accelerate MBRL with demonstrations in long-hon Long-horizon multi-stage stage tasks?
2. What is the relative importance of each algorithmic com- Reward, Policy,
4.2. Benchmark Results
ponent of Demonsttration-Augmeenn and World Model Learning (DEMO*), and how does it scale with the amount of demonstration data?
The main result of our evaluations (see Figure 5) compares the data-efficiency of our method against the proposed base- lines. On average, our method achieves 40% better data- efficiency than the proposed baselines. Notably, Man- iSkill3, our most difficult domain, our method averages a 75% success rate after only 500k steps, performing 50% bet- ter than the second- best baseline. Furthermore, our 3EMO3 can deal better with the high-dimensional action space in ManiSkill Humanoids. Interestingly, while our 3 does better on average, thrives in the most difficult tasks where the horizon and precision of the task are the highest. While the results of LaNE in the Robosuite domain are cer-
3. How do sparse reward functions compare our learned rewards at different levels of stage granularity?
4.1. Baselines
To assess our method's effectiveness, we compare it against three relevant approaches. A complete comparison of other methods can be found in Appendix B.
MoDem (Hansen et al., 2023) is a MBRL algorithm de- signed enhance data-efficiency in visual control tasks
6

---

Demonstration-Auugeene Reward, Policy, and World Model Learning
Average
Stack Cube
Insertion
100%
500
Steps reach 30% success
400
50%
300
Lift Peg Upright
Poke Cube
Pick Place
100%
200
50%
100
0%0
250K
500K6
250K
500K
250K
500K
Number of
Interaction Steps
Ours
TD-MPC2 + learned reward
-uuss
MoDem TD-MPC2
no learned reward
MoDem
no learned reward TD-MPC2 learned reward
Figure 8. Demonstration efficiency. Number of steps to reach a
Figure 7. Ablations. Success rate function of interaction steps for variations of our method on all 5 ManiSkill manipulation tasks.
critical success rate (30%) as function of demonstration count. Data points that did not converge are assigned a 500k step count. 2 challenging manipulation tasks (Stack Cube and Peg Insertion) and averaged across 5 seeds. Shaded areas correspond a 95% confidence interval.
Averaged across random seeds. Baselines included for complete-
Results aggregated
ness. Shaded arcas correspond to 95% confidence intervals.
tainly impressive, the same approach doesn't 't transfer to the rest of the domains.
DEMO shows the strongest performance in the lowest regime of demonstrations. This is particularly evident in challenging tasks such Peg Insertion and Stack Cube, where DEMO is the only method capable of reaching meaningful performance within the interaction budget (see Appendix A.2) with only 5 demonstrations.
Figure 6 shows the learning curves for the most challenging task of each benchmark. Particularly, the ManiSkill tasks, Peg Insertion, and Stack Cube, require a very high level of precision and a long horizon. As shown in 6, DEMO is the only algorithm to reliably solve both tasks in the interaction budget. While performance quite matched with LaNE (Zhao et al., 2024) in Robosuite, LaNE uses a pre-trained encoder preprocess image observations while our method is completely learned from scratch. Finally, TD-MPC2 struggles to get any performance as is typical for pure RL algorithms learning from sparse rewards. Overall, 3EMO shows the highest degree of robustness and efficiency on the proposed long-horizon tasks.
Wall-time comparison Our method achieves competitive wall-time performance, ranking as the second fastest among all evaluated algorithms (Table 2). While slightly slower than TD-MPC2, we attribute the overhead to the additional computation required for reward learning. Importantly, DEMO performs much faster than other demonstration- augmented RL approaches, such as Modem and LaNE
Table 2. Wall-time. Hours per 100k interaction steps, averaged across seeds and all tasks in Robosuite. Lower is better 4.
4.3. Analysis
Relative importance of each component Figure shows the effect of removing dense reward learning, policy pretraining, and demonstration oversampling in the 5 manip- ulation tasks from ManiSkil13. Interestingly, considerable jump in performance is brought by pretraining and oversam- pling from the demonstration dataset with the TD-MPC2 backbone (no learned reward). The effect of reward learn- ing becomes evident in long-horizon tasks where advancing stages can be very challenging without rewards (e.g., Peg Insertion, Pick Place), especially when reducing the number of demonstrations (see Appendix A.2).
Algorithm Time (hours) 1
LaNE
20.40
MoDem
8.37
TD-MPC2
4.84
Ours
5.19
Reward granularity Figure illustrates the data- efficiency of our method across the same tasks under varying reward granularities: stage, 2 stages, 3 stages, and dense rewards. Please refer to D for details about the stage def- initions. As expected, dense rewards provide the optimal learning signal, resulting in the fastest task completion. Re- markably, our method achieves near-equal performance with only 2 stages, demonstrating that stage indicators combined
Demonstration efficiency In Figure 8, we experiment with different dataset sizes and evaluate the dffa-efficiency of different baselines using demonstrations. While most methods scale well with the number of demonstrations,

---

Demonstration-Augmeer Reward, Policy, and World Model Learning
with minimal set of demonstrations are sufficient to guide the learning process effectively. This result shows that over- engineered dense reward functions can easily be avoided by only providing a small number of demonstrations. By relying solely on stage indicators and demonstrations, our method simplifies the reward design process while maintain- ing high performance in complex long-horizon tasks.
et al., 2020; Singh et al., 2019). In this work, we leverage demonstrations in multiple ways simultaneously: learning initial policy, the world model, and reward function.
Reward Learning Designing rewards is challenging due to the need for extensive domain knowledge, prompting the development of data-driven reward learning methods. Rewards can be learned from offline datasets by classifying goals (Smith et al., 2019; Kalashnikov et al., 2021; Du et al., 2023) estimating goal distances (Zakka et al., 2022). Al- ternatively, inverse RL approaches (Ng al., 2000; Ziebart et al., 2008; Ho & Ermon, 2016; Fu et al., 2017) leverage online interactions to infer reward function from expert demonstrations. Additionally, reward shaping techniques (Trott et al., 2019; Wu et al., 2021; Memarian et al., 2021; Escontrela et al., 2022) transform sparse rewards into dense rewards using specific domain knowledge. The reward learn- ing in this work builds upon DrS (Mu et al., 2024), with modifications tailored for compatibility with MBRL.
100%
80%
60%
40%
20%
0%
100K
200K
Interaction Steps
1 stage Figure 9. Reward granularity. Success rate of our method with increasing granularity in the stage division of task. Results are aggregated over 2 challenging manipulation tasks (Stack Cube and Peg Insertion) and averaged across seeds. The shaded area corresponds to a 95% confidence interval.
Dense 3 stages
2 stages
6. Conclusions and Future Directions
In this work, we tackle the challenge of learning long- horizon manipulation skills with sparse rewards using only proprioceptive and visual feedback. We propose DEMO demonstration--uugnnn auggeented MBRL algorithm that simulta- neously learns a reward, policy and world-model for multi- stage manipulation. Our experiments (Section 4) show that our method achieves 40% better performance than the cur- rent state-of-the-a Additionally, DEMO3 excels at the most difficult tasks, converging up to 4x faster than current methods. We propose possible future research directions:
5. Related work
Model-based RL Model-based Reinforcement Learning (MBRL) improves data-efficiency by leveraging a model of the enviromment to guide decision-make making. These models can be either prior-based, such physics-based simulators, or learned, where the agent approximates a dynamic model of the world from data. World models (Schmidhuber, 1990; Ha& & Schmidhuber, 2018) are an internal representation of the environment that enables planning and policy learning without direct interaction. A notable example is MuZero (Schrittwieser& et al., 2020), which extends value-based plan- ning by implicitly learning environment dynamics. Recent advances, such as Dreamer (Hafner et al., 2020; 2022; 2024) and TD-MPC (Hansen et al., 2022: 2024), improve learn- ing in high-dimensional spaces, allowing MBRL to scale to complex visual and continuous control tasks.
Firstly, although our experiments show that small number of demonstrations is sufficient for accelerated learning, we do not explore the effect of using different sources, e.g., data collected via human telcoperation. While prior work suggests such sources have limited impact on RL perfor- mance (Hansen al., 2023). improving robustness to di- verse demonstration types remains relevant for deployment.
Another limitation is that our experiments are limited to simulation. We aim to deploy DEMO* on real robot hard- ware to evaluate its generalization under domain shift and whether the same constrained regime of demonstrations and interaction samples can yield reliable real-world policies. Notably, Hansen et al. (2023) shows that a similar pipeline transfers well to physical systems. Since DEMO3 improves these components in simulation, we are optimistic its benefits will carry over.
inmonstration-Augmeettte RL Learning policies purely through trial and error can be inefficient and unstable, prompting research into leveraging demonstrations to en- hance RL. During online interactions, demonstrations can serve off-policy experience (Hester et al., 2018; Kaptur- owski et al., 2018; Ball et al., 2023; Nair et al., 2018; Escon- trela et al., 2022) or for on-policy regularization (Kang et al., 2018; Rajeswaran et al., 2017). Alternatively, demonstra- tions can be used to estimate reward functions for RL (Xie et al., 2018; Aytar et al., 2018; Vecerik et al., 2019; Zolna
Finally, while our 50% demonstration sampling ratio yields strong results, more sophisticated strategies (Schaul et al., 2016) could further enhance learning efficiency.
8

---

Demonstration-Auumeene Reward, Policy, and World Model Learning
Acknowledgements
Duan, Y., Andrychowicz, M., Stadie, B. C., Ho, J., Schnei- der, J., Sutskever, I., Abbeel, P., and Zaremba, W. One- shot imitation learning, 2017. URL pestps://arxio org haaabb//1003 07326.
Nicklas Hansen supported by NVIDIA Graduate Fellow- ship. Stone Tao supported in part by the NSF Gradu- Research Fellowship Program grant under grant No. DGE2038238.
Escontrela, A., Peng, X. B., Yu, W., Zhang, T., Iscen, A., Goldberg, K., and Abbeel, P. Adversarial mo- tion priors make good substitutes for complex reward functions, 2022. URL htttt:t//axxxv..oo http://xxxv.../aaaaba/
Impact Statement
This paper presents work whose goal is to advance the field of robot learning. While there are many potential societal consequences in developing embodied intelligence for the real world, we do not feel our work presents any particular implications that should be highlighted here.
Fu, J., Luo, K., and Levine, S. Learning robust rewards with adversarial inverse reinforcement learning. arXiv preprint
2017.
Ha, D. and Schmidhuber, J. World models. 2018. doi:
References
10.528 11.288//ZNNOD0.12277333 URL org/ wttp////oood/223////
Atkeson, C. G. and Schaal, S. Robot learning from demon- stration. In Proceedings of the Fourteenth International Conference on Machine Learning, ICML '97, pp. 12-20, San Francisco, CA, USA, 1997. Morgan Kaufmann Pub- lishers Inc. ISBN 1558604863.
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination, 2020. URL htwps:tt/////... httpaaaxvvvvroo.. abs / 91 2. 01 603.
Hafner, D., Lillicrap. T., Norouzi, M., and Ba, J. Mastering atari with discrete world models, 2022. URL https: / / //aaxiiv..rr. www.aas//////o.oooo.o...oe..aeat 2010 02193.
Aytar, Y., Pfaff, T., Budden, D., Paine, T., Wang, Z., and De Freitas, N. Playing hard exploration games by watch- ing youtube. Advances in neural information processing systems, 31, 2018.
Hafner, D., Pasukonis, J., Ba, J., and Lillicrap. T. Master- ing diverse domains through world models, 2024. URL https: /ttpxiaxvvv..//.. abs / 2/0//33111000 00233.044100 04104.
Ball, P. J., Smith, L., Kostrikov, I,, and Levine, S. Ef- ficient online reinforcement learning with offline data. In International Conference on Machine Learning, pp. 1577-1594. PMLR, 2023.
Hansen, N., Wang, X., and Su, H. Temporal difference
learning for model predictive control. In ICMI. 2022.
Hansen, N., Lin, Y., Su, H., Wang, X., Kumar, V., and Rajeswaran, A. Modem: Accelerating visual model- based reinforcement learning with demonstrations. 2023.
Baumli, K., Baveja, S., Behbahani, F., Chan, H., Comanici, G. Flennerhag, S., Gazeau, M., Holsheimer, K., Hor- gan, D., Laskin, M., Lyle, C., Masoom, H., McKinney, K., Mnih, V.. Neitz, A., Nikulin, D., Pardo. F., Parker- Holder, J., Quan, J.. Rocktäschel, T., Sahni, H.. Schaul, T., Schroecker, Y., Spencer, S., Steigerwald, R., Wang, L., and Zhang, Vision-language models as a source of rewards, 2024. URL https: /arxivi.vvvv./////// org/ abs/ 2312. 09187.
Hansen, N., Su, H., and Wang, X. Td-mpc2: Scalable, robust world models for continuous control, 2024.
Hester, T., Vecerik, M., Pietquin, O., Lanctot, M., Schaul, T., Piot, B., Horgan, D., Quan, J., Sendonaris, A., Osband, I., et al. Deep q-learning from demonstrations. In Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
Bellman, R. A markovian decision process. Indiana Univ.
Math. J., 6:679-684 684, 1957. ISSN 0022-2518.
Ho, J. and Ermon, S. Generative adversarial imitation learn- ing. Advances neural information processing systems, 29, 2016.
Clark, J. and Amodei, D. Faulty reward functions in the wild. OpenAl Blog, 2016.
Di Palo, N. and Johns, E. Learning multi-stage tasks with one demonstration via self-replay. In Conference on Robot Learning (CoRL), 2021.
Kalashnikov, D., Varley, J., Chebotar, Y., Swanson, B. Jonschkowski, R., Finn, C., Levine, S., and Hausman, K. Mt-opt: Continuous multi-task robotic reinforcement learning scale. arXiv preprint arXiv:2104.082. 2021.
Du, Y., Konyushkova, K.,
Denil, M., Raju, A., Landon, J., Hill, F. de Freitas, N., and Cabi, S. Vision-language mod- els as success detectors. arXiv preprint arXiv:2303.0728.. 2023.
Kang, B., Jie, Z., and Feng, J. Policy optimization with demonstrations. In International conference on machine learning, pp. 1469-2478, PMLR, 2018

---

Demonstration-Augmeeie Reward, Policy, and World Model Learning
Kapturowski, S., Ostrovski, G., Quan, J., Munos, R., and Dabney, W. Recurrent experience replay in distributed reinforcement learning. In International conference learning representations, 2018.
Pomerleau, D. A. Alvinn: An autonomous land
vehicle in neural network. In Touretzky, D. (ed.), Advances in Neural Information Process- ing Systems, volume 1. Borgan-Kaufmamm 1988. URL. cc 8 2122laaa2875eeeeb33444bbbrbbaaaa pdf.
Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims, T. Morel Model-based offline reinforcement learning. ArXiv, abs/2005.05951 2020.
Paper.
Kumar, S., Zamora, J., Hansen, N., Jangir, R., and Wang, X. Graph inverse reinforcement learning from diverse videos. Conference on Robot Learning (CoRL), 2022.
Rajeswaran, A., Kumar, V., Gupta, A., Vezzani, G., Schul- man, J., Todorov, E., and Levine, S. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. arXiv preprint arXiv:1709.1000 2017.
Kwon, M., Xie, S. M., Bullard, K., and Sadigh, D. Re- ward design with language models. In The Eleventh International Conference on Learning Representations, 2023. URL https: http://pppeeeeeeeeeeeeeeeeeeeeee id-1OuNUg15K1.
Rocamonde, J., Montesinos, V., Nava, E., Perez, E., and Lindner, D. Vision-language models are zero-shot reward models for reinforcemeet learning, 2024. URL https:
www.x/..wrg/abb//333332 12921.
Lancaster, P., Hansen, N., Rajeswaran, A., and Kumar, V. Modem-v2: Visuo-motor world models for real-world robot manipulation. In International Conference on Robotics and Automation (ICRA), 2024.
Schaul, T., Quan, J., Antonogiou, I., and Silver, D. Priori- tized experience replay, 2016. URL https: //axxiv. org abs///a////// /111111111. 1511.05952.
05952.
Ma, Y. J., Liang, W., Wang, G., Huang, D.-A., Bastani, O., Jayaraman, D., Zhu, Y., Fan, L., and Anandkumar, A. Eureka: Human-level reward design via coding large lan- guage models. arXiv preprint arXiv: 2023.
Schmidhuber, J. Making the world differentiable: on us- ing self supervised fully recurrent neural networks for dynamic reinforcement learning and planning non- stationary environments. Forschungsberichte Künstliche Intelligenz. Inst. für Informatik, 1990. URL https: / books google..oo..oooooooo ch/bockssss....ccc ?1d-9c2HMMAAAAAA
Memarian, F., Goo, W., Lioutikov, R., Niekum, S., and Topcu, U. Self-supervised online reward shaping in sparse-reward environments. In 2021 IEEE/RSI Inter- national Conference on Intelligent Robots and Systems (IROS), pp. 2369-2377 IEEE, 2021.
Schrittieserr, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre. L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., Lillicrap, T., and Silver, D. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839)6044609 doi: 10.1038/
Mu, T., Liu, M., and Su, H. Drs: Learning reusable dense re- wards for multi-stage tasks. In The Twelfth International Conference Learning Representations, 2024.
December 2020. ISSN 1476-4687. s41586- 03051-4. URL http: http::::/xxxdx...pooo doi.org/ 10. 10.1038/20 s 00044556000000011111 03051-44
Nair, A., McGrew, B., Andrychowicz, M., Zaremba, W., and Abbeel, P. Overcoming exploration in reinforcement learning with demonstrations. In 2018 IEEE international conference robotics automation (ICRA), pp. 6292- 6299. IEEE, 2018.
Sferrazza, C., Huang, D.-M., Lin, X., Lee, Y., and Abbeel, P. Humanoidbench: Simulated humanoid benchmark for whole-body locomotion and manipulation, 2024.
Singh, A., Yang, L.. Hartikainen, K., Finn, C.. and Levine, S. End-to-end robotic reinforcement learning without reward engineering. arXiv preprint arXiv: urXiv:1904.0785. 2019.
Ng, A. Y., Russell, S., et al. Algorithms for inverse rein- forcement learning. In leml, volume pp. 2, 2000.
Parisi, S., Rajeswaran, A., Purushwalkam, S., and Gupta, A. The unsurprising effectiveness pre-trained vision mod- els for control, 2022. URL https:ttp//........ hwt.t::ttt//attttieee htttps://arxiv.rrgg abs/ 2203/220..0 03580.
Smith, L., Dhawan, N., Zhang, M., Abbeel, P., and Levine, S. Avid: Learning multi-stage tasks via pixel-level trans- lation of human videos. arXiv preprint arXiv:1912.04443, 2019.
Peters, J., Mülling, K., and Altün, Y. Relative entropy policy search. In Proceedings the Twenty-Foure AAA/ Confer- ence on Artificial Intelligence, AAAI 10, pp. 1607-1612. AAAI Press, 2010.
Smith, L., Dhawan, N., Zhang, M., Abbeel, P., and Levine, S. Avid: Learning multi-stage tasks via pixel-level trans- lation of human videos, 2020. URL https: /arxiv. org has//ab///aa 1912. 04443.
10

---

Demonstration-Augmeer Reward, Policy, and World Model Learning
Sutton, R. S. and Barto, A. G. Reinforcement Learn- ing: An Introduction. The MIT Press, second edition, 2018. URL http:/ www.coopleeeddes..eeee net/ book/the-book-220 httm......d.d.......
Yu, T., Quillen, D., He, Z.,, Julian, R., Narayan, A., Shively, H., Bellathur, A., Hausman, K., Finn, C., and Levine, S. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning, 2021. URL https: / //arxiv. www.xiv..rg.ee.. www. abs////////// 1 10897.
Tao, Shukla, A., Chan, T.-k., and Su, H. Reverse forward curriculum learning for extreme sample and demonstra- tion efficiency in rl. 2024a.
Zakka, K., Zeng, A., Florence, P., Tompson, J., Bohg, J., and Dwibedi, D. Xirl: Cross-embodiment inverse rein- forcement learning. In Conference on Robot Learning, pp. 537-546. PMLR, 2022.
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., Huang, Z., Calandra, R., Chen, R., Luo. S., and Su, H. Man- iskill3: Gpu parallelized robotics simulation and ren- dering for generalizable embodied ai, 2024b, URL https: /awxivvivv//./////////a////aaall org/ http//b//22/ 2000/2210..0 00425.
Zhan, A., Zhao, R., Pinto, L.., Abbeel, P., and Laskin, M. Learning visual robotic control efficiently with con- trastive pre-training and data augmentation, 2022. URL https:/ http://axxvvv.oooo https/oorrrrbbb//2..70277
Zhang, M., Vikram, S., Smith, L., Abbeel, P., Johnson, M. J., and Levine, S. Solar: Deep structured latent represen- tations for model- reinforcement learning. ArXiv, abs/1808.099 0910,, 2018.
Trott, A., Zheng, S., Xiong, C., and Socher, R. Keeping your distance: Solving sparse reward tasks using self- balancing shaped rewards. Advances in Neural Informa- tion Processing Systems, 32, 2019.
Zhao, R., ufuk topcu, Chinchali, S. P., and Phiclipp, M. Accelerating visual sparse-rewand reward learning with latent nearst-demmonstaaioon-guuud explorations. In 8th Annual Conference on Robot Learning, 2024. URL https:
Vecerik, M., Sushkov, O., Barker, D., Rothörl, T., Hester, T., and Scholz, A practical approach to insertion with variable socket position using deep reinforcement learn- ing. In 2019 international conference on robotics and automation (ICRA), pp. 754-760. IEEE, 2019.
/ //openreview. net/ ott.aa/w..e....emmoooooooooossse
Zhu, Y., Wong, J., Mandlekar, A., Martín-Martin, R., Joshi, A., Nasiriany, S., and Zhu, Y. robosuite: A modular simu- lation framework and benchmark for robot learning, 2022. URL https::tp/..... /tap:s///rreee arxiv. .rgg abs / 20w. 12293.
Williams, G., Aldrich, A., and Theodorou, E. Model pre- dictive path integral control using covariance variable importance sampling, 2015. URL https://arxi:: org/abs/1500 01149.
Ziebart, B. D., Maas, A. L.., Bagnell, J. A., Dey, A. K., et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pp. 133-- Chicago, IL. USA, 2008.
Wu. Z., Lian, W., Unhelkar, V., Tomizuka, M., and Schaal, S. Learning dense rewards for contact-rich manipula- tion tasks. In 2021 IEEE International Conference Robotics and Automation (ICRA), pp. 2214-6221 IEEF, 2021.
Zolna, K., Novikov, A., Konyushkova, K., Gulcehre, C., Wang, Z.,, Aytar, Y., Denil, M., de Freitas, N., and Reed, S. Offline learning from demonstrations and unlabeled experience. arXiv preprint arXiv:2011. 2020.
Xie, A., Singh. A., Levine, S., and Finn, C. Few-shot goal inference for visuomotor learning and planning. In Conference on Robot Learning, pp. 40-52. PMLR, 2018.
Xie, T., Zhao, S., WU, C. H., Liu, Y., Luo, Q., Zhong, V., Yang, Y., and Yu, T. Text2reward: Reward shaping with language models for reinforcement learning. In The Twelfth International Conference on Learning Represen- tations, 2024. URL h http///pppeeeeeee////e / forum? t= RY
Yao, Y., Rosasco, L., and Caponnetto, A. On early stop- ping in gradient descent learning. Constructive Ap- proximation, 26:289-315, 2007. URL https://aaat semanticschola:: wttp:laarrrgc/ooouaaeeeeeeeeeeee : 8323954.
Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J. Y., Levine, S., Finn, C., and Ma, T. Mopo: Model-based offline policy optimization. ArXiy, abs/2005. 13239, 2020.

---

Demonstration-Auumeenee Reward, Policy, and World Model Learning
A. Additional Results
A.I. Single Task Experimental Results
Average
Stack Cube
Peg Insertion
100%
50%.
0%
Lift Peg Upright
Cube
Place
100%
50%.
0%
250K
250K
250K
Interaction Steps
Ours
MoDem
TD-MPC2
LaNE
Figure 10. ManiSkill Manipulation results. Results averaged across seeds. The shaded area corresponds to a 95% confidence interval.
Average
Assembly
Peg Insert Side
100%
50%
NY M
0%
Pick Place
Stick Push
Stick Pull
100%
50%
0%
250K
250K
500K
250K
Interaction Steps
Ours
MoDem
TD-MPC2
LaNE
Figure 11. Meta- World results. Results averaged across 5 seeds. The shaded area corresponds to a 95% confidence interval.
Average
Place Apple
Transport Box
100%
50%
0%
250K
500K
250K
500K
250K
500K
Interaction Steps
Ours
MoDem
TD-MPC2
LaNE
Figure 12, ManiSkill Humanoids results. Results averaged across 5 secds. The shaded area corresponds to a 95% confidence interval.
12

---

Demonstration-AAgmennee Reward, Policy, and World Model Learning
Lift
Door
100%
0%
50K
100K
0
50K Stack
100K
Pick Place Can
100%
50%
50K
100K
50K
100K
Interaction Steps
Ours
MoDem
TD-MPC2
LaNE
Figure 13. Robosuite results. Results averaged across 5 seeds. The shaded corresponds 95% confidence interval.
A.2. Additional Ablations
demos
10 demos
25 demos
100%
MMM
50%
MmM
0%
50 demos
100 demos
200 demos
100%
50%
0%
0
250K
500K
250K
500K
250K
500K
Interaction Steps
Ours
MoDem
no learned reward
TD-MPC2 learned reward
Figure 14. Demonstration ablation. Success rate on an increasing number of demonstrations (5-200) in the 2 most challenging manipulation tasks in ManiSkill (Stack Cube and Peg Insertion). DEMO the only method that has relative success with only demonstrations. Results are aggregated over both tasks and averaged across seeds.
13

---

Demonstratioon-Augeente Reward, Policy, and World Model Learning
A.3. Learned Reward
33
2.5
1.5
00
6.3
20
2.3
faandd
1.0
20
20
3.0
3.5
0.8
Figure 15. Learned reward animation. Visualization of the learned dense reward across three representative rollouts successful, failed, and successcce,,ul, The reward curves closely track task progress.
14

---

Demonstratioon-Augmentee Reward, Policy, and World Model Learning
B. Baselines
B.1. Comparison to prior work
In Table 3, we compare key components DEMO' to relevant prior methods on demonstration -aggrooteed RL. Our approach is the only one incorporating online reward learning in multi-stage settings for visual inputs and sparse rewards.
Table Comparison prior work. We compare DEMO relevant approaches and ablations. Selected baselines highlighted
Method
Visual Inputs
Sparse Rewards
Multi-Stage Online Reward Learning
Ours
LaNE (Zhao et al., 2024)
MoDem (Hansen et al.. 2023)
X
CoDER (Zhan et al., 2022)
X
SAC DrS (Mu et al., 2024)
X
X
AMP (Escontrela al., 2022)
X
X
V
B.2. Baseline Implementations
TD-MPC2 We use the official implementation' with default parameters. We add two extra layers to the convolutional encoders to handle the higher image resolution of the Meta-Word World and Robosuite benchmarks.
MoDem We use the official implementation with default parameters. To process observations containing multiple images, we add an extra encoder to the world model and average all the embeddings.
We use the code from the official implementation' with default parameters. To adapt it to our experimental setup, we add extra encoders handle the additional observations. The proprioceptive state is passed through an MLP, and its embedding is averaged with the other inputs. Additionally, adapt the algorithm to handle infinite-horizon MDPs by removing value function bootstrapping from the MDP.
LaNE
15

---

Demonstratiion-Augmentee Reward, Policy, and World Model Learning
C. Demonstrations
of our demonstrations are obtained by training a TD-MPC2 model with dense rewards and state observations. The model trained on state observations is then used to rollout N episodes in the stage-based environment, from where we query image observations, proprioceptive states, and sparse stage rewards. Please find below a detailed table the number of demonstrations used per task.
Table 4. Number of demonstrations for each task. We use the minimum amount of demonstrations (empirically determined) to ensure best-performing algorithm can solve task the given interaction budget.
Domain
Task
Number of Demonstrations
Peg Insertion
100
Pick Place
100
ManiSkill Manipulation
Stack Cube
25
Poke Cube
Lift Peg Upright
Assembly
5
Peg Insert Side
5
Meta-World
Stick Push
Stick Pull
Pick Place
Place Apple
ManiSkill Humanoids
Transport Box
Lift
5
Door
10
Robosuite
Pick Place Can
10
Stack Blocks
20
16

---

Demoustration-Augmeente Reward, Policy, and World Model Learning
D. Experiment Details
Stack Cube
Peg Insertion
Poke Cube
Pick Place
Lift Peg Upright
Assembly
Peg Insert
Stick Push
Slick Pull
Pick Place
Place Apple
Transport Box
Lift
Door
Pick Place Can
Stack Blocks
Figure 16. All tasks. Visual description of all tasks organized by domains. In descending order: ManiSkill Manipulation, Meta- World, ManiSkill Humanoids, Robosuite.
D.1. Domain Implementation
Table 5, Implementation details for each of four domains. Time horizon is measured in agent steps (policy forward passes). Proprio, stands for Proprioceptive State. Each domain uses different image resolutions according to the detail the scene.
ManiSkill Manipulation
Meta-World
ManiSkill Humanoids
Robasuite
Time Horizon
100
100
100
100
Image Size
128 128
224 x 224
128 128
128 x 128
Observations
RGB(x2) + Proprio.
RGB + Proprio.
RGB(x2) + Proprio.
RGB(x2)
Cameras
Hand Front
Front
Head + Front
Hand + Front
Action Repeat
2
1
Action Dim
17

---

Demonstratioo--Augmentee Reward, Policy, and World Model Learning
D.2. Stage Definitions
Table ManiSkill Manipulation stage definitions.
Task
Stage 1
Stage 2
Success Criteria
Stack Cube
Cube A grabbed.
Cube A above Cube B.
Cube A stacked on top of Cube B. Peg inserted in hole. Cube at goal.
Peg Insertion
Peg grabbed.
Peg aligned with hole. Cube close to goal.
Pick Place
Cube grabbed.
Poke Cube
Peg grabbed.
Peg touching Cube.
Cube at goal.
Lift Peg Upright
Peg grabbed.
Peg upright.
Peg upright on desk.
Table Mlea-WWrlld stage definitions.
Task
Stage
Success Criteria
Assembly
Grab hook.
Pass nut through pole.
Peg Insert
Peg grabbed.
Peg inserted in hole.
Stick Push
Stick grabbed.
Object pushed to goal location.
Stick Pull
Stick grabbed.
Object pulled to goal location.
Pick Place
Cube grabbed.
Cube is static goal.
Table 8. ManiSkill Humanoids stage definitions.
Task
Stage 1
Stage
Success Criteria
Place Apple
Apple is grabbed.
Apple is above bowl.
Apple is inside bowl.
Transport Box
Box grabbed with 2 hands.
Box above table 2.
Box is on table 2.
Table 9. Robosuite stage definitions.
Task
Success Criteria
Lift
Block lifted above the desk.
Door
Door is open.
Pick Place Can
Can is at goal location.
Stack
Block A is in contact with Block B and above the ground.
18

---

Demonstratiion-Augmentee Reward, Policy, and World Model Learning
D.3. Difficulty Categorization
Across this paper, we often refer to some tasks as more difficult than others. To characterize task difficulty, we follow (Tao reinforcement learning (RL), observe that environments
al., 2024a). As in previous work Cemonstration-augmeeeen with high complexity and substantial initial state randomization are typically more difficult to solve and require larger number of demonstrations. For example, the Peg Insertion task from ManiSkill exhibits significant variability in the peg's s position, orientation, and the hole's size. Consequently, around 100 demonstrations are needed to solve the task from visual inputs. In contrast, task like Meta- Sora- Assembly requires only demonstrations to be successfully solved. Figure 17 qualitatively compares the different initial states of these two tasks. We hypothesize that this effect is related to the distributional coverage of the demonstration dataset: higher randomization reduces the likelihood that the agent encounters familiar states during training if the dataset is limited. Therefore, as the variability in task's initial state increases, this variability must also be well-represented in the dataset ensure effective learning.
Figure 17. Randomization comparison. Qualitative comparison of both Peg Insertion tasks in the ManiSkill and Meta-World domain and Door task in Robosuite. Visibly, ManiSkill presents the highest level of randomization, not only varying the initial state at reset but also changing the geometric properties of the objects.
19

---

Demoustration-Auumeenee Reward, Policy, and World Model Learning
E. Implementation Details
E.1. Model Architecture
Following TD-MPC2, all modules are implemented as MLPs. Here, as an example, we summarize our architecture for a single-camera Meta- World task using Py PyTorch-like notation:
Architecture: ID-M2C2 World Model
Encoder: Moduleticz
1zgh_frontview): Sequentisl
: ShiftAug
PixalPreprocesa
(2): (3, 32, kernel_aise-(1 (?, Th, stride= (2, 21)
(inplace-True)
(4): (32, 32, kernel_size= 5), szridem 2))
(5): (tplaao-rrr))
(6) (32, 32, kernel_size=13, scridem 2))
(7): (inpaanao
Conv2d132,
kernel_size=p
s2ride=12,
(9) ReLU (einlaaecee)Tr
(10): Conv2d(32, 32, (i 3),, stride=(L) 11)
[11) Flazten (starr_ddmm_ end_dim=-l)
(12) (e__rrtt_fecueeeee
aut_features=612, blag=True)
(13): simNorm(dinee
Dynemics: Sequential
: (tr.cliieerrt__faatueeeeeeee
Stt_featuree=512, biaseTzue, actemish)
(1): Normedlincor (st.aiocrrrsn_fooourrrooosa sit_features=512, biaseTrue, act=Mish)
( Norredtinear (tna out out_reatures=512, blas=True, act=SimMorm)
Reward: Sequantial
(0): NosmedLineat (el_uueu,,,,-51,, tellurees,-- blas-True, act-Mish)
(1): NormedLineer (A (ea_ out Seatureesss bias-True, act-Misk)
(2): Linear (in_frrrn_
feat bisssIrue)
Policy prior: Sequentiall
(0): Normediinear (anrd///eerrssfffuuuurrsss/
f02turss555,, blas=True, act-Mish)
(1): No:medtiinear fet_reetu,,--,12, ret__reaure blas-True, act-Maat)
: (in_feetures=112, out_features=1d, bias-True)
G-functiona: Vectorized [Sequential
out tet_teature--112, blas-True, dropout=0.01, 01, act-Miah) (1): No:medLinear (neeeuuueeeet out retureeee,,-- blas-T:ue, sct-Mist.
(0): (ireettttenrrtsneetttrrreee
(2): Linear (infeeturll out_festuues=10l, biss-Irue)
,, Sequential
(0): NozmesLinear (tt.lliiearrin_aaaaaaaee Seaturreess,, fet_ieature=-112, bias-True, dropouz=0, 01, act-Kish)
(1): (o (in (i__feturrrss out sat_features:512, biassTrue, actaMish)
(2) finear (inoarr(t oot_festurea=101, biaseTrue)
Sequential
out S12u bias-True, dropout=0, 01, act-Kish) (in_ fet_feaaures:51, biaseTrue, acteMish) (2): Linear (11 out fet_uessures=101, biss-True)
(0):
(1):
tearnable paramsters: 748,
Discrtminator Architecture: Diacriminator!
(nets): KoduleLiat
:0): Sequentiel
(in fot_footures=]2, bias=Irue)
(11: Sigmeid()) (inaariin_foatureee222, fat_feauur,,n,, biaseTsue)
20

---

Demonstration-Auugeentt Reward, Policy, and World Model Learning
E.2. Hyperparameters
Most of the hyperparameters remain unchanged from our backbone algorithm, TD-MPC2. Here, we list some of the most relevant method and highlight the ones that are unique to approach. Please refer to the TD-MPC2 paper (Hansen et al., 2024) for complete list of hyperparameters.
Table 10. Hyperparameters used in the training setup.
Hyperparameter
Value
Replay buffer
Capacity
300, 000
Sampling
Uniform
Architecture (5M)
Encoder arch.
ConvNet (image inputs)
MLP (state inputs)
Conv. layers
7 (Meta- (Meta-World)
5 (Otherwise)
Encoder MLP dim Dynamics MIP dim
256
512
Latent state dim
512
Task embedding dim
96
Optimization
Update-to-data ratio Batch size
256
Joint-embedding coef.
20
Reward prediction coef.
0.1
Value prediction coef.
0.1
Temporal coef. (A)
0.5
Q-fn. momentum coef.
0.99
Policy prior entropy coef.
1 10
Policy prior loss norm.
Moving (5%, 95%) percentiles
Optimizer
Adam
Learning rate
3 x 10-4
Encoder learning rate
1 1 10
Pretraining
Pretraining loss
Behavioral cloning
BC Policy Architecture
MLP
MLP dim
512 Adam
Optimizer
Learning rate
3 x 10
Encoder learning rate
x 10
0g
x 10-5
Reward learning
Discriminator architecture
MLP
MLP dim
32
Discriminator learning rate
x 10
Discriminator optimizer
Adam
Batch size
256
B
1/3
Demo. sampling ratio
50%
21

---

Demoostratiion-Augmentee Reward, Policy, and World Model Learning
E.3. Computational Resources
All our experiments run single NVIDIA GeForce RTX 3090 GPU and 32GB of RAM to store collected samples.
22


Reviews:
[{"review_id": "XhYlsrvq3Z", "reviewer": "Reviewer_wCad", "summary": "This paper proposes DEMO, a framework that learns dense rewards from demonstrations and interactions with the environment to aid model-based RL learning. DEMO uses a multi-stage paradigm where it learns dense rewards from sparse reward signals (from stage indicators) to indicate \"progress\" and uses the learned dense reward to better learn policies. DEMO learns the reward model from visual signals and shows that incorporating the DEMO reward can help policies learn faster and converge to better performances. Experiments on a number of popular benchmarks (ManiSkill, Meta-World, Robotuite) show that the proposed method achieves SOTA performance using fewer samples.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback. We address your comments in the following.\n\n---\n\n**Q: Suggestion about live plots with policy rollout (1 to 1 video with reward evolution and policy rollout)**  \n**A:** We fully agree with the reviewer that it is important to verify whether the learned dense reward accurately reflects task progress, especially since this is central to our method. Following this suggestion, we have added live plots of the learned reward function aligned with video rollouts, now available in the **“Learned Dense Reward”** section of the project website: [https://sites.google.com/view/icml2025demo3](https://sites.google.com/view/icml2025demo3).\n\nTo provide a comprehensive overview, we include three representative cases:\n- a **successful rollout**,  \n- a **failed rollout**, and  \n- a **semi-successful rollout** that illustrates fluctuating behavior, alternating between progress and regression in the task.\n\nThese visualizations clearly demonstrate that the learned reward function captures nuanced task progression, aligning with intuitive notions of success, failure, and partial recovery. We believe this addition strongly supports the claim that our reward model tracks progress in long-horizon, multi-stage tasks.\n\n---\n\n**Q: Real robot experiments**  \n**A:** We acknowledge that validating DEMO3 in real hardware would further strengthen our contributions, and we plan to pursue this in future work. We are particularly interested in testing whether DEMO3 retains its strong sample efficiency and robustness under real-world conditions. Notably, prior work such as **MoDem-v2** has already demonstrated that a similar demonstration-augmented model-based pipeline can be successfully transferred to real robots. Given that DEMO3 builds on and improves these components, we are optimistic that its benefits will carry over to real-world deployment as well.\n\n---\n\nPlease do not hesitate to let us know if you have any additional comments.", "claims_and_evidence": "This work claims that using a small number of samples and learned dense rewards can speed up and help learning long-horizon tasks using MBRL. The results on the benchmarks and ablations verify these claims. Especially on the harder tasks (in ablation A.2), DEMO achieves the best results with a small number of demonstrations.", "methods_and_evaluation_criteria": "Yes, the proposed benchmarks and evaluation criteria are suitable for the application and problem setup. ManiSkill and MetaWorlds are challenging and popular benchmarks in this space, and the proposed method achieves better performance than baselines.\n\nNo real-robot experiments are conducted, though, which would help bolster the claims of the method.", "theoretical_claims": "No theoretical claims are proposed in this work.", "experimental_designs_or_analyses": "This work showcases results on long-term tasks such as \"stack cube\" and \"stick pull\" that best demonstrate the benefit of having multi-stage rewards and learned dense rewards. DEMO achieves the best results. These tasks can best demonstrate the benefit of using learned dense rewards as they are challenging and each step leads to another.\n\nIn terms of analysis, I feel like adequate analysis is provided in terms of performance of the overall results. Ablation also shows that each component is important for achieving the best result. Some missed opportunities in verifying that the learned dense reward actually corresponds to \"task progress\" as claimed in the paper.", "relation_to_broader_scientific_literature": "I think this work fits nicely in the model-based RL framework and the behavior cloning literature where the strengths of both fields are combined. If a few demonstrations can be effectively used to learn dense rewards for multi-stage and long-horizon tasks, then the benefit of using RL and interactive labels can be maximally leveraged for future results.", "essential_references_not_discussed": "References are adequate.", "other_strengths_and_weaknesses": "## Additional Strength\n- Overall, I think this paper provides an interesting and novel idea in dense reward learning from demonstrations and interactions with the environment.\n- Analysis of the reward granularity shows that the learned reward signals are near optimal compared to dense rewards engineered by humans.", "other_comments_or_suggestions": "I would suggest including live plots of the learned reward function and videos of the policy rollout to better showcase the capabilities of the learned reward function and policy.", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "None"}}, {"review_id": "vqVkguVBty", "reviewer": "Reviewer_xwA3", "summary": "This paper introduces a demonstration-augmented reinforcement learning method to solve data-efficient manipulation tasks with sparse rewards. By utilizing limited demonstrations, the policy, world model, and dense reward are effectively modeled, thus the long-horizon tasks can be solved in a multi-stage manner. Experiments on four benchmarks demonstrate the effectiveness of the proposed method.\n\n## update after rebuttal\nThanks to the authors for your efforts during rebuttal. Most of my concerns are resolved. I have raised my rating.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback. We address your comments below.\n\n---\n\n**Q: Benchmarks/tasks have a limited number of stages**  \n**A:** We appreciate the reviewer’s concern. While many tasks in our suite contain 2–3 defined stages, we argue that this already reflects significant temporal and compositional complexity, particularly in sparse reward visual manipulation. Prior works such as DrS, MoDem, and LaNE often evaluate on shorter or single-stage tasks. In contrast, DEMO3 handles longer-horizon behaviors with similarly sparse supervision, making our setup at least as challenging—if not more—than existing literature.  \nThat said, we agree that scaling DEMO3 to more complex, multi-stage tasks is a valuable next step. We are actively working on tasks like *Stack-N-Cubes*, requiring sequential completion of N > 3 subgoals. While we couldn’t finalize these in time for this submission, we plan to include preliminary results in the camera-ready version.\n\n---\n\n**Q: Improvements over MoDem and TDMPC2 seem trivial since DEMO3 builds upon them**  \n**A:** We thank the reviewer for this comment and appreciate the opportunity to clarify our benchmarking choices.. While DEMO3 does incorporate key ideas from both MoDem and TD-MPC2, we believe it is important—and indeed necessary—to benchmark directly against these methods to quantify the contribution of our core innovation: **dense, multi-stage reward learning**.\n\nTD-MPC2 and MoDem are **strong, state-of-the-art baselines** in visual model-based RL and demonstration-augmented RL, respectively. DEMO3 builds on their foundations but introduces a new training paradigm in which **reward, policy, and world model are trained jointly via online, stage-wise discriminator signals**, rather than using precomputed reward functions or task-specific shaping. This design results in a reward signal that evolves alongside the agent’s experience, improving both learning stability and sample efficiency.\n\nRather than relying on a single benchmark, our results show that introducing learned reward components consistently improves performance across tasks and domains, as illustrated in Figures 5 and 7. Furthermore, our ablation studies demonstrate that the learned reward is particularly important in high-variance, long-horizon tasks where sparse rewards are insufficient for policy discovery.\n\n---\n\n**Q: Comparing with Imitation Learning**  \n**A:** Thank you for the suggestion. To complement comparisons with demonstration-augmented RL methods, we now include *Behavioral Cloning (BC)* results across all tasks. These are visible on the [project website](https://sites.google.com/view/icml2025demo3) and will be added to the final manuscript.  \nAs expected, BC performs better in simpler settings like Robosuite, but fails to succeed in harder domains like ManiSkill Manipulation, which involve longer horizons, multi-stage coordination, and high variability (see Appendix D.3). This supports the view that stronger supervision—via dense rewards or interaction—is needed in these benchmarks.\n\n---\n\n**Q: Why does LaNE perform badly in ManiSkill and Meta-World?**  \n**A:** LaNE was originally evaluated on Robosuite and may be tuned to that benchmark. We attempted to adapt its hyperparameters to all environments, but performance still varied widely. DEMO3, by contrast, uses the same hyperparameters across all benchmarks, demonstrating stronger out-of-the-box generalization.  \nAs shown in Appendix D.3, Robosuite is comparatively easier, while ManiSkill and Meta-World involve longer horizons, greater variability, and multimodal interactions—conditions under which LaNE’s nearest-neighbor mechanism struggles to scale.  \nWe also observed that LaNE performs well in early stages but often fails to progress further, aligning with its reliance on latent-space similarity. Reward plots (see LaNE Task Progress in [Additional Ablations](https://sites.google.com/view/icml2025demo3/additonal-ablations)) illustrate this behavior. We will clarify these findings in the final manuscript.\n\n---\n\n**Q: ManiSkill ablations are not convincing enough**  \n**A:** While main ablations focus on ManiSkill, *per-task results across all domains* are provided in Appendix A.1. We chose ManiSkill because it is the most challenging domain in our suite, requiring precise control, long-horizon planning, and generalization under randomization.  \nBC performance also reflects this difficulty, performing modestly in Robosuite but struggling in ManiSkill. These characteristics make ManiSkill a valuable testbed for evaluating the impact of each component in DEMO3.  \n\nThat said, we agree that broader ablation coverage is useful. Additional ablations on Meta-World are available here: [Meta-World Ablations](https://sites.google.com/view/icml2025demo3/additonal-ablations). These show trends consistent with those in ManiSkill and support the generality of our framework.\n\n---\n\nPlease don’t hesitate to reach out with further comments.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Not sufficiently. One main goal of this paper is to solve the sparse reward problem in long-horizon tasks. However, the reviewer finds the used benchmarks/datasets have a limited number of stages (i.e., <3).", "theoretical_claims": "No theoretical claims are made.", "experimental_designs_or_analyses": "My major concerns lie in the experimental parts,\n\ni) Only three existing works (i.e., TD-MPC2, Modem, LaNE) are used for comparisons,\n  - the proposed method leverages the strength of Modem and TD-MPC2, thus the improvements over these two methods are a little bit trivial;\n  - Why the previous SoTA (ie, LANE) performs well on Robosuite while performing extremely badly on the other three benchmarks (Fig. 1)? \n  - since the method uses demonstrations for augmentation, existing imitation learning based methods should be compared;\n\nii) The tasks used for ablation are mainly chosen from ManiSkill, the environment where the method achieves the largest improvement over other methods. It is not convincing enough to evaluate the effectiveness of the proposed method;", "relation_to_broader_scientific_literature": "Related.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "No.", "other_comments_or_suggestions": "No.", "questions_for_authors": "Please see the experiment section.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "exu8dliVKA", "reviewer": "Reviewer_di7L", "summary": "This paper introduces DEMO3 (Demonstration-Augmented Reward, Policy, and World Model Learning), a novel framework for solving long-horizon, multi-stage manipulation tasks with sparse rewards. The authors address the challenge of designing dense reward functions and effectively exploring large state-action spaces by leveraging a small number of demonstrations for three key purposes: learning a policy, a world model, and a dense reward function. The approach incorporates multi-stage dense reward learning, a bi-phasic training scheme, and world model learning into a demonstration-augmented reinforcement learning framework. The method is evaluated across 16 sparse-reward tasks spanning four domains, including challenging humanoid visual control tasks, demonstrating improved data-efficiency by an average of 40% and by 70% on particularly difficult tasks compared to state-of-the-art approaches.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback. We address your comments in the following.\n\n---\n\n**Q: Discussion on limitations and future work**  \n**A:** Thank you for pointing this out. In response, we have prepared an expanded discussion on limitations and future work. We will incorporate this expanded discussion into the final manuscript upon acceptance.\n\n*We acknowledge that our current implementation relies on demonstrations collected via RL-trained agents, which are typically well-aligned with task objectives but often unimodal and less representative of human-like variability. A key future direction is to extend DEMO3 to handle more **diverse and multimodal demonstrations**, such as those arising from motion planning, human teleoperation, or video sources. Robustness to such sources would broaden the applicability of our method to real-world data collection settings.*\n\n*Another limitation is that our current experiments are confined to **simulation environments**. We aim to deploy DEMO3 on **real robot hardware** to evaluate how well our dense reward learning and world model generalize under domain shift, sensor noise, and actuation delays, and whether the same constrained regime of demonstrations and interaction samples can yield reliable real-world policies. Notably, related work such as **MoDem** has already demonstrated that a similar pipeline can transfer successfully to real robots. Given that DEMO3 builds on and improves these components in simulation, we are optimistic that its benefits will carry over to physical systems.*\n\n*Finally, while our current **50% demonstration sampling ratio** already yields strong results, we believe that more **sophisticated sampling strategies** could further enhance learning efficiency. Inspired by prioritized replay buffers, we plan to explore **priority-based demonstration sampling** that adaptively focuses on more informative or rare transitions during training.*\n\n---\n\n**Q: Thorough analysis of computational complexity**  \n**A:** We appreciate the reviewer’s request for a clearer discussion of computational complexity. As shown in **Table 2** of the main paper, DEMO3 introduces **minimal computational overhead** compared to TD-MPC2. Specifically, our method adds only a lightweight MLP per stage (i.e., the **stage discriminators**), which are trained jointly with the world model. This design keeps the additional parameter count and forward/backward pass time inconsequential relative to the backbone model.\n\nEmpirically, DEMO3 increases per-100k training time by only **6.7%** compared to TD-MPC2 (5.19h vs. 4.84h), and remains significantly faster than other demonstration-augmented baselines like **MoDem (8.37h)** and **LaNE (20.40h)**. These results suggest that DEMO3 offers strong performance gains with **minimal added training cost**, making it practical for real-world deployment scenarios where wall time is a limiting factor.\n\nWe will make this comparison more explicit in the final version of the paper.\n\n---\n\n**Q: Applying the method to a varying number of demonstrations**  \n**A:** We thank the reviewer for raising this important point. We address this question in our **demonstration efficiency analysis** (Figure 8), where we evaluate performance as a function of demonstration count (5, 10, 25, 50, 100, 200) on two of the most challenging tasks: **StackCube** and **PegInsertion**. DEMO3 is the only method that consistently reaches the target success threshold with as few as **5 demonstrations**, and it continues to improve steadily with more data, **outperforming all baselines across the range**.\n\nFor a more detailed view, we refer the reviewer to **Appendix A.2**, which includes **full learning curves under different demonstration regimes**. These results show that DEMO3 gracefully scales with more demonstrations while also being **highly effective in the low-data regime**.\n\n---\n\n**Q: How sensitive is the approach to stage definition? How does it transfer to more “unstructured” tasks?**  \n**A:** This is a great question. While DEMO3 is evaluated on tasks with explicit stage-wise structure, we believe that many real-world tasks — even those considered “unstructured” — can often be approximated with a **coarse stage decomposition**. For example, a task like human locomotion might involve an **initial acceleration phase** followed by a **steady gait phase**, which can naturally map to 2 distinct stages for learning purposes.\n\nTo better understand the effect of stage granularity, we include a **reward granularity ablation (Figure 9)**. This experiment compares performance across **1, 2, 3, and fully dense stage definitions**. While performance improves with finer-grained supervision, **DEMO3 remains surprisingly robust even under coarse or minimal stage labeling**, performing competitively with human-engineered dense rewards using only 1-stage definitions.\n\n---\n\nPlease do not hesitate to let us know if you have any additional comments.", "claims_and_evidence": "The claims made in the submission are well-supported by evidence. The authors claim that their method improves data-efficiency compared to state-of-the-art approaches, which is substantiated by comprehensive experimental results across 16 tasks in four domains. The learning curves in Figure 5 clearly demonstrate that DEMO3 consistently outperforms baseline methods (TD-MPC2, MoDem, and LaNE) in terms of success rate as a function of interaction steps. The claim that the method is particularly effective on difficult tasks is supported by the 70% improvement in performance on complex tasks. The authors also claim that their approach requires only a small number of demonstrations (as few as five), which is validated by their experimental setup described in Table 1.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are appropriate for the problem at hand. The authors use success rate as the primary evaluation metric, which is a clear and relevant measure for manipulation tasks. The experimental setup includes a diverse set of 16 tasks across four domains (ManiSkill Manipulation, Meta-World, Robosuite, and ManiSkill Humanoids), providing a comprehensive evaluation of the method's generalizability. The comparison against strong baselines (TD-MPC2, MoDem, and LaNE) ensures a fair assessment of the method's performance. The authors also conduct ablation studies to analyze the relative importance of each component of their framework, which helps to validate the design choices.", "theoretical_claims": "The paper makes several theoretical claims about the benefits of their approach, particularly regarding the use of stage-specific discriminators for dense reward learning and the bi-phasic training scheme. These claims are well-founded in reinforcement learning theory, particularly in the context of model-based RL and learning from demonstrations. The authors provide a clear mathematical formulation of their approach, including the loss functions for the discriminators (Equation 2) and the dense reward formulation (Equation 4). The theoretical justification for the bi-phasic training scheme is also well-articulated, explaining how it helps to overcome the exploration challenges in sparse reward settings.", "experimental_designs_or_analyses": "The experimental design is robust and comprehensive. The authors evaluate their method on 16 tasks across four domains, with varying levels of complexity and different numbers of stages. The experiments use 5 random seeds to ensure statistical significance, and the results are presented with 95% confidence intervals. The learning curves in Figure 5 provide a clear visualization of the method's performance over time, and the summary of results in Figure 1 offers a concise comparison with baselines. The ablation studies help to isolate the contributions of different components of the framework. The authors also provide details on the number of demonstrations used for each domain and the interaction budget, which helps to contextualize the results.", "relation_to_broader_scientific_literature": "The authors acknowledge prior work on model-based RL (Ha & Schmidhuber, 2018; Zhang et al., 2018; Kidambi et al., 2020; Hafner et al., 2020; Yu et al., 2020; Hansen et al., 2022; 2024; Sferrazza et al., 2024) and learning from demonstrations (Zhan et al., 2022; Hansen et al., 2023; Lancaster et al., 2024). They also discuss the limitations of existing approaches, such as the challenges of designing dense reward functions and the exploration problems in sparse reward settings. The paper builds upon TD-MPC2 (Hansen et al., 2022; 2024) as its backbone for model-based RL, and draws inspiration from MoDem (Hansen et al., 2023) for its bi-phase training scheme, clearly acknowledging these influences.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n- The paper addresses a significant challenge in reinforcement learning: solving long-horizon, multi-stage manipulation tasks with sparse 2. rewards.\n- The proposed method is data-efficient, requiring only a small number of demonstrations (as few as five).\n- The approach is evaluated on a diverse set of tasks, demonstrating its generalizability.\n- The bi-phasic training scheme is a clever way to leverage demonstrations for both initialization and online learning.\n\nWeaknesses:\n- The paper does not extensively discuss the limitations of the approach or potential failure cases.\n- While the method is shown to work with as few as five demonstrations, it's not clear how the performance scales with even fewer demonstrations or how it compares to methods that don't use demonstrations at all.\n- The computational complexity of the approach is not thoroughly discussed, which is important for practical applications.\n- The paper focuses primarily on simulation environments, and it's not clear how well the approach would transfer to real-world robotic systems.", "other_comments_or_suggestions": "- The paper would benefit from a more detailed discussion of the limitations of the approach and potential directions for future work.\n- A more thorough analysis of the computational requirements of the method would be valuable, particularly for real-time applications.\n- It would be interesting to see how the method performs with varying numbers of demonstrations, to better understand the trade-off between demonstration quantity and performance.\n- The paper could discuss more explicitly how the approach might be adapted for real-world robotic systems, addressing challenges such as sensor noise and actuation delays.", "questions_for_authors": "- How does the performance of DEMO3 scale with the number of demonstrations? Is there a minimum number of demonstrations required for the method to work effectively, and how does the performance improve with additional demonstrations?\n- The paper focuses on simulation environments. Have you explored how the approach might transfer to real-world robotic systems, and what additional challenges might arise in that context?\n- The method relies on stage indicators for multi-stage tasks. How sensitive is the approach to the definition of these stages, and how might it be extended to tasks where the stage structure is less clear or where the number of stages might vary between demonstrations?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}]

Author rebuttal:

Meta-review: The paper proposes demo-augmented reward, policy and world model learning. This framework exploits the coherent learning among these elements for efficient learning from visual inputs. The paper is well motivated, clear written and with comprehensive experiments. All three reviewers reach concensus that the paper is ready for publication. 

Please incorporate all the necessary revisions into the camera ready version, including:
- technical details / clarification
- limitation on failure cases / real-robot experiments
- ablations as per requested by reviewers in the rebuttal.
Decision: Accept (poster)