Dreaming Of Others: Latent Teammate Modeling In World Models For Multi-Agent Reinforcement Learning

Technical Summary: Dreaming of Others – Latent Teammate Modeling in World Models for Multi-Agent Reinforcement Learning

Problem Statement
In cooperative Multi-Agent Reinforcement Learning (MARL), agents must coordinate with partners whose internal policies and intentions are not directly observable. While world models, particularly the Dreamer family, have demonstrated strong generalization and sample efficiency in single-agent settings by learning compact latent dynamics, their application to MARL is currently limited. The core challenge is that teammate behavior introduces non-stationary, partially observable structure that is typically compressed into the model as undifferentiated noise. This prevents agents from effectively anticipating partner variability or adapting to unseen collaborators without access to their internal observations or policies. Furthermore, standard cooperative MARL approaches often rely on simplifications such as centralized training with privileged information or explicit communication channels, which are unavailable in realistic human-AI collaboration scenarios.

Methodology
The paper proposes a "Teammate-Conditioned World Model" that reframes teammates not as exogenous noise, but as structured, learnable components within the agent's world model. The methodology is built upon a Dreamer-style Recurrent State-Space Model (RSSM) with the following architectural and objective innovations:

• Factorized Latent State: The stochastic latent state $z_t$ is factorized into two distinct components: $z_t = [z^{env}_t, z^{team}_t]$.
  • $z^{env}_t$ captures the deterministic and stochastic dynamics of the physical environment.
  • $z^{team}_t$ encodes the inferred behavior, character, and intent of the teammate.
• Architecture:
  • Encoder: Consumes the controlled agent's observation $x_t$ and action $a^0_t$ to produce a deterministic hidden state $h_t$.
  • Decoders: Two decoders operate on the factorized latents. An observation decoder reconstructs the environment state $\hat{x}_t$ from $z^{env}_t$. A dedicated "teammate-policy decoder" maps $z^{team}_t$ to a predictive distribution $\hat{\pi}^j_t(\cdot)$ over the teammate's next action.
  • Actor-Critic Conditioning: During the imagination phase, the hidden state and the sampled teammate latents condition the actor and critic heads, allowing the agent to plan policies that account for diverse partner behaviors.
• Theory-of-Mind (ToM) Objective: A dedicated loss function, $\mathcal{L}_{ToM}$, is introduced to supervise the inference of teammate behavior. This objective minimizes a calibrated cross-entropy between the empirical teammate action distribution and the model's prediction, augmented by a Kullback-Leibler (KL) regularization term to ensure temporal consistency of the inferred teammate latent:
  $$\mathcal{L}_{ToM} = \mathbb{E}_t \left[ -\sum_a \pi^j_t(a) \log \hat{\pi}^j_t(a) \right] + \alpha \text{KL}\left( q(z^{team}_t | h_t) \parallel p(z^{team}_t | h_{t-1}, a_{t-1}) \right)$$
  This is added to the standard Dreamer training objectives with a weighting hyperparameter $\lambda_{ToM}$.

Key Contributions

1. Conceptual Framework: The paper proposes a new perspective for MARL where world models serve as simulators of social behavior, not just environmental dynamics. It advocates for modeling teammates as structured latent processes to reduce apparent non-stationarity.
2. Novel Architecture: It introduces a specific RSSM architecture that explicitly separates environmental and teammate latents, coupled with a ToM head to infer partner intent from partial trajectories.
3. Training Protocol: The work outlines a training procedure that enables zero-shot and few-shot coordination without requiring shared imagination, centralized policies, or explicit communication channels between agents.
4. Evaluation Proposal: The authors propose a rigorous evaluation protocol utilizing Multi-Agent Particle Environments, Overcooked-AI, and Melting Pot to assess generalization, sample efficiency, and robustness across diverse partner populations and social contexts.

Results
The paper explicitly states that it is a proposal intended to catalyze discussion and does not report empirical results. Consequently, there are no experimental findings, performance metrics, or comparative benchmarks presented in this work. The contribution is purely architectural and theoretical, outlining how such a system would function and how it should be evaluated.

Significance and Claims
The paper positions this work as a step toward "human-compatible AI" where partners are partially observable and heterogeneous. The authors claim that by treating teammates as structured latent processes rather than noise, the approach offers a path to:

• Reduced Non-Stationarity: Separating social dynamics from physical dynamics may stabilize learning in cooperative settings.
• Improved Coordination: Enabling agents to imagine and adapt to unseen collaborators (zero-shot) and improve rapidly with limited interaction (few-shot).
• Interpretability: Providing latent representations that explicitly encode partner character and intent.

The ultimate significance, as framed by the authors, is the potential for world models to evolve from tools that "dream of the worlds they inhabit" to tools that "dream of the minds that share them," facilitating more robust and adaptable human-AI collaboration.