Subtask-Aware Visual Reward Learning from Segmented Demonstrations
Authors: Changyeon Kim
Minho Heo
Doohyun Lee
Honglak Lee
Jinwoo Shin
Joseph J Lim
Kimin Lee
Venue: ICLR 2025 Poster (2025)
Keywords: Reinforcement Learning
Reward Learning
Robotic Manipulation
Primary area: applications to robotics, autonomy, planning
TL;DR: We propose a novel reward learning framework utilizing action-free videos with minimal guidance for long-horizon complex robotic tasks.

Abstract: Reinforcement Learning (RL) agents have demonstrated their potential across various robotic tasks. However, they still heavily rely on human-engineered reward functions, requiring extensive trial-and-error and access to target behavior information, often unavailable in real-world settings. This paper introduces REDS: REward learning from Demonstration with Segmentations, a novel reward learning framework that leverages action-free videos with minimal supervision. Specifically, REDS employs video demonstrations segmented into subtasks from diverse sources and treats these segments as ground-truth rewards. We train a dense reward function conditioned on video segments and their corresponding subtasks to ensure alignment with ground-truth reward signals by minimizing the Equivalent-Policy Invariant Comparison distance. Additionally, we employ contrastive learning objectives to align video representations with subtasks, ensuring precise subtask inference during online interactions. Our experiments show that REDS significantly outperforms baseline methods on complex robotic manipulation tasks in Meta-World and more challenging real-world tasks, such as furniture assembly in FurnitureBench, with minimal human intervention. Moreover, REDS facilitates generalization to unseen tasks and robot embodiments, highlighting its potential for scalable deployment in diverse environments.

