LUMOS: Language-Conditioned Imitation Learning with World Models

Innovations

• Language-conditioned imitation learning framework that learns skills by practicing in the latent space of a learned world model
• On-policy learning in latent space to mitigate policy-induced distribution shift common in offline imitation learning
• Uses fewer than 1% hindsight language annotations from unstructured play data while remaining steerable with language commands at test time
• Combines latent planning with both image- and language-based hindsight goal relabeling during training
• Optimizes an intrinsic reward defined in the latent space of the world model over multiple time steps to reduce covariate shift
• First to learn a language-conditioned continuous visuomotor control for a real-world robot within an offline world model

Methodology

LUMOS learns a world model from offline unstructured play data with sparse language annotations. During training, it performs on-policy rollouts in the latent space of the world model, using hindsight goal relabeling (image and language) and an intrinsic reward to mitigate distribution shift. The learned policy is then transferred zero-shot to a real robot for language-conditioned control.

Key Results

On the long-horizon CALVIN benchmark, LUMOS outperforms prior learning-based methods with comparable approaches on chained multi-task evaluations. It also demonstrates zero-shot transfer of language-conditioned visuomotor control to a real robot.

Limitations

• Relies on learning a world model from offline data, which may be difficult to obtain or generalize for highly diverse environments
• Real-world validation is limited to a single robot setup and specific tasks, leaving broader generalization unverified
• The requirement for less than 1% language annotations may still depend on the quality and coverage of the unstructured play data