Contextual Latent World Models for Offline Meta Reinforcement Learning

This paper introduces Contextual Latent World Models, a self-supervised approach that jointly trains task inference and latent dynamics to learn expressive task representations, thereby significantly improving generalization to unseen tasks in offline meta-reinforcement learning across multiple benchmarks.

The Gist

Imagine you are trying to teach a robot to play a video game. But there's a catch: you can't let the robot play the game in real-time to learn. Instead, you have to give it a giant library of video recordings of other people playing different versions of the game.

This is the challenge of Offline Meta-Reinforcement Learning. The robot needs to learn a "meta-skill" from these static videos so that when it faces a new version of the game it has never seen before, it can adapt instantly.

The problem? Most robots are terrible at figuring out what is different about the new game just by watching a few seconds of gameplay. They get confused.

This paper introduces a new method called SPC (Self-Predictive Contextual Offline Meta-RL) to fix this. Here is how it works, explained through simple analogies.

1. The Problem: The "Blindfolded Chef"

Imagine a chef who has watched thousands of videos of people cooking different types of pasta.

• The Old Way: The chef tries to memorize the look of the ingredients (the "observation"). If the new pasta looks slightly different (maybe the sauce is a different shade of red), the chef panics because they are trying to match the exact visual details. They fail to realize the rules of cooking have changed, not just the colors.
• The Result: The chef can't cook the new pasta because they are too focused on the surface details.

2. The Solution: The "Storyteller" (Context Encoder)

The paper's method introduces a Context Encoder. Think of this as a Storyteller sitting next to the chef.

• Instead of just looking at the ingredients, the Storyteller watches the first few seconds of the cooking video and says, "Ah, this is a Spicy Tomato recipe," or "This is a Creamy Mushroom recipe."
• The Storyteller creates a Task Representation (a mental label) for the specific game or recipe. This label tells the robot, "Hey, in this specific version, the rules are X, Y, and Z."

3. The Secret Sauce: The "Crystal Ball" (Latent World Model)

Here is where the paper gets clever. How do you train the Storyteller to be accurate without a teacher telling them the recipe name?

The authors use a Latent World Model, which acts like a Crystal Ball.

• The Old Way (Reconstruction): Previous methods tried to train the Storyteller by asking them to "draw a picture" of what the next frame of the video would look like. This is hard and often leads to the Storyteller just memorizing the background scenery (like the kitchen tiles) instead of the actual cooking rules.
• The New Way (Temporal Consistency): The authors say, "Don't worry about drawing the picture. Just predict the future."
  • The robot asks: "If I am in this state and I do this action, what will happen next?"
  • The Storyteller must predict the future state of the game based on the current state.
  • The Magic: To predict the future accurately, the Storyteller must understand the underlying rules (dynamics) of the specific task. If the Storyteller doesn't know the task is "Spicy Tomato," they can't predict that the sauce will boil over differently.

By forcing the Storyteller to be a good predictor of the future, they accidentally become excellent at identifying the task. They learn the "soul" of the game, not just its "skin."

4. The Discrete Codebook: The "Library Index"

The paper also uses a technique called Finite Scalar Quantization (FSQ).

• Imagine the robot's brain is a massive library. Instead of trying to remember every single book (every possible continuous number), the robot organizes books into shelves with specific numbers.
• Instead of saying "The temperature is 23.456 degrees," the robot says, "The temperature is on Shelf 4."
• This makes the robot's brain much more efficient and less prone to getting confused by tiny, irrelevant details. It forces the robot to group similar situations together, making it easier to generalize.

5. The Result: The "Super-Adaptive Robot"

When you combine the Storyteller (who identifies the task) with the Crystal Ball (who predicts the future based on that task), you get a robot that:

1. Watches a few seconds of a new game.
2. The Storyteller instantly figures out, "This is a high-speed, slippery version of the game."
3. The robot uses this label to adjust its strategy immediately.
4. It performs better than any previous method at adapting to new, unseen tasks.

Summary Analogy

Think of learning to drive in different cities.

• Old Methods: You try to memorize the exact color of every building and the specific shade of the sky in every city. When you go to a new city with slightly different buildings, you crash because you are looking for the exact colors you memorized.
• SPC Method: You learn to recognize the traffic patterns and road rules (the "latent dynamics"). You realize, "Ah, this city drives on the left and has narrow streets." Once you understand the rules (the task), you can drive anywhere, even if the buildings look completely different.

The paper proves that by training the AI to be a good predictor of the future, it naturally becomes a master at understanding the context, leading to a robot that can learn new skills from old videos faster and more reliably than ever before.

Technical Summary

1. Problem Statement

Offline Meta-Reinforcement Learning (OMRL) aims to learn a policy that can generalize to unseen tasks using only fixed, offline datasets collected from related tasks, without further environment interaction.

