Achieve What You Imagined: Learning to Align Actions with Visual Plans

I. INTRODUCTION

• The framework treats WAM visual predictions as goal-conditioned proposals and uses a frozen action-conditioned world model to score agreement with predicted consequences and the terminal goal.The resulting rollout-level feedback supports policy optimization without online robot interaction or a separately trained success detector.Source: §I · pp. 1–2
• WAM visual predictions can show task completion even when the accompanying action chunks do not reliably produce those outcomes, creating an imagination–consequence gap.The visual predictions do not encode physics or validate feasibility, so they are not necessarily executable motor plans.Source: §I · p. 1, Figure 1 · p. 1
• FPO adapts to Cosmos3’s high-order UniPC sampler by using an advantage-weighted surrogate from conditional flow-matching losses, treating the sampler as a black box.This enables action-head-only optimization through the native sampling procedure.Source: §I · p. 2

II. RELATED WORK

• Existing joint world-action models separately optimize visual and action generation, while explicit alignment between predicted outcomes and action-conditioned consequences remains insufficiently studied.The paper builds on the joint-modeling paradigm and Cosmos3.Source: §II · p. 2
• Prior reinforcement-learning approaches for generative policies often depend on sampling formulations unavailable for high-order solvers, while learned evaluators require additional preparation, training, calibration, and deployment.Binary terminal rewards can also provide limited task-directed feedback when a weak policy fails sampled rollouts; task-matched demonstrations shift the burden to expert-data collection.Source: §I · p. 1, §II · p. 3

A. Flow-matching generative policies

• A learned time-dependent vector field defines an ordinary differential equation that generates an endpoint from an initial sample under condition c.The trajectory runs from noise at τ = 0 to the generated endpoint at τ = 1.Source: A. Flow-matching generative policies · p. 3

B. Flow Policy Optimization

• FPO constructs a surrogate policy ratio from conditional flow-matching losses, avoiding the need to evaluate a flow policy’s exact likelihood.The ratio compares losses for the same sample, condition, and time–noise pairs under the current and generating parameters.Source: B. Flow Policy Optimization · p. 3

IV. METHOD

• The WAM proposes a goal-conditioned visual prediction and action chunk; a frozen action-conditioned world model predicts the action’s visual consequence and supplies the next observation.The framework uses the discrepancy between these predictions and terminal-goal alignment as feedback, avoiding additional real-robot interaction or task-specific reward-model training.Source: §IV · p. 3

A. Problem Formulation

• At each rollout step, the WAM samples a latent endpoint conditioned on the current observation and instruction, then decodes it into a visual prediction.The prediction and action chunk are generated jointly; the supplied passage defines the endpoint distribution and decoding mapping.Source: A. Problem Formulation · p. 3
• The action chunk is converted to the frozen world model’s control representation, and the model predicts the corresponding action-conditioned visual consequence.The conversion interface includes action conversion and temporal resampling; the visual predictions are described as complementary.Source: A. Problem Formulation · p. 3
• Policy improvement uses model-mediated trajectories rather than real-robot transitions or task-success labels, with video-generation parameters and the action-conditioned model fixed.The WAM parameters are partitioned into frozen video-generation parameters and trainable action-related parameters.Source: A. Problem Formulation · p. 3

B. Method Overview

• The method combines trajectory consistency and terminal visual-goal alignment, then uses group-relative advantages and an FPO surrogate to optimize only the action head.The visual path and action-conditioned world model remain frozen.Source: B. Method Overview · p. 3

C. Closed-Loop Cross-Model Rollout

• The WAM jointly generates a goal-conditioned visual prediction and action chunk; converted actions condition the frozen Ctrl-World model’s consequence video.The conversion restores physical action scale, maps to the world-model control interface, and resamples commands to the world-model rate.Source: C. Closed-Loop Cross-Model Rollout · pp. 3–4
• Each rollout feeds the predicted action consequence into the next policy step, while the WAM-generated visual prediction serves only as a feedback target.Thus, subsequent actions are conditioned on predicted consequences of earlier actions, not on the WAM’s visual prediction.Source: C. Closed-Loop Cross-Model Rollout · p. 4
• Group members use independent latent draws to generate different imagined futures and action proposals under the same initial task condition.Groups share a reset observation and instruction; the rollout procedure stores chunks and computes returns and advantages for optimization.Source: C. Closed-Loop Cross-Model Rollout · p. 4
• The action-head parameters are updated once using each minibatch average of stored chunks.The procedure retains the collected trajectories, probe pairs, advantages, and behavior-policy losses during optimization.Source: C. Closed-Loop Cross-Model Rollout · p. 4

D. Cross-Model Feedback and Temporal Credit

• Trajectory feedback compares corresponding frames from the two predicted videos at common timestamps and matching resolution, using pixel mean squared error.The compared pixel values are on a 0–255 scale.Source: D. Cross-Model Feedback and Temporal… · p. 4
• The goal term compares the world model’s final predicted frame with a predefined visual goal reference and is nonzero only at the final chunk.Nonnegative weights balance trajectory consistency and terminal-goal alignment; higher feedback scores indicate smaller pixel discrepancies.Source: D. Cross-Model Feedback and Temporal… · p. 4
• The trajectory term encourages agreement with the WAM’s imagined motion, while the goal term encourages progress toward the desired outcome.The two terms are combined with nonnegative weights.Source: D. Cross-Model Feedback and Temporal… · p. 4
• Discounted suffix returns assign earlier chunks credit for consequences predicted at later rollout steps, then normalize returns across trajectories at each chunk position.Normalization uses the group mean and population standard deviation, with a positive constant for numerical stability.Source: D. Cross-Model Feedback and Temporal… · p. 4

