Disentangling Task-Relevant Latent Dynamics for Reward Modeling

Dayou Li; Jionghao Wang; Peihao Wang; Jiuzhou Lei; Hao Wang; Jichen Chai; Hezhen Hu; Minghui Zheng

cs.ROOct 9, 2026

Disentangling Task-Relevant Latent Dynamics for Reward Modeling

Authors: Dayou Li, Jionghao Wang, Peihao Wang, Jiuzhou Lei, Hao Wang, Jichen Chai, Hezhen Hu, Minghui Zheng,

Organizations: Texas A&M University · University of Texas at Austin · Google DeepMind · Georgia Institute of Technology

Abstract

Reward design remains a major bottleneck for reinforcement learning in robotic manipulation. Handcrafted dense rewards can compensate for the limited guidance provided by sparse rewards, but require substantial task-specific reward engineering. Learning rewards from videos provides a scalable, data-driven alternative to manual reward design. However, existing video-based rewards are often computed in entangled visual spaces where task-relevant state changes are confounded with extraneous visual variation. As a result, the learned reward may respond to visual distractors instead of actual task execution. To address this limitation, we propose DiTaR (Disentangling Task-Relevant Latent Dynamics for Reward Modeling), a framework for learning dense rewards from heterogeneous videos. DiTaR learns a continuous task-relevant representation from multi-frame video latents through an asymmetric two-stage procedure that suppresses task-irrelevant visual variation. In the task-relevant latent space, we train a history- and language-conditioned rectified-flow model to capture task dynamics and derive a directional dense reward from temporal changes in dynamics compatibility. Together, these components enable reward modeling directly over task-relevant dynamics rather than generic visual variation. DiTaR substantially improves policy learning on MetaWorld-v3 and ManiSkill over video-based reward baselines, and transfers to real-world offline RL through reward relabeling. Reward analyses further show that DiTaR provides more reliable rewards.

Figures & tables

Figure 1: DiTaR learns dense rewards from task-relevant latent dynamics. The DiTaR reward accelerates RL convergence while remaining stable under nuisance visual variations, enabling robust policy learning across diverse tasks.

Figure 2: Pipeline overview. DiTaR disentangles task-irrelevant and task-relevant latents through an asymmetric two-stage representation learning procedure. A history- and language-conditioned rectified-flow model captures task-relevant transition dynamics, and the temporal difference of its residual-based compatibility score provides the dense reward.

Benchmark	Method	Final SR $\uparrow$	SR $\geq$ 80 $\uparrow$	Reach 0.8 $\uparrow$	NormT $@$ 0.8 $\downarrow$
MetaWorld	Env-only	56.8	50.0	50.0	0.68
	VIPER	39.9	33.3	41.7	0.73
	Diff. Reward	13.1	8.3	8.3	0.92
	ReWiND	65.2	33.3	41.7	0.61
	DiTaR	94.3	91.7	91.7	0.24
ManiSkill	Env-only	74.6	50.0	75.0	0.58

Table 1: Aggregate RL performance on MetaWorld and ManiSkill. We report final success rate, task coverage, and normalized time-to-threshold.

\uparrow

higher is better,

\downarrow

lower is better.

Figure 3: Main RL performance. DiTaR achieves strong and consistent policy learning across 12 MetaWorld and 4 ManiSkill tasks, outperforming Env-only, VIPER, Diffusion Reward, and ReWiND in aggregate performance. Panels use task-specific x-axis ranges for readability.

Figure 4: Ablation on history conditioning. Removing visual history from the conditional dynamics model reduces learning efficiency and final performance across the evaluated MetaWorld tasks.

Figure 5: Ablation on temporal-difference reward. Temporal differencing improves policy learning over directly using the absolute compatibility score

\phi_{t}

across MetaWorld and ManiSkill tasks.

(a) Temporal sensitivity: $\Delta$ (%) $\downarrow$
	MetaWorld				Real-world
