Tsinghua University Researchers Propose Utility-WAM for On-Demand Future Video Generation in Robot Action Planning

Predicting future video frames for action planning offers dynamic foresight for long-horizon tasks in embodied AI, but the mixed effects on action outputs yield negligible net gains in average success rates while incurring full computational costs.

Incorporating the paradigm of "imagining a few future frames before executing an action" into policy learning has been a highly regarded approach in embodied AI over the past two years. By internally simulating subsequent scene frames, robots can equip their action planning with dynamic foresight, theoretically making it better suited for long-horizon, multi-stage tasks.

However, predicting future frames alters the actions output by the model. These alterations yield both positive and negative effects. After the positive and negative impacts cancel each other out, there is virtually no net gain in the average success rate, yet the computational cost of generating these predicted frames is fully incurred.

On October 29, Tsinghua University, in collaboration with Johns Hopkins University, Northeastern University, Southeast University, the Institute of Automation, Chinese Academy of Sciences, and Yinwang Intelligent Technology, published the paper Utility-WAM: Correcting Future-Induced Action Changes. Feihong Zhang, Xi Lin, and Yulong Shi serve as co-first authors. Feihong Zhang, affiliated with Tsinghua University, has also contributed to the UniDex robotic dexterous hand project. Corresponding authors Zihao Liang, Yinuo Wang, and Zhaowen Li are based in relevant laboratories at Yinwang Intelligent Technology Co., Ltd.

The fundamental architecture of this approach involves sharing a single diffusion backbone for both video and action. The video branch generates future frames, while the action branch denoises them into action chunks accordingly. Whether the action branch can "see" the future has become the most expensive switch to toggle during deployment.

To address this black-box issue, the team proposed Utility-WAM, a method that corrects action deviations induced by future video predictions and determines whether to generate these videos based on utility. It constructs two parallel action streams on a single shared-weight snapshot, using residual adapters to counteract harmful action offsets. A gating module that reads only the current observation then decides whether to expend computational resources to generate future videos for that specific step.

On the LIBERO and RoboTwin 2.0 datasets, enabling future prediction originally yielded average performance fluctuations of just 0.8 and -0.8 percentage points. After correction and dynamic selection, the success rates for the two tasks reached 98.9% and 93.3%, respectively. In 60 real-robot trials across three tasks using a dual-arm setup, Utility-WAM succeeded 53 times; under identical conditions, FAST-WAM succeeded only 43 times.

The scalability of this approach depends on two factors:

• Whether the corrected action offset patterns remain stable and effective when switching the backbone model or robot hardware.
• Whether the gating mechanism can be upgraded from relying on offline action error proxy metrics to making decisions based on closed-loop task rewards.

Deconstructing future video to quantify action increments

Over the past two years, discussions surrounding world-action models have largely focused on whether it is worth paying for imagination. Utility-WAM takes a different approach: it examines exactly how much a predicted future frame alters the resulting action.

Using a shared Stage-1 checkpoint for comparison, the team established two distinct paths: one where the action branch can access the future video, and one where it cannot. They defined the difference between the actions output by these two paths as the action increment, Δa. This increment might improve task performance, introduce negative effects, or have almost no impact. However, whenever the query is triggered, the latency overhead incurred by video generation is unavoidable.

Adopting the classic definition of negative transfer, the team terms the performance drop observed compared to the no-future path "future-induced negative transfer." This differs from traditional negative transfer in that both policies share parameters and training data; the sole additional information source is the synthesized future frames. The change in closed-loop success rate, ΔJ, is derived from the continuous accumulation of individual Δa values across the entire task trajectory. Negative values in the experimental tables indicate that introducing this future prediction actually degrades policy performance.

These two action streams can be compared in parallel through the use of attention masks.

In the first stage, the shared video DiT and action DiT are jointly trained: the Mwo mask blocks the action branch from reading future video tokens, while the Mw mask grants reading access. Both paths retain current observation information and share parameters completely. The weights remain identical, with the sole variable being whether the action branch can access the future video tokens. Consequently, all subsequent comparisons of Δa and ΔJ can be attributed strictly to the addition of future prediction within the same model and dataset, eliminating the interference of random variations that would arise from training two independent policies.

