Mental-Models for Multi-Agent Systems

AI summary

Overview

Research area: Multi-agent systems, Theory of Mind (ToM), and foundation-model agent training under partial observability.

Technical level: Advanced. The paper combines latent-variable modeling (VAE-style Gaussian posteriors), recursive belief factorization, reward modeling, and policy optimization (GRPO), so it assumes familiarity with reinforcement learning and generative modeling.

Scope: The paper proposes a framework that learns an explicit, amortized latent mental model of a partner agent and turns it into a training signal for a deployable policy, evaluated on language-only and multimodal multi-agent benchmarks. Note that the provided content is truncated mid-conclusion, and appendix material referenced by the paper (e.g., Appendix A.1.3, Table 10) is not included, so those details are not reported here.

What This Paper Is About

Agents built on large foundation models can converse well, but they typically do not maintain an explicit representation of what a partner knows, wants, or intends, which leads to misread intent and brittle plans when interactions unfold unexpectedly. Existing fixes — prompt engineering, memory, self-critique, instruction tuning, preference RL, synthetic self-play — shape behavior without ever embedding a partner's state into the policy's latent dynamics. This paper asks whether an agent can learn an amortized latent belief state from interaction history and use it as a stable, reusable training signal for multi-agent policy learning.

Key Contributions

1. A general framework for ToM-guided multi-agent decision making under partial observability, in which an agent learns an explicit latent belief state over other agents from observed history and uses it to support downstream action selection.
2. A recursive mental-state architecture plus a coupled mental-state/reward learning objective, factorizing the latent into first- and second-order states, each split into belief, intent, and thought (BIT) subspaces, so that partner representations become directly useful for scoring candidate actions.
3. A distillation procedure that freezes the mental-state encoder and reward model and trains a standalone policy under the resulting belief-conditioned reward, so explicit partner modeling improves training while inference-time deployment stays efficient.
4. Broad evaluation plus enriched resources: language-only benchmarks (SOTOPIA, Craigslist-bargain, ToMi, BigToM) and a multimodal benchmark (MMRole), with the existing resources augmented with ToM supervision, utility signals, rationales, and hard negatives.

Main Findings

• SOTOPIA overall gains across three model families. Under a fixed GPT-4o-mini partner on Sotopia-all, overall scores improve for every open-source backbone: LLaMA 3.330 → 3.541, Mistral 3.511 → 3.784, and Qwen 3.290 → 3.856. The paper reports that Qwen2.5-7B achieves an average relative gain of 11.5% over the base model across two partner settings.
• Gains concentrate on task-facing dimensions. Improvements are pronounced on Knowledge and Goal while Believability stays high and stable, which the authors read as better goal-directed reasoning without a loss of plausibility.
• Harder scenarios still benefit. On Sotopia-hard, average scores improve for all three LLMs: LLaMA 5.880 → 6.166, Mistral 6.714 → 6.738, Qwen 6.381 → 6.785.
• Partner-agnostic improvement. Averaging over the GPT-4o-mini and self-partner settings, LLaMA rises 6.588 → 7.245, Mistral 7.109 → 7.552, and Qwen 7.087 → 7.633, suggesting the method is not tied to a particular counterpart.
• Out-of-domain negotiation transfer. Evaluated on a random subset of 100 Craigslist Bargain negotiation examples with no task-specific fine-tuning, LLaMA2 improves deal rate 93.0% → 94.0% and cuts walkaways 7.0% → 1.0%; Qwen2.5 improves deal rate 91.0% → 94.0%. GPT-4o-mini is reported at 96.0% deals, 3.0% walkaways, 1.0% no agreement.
• BigToM paired-branch consistency. The paper reports that its approach improves paired TB/FB accuracy by 26.7 pp on average across the six reported settings. Base models handle the easier true-belief branch but lose consistency on the false-belief branch or when the initial belief changes. On Backward Belief without an initial belief sentence, the base model scores TB 85.0 / FB 92.5 / TB∧FB 77.5 versus 97.0 / 98.5 / 95.5 for the method.
• ToMi generalization from BigToM-trained policies. Without task-specific fitting, average rises 87.92 → 89.16 and joint score 55.16 → 60.46 over the base model; the authors note classic task-specific ToMi baselines remain strong on that synthetic state-tracking format.
• Multimodal role-play gains. On MMRole, overall score improves from 0.994 (MMRole-Agent) to 1.023, with the Out split moving 0.981 → 1.019 and the In split 0.999 → 1.024. The largest gains appear on Coherence, Response Accuracy, Personality Consistency, Knowledge Consistency, and Tone Consistency.
• The latent is not merely a compression code. Scrambling a flat compression VAE latent has only a weak effect on reward-vector MSE, whereas scrambling the mental-model latent sharply disrupts joint reward prediction — evidence that the mental latent is coupled to decision-relevant structure rather than reconstructive summary.
• BIT slices are complementary, not redundant. Pairwise CKA shows belief, intent, and thought subspaces are separated yet related, while a flat mental-summary variant shows significant inter-slice similarity. Linear probes recover ToM-relevant MMRole factors more accurately from BIT than from a flat latent of the same combined dimensionality.
• The recursive z₁ → z₂ link carries nested ToM structure. Shuffling z₂ while preserving its marginal distribution loses substantially more recoverable second-order mental-state information than reward information.
• Partial supervision is enough. Training Qwen2.5-7B with nested subsets of SOTOPIA annotations gives overall scores of 3.300 (0%, mental-free reward baseline), 3.650 (25%), 3.750 (50%), and 3.856 (100%) under the same evaluation seed, partner, judge, and scoring protocol.
• A richer posterior family buys nothing measurable. On a held-out scenario-level BigToM split (n = 5,976), a full low-rank covariance posterior (diag + UUᵀ) gives reward accuracy 96.8 versus 96.6 for the diagonal Gaussian, identical margin 1.05, and belief probe 99.9 versus 99.8. The 0.2-point difference sits within a reported ±0.65-point sampling uncertainty.
• Latent directions correspond to mental and strategic variables. Traversing a probed belief direction monotonically shifts both a branch classifier and a z₁-only reward head that was never trained on the branch label, and traversing an offer/proposal strategy direction raises the strategy probe score and shifts nearest-neighbor labels from questioning/other toward offer/proposal, with more squared-weight mass in second-order ToM blocks.

