Beyond Future Prediction: Denoising as Generative Adaptation for Robot Control

Zanyi Wang; Yuheng Lei; Dengyang Jiang; Ping Luo; Mengdi Wang; Zhixuan Liang; Shilong Liu

cs.ROSep 23, 2026

Beyond Future Prediction: Denoising as Generative Adaptation for Robot Control

Authors: Zanyi Wang, Yuheng Lei, Dengyang Jiang, Ping Luo, Mengdi Wang, Zhixuan Liang, Shilong Liu

Organizations: UC San Diego · HKU · HKUST · Princeton University

Abstract

Pretrained generative Diffusion Transformers (DiTs) capture rich pixel-level visual and language-conditioned structure through large-scale image and video generation training. A growing line of robot policies builds on this generative prior, but how it should be transferred to control remains unclear, and existing approaches commonly instantiate this transfer through future visual prediction. We ask a more basic question: what a pretrained generative DiT actually contributes to action learning, and how this prior should be adapted for control. We introduce NowWAM, a future-target-free co-training formulation that denoises the current observation and predicts robot actions from the same visual stream, directly coupling the native generative objective to the action-facing representation across the denoising trajectory. Under matched controlled settings, past and future visual targets perform comparably, while restricting training to the clean endpoint substantially reduces robustness, suggesting that a separate future target is not essential for generative adaptation, while the denoising trajectory remains an effective interface for control. On LIBERO-Plus, NowWAM reaches 87.7% with FLUX2-Klein, improving over the future-target co-training baseline by 6.1 points while halving training visual tokens (784 to 392) and reducing step time from 2.85 s to 1.63 s, a 1.8x speedup. With the pure text-to-image Z-Image backbone, NowWAM further reaches 87.8%, showing that strong control adaptation is not tied to video generation or image-editing backbones.

Figures & tables

Figure 1: From future prediction to current denoising. (a) Future prediction is used for action prediction. (b) Future targets are used for generative co-training. (c) NowWAM denoises the current visual stream and predicts actions directly. (d) NowWAM improves robustness while reducing training cost.

Figure 2: NowWAM. During training, the current visual stream is sampled along the pretrained denoising trajectory and jointly supports generative prediction and action learning. At inference, the same policy operates at the clean endpoint ( $\sigma=0$ ) in a single forward pass.

Method	Spatial	Object	Goal	Long	Avg.
$\pi_{0.5}$	98.8	98.2	98.0	92.4	96.9
X-VLA	98.2	98.6	97.8	97.6	98.1
LingBot-VA	98.5	99.6	97.2	98.5	98.5
Motus	96.8	99.8	96.6	97.6	97.7
Fast-WAM	98.2	100.0	97.0	95.2	97.6
ImageWAM	97.2	99.2	98.8	98.4	98.4

Table 1: Standard LIBERO and RoboCasa GR1 manipulation.

	Perturbation category							Overall
Method	Camera	Robot	Lang.	Light	Bkg.	Noise	Layout	SR (%)
$\pi_{0}$	13.8	6.0	58.8	85.0	81.4	79.0	68.9	53.6
$\pi_{0.5}$	78.4	73.6	80.8	96.2	94.1	89.0	84.5	84.4
StarVLA	52.5	49.8	88.5	95.7	95.7	73.0	76.9	74.1
OpenVLA-OFT	56.4	31.9	79.5	88.7	93.3	75.8	74.2	69.6
ABot-M0	60.4	67.9	86.4	96.2	91.6	86.4	82.6	80.5

Table 2: Success rates (%) on LIBERO-Plus across seven perturbation categories.

Figure 3: Qualitative rollouts under distribution shift. We show representative trajectories under RobotInit and camera-view perturbations in LIBERO-Plus. Fast-WAM and ImageWAM exhibit distinct grounding and execution failures, whereas NowWAM completes the corresponding tasks. These examples qualitatively illustrate the robustness gains observed in Table 2 .

Init.	Visual loss	SR
Random	Edit	58.14
Pretrained	Edit	83.05
Random	None	50.55
Pretrained	None	75.81

Table 3: All rows are strictly matched; only the stated intervention changes. Results use non-EMA checkpoints on the same fixed 1,923-episode subset.

Method	Training proxy	Vis. tokens $\downarrow$	Step time (s) $\downarrow$	Peak mem. (GiB) $\downarrow$
ImageWAM	Future image	784	2.85	78.9
NowWAM	Current denoising	392 (50% fewer)	1.63 (1.8 $\times$ faster)	74.3 (5.8% lower)

Table 4: Training efficiency. Same

2\times

H200 setup with global batch size 64. Both methods use FLUX2-Klein-4B as the backbone.

Figure 4: Task-conditioned attention across robustness and transfer settings. Each row starts with the RGB observation, followed by attention maps over the rollout.

Robotic Control

Robotic control studies how machines use sensing and computation to produce reliable movement and interaction in changing environments.

See all 571 papers

An Action Is Worth One Patch: Unified World-Action Modeling with PatchWAM
Sep 22, 2026Tianheng Wang, Zhou Xie, Heng Jia +3World Models for Robotics Visuomotor Policy Learning
Westlake University · Lanzhou University · Zhejiang University +2
MT-WAM: Reorienting the One-Pass Predictive Representation Toward Action Generation
Sep 18, 2026Yiguang Yang, Jiankun Peng, Xiaoming Wang +2World Action Models Motion Representation Learning
Aerospace Information Research Institute, Chinese Academy of Sciences · University of Chinese Academy of Sciences
JEPA-WAM: Learning Vision-Language-Action Policies with Joint-Embedding World Modeling
Aug 10, 2026Yihan Lin, Jiawei He, Shifeng Bao +6Joint-Embedding Predictive Architecture Robotic Manipulation
School of Information, Renmin University of China, Beijing, China · XYZ Embodied AI, Beijing, China · Key Laboratory of Data Engineering and Knowledge Engineering, Beijing, China +3

Sep 22, 2026J/K move · Enter open · S save

An Action Is Worth One Patch: Unified World-Action Modeling with PatchWAM

Tianheng Wang, Zhou Xie, Heng Jia, Jianhua Xu+2

Westlake University · Lanzhou University · Zhejiang University+2

World Models for Robotics Visuomotor Policy Learning

Open paper AskPDF arXiv

Generative visual models offer a foundation for learning representations of physical dynamics, yet their extension to continuous control raises a fundamental question: do visual prediction and action generation require separate computational pathways? Existing approaches usually introduce trainable action heads or separate action experts to bridge low-dimensional states and high-dimensional visual representations. In this work, we explore whether the visual backbone's existing capacity can also support control when actions are expressed in a compatible representation. Thus, we introduce PatchWAM (Patch World-Action Model), which treats continuous actions as another type of patch through a fixed mapping called Action-as-Patch. This allows a single model to predict both how the robot should move and what the scene may look like afterward. Visual prediction and action generation become parts of the same generative process, without a dedicated action head or separate action expert. Experiments with subsampled training windows show gains over a matched dual-expert control, while benchmark evaluations reach 91.8% success rate on LIBERO-Plus and 96.12% on RoboTwin 2.0 in a full-data setting with additional augmented demonstrations. More broadly, the result suggests that capability need not be added where it can be inherited: the constraint on extending a generative backbone is the interface a new signal is written in, not the capacity to model it.

Abstract

Robotic Control

Explore similar work

An Action Is Worth One Patch: Unified World-Action Modeling with PatchWAM

MT-WAM: Reorienting the One-Pass Predictive Representation Toward Action Generation

JEPA-WAM: Learning Vision-Language-Action Policies with Joint-Embedding World Modeling