Source-linked AI summary
NoiseGate: Learning Per-Latent Timestep Schedules as Information Gating in World Action Models
Wen Huang, Haoran Sun, Yongjian Guo, Yunxuan Ma, Haoran Li, Jing Long, Zhouying Mo, Zhong Guan, Yucheng Guo, Shuai Di, Junwu Xiong
TL;DR
Joint video–action WAMs typically collapse per-latent timesteps into one scalar, imposing equal reliability across predicted frames. NoiseGate learns these schedules as information gates using independent timestep training, a GPN, and GRPO, improving RoboTwin performance including a +10.0 point gain over the shared-t baseline. Its main limitation is the simulation rollout cost of sparse-reward schedule-policy training.
Problem
Joint video–action WAMs collapse potentially heterogeneous per-latent timesteps into a shared scalar, imposing an equal-reliability prior despite differing frame importance for action generation.
Source: §1 · p. 2, §3.2 · p. 3Method
NoiseGate trains a joint MoT backbone with independent per-latent timesteps, then uses a GPN and GRPO task rewards to learn per-latent time increments without hand-crafted shape priors.
Source: §1 · p. 2Results
+10.0 point gain over the shared-t baseline on RoboTwin random-scene evaluation, while improving a strong Stage-1 WAM on the primary comparison.
Source: §6 · p. 9Takeaways & Limitations
The per-latent timestep schedule is a first-class design object for joint video–action WAMs rather than an implicit hyperparameter.
Source: §6 · p. 9Sparse task-success rewards make simulator rollout collection the dominant training cost, and generalization across large-scale multi-task suites remains future work.
Source: A Limitations · p. 13
Abstract
from arXiv · showhide
World Action Models (WAMs) are an emerging family of policies that tie robot action generation to future-observation modeling. In this work, we focus on the joint video--action modeling paradigm, where actions and imagined future observations are co-generated along a shared denoising or flow trajectory, so that perception, prediction, and control are coupled within one generative process. Existing WAMs typically realize this paradigm with a Mixture-of-Transformers (MoT), where video and action tokens interact through shared self-attention. This architecture can in principle assign a separate timestep $t_f$ to each predicted latent frame, yet current systems collapse this degree of freedom onto a single shared scalar $t$. Under the noise-as-masking view of Diffusion Forcing, this shared schedule imposes the unjustified prior that every predicted latent is equally reliable for action generation. We instead view the per-latent schedule as a \emph{learnable information-gating policy}: by changing a latent frame's noise level, the policy modulates the reliability of its Key/Value contribution to the action tokens. We propose \textbf{NoiseGate}, which combines independent per-latent timestep sampling during backbone training, a lightweight Gating Policy Network that emits per-latent time increments during denoising, and task-reward optimization that trains the schedule policy without hand-crafted shape priors. Built on a joint video--action MoT backbone, NoiseGate delivers consistent gains on diverse RoboTwin random-scene manipulation tasks.
1 Introduction
NoiseGate addresses the shared-timestep prior in joint video–action WAMs by learning per-latent schedules as task-adaptive information gates. It combines independent per-latent training noise, a GPN, and task-reward optimization, outperforming shared and handcrafted schedules on RoboTwin random scenes.
Motivation
- Joint video–action WAMs co-generate future latents and actions through a shared denoising or flow trajectory, coupling perception, prediction, and control.This paradigm models p(v, a | v0, l) rather than action-only prediction.Source: §1 · p. 1
- A shared scalar timestep assumes every predicted latent is equally visible to action generation, despite heterogeneous and task-dependent influence across frames.Handcrafted timestep shapes impose another prior and may be suboptimal for critical events such as grasping or placing.Source: §1 · p. 2
Core idea
- NoiseGate interprets per-latent timesteps as continuous information gates because noise modulates each latent’s Key/Value reliability in shared self-attention.The optimal gating pattern is learned from data rather than prescribed by a fixed schedule.Source: §1 · p. 2
Method
- NoiseGate combines independent per-latent timestep sampling, a GPN emitting per-latent increments during denoising, and GRPO trained with sparse task-success rewards.Independent sampling is applied to a chunk-bidirectional backbone so arbitrary per-latent profiles are valid at inference.Source: §1 · p. 2
Results
- On RoboTwin random-scene tasks, NoiseGate outperforms both shared-scalar and handcrafted per-latent schedules.The contribution passage attributes this result to learning the gating policy rather than relying on a fixed schedule.Source: §1 · p. 2
2 Related Work
Prior work establishes joint video–action WAMs, independent noise as a masking interpretation, learned diffusion schedules, and reward optimization for embodied policies. NoiseGate combines these ideas to learn a task-reward-optimized per-latent schedule inside a joint video–action backbone.
Video–action generative policies
Per-token noise levels and adaptive diffusion schedules
- Diffusion Forcing [9] treats noise as partial masking and independently samples per-token noise, while NoiseGate applies the idea to a chunk-bidirectional backbone for inference-time schedule flexibility.Independent noise is used as a training substrate, not as a rollout mechanism.Source: §2 · pp. 2–3
- NoiseGate predicts a vector t = (t1, . . . , tF) whose elements control future latents’ Key/Value reliability in action-token attention, unlike scalar timestep schedulers.Classical samplers optimize global denoising trajectories, whereas NoiseGate schedules video latents inside a joint video–action model.Source: §2 · p. 3
Reinforcement learning for VLA policies
- Existing VLA reinforcement learning optimizes action-producing policies, whereas NoiseGate freezes the video–action backbone and optimizes only a lightweight per-latent scheduler.The scheduler controls how future latent frames are revealed to action tokens during denoising.Source: §2 · p. 3
Synthesis
- NoiseGate combines joint WAM generation, independent noise levels, learned timestep selection, and task-reward fine-tuning in a distinct scheduling setting.Its learned object is a task-reward-optimized information-gating policy inside the generative process.Source: §2 · p. 3
3 Problem Formulation
The problem formulation represents WAM generation with per-frame video timesteps and a separate action timestep, then interprets heterogeneous video noise as information gating. A policy adapts video timesteps while the action schedule remains a fixed global clock.
3.1 Preliminaries: Video–Action Generation in WAMs
- A WAM models pθ(v, a | v0, l) for future latent frames v, current observation v0, language l, and a co-generated action chunk a.The future video and action are generated conditionally on the current observation and instruction.Source: §3.1 · p. 3
- The MoT backbone nominally supports per-token timesteps, but standard WAM practice collapses video timesteps into global scalars alongside the action timestep.The video timestep vector is tv = [t1, . . . , tF], while ta denotes the action timestep.Source: §3.1 · p. 3
- Existing WAMs set t1 = · · · = tF = t while retaining a synchronous action schedule to support a stable execution trajectory.The formulation focuses on relaxing video-timestep flexibility without changing the action schedule.Source: §3.1 · p. 3
3.2 Per-Latent Timesteps as Information Gates
- The shared scalar assumes every future frame and action token is equally visible or reliable, despite differing temporal dependencies, prediction difficulty, and action relevance.This is the hidden prior that motivates heterogeneous per-latent schedules.Source: §3.2 · p. 3
- Under the noise-as-masking view [9], higher tf makes frame vf less reliable in action-token representations, allowing denoising priority to vary across informative frames.The action timestep ta need not change when video timesteps are heterogeneous.Source: §3.2 · p. 3
- Figure 2 couples the joint-sequence MoT backbone with a GPN that reads predicted-chunk latents and per-latent times, emitting Δtf increments while v0 remains pinned at t0 = 0.The action follows its own global schedule.Source: Figure 2 · p. 4
3.3 Learning the Per-Frame Schedule as a Policy
- The scheduling policy πϕ observes the current state at each denoising step and determines individual time increments for the video latent frames.The action schedule remains a fixed linear or cosine decay, while tv adapts to maximize task-success reward R.Source: §3.3 · p. 4
4 NoiseGate: Per-Latent Schedule as a Learned Gating Policy
NoiseGate makes per-latent timestep scheduling a learned control variable in joint video–action WAMs. It combines a MoT backbone, independent timestep training, a Gating Policy Network, and GRPO-based reward optimization.
4 NoiseGate: Per-Latent Schedule as a Learned Gating Policy
- NoiseGate trains a frozen-backbone scheduler policy with GRPO to emit per-latent timestep increments during denoising.The two-stage procedure first trains the WAM under independent per-latent timesteps, then optimizes only the GPN from task reward.Source: §4 · p. 4, §4.4 · pp. 5–6
4.1 World Action Model and MoT Backbone
- The MoT backbone jointly processes future video latents and actions, using shared self-attention with modality-specific feed-forward layers.Each predicted video latent has its own timestep, while action tokens retain a shared global timestep.Source: §4.1 · pp. 4–5
4.2 Training Substrate: Independent Per-Latent Timestep Sampling
- Independent per-latent timestep sampling trains the backbone to process heterogeneously noised video latents and makes arbitrary schedules usable at inference.Training draws separate video timesteps and an action timestep, then applies the corresponding diffusion or flow-matching target.Source: §4.2 · p. 5
4.3 Gating Policy Network (GPN)
- The GPN observes heterogeneously noised predicted latents, the clean observation, and current timesteps, then outputs decrements only for predicted video latents.The observation stays pinned at t0 = 0 and action tokens follow the original global schedule, keeping action denoising synchronized.Source: §4.3 · p. 5
- Bounded relative timestep scales preserve the sampler’s global pace while letting the policy denoise different predicted latents faster or slower.This avoids requiring the policy to learn both absolute decrement magnitudes and frame priorities from sparse reward.Source: §4.3 · p. 5
4.4 Training with GRPO
- GRPO optimizes the scheduler from binary episodic task rewards using a group-relative baseline, clipped updates, and an entropy bonus.The WAM backbone remains frozen, and no KL term is used because the scheduler is trained from scratch without a meaningful reference policy.Source: §4.4 · pp. 5–6
5 Experiments
Experiments evaluate NoiseGate on RoboTwin random-scene manipulation using matched WAM/MoT backbones, benchmark comparisons, schedule ablations, mechanism probes, and a qualitative case study. Results show task-selective gains and support per-latent scheduling as learned information gating.
5.2 Main Results
- NoiseGate improves tasks with remaining headroom by 6–9 success points on four displayed tasks and 2–5 points on many medium-difficulty tasks, while mostly preserving saturated performance.The only small regressions shown are on Place Object Scale and Open Microwave.Source: §5.2 · pp. 6–7, Table 1 · p. 7
5.3 Ablation Studies
- Removing independent-noise training and the GPN reduces overall success from 67.5 to 57.5, while independent-noise training alone reaches 61.3 and a hand-crafted schedule reaches 63.4.The full NoiseGate configuration closes the gap at 67.5%, showing that both the training substrate and learned schedule contribute.Source: §5.3 · pp. 7–8, Table 2 · p. 7
5.4 Qualitative Analysis: Gating, Schedules, and Noise-as-Masking
- Mean action-to-video attention decreases monotonically as a predicted frame’s noise level t_f increases, empirically making per-latent noise a continuous reliability gate for action tokens.The effect persists in early and late layers but is near-flat in middle layers.Source: §5.4 · p. 8, Figure 3 · p. 8
- GPN-generated timestep trajectories separate across future latents and change their ordering across tasks, selectively exposing useful frames earlier while keeping ambiguous frames partially masked.This differs from both shared-scalar sampling and fixed monotone schedules, which cannot adapt per-frame ordering to task context.Source: §5.4 · p. 8, Figure 4 · p. 8
5.5 Case Study
- In a representative grasping case, NoiseGate keeps the critical grasp frame more uncertain, avoids over-reliance on an unreliable prediction, and succeeds where Stage-1 WAM fails prematurely.The comparison shows executed observations in the top row and predicted future frames in the bottom row for each method.Source: §5.5 · p. 9, Figure 5 · p. 9
6 Conclusion
NoiseGate reframes WAM per-latent timestep scheduling as learnable information gating rather than an implicit hyperparameter. On RoboTwin random-scene evaluation, it improves the Stage-1 WAM and gains +10.0 points over the shared-t baseline in controlled ablations.
A Limitations
The current schedule-policy training is limited by simulation data-collection efficiency, while the information-gating formulation may extend beyond robotics to other jointly denoised generative systems.
A Limitations
- Sparse-reward GRPO requires simulator rollouts for each update, making rollout collection costly and potentially slowing convergence when successful trajectories are rare.Scaling a general schedule policy across large multi-task suites is left to future work.Source: A Limitations · p. 13
B Broader applicability
- The information-gating view can apply beyond robotics when downstream consumers read jointly denoised tokens with heterogeneous utility through shared attention.The passage names video prediction, autonomous driving, and embodied navigation as possible settings.Source: B Broader applicability · p. 13
C GPN Architecture and Inference Details
The GPN maps heterogeneously noised video latents and their timestep vector to per-latent scheduling actions during denoising. Its policy samples relative scales, updates video times, advances the action schedule separately, and runs the frozen MoT backbone.
Interface and notation
- At each denoising step, the GPN consumes the current latent stack and per-latent video times to produce policy parameters for scheduling actions.The observation latent is included in the latent stack, while predicted-frame tokens receive learned position and time-conditioned representations.Source: C GPN Architecture and Inference Details · p. 13
Network architecture
- The network converts each video latent into a token using strided 3-D convolutions, spatial pooling, and a 256-dimensional projection.Predicted-frame tokens are reweighted and fused with the observation token through a gated residual update.Source: C GPN Architecture and Inference Details · p. 13
- A three-layer Transformer combines latent-token and sinusoidal timestep features before attention pooling produces a global policy summary.The Transformer uses eight attention heads and GELU feed-forward blocks of width 4D.Source: C GPN Architecture and Inference Details · p. 13
Squashed-Gaussian policy
- The actor is a three-layer MLP over the global and per-latent representations, with learned log standard deviations clamped to [−5, 2].Source: C GPN Architecture and Inference Details · p. 13
- GRPO stores the sampled action log-probability, with the fixed step-dependent scale canceling the extra log-Jacobian in the policy ratio.Source: C GPN Architecture and Inference Details · pp. 13–14
Inference loop
- The policy samples Gaussian variables, squashes them into relative scales, and converts those scales into per-latent timestep increments.The resulting increments are applied to update the future video times.Source: C GPN Architecture and Inference Details · p. 14
- The action timestep is advanced by a fixed global schedule rather than being supplied to the GPN policy head.Source: C GPN Architecture and Inference Details · pp. 13–14
- The updated video and action times are passed to the frozen MoT backbone to produce the next denoising state.During GRPO, only the GPN parameters are optimized; the Wan 2.2 video and Action-Expert DiTs remain frozen.Source: C GPN Architecture and Inference Details · p. 14
D GRPO Hyperparameters
The appendix provides the GRPO hyperparameters used in the experiments.
D GRPO Hyperparameters
- Table 3 lists the GRPO hyperparameters used across all experiments.Source: D GRPO Hyperparameters · p. 14, Table 3 · p. 14
E Detailed RoboTwin Results
The detailed RoboTwin results document the schedule-ablation setup and show that gating is depth-dependent rather than uniform across the backbone.
E Detailed RoboTwin Results
- The controlled schedule ablation compares Shared-t, Stage-1 WAM, Hand-crafted, and NoiseGate on RoboTwin random-scene tasks using a common backbone and protocol.Table 4 reports per-task results for 100 episodes per task, while the primary comparison uses a scaled Fast-WAM-style configuration.Source: E Detailed RoboTwin Results · p. 14, Table 4 · p. 15
F Additional Qualitative Analyses
- The appendix adds layer-stratified attention statistics and broader schedule summaries to complement the main qualitative probes.These analyses extend the compact main-text visualization with additional diagnostic views.Source: F Additional Qualitative Analyses · p. 15
F.1 Mechanism: layer-stratified action→video attention vs. noise level
- Layer-stratified attention preserves monotone decay with frame noise in early and late layers but is near-flat in middle layers.The gating effect is strongest where action tokens directly read video features and weakest in intermediate modality-agnostic layers.Source: §F.1 · p. 15
F.2 Cross-task structure of final-step residual noise
Final-step residual noise varies across future frames and tasks rather than collapsing to a shared endpoint, revealing task-dependent masking choices by the learned schedule.
F.2 Cross-task structure of final-step residual noise
- Frame 1 is nearly denoised while frame 2 retains more residual noise in several tasks, especially place_a2b_right and place_can_basket.The farther future latent can remain partially masked when uncertainty over future contact or placement outcomes is higher.Source: §F.2 · p. 15
- The learned GPN produces distinct, task-dependent residual-noise levels for the two future frames instead of driving all frames to one clean endpoint.The reported means include ±1 standard deviation across evaluated chunks, with tasks sorted by success rate.Source: §F.2 · p. 15, Figure 7 · p. 17
- The appendix situates these observations within the full 50-task RoboTwin random-scene evaluation context, where unavailable full-method entries are marked “–”.Table 5 reports success rates in percent and includes the Stage-1 WAM random-scene evaluation.Source: Table 5 · p. 16