WMPO: World Model-based Policy Optimization for Vision-Language-Action Models
Authors: Fangqi Zhu
YAN Zhengyang
Zicong Hong
Quanxin Shou
Xiao Ma
Song Guo
Venue: ICLR 2026 Poster (2026)
Keywords: World Models; Vision-Language-Action Models; Reinforcement Learning
Primary area: applications to robotics, autonomy, planning
TL;DR: WMPO is a world-model-based framework that conducts reinforcement learning for Vision-Language-Action policies entirely in imagination.

Abstract: Vision-Language-Action (VLA) models have shown strong potential for general-purpose robotic manipulation, but their reliance on expert demonstrations limits their ability to learn from failures and perform self-corrections. 
Reinforcement learning (RL) addresses these through self-improving interactions with the physical environment, but suffers from high sample complexity on real robots.
We introduce World-Model-based Policy Optimization (WMPO), a principled framework for on-policy VLA RL without interacting with the real environment.
In contrast to widely used latent world models, 
WMPO focuses on pixel-based predictions that align the "imagined" trajectories with the VLA features pretrained with web-scale images.
Crucially, WMPO enables the policy to perform on-policy GRPO that provides stronger performance than the often-used off-policy methods.
Extensive experiments in both simulation and real-robot settings demonstrate that WMPO (i) substantially improves sample efficiency, (ii) achieves stronger overall performance, (iii) exhibits emergent behaviors such as self-correction, and (iv) demonstrates robust generalization and lifelong learning capabilities.

Published as a conference paper at ICLR 2026
WMPO: WORLD MODEL-BASED POLICY OPTIMIZA- TION FOR VISION-LANGUAGE-ACTIII MODELS
Fangqi Zhu', Zhengyang Yan', Zicong Hong's Quanxin Shou', Xiao Ma Song Guo 'Hong Kong University of Science and Technology *ByteDance Seed
http:/uuuconnnc.uss.h. xiao.ma @ bytedance.comm com, songguo @egquuo@usttr @ust.hk Project Page: https: /tw////mpp////////////////////// g thub. 10/
ABSTRACT
Vision-languagge-Actionn (VL.A) models have shown strong potential for general- purpose robotic manipulation, but their reliance on expert demonstrations limits their ability learn from failures and perform self-corrections. Reinforcement learning (RL) addresses these through self-improving interactions with the physical environment, but suffers from high sample complexity on real robots. We introduce Woded-Model- Policy Optimization (WMPO), principled framework for on-policy VLA RL without interacting with the real environment. In contrast to widely used latent world models, WMPO focuses on pixel-based predictions that align the "imagined" trajectories with the VLA features pretrained with web-scale images. Crucially, WMPO enables the policy to perform on-policy GRPO that provides stronger performance than the often-used off- off-policy methods. Extensive experiments in both simulation and real-robot settings demonstrate that WMPO (i) substantially improves sample efficiency, (ii) achieves stronger overall performance, (iii) exhibits emergent behaviors such self-correction, and (iv) demonstrates robust generalization and lifelong learning capabilities.
INTRODUCTION
(VLA) models (Brohan et al., 2023; Kim et al., 2024; Intelligence et al., 2025) have emerged as promising paradigm for general- pprppose robotic manipulation, enabling robots to follow natural language instructions in complex, unstructured environments. The dominant approach for training these models is imitation learning (IL) from large-scale human demonstra- tions (Lin et al., 2024). While effective in mimicking demonstrated behaviors, IL-trained policies often brittle. When encountering out-of- iitoo-ddistribttiin states not seen during training, they can take suboptimal actions that lead to compounding errors (Ross al., 2011), 1), making task completion or recovery nearly impossible (Fig. la).
Vision-Language-A
Reinforcement learning (RL) (Li et al., 2025b) offers natural solution this brittleness by allowing an agent to learn and improve through active interaction with its environment. This eeff-improvement process can lead to policies that are more robust and capable of recovering from failure. However, applying RL directly to real robots is notoriously sample-inefficient, requiring millions of interactions that are impractical, unsafe, and tone-consuming to collect (Fig. 1b). Recent efforts improve sample efficiency fall into two main strategies. The first leverages human intervention to guide learning (Luo al., 2024b; Chen et al., 2025; Xia et al., 2025), which reduces exploration cost but is labor-intensive and hard to scale. The second relies on simulation to reduce real-world interactions (Lu et al., 2025; Li et al., 2025a), but is limited by the difficulty of building accurate simulators for diverse scenarios.
The advent of large-scale generative models, particularly video-generative world models (NVIDIA et al., 2025; Ball et al., 2025), presents a compelling new frontier for leveraging model-based RL (Finn & Levine, 2017) to mitigate the sample inefficiency challenge in VLA RL. These models can learn the dynamics of the world from data and simulate future transitions, offering a path to scalable RL without costly real-world explorations. Nevertheless, integrating these models with existing VLAs remains a challenge. Many classical model-based RI. approaches (Hafner et al., 2019; 2020a;b;
Corresponding authors.

---

Published as a conference paper at ICLR 2026
Learn From Fallures Self-coorretton Sample Efficiency On policy RL
Traject tories
World Model
Human Operation
Human Demonstration
Real World
Polley Behavier Data
Initation Learning
Reward
Action
8
Learn From Failunes
Self salf-correction
Learn From Failures
Sample Efficiency
VLA
VLA
On-policy RL
VLA
Imitation Learning
(b) Real World
(c) WMPO
Figure 1: Three different VLA training paradigms: (a) Imitation learning learns from human demonstrations but lacks the ability for learning from failures and self-correction; (b) Real- world RL improves policy through direct interaction but suffers from high sampling costs and difficulty in achieving on-policy RL; (c) WMPO pretrains world model on large-scale robotic trajectories and fine-tunes it with limited policy behavior data, enabling samplc-efficiet on on-policy RL for VLA without real-world interaction.
2023; Ma et al., 2022; Frauenknecht et al., 2025) operate in an abstract latent space, which creates fundamental mismatch with powerful VLA foundation models that are pretrained on real-world images. The rich, pretrained visual understanding of VLAs cannot be directly applied within a mismatched latent dynamics model. We argue that leveraging pixel-space video- generative world model crucial, as allows the VLA policy to operate on generated visual data that is consistent with its pretraining, effectively bridging the world model with the policy's pre-trained knowledge.
To this end, we propose World Model- boaeel-based Policy Optimization (WMPO), as illustrated in Fig. 1c, a principled framework that grounds VLA RL entirely in action-conditioned video world model. By pretraining a high-fidelity, pixel-space video-generative model on millions of robotic trajecto- risaaaalla,rratinn, 2023), WMPO leverages realistic visual dynamics create scalable training environment for downstream tasks.
WMPO incorporates several key innovations. First, to mitigate the state-distribution mismatch between expert demonstrations and policy rollouts, we introduce policy behavior alignment, finetuming the world model with behavioral data collected by the policy itself. Second, short-horizon prediction makes it difficult to define accurate rewards and is prone to reward hacking. To address this, WMPO generates complete trials through clip-level autoregressive video generation, enabling more reliable outcome-based reward assignment. While this design supports long-horizon rollouts, it also introduces challenges such as visual distortion and action-frame misalignment. WMPO addresses these with noisy-frame conditioning and frame-level action control, ensuring robustness and accurate trajectory simulation. Third, we train a lightweight reward model that predicts task success or failure, providing learned sparse reward signal and avoiding both complex reward shaping and reward hacking. Together, these components form a self-contained environment that enables on-policy RL. entirely "in imagination".. Specifically, we adopt Group Relative Policy Optimization (GRPO) (Shao et al., 2024), whose robustness and scalability have been demonstrated in DeepSeek-Al et al. (2025). Notably. WMPO naturally supports repeated rollouts from the same initial state, which is difficult realize in the physical world but crucial for large-scale GRPO training.
We conduct extensive experiments in both Mimicgen simulation environments (Mandlekar al.. 2023) and real-woris environments to validate the effectiveness of WMPO, Our results show that WMPO achieves substantially higher sample efficiency and consistently outperforms VLA RL methods that directly optimize with real trajectories. Crucially, we provide qualitative evidence of emergent behaviors, where the WMPO-trrinee policy demonstrates self-correction strategies not present in the demonstration data and often completes tasks faster and more smoothly, without noticeable stalls. We further demonstrate WMPO's 's strong generalization compared to offline RL methods, well its capacity for lifelong learning through alternating updates between the VLA policy and the world model. Taken together, these findings highlight WMPO as a scalable and generalizable paradigm for advancing VL. RI.
2