E. Action-Head Policy Optimization

• FPO evaluates the WAM’s joint latent endpoint with a surrogate that retains the native UniPC rollout without requiring likelihoods over the solver trajectory.The clipped objective uses action-head conditional flow-matching losses and group-relative advantages, without gradients through reward computation or the action-conditioned world model.Source: E. Action-Head Policy Optimization · p. 4
• The imagination–execution gap narrows after post-training; each plotted point represents a task, and distance below the diagonal marks the gap.The figure compares success rates before and after RL.Source: Figure 3 · p. 5
• Only action-head parameters are updated; the WAM’s visual-generation parameters and the action-conditioned world model stay frozen.Sampled trajectories, probe pairs, advantages, and behavior-policy losses remain fixed during optimization.Source: E. Action-Head Policy Optimization · p. 5

• The evaluation asks whether WAM visual predictions match consequences predicted from generated actions.Source: §V · p. 5
• The experimental tasks include laptop closing, placing a specified colored duck, pushing a T-shaped block, and nesting four cups.Cup nesting has the longest horizon; the tasks use 40, 200, 150, and 100 demonstrations, respectively.Source: Figure 4 · p. 5
• The experiments also test whether model-derived feedback improves real-world performance and how trajectory-level and terminal-goal feedback affect optimization.Source: §V · p. 5
• Four aspects organize the evaluation: model consistency, real-robot performance, feedback mechanisms, and reward effectiveness.Source: §V · p. 5

1) Tasks:

• The four UR5 tasks are laptop closing, instruction-specified duck placement, Push-T pose matching, and stacking four cups.The tasks use 40, 200, 150, and 100 demonstrations, respectively; cup nesting requires four successive grasp-and-place motions.Source: 1) Tasks · p. 5
• During SFT, imagined-video task completion rises from 18% to 50%, while action-conditioned predictions remain at 6–13%, leaving a 41-point gap at 15k steps.Source: Figure 5 · p. 6

2) Compared methods:

• The comparison includes the full framework, its SFT initialization, Fast-WAM, VLA-JEPA, and π0.5, adapted and evaluated using the same demonstrations and task protocol.Each external system retains its official model, training procedure, observation and action interface, and default execution horizon.Source: 2) Compared methods · p. 6

3) Training and checkpoint selection:

• All methods use the same task-specific UR5 demonstration split, divided 4:1 into training and testing sets.For each method, the selected checkpoint has the lowest validation loss during adaptation.Source: 3) Training and checkpoint selection · p. 6

4) Implementation details:

• FPO uses group-relative advantages over batches of eight groups of 16 trajectories, each with five action chunks.Trajectories within a group share an initial observation and instruction and differ by sampling noise; discounted suffix returns use γ = 0.95.Source: 4) Implementation details · p. 6

B. Model-Space Video–Consequence Gap During SFT

• During SFT, imagined-video success rises from 18% to 50%, while action-conditioned predictions remain at 6–13%, creating a 41-percentage-point gap at 15k steps.The paired rollouts use the same initial observations and action chunks, and a frozen instruction-conditioned classifier scores their terminal predictions.Source: B. Model-Space Video–Consequence Gap… · p. 6

C. Overall Real-Robot Performance

• 75.1% average real-robot success after RL exceeds 43.4% for the SFT initialization and 61.4% for π0.5 across four equally weighted tasks.The improvement over SFT is 31.7 percentage points, and the improvement over π0.5 is 13.7 percentage points.Source: Table I · p. 6, C. Overall Real-Robot Performance · p. 6
• Success improves on all four tasks: laptop 22.0% to 68.0%, duck 95.0% to 100.0%, Push-T 23.3% to 48.9%, and cup 33.3% to 83.3%.Source: C. Overall Real-Robot Performance · p. 6

D. Ablation Study of Reward

• On the close-laptop task, trajectory-only, terminal-only, and combined rewards each improve under their own FPO objective, while model-space success peaks early and later declines.The ablation holds the SFT initialization, Ctrl-World model, UniPC sampler, FPO objective, and training budget constant while varying reward composition; the two feedback components address intermediate evolution and the terminal goal, respectively.Source: D. Ablation Study of Reward · pp. 6–7, Figure 6 · p. 7

C. Overall Real-Robot Performance

• Mean completion time and control steps decrease on all four tasks after RL, including reductions of 16.8% and 25.1% on the long-horizon cup task.Medians follow the same trend except on duck; statistics use each policy’s successful trials, which are different trial sets.Source: C. Overall Real-Robot Performance · p. 6, Table II · p. 7

E. Correlation Between Reward and Predicted Success

• Successful rollouts score higher than failed ones on Laptop and Cup, while reward distributions overlap substantially on Duck and Push-T.The reward–success relationship is derived from Ctrl-World predictions and varies across tasks.Source: E. Correlation Between Reward and… · p. 7, Figure 7 · p. 7
• Across training and evaluation rollouts, successful predictions cluster at higher combined reward and predicted success probability than failures.Both quantities come from Ctrl-World, so this is a within-model-space association.Source: E. Correlation Between Reward and… · p. 7, Figure 8 · p. 8
• Trajectory feedback measures consistency between predicted action evolution and target progression, while terminal feedback emphasizes final-state completion.Combining them supervises both the intermediate process and final outcome.Source: E. Correlation Between Reward and… · p. 7

VI. CONCLUSIONS

• Post-training on the visual prediction–consequence gap raises average success from 43.4% to 75.1% across four real-world UR5 tasks.The conclusion also reports complementary roles for trajectory-level and terminal-goal feedback.Source: §VI · p. 8