Technical Summary: In-Context Reinforcement Learning via Communicative World Models (CORAL)

1. Problem Statement

Reinforcement Learning (RL) agents frequently struggle to generalize to new tasks and contexts without updating their parameters. This limitation stems from learned representations and policies that are overfit to the specificities of training environments. While In-Context Reinforcement Learning (ICRL) offers a pathway to context-driven adaptation by conditioning policies on past interactions, existing approaches often rely on self-supervised transformers that lack an understanding of underlying task dynamics. Consequently, their generalization is fragile and heavily dependent on the quality of offline datasets.

Conversely, World Models (WMs) provide agents with structured understanding of environment dynamics and reward feedback, enabling prediction of future states. However, in standard WMs, the learned representations are often entangled with task-specific policy learning. This entanglement leads to representations that are overly specialized, less transferable, and suffers from objective mismatch between representation learning and policy optimization.

The core problem addressed is how to leverage the dynamic understanding of World Models for ICRL without the entanglement of policy learning, thereby creating a transferable, communicative context that accelerates adaptation in new environments.

2. Methodology: The CORAL Framework

The authors propose CORAL (Communicative Representation for Adaptive RL), a framework that formulates ICRL as a two-agent emergent communication problem. CORAL functionally separates latent representation learning from control, decoupling the roles of world modeling and reward maximization.

2.1 Agent Architecture

The system consists of two distinct agents interacting at each timestep $t$:

• Information Agent (IA): Parameterized by $\phi$ (a Transformer-based architecture), the IA acts as a world model. It processes a context of recent embedded observations and aims to understand environment dynamics and task rewards. Crucially, the IA is not trained on direct return maximization. Instead, it compresses its understanding into concise latent messages ($m_t$).
• Control Agent (CA): Parameterized by $\theta$ (typically a PPO policy), the CA receives the observation $o_t$ and the message $m_t$ from the IA. Its sole objective is to maximize task rewards by selecting actions $a_t$ based on the communicative context.

2.2 Training Objectives

The agents are jointly pre-trained on a diverse distribution of tasks ($\mathcal{T}$) using a composite loss function for the IA, while the CA is optimized via Proximal Policy Optimization (PPO).

The Information Agent's Loss ($\mathcal{L}(\phi)$):
The IA is trained via a self-supervised composite loss consisting of three components:

1. Dynamics Awareness Loss ($\mathcal{L}_{Dyn}$): Anchors the IA's representations by requiring it to predict the next observation ($\hat{o}_{t+1}$), reward ($\hat{r}_{t+1}$), and termination probability ($\hat{d}_{t+1}$) based on the current message and the CA's action. This ensures the internal state captures physical truths about the environment.
2. Temporal Coherence Loss ($\mathcal{L}_{Coh}$): Encourages consistency in communication by training the IA to predict the next message ($\hat{m}_{t+1}$) from the current message ($m_t$) and action. This promotes messages as compact representations of the state.
3. Causal Influence Loss ($\mathcal{L}_{Causal}$): A novel objective that shapes the communication protocol. It measures the Instantaneous Causal Effect (ICE), defined as the Reverse KL Divergence between the CA's policy with the message versus a zero-vector message. This loss is weighted by a hybrid utility signal (combining Generalized Advantage Estimates and immediate value changes) to incentivize the IA to send messages that cause beneficial, decisive shifts in the CA's policy.

The Control Agent's Objective:
The CA maximizes the expected discounted return $E[\sum \gamma^t r_t]$ using standard PPO, conditioned on the observation and the IA's message.

2.3 Training Regime

To ensure generalization, the pre-training regime utilizes multi-environment training. The agents are trained on a diverse, discrete set of tasks sharing common entities and dynamics but requiring different solutions. This is implemented using vectorized training (via JAX) to run parallel environments, preventing catastrophic forgetting and forcing the IA to learn abstract dynamics rather than memorizing specific instances.

2.4 Deployment

During deployment, the IA's weights are frozen. It serves as a fixed "contextualizer" or communicative prior for a new, randomly initialized CA. The CA learns to solve unseen tasks by interpreting the messages provided by the frozen IA, enabling rapid in-context adaptation without updating the world model.

3. Key Contributions

1. Formulation: The paper models single-agent ICRL as a two-agent emergent communication problem where a World Model (IA) learns to communicate dynamics to a Control Agent (CA).
2. Hybrid Training Strategy: The framework decouples representation and control learning. The IA adopts self-supervised training with three specific heads (dynamics, coherence, influence), while the CA uses reinforcement pre-training. This separation avoids the objective mismatch found in monolithic WMs.
3. Causal Influence Loss: The introduction of a novel loss term that explicitly incentivizes the IA to generate messages that causally influence the CA's policy toward high-utility actions, ensuring the communication is pragmatic and effective.
4. Empirical Validation: Extensive experiments demonstrate that CORAL improves sample efficiency in online sparse-reward environments and achieves superior zero-shot generalization compared to standard PPO, World Models, and state-of-the-art In-Context RL methods (e.g., Decision Transformer, Agentic Transformer).

4. Experimental Results

The authors evaluated CORAL across three regimes:

• Online Adaptation (MiniGrid Suite):

  • CORAL demonstrated significantly higher sample efficiency than baselines (PPO, PPO-Transformer, World Model).
  • In tasks like DoorKey 8x8 and Crossings S11N5, CORAL achieved near-optimal rewards roughly twice as fast as PPO.
  • In harder, larger variants (e.g., DynObs 16x16), CORAL successfully bootstrapped the CA to solve tasks where baselines failed or performed near random.
  • Time-to-Threshold (TTT): CORAL reached 90% of maximum performance significantly faster (e.g., 2.3x faster than PPO in DoorKey 8x8).

• Zero-Shot Generalization:

  • When evaluated on unseen, more complex environments without further training, CORAL agents pre-trained on simpler tasks outperformed PPO, World Models, and Decision Transformers.
  • This indicates the communicative prior generalizes better to dynamic shifts than purely sequence-based cloning of demonstrations.

• Offline Generalization (D4RL):

  • In fully offline settings, CORAL achieved a total average normalized score of 85.47, outperforming Decision Transformer (77.49) and matching TD3+BC (84.83).
  • Notably, CORAL excelled on difficult Medium-Replay datasets (e.g., 98.33 on Hopper-M-Replay), suggesting its communicative prior effectively filters noise from sub-optimal demonstrations.

• Ablation Studies:

  • Removing Temporal Coherence Loss led to high variance and slower learning.
  • Removing Causal Influence Loss reduced sample efficiency and final performance, confirming that a passive world model is insufficient; the IA must be incentivized to influence the CA's policy.
  • Analysis of the Instantaneous Causal Effect (ICE) showed that the protocol provides strong influence during learning (high uncertainty) and stabilizes as the policy converges, acting as a learning regularizer.

5. Significance and Claims

The paper claims that CORAL provides a principled middle ground between context-driven adaptation and world modeling. By separating the learning of dynamics (via the IA) from control (via the CA), the framework learns a transferable communicative representation that is not overfit to specific task rewards.

The authors assert that this approach validates the efficacy of learning a transferable communicative prior. The results suggest that emergent communication is not only useful for multi-agent coordination but serves as a powerful general method for rapid in-context adaptation in single-agent RL. The framework enables agents to achieve significant gains in sample efficiency and successful zero-shot adaptation in diverse online and offline environments, addressing the fragility of current ICRL approaches and the entanglement issues of standard World Models.