Published as a conference paper at ICLR 2025
SUBTASK-AWARE VISUAL REWARD LEARNING FROM SEGMENTED DEMONSTRATIONS
Changyeon Kim
Minho Heo
Doohyun Lee'
Honglak Lee Jinwoo Shin'
Joseph J. Lim Kimin Lee LG Al Research
KAIST University of Michigan
ABSTRACT
Reinforcement Learning (RL) agents have demonstrated their potential across var- ious robotic tasks. However, they still heavily rely human- eegaa-engineered reward functions, requiring extensive trial-and-error and access to target behavior infor- mation, often unavailable in real-world settings. This paper introduces REDS: RE- ward learning from Demonstration with Segmentations, a novel reward learning framework that leverages action-free videos with minimal supervision. Specifi- cally, REDS employs video demonstrations segmented into subtasks from diverse sources and treats these segments as ground-trute rewards. We train a dense re- ward function conditioned on video segments and their corresponding subtasks to ensure alignment with ground-truth reward signals by minimizing the Equivalent- Policy Invariant Comparison distance. Additionally, we employ contrastive learn- ing objectives to align video representations with subtasks, ensuring precise sub- task inference during online interactions. Our experiments show that REDS sig- nificantly outperforms baseline methods on complex robotic manipulation tasks in Meta-World World and more challenging real-wood wallwwrldd tasks, such furniture assembly in FurnitureBench, with minimal human intervention. Moreover, REDS facilitates generalization unseen tasks and robot embodiments, highlighting its potential for scalable deployment in diverse environments.
Meta-
INTRODUCTION
Reinforcement Learning (RI.) has demonstrated significant potential for training autonomous agents in various real-world robotic tasks, provided that appropriate reward functions are available (Levine al.. 2016; Gu al.. 2017: Andrychowicz However, reward engineering typically requires substantial al., 2023) and extensive task knowledge, often necessitating specialized instrumentation (e.g., motion trackers (Peng et al., 2020) or tactile sensors (Yuan et al., 2023)) or detailed information about target objects (James et al., 2020; Zhu et al., 2020; Yu et al., 2020; Mu et al.. 2021; Gu et al., 2023; Sferrazza et al., 2024), which are difficult to obtain in real-world settings. Learning reward functions from action-free videos has emerged as a promising alternative, as it avoids the need for detailed action annotations precise target behavior information, and video data can be easily collected from online sources (Soomro et al., 2012; Kay et al., 2017: Damen et al., 2018). Approaches in this domain include learning discriminators between video demonstrations and policy rollouts (Chen et al., 2021; Yang et al., 2024), training temporally aligned visual representations from large-scale video datasets (Sermanet et al., 2018: Zakka al., 2021; Kumar et al., 2022: Ma et al., 2023b;a) to estimate reward based on distance to a goal image, and using video prediction models to generate reward signals (Escontrela al., 2023; Huang et al., 2024).
al., 2020; Smith al., 2023; Handa et al., 2023). (Booth et al., 2023; Knox
Despite this progress, existing methods often struggle with long long-horizon, complex robotic tasks that involve multiple subtasks. These approaches typically fail provide context-aware reward signals, relying only on few consecutive frames or the final goal image without considering subsequent subtasks. For example, One Leg task (see Figure 2d) from FurnitureBench (Heo et al., 2023), prior methods often overemphasize the reward for picking up the leg while neglecting crucial steps such as inserting the leg into a hole and tightening it. Recent work (Mu et al., 2024) proposes a
"Equal advising.
Project page: https: ptttss:/cchanee //ww/////aaaaa.caca/a reds

---

Published as a conference paper at ICLR 2025
discriminator- based approach that treats complex tasks sequence of subtasks. However, it as- sumes that the environment provides explicit subtask identification, which often demands significant human intervention in real-world scenarios. Moreover, discriminator- based methods are known to be prone to mode collapse (Wang et al., 2017; Zolna et al., 2021) (please refer to Figure 11 for em- pirical evidence over prior work). Consequently, designing an effective visual reward function for real-world, long-horizon tasks remains an open problem.
Our approach To address the aforementioned limitations, we propose a novel reward learning framework, REDS: REward learning from Demonstration with Segmenlations, which infers subtask information from video segments and generates corresponding reward signals for each subtask. The key idea is to employ minimal supervision produce appropriate reward signals for intermediate subtask completion. Specifically, REDS utilizes expert demonstrations, where subtasks are anno- tated at each timestep by various sources (e.g., human annotators, code snippets, vision-language models; see the left figure of Figure 1). These annotations serve as ground-truth rewards. For train- ing, we introduce new objective function minimizing the Equivalent-Pole Policy Invariant Comparison (EPIC) (Gleave et al., 2021) between the learned reward function and the ground-truth rewards, guaranteeing theoretical upper bound on regret relative to the ground-truth reward function. Addi- tionally, to correctly infer the ongoing subtask in online interactions, we adopt contrastive learning objective to align video representations with task embeddings. In terms of architecture, our reward model is designed to capture temporal dependencies in video segments using transformers (Vaswani et al., 2017), leading to enhanced reward signal quality.
We find that REDS can generate appropriate reward signals to solve complex tasks by recognizing subtask structures, enabling the agent to efficiently explore and solve tasks through online interac- tions, using only expert demonstrations and subtask segmentations. Our experiments show that RI. agents trained with REDS achieve substantially improved sample efficiency compared to baseline methods on various robotic manipulation tasks from Meta-World (Yu et al., 2020). Additionally, we show that REDS can effectively train agents to perform long-horizon, complex furniture as- sembly tasks from FurnitureBench (Heo al., 2023) using real-world online RL. Moreover, REDS facilitates RL training in unseen environments involving new tasks and embodiments, which would otherwise require significant effort in prior reward-shaping methods.
Contributions We highlight the key contributions of our paper below:
We present a novel visual reward learning framework REDS: REward learning from Demon- stration with Segmentations, which can produce suitable reward signals aware of subtasks in long-horizon complex robotic manipulation tasks.
We show that REDS significantly outperforms baselines in training RL agents for robotic manip- ulation tasks in Meta-world, and even surpasses dense reward functions in some tasks.
We demonstrate that REDS can train real-world RL agents perform long-horizen complex furniture assembly tasks from FurnitureBench.
We demonstrate that our approach shows strong generalization across various unseen tasks, em- bodiments, and visual variations.
RELATED WORK
Reward learning from videos Learning from observations without expert actions has been promising research area because it does not require extensive instrumentation and allows for the easy collection of vast amounts of video from online sources. Notably, several studies have pro- posed methods for learning rewards directly from videos and using the signal to train RL agents. Previous work has been focused on learning a reward function by aligning video representations in temporal order (Sermanet et al., 2018: Zakka et al., 2021; Kumar et al., 2022) while others train re- ward function for expressing the progress of the agent towards the goal (Hartikainen et al., 2020; Lee et al., 2021; Yang et al., 2024). Most recent work (Escontrela et al., 2023) inspired by the success of video generative models (Yan et al., 2021; Ho et al., 2022) utilizes the likelihood of pre- pre-trained video prediction models reward. To effectively utilize video for long-horize tasks, propose a new reward model conditioned both on video segments and corresponding subtasks trained with subtask segmentations.
2

---

Published as a conference paper at ICLR 2025
1. Subtask segmentation on video
2. Reward learning
3. Online RL with learned reward
Segmented video /y
Predict subtask U/
REDS
Replay Buffer
2
3
Label ("
Subtask Instruction
Vidao Encoder
Subtask Embescer
(8, 5
box
Lift box
Put the box the table
Vi,K
Contrastive Loss
Agent
R
EPIC
Figure 1: Illustration of REDS. Our main idea to leverage expert demonstrations annotated with the ongoing subtask as the source of implicit reward signals (left). We train a reward model con- ditioned on video segments and corresponding subtasks with contrastive loss to attract video segments and corresponding subtask embeddings and 2) EPIC (Gleave et al., 2021) loss to generate reward equivalent to subtask segmentations (middle). In online RL, REDS infers the ongoing sub- task using only video segments at each timestep and computes the reward with that (right).
Inverse reinforcement learning Designing an informative reward function remains a long- standing challenge for training RL agents. achieve this, Inverse Reinforcement Learning (IRL) (Ng & Russell, 2000; Abbeel & Ng, 2004; Ziebart et al., 2008) aims to estimate the underlying reward function from expert demonstrations. Adversarial imitation learning (AIL) approaches (Ho & Ermon, 2016; Fu et al., 2018; Zolna et al., 2020; 2021; Mu et al., 2024) address this by training a discriminator network to discriminate transitions from expert data or policy rollouts and using the output from the discriminator ours is DrS (Mu et al., 2024), which also utilizes subtask information the multi-stage task. While DrS assumes that the information ongoing subtasks can be obtained from the environment during online interaction, our method has no such assumption, so can be applied in more general cases when the segmenting of the subtask is hard in automatic ways (e.g., Heo al. (2023)).
reward for training agents with RL. The most similar work to
3 PRELIMINARIES
Problem formulation We formulate a visual control task as a Markovian Decision Process (MDP) (Sutton & Barto, 2018). As single image observation is not sufficient for fully describ- ing the underlying state of the task, we use the set of consecutive past observations to approximate the current state following common practice (Mnih et al., 2015; Yarats al., 2021; 2022). Taking this into account, we define MDP as a tuple M (S, A, R. po. po.r). 7). S is a state space consisting of stack of K consecutive images,
A is an action space, R is the sparse reward function which outputs when the agent makes suc- cess; otherwise, p(s (s, a) is the transition function, Po is the initial state distribution, and 7 is the discount factor. The policy
S + ((A) is trained to maximize the expected sum of dis- counted rewards =. > 7¹ R(se, st, dt)]. Our goal is to find dense reward function R((s") only conditioned on visual observations, from which we can get the optimal policy n² for M.
EPIC Equivalent-Policy Invariant Comparison (EPIC) (Gleave et al., 2021) is a pseudometric for quantifying differences between different reward functions, which is designed to ensure the invari- ance on the equivalent set of reward functions inducing the same set of optimal policies. To this end, EPIC first canonicalizes potential shaping the reward function R with some arbitrary distribution Ds E A(S) over states S, which is to be invariant to potential shaping below
(R)) = R((s) Es-Ps (R(S) - R(S) - R(()) = R(s) Es-Ps [[(S)),
(I)
where S denotes set of batches independently sampled from the arbitrary distribution Ds, and the discount factor. EPIC is then defined by the Pearson distance between canonically shaped rewards scale-invariant manner:
De,Ds (RA, RB) = EL-De [Dp(Cos (s), Cp (s)),
(2)
We only consider the action-independdnt reward functions, and omit the prime notation on s' for the simplicity of notation.
3

---

Published as a conference paper at ICLR 2025
/ = is the Pearson distance
where is from the coverage distribution De. D)(X,Y) between two random variables X and Y, and p(X; () is the Pearson correlation between X and Y. Please refer to Gleave et al. (2021) for more details.
METHOD
REward learning from Demonstration with Segmentations, a visual reward learning framework designed for long-horizon tasks involving multiple subtasks. To gen- erate proper reward signals for solving intermediate subtasks, we utilize segmentations identifying ongoing subtasks in demonstrations. In Section 4.1, we explain our intuition and formal definitions behind subtask segmentation. Section 4.2 outlines the reward model architecture, and Section 4.3 describes our training objective. Finally, in Section 4.3, we elaborate on the details of training and inference of REDS. For an overview, see Figure 1.
This section presents REDS:
4.1 SUBTASK SEGMENTATION
The sparse reward function R provides feedback only on the overall success or failure of a task, which insufficient for guiding the agent through intermediate states. To address this, drawing inspiration from previous work on horgoooizonn robotic manipulation tasks (Di Palo & Johns, 2022; Mandlekar et al., 2023; Heo et al., 2023; Mu et al., 2024), we decompose a task into m object- centric subtasks, denoted as M = (U ..., }. Each subtask U; represents a distinct step in the task sequence and is based on the coordinate frame of a single target object. 2 and (ii) pulling the door to the goal position (which involves motion relative to the green sphere-shaped goal). This approach intuitive because humans naturally perceive tasks as sequences of discrete object interactions, be generally applied to different manipulation skills (e.g.. and-place,
and this assumption inserting) with diverse objects. Additionally, we provide text instructions X' = {x; that describe how to solve each subtask, which helps guide the agent more effectively.
To obtain subtask segmentations, we map each observation at timestep in the trajectory -r) to its corresponding subtask using segmentation function s O - U. Specifi- cally, z0 outputs the index of the ongoing subtask based on the observation at each timestep, with the output value increasing as the number of completed subtasks increases (refer the graph of 1o in the center of Figure 1). The function z can be derived from various sources such as code snippets based on domain knowledge (James & Davison, 2022; James et al., 2022: Mees et al., 2022), guidance from human teachers (Heo et al., 2023), or vision-language models (Zhang et al., 2024; Kou et al., 2024). In our experiments, we use the predefined codes in Meta-world and human annotators in Furniture to collect subtask segmentations. Note that these segmentations are only used dur- ing training: our framework is designed to automatically infer subtasks during online interactions without external annotations.
4.2 ARCHITECTURE
As mentioned in Section 1, previous reward learning methods generate rewards only by a single frame or consequent frames, not taking into account the order subtasks. To resolve the issue, we propose a new reward predictor nt = i((;; U)) conditioned on each subtask. To efficiently process visual observations, we first encode each image into www-dimeniional representations using pre- trained visual encoder E,. To capture temporal dependencies, these representations are processed through causal transformer (Vaswani et al., 2017). We add positional embeddings for each image in the sequence 8 and pass them through the transformer network, producing the output representation K.KK fu,-k-1, 1t-1, 2,} such that t-th output depends on input up to t. To embed subtask Us. we encode æ; with pre-trained text encoder E, and project it to a shallow MLP to earn e;- This design allows REDS generate rewards for unseen tasks when 24 and X' provided. (see Section 5.4 for supporting experiments). Finally, we concatenate a sequence of video representations V..K and subtask embedding e; Pi-K-1....; e;] and project it another shallow MLP f to obtain Rg
(s;; U?) = f (veKkee)). e.).
For instance, Door Open can be divided into (i) reaching the door handle (which involves motion relative the door handle.)
4

---

Published as a conference paper at ICLR 2025
4.3 REWARD MODELING
Reward equivariance with subtask segmentation Our key insight is that the subtask segmen- tation function & can be thought of as the ground reward function, providing implicit signals for solving intermediate tasks. To ensure our reward function induces the same of optimal poli- cies as ei we train minimize EPIC (Gleave al., 2021) distance between our reward model R. parameterized by 0 and v for all subtasks:
(0)
EPIC De.Ds (1gg..))
(3)
Progressive reward signal However, minimizing EPIC with v alone can lead to overfitting and the inability provide progressive signals within each subtask. To mitigate this issue, we propose an additional regularization term to enforce progressive reward signals. Inspired by previous work (Lee al., 2021: Hartikainen al., 2020; Wu al., 2021), we view the reward function as a progress indicator for each subtask, and we regularize the reward function output to be higher in later states of expert demonstration as follows:
Crgg (0) max (icc(Ro(sstjiiie
Ro(se: ((i))))
where j is randomly chosen from fixed set of values, and is hyperparameter. Note that apply this objective only for the expert demonstrations and not suboptimal demonstrations collected in iterative processes (please refer to Section 4.4).
Aligning video representation with subtask embeddings As the reward model lacks information about the ongoing subtasks in online interactions, it must infer the agent's 's current subtask. To achieve this, train the video representation to be closely aligned with the corresponding subtask embedding by adopting a contrastive learning objective to make the model select the appropriate subtask embedding only by the video segment.
Loont (0) -olog (u, (o(o)) (ie.)....)) simm ei)
(5)
where sim represents a cosine similarity. In summary, all components parameterized by 0 are jointly optimized to minimize our total training objective:
[()) (f) + (eeg()) Leoat (6).
4.4 TRAINING AND INFERENCE
Inference To compute the reward at timestep 4, REDS first infers the agent's current subtask from sequence of recent visual observations. Specifically, REDS selects the subtask index by choosing the subtask embedding e; that has the highest cosine similarity with the final output of the causal transformer, denoted vt. The final reward is then computed using video embedding and text embedding of the inferred subtask as Ro(s i U2) (VV eq). Please refer Appendix A for more details.
Training We outline the training procedure for REDS. First, we collect subtask segmentations from expert demonstrations, creating dataset D and use train the initial reward model, M However, reward models trained solely on expert data are susceptible to reward misspecifica- tion (Pan al., 2022). To address this, we iteratively collect suboptimal demonstrations and fine- tune the reward model using expert and suboptimal data. Unlike expert demonstrations, suboptimal demonstrations cover broader range of states and more diverse observations, making manual seg- mentation
and error- aroorprone. To reduce the burden human annotators, develop an automatic subtask inference procedure, avoiding the need for manual segmentation.
Before the iterative process, we compute similarity scores for all states in the expert demonstrations using the initial reward model M For each subtask U,, we calculate a threshold Tr, based on the similarity scores between the expert states and the corresponding instructions, ensuring Tv, represents the minimum similarity required for successful subtask completion. In each iteration i {1,...,)) nI,, proceed as follows:
5

---

Published as a conference paper at ICLR 2025
(a) Door Open (b) Peg Insert Side
(c) Sweep Into
(d) One Leg
Figure 2: Examples of visual observations used in our experiments. We consider a variety of robotic
manipulation tasks from Meta-world (Yu et al., 2020) and FurniturBennch (Heo et al., 2023).
Step (Suboptimal data collection): We train an RI. agent using the reward model M' and collect suboptimal demonstrations D Piplla from the agent's 's replay buffer.
Step 2 (Subtask inference for suboptimal data): For each timestep in the suboptimal trajectory, we infer the subtask index i using the same procedure as in inference and compute sim (v, e;). If the similarity falls below the threshold Ty, at any timestep, we mark the subtask as failed and assign the remaining timesteps to that subtask.
Step 3 (Fine-tuning): We fine-tune the reward model M using the combined dataset D = D' U Pieplay obtain Mr.
We use the final reward model M for downstream RL training.
5 EXPERIMENTS
We design our experiment to evaluate the effectiveness of REDS on providing useful reward signals in training various RL algorithms (Hafner et al., 2023; Kostrikov et al., 2022). We conduct extensive experiments robotic manipulation tasks from Meta-world (Yu et al., 2020) (see Section 5.1) in simulation and robotic furniture assembly tasks from FurnitureBench (Heo et al., 2023) (see Sec- tion 5.2) in the real-world. We also conduct in-depth analyses to validate the effectiveness of each component and how reward function aligns with subtask segmentations (see Section 5.5).
Implementation and training details We used the open-source pre- pre-trained CLIP (Radford al., 2021a) with ViT-B/16 architecture to encode images and subtask instructions for all experiments. We adopt a GPT (Radford et al., 2018) architecture with 3 layers and 8 heads for the causal transformer. To canonicalize our reward functions, we use the same 'D for both coverage distribution De and potential shaping distribution Ds. and we estimate the expectation over state distributions using a sample-based average over additional samples from D per sample. All models are trained with AdamW W (Loshchilov & Hutter, 2019) optimizer with a learning rate of 1 x 10 and a mini-batch size of 32. To ensure to visual distractions, we apply color jittering and random shift (Yarats et al., 2021) to visual observations in training REDS. Please refer to Appendix A for more details.
Baselines We consider the following baselines: (1) human- reward functions provided in the benchmark, (2) ORIL (Zolna al., 2020), an adversarial imitation learning (AIL) method trained only with offline demonstrations, (3) Rank2Reward (R2R) (Yang et al., 2024), an AlL method which trains a discriminator weighted with temporal ranking of video frames to reflect task progress. (4) VIPER (Escontrela et al., 2023), reward model utilizing likelihood from pre- trained video prediction model as a reward signal, and (5) DrS (Mu et al., 2024), an AIL method that assumes subtask information from the environment and trains separate discriminator for each subtask. We provide additional details on baselines in Appendix C.
5.1 META-WORLD EXPERIMENTS
Setup We first evaluate our method on 8 different visual robotic manipulation tasks from Meta- world (Yu et al., 2020). As a backbone algorithm, we use DreamerV3 (Hafner et al., 2023), a state-of-the-art visual model-based RL algorithm that learns from latent imaginary rollouts. For collecting subtask segmentations, we utilize a scripted teacher in simulation environments for scal- ability. Specifically, we use the predefined indicator for subtasks provided in the benchmark for all subtask segmentations (see Appendix D for the list of subtasks and corresponding text instructions for each task). We do not these indicators when training/evaluating RI. agents. For training

---

Published as a conference paper at ICLR 2025
REDS (Ours)
Humian-Engineered
ORIL
R2R
VIPER
DrS
Faucet Close
Drawar Open
Lever Pull
Door Open
Suspoea falm (%)
60
Sucocds Rote (%5
60
m
Suoess Rata (N)
60-
60
Rate
10
10
20
20
22
Eniirmmnn Step (+10))
Envirommnt Stop (x100)
Ervirnment Step (x10
Envitonmnt Stop [+110)
Coffee Pul
Peg Insert Side
Push
Sweep Info
Statteod Foom (N)
60
SUccoss Rate (3)
20.
800
62-
60
20
ty
Success Rat
80
40
20
20
200
0
16
10
20
30
11
Enviroment Step 1x
Emvrromeen Step (+10")
Envtronment Step 10
Emvironmen Step [+100]
Figure 3: Learning curves of D3eamerV3 (Hafner et al., 2023) agents trained with different reward functions for solving eight robotic manipulation tasks from Meta-world (Yu et al., 2020), measured by success rate (%). The solid line and shaded regions represent the mean and stratified bootstrap interval across runs.
REDS, we first collect subtask segmentations from 50 expert demonstrations for initial training and train D3eamerV3 agents for 100K environment steps with the initial reward model to collect subop- timal trajectories, which is used for fine-tuning. In evaluation, we measure the success rate averaged over 10 episodes in every 20K steps. Please refer to Appendix A for more details.
Results Figure 3 shows that REDS consistently improves the sample-effciincy of DreamerV3 agents by outperforming all baselines. While baselines exhibit non-zero success rates in simple tasks like Faucet Close, their performance significantly deteriorates in more complex tasks, such as Peg Insert Side. On the other hand, our method maintains non-zero success rates across all tasks and even surpasses human- auman-engineered reward functions in some tasks (e.g., Drawer Open, Push, Cof- fee Pull) without requiring task-specifi reward engineering. These results show that REDS effec- tively generates appropriate rewards for solving intermediate tasks by leveraging subtask -eegteetede demonstrations. A key advantage of REDS is that it relies solely visual observations for gener- ating rewards during online interaction, whereas DrS and human-engineered rewards require ad- ditional information from the environment, such as the position and reachability of target objects. This result underscores REDS's potential for application in environments where reward engineering challenging or additional sensory information is unavailable.
5.2 FURNITUREBENCH EXPERIMENTS
Setup We further evaluate our method on real-world furniture assembly tasks from Fur- nitureBench (Heo et al., 2023), specifically focusing One Leg assembly. This task involves a sequence of complex subtasks such as picking up, inserting, and screw- ing (see Figure 2d). For training REDS, we use 300 expert demonstrations with subtask segmentations provided by BennitureBench, along with additional 200 rollouts from IQL (Kostrikov et al., 2022) policy trained with expert demonstrations in a single training iteration. To prevent misleading reward signals stemming from visual occlusions, we utilize visual observations from the front camera and wrist cameras in training REDS. For downstream RL, we first train offline RL agents using 300 expert demonstrations labeled with each reward model, followed by online fine-tuning to assess improve- ments. For baselines, we compare against VIPER and DrS. We emphasize that our method enables fully autonomous training in online RL sessions, in contrast to DrS. which relies on subtask indi- cator provided by humans. In our DrS experiments, subtasks were manually identified by a human. We measure the average number of completed subtasks over 20 rollouts for evaluation. We provide more details in Appendix A.
Table I: Online fine-tuning results of IQL agents in One Leg from FurnitureBench. We report the initial performance after offline RL (left) and the perfor- mance after 150 episodes of online RL (right).
Completed Subtasks (Offline Online)
Method
Expert Demos
Sparso (Ofllin)) (Hico et al., 2023) VIPER DrS REDS (Ours)
500
1.8
300
1.10 1.25
300
1.05 1.10
300
1.10 - 2.45
7

---

Published as a conference paper at ICLR 2025
2
REDS Hy man engghoordd
-125
-0.50
-125
Huumon
-150
175
(a) Door Open
(b) One Leg
Figure 4: Qualitative results REDS in Door Open in Meta-world (Yu et al., 2020) and One Leg from FurnitureBench (Heo et al., 2023). We observe that REDS produces suitable reward signals aligned with ground-trun reward functions by predicting ongoing subtasks effectively and providing progressive reward signals.
As shown in Table 1, REDS achieves significant performance improvements through on- line fine-tuning, whereas the improvements from baselines are marginal. These results indicate that our method produces informative signals for solving a sequence of subtasks, while baselines either fail to provide context-aware signals or dense rewards for better exploration (see Appendix F for qualitative examples). Moreover, we note that our method outperforms the IQL trained with 500 ex- pert demonstrations, achieving score of 2.45 compared to 1.8 reported by FurnitureBench, despite using only 300 expert demonstrations. Considering that REDS does not require additional human interventions beyond resetting the environment, these results highlight the potential to extend our approach to wider range of real-world robotics tasks.
Results
5.3
ALIGNMENT WITH GROUND-TRUTH REWARDS
EPIC measurement To quantitatively validate the alignment of our method with ground-truth reward functions, we mea- sure the EPIC distance with a set of unseen demonstrations during training. Specifi- cally, use rollouts from the reference policy trained with expert demonstrations for state distribution. In Table 2, we observe that REDS exhibits significantly lower EPIC distance than baselines across all tasks. Particularly, the difference be- tween REDS and baselines more pronounced in complex tasks like One Leg. This result consis- tently supports the empirical findings from previous sections.
Table 2: EPIC (Gleave al., 2021) distance (lower is better) between learned reward functions and hand- engineered reward functions (Meta-world) / subtask
segmentations (FurnitureBench) in unseen data.
Task
VIPER
R2R
ORIL
REDS (Ours)
Meta world Door Open
0.5934
0.5649
0.7071
0.4913
Metn-world Push
0.6144
0.6838
0.7073
0.5381
Meta-world Pog Insert Side
0.5974
0.5806
0.6989
0.4674
Meta-world Sweep Into
0.6248
0.6413
0.7001
0.4673
FurnitueBench One Leg
0.7035
0.6001
0,7014
0.0713
Qualitative analysis We provide the graph of computed rewards from REDS in Figure 4. We observe that REDS can induce suitable reward signals aligned with ground-truth reward functions. For example, REDS provides subtask-aware signals in transition states (e.g., between 2 and and between 4 and 5) and generates progressive reward signals throughout each subtask. Please refer to Appendix F for the extensive comparison between REDS and baselines.
5.4 GENERALIZATION CAPABILITIES
Transfer to unseen tasks As mentioned in Section 4.2, our model can be applied reward function unseen tasks. To validate this, we conduct additional experiments by training REDS with segmentation data from 3 tasks (Door Open, Drawer Open/Close) and using the reward model to train RL agents in two unseen tasks. In Door Close, we aim to validate that REDS can provide informative signals for Close, we aim to determine whether REDS can provide suitable reward signals for familiar behaviors (closing) with an unseen object (window). In evaluation, change the text instruction following
task involving previously seen object and behaviors. In Window
8

---

Published as a conference paper at ICLR 2025
Door Close
Window Close
Oper
Draws Clace
Own
80
Success Rale (%)
50
Success 00
60
40
Hom ne ginweee (Onse)) 1ED0
8
IC000
(x10")
Environment Slep (x10")
Figure 5: We train REDS with different tasks from Meta-world (Yu et al., 2020) and this model to train RL agents in 2 unseen tasks (left). We present learning curves on Door Close (center) and Window Close (right), as measured by success rate (%). The solid line and shaded regions represent the mean and stratified bootstrap interval across runs.
100-
Success Rate (N6
80-
80-
60
20
pos:
0.71
Mebe pos:
20
Environment Step 1-10
(a) Original
(b) Examples of visual distractions
(c) Learing curve
Figure 6: We provide visual observations from (a) the original environment and (b) unseen environ- ments with visual distractions used our experiments Section 5.4.
the target object (as shown in Table D), and we do not fine-tune the reward model. Figure shows that REDS provides effective reward signals unseen tasks and achieves comparable even better RL performance than REDS trained on the target task. This result demonstrates that REDS can be applied to RL training in unseen tasks that share properties with training tasks.
Robustness to visual distractions To prove the robust performance REDS against visual dis- tractions, we train RL agents with reward model in new Meta-world environments incorporating visual distractions, such as varying light and table positions following Xie et al. (2024) (see Fig- ure 6b). Note that the reward model was trained using demonstrations only from the original envi- ronment. As Figure 6e shows, REDS can generate robust reward signals despite visual distractions and train RI agents solve the task effectively.
Transfer to unseen embodiments Since our framework lever- ages only action-free video data, we hypothesize that transferring to other robot embodiments with similar DoF' is feasible. To sup- port this claim, we train REDS with demonstrations of the Franka Panda Arm and then compute the reward of an unseen demonstra- tion of the Sawyer Arm in Take Umbrella Out of Stand from RL- Bench (James et al., 2020). Figure shows that REDS generates informative reward signals even with the unseen embodiment. For instance, REDS can capture the behavior of taking the umbrella out of the stand, as indicated by the increased reward signals between 6 and 7. Additionally, Figure shows that REDS trained only with the Panda Arm can be used to train downstream RL agents in the environment with the Sawyer Arm.
Figure 7: Learning curve for DreamerV3 agents in envi- ronments of the Sawyer Arm.
Furmaeenn neaerd (Cnaca)
Success Pate (%))
201
0
Environment Step (x10
5.5 ABLATION STUDIES
Effeet of training objectives We investigate the effect of each training objective Figure 9a. Specifically, we compare REDS with 1) baseline trained with regression subtask segmentation instead of EPIC loss LEPICS 2) a baseline that utilizes only video representations without subtask embeddings, and 3) baseline trained without the regularization loss Lee. We observe that RL performance significantly degrades without each component, implying that our losses synergistically improve reward quality.

---

Published as a conference paper at ICLR 2025
(a) (Train) Panda
(b) (Unseen) Sawyer
Figure 8: Qualitative results of REDS with different robot embodiments. REDS was trained using demonstrations from the Panda Arm and evaluated on an unseen demonstration from the Sawyer Arm in Take Umbrella Out of Stand from RLBench (James et al., 2020). We visualize several frames above the graph and mark them with a diamond symbol.
RODT PoR Term
Ricoe (retting-))
Sucome Ratts 2N1)
80
pe:n6 aiceeates.... X1
Sucoee Plare (46)
c0
8
M/
20
0
Enshrommet Step
Endirommen Step (=10))
Endereeenn Step (+10)
Endrrrment Step (-10
(a) Training objectives
(b) Architecture
(c) Fine-tuning
(d) Expert demonstrations
Figure 9: Learning curves for two Meta-world (Yu et al., 2020) robotic manipulation tasks, measured by success rate (%), to examine the effects of (a) training objectives, (b) architecture, (c) fine-tuning, and (d) the number of expert demonstrations. The solid line and shaded regions show the mean and stratified bootstrap interval across runs.
Effect of architecture To verify the design choice proposed in Section 4.2, we compare REDS with 1) a baseline using CNN for encoding images instead of pre-trained visual representations (PVR) and 2) a baseline simply concatenating pre-trained visual representations without causal transformer. Figure 9b shows that both baselines show worse performance compared ours. No- tably, detaching a causal transformer significantly degrades RI. performance, implying that temporal information essential for providing suitable reward signals robotic manipulation.
Effect of fine-tuning In Figure 9c, we compare REDS trained only with the expert demonstrations in the initial phase to REDS fine-tuned with additional suboptimal demonstrations as described in Section 4.4. REDS shows improved RL performance when trained with additional suboptimal demonstrations, indicating that the coverage of state distribution impacts the reward quality. Further investigation on how to efficiently collect suboptimal demonstrations to enhance the performance of leamed reward function is promising future direction.
Effect of the number of expert demonstrations We investigate the effect of the number of expert demonstrations by measuring the RL performance of DreamerV3 agents with REDS trained with different numbers of expert demonstrations in 2 tasks (Door Open, Drawer Open) from Meta-world. Figure 9d shows that the agents' RL performance positively correlates with the number of expert demonstrations trained for reward learning.
6 CONCLUSION
We proposed REDS, a visual reward learning framework considering subtasks by utilizing subtask segmentation. Our main contribution is based on proposing new reward model leveraging minimal domain knowledge as a ground-truth reward function. Our approach is generally applicable and does not require any additional instrumentations in online interactions. We believe REDS will sig- nificantly alleviate the burden of reward engineering and facilitate the application of RL to a broader range of real-world robotic tasks.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This research is supported by Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (RS-2022-1122955, Self-directed AI Agents with Problem-solving Capability; RS-2019-11190099 Artificial Intelligence Graduate School Program (KAIST); No. RS-2024-005092279, Global Al Frontier Lab).
LIMITATION AND FUTURE DIRECTIONS
One limitation of our work is that assume the knowledge of the object-centric subtasks in lask. For automating subtask definition and segmentation in new tasks and new domains other than robotic manipulations, investigating the planning and reasoning capabilities of pre-trained Multi- modal Large Language Model (MLLM) (Park et al., 2023; Honerkamp et al., 2024; Zawalski et al., 2024: Shah et al., 2024; Liu al., 2024) would be an intriguing research direction.
Additionally, the performance of REDS relies on pre- representations trained with natural im- age/text data for encoding videos and subtask instructions. Although REDS proves its effectiveness in various robotic manipulation tasks, we observe that REDS struggles to distinguish subtle changes (e.g., screwing the leg in One Leg) even with pre-trained representations trained on ego-centric mo- tion videos (Ma et al., 2023a). We believe that the quality of rewards can be further improved by utilizing 1) pre-trained representations with large-scale data with diverse robotic tasks (Padalkar al., 2023; Khazatsky et al., 2024) and 2) representations trained with objectives considering affor- dances (Bahl et al., 2023) or object-centric methods (Devin et al., 2018).
Furthermore, there is room for improvement to enhance generalization and robustness. Although our experiments are designed to evaluate generalization in unseen environments, they may face challenges in ant-of-distribution environments, such significant changes in the background or camera angles. Future work could address these challenges through data augmentation or domain adaptation techniques. while our contrastive learning objective currently focuses minimizing dis- tances between relevant video and text embeddings, the reward model may generate inappropriate reward signals for semantically different subtasks that share similar video content and text instruc- tions. Incorporating loss term to maximize distances between irrelevant embeddings could further improve robustness in tasks with similar subtasks, and we will explore this enhancement future work.
Finally, the number of expert demonstrations and the number of iterations for fine-tuning REDS are determined by empirical trials. Investigating how to collect failure demonstrations to mitigate reward misspecification efficiently an interesting future direction.
ETHIC STATEMENT
Video demonstrations and subtask segmentations used in the experiments were sourced from pub- licly available benchmarks (Meta-world, RLBench, FurnitureBench), ensuring no personal or sen- sitive information is involved. Potential risks could arise when training and deploying RL agents directly in real-world scenarios, particularly in human-robot interactions. Ensuring the safety and reliability of these agents before deployment is essential to prevent harm.
REPRODUCIBILITY STATEMENT
For the reproducibility of REDS, we have provided detailed explanation of implementation details and experimental setups in Section 4.4, Section and Appendix A. In addition, further facilitate the reproduction, we release the open-sourced implementation through the project website.
REFERENCES
Pieter Abbeel and Andrew Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference Machine Learning, 2004.