Perturbation	DiTaR	Diff. R.	VIPER	ReWiND	DiTaR	Diff. R.	VIPER	ReWiND
Locally shuffled	$\mathbf{-22.2}$	$-0.6$	$-10.2$	$+0.1$	$-3.2$	$+0.1$	$\mathbf{-10.0}$	$+0.2$
Temporal jitter $\pm 2$	$\mathbf{-19.1}$	$-0.3$	$-4.7$	$+1.5$	$-4.2$	$-0.2$	$\mathbf{-5.5}$	$+0.6$
Shuffled chunks	$\mathbf{-28.9}$	$+0.5$	$-9.3$	$-15.9$	$-7.0$	$-1.6$	$\mathbf{-11.8}$	$+19.1$
(b) Visual robustness: $\|\Delta\|$ (%) $\downarrow$

Table 2: Reward sensitivity and robustness. For temporal perturbations,

\Delta

denotes the signed reward change from clean trajectories, with more negative values indicating stronger temporal sensitivity. For visual perturbations,

|\Delta|

denotes the reward-change magnitude, with smaller values indicating greater robustness. Diff. R.: Diffusion Reward.

Figure 6: Reward robustness to visual perturbations. Per-frame reward trajectories on MetaWorld expert rollouts under lighting changes and $1.25\times$ camera zoom. $\Delta$ denotes the relative mean reward change from the clean rollout. DiTaR shows substantially smaller reward shifts than VIPER, Diffusion Reward, and ReWiND under both perturbations.

Figure 7: Real-world task suite. Three tabletop manipulation tasks used for offline RL evaluation: (a) object-to-basket placement, (b) bottle-to-plate placement, and (c) knife-to-glove placement.

Metric	VIPER	Diffusion Reward	ReWiND	DiTaR
Success Rate	16.7%	26.7%	43.3%	83.3%

Table 3: Real-world offline RL results. All methods relabel the same offline demonstrations and use the same IQL training pipeline. We conduct 10 evaluation trials per task.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Ratio (%)	View	Frame stride
AgiBotWorld ( Bu et al., 2025a )	38.0	Head	$\{1,2,3,4\}$
DROID ( Khazatsky et al., 2024 )	38.0	Exterior	$\{1,2,3,4\}$
Bridge ( Walke et al., 2023 )	19.0	Exterior	$\{1,2,3\}$
MetaWorld ( McLean et al., 2025 )	2.5	Left	$\{1,2\}$
ManiSkill ( Gu et al., 2023 )	2.5	Default	$1$

Appendix

Table 4: Video sources, proportions, camera views, and frame strides used for DiTaR pretraining.

Signed reward change: $\Delta$ (%) $\downarrow$
	Bridge			Real-world
Corruption	DiTaR	Diff. R.	VIPER	DiTaR	Diff. R.	VIPER
Backward drift	$\mathbf{-2.4}$	$-1.1$	$-0.5$	$\mathbf{-1.6}$	$+7.0$	$+0.6$
Looped segment	$\mathbf{-3.5}$	$-0.6$	$-1.7$	$\mathbf{-1.0}$	$+7.4$	$+0.7$
Shuffled chunks	$\mathbf{-16.8}$	$-0.7$	$-6.0$	$-9.2$	$-2.5$	$\mathbf{-11.8}$
Duplicate frames	$\mathbf{-21.1}$	$+0.9$	$+7.0$	$\mathbf{-26.4}$	$+6.8$	$+5.3$

Appendix

Table 5: Temporal sensitivity. Signed reward changes

\Delta

(%) over 16 expert episodes per dataset. More negative values indicate stronger temporal sensitivity. The largest decrease per row and dataset is bold. Diff. R.: Diffusion Reward.

Figure 8: Temporal corruption on Bridge. Clean and corrupted rewards over 16 episodes. Bands indicate one standard deviation; $\Delta$ is the mean signed reward shift.

Figure 9: Temporal corruption localization. Temporal jitter of $\pm 2$ frames is applied to the shaded middle third. Left two columns: MetaWorld, “take the object out of the hole.” Right: ManiSkill, “push the red T-shaped block into place.” Solid and dashed curves show clean and corrupted rewards.