Methodology in Plain English

The framework treats interaction under partial observability as a three-stage pipeline.

Stage 1 — learn the mental model and what it is good for. The agent sees a history of observations and its own past actions, and cannot observe the partner's goals, beliefs, knowledge, capabilities, or strategy. Exact Bayesian inference over those hidden variables is intractable in open-ended interaction, so the method learns an amortized approximation: a network that maps history directly to a compact latent belief state. Crucially, this latent is recursive rather than flat. A first-order state represents the partner's mental state; a second-order state represents the partner's model of the agent or environment and is conditioned on the first-order inference. Each level is split into three semantic slices — belief, intent, and thought process — modeled as multivariate Gaussian distributions parameterized by an encoder, with a standard normal prior on the first-order latent and a conditional Gaussian prior on the higher-order latent. Uncertainty is carried in the posterior variance rather than collapsed to a point estimate.

A latent alone cannot tell the agent what to do, so a reward/utility head is learned alongside it that scores candidate actions conditional on the inferred mental state, averaging over the mental states still plausible given the history. The two are trained jointly with a combined objective: a candidate-reward term (including hard-negative ranking when contrastive responses exist), a mental-state inference term that anchors the latent to first- and second-order mental-state evidence, and auxiliary terms such as rationale-conditioned utility prediction and future-action prediction. The authors motivate the joint training by two failure modes of training separately — a mental model that is interpretable but decision-irrelevant, and a reward model that scores actions from surface features while ignoring partner state.

Stage 2 — distill the signal into a policy. The mental-state encoder and reward model are frozen. The policy proposes candidate actions from the observable history; the frozen scorer evaluates them using the mental-model-conditioned utility, scalarized by evaluation-specific weights; and the policy is updated to raise the likelihood of high-scoring actions relative to other plausible ones in the same context. This is instantiated as mental-reward-guided GRPO, with the frozen reward used to form group-relative advantages, plus a KL term that keeps the update near the reference policy. Because the reward was computed by marginalizing over the latent, the gradient carries information about hidden partner state, and the policy absorbs regularities linking observable cues to partner beliefs, intents, and goals.

Stage 3 — deploy the policy alone. At test time only the trained policy is used, conditioned on the history-derived prompt; no mental model or reward model is needed at inference.

Experimental setup. All trained open-source agents use LoRA adapters on 7B-scale backbones: Qwen2.5-7B-Instruct as the primary language backbone, with LLaMA2-7B-Chat and Mistral-7B-Instruct-v0.3 as additional SOTOPIA comparisons, and Qwen2.5-VL-Chat for MMRole. SOTOPIA-π episodes are augmented with ToM labels and utility annotations to train the mental-state/reward model; BigToM and ToMi training uses 5K synthetic BigToM-style samples with evaluation on the official BigToM and ToMi splits; MMRole's released training set is augmented with mental labels. Craigslist-bargain and FANToM are used only as transfer evaluations. Baselines include zero-shot/base agents, SFT and GRPO variants trained on the generated data, closed-source GPT-4o-mini/Gemini, GPT-4 and GPT-4 with chain-of-thought on BigToM, classic ToMi models (MemNN, RelNet, EntNet), and MMRole's official VLM and MMRole-Agent baselines.