The sequence of correcting first and then deciding whether to generate is dictated by the inherent characteristics of the action stream. Once a query activates the future-conditional path, any detrimental action increment has already entered the control loop. Even if the system switches back to the no-future path later, it can only affect subsequent steps and cannot reverse the deviation introduced at the current step. Treating this merely as a routing problem would mean tolerating the negative impact already incurred and deferring its resolution.

Shifting to this perspective allows many previous discussions to be reevaluated.

Always generating future frames assumes that the expected return of Δa can offset the latency cost. Directly disabling future prediction during deployment assumes that the action increment is not worth the computational cost, while also forfeiting its positive value in long-horizon tasks. Routing based on task phases assumes that the task phase itself can serve as a reliable indicator of utility.

The core idea behind Utility-WAM is to evaluate utility directly: it first determines whether the corrected action increment can reduce the action error, and then decides whether to trigger video generation based on current observations. This approach no longer relies on phase labels but must retain the utility signal.

Compared to similar approaches, the difference lies in how information is utilized.

FAST-WAM accesses future information during training but directly disables video generation during deployment. AdaWAM decides whether to run visual reasoning based on the execution phase, with routing rules relying on phase supervision signals. PFD distills privileged future information into a policy that observes only the current frame, eliminating visual prediction during inference. This line of research can be traced back to earlier model-based reinforcement learning methods like I2A, STEVE, and MBPO, which similarly had to address the challenge that simulation rollouts are valuable but prone to errors.

The unique aspect of Utility-WAM is that, keeping parameters and data fixed, it treats the synthesized future itself as an additional source of information that induces negative transfer, and directly corrects the action bias it introduces.

Three-stage training and selective inference

The first stage introduces attention masks only during joint training, without separately training an independent policy. The total loss comprises three components: the video loss and the action losses under the two masking conditions, weighted by λwo and λw. After training, the model naturally forms two internal action streams: πwo, which does not read future information, and πw, which does; this stage adds no extra parameters.

The second stage freezes the shared WAM and trains only the residual adapter ϕ. During inference, the adapter receives model-generated future features, and during training, it aligns with the same data distribution. Under a unified noise level, the frozen WAM outputs three sets of velocities: the velocity without future conditioning vwo, the velocity based on self-generated future vw, and the velocity with noisy ground-truth future conditioning vGT. The adapter predicts the velocity residual based on the action query and pooled future features, with the corrected velocity expressed as v̂w = vw + Δvϕ.

The training supervision signal comprises three terms. The action term continuously fits the expert demonstration velocity. The alignment term constrains the difference between the corrected future-conditioned velocity and the unconditioned velocity, steering it toward the effective correction magnitude δ provided by the ground-truth future. If the ground-truth future yields no significant improvement, δ is set directly to zero when the relative improvement margin ρ is 0.05, preventing the adapter from drifting arbitrarily. Finally, the excess error term specifically penalizes the portion of the local flow-matching error that exceeds the unconditioned baseline.

Each of the three supervision terms serves a distinct purpose: the alignment term dictates the correction direction, while the excess error term establishes an error ceiling, together forming an asymmetric safety valve. This mechanism remains inactive during most routine tasks, but if destructive drift occurs, it pulls the error back within the unconditioned baseline. Neither term can achieve the full effect on its own; only their combined operation restores performance to a usable range for precision tasks.

In simple terms, the adapter performs error correction using a reference answer. The ground-truth future trajectory acts as an open-book reference: if it significantly improves velocity prediction, this effective correction becomes the learning target; if it offers no benefit, the future-conditioned path is constrained from deviating too far from the unconditioned baseline. Because self-generated futures inherently contain sampling noise and distribution shifts, directly adopting vw − vwo would inject noise into the control logic. Pairing the alignment term with the excess error term filters out the few genuinely useful correction signals from this noise. The team noted that this approach builds on the core concept of professor forcing.

In the third stage, the WAM and the adapter are frozen to train a lightweight gating module ω. The gating labels are derived from the adapter's corrected paired outputs: if the corrected action increment reduces the error relative to the expert action, it is labeled as a positive sample. The gating module takes only the current observation and proprioceptive state as inputs, learning to predict whether the correction will be valuable before the video is generated. Since Δa can only be computed after video generation is complete, the gating mechanism essentially estimates the benefit in advance.