Dataset	Method	Temporal jitter $\uparrow$	Local shuffling $\uparrow$
MetaWorld	VIPER	0.91	0.96
MetaWorld	DiTaR	0.95	1.00
ManiSkill	DiTaR	0.83	0.94

Appendix

Table 6: Temporal corruption localization. Mean localization AUCs on MetaWorld and ManiSkill; higher values indicate better localization.

Camera yaw	$4^{\circ}$	$8^{\circ}$	$16^{\circ}$	$24^{\circ}$	$32^{\circ}$
Relative change ratio	$3.8\times$	$5.2\times$	$4.8\times$	$10.7\times$	$6.6\times$
Smaller DiTaR change (% of videos)	82	86	86	100	100

Appendix

Table 7: Viewpoint robustness on ManiSkill. Relative reward-change ratio (ReWiND/DiTaR) and percentage of videos with a smaller reward change for DiTaR.

Figure 10: Viewpoint robustness on ManiSkill. Relative reward changes under synthetic camera yaw (circles) and a second rendered view of the same rollout (stars). The vertical axis is logarithmic; smaller changes indicate greater robustness.

Figure 11: Rewards under visual perturbations. Left two columns: MetaWorld, “press the button,” under $1.25\times$ camera zoom. Right column: ManiSkill, “push the red T-shaped block into place,” under a second camera view. Solid and dashed curves show clean and perturbed rewards.

Figure 12: Reward evolution during policy learning. DiTaR rewards at a fixed task-specific timestep across training stages on 12 MetaWorld-v3 tasks. Solid curves show smoothed traces and dashed lines show linear fits.

Method	Final SR (%) $\uparrow$	SR $\geq$ 80 (%) $\uparrow$	Reach 0.8 (%) $\uparrow$	[email protected] $\downarrow$
Env-only	56.8	50.0	50.0	0.678
VIPER	39.9	33.3	41.7	0.734
Diffusion Reward	13.1	8.3	8.3	0.922
ReWiND	65.2	33.3	41.7	0.609
w/o TR latents	67.2	58.3	66.7	0.458
DiTaR	94.3	91.7	91.7	0.244

Appendix

Table 8: Task-relevant latent ablation on MetaWorld-v3, with baseline results for comparison. Metrics are defined in Sec. F.3 .

Figure 13: Task-relevant latent ablation. Learning curves on 12 MetaWorld-v3 tasks comparing DiTaR’s task-relevant latents with latents from a single-stage reconstruction-based encoder trained without language conditioning. Both variants use the same conditional rectified-flow architecture and TD reward formulation, with the dynamics model trained separately for each latent representation. DiTaR achieves higher overall success and learns more efficiently.

Figure 14: Cross-dataset latent grouping. Groups of five-frame clips with pairwise cosine similarity of at least 0.90 in the task-relevant latent space. The left and right groups show placing and picking behaviors across different robots and visual settings.

Hyperparameter	Value
Input latent dimension	48
Frames per training window	5
Model dimension	768
Attention heads	12
Encoder blocks	4
Decoder blocks	12

Appendix

Table 9: Architecture and training settings for the latent encoders. The pretrained Wan VAE and T5-base text encoder are frozen throughout representation learning.

Hyperparameter	Value
History input	Five-frame history encoded by the Wan VAE
History tokens / token dimension	512 / 48
Target input	Task-relevant latent of the target clip
Target tokens / token dimension	4 / 32
Text conditioning	Target-only cross-attention; scalar gate initialized to 0
Model dimension	512

Appendix

Table 10: Architecture and training settings for the conditional rectified-flow model.

Hyperparameter	Value
Environment
Action repeat	1
Frame stack	1
Observation size	$64\times 64\times 3$
Episode length	250 environment steps
Per-task training budget	$3\times 10^{6}$ environment frames

Appendix

Table 11: Online visual RL settings for MetaWorld-v3 and ManiSkill.

