SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation
Authors: Qianzhong Chen
Justin Yu
Mac Schwager
Pieter Abbeel
Fred Shentu
Philipp Wu
Venue: ICLR 2026 Poster (2026)
Keywords: Imitation Learning
Reward Modeling
Robotics Manipulation
Primary area: applications to robotics, autonomy, planning

Abstract: Large-scale robot learning has made progress on complex manipulation tasks, yet long-horizon, contact-rich problems—especially those involving deformable objects—remain challenging due to inconsistent demonstration quality. We propose a stage-aware, video-based reward modeling framework that jointly predicts task stage and fine-grained progress, using natural-language subtask annotations to derive consistent labels across variable-length demonstrations. This avoids the brittleness of frame-index-based labeling and provides stable supervision even in tasks like T-shirt folding. Our reward model is robust to demonstration variability, generalizes to out-of-distribution scenarios, and improves downstream policy training. Building on it, we introduce Reward-Aligned Behavior Cloning (RA-BC), which filters and reweights demonstrations based on reward estimates. Experiments show that our method significantly outperforms baselines in both real-world rollouts and human validation. On T-shirt folding, we achieve 83\% success from the flattened state and 67\% from the crumpled state, compared to 8\% and 0\% with vanilla BC. Overall, our results highlight reward modeling as a scalable and annotation-efficient solution for long-horizon robotic manipulation. Project website: https://qianzhong-chen.github.io/sarm.github.io/.

Published as a conference paper at ICLR 2026
SARM: STAGE-AWARE REWARD MODELING FOR LONG HORIZON ROBOT MANIPULATION
Qianzhong Chen 13+1 Justin Yu?.ă Mac Schwager', Pieter Abbeel?, Yide Shentu Philipp Wu? Stanford University. "UC Berkeley, 'xdofaa (gchen23, schwager)@stanforzd edu (yujustin, pabbeel, fredshentu/Bberkelley [davidchen, justinyu, philipp)exdof..i
ABSTRACT
Large-scale robot learning has made progress on complex manipulation tasks, yet long-horizon, contact-rich problems especiallyy those involving deformable objects-remte challenging due to inconsistent demonstration quality. We pro- pose stage-aware, video-based reward modeling framework that jointly predicts task stage and fine-grained progress, using natural-language subtask annotations to derive consistent labels across variable-length demonstrations. This avoids the brittleness Trame-index-brr based labeling and provides stable supervision even tasks like T-shirt folding. Our reward model is robust to demonstration variabil- ity, generalizes to out-of-distribution scenarios, and improves downstream policy training. Building on it, introduce Reward-Aligned Behavior Cloning (RA- BC), which filters and reweights demonstrations based on reward estimates. Ex- periments show that our method significantly outperforms baselines in both real- world rollouts and human validation. On T-shirt folding, we achieve 83% success from the flattened state and 67% from the crumpled state, compared 8% and 0% with vanilla BC. Overall, our results highlight reward modeling as scalable and annotation-efficcent solution for long-horizon robotic manipulation. Project web- site: hhttts: aazhong-chen. gi thub.......... io/sarm. .etrrm.ggttuuie prtthbi..o// Keywords: Imitation Learning, Reward Modeling, Robotics Manipulation
INTRODUCTION
The long-standing vision of enabling robots to seamlessly assist humans in household chores has inspired decades of research in robotics. From tidying living spaces to preparing meals, such capa- bilities hold the promise of freeing up human time, and improving quality of life. Recent progress in foundation models for robotics, more generally robot behavior models (RBMs), has sparked renewed optimism toward this goal. By combining visual perception, motor control, and option- ally language processing in a single framework, RBMs (Chi et al., [023] Zhao et al. 2023 Chen
et al. 2025 Sun et al.) 2024 Huang et al., 2024 Yu et al.) 2024 Wang et al. 2023a Black et al.; Team al. 2024 Zitkovich al. 2023 Shentu al. 2024 Huang al.) 2025a b) enable robots
to perform complex tasks, making it possible to execute these tasks in unstructured household envi- ronments.
Despite their promise, RBMs still struggle with long-horizon, contact-rich manipulation, particularly with deformable objects like T-shirts. Such tasks demand handling changing geometries, occlusions, fabric variations, and error- error-free multi-step planning challenges where current models, often tuned for short-horizon rigid-object tasks, fall short. They fail generalize beyond curated data, lose consistency over time, and misinterpret intermediate states. While many prior works in RBMs have focused on scaling up data (Barreiros et al.) [025]] Lin et al.] 2024), far less attention has been given to data quality. However, high-quality data is difficult to obtain: expert demonstrations are costly and time-intensive, while larger datasets often include noisy or suboptimal trajectories from less experienced operators. Even more challenging, demonstration quality itself is difficult
Conresoonding author.
Work done during internships at xdof.ai.

---

Published as a conference paper at ICLR 2026
PIESO
Sugge Model
D diveense
Suage 4l grabs
Reward Model
Stage e2: Baneen
Reward Model
Image
Joint State
Rewsnf-Aligned Weights
Stage #3: fold
V1:1
T1:N
y1:N
Model
(+) Label progress with annocated data
(b) Train singe-aware reward model
(c) Train robut policy with RA-BC
Figure 1: Overview of method's framework for (a) data processing, (b) reward model training, and (c) policy training with reward signals. Danno denotes the annotated dataset used for training the reward model, with examples shown in Fig. (gg_5]aa and Fig. Daiverse refers diverse expert dataset without annotations, which contains many suboptimal trajectories.
metric to quantify, since it depends on hidden factors such as action consistency and contact stability that cannot be directly measured, aside from simple proxy heuristics like task duration. Although there exist more sophisticated data-modeling approaches for assessing data quality and filtering trajectories after policy training (Belkhale et al.) (2023) Dass et al. 2025 Agia al. 2025). practical evaluation of demonstration quality remains challenging.
In light these challenges, we propose video-based reward modeling framework that leverages natural language annotations to assign progress labels and enable stable reward estimation for multi- step tasks. The learned reward model drives a Reward-Allined Behavior Cloning (RA-BC) frame- work, filtering higher-quality data and improving policy performance in both simulation and the real world. Focusing on the T-shirt folding task, our experiments show that coupling the reward model with RA-BC significantly boosts performance, underscoring the importance of data quality in long-horizon manipulation. Together, these contributions advance scalable and annotation -fffiiitn-efficcentt imitation learning. An overview is shown in Fig.
Our contributions can be summarized as follows:
We present SARM: stage-aware reward modeling framework that automatically derives task progress labels from natural language annotations. Given any subsequence of RGB frames, the model jointly predicts the current task stage and fine-grained progress within that subtask, achieving robustness, generalization out-of-distribution scenarios, and strong utility for down- stream policy learning.
We propose the RA-BC framework, which leverages the learned reward model to identify high- quality demonstrations and reweight training data accordingly.
We validate our approach on the real-world task of T-shirt folding, a challenging long-horize task that requires manipulating deformable objects, where it consistently outperforms strong behavior cloning baselines.
RELATED WORKS
2.1 LEARNED REWARD MODELS FOR ROBOTICS
Prior work on learning reward functions includes inverse reinforcement learning (g et al. 2000
Abbeel & Ng 2004 Ramachandran & Amir 2007 Ziebart al. 2008,; Finn et al.) 2016). which infers rewards from demonstrations but suffers from reward identifiabliity and sensitivity to partial observability that hinder scalability to high-dimensional, long-horizon problems.
2

---

