LUMOS: Language-Conditioned Imitation Learning with World Models

TL;DR

LUMOS introduces a language-conditioned imitation framework trained entirely in the latent space of a learned world model to tackle long-horizon robotic tasks. By combining latent planning, hindsight relabeling, and an intrinsic latent-space reward, it mitigates covariate shift and enables zero-shot transfer to real robots. The approach demonstrates strong performance on CALVIN in simulation and robust real-world transfer on a Franka Panda, outperforming prior methods and providing detailed ablations confirming the value of latent planning, language alignment, and latent rewards. This work advances offline, multi-task, language-conditioned manipulation by showing that dynamics and policies learned in a world model can generalize to real-world tasks without online fine-tuning, with practical implications for scalable robot learning from unstructured data.

Abstract

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.

Figures (6)

• Figure 1: LUMOS learns a general-purpose language-conditioned visuomotor policy within the latent space of a learned world model. By optimizing an intrinsic reward to match expert performance, it recovers from its own mistakes across multiple time steps and reduces covariate shift. As a result, LUMOS can handle complex, long-horizon robot tasks from abstract language instructions in real-world scenarios, without online learning or fine-tuning.
• Figure 2: LUMOS learns a language-guided general-purpose policy within the latent space of a world model. (1) The world model, comprising an image encoder, a Recurrent State-Space Model (RSSM) for dynamics, and an image decoder, transforms play dataset experience into a predictive model that enables behavior learning in the latent state space. (2) The goal-conditioned policy samples latent trajectories and uses either a language annotation or the final latent state as the goal, with plan recognition and proposal networks being trained to identify and organize behaviors in a latent plan space. The action decoder is intrinsically rewarded by matching the expert's latent trajectory. (3) During inference, the policy acts based on the latent state inferred by the world model from the current observation and is guided by a user's language command.
• Figure 3: World Model Long-Horizon Predictions. Using the first five images of a hold-out trajectory as context, the world model predicts the next 195 steps using its latent dynamics, given only the actions. Trained on sequences of horizon 50, our model makes precise long-term predictions, aiding efficient behavior learning in a compact latent space. Only gripper camera reconstructions are visualized.
• Figure 4: Real-world Tasks. Examples from left to right are: placing the carrot onto the pan, lifting the carrot from the table, dropping the carrot into the bowl, storing the carrot in the cabinet, placing the pan into the cabinet, setting the pan onto the stove, adding the eggplant to the pan, picking up the pan from the table, taking the eggplant from the cabinet, and putting the eggplant into the bowl.
• Figure 5: Visualization of the complete real-world robot setup. This figure highlights the static and gripper cameras, as well as the VR controller and tracking system.