---

Published as a conference paper at ICLR 2025
Ademi Adeniji, Amber Xie, and Pieter Abbeel. Skill-based reinforcement learning with intrinsic reward matching. arXiv preprint drXiv:2210.07426, 2022.
OpenAI: Marcin Andrychowicz, Bowen Baker, Maciek Chociej. Rafal Jozefowicz, Bob McGrew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, et al. Learning dexterous in in-hand manipulation. International Journal of Robotics Research, 2020.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 10..1.000.04555 2016.
Shikhar Bahl, Russell Mendonca, Lili Chen, Unnat Jain, and Deepak Pathak. Affordances from human videos versatile representation for robotics. In IEEE Conference Computer Vision and Puttern Recognition, 2023.
Philip Ball, Laura Smith, llya Kostrikov, and Sergey Levine. Efficient online reinforcement learn- ing with offline data. In International Conference on Machine Learning, 2023.
Serena Booth, W Bradley Knox, Julie Shah, Scott Niekum, Peter Stone, and Alessandro Allievi. The perils of tril--and-error reward design: misdesign through overfitting and invalid task speci- fications. In AAAI Conference on Artificial Intelligence, 2023.
Annie S Chen, Suraj Nair, and Chelsea Finn. Learning generalizable robotic reward functions from" in-the-wild" human videos. In Robotics: Science and Systems, 2021.
Dima Damen, Hazel Doughty, Giovanni Maria Farinella, Sanja Fidler, Antonino Furnari, Evangelos Kazakos, Davide Moltisanti, Jonathan Munro, Toby Perrett, Will Price, and Michael Wray. Scal- ing egocentric vision: The epic-kitchens dataset. In European Conference on Computer Vision, 2018.
Coline Devin, Pieter Abbeel, Trevor Darrell, and Sergey Levine. Deep object-centric representations for generalizable robot learning. IEEE International Conference on Robotics and Automation, 2018.
Norman Di Palo and Edward Johns. Learning multi-stage tasks with one demonstration via self- replay. In Conference on Robot Learning. 2022.
Alejandro Escontrela, Ademi Adeniji, Wilson Yan, Ajay Jain, Xue Bin Peng, Ken Goldberg, Young- woon Lee, Danijar Hafner, and Pieter Abbeel. Video prediction models as rewards for reinforce- ment learning. In Conference on Neural Information Processing Systems, 2023.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse rein- forcement learning. In International Conference on Learning Representations, 2018.
Adam Gleave, Michael D Dennis, Shane Legg, Stuart Russell, and Jan Leike. Quantifying differ- ences in reward functions. In International Conference on Learning Representations, 2021.
Jiayuan Gu, Fanbo Xiang, Xuanlin Li, Zhan Ling, Xiqiang Liu, Tongzhou Mu, Yihe Tang, Stone Tao, Xinyue Wei, Yunchao Yao, Xiaodi Yuan, Pengwei Xie, Zhiao Huang, Rui Chen, and Hao Su. Maniskil12: A unified benchmark for generalizable manipulation skills. In International Conference on Learning Representations, 2023.
Shixiang Gu, Ethan Holly, Timothy Lillicrap, and Sergey Levine. Deep reinforcement learning for robotic manipulation with asynchronous off-policy updates. In IEEE International Conference on Robotics and Automation, 2017.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint ar11::2301.04104 2023.
Ankur Handa, Arthur Allshire, Viktor Makoviychuk, Aleksei Petrenko, Ritvik Singh, Jingzhou Liu, Denys Makoviichuk, Karl Van Wyk, Alexander Zhurkevich, Balakumar Sundaralingam, et al. Dextreme: Transfer of agile in hand manipulation from simulation to reality. IEÉE Interna- fional Conference on Robotics and Automation, 2023.
12

