1 Introduction↩︎

World models have emerged as a powerful paradigm for reinforcement learning (RL), enabling agents to learn compact latent dynamics and train decision policies from imagined trajectories [1]. Despite their success in single-agent domains, deploying world models in cooperative multi-agent RL (MARL) remains challenging because the behavior of teammates introduces non-stationary, partially observable structure that is typically compressed as undifferentiated noise.

We advocate extending world models from simulating environments to simulating others within them. Concretely, we propose a teammate-conditioned world model that (i) factorizes the latent state into environment and teammate components and (ii) augments the model with a Theory-of-Mind (ToM) head that infers latent embeddings of partner behavior from partial histories. The resulting teammate latents condition the actor and critic during imagination, allowing agents to anticipate partner variability and adapt to unseen collaborators without access to their internal observations or policies.

This focus on collaboration with humans removes many of the simplifications common in multi-agent reinforcement learning. Agents cannot rely on centralized training with privileged information, explicit communication channels, or access to a partner’s policy, goals, or observations. The human partner instead appears as a partially observed, adaptive process whose behavior must be inferred and integrated into the agent’s predictive model.

This paper is a proposal intended to catalyze discussion in the world-models community; we do not report empirical results. We precisely describe the architecture, objectives, and evaluation protocol, argue why modeling teammates as structured latent processes can reduce apparent non-stationarity, and outline how to assess zero-shot and few-shot coordination. The broader goal is human AI collaboration, where partners are partially observable and heterogeneous [2], [3].

2 Background and Related Work↩︎

2.0.0.1 World Models and Dreamer.

World models learn a generative model of environment dynamics to support prediction, planning, and control in latent space. By compressing high-dimensional observations into a recurrent latent representation, these models allow agents to reason about abstract dynamics and train policies through imagination rather than direct interaction. The Dreamer family [1] exemplifies this paradigm: it combines a Recurrent State-Space Model (RSSM) that captures both deterministic memory and stochastic latent variables with an actor-critic trained entirely through latent imagination rollouts. The approach has demonstrated remarkable generalization and sample efficiency, mastering hundreds of tasks across diverse continuous-control domains with fixed hyperparameters. DreamerV3 in particular shows that large world models can scale and generalize to unseen tasks without task-specific tuning, suggesting that latent imagination provides a powerful inductive bias for learning structured, reusable dynamics.

2.0.0.2 Cooperative Multi-Agent Reinforcement Learning.

Cooperative MARL studies how multiple agents can learn to coordinate toward a shared reward function under partial observability. The most common training paradigm, Centralized Training with Decentralized Execution (CTDE), provides agents with access to joint information during training while requiring independent execution policies at test time. Policy-gradient algorithms such as MAPPO [4] and value-decomposition methods such as QMIX [5] have become standard baselines, offering empirical robustness and scalability across benchmark tasks. In these settings, non-stationarity remains a central challenge: as teammates update their policies, the effective environment observed by each agent changes.

2.0.0.3 World Models for MARL.

Recent works have begun to extend world models to multi-agent coordination, highlighting the promise and limitations of imagination-based approaches in social environments. MA-Dreamer [6] introduces shared imagination between agents, enabling communication through latent rollouts. CoDreamer [7] decentralizes this process, using latent message passing to synchronize multiple local world models. GAWM [8] proposes a global-aware latent representation to capture collective state information across agents, while MACD explores counterfactual imagination to improve credit assignment [9], [10]. Although these studies advance model-based cooperation, they often assume shared access to latent states or centralized imagination, effectively sidestepping the challenge of partial observability.

2.0.0.4 Theory of Mind and Agent Modeling.

Parallel research in agent modeling and Theory of Mind (ToM) explores how artificial agents can infer the beliefs, goals, or strategies of others. ToMnet [11] introduced the idea of learning to predict another agent’s future actions from past trajectories, laying the foundation for meta-learning and inference-based models of social reasoning. Subsequent work formalized intent inference under partial observability [12], demonstrating improved robustness and interpretability in cooperative and competitive settings. In human-AI coordination, explicitly modeling partner preferences and intentions improves adaptability and alignment [2], [3].

3 Teammate-Conditioned World Models↩︎

3.1 Design Rationale↩︎

