# LUMOS: Language-Conditioned Imitation Learning with World Models

arXiv: https://arxiv.org/abs/2503.10370 · HTML: https://mlsanity.com/paper/2503.10370/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Placed in Embodied and robot world models, one of 11 directions in World models.
- Builds on 29 of its 38 references that are in the map.
- First posted on arXiv 13 Mar 2025.

## Placed in

- World models / [Embodied and robot world models](https://mlsanity.com/wm/embodied-and-robot-world-models/)
- Offline RL and imitation learning from data / [Language-conditioned policies and VLAs](https://mlsanity.com/offlinerl/language-conditioned-policies-and-vlas/)

## Builds on (most cited first, 29 of 29)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [CALVIN: A Benchmark for Language-Conditioned Policy Learning for Long-Horizon Robot Manipulation Tasks](https://mlsanity.com/paper/2112.03227/) (2021)
- [Learning agile and dynamic motor skills for legged robots](https://mlsanity.com/paper/1901.08652/) (2019)
- [PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture Likelihood and Other Modifications](https://mlsanity.com/paper/1701.05517/) (2017)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning](https://mlsanity.com/paper/2107.09645/) (2021)
- [Learning Latent Plans from Play](https://mlsanity.com/paper/1903.01973/) (2019)
- [3D Diffuser Actor: Policy Diffusion with 3D Scene Representations](https://mlsanity.com/paper/2402.10885/) (2024)
- [Zero-Shot Robotic Manipulation with Pretrained Image-Editing Diffusion Models](https://mlsanity.com/paper/2310.10639/) (2023)
- [Goal-Conditioned Imitation Learning using Score-based Diffusion Policies](https://mlsanity.com/paper/2304.02532/) (2023)
- [Language Conditioned Imitation Learning over Unstructured Data](https://mlsanity.com/paper/2005.07648/) (2020)
- [Scaling Up and Distilling Down: Language-Guided Robot Skill Acquisition](https://mlsanity.com/paper/2307.14535/) (2023)
- [What Matters in Language Conditioned Robotic Imitation Learning over Unstructured Data](https://mlsanity.com/paper/2204.06252/) (2022)
- [Multimodal Diffusion Transformer: Learning Versatile Behavior from Multimodal Goals](https://mlsanity.com/paper/2407.05996/) (2024)
- [Grounding Language with Visual Affordances over Unstructured Data](https://mlsanity.com/paper/2210.01911/) (2022)
- [DITTO: Offline Imitation Learning with World Models](https://mlsanity.com/paper/2302.03086/) (2023)
- [Language Control Diffusion: Efficiently Scaling through Space, Time, and Tasks](https://mlsanity.com/paper/2210.15629/) (2022)
- [Language-Conditioned Imitation Learning with Base Skill Priors under Unstructured Data](https://mlsanity.com/paper/2305.19075/) (2023)

## Built on by (most cited first, 14 of 14)

- [Vision-Language-Action Models for Robotics: A Review Towards Real-World Applications](https://mlsanity.com/paper/2510.07077/) (2025)
- [A Comprehensive Survey on World Models for Embodied AI](https://mlsanity.com/paper/2510.16732/) (2025)
- [A Survey: Learning Embodied Intelligence from Physical Simulators and World Models](https://mlsanity.com/paper/2507.00917/) (2025)
- [DynaGuide: Steering Diffusion Polices with Active Dynamic Guidance](https://mlsanity.com/paper/2506.13922/) (2025)
- [DiWA: Diffusion Policy Adaptation with World Models](https://mlsanity.com/paper/2508.03645/) (2025)
- [Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling](https://mlsanity.com/paper/2603.04553/) (2026)
- [Language-Conditioned World Modeling for Visual Navigation](https://mlsanity.com/paper/2603.26741/) (2026)
- [Enhancing Policy Learning with World-Action Model](https://mlsanity.com/paper/2603.28955/) (2026)
- [An Anatomy of Vision-Language-Action Models: From Modules to Milestones and Challenges](https://mlsanity.com/paper/2512.11362/) (2025)
- [AnoleVLA: Lightweight Vision-Language-Action Model with Deep State Space Models for Mobile Manipulation](https://mlsanity.com/paper/2603.15046/) (2026)
- [Wavelet Policy: Imitation Learning in the Scale Domain with World Prior Memory](https://mlsanity.com/paper/2504.04991/) (2025)
- [Prompting Robot Teams with Natural Language](https://mlsanity.com/paper/2509.24575/) (2025)
- [BLINK: Behavioral Latent Modeling of NK Cell Cytotoxicity](https://mlsanity.com/paper/2603.05110/) (2026)
- [Dream to Remember: Continual Steering of Generalist Robot Policies with World Models](https://mlsanity.com/paper/2610.13816/) (2026)

## Abstract (arXiv)

We introduce LUMOS, a language-conditioned multi-task imitation learning framework for robotics. LUMOS learns skills by practicing them over many long-horizon rollouts in the latent space of a learned world model and transfers these skills zero-shot to a real robot. By learning on-policy in the latent space of the learned world model, our algorithm mitigates policy-induced distribution shift which most offline imitation learning methods suffer from. LUMOS learns from unstructured play data with fewer than 1% hindsight language annotations but is steerable with language commands at test time. We achieve this coherent long-horizon performance by combining latent planning with both image- and language-based hindsight goal relabeling during training, and by optimizing an intrinsic reward defined in the latent space of the world model over multiple time steps, effectively reducing covariate shift. In experiments on the difficult long-horizon CALVIN benchmark, LUMOS outperforms prior learning-based methods with comparable approaches on chained multi-task evaluations. To the best of our knowledge, we are the first to learn a language-conditioned continuous visuomotor control for a real-world robot within an offline world model. Videos, dataset and code are available at http://lumos.cs.uni-freiburg.de.