---

Published as a conference paper at ICLR 2026
2
RELATED WORK
LsfionLLanguage-Action Models. Vision-Language-Action (VLA) models aim to map visual inputs and natural language instructions into executable robotic actions, enabling general-purpose manipulation. Most VLAs build upon pretrained vision-language models (VLMs) and are further fine- tuned on robotic trajectories, thereby inheriting strong visual grounding and language understanding. This progress has been driven both by large-scale demonstration collections (Collaboration, 2023) and by advances in action parameterization, including discrete next-token prediction (Kim et al., 2024), parallel decoding (Kim et al., 2025), and continuous flow-based heads (Black et al., 2024). Despite these advances in data and model design, existing VLAs largely remain within the imitation learning (IL) paradigm, making them brittle when encountering out-of -uttrf-distribution states and unable to leverage failed exccutions for improvement (Ross et al., 2011).
Reinforcement Learning for VLA Models. RL provides natural complement to IL by enabling policies learn from interaction, thereby improving robustness and recovery capabilittes. However, applying on-policy RL to VLA remains challenging due to severe sample inefficiency and substantial system-level complexity. Prior works can broadly divided into two strategies. The first introduces human intervention to guide exploration, e.g., supplying corrective signals when policies encounter unrecoverable states (Luo et al., 2024a; Chen et al., 2025; Xia et al., 2025). While effective at reducing exploration cost, this approach requires continuous human supervision, making it labor- intensive and difficult to scale. The second attempts to perform RL. directly in simulation or in the real world. For instance, Zhang et al. (2024) adopt trajectory-level DPO (Rafailov et al., 2023), while others apply PPO (Schulman et al., 2017) or GRPO (Shao et al., 2024) to optimize VLA policies in simulation (Lu et al., 2025; Li et al., 2025a). These approaches avoid human supervision but still suffer from poor sample efficiency, and constructing accurate simulators for each real-world scenario introduces prohibitive engineering overhead. In contrast, WMPO enables policy optimization entirely within a learned world model, substantially improving sample efficiency and scalability.
World Models. World models aim mitigate the need for costly real-world interactions by learning generative dynamics and enabling policies to learn from "imagine" trajectories. Early approaches (Hafner et al., 2020a;b; 2023) learned world model in the latent space of recurrent state-space models, which provided efficient but overly abstract rollouts. More recent work intro- duced diffusion-based world models, showing that retaining pixel-level fidelity crucial for RL with Gaussian policies (Alonso et al., 2024; Jiang et al., 2025). Building on this trend, large-scale video world models (NVIDIA et al., 2025; Ball et al., 2025) have demonstrated impressive gen- erality across diverse domains. However, when applied to robotics, they suffer from distribution mismatch-struggling to faithfully reproduce policy rollouts and fine-grained robot-object interac- tions. In contrast, WMPO addresses these challenges through policy behavior alignment and, for the first time, verifies the feasibility of leveraging high-fidelity world models for scalable VLA RL.
WORLD MODEL-BASED POLICY OPTIMIZATION
We introduce World Model-based Policy Optimization (WMPO), a novel framework for learning complex, vision-based robotic manipulation policies. The WMPO framework is grounded in the principles of model-based RL, where policy optimization is performed entirely within learned world model, thereby eliminating the need for costly real-world interactions. WMPO operates directly in the pixel space, instead of latent space, which better bridges the pretrained VLA features from web-scale video-action pairs with the "imagined" trajectories. Specifically, the world model is trained to accurately reflect the policy's 's behavior through process we call Policy Behavior Alignment, where it is fine-tuned on a small set of real trajectories collected from the policy itself. This ensures the model can faithfully simulate diverse scenarios, including failures. We also introduce noisy frame conditioning and frame-level action control techniques to overcome the problems of visual distortion and cetion-frame misalignment in long-horizon video prediction. With these modifications, we are able to perform strong on on-policy Group Relative Policy Optimization (GRPO) using trajectories imagined by the learned world model, which significantly enhances sample efficiency compared to direct RL methods. Fig. 2 illustrates overview of the training procedure of WMPO.
3

---