During inference, the process simplifies to a binary decision. The gating module runs first; if its score falls below the threshold, the system follows the unconditioned action flow, executing a ten-step action update while entirely bypassing video computation. If the score exceeds the threshold, the system generates the future video and performs the adapter-corrected action update after each video step. The thresholds are tuned on the validation set: 0.439 for LIBERO and 0.541 for RoboTwin 2.0. Both branches run during training to construct Δa, but only the selected branch is executed during inference. This design dictates latency: the gating module itself incurs minimal computational overhead, with video generation accounting for the bulk of the cost.

The correction module is attached to the future-conditioned path because this branch relies on synthetic representations derived from noisy sampling. Long-horizon tasks depend on these representations for look-ahead capability, but if the representations drift, precision manipulation is the first to suffer; here, the unconditioned velocity serves as a stabilizing anchor. The gating labels are constructed from paired action errors and do not require task-stage labels, allowing them to adapt directly to real-robot test episodes of variable lengths.

Collapse and Recovery Masked by Average Scores

The experimental backbone adopts Wan2.2-TI2V-5B from FAST-WAM, with each benchmark trained independently. All variants share the same Stage-1 checkpoint, initial state, and random seed. The core value of this design is to isolate the effects introduced by the adapter and the gating module from the inherent stochastic variations of the backbone model itself. The trade-off is the inability to evaluate the variance introduced by retraining the backbone, but it yields much cleaner conclusions for the component ablations.

LIBERO focuses on long-horizon, language-conditioned desktop tasks, whereas RoboTwin 2.0 targets bimanual manipulation. Combining these two benchmarks can easily create a significant disparity between average metrics and tail-case performance.

Using the no-future branch as the baseline, enabling future prediction yields the following results: the average success rate on LIBERO increases from 97.4% to 98.2%, a gain of 0.8 percentage points, while on RoboTwin 2.0, it drops from 91.9% to 91.1%, a decrease of 0.8 percentage points.

The marginal positive gain on LIBERO comes almost entirely from the Long (long-horizon) subset, with fluctuations in the remaining subsets not exceeding 0.4 percentage points. RoboTwin exhibits consistent trends across both clean and randomized environment settings, indicating that random environmental perturbations do not amplify this negative transfer effect. The Joint and IDM future variants of FAST-WAM reproduce this same pattern: long-horizon tasks benefit, while average performance on the bimanual benchmark declines.

Relying solely on average metrics can easily lead to the misconception that the impact of future prediction is marginal, but a task-by-task statistical analysis reveals the exact opposite.

Out of the 50 tasks in RoboTwin, performance drops in 21 tasks when future prediction is enabled. This degradation is highly concentrated: Press Stapler falls by 54.0 percentage points and Place Can Basket by 34.0 percentage points, while the maximum drop for the remaining impaired tasks is only 2.5 percentage points. In contrast, only 7 out of 40 tasks decline on LIBERO, with a relatively mild fluctuation range. Furthermore, these two severely affected tasks experience critical failures across multiple schemes for utilizing future information, indicating a systemic issue.

Another easily overlooked distributional characteristic is that for half of the RoboTwin tasks, the change in success rate before and after enabling future prediction does not exceed 1 percentage point. A subset of tasks achieves positive gains, a subset of precision tasks suffers steep drops, and the majority of tasks remain largely unaffected. Averaging these three categories of samples creates the illusion that enabling or disabling future prediction makes little difference.

The latency cost of generating future videos is substantial. Tested on a single 80 GB H100 GPU using BF16 precision and a batch size of 1, the latency for the path with future prediction is approximately 1.7 to 1.9 times that of the path without it. Regardless of whether the prediction ultimately benefits or harms task performance, this computational overhead is incurred whenever video generation is invoked.

For the Utility-WAM (Correction) variant, which enables only the correction module and generates videos throughout the sequence, the success rates are 98.7% on LIBERO and 92.8% on RoboTwin. The performance drops for the two severely affected tasks are reduced to 6.0 and 14.0 percentage points, respectively. Nearly 80% of the overall gain on RoboTwin comes from these two repaired tasks, while the original positive gains on the long-horizon subset are preserved.

