Paper deep dive
Towards VLA-Dreamer: Refining VLA Behavior Using World Models
Parsa Mastouri Kashani, Jan-Gerrit Habekost, Stefan Wermter
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Vision-Language-Action models (VLAs), while showing strong potential for robot control, require massive amounts of high-quality imitation learning data. Moreover, the absence of an explicit world model casts further doubt on their control capabilities. In this concept paper, we propose a novel architecture that addresses sample efficiency in VLAs by training a predictive world model on the embedding space of the VLA's vision encoder. We hypothesize that these embeddings are action-relevant and usable for future prediction. To this end, we propose using the suggested architecture to investigate how well these embeddings predict the future based on actions, as the inability to do so would mark a key limitation of VLA architectures: the lack of a non-lossy implicit world model to simulate real-world dynamics. The proposed architecture differs from the standard world model dynamics as the loss comes from the embedding space rather than the pixel space, similar to joint embedding predictive architectures. Furthermore, the trained world model can be utilized for short-term planning tasks by sampling VLA actions given goal images. We intend to examine the richness of vision embeddings in VLAs and reduce their high data requirements through a world model that can also generate plans during inference.
Tags
Links
- Source: https://arxiv.org/abs/2609.31313v1
- Canonical: https://arxiv.org/abs/2609.31313v1
Trouble viewing inline? Open PDF directly →
Full Text
19,224 characters extracted from source content.
Expand or collapse full text
Towards VLA-Dreamer: Refining VLA Behavior Using World Models Parsa Mastouri Kashani, Jan-Gerrit Habekost, and Stefan Wermter Knowledge Technology, Dept. of Informatics, University of Hamburg, Germany parsa.mastouri.kashani, jan-gerrit.habekost, stefan.wermter@uni-hamburg.de Abstract. Vision-Language-Action models (VLAs), while showing strong potential for robot control, require massive amounts of high-quality im- itation learning data. Moreover, the absence of an explicit world model casts further doubt on their control capabilities. In this concept paper, we propose a novel architecture that addresses sample efficiency in VLAs by training a predictive world model on the embedding space of the VLA’s vision encoder. We hypothesize that these embeddings are action- relevant and usable for future prediction. To this end, we propose using the suggested architecture to investigate how well these embeddings pre- dict the future based on actions, as the inability to do so would mark a key limitation of VLA architectures: the lack of a non-lossy implicit world model to simulate real-world dynamics. The proposed architecture differs from the standard world model dynamics as the loss comes from the embedding space rather than the pixel space, similar to joint embed- ding predictive architectures. Furthermore, the trained world model can be utilized for short-term planning tasks by sampling VLA actions given goal images. We intend to examine the richness of vision embeddings in VLAs and reduce their high data requirements through a world model that can also generate plans during inference. Keywords: Joint-Embedding Predictive Architecture· Vision-Language- Action Models· World Model· Model-based RL· Robotics 1 Introduction Robotics has seen a recent surge with the introduction of Vision-Language- Action models (VLAs) [4, 7, 9], which are, in essence, Vision-Language Models (VLMs) that are themselves extensions of LLMs equipped with a vision back- bone. However, VLAs remain reactive models with no explicit model of the real world. World-Action Models (WAMs) [13, 14] address this by using video- generator backbones that can predict how the world might evolve. The JEPA [8] line of work, however, argues that because these generative objectives lie in pixel space, the models are forced to “think in pixels”, which is not an optimal surrogate for a mental model of the world, since such a model must account for every detail of a scene rather than reasoning at the right level of abstraction. arXiv:2609.31313v1 [cs.RO] 25 Sep 2026 2P. Mastouri Kashani et al. This forces the model to attend to unnecessary details, such as the leaves on the sidewalk in a self-driving car application. In response, models such as V-JEPA 2-AC [1] embed observations in a rich embedding space free of the aforementioned pixel-space overhead. Their vision encoder is first trained with a mask-denoising objective, which forces the en- coder to retain prediction-relevant information. An action-conditioned predictor is then trained on top of the frozen encoder to predict the next state given the current observations and an action. However, they rely on Model Predictive Control (MPC), an expensive planning procedure that iteratively selects actions whose predicted future embeddings minimize the distance to the goal embedding. Furthermore, their solution requires goal images, or sub-goal frames, of the ex- pected outcome at inference time, which might not be available and which limit the horizon of the tasks the model can perform between consecutive sub-goals. An alternative would be to train a lightweight MLP-based policy with RL, trained from scratch inside that embedding space, to achieve the same goal at a lower computational cost. Here, the world model serves as a simulator, while the agent remains small and simply reacts to the given state. However, such models lack the semantic priors needed for language conditioning and are forced to remain small to be trainable using RL without overfitting. Furthermore, designing a reward function for such a system is challenging as the objective changes from one task to another. We propose a novel approach: reusing the frozen vision encoder of a VLA, we train a world model that predicts the next observation embedding given an action, the last few observation embeddings, and the robot’s proprioception. We then continue training our VLA, which has already been fine-tuned on our embodiment, inside this world model using reinforcement learning. This approach addresses sample efficiency, and we hypothesize that having the VLA explore outcomes inside the world model would induce an efficient inner world model within the VLA agent itself, as it would be aware of the consequences of its actions. Furthermore, using a VLA allows us to gain further insight into the policy through question answering in the embedding space, since the embeddings are similar to the language-aligned vision encoder embeddings. 2 Methodology 2.1 Training and Choice of VLA Choice of VLA Many VLAs, such as SmolVLA [12], keep the vision encoder of their VLM backbone frozen during training. Such decisions, while plausible for training on limited data domains, would rely on the vision encoder to have encoded all necessary action-related information from the image, such as the precise location of each object in the scene. Although the encoder in such models is trained on VQA-style tasks during VLM pre-training, it might not capture all the action-related features needed by the policy. As a result, we pivot to using a VLA whose training has also updated the vision encoder. We hypothesize it would lead to better action-related features for future embedding generation. Towards VLA-Dreamer: Refining VLA Behavior Using World Models3 Based on these considerations, we propose to build onπ 0 -FAST [11] and OpenVLA [7] as our VLAs of choice, since they both operate on token prob- abilities rather than diffusion or flow-matching, which allows for a simpler RL pipeline. Furthermore, DINOv2 [10] has already proven to be able to generate future embeddings in DINO-WM [16], and OpenVLA makes use of a fusion of DINOv2 and SigLIP [15] adapted from the Prismatic line of VLMs [5]. Fig. 1. We fine-tune our pre-trained VLA model on our NICOL robot’s imitation learning dataset before further training the model using RL. VLA Training First, we fine-tune the VLAs given our current imitation learn- ing dataset for NICOL [6], which consists of language commands and teleoper- ation data, including proprioception (joint positions) and the camera frames of the robot at each timestep. To avoid misalignment between the world model em- beddings and the VLA’s vision encoder, we freeze the vision encoder. To avoid losing semantic knowledge, we also apply strong regularization and freeze most weights in the LLM backbone of the VLA. 2.2 World Model To train the world model we use a mixture of random actions and demonstra- tions by semi-expert or expert policies controlling the robot. This ensures good coverage of the world dynamics as training on expert demonstrations or random actions alone might leave the robot environment unexplored. The world model forπ 0 -FAST would be in charge of predicting the next proprioceptive state and the next-timestep frame embeddings conditioned on the current action, the cur- rent robot proprioceptive state, and a history of recent camera observations. The world model for OpenVLA would also predict proprioception, but only as an auxiliary loss, as proprioception is not a necessary part of the OpenVLA training recipe and would not be fed to the VLA model during RL training. Alongside the world model, we train an image decoder, but only for probing and further visualization, as the embeddings output by the world model are not optimal for image reconstruction, and the decoder is only used as a probe to see a lossy picture of what the model is dreaming. A summary of the process can be seen in Figure 2. 4P. Mastouri Kashani et al. Fig. 2. This figure demonstrates how the world model gets trained. A history of em- beddings and proprioception, along with an action, is input to the world model, which predicts the next observation embedding. 2.3 Embedding Assessment To test our central hypothesis that the frozen vision encoder of a VLA model retains enough action-relevant information to allow future prediction, we attach two probes to the vision encoder alongside the image decoder probe: a semantic segmentation probe and a depth estimation probe. We hypothesize that a capable manipulation policy is aware of object masks and scene geometry, and as a result, strong probe performance provides evidence that the embeddings encode them. Weak performance on the two given tasks indicates a limitation for both action generation and future prediction. The two quantitative values would be mIoU for segmentation and RMSE with threshold accuracy for depth. We expect the embeddings from OpenVLA’s Prismatic backbone to already perform well at depth estimation, because of its use of DINOv2. By contrast, the embeddings fromπ 0 -FAST may be less capable in this regard as its vision encoder begins as a SigLIP model and then undergoes VLM training as part of PaliGemma-3B [2], which is further trained under theπ 0 -FAST robot training recipe. This comparison therefore offers a useful demonstration of what addi- tional features the vision encoder acquires during VLA pre-training. Unlike models such as OpenVLA, theπ family of VLAs [3] are not restricted to a single input image. Depth estimation is therefore not constrained to the monocular setting, which can make the task easier for these models. Finally, the same probes apply to the embeddings predicted by the world model. Beyond visualization, this quantifies whether imagined futures preserve object masks and Towards VLA-Dreamer: Refining VLA Behavior Using World Models5 depth as the imagination horizon grows, turning probe accuracy on predicted embeddings into a direct measure of world-model fidelity and drift. 2.4 RL Training After fine-tuning the VLA on our embodiment and training the world model to predict future embeddings, we pick video frames of the robot executing tasks, for example, stacking a square block onto another. First, we embed the initial table configuration and then the goal image. Now we provide the correct prompt to our VLA policy and use the world model as a simulator for our policy to interact with. The negative mean squared distance between the current state embedding and the goal state embedding would be used as a dense reward function for the RL policy. The choice of VLAs, which output token probabilities, allows us to sample actions and to run model-free RL algorithms such as PPO. Since the VLA is already competent, RL serves to refine fine-grained actions rather than discover them, relaxing reward-shaping demands. We restrict the imagination horizon and keep the policy near its pre-trained prior to limit both compounding drift and exploitation of world model error. A summary of our approach can be seen in Figure 3. 3 Discussion and Challenges The major challenge for the system is its reliance on the vision encoder in VLA models to produce next-observation embeddings. Even though VLA models rely on their vision encoder to function and the embeddings are rich enough for action generation, they might not be enough for future prediction. While using mul- tiple observations and camera viewpoints for embedding generation, combined with proprioception data, helps mitigate this, it might not be enough in highly detailed and cluttered scenarios. Another concern is the effect of drift and compounding errors in the world model predictions. The amount of drift indicates the horizon of the imagination. Substantial drift might only allow us to implement an RL policy to refine the fine-grained actions rather than plan over longer trajectories. Our approach can be a better alternative to using simulation for generating trajectories. Generating imagination requires less compute than building and running a full simulator, and it avoids the sim-to-real gap as the world model is trained on real demonstrations and therefore captures task-relevant aspects of the environment that hand-built simulations often lack, especially in manip- ulation scenarios. However, as stated earlier, the performance depends on the quality of the embeddings coming from the vision encoder, and training an ad- equate world model requires substantial compute. Finally, using image goals to produce embeddings for the RL reward might prove difficult, as there may be a multitude of ways to capture visual snippets of the goal state. Furthermore, the mean squared distance between embeddings, depending on the task, might not be monotonic and hence might pose a challenge. 6P. Mastouri Kashani et al. Fig. 3. After training the VLA on our embodiment’s teleoperation dataset and training our world model and freezing it, we proceed to use the presented model for RL. By using the negative mean squared distance between the current state embedding and the goal state embedding, we provide a dense reward function for the VLA to act inside the world model. Towards VLA-Dreamer: Refining VLA Behavior Using World Models7 Our hypothesis is that the semantic abstractions of the embeddings and the language-defined goal given by the prompt would mitigate this phenomenon. Moreover, the VLA is under heavy regularization and includes most of the needed knowledge inside its pre-training recipe, so the reward would help with refining the trajectory. If proven insufficient, the reward can be replaced or augmented with a sparse success signal from a classifier for task success. 4 Conclusion We proposed an architecture that trains a predictive world model using the em- beddings of a VLA’s frozen, language-aligned vision encoder, and that refines the fine-tuned VLA through reinforcement learning inside this model. Because the policy input, the prediction target, and the goal-distance reward all share one representation, a single encoder supports action generation, future prediction, and goal specification. Beyond the architecture itself, the proposed experiments test whether the embeddings a VLA already relies on for action generation are rich and predictable enough to serve as the latent space of a world model. A positive result would suggest VLAs carry an implicit, usable model of world dy- namics, while a negative one would point to a limitation of current architectures, namely the absence of a non-lossy implicit world model. Our next step is to train the world model and probes, enabling us to measure embedding-prediction er- rors and probe fidelity as the imagination horizon grows, quantifying whether imagined futures stay geometrically coherent enough to support planning and RL. Acknowledgments. The authors gratefully acknowledge support from the German Research Foundation DFG under project LUMO (No 551629603) and funding from Horizon Europe under the MSCA grant agreements No 101226624 (GREET), No 101168792 (SWEET), and No 101072488 (TRAIL). Disclosure of Interests. The authors have no competing interests to declare that are relevant to the content of this article. References 1. Assran, M., Bardes, A., Fan, D., Garrido, Q., Howes, R., Komeili, M., Muckley, M., Rizvi, A., Roberts, C., Sinha, K., et al.: V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning. arXiv preprint arXiv:2506.09985 (2025). 2. Beyer, L., Steiner, A., Susano Pinto, A., Kolesnikov, A., Wang, X., Salz, D., Neu- mann, M., Alabdulmohsin, I., Tschannen, M., Bugliarello, E., et al.: PaliGemma: A Versatile 3B VLM for Transfer. (2024). https://doi.org/10.48550/arXiv.2407.07726 3. Black, K., Brown, N., Driess, D., Esmail, A., Equi, M.R., Finn, C., Fusai, N., Groom, L., Hausman, K., Ichter, B., et al.:π 0 : A Vision-Language-Action Flow Model for General Robot Control. In: Proceedings of Robotics: Science and Systems (RSS), Los Angeles, CA, USA (June 2025). 8P. Mastouri Kashani et al. 4. Brohan, A., Brown, N., Carbajal, J., Chebotar, Y., Chen, X., Choromanski, K., Ding, T., Driess, D., Dubey, A., Finn, C., et al.: RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control. In: Proceedings of The 7th Conference on Robot Learning, p. 2165–2183. PMLR (2023). 5. Karamcheti, S., Nair, S., Balakrishna, A., Liang, P., Kollar, T., Sadigh, D.: Pris- matic VLMs: Investigating the design space of visually-conditioned language mod- els. In Forty-first International Conference on Machine Learning, 2024. 6. Kerzel, M., Allgeuer, P., Strahl, E., Frick, N., Habekost, J.-G., Eppe, M., Wermter, S.: NICOL: A Neuro-Inspired Collaborative Semi-Humanoid Robot That Bridges Social Interaction and Reliable Manipulation. IEEE Access 11, 123531–123542 (2023). https://doi.org/10.1109/ACCESS.2023.3329370 7. Kim, M.J., Pertsch, K., Karamcheti, S., Xiao, T., Balakrishna, A., Nair, S., Rafailov, R., Foster, E., Lam, G., Sanketi, P., et al.: OpenVLA: An Open-Source Vision-Language-Action Model. In: Proceedings of the Conference on Robot Learn- ing (CoRL) (2024). 8. LeCun, Y.: A Path Towards Autonomous Machine Intelligence, Version 0.9.2. Open Review (2022). https://openreview.net/pdf?id=BZ5a1r-kVsf 9. Octo Model Team, Ghosh, D., Walke, H., Pertsch, K., Black, K., Mees, O., Dasari, S., Hejna, J., Xu, C., Luo, J., et al.: Octo: An Open-Source Generalist Robot Policy. In: Proceedings of Robotics: Science and Systems, Delft, Netherlands (2024). 10. Oquab, M., Darcet, T., Moutakanni, T., Vo, H.V., Szafraniec, M., Khalidov, V., Fernandez, P., Haziza, D., Massa, F., El-Nouby, A., et al.: DINOv2: Learning Robust Visual Features without Supervision. Transactions on Machine Learning Research (2024). 11. Pertsch, K., Stachowicz, K., Ichter, B., Driess, D., Nair, S., Vuong, Q., Mees, O., Finn, C., Levine, S.: FAST: Efficient Action Tokenization for Vision-Language- Action Models. In: Proceedings of Robotics: Science and Systems (RSS), Paper 12 (2025). 12. Shukor, M., Aubakirova, D., Capuano, F., Kooijmans, P., Palma, S., Zoui- tine, A., Aractingi, M., Pascal, C., Russi, M., Marafioti, A., et al.: SmolVLA: A Vision-Language-Action Model for Affordable and Efficient Robotics (2025). https://doi.org/10.48550/arXiv.2506.01844 13. Wu, H., Jing, Y., Cheang, C., Chen, G., Xu, J., Li, X., Liu, M., Li, H., Kong, T.: Unleashing Large-Scale Video Generative Pre-training for Visual Robot Ma- nipulation. In: Proceedings of the Twelfth International Conference on Learning Representations (ICLR) (2024). 14. Ye, S., Ge, Y., Zheng, K., Gao, S., Yu, S., Kurian, G., Indupuru, S., Tan, Y.L., Zhu, C., Xiang, J., et al.: World Action Models are Zero-shot Policies. (2026). https://doi.org/10.48550/arXiv.2602.15922 15. Zhai, X., Mustafa, B., Kolesnikov, A., Beyer, L.: Sigmoid Loss for Language Im- age Pre-Training. In: Proceedings of the IEEE/CVF International Conference on Computer Vision, p. 11975–11986 (2023). 16. Zhou, G., Pan, H., LeCun, Y., Pinto, L.: DINO-WM: World Mod- els on Pre-trained Visual Features enable Zero-shot Planning (2024). https://doi.org/10.48550/arXiv.2411.04983