Published as a conference paper at ICLR 2026
Imagined Trajectory Generation
Asscct Asaac
Poliyy Mode!
Updeee
Upcons
Policy Model
sefoulos)
N0aid
ru00
Buffer
(5lommo de
3. Policy Update
B
Lakgage 3
Fetiel State
2. Trajectory Sampling
World Model P
TrajecteryT
collnn
aeace
Reward Model Rg.
Computatonn
Trojectory T
Trajectory
Figure 2: WMPO starts from an initial state sg- The overall training procedure consists of three components: (1) Imagined Trajectory Generation, where policy model Tew and world model pa interact alternately to generate full imagined trajectory; (2) Trajectory Sampling, where multiple trajectories are sampled and evaluated by the reward model Re; and (3) Policy Update, where the policy parameters are optimized via Eq. 4. This process is iteratively repeated throughout training.
3.1 PROBLEM FORMULATION
We formulate the VLA manipulation task as a decision-making problem in the form of an MDP M P, R).
G, where I denotes the image observation space, i.e., image sequences Jo.K, and G denotes the language instruction space. Here, we make the assumption that robot states can be solely defined by their image observations. For more complicated setups, e.g., Partially Observable MDPs (POMDPs) (Igl et al., 2018; Ma et al., 2021), we leave them for future studies.
State space. S = I
Action space. A denotes the space of action chunks, i.e., sequences of robot actions. Each action an action chunk of length K, a E IR KR''D, represents a D-DoF (degree of freedom) control vector. For policy optimization, each dimension is discretized into 256 bins. Specifically, we learn to sample action chunks using a parameterized VLA policy, a ~ (a 8).
Transition function. P S x A + S is realized by parameterized world model 8141 Pø(8++1 8t, Qt), which generates future observations conditioned on past observations and actions. In particular, we sample imagined trajectories = {8o, $1,..., ,rr by iteratively sampling from the world model $141 ~ (0 8t, ae) and the VLA policy a; ~ Pa(a; 8t), given an initial state so sampled in the real environment.
Reward function. The reward is given by learned model Ry that determines task success from the full trajectory, assigning binary outcome Rs (T) E 10.
Our objective is to train a policy (a(a) 8) such that the predicted cumulative return of the imagined trajectories will be maximized
max ET-to.ps [R) (T)]
This formulation highlights general paradigm: RL for VLA can be decoupled from real-world interactions by leveraging a generative world model as the imaginary training ground.
3.2 GENERATIVE WORLD MODEL
Imagined Trajectory Generation. Given initial frames Io the policy na takes the most recent m frames and language instruction g as input and predicts an action chunk i.e., axs+ K (o(Ti- m:i g). The world model Ps then conditions on the last observed frames and the predicted action chunk to generate the next K frames:
Litt ~ Pol(t-est k.)
(2)
'To avoid confusion in subscripts. we specify that i is used as the frame-level subscript and t as the state-level subscript, where state often includes multiple frames. N and T denote the maximum length of an imagined trajectory the frame level and state level, respectively.
4

---

Published as a conference paper at ICLR 2026
Repeating this process until maximum length N yields trajectory T = {Iu.N ag:N). A reward model Rg evaluates the frames in T and outputs a binary label y = R E {0, 1}. Thus, each imagined trajectory in the world model is represented as a labeled pair (r, V), which is then used for policy optimization.
Model Architecture. Our world model is based a video diffusion backbone inherited from OpenSora (Zheng et al., 2024), with modifications designed for simulating robot-object interactions. Specifically, we replace the 3D VAE in OpenSora with the 2D VAE from SDXL (Podell al., 2024), which better preserves fine-grained motion details and avoids temporal distortions caused by excessive compression. Consequently, the diffusion process operates in the VAE's latent space. When applying the imagined trajectory to VLA optimization, we decode the images back into pixel space to better leverage the pretrained knowledge, rather than retraining the VLA in latent space such as that from RSSM (Hafner et al., 2023).
Since our world model generates full trajectories in autoregressive manner- previously generated frames as conditioning for future predictions-errort can accumulate, leading to severe degradation over long horizons, ultimately leading to failed prediction. To mitigate this issue, we introduce a noisy-frame conditioning technique: during training, conditional frames 4-m are perturbed with mild diffusion noise corresponding to an early timestep (i.e., the noise level at step 50, where step 1000 would correspond to pure Gaussian noise), rather than kept clean. This improves robustness to imperfect conditioning.
As a result, the world model achieves stable long-horizon generation, producing trajectories of hundreds frames without noticeable quality loss.
To enable precise action conditioning, inspired by Zhu et al. (2025), we incorporate a frame-level action control mechanism for better ection-frame alignment. Specifically, we extend the AdaLN (Xu et al., 2019) block to inject both action signals and diffusion timestep embeddings at the frame level. For each action ag an MLP generates modulation coefficients: scale y and shift 3 for the LayerNorm output, and scale af for the residual connection of either the MHA or FFN block. Let x' denote the feature representation at frame i; the update rule within each transformer block is given as:
x" x + (1 af) Block (1 LayerNorm(x') + 81).
Policy Behavior Alignment. We pretrain the world model on Open X- X-Embodiment (OXE) Collab- oration (2023) trajectories, which offer diverse demonstrations of robot interactions and endow the model with broad knowledge of physical dynamics. However, because OXE trajectories primarily consist of successful executions, failure scenarios are underrepresented in the observation distribution. Likewise, training only expert demonstrations from downstream tasks leaves the model unable to simulate failures, making the imagined trajectories unsuitable for training. To address this mismatch, we fine-tune the world model on real rollout trajectories collected from the policy itself, thereby adapting it to the downstream (state, action) distribution and capturing failure modes more faithfully. Without this adaptation, the model's imagination of failure cases remains brittle and unfaithful.
3.3 REWARD MODEL.
key requirement for scalable policy optimization in the world model is automatically judging whether an imagined trajectory indicates task success. We construct lightweight reward model trained on real trajectories. Given a trajectory 7 = {4o:N] }, we define clip of length L as = 1 La- The terminal clip ey of a successful trajectory serves as a positive sample, whereas negatives are drawn from {c: L < i < N L} of successful trajectories and from arbitrary clips of failed trajectories. To address class imbalance, we balance the number of positive and negative samples within each training batch. The reward model, implemented as a VideoMAE (Tong et al., 2022) encoder with linear head, is trained with binary cross-entropy loss. At inference, the model applies sliding window with stride classified as successful if any clip exceeds a threshold Tehrs which is selected via validation experiments.
7 to compute the success probability of each clip. A trajectory
5

---