• The Challenge: Existing context-based OMRL methods typically infer a task representation (context) from transition histories using contrastive learning (e.g., pulling same-task transitions together and pushing different tasks apart).
• The Limitation: Contrastive learning focuses on task discrimination but fails to explicitly enforce predictive structure over time. Consequently, the learned task representations often fail to capture task-specific dynamics and reward functions, leading to poor generalization. Furthermore, relying on observation reconstruction (as in some world models) is often unnecessary and computationally expensive, and can lead to representation collapse.

2. Methodology: Self-Predictive Contextual OMRL (SPC)

The authors propose SPC, a framework that unifies task inference and predictive modeling by conditioning a Latent World Model on an inferred task representation. The core innovation is training the context encoder and the world model jointly using temporal consistency rather than observation reconstruction.

Key Components

1. Context Encoder ($E_\theta$): Maps a history of transitions (context) to a latent task representation $z$.
2. Observation Encoder & Quantization ($F_\phi, f$): Maps raw observations $s_t$ to discrete latent states $c_t$ using Finite Scalar Quantization (FSQ). This discretization allows for efficient classification-based training.
3. Task-Conditioned Latent World Model:
   • Latent Dynamics ($D_\phi$): Predicts the next discrete latent state $\hat{c}_{t+1}$ given the current state $c_t$, action $a_t$, and task representation $z$.
   • Reward Model ($R_\phi$): Predicts the reward $\hat{r}_t$ given $(c_t, a_t, z)$.
4. Policy Optimization: Uses Implicit Q-Learning (IQL) to learn the policy $\pi$ and value function $V$ directly in the latent space $(c_t, z)$, avoiding out-of-distribution (OOD) action issues common in offline RL.

Training Objectives

The model is trained jointly on two primary objectives:

1. Temporal Consistency Loss ($L_{TC}$): A multi-step self-predictive loss. The model minimizes the cross-entropy between the predicted next latent state and the actual next latent state (derived via an EMA encoder), plus a mean-squared error for reward prediction.
   • Crucial Insight: This enforces that the task representation $z$ contains sufficient information to predict future dynamics and rewards, rather than just reconstructing raw observations.
2. Contrastive Loss ($L_{Contrastive}$): An InfoNCE loss applied to the task representations to ensure distinct tasks occupy different regions in the latent space, enhancing task discrimination.

The total context encoder loss is:
$$L_{Context} = L_{TC} + \beta L_{Contrastive}$$

3. Key Contributions

• C1: Temporal Consistency for Task Inference: The paper demonstrates that enforcing latent temporal consistency during context encoding yields task representations that capture task variation factors (dynamics and rewards) more effectively than reconstruction-based objectives.
• C2: Theoretical Analysis: The authors provide a formal value error bound for the latent MDP. They prove that accurate control is possible without reconstructing observations, provided the latent state and task representation preserve the necessary predictive information. The bound decomposes error into:
  1. Latent abstraction error (Markov property of the code).
  2. World model approximation error.
  3. Task inference error (quality of $z$).
• C3: Discrete Latent Space: The use of FSQ and classification loss (Cross-Entropy) for temporal consistency is shown to be superior to regression-based continuous latent spaces, leading to better handling of stochastic and multimodal dynamics.

4. Experimental Results

The method was evaluated on MuJoCo, Contextual DeepMind Control (DMC), and Meta-World benchmarks, comparing against state-of-the-art baselines like FOCAL, CSRO, DORA, and UNICORN.

• Generalization Performance: SPC significantly outperforms baselines in both few-shot (adaptation with a few new transitions) and zero-shot (no new transitions) settings.
  • On MuJoCo and Contextual-DMC, SPC achieves higher normalized returns and lower optimality gaps.
  • On Meta-World (ML10/ML45), SPC achieves higher success rates in generalizing to unseen environments.
• Representation Quality:
  • Disentanglement: SPC learns task representations with higher disentanglement (DCI metrics) and better alignment with true variation factors compared to reconstruction-based methods (UNICORN-SUP).
  • Robustness: SPC maintains a high matrix rank and low "dormant neuron" ratio, indicating it avoids representation collapse and learns more diverse features.
• Ablation Studies:
  • Latent Space: Discretizing the latent space via classification (Cross-Entropy) is the primary driver of performance gains, outperforming continuous regression or simple bounding.
  • Objectives: Combining temporal consistency with contrastive learning is essential; using only one leads to suboptimal performance (e.g., temporal consistency alone fails to distinguish tasks with similar dynamics but different rewards).
  • Comparison to DreamerV3: SPC significantly outperforms DreamerV3 in the offline setting, where DreamerV3 struggles to generalize without online interaction to correct model inaccuracies.

5. Significance

This paper bridges the gap between latent world models (typically used for model-based planning) and context-based offline meta-RL.

• Paradigm Shift: It challenges the necessity of observation reconstruction for learning effective representations in RL, showing that predictive consistency in a latent space is sufficient and more robust.
• Efficiency: By jointly training the context encoder and world model, the method learns task representations that are inherently optimized for control, leading to faster adaptation and better generalization to out-of-distribution tasks.
• Practical Impact: The approach offers a viable solution for real-world robotics and control applications where online interaction is expensive or impossible, enabling agents to learn from static datasets and adapt to new tasks with minimal data.