Contextual Latent World Models for Offline Meta Reinforcement Learning

ICLR 2026 Conference Submission20756 Authors

Keywords: Meta Learning, Offline Reinforcement Learning, World Modeling, Representation Learning

Abstract: Offline meta-reinforcement learning seeks to overcome the challenges of poor generalization and expensive data collection by leveraging datasets for related tasks. Context encoding is a prevalent approach, where an encoder maps transition histories to a task representation. In parallel, latent world models -- which map observations into temporally consistent latent spaces -- advanced self-supervised representation learning for planning and policy optimization. In this work, we unify these directions by introducing contextual latent world models: world models conditioned on the task representation and trained jointly with the context encoder. Coupling task inference with predictive modeling yields task representations that capture variation factors across tasks and empirically improves generalization to out-of-distribution tasks in diverse benchmarks, including MuJoCo, Contextual-DeepMind Control suite, and Meta-World.

Primary Area: reinforcement learning

Submission Number: 20756