Why This Matters

Impact on research. The paper argues that partner modeling in LLM-based multi-agent systems has largely been implicit — encoded in prompt text or trajectory summaries rather than in the policy's latent dynamics — which limits interpretability, complicates reuse across tasks, and prevents partner modeling from being integrated systematically into different policy learners. By making the belief state an explicit, reusable decision variable with quantified uncertainty, and by showing that it can be distilled away at deployment, the work offers a concrete bridge between the sequential-decision-making literature on beliefs over hidden state and modern foundation-model agents. The ablations also give the community diagnostic tools: latent scrambling, CKA on subspaces, pair shuffling, and functional latent traversals, rather than accuracy alone.

Real-world applications

• Negotiation and bargaining assistants that must infer a counterpart's constraints and reservation points from limited offers — the paper reports transfer to negotiation data without task-specific fine-tuning.
• Customer-facing dialogue and support agents where inferring a user's underlying goal and knowledge state matters more than fluent phrasing.
• Assistive and collaborative robotics, where a robot must model a human partner's beliefs and intentions under partial observability (the paper's framing draws on mental models in robotics and cognitive science).
• Multimodal role-playing and interactive characters, evaluated here via MMRole on persona-grounded images and dialogue, where coherence, knowledge consistency, and personality consistency require inferring who a character is and what they know.

Industry relevance. The method is designed to keep inference cheap: mental modeling happens during training, and only the policy is deployed, which suits practitioners who want behavioral improvement without serving an extra encoder and reward model. The paper reports that limited annotation (25% of SOTOPIA annotations) already yields most of the benefit (3.650 vs. 3.856 overall), which matters for teams that cannot afford exhaustive mental-state labeling. Reusable 7B-scale LoRA adapters also fit within practical serving budgets, and the augmentation recipes (utility signals, rationales, hard negatives) are resources other groups can build on. Code is publicly available at https://github.com/hananshafi/Mental-Models.

Future Directions

• Scaling the supervision question. The paper shows performance rises with annotation coverage (25% → 3.650, 50% → 3.750, 100% → 3.856), but it does not report whether returns continue past full SOTOPIA annotations or how far proxy evidence can substitute for true mental-state labels at larger scale.
• Beyond two agents and two orders. The formulation is presented for the two-agent case with first- and second-order latents; how the recursive factorization extends to many interacting agents and deeper orders is not reported.
• Reconciling generality with task-specific strength. On ToMi, classic task-specific baselines remain strong against the transferred general policy, raising the open question of how to keep broad multi-agent generality while matching specialized state-tracking models.
• Richer posterior and architecture choices. The posterior-family ablation found no measurable benefit from a full low-rank covariance posterior at n = 5,976, leaving open whether other structural or architectural changes could yield gains, and what the limits of the diagonal residual-noise assumption are.

Target Audience

Researchers and graduate students working on multi-agent systems, Theory of Mind in AI, and LLM agent training; reinforcement-learning practitioners interested in belief-state-based policy learning and reward modeling; and applied engineers building interactive agents for negotiation, dialogue, robotics, or multimodal role-play who want a training-time recipe that leaves inference-time cost unchanged. Readers without a background in latent-variable models, variational inference, and policy-gradient methods will find the methodology section dense and should treat the analysis and ablation sections as the accessible entry point.

Authors’ abstract

Large foundation models have accelerated progress toward general-purpose agents that interact with humans and other agents through language and multimodal signals. However, robust multi-agent decision-making requires reasoning about what other agents know, intend, and are likely to do under partial observability. Current agentic systems often operate through prompt design, memory, or end-to-end behavioral shaping, but typically do not learn an explicit partner-state representation that can be reused as a decision variable across tasks. We introduce \emph{mental-model-enabled agents}, a framework that equips an agent with a latent mental model of its counterpart, allowing it to infer hidden beliefs, intentions, and likely reactions from the observed history and use these inferences to guide action selection. Our method learns an amortized recursive Theory-of-Mind representation, with first- and second-order mental-state structure, jointly with a belief-conditioned reward model that evaluates candidate actions relative to the inferred partner state. A policy is then learned under this belief-aware signal, yielding an agent that can act independently at inference time while retaining the benefits of explicit partner modeling. We evaluate the same framework on both language-only and multimodal benchmarks. Across these settings, explicit mental-state modeling consistently improves interaction quality and Theory-of-Mind performance over base agentic systems, showing that structured partner modeling is a useful inductive bias for general multi-agent systems. Our code is publicly available at https://github.com/hananshafi/Mental-Models