Ablation studies rule out the hypothesis that recovery is achieved merely by increasing network capacity. Enabling the alignment term alone yields partial improvements, but the recovery capability only reaches its upper limit once the excess error penalty is applied.

By adding gated selective inference on top of the correction module, the final version, Utility-WAM (Selective), achieves 98.9% on LIBERO and 93.3% on RoboTwin. In RoboTwin, it retains only 24.5% of video generation requests and bypasses the video branch for 75.5% of queries, reducing end-to-end latency by 37.9%. In LIBERO scenarios, the overall positive gains from future prediction remain stable, with the gate mostly choosing to enable future prediction, resulting in only marginal latency reductions.

Under an identical compute budget, the learned gating strategy significantly outperforms two naive baselines: fixed truncation steps and randomly enabling the correction branch.

The source of these additional gains from gating lies in filtering out timesteps where future prediction still fails to provide positive value even after correction. Correction stabilizes the performance floor for the most severely degraded tasks, while gating further eliminates ineffective calls, raising the overall performance ceiling under a lower compute budget.

Compared to similar methods: the future variant of FAST-WAM improves performance on long-horizon tasks but suffers a performance drop on dual-arm tasks; PFD yields considerable results but completely abandons visual prediction capabilities during inference; and AdaWAM also performs routing but relies on manually annotated phase supervision labels.

Ten more successes across three dual-arm tasks

Moving beyond the simulation environment, the team built a dual-arm platform using two AgileX PiPER 6-DOF robot arms, each equipped with a two-finger gripper.

The observation inputs include a fixed front-facing RGB camera and the wrist-mounted RGB cameras for both the left and right arms. The robot's proprioceptive state is a 14-dimensional vector comprising the joint angles of both arms and the gripper opening states.

The experiments evaluate three tasks: Place Empty Cup, Insert Can, and Stack Three Blocks. For each task, 200 teleoperation demonstrations were collected at a sampling frequency of 10 Hz, with a training window of 32 steps corresponding to a 3.2-second action horizon.

Each task was evaluated 20 times, with the initial object placements randomized during testing. Both FAST-WAM and Utility-WAM (Selective) were trained on the same set of demonstration data, with the latter following the three-stage training process described earlier.

Utility-WAM (Selective) achieved 53 successes out of 60 trials, yielding an 88.3% success rate, while FAST-WAM succeeded in 43 out of 60 trials for a 71.7% success rate. This represents a difference of 16.7 percentage points, with a two-sided 95% Newcombe confidence interval ranging from 2.3 to 30.4.

Breaking down the results by individual task:

For Place Empty Cup, both groups tied at 18 out of 20. For Insert Can, performance improved from 14 out of 20 to 18 out of 20—an increase of four successes—with the difficulty concentrated on aligning the can with the groove. For Stack Three Blocks, performance improved from 11 out of 20 to 17 out of 20, an increase of six successes; this is also the longest multi-stage task in terms of temporal duration among the three. Operations such as alignment and stacking are highly sensitive to errors, making them particularly well-suited to utilize calibrated future trajectories at specific time steps.

This hierarchical pattern corroborates the simulation conclusion that long-horizon tasks are well-suited to leveraging future predictions, and that correction mechanisms protect precise operations.

The results from the intermediate control group are equally clear. The shared WAM without a future branch achieved 47 out of 60, for a 78.3% success rate. Enabling uncalibrated future prediction dropped this to 45 out of 60, a 75.0% success rate.

In real-robot scenarios, directly enabling future prediction actually causes a slight drop in performance. Enabling only the correction module achieves 50 out of 60, an 83.3% success rate. Adding the gating mechanism further improves this to 53 out of 60. The benefits accumulate incrementally, validating the approach of applying correction first and then selectively generating video. Simply generating future frames continuously fails to achieve the same results.

However, the experiments reveal a clear hierarchical pattern: some tasks remain unaffected, alignment tasks benefit, and long-horizon stacking tasks show the most significant improvement. This hierarchical logic aligns with the simulation diagnostics. The real-robot tests did not reproduce the extreme failure seen in the Press Stapler task, which suffered a 54.0 percentage point drop. Nevertheless, the underlying trend remains valid: enabling uncorrected future predictions causes a slight performance drop, while combining correction with gating progressively improves performance.