---

Published as a conference paper at ICLR 2025
Kristian Hartikainen, Xinyang Geng, Tuomas Haarnoja, and Sergey Levine. Dynamical distance learning for semi-supervised and unsupervised skill discovery. In International Conference on Learning Representations, 2020.
Minho Heo, Youngwoon Lee, Doohyun Lee, and Joseph J. Lim. Fumnturebench: Reproducible real-world benchmark for long-horizon complex manipulation. In Robolics; Science and Systems, 2023.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Conference on Neural Information Processing Systems, 2016.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Conference on Neural Information Processing Systems, 2022.
Daniel Honerkamp, Martin Büchner, Fabien Despinoy, Tim Welschehold, and Abhinav Valada. Language-grounded dynamic tion. IEEE Robotics and Automation Letters, 2024.
graphs for interactive object search with mobile manipula-
Tao Huang, Guangqi Jiang, Yanjie Ze. and Huazhe Xu. Diffusion reward: Learning rewards via conditional video diffusion. In European Conference on Computer Vision, 2024.
Stephen James and Andrew J Davison. Q-attention: Enabling efficient learning for vision-based robotic manipulation. IEEE Robotics and Automation Letters, 2022.
Stephen James, Zicong Ma, David Rovick Arrojo, and Andrew Davison. Rlbench: The robot learning benchmark & learning environment. IEEE Robotics and Automation Letters, 2020.
Stephen James, Kentaro Wada, Tristan Laidlow, and Andrew J Davison. Coarse-to-fine q-attention: Efficcint learning for visual robotic manipulation via discretisation. In IEEE Conference on Com- puter Vision and Pattern Recognition, 2022.
Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijaya- narasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, et al. The kinetics human action video dataset. arXiv preprint (.XXvv/705.069550 2017.
Alexander Khazatsky, Karl Pertsch, Suraj Nair, Ashwin Balakrishna, Sudeep Dasari, Siddharth Karamcheti, Soroush Nasiriany, Mohan Kumar Srirama, Lawrence Yunliang Chen, Kirsty El- lis, et al. Droid: A large-scale in wi-the-wild robot manipulation dataset. arXiv preprint arXiv:2403. 12945, 2024.
W Bradley Knox, Alessandro Allievi, Holger Banzhaf, Felix Schmitt, and Peter Stone. Reward (mis) design for autonomous driving. Artificial Intelligence, 2023. 2023.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. In International Conference on Learning Representations, 2022.
Longxin Kou, Fei Ni, Yan Zheng, Jinyi Liu, Yifu Yuan, Zibin Dong, and HAO Jianye. Kisa: A unified keyframe identifier and skill annotator for long-horizon robotics demonstrations. In Inter- national Conference Machine Learning, 2024.
Sateesh Kumar, Jonathan Zamora, Nicklas Hansen, Rishabh Jangir, and Xiaolong Wang. Graph inverse reinforcement learning from diverse videos. In Conference on Robot Learning. PMLR, 2022.
Youngwoon Lee, Andrew Szot, Shao-Hua Sun, and Joseph J Lim. Generalizable imitation learning from observation via inferring goal proximity. In Conference on Neural Information Processing Systems, 2021.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuo- motor policies. Journal of Machine Learning Research, 2016.
Xinran Liang, Katherine Shu, Kimin Lee, and Pieter Abbeel. Reward uncertainty for exploration in preference-based reinforcement learning. In International Conference on Learning Representa- lions, 2022.
13

---

Published as a conference paper at ICLR 2025
Fangehen Liu, Kuan Fang, Pieter Abbeel, and Sergey Levine. MOKA: Open- vocabuocay robotic ma- nipulation through mark-based visual prompting. In IEEE International Conference on Robotics and Automation, 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- ence on Learning Reprexentations, 2019.
Yecheng Jason Ma, William Liang, Vaidehi Som, Vikash Kumar, Amy Zhang, Osbert Bastani, and Dinesh Jayaraman. Liv: Language-image representations and rewards for robotic control. In International Conference Machine Learning, 2023a.
Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. VIP: Towards universal visual reward and representation via value-implicit pre-training. In International Conference on Learning Representations, 2023b.
Ajay Mandlekar, Soroush Nasiriany, Bowen Wen, Iretiayo Akinola, Yashraj Narang, Linxi Fan, Yuke Zhu, and Dieter Fox. Mimicgen: A data generation system for scalable robot learning using human demonstrations. In Conference on Robot Learning, 2023.
Oier Mees, Lukas Hermann, Erick Rosete-Beas, and Wolfram Burgard. Calvin: A benchmark for language-conditionee policy learning for long-horizon robot manipulation tasks. IEEE Robotics and Automation Letters, 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 2015.
Tongzhou Mu, Zhan Ling, Fanbo Xiang, Derek Cathera Yang. Xuanlin Li, Stone Tao, Zhiao Huang, Zhiwei Jia, and Hao Su. Maniskill: Generalizable manipulation skill benchmark with large- scale demonstrations. In Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Tongzhou Mu, Minghua Liu, and Hao Su. Drs: Learning reusable dense rewards for multi-stage tasks. In International Conference on Learning Representations, 2024.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: univer- sal visual representation for robot manipulation. In Conference Robot Learning, 2022.
Andrew Y Ng and Stuart Russell. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning, 2000.
Abhishek Padalkar, Acorn Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Anthony Brohan, et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint 10t.:2310.8886.. 2023.
Alexander Pan, Kush Bhatia, and Jacob Steinhardt. The effects of reward misspeciffcaatio: Mapping and mitigating misaligned models. In International Conference on Learning Representations, 2022.
Joon Sung Park, Joseph O'Brien, Carrie Jun Cai, Meredith Ringel Morris, Percy Liang, and Michael $ Bernstein. Generative agents: Interactive simulacra of human behavior. In Proceedings of the 36th annual acm symposium on user interface software and technology, pp. 1-22, 2023.
Xue Bin Peng. Erwin Coumans, Tingnan Zhang, Tsang- Wei Lee, Jie Tan, and Sergey I Levine. Learn- ing agile robotic locomotion skills by imitating animals. arXiv preprint arXiv:2004.0078. 2020.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language under- standing by generative pre- pre-training. 2018.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, 2021a.
14

---

Published as a conference paper at ICLR 2025
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference Machine Learning.
PMIR, 2021b.
Juan Rocamonde, Victoriano Montesinos, Elvis Nava, Ethan Perez, and David Lindner. Vision- language models are zero-shot reward models for reinforcement learning. In International Con- ference Learning Representations, 2024.
Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control. In Conference Robot Learning. PMLR, 2022.
Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain. Time-contrastive networks: Self-supervised learning from video. In IEEE International Conference on Robotics and Automation, 2018.
Carmelo Sferrazza, Dun-Ming Huang, Xingyu Lin, Youngwoon Lee, and Pieter Abbeel. Humanoid- bench: Simulated humanoid benchmark for whole-body locomotion and manipulation. arXiv preprint arXiv:2403. /0506, 2024.
Rutav Shah, Albert Yu, Yifeng Zhu, Yuke Zhu, and Roberto Martin-Martin. Bumble: Unifying reasoning and acting with vision-language models for building- wide mobile manipulation. arXiv preprint arXiv:2410.06237, 2024.
Lucy Xiaoyang Shi, Archit Sharma, Tony Z Zhao, and Chelsea Finn. Waypoint-based imitation learning for robotic manipulation. In Conference on Robot Learning, 2023.
Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Perceiver-actor: A multi-task transformer for robotic manipulation. In Conference on Robot Learning. 2022.
Joar Max Viktor Skalse, Lucy Farnik, Sumeet Ramesh Motwani, Erik Jenner, Adam Gleave, and Alessandro Abate. STARC: general framework for quantifying differences between reward functions. In International Conference on Learning Representations, 2024.
Laura Smith, Ilya Kostrikov, and Sergey Levine. walk in the park:
Learning to walk in 20 minutes with model-free reinforcement learning. In Robotics; Science and Systems, 2023.
Khurram Soomro, Amir Roshan Zamir, and Mubarak Shah. Ucf101: dataset of 101 human actions classes from videos in the wild. arXiv preprint arXiv: 1212.0402. 2012.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Conference on Neural In- formation Processing Systems, 2017.
Ziyu Wang, Josh S Merel, Scott E Reed, Nando de Freitas, Gregory Wayne, and Nicolas Heess. Robust imitation of diverse behaviors. In Conference on Neural Information Processing Systems, 2017.
Zheng Wu, Wenzhao Lian, Vaibhav Unhelkar, Masayoshi Tomizuka, and Stefan Schaal. Learning dense rewards for contact-rich manipulation tasks. IEEE International Conference on Robotics and Automation, 2021.
Blake Wulfe, Logan Michael Ellis, Jean Mercat, Rowan Thomas McAllister, and Adrien Gaidon. Dynamics-aware comparison of learned reward functions. In International Conference Learn- ing Representations, 2022.
Annie Xie, Lisa Lee, Ted Xiao, and Chelsea Finn. Decomposing the generalization gap in imitation learning for visual robotic manipulation. IEEE International Conference on Robotics and Automation, 2024.
15

---

Published as a conference paper at ICLR 2025
Danfei Xu and Misha Denil. desitive-unlabele reward learning. In Conference on Robot Learning, 2021.
Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. Videogpt: Video generation using vq-vae and transformers. arXiv preprint arXiv:2104. 10157, 2021.
Daniel Yang, Davin Tjia, Jacob Berg, Dima Damen, Pulkit Agrawal, and Abhishek Gupta. Rank2reward: Learning shaped reward functions from passive video. In IEEE International Con- ference on Robotics and Automation, 2024.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference Learning Representa- tions, 2021.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous con- trol: Improved data-augmented reinforcement learning. International Conference on Machine Learming, 2022.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: benchmark and evaluation for multi-task and meta reinforcement learning. In Conference Robot Learning. 2020.
Ying Yuan, Haichuan Che, Yuzhe Qin, Binghao Huang, Zhao-Heng Yin, Kang-Won Won Lee, Yi Wu, Soo-Chul Lim, and Xiaolong Wang. Robot synesthesia: In-hand manipulation with visuotactile sensing. arXiv preprint arXiv:2312.01855 2023.
Kevin Zakka, Andy Zeng, Pete Florence, Jonathan Tompson, Jeannette Bohg, and Debidatta Dwibedi. Xirl: Cross-emboding inverse reinforcement learning. In Conference on Robot Learning, 2021.
Michal Zawalski, William Chen, Karl Pertsch, Oier Mees, Chelsea Finn, and Sergey Levine. Robotic control via embodied
reasoning. In Conference Robot Learning, 2024.
Zichen Zhang, Yunshuang Li, Osbert Bastani, Abhishek Gupta, Dinesh Jayaraman, Yecheng Jason Ma, and Luca Weihs. Universal visual decomposer: Long-horizon manipulation made easy, In IEEE International Conference Robotics and Automation, 2024.
Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martín-Martín, Abhishek Joshi, Soroush Nasiri- any, and Yifeng Zhu. robosuite: A modular simulation framework and benchmark for robot learning. arXiv preprint arXiv:2009. 12293, 2020.
Brian Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In AAAI Conference Artificial Intelligence, 2008
Konrad Zolna, Alexander Novikov, Ksenia Konyushkova, Caglar Gulcehre, Ziyu Wang, Yusuf Ay- tar. Misha Denil, Nando de Freitas, and Scott Reed. Offline learning from demonstrations and unlabeled experience. In Conference on Neural Information Processing Systems, 2020.
Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarejo, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang. Task-relevant adversarial imitation learn- ing. In Conference on Robot Learning, 2021.
16

---

Published as a conference paper at ICLR 2025
A EXPERIMENT DETAILS
Training and inference details To ensure robustness against visual changes, we apply data aug- mentations, including random shifting (Yarats et al., 2021; 2022) and color jittering, For optimiza- tion, we train REDS with AdamW (Loshchilov & Hutter, 2019) optimizer with a learning rate of 1 x 10-4, weight decay of 2 10-2, and cosine decay schedule for adjusting the training learning rate. We apply a warm-up scheduling for the initial 500 gradient steps starting from learning rate of O. Note that the parameters for CLIP visual/text encoders have not been updated. For training downstream RL agents, we normalize the reward by dividing it by the maximum value observed in the expert demonstrations. We report the hyperparameters used in our experiments in Table For both coverage distribution De and potential shaping distribution Ds. we use the same dataset with subtask segmentations DD, unlike prior work dealing with arbitrarily random distributions because of the absence of subtask segmentations. To canonicalize our reward functions, we estimate the expectation over state distributions using a sample-based average over 8 additional samples from D per sample. To prevent false positive cases in predicting subtasks in online interactions, we add mar- gins to similarity scores inversely proportional to the subtasks in online interactions. Specifically, we infer the subtask as follows:
argmax (sim( v% e,) + 17 (ls - i)),
where 1y is a hyperparameter for the margin between subtasks. For each subtask U. we compute similarity scores between the visual observations within the subtask from expert demonstrations and their corresponding instructions. The threshold TV, is set to the 75th percentile of these scores to account for demonstration variability while capturing the most relevant matches. Please refer to Figure 12c for supporting experiments.
Table 3: Hyperparameters of REDS used in our experiments.
Hyperparameter
Value
Batch size
32 (Mela-world, RLBench), 8 (FurnitureBench)
Training steps
5000
Learning rate
0.0001
Optimizer
Adam W (Loshchilov Hutter, 2019)
Optimizer momentum
81 - 0.9, 32 0.999
Weight decay
0.02
Learning rate decay
Linear warmup and cosine decay
Warmup steps
500
Context length
4
Causal transformer size
3 layers, 8 heads, 512 units
EPIC canonical samples
8
for progressive reward signal
0.05
for inferring subtasks
0.01 (Meta-world, RIBBench), 0.05 (FumitureBench)
number of training iterations
2 (orla-world, RLBench), (FumitureBench)
Meta-world experiments We use visual observations of 64 x 64 x 3. To consistently use a single camera viewpoint over all tasks, we the modified version of the corner2 viewpoint as suggested by Seo et al. (2022). Expert demonstrations for each task are collected using scripted policies publicly released in the benchmark. We and set the maximal episode length as 250 for all Meta-world tasks. For downstream RL, we use the implementation of DreamerV3 from VIPER*. We report the hyperparameters of DreamerV3 agents used in our experiments in Table 4. Unless otherwise specified, we use the same set of hyperparameters as VIPER. For all ablation experiments (Figure 6, 9, 12), we report results in Door Open and Drawer Open.
action repeat of 2 to accelerate training
RLBench experiments For training both reward models and downstream RL. agents, we utilize 64 x 64 x 3 RGB observations from the front camera and wrist camera. For downstream RL, we don't use any expert demonstrations, and we use the same set of hyperparameters as VIPER.
httpps://oo thub. www./////mootteeeeeeeeoeoe a/ a/viperrrr
17

---

Published as a conference paper at ICLR 2025
Table 4: Hyperparameters of DreamerV3 (Es- contrela et al., 2023) used in Meta-world exper- iments.
Table 5: Hyperparameters of IQL (Kostrikov et al., 2022) used in FurnitureBench experi- ments.
Hyperparameter
Value
General
Hyperparameter
Value
Replay Capacity (FIFO)
x 10
Leaminng rate
3 x 10 256
Start learning (prefill)
5000
Batch size
Policy hidden units
(512, 256, 256)
MLP size
x 512
Critic/value # hidden units
(512, 256, 256)
Image encoder
R3M (Nair al., 2022)
World Model
Discount factor (t)
0,996
RSSM size
512
Expectile (r)
0.8
Inverse Temperature (33
10.0
Base CNN channels
32
Codes per latent
32
BurnitureBench experiments We the implementation of IQL. from FurnitureBench 4 for experiments. We utilize 224 x 224 x 3 RGB observations from the front camera and wrist cam- cras, along with proprioceptive states, represent the current state. We encode each image with pre-trained R3M (Nair et al., 2022) for visual observations. Following Kostrikov et al. (2022), we first run offline RL for 1M gradient steps, then continue training while collecting environment in- teraction data, adding it to the replay buffer, and repeating this process for 150 episodes. Before online fine-tuning, we pre-fill the replay buffer with 10 rollouts from the pre-trained IQL policy. We adopt techniques from RLPD (Ball et al., 2023) for efficient offline-to-online RL training. Specif- ically, we sample 50% the data from the replay buffer and the remaining 50% from the offline data buffer containing 300 expert demonstrations. We also apply LayerNorm (Ba et al., 2016) in the critic/value network of the IQL. agent to prevent catastrophic overestimation. We list the hyperpa- rameters used in experiments in Table 5. For training REDS, collect subtask segmentations for suboptimal demonstrations using the automatic subtask inference procedure described Sec- tion 4.4, and we manually modified some subtask segmentations with false negatives guarantee stable performance.
Computation We use Intel Xeon CPU @ 2.2GHz CPU cores and NVIDIA RTX 3090 GPUs for training reward model, which takes about 1.5 hours in Meta-world and hours in Furni- tureBench due to high-resolution visual observations from multiple views. For training DreamerV3 agents in Meta-world, 3090 GPU, which takes approximately 4 hours over 500K environment steps. For training IQL agents in FurnitureBench, we 24 Intel Xeon CPU @ 2.2GHz CPU cores and single NVIDIA RTX 3090 GPU, taking approximately 2 hours for IM gradient steps in offline RL and 4.5 hours over 150 episodes of environment interactions in online RL.
use 24 Intel Xeon CPU @ 2.2GHz CPU cores and a single NVIDIA RTX
B REDS ARCHITECTURE DETAILS
We encode visual observations with pre-trained CLIP (Radford et al., 2021b) ViT-B/16 visual en- coder, utilizing all representations from the sequence of patches. We adopt 1D learnable parameters with the same size for positional embedding, and we add these parameters to 2D fixed sin-cos em- beddings and add them to features. To encode temporal dependencies in visual observations, a GPT (Radford et al., 2018) architecture with 3 layers and 8 heads. FurnitureBench, we use sequence of images from both the front camera and wrist camera as input. Given s.ront wrrsst from the front/wrist camera, we concatenate visual observations to front -o wrist (t-K-1.... of front or wrisst] add positional embeddings, 2D fixed sin-cos embeddings, and additional 1D learnable parameters for each viewpoint for effectively utilizing images from multiple cameras. We then pass the features to the transformer layer, the same the model with single image. The subtask embedder and final reward predictor are implemented as 2-layer MLPs.
use
'https: / thub. com /lma/ furni http://ww//////ttttteeeeoo/////r learning
18

---

Published as a conference paper at ICLR 2025
C BASELINE DETAILS
ORIL (Zolna et al., 2020) For implementing ORIL with visual observations, we use the CNN architecture from Yarats et al. (2021) to encode image observations. For training data, use the same set of demonstrations as for training REDS. Since our training data are divided into success and failure demonstrations, we do not use positive-unlabeled learning (Xu Denil, 2021) in our experiments. For robustness against visual changes, we apply the same augmentation techniques used for training REDS.
Rank2Reward (R2R) (Yang et al., 2024) To ensure compatibility with backbone RL algo- rithms (Hafner et al., 2023; Kostrikov et al., 2022) implemented in JAX, we reimplement the reward model with JAX following the official implementation of Rank2Reward and use the same hyper- parameters. We first pre-train the ranking network using the same expert demonstrations as REDS, and we then train discriminator for the expert demonstration and policy rollouts, weighted by the output from the pre-trained ranking network. For training efficiency, we use the CNN architecture from Yarats et al. (2021) for encoding visual observations instead of R3M (Nair et al., 2022), find- ing no significant difference when we use the pre-trained visual representations like R3M, but with much slower training in online RL. We observe that our R2R implementation with DreamerV3 in JAX outperforms the original version implemented with DrQ-V2 (Yarats et al., 2022) agents.
DrS (Mu et al., 2024) Similar to R2R, we reimplement DrS with JAX following the official im- plementation of DrS 6, and use the same set of hyperparameters for reward learning. As the original DrS implementation is based on a state-based environment, we switch the backbone RL algorithm from SAC to DrQ-V2 (Yarats et al., 2022) and apply the augmentation technique in the reward learn- ing phase for processing visual observations efficiently. To report the RL performance, learned dense reward model to train new RL agents. In FurnitureBench experiment, we train the reward model with the same expert/failure demonstrations in Section 5.2, without online interac- tion, to avoid unsafe hehaviors and a significant increase in training time from online interactions.
use the
VIPER (Escontrela al., 2023) We use the official implementation of VIPER for our exper- iments. Given the similarities among robotic manipulation tasks, we use the same set of hyperpa- rameters as in RLBench (James et al., 2020) experiments to train VQ-GAN and VideoGPT. We train 100K steps, choosing the checkpoint with the minimum validation loss. In FurnitureBench experi- ment, we use images from the front camera, resized to 64 x 64 x 3, and set the exploration objective O.
D TASK DESCRIPTIONS
In this section, we list the subtasks and corresponding text instructions for each task in Table 6. For Meta-world tasks, we provide the code snippet used to determine the success of each subtask (Please refer the Meta-world (Yu al., 2020) for more details). For the FurnitureBench One Leg task, we outline the criteria used by human experts to assess the success of each subtask based on the metric defined in BurnitureBench (Heo al., 2023).
E EXTENDED RELATED WORK
Quantifying differences between reward functions Previous work has explored methods for measuring the difference between reward functions without relying on policy optimization proce- dures (Gleave et al., 2021; Wulfe et al., 2022; Skalse et al., 2024). In particular, Gleave et al. (2021) introduced the EPIC distance, pseudometric invariant equivalent classes of reward functions. Subsequent work (Rocamonde et al., 2024; Adeniji et al., 2022: Liang et al., 2022) has employed EPIC assess the quality of reward functions. In this paper, we take different approach by using EPIC distance as an optimization objective. While Adeniji et al. (2022) also utilizes EPIC distance
'https: http::::///itttbb...o http:tthh.....mmmmmatttaameaeate http://ww.....oooaaa//aaa/eeeaae http://ggghbb...ommmoooggggu////
Phttps://g http://ggtttubb.oooeeoeeeeoce htt.tt...oomaeleccttteeeeaaaaaee
19

---

Published as a conference paper at ICLR 2025
Table 6: A list of subtasks and language description for each subtask used for REDS in our experi- ments.
Telk Mee Ciee
pagg deertaiee
anteeniing bor Eaorthend hanely
Morevestd Unswer Open
terelggeee 000
Moravesd Loar
coady 2u -83
Litti am coculay lock
encesbrrely 2.10
daaidals Jair una)))
beiling hand:
odid KJ wch.tiitt.ool
irt am liing goon por tves
Mita woild Palk
egcluuyy 6/0
bager ox 8:55
1iti um heikknng grbbod and cate toa grres angst piis.
Mota woild Smmeplics
ed (.U.teliig www.eigggeesetgttttrr. Eoopupooll0
bagan c8) 5:85
artht un amieldnn goaboo cate soa de Muc sign yeas
Neuavesd Baer Ulive
alls
Neuavesd Wadden Olee
cagan butht 555
me kd
helding aad
Peritaslledd
orfore ofil
uoand oxer
caness cple cheeckck
og tacradi
pigpes
BatTEa do
fely
lind
Milkect Unaeai Oavee Lnbeeta Sane
Grappoo DatecteCSeddieeee scdeell
qcape: eddratttaa 1o
1I crobhig te entcis
auaclu unbecta
sagared 211
arth unieking extccl oe antech maa..
Table 7: A list of language description used for CLIP and LIV.
Task
Language description
Meta-world Duor Open Meta-world Drawer Open
a robot arm grabbing the drawer handle and opening the drawer. a rohot arm grabbing the door handle and opening the door to the green target point.
for optimizing intrinsic reward functions in skill discovery, our method applies EPIC distance to train dense reward functions for long-horizon tasks, serving direct reward signal for RL training.
Segmenting demonstrations for long-horizon manipulation tasks Several approaches have been proposed decompose long-horizon demonstrations into multiple subgoals to prevent error accumulation and provide intermediate signals for agent training. These include extracting key points from proprioceptive states (James Davison, 2022; James et al., 2022; Shridhar al., 2022; Shi et al., 2023), employing greedy heuristics on off-the-shelf visual representations pre-trained with robotic data (Zhang al., 2024), and learning additional modules on top of pre-trained visual- language models to align with keyframes (Kou et al., 2024). Our work builds on these efforts by leveraging subtask segmentations but focuses on developing a reward learning framework that explicitly incorporates subtask decomposition to generate suitable reward signals for intermediate tasks. Additionally, we further demonstrate that our model generalizes effectively to unseen tasks and robot embodiments.
20