Published as a conference paper at ICLR 2026
3.4 ON- ON-POLICY REINFORCEMENT LEARNING WITH WMPO
Reinforcement learning in VLA tasks faces two key bottlenecks. (1) Physical interaction bottleneck.
Unlike LLM settings where feedback is cheap, VLA tasks require repeated rollouts in the physical world, which incur high hardware costs, safety concerns, and limited scalability. (2) Value estimation bias by off-policy RL. Given the aforementioned physical constraints, existing real-world RL methods often consider off-policy RL methods (James et al., 2022; Seo et al.. 2024; Wagenmaker et al., 2025; Chen et al., 2025). However, off off-policy methods naturally cause value estimation errors (Park et al., 2025) and on-policy methods are often favorable for better performance.
To overcome these challenges, we optimize policies entirely within world model: replacing costly real-world rollouts with model-generated trajectories alleviates reliance on physical interaction and enables scalable online learning- We adopt Group Relative Policy Optimization (GRPO) as the policy optimization algorithm, since it provides stable and scalable training in settings with sparse rewards. In our case, state transitions are simulated by the world model (Eq. 2), and rewards are binary with (T) E {0, depending task success.
Trajectory Sampling. From each initial frames Io.e sample from real environment D, we sample a group of imagined trajectories (Ti....; TG) from current policy Te inside the world model. The reward model is then employed to predict the success or failure each trajectory. To mitigate vanishing gradients, we adopt Dynamic Sampling strategy (Yu al., 2025): all trajectories in group are predicted to be successful or all unsuccessful, the group is discarded and additional rollouts are sampled until the batch fully populated. The pog-probability of each action chunk under Tew pre-computed as reference:
D log mia (a)" 8x):
log Thaa (at. 8t)
where a, denotes the action chunk time t, and a. represents the i-th action in the j-th DoF.
Policy Update. Following DAPO (Yu et al., 2025), we remove the KL divergence regularization so that no reference model is required during training, thereby reducing memory consumption and encouraging the policy to explore novel behaviors. The final training objective is given by
J(0)= (rr))
min (O), (ip(ra,.()) 1 Clow; + (high A.)
(4)
with
To aa | 8i,e)
R - mean ( A
(f)
(5)
(ai, s2,1)²
(td(
Here re. (0) is the probability ratio between new and old policies at step of trajectory т., R. = R(T)). and A, is the normalized advantage of trajectory T₃ over the horizon N. The overall training pipeline detailed Algorithm Appendix A.
EXPERIMENTS
We conduct extensive experiments to evaluate the effectiveness WMPO, focusing on the following questions: (1) can WMPO outperform online and offline RL in simulation environments; (2) how does the behavior of WMPO differ from imitation learning; (3) can WMPO generalize to unseen settings; (4) can WMPO achieve iterative improvement during deployment; and (5) can WMPO be applied on real robot?
4.1 EXPERIMENT SETTINGS
Implementation Details In this work, fine-tune OpenVLA-OFF (Kim et al., 2025) via imitation learning on target manipulation tasks as our base policy. For simplicity, we omit the robot propri- oceptive state and wrist camera inputs, and set the action chunk length K to 8. We collect P real

---

Published as a conference paper at ICLR 2026
Base Policy
WMPO
Continue moving down
Collision
Self-correction
Figure Behavior analysis of the Square task (inserting the square into the stick) shows that, compared with the base policy, WMPO demonstrates the ability to self-correct.
trajectories using the base policy, and conduct experiments with P = 128 and P 1280 to evaluate the scalability our approach. These trajectories are further used to fine-tune a world model, which predicts the next K frames given 4 conditioning frames and action chunk. In addition, we train reward models on the collected trajectories, using a video clip length of L = 8 and stride of during evaluation. Further training details are provided in Appendix A.
Simulation Environment Details We conduct experiments in the Mimicgen simulation (Man- dlekar et al., 2023). We select four fine-gra aneeggrined manipulation tasks: Coffee DO, StackThree DO, ThreePiiceAssembly DO, and Square DO, and fine-tune the Open VLA-OFT model with 300 expert trajectories per task as the base policy. For evaluation, we test 28 different initial states for each task and report the average success rate.
4.2 COMPARISON EXPERIMENTS
We compare WMPO with two established RL algorithms, GRPO (Shao al., 2024) and DPO (Rafailov et al., 2023), both widely used for optimizing large language models. To ensure fairness, all methods are allocated the same real rollout budget P. We consider both online and offline baselines: GRPO is implemented in an online setting, where the policy is updated directly from trajectories collected the environment; DPO is implemented an offline setting, where the base policy serves as the reference and trajectory pairs (success vs. failure) are constructed for optimization using the standard DPO loss. Unlike GRPO, which discards trajectories after each update, DPO can repeatedly
collected data, but it lacks the ability to update the policy online WMPO does. More implementation details are provided in Appendix B.
As shown in Tab. 1, WMPO consistently outperforms all baselines across all tasks. With small rollout budget of P-128, it already surpasses the strongest baseline by +9.8 points, demonstrating
Table Comparison of policy optimization methods across four manipulation tasks the Mimicgen simulation benchmark. /' denotes the rollout budget, i.e., the number of full real trajectories available for policy optimization. Results show that WMPO consistently outperforms both GRPO and DPO baselines under different budgets. As the rollout budget increases from 128 to 1280, WMPO continues to exhibit substantial improvements, highlighting both its data efficiency and scalability. Performance reported the task success rate (%).
Rollout budget P
Methods
Coffee
StackThree
ThreePicceAssembly Square
Mcan (%)
Base policy
43.8
46.9
19.5
24.2
33.6
GRPO
38.3
52.3
17.2
25.0
33.2
128
DPO
43,8
53,9
23.4
28.1
37.3
Ours
61.7
56.3
37.5
32.8
47.1
GRPO
47.7
54.7
20.3
25.8
37.1
1280
DPO
52.3
57.0
26.7
33.6
42,4
Ours
75.0
64.1
46.1
45.3
57.6

---

Published as a conference paper at ICLR 2026
Fixed Stick
Random Stick
White Bg.
Grey Bg.
Red Base
Black Base
(a) Position Disruption
(b) Background Disruption
(c) Texture Disruption
Figure 4: (a) For the Square task, we vary the stick's position from fixed a random position inside rectangle. (b) For the StackThree task, we substitute the tabletop background with a gray background. (c) For the ThreePieceAssembly task, we substitute the red base with a dark wooden base.
strong data efficiency. When the budget increases to P=1280, the margin further expands to +15.2 points average, indicating that WMPO leverages additional trajectories more effectively than existing methods and scales robustly with more rollouts. contrast, GRPO often underperforms with limited updates, and DPO plateaus due static data reuse, whereas WMPO continues improve steadily as the rollout budget grows. Furthermore, we evaluate the reward model and find that achieves an Fl score above 0.95 across all tasks, reliably distinguishing success from failure and effectively mitigating reward hacking. These results highlight the effectiveness and scalability of WMPO for policy optimization in robotic manipulation.
4.3 EMERGENT BEHAVIOR OF WMPO
To better understand the source of WMPO's strong performance, we conduct a visual comparison of its test-time behavior against the base policy. We identify two emergent behaviors unique to WMPO: (I) the WMPO policy learns to self-correct, recovering from nearly failure states; and (2) the WMPO policy executes tasks more efficiently, as rarely becomes "stuck" in suboptimal states.
First, taking the Square task an illustrative example (see Fig. 3), observe that when both the base policy and WMPO deviate from the correct trajectory due to error accumulation and encounter a collision, their behaviors diverge. The baseline policy, trained only on expert demonstrations, has never observed collisions during training: it continues to push the square against the stick until the maximum time horizon is reached, resulting in failure. In contrast, WMPO benefits from large-scale imagined trajectories generated by the world model, enabling it to learn the self-correction behaviors, which is challenging to obtain through imitation learning alone. Specifically, the policy autonomously learns to lift the square, realign it, and then insert it correctly, ultimately succeeding in the task. Second, we analyzed the lengths of successful trajectories generated by different policies, as shown in Fig. 5. We found that the trajectories of policies trained with WMPO are significantly shorter. This because WMPO discourages stuck behaviors, which often result in failures due to timeouts. As side benefit, WMPO also makes the policy's behavior faster and smoother.
4.4 GENERALIZATION TO NOVEL TASKS
In this section, we evaluate the generalization ability of WMPO across three novel disruption scenarios (Fig. 4), which systematically assess generalization under spatial, background, and texture shifts. As shown in Tab. 2, WMPO consistently achieves the best performance across all disruption types. DPO attains modest improvements in the in-distribution setting compared to the base policy, but its performance degrades significantly under background and texture changes, suggesting reliance on spurious visual cues rather than transferable manipulation skills. GRPO
Table We evaluate each policy in its corresponding disruption scenario and report the success rate (%).
Methods
Pos. Dis.
Bg. Dis.
Tex. Dis.
Miaa
Base policy GRPO
14.1
46.1
10.9
23.7
47.7
DPO
16.4
34.4
7.8
19.5
Ours
22.3
50.0
29.6
exhibits performance similar to the base policy, and both are scenarios. In contrast, WMPO, trained entirely in the world model, captures more generalizable strategies and maintains reliable performance across spatial, background, and texture variations.
than WMPO across all disruption
8

---

Published as a conference paper at ICLR 2026
104
Base Policy
Bate Palicy
Relative Length (%)
102
DPO
WMPO
100
Success Ratt
50
94
45
ThraePlacokssombly
Square
Rollout Budget
Figure 5: Relative average trajectory length of success- ful trials across different policies (Base Policy 100%).
Figure 6: Lifelong learning sults of WMPO and baselines.
Real World Trajectory
Imagined Trajectory
Figure 7: Real-world experiments on the fine-grained manipulation task "Insert the square into the stick" where the clearance between the square and the stick is only 5 mm. The top row shows the real-world trajectory of the base policy executed in the real world, while the bottom row depicts the corresponding imagined trajectory starting from the same initial state within our learned world model. Despite never observing this trajectory during training, the world model accurately predicts the future evolution, demonstrating its ability capture precise task dynamics.
4.5 LIFELONG LEARNING
In this section, we demonstrate that WMPO can continuously improve the performance VLA by iteratively collecting real trajectories from the environment. Specifically, we iteratively collect P= 128 real trajectories, perform WMPO optimize the policy, and then use the updated policy to collect another P real trajectories. We apply the same setting to the DPO baseline. To compare WMPO with an imitation learning-basse policy using more expert demonstrations, we leverage 300, 428, and 556 expert trajectories to train the base policy reference. It is important to note that the base policy requires human-collected trajectories, whereas WMPO only relies on trajectories collected by the policy itself, making it more scalable. The results the StackThree task, shown in Fig. 6, demonstrate that WMPO achieves stable and substantial improvements over both baselines, whereas DPO fails improve iteratively due unstable training.
4.6 REAL-WORLD EXPERIMENTS
In this section, we evaluate the challenging real-world manipulation task, "Insert the square into the stick" (see Fig. 7, more cases including failure could be found in Appendix C), to validate the effectiveness of WMPO. Using the Cobot Mobile ALOHA platform, we collect 200 high-quality expert demonstrations to fine-tune the OLenVLA-OFT model as the base policy. We then deploy this policy collect additional 128 trajectories, which are used to further fine-tune the world model and optimize the policy within it. For comparison, we also train an offline DPO policy using the same dataset. All models are evaluated under identical experimental conditions, and report the average success rate over 30 trials. The results show that the base policy, DPO, and WMPO achieve success rates of 53%, 60%, and 70%, respectively, demonstrating the effectiveness of WMPO real robots.

---

Published as a conference paper at ICLR 2026
5 CONCLUSION
In this work, we introduced WMPO, a novel framework for on-policy RL of VLA models. By grounding policy optimization entirely in video-generative world model, WMPO eliminates the need for costly real- wolld interactions while maintaining consistency with pretrained VLA representations. Through policy behavior alignment, robust autoregressive video generation, and a lightweight reward model, WMPO enables scalable training with strong sample efficiency. Extensive experiments in both simulation and real-worke settings demonstrated that WMPO (i) consistently outperforms state-of-the of-th-rr-tre-tttttt model-free baselines, (ii) exhibits emergent self-correction behavior, (ii) generalizes reliably unseen scenarios, and (iv) supports iterative lifelong learning. Together, these findings highlight WMPO as a scalable and generalizable paradigm for advancing VLA RI..
6 ETHICS STATEMENT
This work uses only publicly available, open-source datasets for training and evaluation. No private or sensitive data is involved. The proposed methods are intended solely for academic research in machine learning and robotics. No foreseeable negative societal impacts are anticipated beyond standard considerations for robotic learning research.
REPRODUCIBILITY STATEMENT
We have made every effort to ensure the reproducibility of our work. All datasets used in this study are publicly available and open source. Detailed descriptions of the experimental settings, hyperpa- rameters, and training protocols are provided in Appendix A. We also release our implementation at https: / //// thub. com /WM-PO. PO/ 1/0/in /MMO, which includes training scripts, evaluation pipelines, and configuration files for all experiments. To further support downstream research, we additionally provide pretrained checkpoints of the world model and task-specific fine-tuned models, enabling researchers to directly reproduce our results and easily adapt the models to new tasks.
ACKNOWLEDGMENTS
We very grateful to Dr. Hongtao Wu for his insightful discussions and valuable suggestions in the early stage of this work. This research was supported by fundings from the Hong Kong RGC General Research Fund (152228/23E, 162161/24E, 1621 16/25,, 162180/25E), National Natural Science Foundation of China (NSFC) Key Program 6N0.62530055, Collaborative Research Fund (No. C1042-23GF, No. 5097-25G), NSFC/RGC Collaborative Research Scheme (Grant No. 62461 160332 & CRS HKUST602/24), Research Impact Fund (No. R5011-23F) Areas of Excellence Scheme (A0F/E--601/22-R), and the InnoHK (HKGAI).
REFERENCES
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Flcuret. Diffusion for world modeling: Visual details matter in atari. In The Thirty- eighth Annual Conference on Neural Information Processing Systems, 2024. URL https:
// //openreview wwt.wwiwwwee..eeee..e her///ormmmm/// iiddadddwooogg.....
Philip J. Ball, Jakob Bauer, Frank Belletti, Bethanie Brownfield, Ariel Ephrat, Shlomi Fruchter, Agrim Gupta, Kristian Holsheimer, Aleksander Holynski, Jiri Hron, Christos Kaplanis. Marjorie Limont, Matt McGill, Yanko Oliveira, Jack Parker-Holder, Frank Perbet, Guy Scully, Jeremy Shar, Stephen Spencer, Omer Tov, Ruben Villegas, Emma Wang. Jessica Yung. Cip Baetu, Jordi Berbel, David Bridson, Jake Bruce, Gavin Buttimore, Sarah Chakera, Bilva Chandra, Paul Collins, Alex Cullum, Bogdan Damoc, Vibha Dasagi, Maxime Gazeau, Charles Gbadamosi. Woohyun Han, Ed Hirst, Ashyana Kachra, Lucie Kerley, Kristian Kjems, Eva Knoepfel, Vika Koriakin, Jessica Lo. Cong Lu, Zeb Mehring, Alex Moufarek, Henna Nandwani, Valeria Oliveira, Fabio Pardo, Jane Park, Andrew Pierson, Ben Poole, Helen Ran, Tim Salimans, Manuel Sanchez, Igor Saprykin, Amy Shen, Sailesh Sidhwani, Duncan Smith, Joe Stanton, Hamish Tomlinson, Dimple Vijaykumar, Luyu Wang, Piers Wingfield, Nat Wong, Keyang Xu, Christopher Yew, Nick Young, Vadim Zubov,
10

---

Published as a conference paper at ICLR 2026
Douglas Eck, Dumitru Erhan, Koray Kavukcuoglu, Demis Hassabis, Zoubin Gharamani, Raia Hadsell, Aäron van den Oord, Inbar Mosseri, Adrian Bolton, Satinder Singh, and Tim Rocktäschel. Genie 3: A new frontier for world models. 2025.
Kevin Black, Noah Brown, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, Lachy Groom, Karol Hausman, Brian Ichter, et al. wo: A vision-language-acionn flow model for general robot control. arXiv preprini urXiv:2410.21164, 2024.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choromanski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, et al. Rt-2: Vision-language-actinn models transfer web knowledge to robotic control. arXiv preprint arXiv:2307.18814 2023.
Yuhui Chen, Shuai Tian, Shugao Liu, Yingting Zhou, Haoran Li, and Dongbin Zhao. Conrft: A rein- forced fine-tuning method for vla models via consistency policy, arXiv preprint arXiv:2502.05450, 2025.
Open X-Embodiment Collaboration. Open X-Embodiment: Robotic learning datasets and RT-X models.
org. abs//3/0..011// 08864, 2023.
DeepSeek-AI et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning, 2025. URL https: / arxiv, org/ abs/2501, 12948.
Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. In 2017 IEEE International Conference on Robotics and Automation (ICRA), pp. 2786-2793 IEEE Press, 2017. doi: 10.110//CRA.2211.7999333 URL httpt:ttt//oi.rrgg/ggeee 1109/ICRA 2017. 7989324.
Bernd Frauenknecht, Devdutt Subhasish, Friedrich Solowjow, and Sebastian Trimpe. On rollouts in model-based reinforcement learning. In The Thirteenth International Conference on Learning Representations, 2025. URL https: h//ppppeevvieee net/ forum? 1i==hhhGRmiive
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pp. 2555-2565, PMLR, 2019.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. ICLR 2020, 2020a.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193 2020b.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint 2301v:2301.0411 2023.
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for pomdps. In International conference on machine learning, pp. 2117- 2126. PMLR, 2018
Physical Intelligence, Kevin Black, Noah Brown, James Darpinian, Karan Dhabalia, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, et al. mn.s: a vision-language-action model with open-world generalization. arXiv preprint arXiv:2504. 16054, 2025.
Stephen James, Kentaro Wada, Tristan Laidlow, and Andrew Davison. Coarse-to-finn fine q-attention: Ef- ficient learning for visual robotic manipulation via discretisation. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 13748, 2022.
Zhennan Jiang, Kai Liu, Yuxin Qin, Shuai Tian, Yupeng Zheng, Mingcai Zhou, Chao Yu, Haoran Li, and Dongbin Zhao. World4rl: Diffusion world models for policy refinement with reinforcement learning for robotic manipulation, 2025. URL https: /axiiiiii..////vi org/ abs/2500 19080.
Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan Foster, Grace Lam, Pannag Sanketi, et al. Openvla: An open-source vision-language-action model. arXiv preprint arXivv2406.09244. 2024.
11

---

Published as a conference paper at ICLR 2026
Moo Jin Kim, Chelsea Finn, and Percy Liang. Fine-tuning vision-language-actiee models: Optimizing speed and success, 2025. URL https://aaxxi org. abs 2502 19645.
Haozhan Li, Yuxin Zuo, Jiale Yu, Yuhao Zhang, Zhaohui Yang, Kaiyan Zhang, Xuekai Zhu, Yuchen Zhang, Tianxing Chen, Ganqu Cui, Dehui Wang, Dingxiang Luo, Yuchen Fan, Youbang Sun, Jia Zeng, Jiangmiao Pang, Shanghang Zhang, Yu Wang, Yao Mu, Bowen Zhou, and Ning Ding. Simplevla-rl: Scaling vla training via reinforcement learning, 2025a. URL https: //axpiiv//rxxv.. org/abs/ 2500. 09674.
Qiyang Li, Zhiyuan Zhou, and Sergey Levine. Reinforcement learning with action chunking, 2025b. 07969.
URL https:tt/s/ett //axxiv///xvvvva org/
Fanqi Lin, Yingdong Hu, Pingyue Sheng, Chuan Wen, Jiacheng You, and Yang Gao. Data scaling laws imitation learning for robotic manipulation, 2024. URL https::/aaat //tp:///aaivv./.eeeeaeeeeeeeevee org/ / 2410 0. 18647
Jie Liu, Gongye Liu, Jiajun Liang, Yangguang Li, Jiaheng Liu, Xintao Wang, Pengfei Wan, Di Zhang, and Wanli Ouyang. Flow-grpo: Training flow matching models via online rl, 2025. URL https: /tle:///rrae /ttp//aavvvor////e orgg aaa//////2555 05470.
Guanxing Lu, Wenkai Guo, Chubin Zhang, Yuheng Zhou, Haonan Jiang, Zifeng Gao, Yansong Tang, and Ziwei Wang. Vla-rl: Towards masterful and general robotic manipulation with scalable reinforcement learning, 2025. URL https: : / //aa//aivvv//////////////////v.a org/ abs/ /5005 18719.
Jianlan Luo, Zheyuan Hu, Charles Xu, You Liang Tan, Jacob Berg, Archit Sharma, Stefan Schaal, Chelsea Finn, Abhishek Gupta, and Sergey Levine. Serl: A software suite for sample-efficient robotic reinforcement learning. In IEEE International Conference on Robotics and Automation (ICRA), pp. 16961-16969. IEEE. 2024a.
Jianlan Luo, Charles Xu, Jeffrey Wu, and Sergey Levine. Precise and dexterous robotic manipulation via human-in-the-loop reinforcement learning. arXiv preprint arXiv:2410.2/845, 2024b.
Xiao Ma, Siwei Chen, David Hsu, and Wee Sun Lee. Contrastive variational reinforcement learning for complex observations. In Conference robot learning, pp. 959-972. PMLR, 2021.
Xiao Ma, David Hsu, and Wee Sun Lee. Learning latent graph dynamics for visual manipulation of deformable objects. In 2022 International Conference on Robotics and Automation (ICRA), pp- 8266-8273. IEEE, 2022.
Ajay Mandlekar, Soroush Nasiriany, Bowen Wen, Iretiayo Akinola, Yashraj Narang, Linxi Fan, Yuke Zhu, and Dieter Fox. Mimicgen: A data generation system for scalable robot learning using human demonstrations. arXiv preprint arXiv: /0. /7596, 2023.
NVIDIA, Niket Agarwal, Arslan Ali, Maciej Bala, Yogesh Balaji, Erik Barker, Tiffany Cai, Prithvijit Chattopadhyay, Yongxin Chen, Yin Cui, Yifan Ding, Daniel Dworakowski, Jiaojiao Fan, Michele Fenzi, Francesco Ferroni, Sanja Fidler, Dieter Fox, Songwei Ge, Yunhao Ge, Jinwei Gu, Siddharth Gururani, Ethan He, Jiahui Huang. Jacob Huffman, Pooya Jannaty, Jingyi Jin, Seung Wook Kim, Gergely Klár, Grace Lam, Shiyi Lan, Laura Laaxe Anqi Li, Zhaoshuo Li, Chen- Hsuan Lin, Tsung-Yi Lin, Huan Ling, Ming-Yu Yu Liu, Xian Liu, Alice Luo, Qianli Ma, Hanzi Mao, Kaichun Mo, Arsalan Mousavian, Seungjun Nah, Sriharsha Niverty, David Page, Despoina Paschalidou, Zeeshan Patel, Lindsey Pavao, Morteza Ramezanali, Fitsum Reda, Xiaowei Ren, Vasanth Rao Naik Sabavat, Ed Schmerling, Stella Shi, Bartosz Stefaniak. Shitao Tang. Lyne Tchapmi, Przemek Tredak, Wei-Cheng Tseng, Jibin Varghese, Hao Wang, Haoxiang Wang, Heng Wang, Ting-Chun Wang, Fangyin Wei, Xinyue Jay Zhangjie Wu, Jiashu Xu, Wei Yang, Lin Yen-Chen, Xiaohui Zeng, Yu Zeng, Jing Zhang, Qinsheng Zhang, Yuxuan Zhang, Qingqing Zhao, and Artur Zolkowski. Cosmos world foundation model platform for physical ai, 2025. URL https:
/tttps:/arxxvv.rxt wttp://.rggaaaa///////o/// als atsr///bb/2501 03575.
Seohong Park, Kevin Frans, Deepinder Mann, Benjamin Eysenbach, Aviral Kumar, and Sergey Levine. Horizon reduction makes rl scalable. arXiv preprint 2500:22406.444 2025.
12

---

Published as a conference paper at ICLR 2026
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. SDXL: Improving latent diffusion models for higl-resolttion image synthesis. In The Twelfth International Conference Learning Representations, 2024. URL
https: http:///peeevvveeeeeeeeeeeeeeeee eww. net/ / forum? d=di R2
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2023.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In AISTATS, pp. 627-635, 2011.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 177.04777. 2017.
Younggyo Seo, Jafar Uruç, and Stephen James. Continuous control with coarse-to-fine reinforcement learning. arXiv preprint DxXix:2407.07788 2024.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024. URL httpp:://axxi......// org/ abs/ 2402. 03300.
Zhan Tong, Yibing Song, Jue Wang, and Limin Wang. VideoMAE: Masked autoencoders are data-efficient learners for self-supervised video pre-training. In Advances in Neural Information Processing Systems, 2022.
Andrew Wagenmaker, Mitsuhiko Nakamoto, Yunchu Zhang, Seohong Park, Waleed Yagoub, Anusha Nagabandi, Abhishek Gupta, and Sergey Levine. Steering your diffusion policy with latent space reinforcement learning. arXiv preprint arXiv:2506. 15799, 2025.
Wenke Xia, Yichu Yang. Hongtao Wu, Xiao Ma, Tao Kong, and Di Hu. Robotic policy learning via human-assisted action preference optimization, 2025. URL https: / aarxi.ii org abs/ 1556607122
Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization, 2019. URL https://arxiv org/ abs 07013.
Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, Haibin Lin, Zhiqi Lin, Bole Ma, Guangming Sheng, Yuxuan Tong, Chi Zhang, Mofan Zhang, Wang Zhang, Hang Zhu, Jinhua Zhu, Jiaze Chen, Jiangjie Chen, Chengyi Wang, Hongli Yu, Weinan Dai, Yuxuan Song. Xiangpeng Wei, Hao Zhou, Jingjing Liu, Wei- Wei-Ying Ma, Ya-Qin Zhang. Lin Yan, Mu Qiao, Yonghui Wu, and Mingxuan Wang. Dapo: An ooen-source IIm reinforcement learning system at scale, 2025. URL https::/p//sss...... http:///axxx...././//bo 14476.
Zijian Zhang. Kaiyuan Zheng, Zhaorun Chen, Joel Jang, Yi Li, Siwei Han, Chaogi Wang, Mingyu Ding. Dieter Fox, and Huaxiu Yao. Grape: Generalizing robot policy via preference alignment. arXiv preprint arXiv:2411.1930% 2024.
Zangwei Zheng, Xiangyu Peng, Tianji Yang, Chenhui Shen, Shenggui Li, Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all, 2024. URL https:/ arxiv. org/ abs 22022.2212.. 20404.
Fangqi Zhu, Hongtao Wu, Song Guo, Yuxiao Liu, Chilam Cheang, and Tao Kong. Irasim: A fine- grained world model for robot manipulation, 2025. URL https: //trxivivv./////////////////atv/ /ttp//aww.g//aabb..../. wtrg/bbb/2400000 14540.
13

---

Published as a conference paper at ICLR 2026
APPENDIX
A TRAINING DETAILS
In this section, we provide the training details of WMPO. The supervised finetuning of VpenVLA- OFT is conducted on 8 H100 GPUs. We exclude robot proprioceptive states and wrist camera inputs for simplicity. Instead of applying L1 regression, action tokens in parallel. The subsequent world model training and policy optimization are performed on 32 H100 GPUs.
adopt output head that predicts discrete
The world model is first pretrained on the Open X-Embodiment (OXE) dataset Collaboration (2023) and then fine-tuned on downstream policy behavior data. Table 3 summarizes the hyperparameters
for training the world model, while the remaining settings follow OpenSora (Zheng et al., 2024).
The hyperparameters for policy optimization with GRPO are provided in Tab. with other settings inherited from Li et al. (2025a).
Algorithm 1: World Model Policy Optimization
Input :orlld model Pæ, reward model Ry, VLA policy Te: group size G; batch size B (trajectories); mini-batch size M (M | By; epochs E; horizon
Output: Uppdttd policy parameters
Gold t
while not converged do
B < Ø
while [B] B do
Sample initial state = (To.e: g) ~ D
for to G
Imagine T2 ((s), at) po( 80; To
R4 + R(TT)) E 10
all R..... Rc) or mone(RR,.... RG) then L continue
1 -. Rr, ot v Σ. 1 (R, 1R,-0 1 p2)
for i == 1 do
A, (R4 -N)/
Precompute {log aema TT
Append (T., A.) to B
if [ > B then break
for e to E do
for j = 1 to B/M do
M + j-th contiguous block of M trajectories from B Update 9 according Eq. 4, where
(0))(0) = exp( log no (ooaa,))ss log To.ia (as,e((s,e))
oold <
B BASELINE DETAILS
We provide the implementation details of our baselines, including online GRPO and offline DPO. For the GRPO baseline, the main hyperparameters are inherited from Tab. 4. We observe that the batch size has a significant impact on performance. Larger batch sizes (e.g.. 64) yield more stable improvements; however, they require substantial number of real trajectories. Specifically, single model update requires at least 64 x 8 = 512 real trajectories. Moreover, dynamic sampling further filters out groups with success rate of such large batch sizes are infeasible, and even with P 1280, the model can only be updated twice. Therefore, additionally experimented with smaller batch size of 8, scaling the
1. As result, when the rollout budget is P 128,
14

---

Published as a conference paper at ICLR 2026
learning rate down proportionally (by factor of 8). We report the best results obtained across both configurations (batch size = and = 64).
For the DPO baseline, follow the standard preference-based offline training paradigm. Specifically, we construct a preference dataset from trajectories collected by the supervised fine-tuned Open VLA- OFT model and use it to optimize the policy with the DPO objective. The model architecture and optimizer settings are kept consistent with those used GRPO for fair comparison. All baselines were trained under the same rollout budgets and evaluation protocols as WMPO ensure fairness.
C REAL WORLD CASES
In this section, provide additional examples of trajectory predictions made by the world model in real-world settings. Fig 9 illustrates cases where the world model successfully predicts failure trajectories: has learned that when the square and the stick are misaligned, the square cannot be inserted into the stick. In contrast, Fig 10 shows failure case where the model does not correctly predict failed trajectory. Although predictions remain accurate until the final frame, subtle perturbations prevent the model from faithfully capturing the moment when the square gets stuck in the stick. Empirically, such failures are relatively rare on the validation set, indicating that the world model can reliably predict both successful and failed outcomes in the vast majority of scenarios, which is crucial for stable policy optimization.
D LIMITATION
While the WMPO framework can in principle support flow-based policies, this work focuses on dis- cretized action representations. As future work, plan to extend WMPO to more expressive policy classes, such as flow-matching based policies (Black et al., 2024), and explore policy optimization with FlowGRPO (Liu et al., 2025), thereby broadening its applicability across diverse action spaces.
E LLM USAGE STATEMENT
GPT-5 used solely assist with language refinement and stylistic polishing of the manuscript. The authors confirm that all scientific ideas, study design, data analyses, and conclusions presented in this work are entirely their own. The LLM did not contribute the generation of research concepts, execution of experiments, or interpretation of findings.
Hyperparameter
Value
Optimizer
AdamW
0.9, B 0.999) 0.0001
Learning rate
Batch size
128
Gradient clip
0.1
Pretrain training steps
12,000,000
Fine-tune training steps
3,000,000
EMA
0.9999
Weight decay Prediction target
0.0
Table 3: Hyperparameters for training the world model.
15

---

Published as a conference paper at ICLR 2026
Hyperparameter
Value
Optimizer
Adam W(3 8 = 0.9, 3 0.999)
Learning rate
x 10
Training batch size
64
Group size G
Mini-batch size
128
Clip ratio Elow
0.20
Clip ratio Ehigh Temperature
0.28
1.6
Table 4: Hyperparameters for the GRPO algorithm.
Coffee
Stackthree
Threepiccassemlll
Square
47-M MMMAA
26
Training Stops
(a) DPO, P = 128
Coffee
Stackthree
Threepiecessoombby
Square
19
42
35
15
50
28
Success
12
Training Steps
(b) WMPO, 128
Coffee
Stackhhree
Success hate
32
15
14
185
Training Steps
(c) DPO, P = 1280
Coffee
Stackthree
Thrrepiccassembly
Square
61.
48
34
38-
35-
Success
104 230
120
Training Steps
(d) WMPO. P 1280
Figure 8: Success rates of DPO and WMPO across four Mimicgen manipulation tasks (Coffee. StackThree, ThreePieceAssembly and Square) under different rollout budgets (P = 128 and P 1280). WMPO consistently achieves more stable improvements and higher final performance compared to DPO, with the performance gap widening under larger rollout budgets.
16

---

Published as a conference paper at ICLR 2026
Real World Trajectory
Imagined Trajectory
Figure 9: Real-world experiments on the fine-grained manipulation task "Insert the square into the stick". The top row shows the rollout trajectory of the base policy executed the real world, while the bottom row depicts the corresponding imagined trajectory starting from the same initial state within our learned world model. The world model successfully predicted failure cases.
Real World Trajectory
Imagined Trajectory
Figure 10: Example of a failure case. Although the predicted trajectory remains accurate until the final frame, the model fails capture the square getting stuck in the stick due to subtle perturbations.
17


Reviews:
[{"review_id": "iI3EcqKQm5", "reviewer": "Reviewer_VLDi", "summary": "WMPO proposes a novel policy finetuning method without using online interaction by using a world model on the pixel space, which allows robust finetuning of vision-language-action models on imagined trajectories. There are a few desirable properties of this method compared to traditional world modeling and finetuning objectives:\n1. The world model can directly be trained in pixel space, which allows a pixel based policy to interface with the world model without extra decoders.\n2. The method can be completely open sourced, and can be adapted into other distributions. \n3. This method allows on policy finetuning of VLAs using GRPO without any real-world demonstrations or preexisting datasets.\nThe authors then used GRPO to perform finetuning of a base OpenVLA-OFT policy, which shows desirable success rate, scalability, and robustness in RoboMimic. On the real world, WMPO also shows better performance compared to other finetuning methods such as direct GRPO and DPO.", "questions": "1. I might have missed this, but when you are increasing the rollout budget, do you use the larger set as well to finetune the world model?\n2. One reason for not using real-world demonstrations at all is due to it being more expensive. Have you considered using a few real-world demonstrations as regularization in your method, and if so, how does it compare to only using imagined trajectories?\n3. It seems that the authors chose to use on policy RL because of flaws with off policy RL [1] and propagating the correct value. Are there going to be any potential concerns when implementing such a method in the off-policy setting?\n\nMinor remarks:\n1. The main figure stated that BC from human demonstrations cannot achieve self-corrective behavior, but this is not correct when the training data contains corrective behavior. Similarly, you can do on-policy RL in real world if you can run it online.\n2. Follow up on q3, it would be better to show differences in scalability if you have to use an off policy setting.\n\nReferences:\n[1] Park, S. et al., 2025. “Horizon Reduction Makes RL Scalable.” NeurIPS.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The method section of the paper is concise and informative. \n2. I believe that there are sufficient ablations being done on the method, and I like that the authors have demonstrated good robustness of the method when dealing with OOD settings and desirable scalability.", "weaknesses": "1. I believe that the paper did not address how this method can be extended into a generalist setting. The scope of the environment is also rather limited, albeit there are adequate ablations being conducted.\n2. In addition, the paper used OpenVLA-OFT as the base policy. This again limits how much promise the method can bring to generalist policies. If the authors can provide additional ablations without using OpenVLA-OFT, I believe this can strengthen the paper.\n3. I believe that the paper did not address the question of how adaptable this method is when concerned with language-conditioned settings (even though OpenVLA is language conditioned, there are no mention of how to use language labels in this paper), furthering this discussion can also be beneficial to the rating of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ffCLMQTSmA", "reviewer": "Reviewer_Wgig", "summary": "This paper proposes WMPO (World-Model-based Policy Optimization), an on-policy RL framework for VLA models that replaces costly real-robot rollouts with trajectories “imagined” by a pixel-space video world model. Key ingredients are: (1) a diffusion video world model with noisy-frame conditioning for robustness and frame-level action control for precise action–frame alignment; (2) policy behavior alignment, i.e., fine-tuning the world model on the policy’s own rollouts to match failure as well as success states; (3) a lightweight clip-based reward model for outcome (0/1) labeling; and (4) on-policy GRPO training with dynamic sampling and no KL term (no reference model). The method aims to keep the visual state space aligned with pretrained VLA encoders by decoding back to pixels. Experiments on Mimicgen show consistent gains over GRPO and DPO baselines at two real-rollout budgets (e.g., mean SR 47.1% vs. 37.3% at P=128 and 57.6% vs. 42.4% at P=1280). WMPO also reports emergent self-correction, improved generalization under spatial/background/texture shifts, lifelong learning improvements via alternating policy/world-model updates, and a real-robot result on “Insert the square into the stick”", "questions": "1. please try to handle these weaknesses\n\n2. Do you try to use some PeFT methods for low-resource environments? Maybe these discussions will help", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Clear, modular recipe: world-model rollouts + outcome classifier + GRPO, with practical choices (pixel-space decoding to match VLA features; noisy-frame conditioning; frame-level action injection). The write-up is concrete and reproducible. \n\n\nConsistent empirical gains over strong baselines across four Mimicgen tasks and two rollout budgets; improvements grow with budget (data-efficiency + scaling). \n\n\nBehavioral insights: convincing qualitative evidence of self-correction and reduced “getting stuck,” with trajectory-length analysis. \n\n\nGeneralization: better robustness under spatial/background/texture disruptions than baselines.", "weaknesses": "Heavy compute / practicality: training uses 32× H100 for world-model/WMPO phases (plus 8× H100 for SFT). The paper would benefit from wall-clock, throughput, and ablations on smaller budgets/hardware. \n\nModel-world fidelity & safety: while qualitative results are strong, there’s limited quantitative assessment of rollout fidelity (e.g., per-step action-conditioned metrics), failure taxonomy, or safety constraints—especially since outcome-only rewards can reward shortcuts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BTuHk3d7dR", "reviewer": "Reviewer_vH24", "summary": "The paper proposes World Model-based Policy Optimization (WMPO), a framework for improving Vision-Language-Action (VLA) models via reinforcement learning inside a learned pixel-level video world model. The core idea is to avoid costly real-world robot interactions by performing on-policy RL (specifically GRPO) entirely within an autoregressive video diffusion model pretrained on large-scale robotic data and fine-tuned on downstream policy rollouts. The framework includes: A pixel-space world model for visual fidelity; Policy Behavior Alignment to adapt the world model to policy-induced state/action distributions; A learned binary reward model for task success; Noisy-frame conditioning and frame-level action control to mitigate long-horizon prediction drift; Dynamic Sampling during GRPO to avoid vanishing gradients in sparse-reward settings. Experiments in MimicGen simulation and real-world tasks (e.g., “insert square into stick”) show consistent improvement over imitation learning baselines and offline RL methods such as DPO and limited online GRPO.", "questions": "Please address the aforementioned weakness as thoroughly as possible.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The paper addresses a critical bottleneck in VLA RL: sample inefficiency and brittleness of imitation learning. While prior works (e.g., RT-2, OpenVLA) have shown impressive generalization, they remain confined to IL and struggle to recover from failures. WMPO’s goal — learning to self-correct via on-policy RL in a world model — is both ambitious and well-justified.\n\n2.  The authors’ fine-tuning of the world model on policy-generated trajectories is a principled way to close the distribution shift between expert demonstrations and actual policy rollouts.\n\n3. The qualitative results showing self-correction are compelling. This is not merely improvement in success rate but evidence of learning novel recovery strategies absent in demonstrations.", "weaknesses": "1. Overclaiming of “On-Policy” Scalability Without Real Costs: While WMPO avoids real-world rollouts during optimization, it still requires 128–1280 real trajectories to fine-tune the world model and initialize policy behavior alignment. This is not zero-shot or low-data RL — it is offline world-model RL with modest real data.\nRecent works like **IRASim** [1] and **World4RL** [2] also use diffusion world models but start from far fewer real trajectories (e.g., 50–100).\nThe paper does not compare directly to these methods, making it unclear whether the gains come from GRPO or from the specific world-model design.\n\n2. Reward Model is a Black Box with No Robustness Checks. The reward model is trained on binary success labels from real trajectories. But: What if the world model generates semantically correct but visually shifted outcomes (e.g., object slightly displaced)?\nThere is no evaluation of reward hacking — e.g., does the policy learn to “fool” the reward model by generating plausible but incorrect motions?\n\n3. Comparison to Offline RL is Weak. The DPO baseline is implemented naively: it uses trajectory-level preferences but does not leverage recent advances. Why not compare to IQL + diffusion policy? Also, DPO’s poor performance on background disruption (Tab 2) may reflect overfitting to visual cues, not the inherent weakness of offline RL.\n\n4. The quality of experiments: One task, 30 trials for the real robot task, underpowered for statistical significance. Missing key ablations (e.g., world model w/o policy alignment, latent vs. pixel space, reward model threshold sensitivity).\n\n[1] IRASim: A Fine-Grained World Model for Robot Manipulation \n\n[2] World4RL: Diffusion World Models for Policy Refinement with Reinforcement Learning for Robotic Manipulation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g15qTAOpsw", "reviewer": "Reviewer_eaSU", "summary": "The authors propose a system to train a VLA via RL on trajectories from a world model. A video diffusion model is transformed into a world model by additionally conditioning upon actions. For the reward model, a videoMAE model is finetuned on binary classification on successful/failed trajectories, and the model output probability is used as a reward. The used RL algorithm is GRPO, where the initial frame comes from a real dataset, and the trajectories in the cohort are different sampled videos starting from the same frame. For baseVLA, openvla-OFT is chosen. Across both sim and real, this method beats baselines, and the authors also show that the method improves generalization.", "questions": "1. What is meant by 50/1000 steps?\n2. Can you compare with PPO as well?\n3. Can you try multiple different VLAs?\n4. Can you do more real-world experiments?\n5. Can you add more details regarding how the world model was trained? Hyperparameters, datasets and so on.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The idea is simple and clearly explained\n- The authors show solid improvement on sim and real", "weaknesses": "- only a single base-VLA is used (openvla-OFT)\n- More real-world experiments would be good\n- missing comparison to DreamGen\n- the training of the world model lack details", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)