Benchmark	$(\lambda_{D},\lambda_{\mathrm{expl}})$	Tasks
MetaWorld-v3	$(0.05,0.95)$	button-press , coffee-pull , coffee-push , door-open , drawer-open , faucet-close , faucet-open , hammer , handle-press , handle-pull
	$(0.95,0.05)$	drawer-close
	$(0.50,0.50)$	dial-turn
ManiSkill	$(0.05,0.95)$	PullCube , PushCube
	$(0.95,0.05)$	PokeCube
	$(0.50,0.05)$	PullCubeTool

Appendix

Table 12: Task-specific reward and exploration coefficients

(\lambda_{D},\lambda_{\mathrm{expl}})

in Eq. ( 15 ), shared across learned-reward methods. Following the reward-mixing setup of Diffusion Reward ( Huang et al., 2024 ) , we use task-specific reward and exploration coefficients. Tasks sharing the same coefficient pair are grouped in a single row.

Dense Reward RL

Dense reward reinforcement learning uses intermediate feedback to guide agents when task outcomes alone provide too little learning signal.

See all 28 papers

Rewarding DINO: Predicting Dense Rewards with Vision Foundation Models
Mar 17, 2026Pierre Krack, Tobias Jülg, Wolfram Burgard +1Reward Modeling Reward Design for RL
Department of Computer Science & Artificial Intelligence, University of Technology Nuremberg, Germany. · TUM School of Computation, Information and Technology, Technical University of Munich, Germany
Stage-Transition Dense Reward Modeling for Reinforcement Learning
Jun 30, 2026Yang Yang, Bingjie Chen, Zihan Wang +4RL for Robotics Long-Horizon Robotic Manipulation
Tsinghua Shenzhen International Graduate School, Tsinghua University, Shenzhen 518055, China. · Zerith Robotics, Shenzhen 518055, China.
DenseReward: Dense Reward Learning via Failure Synthesis for Robotic Manipulation
Jul 14, 2026Yu Fang, Wanxi Dong, Jiaqi Liu +7RL for Robotics Reward Modeling
University of North Carolina at Chapel Hill · Carnegie Mellon University · Shanghai Jiao Tong University +1

Mar 17, 2026J/K move · Enter open · S save

Rewarding DINO: Predicting Dense Rewards with Vision Foundation Models

Pierre Krack, Tobias Jülg, Wolfram Burgard, Florian Walter

Department of Computer Science & Artificial Intelligence, University of Technology Nuremberg, Germany. · TUM School of Computation, Information and Technology, Technical University of Munich, Germany

Reward Modeling Reward Design for RL Dense Reward RL Multimodal Reward Modeling

Open paper AskPDF arXiv

Well-designed dense reward functions in robot manipulation not only indicate whether a task is completed but also encode progress along the way. Generally, designing dense rewards is challenging and usually requires access to privileged state information available only in simulation, not in real-world experiments. This makes reward prediction models that infer task state information from camera images attractive. A common approach is to predict rewards from expert demonstrations based on visual similarity or sequential frame ordering. However, this biases the resulting reward function towards a specific solution and leaves it undefined in states not covered by the demonstrations. In this work, we introduce Rewarding DINO, a method for language-conditioned reward modeling that learns actual reward functions rather than specific trajectories. The model's compact size allows it to serve as a direct replacement for analytical reward functions with comparatively low computational overhead. We train our model on data sampled from 24 Meta-World+ tasks using a rank-based loss and evaluate pairwise accuracy, rank correlation, and calibration. Rewarding DINO achieves competitive performance in tasks from the training set and generalizes to new settings in simulation and the real world, indicating that it learns task semantics. We also test the model with off-the-shelf reinforcement learning algorithms to solve tasks from our Meta-World+ training set.

Abstract

Dense Reward RL

Explore similar work

Rewarding DINO: Predicting Dense Rewards with Vision Foundation Models

Stage-Transition Dense Reward Modeling for Reinforcement Learning

DenseReward: Dense Reward Learning via Failure Synthesis for Robotic Manipulation