---

Published as a conference paper at ICLR 2025
F EXTENDED QUALITATIVE ANALYSIS
VPER
12.
VIPE
Haman-Erginend
l'amanCrgreend
15.
-2000
VyPER Roword
2000
-2000
2200
-0000
72
-000
-1500
80
Tireskop
15
Tireeko
2.8
OHL Haman-engseend
7.8
lnghon-nggeeeedd
160
Hutte www.-annteeeeeee hewerd
8.
CRI 32400
7.8
100
2.5
7.8
2.4
Timaeda::
Timestoo
Humar--ngiaan onde-dd
Ore
17.5
15.2
Homeraa cwweedd reneerd
-1.50
Hursss-ngrreee
DD5
-1455
Dr Roward
-1.70
Rewore
10.2
-1.79
-1.08
-150
-1.90
Tirendep
Timestop
0.6
REDS
-0.2
REDS
Fursas-nnd neaned
15.3
Hamateet ggg www.eedd heskedd
-0.4
REDS Rovord
-1.0
12.5
-0.2
103
2.5
+1.
25
Timestec
Timssdep
Figure 10: Qualitative results of VIPER (Escontrela et., 2023). ORIL (Zolna et al., 2020), DrS (Mu et al., 2024), and REDS (Ours) in Peg Insert Side (left), and Sweep Into (right) from Meta-world Yu al. (2020). We visualize several frames above the graph and mark them with a diamond symbol.
21