In cooperative settings, observations reflect both physical dynamics and the evolving behavior of teammates. Treating other agents as exogenous noise induces apparent non-stationarity as partners change policies. We posit that modeling teammates as structured latent processes, separate from environmental latents, can reduce non-stationarity, produce socially consistent imagined rollouts, and enable rapid adaptation to new partners.

3.2 Architecture↩︎

At each time step \(t\), the encoder consumes the controlled agent’s observation \(x_t\) and action \(a_t^0\) to produce a deterministic hidden state \(h_t\). We posit a factorized stochastic latent \[z_t = \big[z_t^{\text{env}},\; z_t^{\text{team}}\big],\] where \(z_t^{\text{env}}\) captures environment dynamics and \(z_t^{\text{team}}\) encodes inferred teammate behavior.

Two decoders operate on these latents. An observation decoder maps \(z_t^{\text{env}}\) to \(\hat{x}_t\). A teammate-policy decoder maps \(z_t^{\text{team}}\) to a predictive distribution \(\hat{\pi}_t^j(\cdot)\) over the teammate’s next action. Actions from all agents \((a_t^0, a_t^j)\) drive the transition to \(h_{t+1}\). During actor-critic learning, the hidden state and teammate latents condition policy and value heads to generate \((a_t^0, v_t, r_t)\) and support imagination with sampled \(z_t^{\text{team}}\).

We retain standard Dreamer training for the world model and control components [1] and add a dedicated ToM objective for teammate modeling, without requiring shared imagination, centralized policies, or explicit communication [6]–[8].

3.3 Teammate-Modeling Objective↩︎

Let \(\pi_t^j(\cdot)\) denote the empirical or behavior distribution of the teammate’s next action at time \(t\) obtained from data, and \(\hat{\pi}_t^j(\cdot)\) the model’s prediction from \(z_t^{\text{team}}\). We minimize a calibrated cross-entropy with temporal regularization: \[\label{eq:tomloss} \mathcal{L}_{\text{ToM}} = \mathbb{E}_{t}\Big[ -\sum_{a}\pi_t^j(a)\,\log \hat{\pi}_t^j(a) \Big] + \alpha\,\mathrm{KL}\!\left( q\!\left(z_t^{\text{team}}\mid h_t\right)\,\middle\|\,p\!\left(z_t^{\text{team}}\mid h_{t-1},a_{t-1}\right) \right),\tag{1}\] where the KL term stabilizes temporal consistency of the inferred teammate latent. In practice, \(\pi_t^j\) can be a one hot target from the observed \(a_t^j\) or a smoothed label. The total objective augments standard Dreamer training with \(\lambda_{\text{ToM}}\mathcal{L}_{\text{ToM}}\).

3.4 Deployment and Adaptation↩︎

At test time, the model infers \(z_t^{\text{team}}\) online from observed teammate actions and conditions the actor and critic on this embedding. Imagined rollouts sample plausible teammate trajectories, enabling zero-shot coordination with unseen partners and few-shot improvement as more behavior is observed, without access to the partner’s policy or observations.

4 Evaluation Protocol↩︎

We propose an evaluation protocol that emphasizes generalization and adaptability across partners and settings. Multi-Agent Particle Environments [13] provide lightweight, controllable scenarios for diagnosing whether the model learns identifiable teammate latents and separates social from physical dynamics. Overcooked-AI [2] serves as the main testbed for zero-shot coordination, using standard partner splits to assess how well the inferred teammate embedding supports immediate coordination with unseen collaborators. Melting Pot [14] evaluates robustness across diverse social contexts and partner populations at scale, probing how the method copes with distributional shift in norms and conventions. Across all environments, we plan to leverage unified tooling from JaxMARL [15] and BenchMARL [16], which make significant strides toward standardizing environments, metrics, and baselines. We will report episodic return and sample efficiency, a zero-shot coordination score with unseen partners, few-shot improvement after a small number of interactions and cross-play robustness across partner pairings.

5 Discussion and Outlook↩︎

We present a conceptual architecture that integrates teammate modeling into world models for MARL. By treating other agents as structured latent processes rather than stochastic noise, the proposed design provides a path toward socially grounded world models. This formulation offers several potential benefits: reduced non-stationarity, improved zero-shot coordination, interpretable latent representations, and greater compatibility with human partners.

Ultimately, teammate-conditioned world models may help agents not only dream of the worlds they inhabit, but also of the minds that share them.

References↩︎