Published as a conference paper at ICLR 2026
Learning from human feedback (e.g., preference rankings, scaled preferences, interventions) has proven effective in training large language models (LLM) (Christiano et al. 2017: Ziegler et al. 2019). Recently, RLHF has also gained increasing interest in robotics but still requires substantial task-specific input and suffers from annotator inconsistency (Sadigh et al. 2017 Liu et al. 2023).
A complementary direction uses LLM to synthesize reward functions or shaping code (Ma al., 2024a Shentu et al. 2024). which can accelerate bootstrapping but often assumes privileged or structured state information that rarely available outside simulation and can degrade under sensor noise and domain shift.
Several prior works (Lee et al. 202T by computing the feature distance to goal state, enabling self-supervised reward model training without manual annotation. While effective for simple tasks with single objective, such approaches struggle in long-horizon settings where the task naturally decomposes into multiple subtasks or stages. these cases, a single goal distance fails to capture intermediate progress, often causing the reward signal to become uninformative or misleading.
Ma al., 2022 Escontrela al. 2023) estimate rewards
Another line computes rewards directly from visual observations combined with task text using vision-language models (VLM). Among these, LIV (Ma et al., 2023). VLC (Alakuijala et al. 2024). GVL (Ma al., 2024b). VICtoR (Hung et al. 2024), REDS (Kim al. 2025), Re WiNDD ZZang robot manipulation tasks directly from visual perceptions. In practice, many VLM based reward models struggle on long-horizon, highly dynamic, and contact- rich manipulation tasks because they process entire trajectories from the initial frame to resolve temporal dependencies, which increases data and computation demands and impedes scaling.
al. 2025) and
There are prior works such DrS (Mu al. 2024) and REDS (Kim et al. 2025) that stage- aware reward models for long-horizon tasks. However, DrS is fundamentally different from visual- based SARM: it is purely state-based, depends on full simulator states stage indicators, and quires training separate discriminator for each stage, making it difficult to scale. REDS also differs from SARM: instead of modeling continuous frame-wise progress curve, it learns semi-sparse step-shaped reward with monotonicity regularization, which struggles to generalize when trajec- tories progress at different speeds. addition, REDS infers stage via image- subtask embedding similarity rather than a dedicated stage-estimation network, which becomes unreliable when sub- task descriptions are semantically similar.
2.2 IMITATION LEARNING WITH SUBOPTIMAT. DEMONSTRATIONS
Prior works have explored imitation learning under suboptimal datasets. One direction adopts boot- strapped frameworks (Sasaki& & Yamashina| 2020: Belkhale et al. 2023 Dass et al.) 2025 Agia et al.] 2025), which actively change the dataset distribution based on analyzing current policy's gradient, rollout, or learning objective during training. While effective, such methods are computationally expensive and require extensive hyperparameter tuning. Another line of research focuses on explic-
itly labeling and classifying demonstrations (Wu et al. 2019 Wang et al. 2023b). but this approach depends on small, high quality dataset as prior knowledge.
An alternative direction investigates weighted BC through offline reinforcement learning (RL) tech- niques (Wang et al., 2018 Chen al. 2020 (Siegel et al.. 2020 Xu al. 2022). where estimates of the advantage function are used to prioritize actions in the dataset. These methods, however, often assume access to full-state feedback and a well-trained critic, and have not been validated on real- world, vision-based, long-horizon manipulation tasks. In contrast, our RA-BC framework leverages a pre-trained, vision-based reward model to generate robust and accurate K-step advantage esti- mates, which then guide weighted BC training.
3 METHOD
3.1 REWARD MODEL TRAINING
Data Processing. Extracting dense reward labels remains challenge, especially in long-horizer complex tasks. Prior work often relies on frame indices as labels (Zhang et al.) (025). While this may suffice for short tasks with fixed duration, such as "pick up the cup." it fails for tasks like
3

---

Published as a conference paper at ICLR 2026
E1:N
Stage Estimator
Classification Head
MLP Fusion Net
Transformer Soquential Aggregator
PR-LEN
Image Frames
Joint State
Lineur Projector
Linear Projector
+
CLIP
State Normalizer
Subtask
Regression
Fi-x
Estimator
Head
Figure 2: Overview of SARM, stage-aware reward modeling. Left:: SARM overview, which in- cludes both a stage estimator and subtask estimator. First the task stage is predicted from the ob- servations. This prediction is additionally passed into the subtask estimator which predicts scale value of the progress within the stage. Right: An overview of the estimator architecture which is replicated for both the stage estimator and the subtask estimator.
"fold the T-shirt," where trajectories vary greatly, task duration not fixed, and motion sequences differ across demonstrations. For example, in T-shirt folding, the flattening phase may require more or fewer motions depending on shirt placement or fabric configuration, yet frame- based labeling only reflects elapsed time. As result, identical task states (e.g., fully flattened shirt) can receive progress values ranging from 0.2 to 0.8, introducing severe label noise that harms reward model learning and downstream policy training.
To resolve this, we leverage subtask annotations on the robot trajectory data. The collected trajec- tories consist of three video streams (top, left wrist, and right wrist), joint states, and joint actions. Before annotation, we designed annotation protocols by decomposing each task into semantically meaningful subtasks. For T-shirt folding. developed notation and another for dense annotation, as illustrated in Table A.3 and Fig. and During annotation, only the subtasks defined by the protocol were labeled, and any trajectory that did not contain the complete sequence of subtasks specified by the protocol was discarded. Annotators watched the top-view video and segmented each trajectory into subtasks by recording the start and end frame indices. If a serious mistake (e.g. the manipulator hitting the table heavily executing completely reversed motion sequence) occurred during execution, its start and end frames were also labeled; trajectories containing mistakes were excluded from subsequent model training.
distinct protocols: one for sparse an-
Using the annotated data, we computed the average temporal proportion of each subtask across the dataset to automatically assign progress values the start and end frames of each subtask. Within each subtask, finer-grained progress labels were generated by lincarly interpolating over frame indices. This procedure ensures that progress labels remain closely aligned with the semantic meaning of the motions while maintaining consistency across the entire dataset.
Labeling by subtask priors: Let trajectory have total length T; and be segmented into K subtasks with lengths {Li,k 1 We estimate a dataset-level prior proportion for each subtask
1
M
Li
K
ās 2
O,
1,
M
T
1=1
A=
where M is the number of trajectories.
Frame-wise progress turgets: For frame that lies inside subtask & with local bounds [SA, en). 1- c [0, 1] and the cumulative prior Ps
define the within-subtask normalized time
Σ 0g (with Po 0). We assign the normalized progress target
Dt = Pa-1 + a Te [0. 1).
(2)
that Us, Pk-1 and you Pk-
4

---

Published as a conference paper at ICLR 2026
Model Architecture. We adopt dual reward- -moeel architecture with shared backbone archi- tecture and two task-specific heads. The stage model predicts the current high-level stage, while the subtask model estimates fine- fine-graine grained progress conditioned on the stage prediction. An overview of SARM architecture is demonstrated in Fig. These models operate sequentially: the subtask model uses the predicted stage prior context refine the final progress estimate. The stage model out- puts probability distribution over discrete task stages, providing a coarse localization of the robot's s progress, while the subtask model leverages the stage embedding to produce a continuous progress value in [0, 1]. Together, they provide both high-level stage classification and fine-grained progress estimation, enabling stable reward modeling in long- nong-ooizonn manipulation tasks. An overview of the reward-model architecture is shown in Fig. 2
The input pipeline proceeds as follows: (1) sequence of N images is encoded by frozen CLIP encoder, producing visual embeddings shared across both models; (2) visual embeddings and joint states are projected into a common dinee--iieenional space, where only the first frame receives an explicit positional bias to prevent absolute temporal leakage, following ReWiND Zhang et al., 2025;; (3) the multimodal sequence then processed by transformer encoder to capture tem- poral dependencies and cross-modal interactions; (4) lightweight MLP head fuses the aggre- gated features and outputs either stage logits VI.N R (stage model) or scalar progress predictions i1:N E [0, 1j (subtask model), where the latter is explicitly conditioned on the predicted stage to refine the progress estimate. Stage probabilities are obtained via a softmax 11 softmax 11.N) E [0, 1] from which the discrete stage prediction and normalized progress are calculated as
SI.N arg max IIN, S. E (k.....,k) VE(L....kk
(3)
VIN = PR-I,IN + Cg, IN PIN, DIN E [0.1]. 1].
(4)
3.2 REWARD-ALIGNED BEHAVIOR CLONING (RA-BC)
Behavior Cloning (BC) trains policy mo to imitate actions from demonstrations by minimizing supervised loss on state-action pairs (o;, a;). The standard BC objective averages per-sample losses,
N
1
(0)
((To(o:), az),
(5)
N
where é is mean squared error for continuous actions or cross-entropy for discrete actions.
RA-BC objective. RA-BC replaces the uniform prior in equation with reward-aligned weight- ing that emphasizes demonstrations predicted to make progress. For each training item i, we sam- ple a current window (anchor) and its next window obtained by advancing one action chunk. Let <(-) E [0, 1] denote the normalized progress score produced by the reward model (Sec. 3.T.. If the anchor window ends at time t and the chunk length (stride) is A, we form a per-item progress delta
= ((o - ((.)).
(6)
which serves scalar signal of expected improvement. This Pr is then mapped to weight w, [o, 1] (see weighting rules below), and RA-BC minimizes the normalized weighted objective
N wA ((TT(O)), a.)
CRA-BC (0)
ΣN w/ +E
with a small & > 0 to avoid division by zero.
Weighting from running statistics. To calibrate w, without fixed heuristics, RA-BC maintains online running statistics (mean and standard deviation σ) of the raw progress deltas (fy} via numerically stable estimator (Welford). We clamp the running mean be nonnegative, p + max(u, 0), to avoid centering weights around negative progress in early training. Each F; is mapped soft weight by a linear ramp between (μ 220) and (μ + 20);
P: (μ 20)
clip
0,
(8)
4a +
where clip(ar, 0, 1) min(max(r, 0), 1) and
0 guards small variances.
5

---

Published as a conference paper at ICLR 2026
1.0
Predicted Progress
0.5
Ground Truth
Proposed Model
ReWiND Model
0.0
5
10
15
20
25
30
35
Time (8)
Figure 3: A visualization of the predicted task progress for T-shirt folding demonstrations. Com- pared with ReWiND, SARM provides more accurate and calibrated estimates.
Prior overrides and validity mask. We incorporate lightweight prior knowledge via threshold k to make weights decisive for clearly good/bad items:
Wi (f:2 + [ogRgs) Wi.
Granularity and implementation. RA-BC is architecture-agnosti and can be applied at the sam- ple or sequence level. our implementation, losses are first averaged over a temporal chunk to ob- tain per-item loss, after which Eqs. equation equation8-equatioon makes RA-BC a drop-in replacement for Eq. equation 5 that softly filters noisy or non non-proge proggreesssin data while preserving training stability via normalization. In practice, RA-BC selectively empha- sizes high-quality segments and down-weights suboptimal ones, enhancing policy learning espe- cially when the dataset is diverse and contains imperfect demonstrations.
produce w, used in Eq. equation] 7 This
RESULTS
In this section, we answer three questions:
Q1. How does SARM lead to more robust reward model when faced with long horizon, complex manipulation tasks?
Q2. How can RA-BC enhance policy training when faced with diverse datasets?
Q3. How does the quality the reward model affect RA-BC performance?
4.1 Q1: REWARD MODEL EVALUATION
We evaluate SARM two tasks, (1) T-shirt folding: a long-horizon, multi-stage stage, contact-rich ma- nipulation problem, (2) unload dishes from a rack: shorter-horizon multi-stage task with high vari- ation due varying dish counts, orientations, rack positions, and optional handovers. We utilized these two tasks to demonstrate the effectiveness of the proposed reward model training framework.
Baselines. We compare SARM against LIV (Ma al. 2023). reward model pre-trained EpicKitchens (Damen et al., 2022); VLC (Alakuijala et al. 2024), which fine-tunes a VLM via sequential ranking to encourage monotonically increasing rewards: GVL (Ma et al.) 2024b), which prompts pre-trained VLM with shuffled frames to predict per-frame progress; VICToR (Hung et al., 2024). which determines the motion class and motion progress by evaluating the similarity between the vision and text embeddings encoded from current frame and language instruction, respectively; REDS (Kim et al. 2025), which stage-aware reward model learning from the stage segmenta- tion; and ReWiND (Zhang al.) 2025), which augments input sequences with rewound frames to improve robustness against failure cases. All applicable baselines use transformer encoders matched in model size SARM; implementation details provided in A.4 All baselines trained on the union of the dense and sparse datasets. Ground-Fe truth rewards are normalized to the range [0, 1] for both annotation types.

---

Published as a conference paper at ICLR 2026
Table 1: Evaluation of reward models. "Demo _ denotes the single-step MSE of reward models on the validation set. All models are evaluated on 70 trajectories (50 from D sparse and 20 from Pdense), where both ground-truth progress and model predictions are normalized to the [o, 1] range. The two- scheme models (last two columns) are evaluated in "sparse mode." " "Rollout p reports performance on real policy rollouts. Visualization examples of reward model predictions on both demonstration data and policy rollouts are provided in Appendix A.5
Metries
Baseline Methods
Ablation Studies
SARM
GVL
VLC
LIV
REDS
VICtoR
ReWIND
Dense
Sparse
w/o
SARM (VB)
Demo C. 4.
0.064
0.083
0
0.036
0.079
0.019
0.027
0.013
0.008
0.015
0.019
Rollout p
-0.39
-0.33
0.33
0.16
0.00
0.50
0.11
0.78
0.67
0.78
0.94
Classification SR Breakdoww
5B
0/12
12/12
6/12
12/12
0/12
12/12
12/12
10/12
12/12
12/12
12/12
6/12
0/12
12/12
8/12
3/12
8/12
11/12
9/12
10/12
11/12
F2
5/12
0/12
6/12
7/12
12/12
7/12
5/12
11/12
9/12
10/12
12/12
Ablation Studies. We additionally perform ablations by (1) Dense: training a single-scheme model only on the dense annotation dataset, (2) Sparse: training a single-scheme model only on the sparse annotation dataset, (3) w/o R: removing rewinding frames augmentation, and (4) SARM (VB): evaluating SARM under varied brightness, where each frame's 's brightness is perturbed by up to 0.3 from its original value.
Evaluation consists of two parts. First, for human demonstration progress estimation, models are evaluated on unseen testing data. For all baselines and two-scheme models, the validation set is the union of the 10% hold-out data from both datasets. For the single-scheme ablations, we apply cross-dataset validation: models trained on dense annotations are evaluated on the sparse hold-out set, and vice versa. We report single-step mean squared error (MSE) loss L on the validation set. Second, for robot rollout progress estimation, we fine-tune a PiO policy (Black) et al.) on the datasets we mentioned above using RA-BC, and then deploy the policy at different training stages on a real robot to collect 36 trajectories. These trajectories include 12 successful episodes (SE), 12 partially successful episodes (PSE), and 12 failed episodes EF rollouts. Reward models are evaluated on these rollout trajectories according to the following classification protocol:
Evaluation Protocol.
SE,
if P Pinal > 0.8 A 1 Ei Pt > 0.6,
T/3
Label
PSE, if iT P. P.
(10)
FE, otherwise.
where P, is the predicted progress at frame t, T is the trajectory length, and & is the median of the average progress over the non-successful rollouts, ensuring an equal split between PSE and FE. By doing so, we avoid the bias introduced by manually setting thresholds for distinguishing between PSE and FE. We further compute a classification score p by assigning +1 for each correct prediction and -1 for each incorrect one, normalized by the total number of rollouts (36), i.e., #conect- #wrong We additionally report breakdown of the estimation success rate (SR) for each category (SE, PSE, FE) across models, in order to highlight their distinct behaviors, such as being overly optimistic assigning zero progress universally.
36
Analysis on T-shirt folding task. We prepare two datasets: Ddense with dense annotations con- taining 200 trajectories, and Daparse with sparse annotations containing 500 trajectories, examples of dense and sparse annotated demonstration data can be found at Fig. [5] [6 Importantly, these are distinct demonstrations rather than the same trajectories annotated differently. We reserve 10% of each dataset for testing and use the remainder for training. D Daparse. due to its larger size, cov- ers a wider range of scenarios, whereas Daeuse provides more detailed per-trajectory labeling. We trained SARM reward model both D and D Dopusee the details model training are provided in Appendix A.4
The detailed comparison results of the reward models are presented in Table Among the baselines, GVL, VLC, and VICtoR failed to provide reliable reward signals: classification breakdowns reveal that GVL and VICtoR tend to be overly pessimistic, while VLC is excessively optimistic. LIV,
7

---

Published as a conference paper at ICLR 2026
ReWiND and REDS achieve stronger performance, delivering more accurate classifications of policy rollouts. However, due to the unstable reward labeling issues discussed in Section (.11 their effectiveness remains limited on both human demonstrations and robot rollouts when compared with SARM. Although its regularization loss introduces a mild progressive trend, the estimation backbone of REDS remains step-shaped and semi-sparse, making it difficult to produce dense and accurate reward estimates.
For the ablation studies, both single-scheme variants underperform relative our two-scheme model, highlighting the advantage of leveraging larger datasets with heterogeneous annotation pro- tocols. Notably, the model trained solely on the dense annotation dataset performs poorly on unseen scenarios. Since this dataset is smaller and less diverse, it fails to capture the wide range of situations present in real-world rollouts, which often involve complex patterns such as back-and-forth motions, misgrasps, and recovery struggles. Training without rewind augmentation also results in degraded performance. While human demonstration evaluation remains largely unaffected since the dataset does not contain deliberate failure cases and progress is generally monotonic- the performance real robot rollouts drops substantially. In this case, the model becomes overly optimistic, failing to recognize regressions failures. These findings demonstrate that rewind augmentation is essential for building reward models that generalize to real-world policies. SARM (VB) showcased SARM's robustness under varied lighting conditions, which is important when faced with diverse dataset.
In summary, our method consistently outperforms all baselines, achieving more than 50% relative improvement on human demonstration benchmarks and over 80% improvement real robot roll- outs compared to the strongest baseline, ReWiND.
Analysis on dish unloading task. Dish unloading is a multi-stage task where the robot removes dishes from rack and places them flat on table. The process has three stages: (1) grasp and lift, (2) optionally hand over to the other arm, and (3) place on the table. Unlike T-shirt folding, it shorter- horizon and does not involve deformable objects, but introduces greater execution diversity: varying dish counts, orientations, rack positions (left vs. right arm use), and whether handovers are needed. This variability makes reward modeling challenging, especially for visual understanding. We a single dataset, Daish. with the same annotation protocol and training setup as T-shirt folding. As shown in Table 2 SARM consistently outperforms all baselines. The visualization results can be
found at Fig. 13 and 14
Table 2: Evaluation of reward models for "unloading dishes" task. "Demo f denotes the single-
step MSE of reward models on the validation set. All models are evaluated on 30 trajectories which are not included in training set. Both ground-truth progress and model predictions are normalized to the [0, 1] range. "Rollout p" reports performance on real policy rollouts.
Metries
GVL
VLC
LIV
ReWiND
w/o
SARM
Demo L1
0.089
0.045
0.042
0.018
0.013
0.013
Rollout p
-0.33
0.39
0.55
0.50
0.67
Classification SR Breakdown
SE
0/12
12/12
7/12
12/12
12/12
10/12
PSE
6/12
0/12
12/12
9/12
7/12
9/12
FE
12/12
0/12
6/12
7/12
8/12
11/12
4.2 Q2: POLICY LEARNING WITH RA-BC
Folding a crumpled T-shirt is among the most challenging robotic manipulation tasks, as it requires robust visual understanding, long-horizon planning, and the ability to handle deformable objects. PiO (Black et al.) a RBM demonstrated the capabilities of completing this task. Although the policy weights have been open-sourced, the embodiment gap and the lack of high-quality datasets still makes it difficult to reproduce or further improve upon their results, For the remainder of this section, all policies we report are fine-tuned from Pio.
Dataset.
We collect large dataset Dau comprising 200 hours of T-shirt folding demonstrations using the GELLO teleoperation system (Wu et al. 2024) with YAM 7-DoF bimanual robotic From this we derive smaller subset Dzmin by filtering trajectories based task duration,
8

---

Published as a conference paper at ICLR 2026
Table 3: Success rates (SR) of T-shirt folding policies 20K and 40K training steps. Each block
reports the overall SR for each task. Detailed per-color results are provided in Table
Training Steps
Tasks
(1) Da
(2) Dzmin
(3) ReWiND
(4) SARM
Simple
12/12
12/12
12/12
12/12
20K
Medium
0/12
4/12
1/12
7/12
Hard
0/12
1/12
1/12
6/12
Simple
12/12
12/12
12/12
12/12
40K
Medium
1/12
7/12
6/12
10/12
Hard
0/12
0/12
3/12
8/12
retaining only those completed within minutes, resulting in a 20-hour dataset. Each demonstration follows structured procedure: (I) picking a randomly crumpled T-shirt from box, (2) flattening the T-shirt, (3) folding the T-shirt, and (4) placing it neatly in the corner. To encourage general- ization, randomize T-shirt color and texture as well as the background environment. Aside from trajectory duration, however, no direct quantitative index is available to measure demonstration qual- ity. Furthermore, annotations not explicitly incorporated during policy training.
Tasks and Evaluation Protocol. To conduct a more detailed evaluation of the trained T-shirt fold- ing policy, we decompose the task into three sub-tasks of increasing difficulty: (1) Easy: picking the shirt from the box and placing it at the center of the table, (2) Medium: folding the T-shirt start- ing from a flattened state, and (3) Hard: completing the full pipeline from a crumpled initial state. Task is relatively simple, requiring only picking and placing skills, with human demonstrations typically lasting within seconds. Task 2 demands contact-rich manipulation of deformable objects and long- aag-horrzon planning, with human demonstrations lasting 30 seconds to minute. Task 3 is the most challenging, as it includes the flattening stage. Here, the robot must rely on strong visual un- derstanding handle occlusions and uncertainties inherent in deformable object manipulation. The policy must judge whether the T-shirt is sufficiently flattened, as failure to do so would compromise the subsequent folding step. Human demonstrations for Task typically range from 1 to minutes. For evaluation, we test each task using three different colored T-shirts (red, black, and blue). Each defined as
task is rolled out 4 times per color, for total of 12 trials per task. Success criteria follows: for Task 1, the T-shirt must be picked from the box and placed at the table center within minute; for Task 2, the T-shirt must be neatly folded and placed the corner within 3 minutes; and for Task 3, the T-shirt must be neatly folded and placed at the corner within 5 minutes.
Training Methods. We fine-tune four policies total: (1) BC-All, trained on the full dataset Dan using standard behavior cloning; (2) BC-2min, trained on the filtered high-quality subset Dzmin using standard behavior cloning, (3) RA-BC-ReWiND
trained on Dan using RA-BC with a reward model trained by the WeWIND baseline (Zhang et al. 2025); and (4) RA-BC-SARM, trained on Dan using RA-BC with our proposed reward model, SARM. It is to be noted that (1) and (2) are trained without any reweighting. They serve baselines/ablations plain behavior cloning on the diverse dataset Pall and on the naively filtered subset D2min The policy training details are listed in Appendix A.7 We evaluate policies both 20k and 40k training steps and report their success rates for each task. The experiment results can be found at Table 4.2
illustrate the performance of
All policies achieve high success rates on the easy task (picking and placing the T-shirt), indicating that both the dataset and training procedure are sufficient for learning basic manipulation skills. On the medium task (folding from flattened state), the policy trained on Dzmin substantially outper- forms the one trained on Dan. with its success rate at 40k steps improving from near 0% to over 50%. This highlights the importance of carefully filtered, high-quality data for learning more complex ma- nipulation behaviors. Nevertheless, this policy still fails the hard task (folding from crumpled state), indicating that filtering by duration alone is insufficient. Such a naive strategy cannot empha- size demonstrations that require advanced perception and decision making, such judging whether the T-shirt has been adequately flattened before folding. As a result, it fails to deliver the dynamic and contact-rich manipulation needed for tasks that demand seamless coordination of both arms.
By leveraging SARM, the RA-BC policy surpasses both BC baselines by a significant margin on medium and hard tasks both 20k and 40k steps. In particular, the RA-BC policy at 40k steps

---

Published as a conference paper at ICLR 2026
achieves 83% success rate the medium task and 67% success rate the hard task. These results demonstrate that RA-BC effectively exploits diverse datasets by filtering high-quality data frames, enabling the policy to learn robust long- long-horizon manipulation strategies.
4.3
Q3: EFFECT REWARD MODEL QUALITY IN RA-BC
To investigate how the quality of the reward model influences RA-BC performance, we conduct an ablation study by training the T-shirt folding policy using RA-BC with the baseline WeWiND (Zhang et al. 2025) reward model. The evaluation results of reward models are presented in Table and the corresponding policy performance is summarized in Table 4.2 Compared to RA-BC with SARM, RA-BC with the ReWiND reward model achieves substantially lower success rates both the medium (83% v.s. 50%) and hard tasks (67% v.s. 25%). In particular, on the medium task, its performance drops to the level of the vanilla BC baseline trained on the filtered dataset Dzmin-
These results highlight the central role of reward model quality in RA-BC. A reliable model accu- rately captures task progress, enabling effective filtering of demonstrations and consistent supervi- sion for policy learning. In contrast, a poor model misjudges progress, misweights data, and weakens the benefits of filtering- This reliability especially crucial in long-horizon, multi-stage tasks like T-shirt folding, where failures and partial progress are common.
We further explore the of SARM in reinforcement learning (RL) by training DiffQL (Wang et al. 2022) manipulation policy with reward signals provided by SARM. The detailed methodology, experimental setup, and results are presented in Appendix
A.8 A.8
CONCLUSION
This paper explored how demonstration quality shapes the effectiveness of RBMs when tackling complex, long-horizon manipulation. Using T-shirt folding as demanding case study, we showed that naively scaling dataset size is insufficient, and that progress-aware supervision needed to guide learning. To this end, we designed SARM, a stage-aware, video-based reward modeling framework that transforms natural language annotations into structured progress signals, enabling more reliable estimation of task advancement across diverse demonstrations. We further introduced the RA-BC framework, which incorporates these signals to emphasize higher-value trajectories dur- ing training. Our empirical evaluation revealed clear benefits: SARM consistently surpassed prior baselines, and policies trained with RA-BC achieved strong performance on real robots, including an 83% success rate when folding T-shirts from flattened configuration and 67% when starting from crumpled state. Additional analysis showed that the accuracy of the reward model is pivotal-when when the reward signal is weak, RA-BC loses its ability properly weight samples and overall policy performance degrades. We also demonstrate that SARM can be incorporated into reinforcement learning (RL) further improve policy performance by modifyyng DifiQL Wang et al. (2022) on a pick-and-place task in the MuJoCo Todorov et al.] (2012) simulation environment. Details are provided in A.8 These findings underscore that high-quality reward modeling, combined with se- lective data filtering, is a powerful path forward for building robust and scalable RBMs capable of addressing long-horizon manipulation challenges.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Pieter Abbeel and Andrew Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference Machine learning, pp. 2004.
Christopher Agia, Rohan Sinha, Jingyun Yang, Rika Antonova, Marco Pavone, Haruki Nishimura, Masha Itkina, and Jeannette Bohg. Cupid: Curating data your robot loves with influence func- tions. arXiv preprint 2rX6v:2506.191 2025.
Minttu Alakuijala, Reginald McLean, Isaac Woungang, Nariman Farsad, Samuel Kaski, Pekka Marttinen, and Kai Yuan. Video-language critic: Transferable reward functions for language- conditioned robotics. arXiv preprint arXiv:2405. 19988, 2024.
Jose Barreiros, Andrew Beaulicu, Aditya Bhat, Rick Cory, Eric Cousineau, Hongkai Dai, Ching- Hsin Fang. Kunimatsu Hashimoto, Muhammad Zubair Irshad, Masha Itkina, et al. care- ful examination of large behavior models for multitask dexterous manipulation. arXiv preprint arXiv:2507.05331, 2025.
Suneel Belkhale, Yuchen Cui, and Dorsa Sadigh. Data quality in imitation learning. Advances in neural information processing systems, 36:80375-80395, 2023.
Kevin Black, Noah Brown, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, Lachy Groom, Karol Hausman, Brian Ichter, et al. T0: A Aisionlanguage-action flow model for general robot control. corr, abs/2410.24164, 2024. doi: 10.48550. arXiv preprint ARXIV.2410.24166
Qianzhong Chen, Naixiang Gao, Suning Huang, JunEn Low, Timothy Chen, Jiankai Sun, and Mac Schwager. Grad-nav++: Vision-language model enabled visual drone navigation with gaussian radiance fields and differentiable dynamics. arXiv preprint arXiv:2506.14000 2025.
Xinyue Chen, Zijian Zhou, Zheng Wang, Che Wang, Yanqiu Wu, and Keith Ross. Bail: Best- action imitation learning for batch deep reinforcement learning. Advances in Neural Information Processing Systems, 33:18353-18363 2020.
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The Inter- national Journal of Robotics Research, pp. 02783649241273666 2023.
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances neural information processing sys- tems, 30, 2017.
Dima Damen, Hazel Doughty, Giovanni Maria Farinella, Antonino Furnari, Evangelos Kazakos, Jian Ma, Davide Moltisanti, Jonathan Munro, Toby Perrett, Will Price, et al. Rescaling egocen- tric vision: Collection, pipeline and challenges for epic-kitchens-100 International Journal of Computer Vision, 130(1):33-55, 2022.
Shivin Dass, Alaa Khaddaj. Logan Engstrom, Aleksander Madry, Andrew Ilyas, and Roberto Martín-Martin Datamil: Selecting data for robot imitation learning with datamodels. arXiv preprint ar (xXx:25055096// 2025.
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. Advances in Neural Information Processing Systems, 36:68760-688 68783, 2023.
Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. In International conference on machine learning, pp. 49-58. PMLR, 2016.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3 2022.
Huang Huang, Fangchen Liu, Letian Fu, Tingfan Wu, Mustafa Mukadam, Jitendra Malik, Ken Gold- berg, and Pieter Abbeel. Otter: A vision-land aaaioon-annugg--ctttin model with text-aware visual feature extraction. arXiv preprint arXin:2503.03734 2025a.

---

Published as a conference paper at ICLR 2026
Suning Huang, Zheyu Zhang, Tianhai Liang, Yihan Xu, Zhehao Kou, Chenhao Lu, Guowei Xu, Zhengrong Xue, and Huazhe Xu. Mentor: -ixture-of-experts network with task-oriented pertur- bation for visual reinforcement learning, arXiv preprint arXlv:2410.14972, 2024.
Suning Huang, Qianzhong Chen, Xiaohan Zhang, Jiankai Sun, and Mac Schwager. Particleformer: A 3d point cloud world model for multi-object, multi-material robotic manipulation. arXiv preprint 2..i:::200..322222 2025b.
Kuo-Han Hung, Pang-Chi Lo, Jia-Fong Yeh, Han-Yuan Hsu, Yi-Ting Chen, and Winston H Hsu. Victor: Learning hierarchical vision-instruction correlation rewards for long-horizon manipula- tion. arXiv preprint arXiv:2405.16555 2024.
I2RT-Robotics. Yam
6-dof robotic arm
attps://2rtt.com/prroom
yam- manipulator 2025.
Changyeon Kim, Minho Heo, Doohyun Lee, Jinwoo Shin. Honglak Lee. Joseph J Lim, and Kimin Lee. aawre visual reward learning from segmented demonstrations. arXiv preprint arXiv:2502.206600 2025.
Youngwoon Lee, Andrew Szot, Shao-Hua Sun, and Joseph Lim. Generalizable imitation learn- ing from observation via inferring goal proximity. Advances in neural information processing systems, 34:16118-16100, 2021.
Fanqi Lin, Yingdong Hu, Pingyue Sheng, Chuan Wen, Jiacheng You, and Yang Gao. Data scaling laws imitation learning for robotic manipulation. arXiv preprint arXiv:2410. 18647, 2024.
Runze Liu, Yali Du, Fengshuo Bai, Jiafei Lyu, and Xiu Li. Pearl: Zero-shot cross-task preference alignment and robust reward learning for robotic manipulation. arXiv preprint arXiv:2306.03615. 2023.
Jason Ma, William Liang, Hung-Ju Wang, Yuke Zhu, Linxi Fan, Osbert Bastani, and Dinesh Jayara- man. Dreureka: Language model guided sim-to-real transfer. RSS, 2024a.
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. arXiv preprint arXiv:2210.0005: 2022.
Yecheng Jason Ma. Vikash Kumar, Amy Zhang, Osbert Bastani, and Dinesh Jayaraman. Liv: Language-iman image representations and rewards for robotic control. In International Conference on Machine Learning. pp. 23301-23320. PMLR, 2023.
Yecheng Jason Ma, Joey Hejna, Chuyuan Fu, Dhruv Shah, Jacky Liang, Zhuo Xu, Sean Kirmani, Peng Xu, Danny Driess, Ted Xiao, et al. Vision language models are in-context value learners. In The Thirleenth International Conference Learning Representations, 2024b.
Tongzhou Mu, Minghua Liu, and Hao Su. Drs: Learning reusable dense rewards for multi-stage tasks. arXiv preprint arXiv:2404. 16779, 2024.
Andrew Y Ng, Stuart Russell, et al. Algorithms for inverse reinforcement learning. In Ieml, vol- ume 1, pp. 2. 2000.
Deepak Ramachandran and Eyal Amir. Bayesian inverse reinforcement learning. In IJCAI, vol- ume 7, pp. 1586625591 2007.
RealSense.
Realsensedepth camera d405.
com/
buy happ-iittt---ealeee---ed-------- cane ra d405. html 2025.
Dorsa Sadigh, Anca Dragan, Shankar Sastry, and Sanjit Seshia. Active preference-based learning of reward functions. 2017.
Fumihiro Sasaki and Ryota Yamashina. Behavioral cloning from noisy demonstrations. Interna- fional Conference on Learning Representations, 2020.
12

---

Published as a conference paper at ICLR 2026
Yide Shentu, Philipp Wu, Aravind Rajeswaran, and Pieter Abbeel. From Ilms to actions: Latent codes as bridges in hierarchical robot control. In 2024 JEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 8539-8548 8546. IEEE, 2024.
Noah Y Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Ne- unert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavioral modelling priors for offline reinforcement learning. arXiv preprint arXiv:2002.08336, 2020.
Jiankai Sun, Aidan Curtis, Yang You, Yan Xu, Michael Koehle, Leonidas Guibas, Sachin Chitta, Mac Schwager, and Hui Li. Hierarchical hybrid learning for long-horizon contact-rich robotic assembly. arXiv preprini arXiv:2409. 16451, 2024.
Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et al. Octo: An open-source generalist robot policy, arXiv preprint 24Xx::2405.12215 2024.
Emanuel Todorov, Tom Frez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 JEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033.
IEEE, 2012.
Chen Wang, Linxi Fan, Jiankai Sun, Ruohan Zhang, Li Fei-Fei, Danfei Xu, Yuke Zhu, and An- ima Anandkumar. Mimicplay: Long-horizon imitation learning by watching human play. arXiv preprint arXiv:2302.12-422 2023a.
Qiang Wang, Robert McCarthy, David Cordova Bulens, Kevin McGuinness, Noel E O'Connor, Francisco Roldan Sanchez, Nico Gürtler, Felix Widmaier, and Stephen J Redmond, Improving behavioural cloning with positive unlabeled learning. In Conference on robot learning. pp. 3851- 3869. PMLR, 2023b.
Qing Wang, Jiechao Xiong, Lei Han, Han Liu, Tong Zhang, et al. Exponentially weighted imitation learning for batched historical data. Advances in Neural Information Processing Systems, 31, 2018.
Zhendong Wang, Jonathan Hunt, and Mingyuan Zhou, Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.00 06193, 2022.
Philipp Wu, Yide Shentu, Zhongke Yi, Xingyu Lin, and Pieter Abbeel. Gello: A general, low-cost, and intuitive teleoperation framework for robot manipulators. 2024 IEEE/RS. International Conference on Intelligent Robots and Systems (IROS), pp. 12156-12163, IEEE, 2024.
Yueh-Hua Wu, Nontawat Charoenphakdee, Han Bao, Voot Tangkaratt, and Masashi Sugiyama. Imi- tation learning from imperfect demonstration. International Conference on Machine Learning. pp. 0818-6827. PMLR, 2019.
Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator- offline imitation learning from suboptimal demonstrations. In International Conference on Machine Learning, pp. 14725-24744 PMLR, 2022.
Justin Yu, Tara Sadjadpour, Abby O'Neill, Mehdi Khfifi, Lawrence Yunliang Chen, Richard Cheng, Muhammad Zubair Inshad, Ashwin Balakrishna, Thomas Kollar, and Ken Goldberg. Manip:
A modular architecture for integrating interactive perception for robot manipulation. In 2024 JEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 1283-1289. IEEE, 2024.
Jiahui Zhang, Yusen Luo, Abrar Anwar, Sumedh Anand Sontakke, Joseph J Lim, Jesse Thomason, guided rewards teach robot policies without
Erdem Biyik, and Jesse Zhang. Rewind: new demonstrations. arXiv preprini arXiv:2505. 10911. 2025.
Tony Z. Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:2304. 13700, 2023.
Brian Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pp. 1433-1438. Chicago, II., USA, 2008.
13

---

Published as a conference paper at ICLR 2026
Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford. Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.0855 2019.
Brianna Zitkovich, Tianhe Yu, Sichun Xu, Peng Xu, Ted Xiao, Fei Xia, Jialin Wu, Paul Wohlhart, Stefan Welker, Ayzaan Wahid, et al. Rt-2: Vision-language-actioo models transfer web knowledge to robotic control. In Conference on Robot Learning, pp. 2165-2183, PMLR, 2023.
14

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 LLM USAGE DISCLOSURE.
In accordance with the submission policy, we disclose that large language models (LLMs) were used only for minor wording improvements, grammar checking, and proofreading. All ideas, technical content, experimental design, and analysis were solely developed by the authors, who take full responsibility for the final manuscript.
A.2 HARDWARE SETUP
For our real world experiments leverage a bimanual robot table top platform. The system consists of:
Two 6 DOF YAM robot arms, built by the manufacturer 12RT (12T Robotics 2025)_
Three RealSense D405 cameras. One for each wrist and third stat- ically mounted above for viewing the scene (RealSense) 2025).
Data is collected using leader follower system GELLO teleoperation system Wu al. 2024). The environment run recorded at 30 fps and includes synchro- nized streams from three cameras (left wrist, right wrist, and fixed top view) along with robot joint angles and action joint angle commands.
Figure The physical station used for data collection and policy evaluation.
A.3 T-SHIRT FOLDING EXPERT DEMONSTRATION DATA
We provide examples of demonstration trajectories collected with the GELLO system in Fig. and Fig. 6 For clarity, the annotation corresponding to each motion stage is shown above every frame. Compared Dssarse. the dense annotated dataset Deense further decomposes the overall "fold the T-shirt" stage into five fine-grained motions. A comparison of the average temporal portion of each task in the two datasets summarized Table A.3
important note that Dspars and Doense are distinct datasets collected from different trajectories; therefore, even for the same task (e.g., "flatten the T-shirt out"), the average temporal portion differs a little across datasets.
Table Average temporal portion of each task in two dataset.
Sparse Annotated Dataset D sparse Task
Dense Annotated Dataset D dense Task
Portion (%)
Portion (%)
Grab the T-shirt from the pile
Grab T-shirt and move to center
9
Move the T-shirt to the center
Flatten out the T-shirt
26
Flatten the T-shirt out Fold the T-shirt Put folded T-shirt into corner
25
Grab near side and fold
15
55
Grab far side and fold
13
10
Rotate the T-shirt 90 deg
Grab bottom and fold
9
Grab 2/3 side and fold
Put folded T-shirt into corner
11
15

---

Published as a conference paper at ICLR 2026
grab the tshirt from the pile
move tshirt to center
flatten tshirt out
fold the tshirt
put folded tshirt into comer
Figure Expert demonstration with sparse annotation.
grah move
grah side
grab fold
rotate the tshirt 90 deg
gmb bottom and Fold
grab 2/3 side and fold
put folded tshirt into corner
Figure 6: Expert demonstration with dense annotation.
16

---

Published as a conference paper at ICLR 2026
A.4 REWARD MODEL TRAINING
Implementation Details. We employ a frozen cllp-vit--oe - encoder to process both RGB image sequences and task descriptions. The dataset is recorded at fixed frame rate of 30 fps, and each input sequence consists of images: the first is always the initial frame of the episode, while the remaining 8 are sampled consecutively from the same episode with a fixed interval of 30 frames, resulting in a temporal span of approximately 8 seconds. To enhance temporal diversity and better capture failure scenarios, we follow the rewind augmentation strategy (Zhang et al. 2025). appending up to four frames from earlier timestamps with reversed order to the end of each training sequence, Additionally, to improve video-language alignment, the task descriptions are occasionally perturbed with randomly generated incorrect instructions.
The backbone is a transformer- based temporal aggregator with layers, attention heads, and a hidden dimension of 768. To mitigate information leakage, positional embeddings applied only to the first frame, corresponding to the episode start. On top of the backbone, we incorporate twin MLP-based output heads tailored for different annotation types, namely dense and sparse labels. This design enhances the flexibility of the reward model, allowing it to effectively utilize hetero- geneous supervision and remain compatible with multiple annotation protocols. Each output head comprises 2 layers with a hidden dimension of 512. The stage model is trained with cross-entropy loss, whereas the subtask model optimized with mean squared error loss.
Optimization is performed with the Adam optimizer, using a learning rate of 5 x 10-5 and a weight decay of 10-3 Models are trained for 2 epochs with batch size of 64 on a single NVIDIA RTX 4090 GPU.
Scale Analysis. We study the effect of model scale on reward model performance by varying the number of transformer layers in the temporal aggregator from 4, 8, to 12, corresponding to models with 30M, 60M, and 90M parameters, respectively, while keeping all other hyperparameters fixed. The results summarized in Table 5 and Fig. The smallest model (30M) exhibits clear underfitting, with poor performance across both evaluation metrics. Increasing the size from 30M to 60M leads to substantial gains, but further scaling from 60M to 90M yields only marginal or negligible improvement. This suggests that a 60M-parameter model is sufficient to capture the task dynamics of T-shirt folding. Larger models risk overfitting, particularly given the limited size of the training data. Overall, the chosen 60M configuration strikes an effective balance between model capacity and computational efficiency, and is therefore adopted throughout the paper.
0.5
Ieyer
0.4
Jayens
Jayers 12
0.3
1,08
0.2
0.1
0.0
10k
20k
30k
40k
50k
Stup
Figure 7: Scale analysis plots of reward models with various layers.
Table 5: Scalability analysis of reward model on T-shirt folding task
Layer Number 8
Metrics
Demo L
0.015
0.009
0.007
Rollout pt
0.72
0.94
0.88
Classification SR Breakdown
SE
10/12
12/12
12/12
PSE
10/12
11/12
11/12
FE
11/12
12/12
11/12
17

---

Published as a conference paper at ICLR 2026
Ablation Study. We perform comprehensive ablation study to evaluate the impact of different design choices in reward model training. Specifically, we examine: (1) the use of joint state as an additional input, (2) the inclusion of wrist cameras, (3) the number of observation steps, and (4) the frame gap between observation steps. The corresponding training loss curves are shown in Fig. g
Joint state imput: We compare variants of the reward model: joint state (SARM's default configuration) and one that relies only on visual observations. As shown in T including joint state leads to more accurate estimation on both human demonstrations and policy rollouts. The improvement, however, is relatively modest, suggesting that visual input already contains most of the task-relevant information.
that incorporates the robot's
Wrist cameras: We evaluate the effect of adding wrist camera views in addition to the fixed top-down camera. The results in Table show little no benefit, likely because the top-down perspective already provides sufficient task coverage, while wrist cameras contribute redundant information. Moreover, incorporating wrist views increases system complexity and introduces threefold I/O cost. For these reasons, we exclude them from the final design.
Number of observation steps: We vary the number of observation steps from 4, 8 (SARM's default), to 12, while keeping the temporal horizon fixed at approximately 8 seconds. Results in Table indicate that too few steps (4) limit the model's 's ability to capture temporal dynamics, resulting in underfitting. On the other hand, too many steps (12) introduce redundancy and additional compu- tational overhead without clear gains. An intermediate choice of steps provides a good balance between temporal coverage and efficiency.
Frame gap between steps: We also assess the effect of varying the frame gap between consecutive observation steps at 15, 30 (SARM's default), and 60 frames. As shown in Table p small gap of 15 frames (0.5s) produces highly correlated inputs, limiting temporal diversity and shortening the effective temporal span, which leads to underfitting. A large gap of 60 frames (2s) risks missing important intermediate states, thereby confusing the model and degrading performance. A moderate gap of 30 frames (Is) achieves the best trade-off by capturing meaningful temporal transitions while avoiding redundancy.
In summary, these results underscore the importance of balancing model capacity, input modalities, and temporal resolution. Our final design choices susiice--ssingg joint state input, excluding wrist cameras, adopting observation steps, and a frame gap of 30 caanes-reflee the insights gained from this ablation analysis and provide robust configuration for long-horizon reward modeling.
Joint State Ablation
Wrist Camera Ablaton
0.5
0.5
wio Stute
Canera
Bandin:
Btaclins
0.1
0.1
0.3
0.3
0.2
0.2
0.1
0.1
0.0
0.0
10k
20%c
30k
40k
50k
10k
20kk
30k
40k
50k
Step
Step
10bbs Steps Ablation
Giap Ablition
0.5
0.5
fomee 11
Sune. .xp
0.4
13
Suan ssp-60
0.3
0.0
0.0
10%
20%
306
-0k
50k
10k
20k
30k
40k.
50k
Stup
Step
Figure 8: Ablation study training loss curves
18

---

Published as a conference paper at ICLR 2026
Table 6: Ablation study of T-shirt folding reward model on using joint state (our choice: Yes).
Using Joint State (our choice: Yes).
Metrics
Yes
No
Demo L1
0.010
0.009
Rollout p
0.72
0.94
Classification SR Breakdown
SE
12/12
12/12
PSE
9/12
11/12
FE
10/12
12/12
Table 7: Ablation study of T-shirt folding reward model on using wrist cameras (our choice: No).
Using Wrist Cameras
Metrics
Yes
No
Demo L
0.008
0.009
Rollout p
0.94
0.94
Classification SR Breakdown
12/12
12/12
PSE
11/12
11/12
12/12
12/12
Table 8: Ablation study of T-shirt folding reward model on observation steps number (our choice: 8).
Observation Step Number
Metrics
12
Demo L 0.013
0.009
0.009
Rollout p1 f
0.67
0.94
0.89
Classification SR Breakdown
SE
12/12
12/12
12/12
PSE
8/12
11/12
11/12
FE
10/12
12/12
11/12
Table 9: Ablation study of T-shirt folding reward model on sequence frames gap (our choice: 30).
Frame Gap
Metrics
30
60
Demo L
0.015
0.009
0.022
Rollout pt
0.50
0.94
0.56
Classification SR Breakdown
SE
9/12
12/12
12/12
PSE
8/12
11/12
8/12
FF
10/12
12/12
8/12
19

---

Published as a conference paper at ICLR 2026
A.5 REWARD MODEL EVALUATION RESULTS VISUALIZATION
Demo Data Estimation. We present two visualization examples of reward predictions from SARM and the ReWIND baseline in Fig- 9 and Fig. TO using trajectories from the validation of human demonstration data. Compared with SARM, ReWiND exhibits several notable shortcom- ings: (1) relies solely on direct regression, fails capture the full progression of long-horizon tasks-foo for instance, its predictions do not start at zero even at the beginning of trajectory; and (2) its estimates are highly unstable, with frequent oscillations and even negative spikes, which should not occur in human demonstration data. These issues prevent ReWiND from producing consistent long-horizon reward signals and ultimately limit its effectiveness for downstream policy learning.
1.0
Predicted Progress
0.5
Time 20
30
1.0
Predicted Progress
0.0
30
Time (s)
Figure 9: Examples of proposed reward model prediction on demonstration data.
1.0m
Predicted Progress
0.5
0.0
30
(6)
1.00
Predicted Progress
0.5
Figure 10: Examples of ReWiND reward model prediction on demonstration data.
20

---

Published as a conference paper at ICLR 2026
Policy Rollout Estimation. We present two visualization examples of reward predictions from SARM and the ReWiND baseline Fig. and Fig. 12 using trajectories from real robot policy rollouts. Compared with human demonstration data, policy rollouts more challenging because they often include failure modes that are out-of-distribution (OOD), such as misgrasps, recovery attempts, and back-and-forth motions. In the first example, the trajectory corresponds to successful rollout where the robot folds the T-shirt correctly, with only minor struggles and misgrasps in the first ten seconds. In this case, SARM remains stable, keeping the estimated progress near zero during these OOD motions, whereas ReWiND is easily triggered and produces noisy, unstable estimates. The second example highlights failed rollout, with four key frames: (I) the T-shirt is flattened after struggling, (2) folding is nearly complete, (3) the robot suddenly fails and crumples the T-shirt on the table, and (4) the unfolded T-shirt is placed in the corner. SARM provides reasonable progress estimates across all four stages, reflecting the actual task status. By contrast, ReWiND continues to exhibit high noise and spurious spikes, and even assigns a high progress score to the final "fake finish" state, effectively being misled by the failed outcome. These results further emphasize the robustness and reliability of SARM framework for real-world robotic applications.
Pholicad Prognce
0.
30
Progecss
0.5
50
150
Time
Figure 11: Examples of proposed reward model prediction on policy rollouts.
Paaliced Pangnsh
ww
30
50
Podiiced Progros
Figure 12: Examples of ReWiND reward model prediction on policy rollouts.
21

---

Published as a conference paper at ICLR 2026
A.6 TRAINING SARM FOR DISH UNLOADING
Demo Data Estimation. The visualization of reward predictions from SARM and the ReWiND baseline two example trajectories from the validation set of human demonstration data is shown in Fig. 13 and Fig. [4] SARM produces consistent and robust progress estimates, maintaining sta- ble predictions for sequences as long as unloading eight dishes consecutively, which corresponds to over 1.5 minutes of execution. This demonstrates the effectiveness of SARM handling diverse and highly dynamic tasks. By contrast, ReWiND exhibits similar shortcomings folding experiments: it fails to capture the full progression of the task (in this case never estimat- ing completion, with peak values below 0.75) and generates unstable predictions with noticeable fluctuations.
in the T-shirt
Predicted Progress
0.0
10
Time (4)
1.00
Predicted Progress
0.50
0.00
100
Time (6)
Figure 13: Examples of proposed reward model prediction demonstration data.
22

---

Published as a conference paper at ICLR 2026
1.0
Predicted Progress
0.5
0.0
10
20
30
40
Time (5)
Presicted Progress
0.75
0.50
0.00
100
Time (6)
Figure 14: Examples of ReWind reward model prediction on demonstration data.
23

---

Published as a conference paper at ICLR 2026
A.7 TRAINING MANIPULATION POLICY WITH RA-BC
Implementation Detail. All policies are fine-tuned with low rank adaptation (LoRA) (Hu et al, 2022) for 40k steps using batch size of 32 a dual NVIDIA RTX 4090 machine. The RA-BC hyperparameters are set to K = 0.01 and e = 10 with a chunk length of A 25 actions to align with the policy's action chunking. Since the dataset is recorded 30 fps, K 0.01 corresponds roughly to task duration of minute 30 seconds, which represents the threshold for the top 5% of demonstrations. For data points better than this threshold, we assign weight of For data points that are worse than this threshold but still demonstrate positive progress, we assign a soft weight between and according to Eq. g For data points exhibiting negative progress, the assigned weight is O.
Loss Curve.
The policy training loss curves for all four methods are shown in Fig. 15 We observe that the two pure BC methods initially exhibit faster decrease in loss compared to RA-BC, but they plateau early and converge to a higher final loss value. In contrast, RA-BC methods display slower but more consistent reduction loss and ultimately achieve lower convergence values. This phenomenon can be explained by the data distribution: pure BC leverages a broader set of demonstrations, which allows an unconverged policy quickly match parts of the dataset, resulting in smaller loss at the early stages of training. However, this diversity also introduces conflicting gradient signals that prevent further improvement, causing convergence at a suboptimal plateau. the other hand, RA-BC employs more focused learning objective that emphasizes high-quality data. Although such data is harder to fit initially, the targeted supervision enables the policy to continue improving and avoid stagnation. The final converged loss values are consistent with the policy evaluation results in Table A.7 where RA-BC with our SARM framework delivers the best overall performance.
0.10
BC-All
0.09
BC-2min
0.08
RA-BC-ReWiND
RA-BC-SARM
0.07
Loss
0.06
0.05
0.04
0.03
0.02
5k
10k
15k
20k
25k
30k
35k
40k
Step
Figure 15: Trained T-shirt folding policies loss curves.
24

---

Published as a conference paper at ICLR 2026
Table 10: Success (SR) of T-shirt folding policies 20K and 40K training steps. For each method, the first column reports SR for each T-shirt color (R = red, B blue, K = black), and the last row in each block shows the overall SR for that task.
Training Steps
Tasks
Color
(1) Da
(2) Damin
(3) ReWiND
(4) SARM
R
4/4
4/4
4/4
4/4
B
4/4
4/4
4/4
4/4
Simple
4/4
4/4
Overall
12/12
12/12
12/12
12/12
0/4
0/4
20K
B
0/4
1/4
1/4
2/4
Medium
K
0/4
0/4
0/4
2/4
Overall
0/12
4/12
1/12
7/12
R
0/4
1/4
1/4
2/4
B
0/4
0/4
0/4
4/4
Hard
0/4
0/4
Overall
0/12
1/12
1/12
6/12
4/4
4/4
B
4/4
4/4
4/4
4/4
Simple
4/4
4/4
4/4
4/4
Overall
12/12
12/12
12/12
12/12
R
0/4
4/4
2/4
4/4
40K
B
1/4
1/4
3/4
4/4
Medium
K
0/4
2/4
1/4
2/4
Overall
1/12
6/12
10/12
R
0/4
0/4
2/4
2/4
0/4
0/4
Hard
0/4
0/4
1/4
2/4
Overall
0/12
0/12
3/12
8/12
Rollout Example. An example policy rollout is shown in Fig. 16 where the robot successfully
folds T-shirt from a crumpled state into neat configuration within 90 seconds, without any mis- grasps.
2
3
4
7
8
Figure 16: Example of RA-BC trained T-shirt folding policy rollout.
25

---

Published as a conference paper at ICLR 2026
A.8 REINFORCEMENT LEARNING EXAMPLE.
Beyond RA-BC, we also explored integrating SARM with reinforcement learning (RL) to further improve policy performance. We adopt two-stage training scheme: (1) pre-training, where the policy is trained with pure behavior cloning (BC) using the diffusion policy (Chi et al. 2023); and (2) fine-tuning, where the policy is refined with DiffQL (Wang et al. 2022). a Q-learning method specifically designed for diffusion-based policies. We refer to this approach as RA-QL.
We evaluate RA-QL on a simple two-stage task: pick up cube. In this task, the robot arm must first reach toward the cube on the desktop, then grasp and lift it into a goal region, which is fixed designated box space. An illustration of the task provided in Fig. 17 We collected 300 expert demonstrations in the MuloCo simulation environment (Todorov et al. 2012) with randomized cube initial position, denoted Deube. From these, 100 trajectories were annotated train reward model with the same architecture used for T-shirt folding.
Figure 17: Expert demonstration of the cube-picking task. The desired goal region highlighted in green.
We first pre-trained diffusion transformer policy on Peube using pure BC. While the policy achieved high success rates, the motions were often inefficient and imprecise, especially during the reaching and grasping phases. We then fine-tuned this BC-trained policy with DiffQL. for an additional 10k steps. Our implementation closely follows DifiQL (Wang al. 2022), with the key modification that the critic network receives rewards from our learned reward model, SARM, rather than hand- crafted signals. For ablation, we also continued training the diffusion policy with pure BC for another 10k steps under the same conditions.
During fine-tuning, evaluated both the BC and RA-QL policies every 500 steps. Each policy was rolled out 10 times with randomized cube positions, and we report the average success rate (SR) and average discounted return. The experiment results are demonstrated in Fig. T8 The reward function automatically judged by the simulator: a step reward of 1 given if the eube is lifted to the desired height, and 0 otherwise. An episode terminates either when the cube is successfully placed in the goal region (labeled as success) or when the horizon of 1000 steps is reached (labeled as faill 1).
T-1-1
G
Tetk:
(11)
where G, denotes the discounted return from time step and y 0.995 the discount factor.
Suceess Rate Comparison
Disevanted Return Comparison
1.1
0.85
RA-QI.
RAA-QI.
1.0
BC
0.80
pure
(2)0
Return
0.75
0.70
0.8
0.65
0.7
0.60
2000
4000
60000
18000
2000
4000
6000
8000
10000
Step
Step
Figure 18: Comparison of RA-QL and pure BC on picking up cube policy training.
26


Reviews:
[{"review_id": "aWBYaX7Tm8", "reviewer": "Reviewer_pgZ9", "summary": "In this work, the authors introduce a stage-aware, video-based reward modeling framework that jointly predicts high-level task stages and fine-grained progress within each stage. This design effectively addresses the limitations of frame-index-based labeling, which often fails in long, variable-duration tasks such as T-shirt folding. The learned reward models are further utilized to filter high-quality data and reweight training samples based on the estimated rewards. Extensive experiments demonstrate that the proposed reward model outperforms existing baselines in out-of-distribution real-robot rollouts and human demonstration evaluations.", "questions": "1. Generalization: How well does the proposed reward model generalize when applied to large-scale, real-world cloth manipulation datasets such as DROID and Flat’n’Fold?\n\n2. Robustness: How robust is the model to lighting and illumination changes, which commonly occur in real-world robotic settings?\n\n3. Comparative Evaluation: How does the proposed method compare with VICtoR (ICLR 2025) in terms of reward modeling effectiveness and policy performance on long-horizon, contact-rich manipulation tasks?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. The authors investigate a very important and timely direction, i.e., how to acquire high-quality data for training imitation learners on long-horizon, contact-rich manipulation tasks such as T-shirt folding. Quantifying demonstration quality is challenging because it depends on latent factors such as action consistency and contact stability, which cannot be directly measured. In this work, the authors propose a video-based reward modeling framework that leverages natural language annotations to assign progress labels and enable stable reward estimation, serving as a mechanism to filter high-quality data and improve policy performance in both simulation and real-world settings. The proposed method is timely and presents a novel idea for learning from suboptimal demonstrations, making it a valuable contribution to the community.\n\n2. The authors provide a comprehensive analysis of the proposed SARM reward model, comparing it against existing methods. The proposed approach demonstrates superior performance across multiple benchmarks. Furthermore, the authors evaluate the framework on  dish unloading, which further validates its effectiveness and generalization capability.\n\n3. By leveraging SARM, the RA-BC policy surpasses both BC baselines by a significant margin on medium and hard tasks. The experiments show that RA-BC effectively exploits diverse datasets by filtering out high-quality data frames, enabling the policy to learn robust, long-horizon manipulation strategies.\n\n4. The authors empirically demonstrate that RA-BC with SARM achieves substantially higher success rates than RA-BC with ReWiND, achieving 83% vs. 50% on medium tasks and 67% vs. 25% on hard tasks, clearly highlighting the advantage of the proposed reward model.", "weaknesses": "**Overall Assessment:**\n\nThe paper provides a well-justified motivation for the proposed reward model and demonstrates its utility in selecting high-quality data for policy learning. However, several aspects remain unclear and merit further clarification.\n\n1. **Applicability to existing datasets:**\n    \nThe proposed reward model has shown promise in filtering high-quality data, but it remains unclear how well it generalizes to existing cloth manipulation datasets, such as **DROID* [1] and **Flat’n’Fold** [2]. How would the model perform if applied to these datasets for filtering or reweighting demonstrations?\n \n[1] DROID: A Large-Scale In-the-Wild Robot Manipulation Dataset, 2024\n[2] Flat’n’Fold: A Diverse Multi-Modal Dataset for Garment Perception and Manipulation, 2024\n    \n2. **Robustness to visual variations:**\n    \nHow does the proposed reward model handle changes in lighting and illumination conditions? Since real-world environments are often visually diverse and dynamic, it would be valuable to understand the model’s robustness and reliability under such variations.\n    \n3. **Baseline comparison:**\n    \nThe recent work by **Hung et al. (ICLR 2025)** [3] introduces an alternative strategy for learning reward models in long-horizon manipulation tasks. Although the learning objectives differ from those of SARM, a direct comparison between the two approaches under similar long-horizon, contact-rich settings would provide stronger empirical grounding and highlight the unique advantages of the proposed framework.\n\n[3] Hung et al., VICtoR: Learning Hierarchical Vision–Instruction Correlation Rewards for Long-Horizon Manipulation, ICLR 2025", "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": "idT70dAByO", "reviewer": "Reviewer_fc5R", "summary": "A core challenge in imitation learning literature (e.g. behavior cloning) is learning from diverse demonstration data: multi-modal action distributions, varying teleoperation strategies and proficiency, human errors, as well as sometimes wildly different trajectory lengths depending on the teleoperator and initial/goal configurations. This paper aims to automatically estimate task progress in demonstration data regardless of factors such as trajectory length (recent method ReWiND is prone to fail in this setting), and then subsequently using estimated task progress (reward) to weigh samples during BC policy learning. To do so, the proposed method estimates stage and subtask progress from image observations + joint states + language annotations using a transformer backbone, as well as a number of engineering decisions and heuristics detailed in the paper. Experiments are conducted on a real-world T-shirt folding task using a dual-arm setup, and results indicate that the proposed framework (SARM / RA-BC) improves task progress estimation and downstream policy performance compared to strong baselines.", "questions": "I would really appreciate it if the authors can address my comments in the \"weaknesses\" section above using written arguments and potentially additional experimental results. My main concerns / questions pertain to writing and the data filtering process.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "My initial assessment of this paper leans positive overall; it is generally well written, timely, and technically sound. Specifically:\n- I believe that this paper studies a relevant and timely problem (learning from \"diverse\" demonstrations wrt. manipulation strategy, trajectory length etc.), and is likely to be of interest to the community. The problem is clearly motivated in the introduction, and shortcomings of existing work (ReWiND in particular) is described in the related work section. The paper is generally well written and easy to follow, and the illustrations (Figure 1, 2) are helpful for understanding the technical contributions.\n- The proposed method seems technically sound and appears to rely fairly little on engineering or manual labor which is definitely a plus (besides language annotation of the demonstration data, I didn't see much discussion on this in the paper but I imagine it may be rather time-consuming). The approach is rather simple and intuitive (I consider this a strength), which makes me believe that it could feasibly be transferred to other tasks and potentially entirely different task domains provided that they can be separated into stages / subtasks.\n- I appreciate the focus on a difficult real-world task such as T-shirt folding, and the empirical results on this task are rather compelling. In particular, the qualitative comparison with ReWiND in Figure 3 clearly demonstrates the shortcoming of ReWiND: it is not very robust to varying trajectory lengths and, as a result, struggle with predictions at either end of the completion spectrum. The paper also includes a limited set of results on a dish unloading task (appendix A.6) as well as RL in simulation (pick cube) and the conclusions seem to be rather similar across these tasks and problem settings.", "weaknesses": "I do not have any major concerns with the paper in its current form. However, I do believe that there is room for improvement in several ways:\n- The writing is a bit odd, with large parts of the experiment section moved to the appendices entirely. It would be helpful if the authors can be more clear about the experimental setup (T-shirt folding, dish unloading, RL finetuning in simulation) and perhaps summarize the results across these different settings in the main paper rather than only discussing them in the appendices; it is easy for readers to miss otherwise.\n- I am a little concerned about this particular engineering choice: *\"During annotation, only the subtasks defined by the protocol were labeled, and any trajectory that did not contain the complete sequence of subtasks specified by the protocol was discarded. Annotators watched the top-view video and segmented each trajectory into subtasks by recording the start and end frame indices. If a mistake occurred during execution, its start and end frames were also labeled; trajectories containing mistakes were excluded from subsequent model training\"* (L184). If one of the main claims of the paper is that SARM / RA-BC can learn from diverse demonstration data, why exclude trajectories that *e.g.* contain a human error somewhere in the trajectory? If the approach was truly robust, it should be able to infer that such actions are indeed errors that then subsequently are corrected by the operator. It would be helpful if the authors can clarify this particular aspect of their approach, and potentially include some experimental results in which data is not filtered as aggressively before any learning; that should give readers a better indication on what the key steps in the data pipeline are.\n\nMinor comments that need to be addressed but are insignificant wrt my overall assessment of the work:\n- The abstract is quite long which feels unnecessary, I would suggest revising it to be more concise.\n- It would be helpful if the authors could introduce the considered tasks somewhere in the main paper along with e.g. illustrations of what those tasks look like.\n- There are a few typos and grammatical errors throughout the paper. While they do not affect understanding, it would be a good idea to proof-read the paper before publication.", "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": "ZoMsII01t7", "reviewer": "Reviewer_3zBM", "summary": "This paper introduces SARM, which introduces stage information about a task to improve reward predictions given a real world robot trajectory/video. The results show this outperforms prior work like ReWiND and aligns closer to the ground truth progress rewards. The learned reward model in SARM further is used to reweight training samples for better performance.", "questions": "See above\n\nGenerally think this work is solid and should be accepted. Happy to raise score if the question about the reward reweighting is addressed.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Introduces a novel use of stage prediction to get improved auto generated reward labels of a dataset. This contrasts with prior work that simply used frame indices of a video as the label of the reward progress, which is limited since different stages of the task may proceed at different speeds.\n- This work further introduces some useful tricks that improve performance and leads to easier to learn reward models that are much more smooth.\n- Strong improvement upon prior work is shown with ablations indicating how the different proposed ideas improve SARM.", "weaknesses": "- Compared to prior work this method requires more annotation of datasets (subtask labels).\n- There doesn't seem to be an ablation on the reward-model based weighting of training samples, which I believe would be fairly important given the amount of space allocated to describe the RA-BC method. The discussion in 4.3 suggests that a bad reward model would cripple RA-BC results since it may lead to poor reweighting, it may be possible that not doing any re-weighting may end up working better.", "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": "ijFBzvCc3R", "reviewer": "Reviewer_WSRa", "summary": "The paper proposes a reward-modeling approach that supplies stable, dense feedback for long-horizon manipulation. It decomposes progress estimation into two parts: (1) a high-level predictor that identifies the current stage of the task and (2) a fine-grained regressor that measures progress within that stage. This two-stage design reduces sensitivity to variable demonstration lengths (e.g., in cloth folding) and yields more reliable rewards. Building on this signal, the authors also introduce an imitation-learning scheme that up-weights higher-quality demonstrations. Experiments show consistent gains over single-stage progress-prediction baselines, with real-robot results on cloth-folding validating the approach.", "questions": "- [Q1] Is the low performance due to underfitting?\n    + In Table 2, ReWiND improves with longer training (e.g., 20K Medium: 1/12 → 40K Medium: 6/12). Does this indicate the model hasn’t converged yet and could achieve higher performance with more training? (possibly outperform your model as well?)\n\n- [Q2] Model architecture parity\n    + Are ReWiND and SARM matched in model size? Since original ReWiND leverages large-scale data (OXE), it may use a larger model, which could overfit on your 200-hour dataset.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "- Strong real-world gains: On a challenging cloth-folding task, a naive fine-tuning SOTA behavior policy (Pi0) fails 0/12 trials, while the proposed method succeeds in 8/12—highlighting the efficacy of their good reward modeling.\n- The paper analyzes performance across (i) high-level stage-prediction accuracy, and (ii) end-to-end policy rollouts, helping readers to understand where the method adds value.\n- Consistent improvements: The performance improvement holds across different rollout outcomes (success/partial/failure progress prediction) and task difficulty, from simple cloth pick-and-place to full folding.", "weaknesses": "**Novelty of the stage-aware reward modeling.**\n  + Stage-aware reward modeling for long-horizon manipulation is not new per se. Please clarify what is distinct, and compare directly to prior work (e.g., Drs [1], REDS [2]).\n\n**Clarity on data-quality literature.**\n\n  + L76–79 state data quality is hard to assess beyond simple heuristics (e.g., duration). This could be misleading: recent work proposes stronger proxies. Please cover this literature (like [3] and [4]) and briefly position your method relative to these proxies.\n\nPlease clarify the literature on stage-aware reward modeling and data quality, state how your approach differs from prior work, and—if possible—add experiments showing superiority over these to strengthen your paper’s contribution and rigor.\n\n### References\n- [1] Mu et al., \"Drs: Learning reusable dense rewards for multi-stage tasks\", ICLR 2024.\n- [2] Kim et al., \"Subtask-Aware Visual Reward Learning from Segmented Demonstrations\", ICLR 2025.\n- [3] Belkhale et al., \"Data Quality in Imitation Learning\", NeurIPS 2023\n- [4] Dass et al., \"DataMIL: Selecting Data for Robot Imitation Learning with Datamodels\", arXiv preprint arXiv:2505.09603 2025", "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)