Contextual Latent World Models for Offline Meta Reinforcement Learning
This work addresses the problem of task representation learning in offline meta-reinforcement learning for researchers and practitioners, offering a novel method that improves generalization but is incremental in building on existing latent world model approaches.
The paper tackled the challenge of learning effective task representations for offline meta-reinforcement learning without supervision by introducing contextual latent world models, which condition latent world models on inferred task representations and train them jointly, resulting in significant improvements in generalization to unseen tasks across benchmarks like MuJoCo, Contextual-DeepMind Control, and Meta-World.
Offline meta-reinforcement learning seeks to learn policies that generalize across related tasks from fixed datasets. Context-based methods infer a task representation from transition histories, but learning effective task representations without supervision remains a challenge. In parallel, latent world models have demonstrated strong self-supervised representation learning through temporal consistency. We introduce contextual latent world models, which condition latent world models on inferred task representations and train them jointly with the context encoder. This enforces task-conditioned temporal consistency, yielding task representations that capture task-dependent dynamics rather than merely discriminating between tasks. Our method learns more expressive task representations and significantly improves generalization to unseen tasks across MuJoCo, Contextual-DeepMind Control, and Meta-World benchmarks.