---

Published as a conference paper at ICLR 2025
2
3
4
8
9
600
Subtask segmenattin
Sngmentation
WPEER Rovire
-200
403
1193
100
300
Timestep
5
89
segmentation
06 Rowar
-1.5-
-..6
Timestap
5
89
10
Bubrodk segmentaiin
REDI Rlvvard
Suttark
50
300
Subtask 3
Subtask
Subtesk
Subtask
Subtask
-684
-2.50
11.0
-6.6
-100
-3.00
11.2
-1.0
-325
-004
-12
-2
10.4
10
20
50
-8.75
60
80 100
120
160
200
250
Figure 11: Qualitative results VIPER (Escontrela et al., 2023), DrS (Mu et al., 2024), and REDS (Ours) in One Leg from FumnttreBennc (Heo et al., 2023). We visualize several frames above the graph and mark them with a diamond symbol. VIPER, which does not utilize subtask information, assigns lower rewards to later subtasks, making agents stagnate in earlier phases. While DrS uses ground-truth subtask information from the environment, produces sparse reward signals within each subtask. In contrast, REDS provides subtask-aware signals in subtask transitions and generates progressive reward signals (see the bottom figure zoomed for each subtask).
22

---

Published as a conference paper at ICLR 2025
100
REDS CLIP
100
c00.
100
percerril
Success Ratis (%)
UV
Sucosss Rats (96)
--0.
+-0.1
Success Rate (1%
50% percental
25th
perceerde
DR
6-1.0
60
40
40
40
20
20
8
12
16
23
8
12
16
12
16
20
Environment Step (x<10)
Environment Step (x10")
Environment Step (x10")
Additional baselines
(b) hyperparameter e
(c) Threshold Tu
Learning curves for two Meta-world (Yu et al., 2020) robotic manipulation tasks, mea- sured by success rate (%). The solid line and shaded regions show the mean and stratified bootstrap interval across 4 runs.
Figure 12:
G ADDITIONAL EXPERIMENTS
Comparison with additional baselines We first compare REDS with additional baselines, CLIP (Radford al., 2021a) and LIV (Ma al., 2023a), utilizing the distance between visual obser- vation and text instructions for generating rewards. It's important to note that these models cannot infer subtasks in online interaction unlike REDS; therefore, we use other text instructions describing how to solve the whole task. (refer to Table for details). Figure 12a shows that REDS significantly outperforms baselines, indicating that providing detailed signals aware of subtasks crucial for bet- ter RI. performance. Additionally, we compare REDS with Diffusion Reward (DR) (Huang et al., 2024), which utilizes conditional entropy from a video diffusion model as a reward signal. Our findings indicate that REDS also significantly outperforms DR. This is attributed to the fact that DR does not explicitly incorporate subtask information, which essential for generating context-aware rewards in Long-horizon tasks. These results further emphasize the advantage of REDS in handling tasks requiring precise subtask guidance.
Effect of scaling progressive reward signals In Figure 12b, examine the effect of scaling
the regularization for progressive reward signals in Equation 4. We observe that e 0.5 shows the best performance, while smaller values relatively weaken helpful progressive signals, and larger values degrade the reward function by reducing the accuracy in inferring subtasks.
Effeet of threshold Ty for subtask inference We present experimental results using various threshold Ty in Figure 12c. We observe that lower percentile threshold exhibits lower RL per- formance. These results indicate that a lower percentile threshold allows more observations to be classified as successful; however, this can lead misleading subtask identification, resulting in decreased RL performance.
Subtask identification ability of REDS To assess the subtask identification capability of REDS, we measure its precision before and after fine-tuning with additional suboptimal demonstrations. This evaluation involves us- ing 50 unseen expert demonstrations and 50 suboptimal demonstrations sampled from the replay buffer of Dream- erV3 agents that were trained with a numan-engineeree reward. Table 8 shows that the precision is comparable for expert demonstrations for both agents; however, there
Table 8: Precision in identifying sub- tasks of REDS on 50 unseen expert demonstrations and 50 unseen subopti- mal demonstrations.
Fine- tunittuning
Expert
Suboptimal
Total
X
94.49%
70.90%
82.70%
92.56%
91.49%
92.03%
significant increase in precision for suboptimal demonstrations after fine-tuning. This improve- ment in precision results in enhanced RL performance, as illustrated in Figure 9c.
23

---

Published as a conference paper at ICLR 2025
Additional EPIC measurements To further validate the efficacy of our method, we mea- sure the EPIC distance between learned reward functions and subtask segmentations in Meta- world environments. Note that we report the sult with the same set of unseen demonstrations used in Section 5.3. In Table we observe that REDS exhibits significantly lower EPIC distance than baselines across all tasks, consis- tently supporting the claims from the experiments in the main text.
Table EPIC (Gleave et al., 2021) distance (lower is better) between learned reward functions
and subtask segmentations in unseen data.
(Ours)
Meta-world Door Open
0.5017
0,5751
0.7017
0.4870
Meta-world Push
0.6293
0.7014
0.7094
0.5129
Metn-world Peg
Metn-world Sweep Into
0.6179
0.6584
0.7011
0.4293
24


Reviews:
[{"review_id": "Ea7hwyWots", "reviewer": "Reviewer_mmyJ", "summary": "Demonstration with Segmentations (REDS), for long-horizon robotic tasks that are usually composed of numerous subtasks. The key intuition is to use subtask annotations or descriptions for expert demonstrations to aid in learning a dense reward function. The reasoning behind this approach is that the subtasks occur in a particular order in expert demonstrations, and this order can be used as ground truth to learn the reward function. \n\nThe reward model or the reward predictor is defined as a function of the sequence of prior observations and the current subtask. The model is trained using the Equivalent-Policy Invariant Comparison (EPIC) metric. The goal is to match the learned reward function to the ground-truth reward function obtained from the ordering of subtasks in expert demonstrations. To identify the current subtask, the paper uses the cosine similarity between the CLIP image embeddings of the observations and the CLIP text embeddings of the textual descriptions of the subtasks required to achieve the goal. Therefore, before the execution of an episode, the subtasks have to be defined in natural language.\n\nAdditionally, the authors propose to continually learn the reward function from suboptimal demonstrations obtained from an RL agent. Here, the RL agent is trained using the reward function obtained from expert demonstrations. This process is repeated multiple times.\n\n The proposed approach outperforms other reward learning baselines for manipulation tasks from the Meta-world and the FurnitureBench environments.", "questions": "Questions:\nIs there a way to quantify the subtask identification performance? This could be useful as the identified subtask affects the reward prediction. \nAn ablation study depicting the effect of epsilon, which is used to enforce progressive reward signals, will be useful. CLIP-based subtask identification and ground-truth rewards from subtask ordering are not sufficient to enforce the required expert behavior. From my understanding, this is implicitly enforced using progressive reward signals. Therefore, an ablation study for epsilon will be crucial.\nThe authors need to explain how their approach is different from [1] and [2].\n\n[1] Vision-Language Models are Zero-Shot Reward Models for Reinforcement Learning. \n[2] LIV: Language-Image Representations and Rewards for Robotic Control.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, the paper is well written, and the problem is very well explained. The authors have clearly delineated their contributions from previous works. The experiments and ablations are detailed and informative. With respect to the proposed approach, the strengths are:\nThe approach requires minimal human supervision in terms of defining the subtasks accurately.\nThe proposed approach generalizes well for manipulation tasks with unseen objects. Additionally, since the reward model does not depend on actions, the authors show generalization capabilities to different robots. For example, the reward model trained only with the Panda arm can be used to train RL agents for the Sawyer arm.\nDependence on foundation models like CLIP for subtask identification reduces the adversarial effects of visual artifacts like varying lighting conditions, locations of objects, etc.", "weaknesses": "The expert demonstrations would always contain the subtasks in a particular order. This might lead to poor reward signals when the subtask estimation turns out to be incorrect. Such instances could occur when the RL agent is exploring. \nThe effect of the hyperparameter epsilon, which is used to enforce progressive reward signals within each subtask, is not clearly explained. The authors show an ablation for the cases with and without regularization loss. But, the effect of epsilon is not clearly described. \nFrom Fig. 3, it appears that the baselines have not converged for the Lever Pull, Sweep Into, and Push cases. \nThe proposed approach is not compared against CLIP-based zero-shot reward models like [1].\nLIV [2] also learns dense rewards from just observations and text descriptions. The proposed approach has to be either differentiated from or compared against this paper.\n\n[1] Vision-Language Models are Zero-Shot Reward Models for Reinforcement Learning. \n[2] LIV: Language-Image Representations and Rewards for Robotic Control.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "P3iBvhFhg0", "reviewer": "Reviewer_j9z4", "summary": "This paper addresses the challenge in reinforcement learning (RL) of relying heavily on human-designed reward functions, especially for long-horizon tasks. It introduces REDS (Reward Learning from Demonstration with Segmentations), which infers subtask information from video segments and generates corresponding reward signals for each subtask, utilizing minimal supervision. The framework employs contrastive learning objectives to align video representations with subtasks and uses the EquivalentPolicy Invariant Comparison (EPIC) distance to minimize the difference between the learned reward function and ground-truth rewards.", "questions": "1. Note that in the experimental setting of this article, expert demonstration videos of the current task are already available. Should imitation-learning-from-observation-based methods also be used as a baseline, for example, to directly derive rewards based on the similarity between the agent trajectory and the expert demonstration video?\n2. When learning a new task, how can we determine the appropriate way to divide it into subtasks? For instance, should it be broken down into 2, 3, or 4 subtasks?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This paper is written clearly and highlights an important challenge for long-horizon reinforcement learning.\n2. This paper provides a reward model training method that utilizes both expert demonstration videos and suboptimal videos.", "weaknesses": "- The method seems to rely heavily on carefully predefined subtasks or key completion points (like in table 6). This may limit the generalizability of the method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "rOpOJkaiLs", "reviewer": "Reviewer_Vi9Z", "summary": "This paper introduces an inverse reinforcement learning approach that enables the extraction of a dense reward signal from video demonstrations with minimal human intervention. The reward model is designed to consider both the observation and subtask information. This connection between the reward and subtask information intuitively breaks down the intricate long-term task into manageable subtasks. Empirical findings further validate the benefits of integrating subtask information into the reward learning process.", "questions": "What if we utilize the text embedding directly as input during the inference phase? Does this approach exhibit a substantial performance advantage over the method that employs video embedding?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "i) In this paper, subtask information is integrated into the reward learning process. During training, the subtask is included in the reward function input as a text embedding that provides instructions on completing a specific subtask. In the inference phase, the text embedding is substituted with a video embedding as an additional input.\n\nii) The approach employs the EPIC loss function to reduce the disparity between the predicted reward sequence and the ground truth reward. Experimental results demonstrate superiority of this approach.", "weaknesses": "i) In the training phase, this paper decomposes the overall task into multiple subtasks based on domain knowledge. However, the reliance on predefined instructions from the environment for task decomposition raises concerns about practical applicability. Some environments may lack such predefined knowledge, necessitating human annotations or the need for a learned task decomposition model when extending this approach to new environments. This raises doubts about the novelty and scalability of this methodology. \n\nii) This paper introduces the function $\\psi$, which maps observations to subtasks. However, providing more details about this function is essential. Specifically, elucidating the form of the subtask output by $\\psi$ is crucial since it serves as the ground truth for reward learning, a key component for the success of this method. Therefore, the authors should elaborate on the detail of the subtasks produced by $\\psi$.\n\niii) During the reward learning, this method utilizes the EPIC loss function to quantify the disparity between the estimated reward and the proposed ground truth. While EPIC was initially introduced in a prior work, the rationale behind its selection and comparisons with previous loss functions for reward learning should be further discussed by the authors. This explanation should include both intuitive reasoning and empirical evidence.\n\niv) During the inference phase, the substitution of text embedding with video embedding for subtask information is addressed. Although an alignment mechanism is employed to bring video and text embeddings closer in the latent space, the loss function appears to overlook the maximization of distances between irrelevant embeddings. This oversight could be crucial when dealing with subtasks that share similar text descriptions or video content. It is advisable for the authors to investigate the effectiveness of learned embeddings within tasks involving similar subtasks, although this aspect is not a primary concern.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "0QlwExXiVy", "reviewer": "Reviewer_VAFb", "summary": "This paper presents a novel visual reward learning framework that aligns video demonstration segments with learned rewards, effectively indicating progress in subtask completion. The performance improvements across eight simulation tasks and a real-world assembly task demonstrate the method's efficacy. The visualizations of the rewards, the framework's generalization capabilities, and the impacts of various design choices significantly enhance the comprehension of the learned reward.", "questions": "- What rules are used for the division of each task? Could different rules yield varying downstream performance?\n- In Table 1, why do DrS, VIPER, and REDS perform worse than Sparse Reward after the offline phase (approximately 1.1 vs. 1.8)?\n- What is the effect of $T_U$? Could the authors conduct experiments to explore this?\n- In Table 2, while REDS achieves the highest EPIC score, what happens if hand-engineered rewards are replaced with sparse ones (e.g., step-style rewards like in the One Leg Task)? Would the conclusion remain valid?\n- In Figure 4, REDS appears continuous in (a) but step-like in (b), despite both rewards being learned to minimize the distance between $\\phi$. What accounts for this difference in appearance of the learned rewards?\n- Why is Diffusion Reward not included in the experiments? It may outperform VIPER.\n- Does the format of the language instruction influence the results?\n- What explains the sudden drops for VIPER and DrS and flat curve of ORIL in Figure 10?\n\n[1] Sermanet, P., Xu, K., & Levine, S. (2016). Unsupervised perceptual rewards for imitation learning. RSS, 2017.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-organized, providing clear descriptions of the proposed approach, including effective reward visualizations and experimental setups.\n- Experiments are thorough and convincing. Results are favorable compared with SOTA approaches like VIPER and DrS. The comprehensive ablation study clearly delineates the impact of each component.\n- The experiments on generalization ability and real robots are impressive. It is good to see these properties are held in the proposed method.", "weaknesses": "While I think this paper deserves acceptance based on its strong results and presentation, I have several concerns that should be addressed in the next revision:\n\n- **Context-Aware Reward Signals**: While intuitively beneficial, learning context-aware reward signals may require extensive annotation, potentially weakening the motivation for this work.\n    - Existing studies, such as VIP and XIRL, have successfully learned context-aware reward signals without context labels but are not compared in this paper.\n    - The method bears resemblance to that presented in [1], which employs purely unsupervised learning without requiring subtasks; this similarity should be discussed and compared.\n    - Although the authors suggest that large models could be used to alleviate annotation burdens, this claim is not substantiated in the current version. I recommend either removing this mention or including an experiment that compares large-model-based annotations with human-annotated ones.\n- **Clarification of Motivation**: The motivation behind the proposed method needs more clarity. In Line 049, the authors reference prior approaches in the \"One Leg Task,\" but this is not empirically supported in Figure 2 or subsequent sections.\n- **Extension to Other Domains:** The applicability of the proposed method to other domains is unclear. Specific rules for subtask segmentation and the threshold $T_U$ should be explicitly described. The authors are encouraged to discuss how to adapt existing methods to other domains, what components need more attention, and potential failure scenarios.\n- **More Experiments**: The experiments could be more convincing. While the authors claim that the proposed reward can better handle long-horizon tasks, the selected tasks are not too long-horizon, in my perspective. It is suggested to extend One Leg task to Two/Three Leg Tasks, and see the changes in learned reward and the corresponding performance.\n- **Generalization Failures**: It is impressive to see experiments on generalization ability. However, it would be beneficial to present cases where the proposed method struggles. For instance, the authors could include discussions on scenarios such as changes in camera view or background.\n- I also include several smaller questions below that should be addressed in the next revision as well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "B488Zdye5v", "reviewer": "Reviewer_7C5k", "summary": "This paper introduces REDS, a novel framework aimed at improving reward learning for reinforcement learning (RL) in complex robotic tasks. The framework tackles the problem of designing reward functions, which typically require extensive human effort and domain-specific knowledge. Instead of relying on predefined reward functions, REDS uses segmented video demonstrations where subtasks are labeled, allowing the model to learn from these segmentations with minimal supervision.\n\nThe key innovation lies in aligning video segments with subtasks using contrastive learning and employing the EPIC (Equivalent-Policy Invariant Comparison) distance to ensure that the learned reward function is consistent with ground-truth subtasks. The model is tested on both simulated tasks (Meta-World) and real-world tasks (FurnitureBench) and demonstrates superior performance in long-horizon tasks with multiple subtasks.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. A novel framework for reward learning.\n2. Experiments are very extensive, including both simulated and real-world tasks (FurnitureBench).", "weaknesses": "1. The method relies on pre-trained visual and text encoders (e.g., CLIP), which may not be optimal for robotic tasks. The paper acknowledges that subtle visual changes are not well-handled, and reward quality could be further improved with better pretraining.\n2. While REDS reduces the need for hand-crafted reward functions, it still requires segmented demonstrations. The quality and quantity of these demonstrations could impact the performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": {}}]

Author rebuttal:

Decision: Accept (Poster)