Mental-Models for Multi-Agent Systems

Hanan Gani; Lulu Shao; Manmohan Chandraker

cs.MAOct 8, 2026

Mental-Models for Multi-Agent Systems

Authors: Hanan Gani, Lulu Shao, Manmohan Chandraker

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

Figures & tables

Figure 1: Mental-model-enabled decision making under partial observability. Observed behavior only partially reveals another agent’s hidden goals, beliefs, constraints, and likely strategy. By inferring these latent mental states, a mental-model-enabled agent can choose more adaptive actions, improving outcomes while preserving cooperation across multi-agent settings.

Figure 2: Method overview. Phase 1 jointly learns an amortized recursive mental-model and a mental-conditioned reward model from multi-agent context. The latent contains first- and second-order belief, intent, and thought subspaces, which condition reward prediction. Phase 2 freezes the mental-state encoder and reward model, scores sampled policy actions with the learned mental reward, and optimizes the policy with mental-reward guidance. Phase 3 deploys only the trained policy in the multi-agent setting.

Table 4

Table 5

Model	Average	Joint	w/o FB 1st	w/o FB 2nd	w/o FB Reality	w/o FB Memory	FB 1st	FB 2nd
Classic TOMi baselines
MemNN	77.20	44.30	85.45	82.67	93.39	98.90	12.62	17.27
RelNet	86.00	57.40	96.42	95.37	100.0	99.90	10.40	17.81
EntNet	90.60	66.80	94.29	85.08	100.0	100.0	54.95	36.55
Closed source LLMs
GPT-3.5-Instruct	79.78	36.94	94.42	84.05	97.30	94.79	34.26	7.82

Table 7: ToMi Generalization from the BigToM -trained policy. Ours improves the aggregate Average and Joint scores over the base model and strengthens non-false-belief first-/second-order belief tracking, highlighting OOD generalization of our mental-model trained policy.

Table 7

Panel of Figure 4 — Figure 4: a) $z$ -scrambling barely affects the compression VAE but sharply increases BIT reward-vector MSE, showing stronger reward-causal structure. b) Pairwise CKA shows complementary belief, intent, and thought slices. c) Frozen-latent probes recover ToM-relevant MMRole factors more separably from BIT than from a flat latent. d) Shuffling $z_{2}$ weakens second-order probes, supporting the recursive $z_{1}\rightarrow z_{2}$ structure.

Panel of Figure 5 — Figure 5: Latent traversal analyses. a) Belief-axis traversal shifts both the branch classifier and frozen reward head beyond random controls. b) Offer/proposal traversal increases strategy signal and local label enrichment, with stronger second-order contribution.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Multi- modal	Episode Reward	Per-turn Reward	Intent Annot.	1st-order Belief	2nd-order Belief	Rationale / CoT	Hard Negatives
BigToM (original)	✗	✗	✗	✗	✓	✗	✗	✗
BigToM (ours)	✗	✗	✗	✓	✓	✓	✓	✓
Sotopia - $\pi$ (original)	✗	✓	✗	✗	✗	✗	✗	✗
Sotopia - $\pi$ (ours)	✗	✓	✓	✓	✓	✓	✓	✓
MMRole (original)	✓	✓	✗	✗	✗	✗	✗	✗
MMRole (ours)	✓	✓	✓	✓	✓	✓	✗	✓

Appendix

Table 10: Comparison of original and augmented datasets across annotation dimensions.

Subset	$n$	Mean ( $1$ – $5$ )	Faithfulness (%)
MMRole	$50$	$4.30$	$86.0$
Sotopia	$50$	$4.60$	$92.0$
Overall	$\mathbf{100}$	$\mathbf{4.45}$	$\mathbf{89.0}$

Appendix

Table 11: Blinded human audit of generated mental-state annotations. Faithfulness linearly rescales the mean rating to

0

–

100\%

Model	Belief choice	Belief distance	First order	Second order
Base	$25.4$	$35.2$	$26.8$	$49.7$
SFT	$26.5$	$33.8$	$29.4$	$41.4$
Ours	$\mathbf{27.9}$	$\mathbf{38.0}$	$\mathbf{29.7}$	$\mathbf{52.2}$

Appendix

Table 12: Zero-shot BigToM

\rightarrow

FANToM transfer on

5{,}000

probes using the official scorer. Scores are percentages; higher is better.

Figure 6: Mental-latent geometry from an unsupervised UMAP of latent $z$ . Left: coloring by task shows task/scenario macro-clusters. Right: coloring by branch shows aware and not-aware samples remain near paired scenarios but are displaced within task regions, consistent with branch F1 $=0.989$ .

		Without Initial Belief			With Initial Belief
Task	Model	TB	FB	TB $\land$ FB	TB	FB	TB $\land$ FB
Forward Belief	GPT-4 (0-shot)	99.0	98.0	97.0	91.0	99.0	90.0
	GPT-4 (0-shot-CoT)	99.0	99.0	98.0	99.0	99.0	98.0
	GPT-4 (1-shot)	99.0	99.0	97.0	97.0	99.0	96.0
	GPT-4 (1-shot-CoT)	100.0	99.0	99.0	97.0	99.0	96.0
	Base	83.0	94.0	77.0	49.5	98.5	48.0

Appendix

Table 13: BigToM official-protocol results. TB denotes true-belief accuracy, FB denotes false-belief accuracy, and TB

\land

FB denotes paired accuracy, where both the true-belief and false-belief versions of the same scenario must be answered correctly.

Figure 7: Range-normalized lift from Qwen2.5-7B (ours w/ 1st order) to Qwen2.5-7B (ours w/ 2nd order) on Sotopia . Petal lengths encode percentage-of-range changes for each official dimension rather than raw scores. The center reports the relative lift in the average score, highlighting that second-order ToM improves the aggregate profile while keeping the comparison scale-free.

Figure 8: Preference-margin distribution. BIT decomposition produces sharper preference margins than flat or shuffled mental supervision. Kernel density of the per-sample margin $r^{+}-r^{-}$ on 1,508 Sotopia validation turns. Dotted vertical lines mark medians; the dashed line marks zero. BIT ranks positives over negatives with roughly $2{\times}$ higher confidence (median $+25.06$ vs. flat $+12.76$ ), and every BIT sample is correctly ranked with margin $\geq+1.40$ . Flat overlaps the shuffled-label control, indicating that mental supervision improves calibration when belief, intent, and thought are routed to dedicated latent sub-blocks.

Figure 9: Cross-agent transfer of the mental-reward model on Sotopia . Bars compare the average scores of LLaMA2-7B (ours) with its same-architecture mental-reward model and LLaMA2-7B (ours w/ r-Qwen) with the Qwen mental-reward model. The transferred Qwen mental-reward model preserves overall performance, supporting its generalization across agents.

Figure 10: Role-swap latent intervention. Swapping a BIT role-specific block moves the corresponding role prediction more than unrelated roles. Flat and Shuffled variants do not show this separation, yielding more diffuse effects.

Figure 11: t-SNE visualization The concatenated latent representation [z1; z2] reaches 0.662 5-NN accuracy (silhouette +0.046) on the 5-way emotion task compared to 0.529 for Qwen without the adapter, and a chance baseline of 0.200. The +0.13 gap over the no-adapter baseline indicates that the learned mental latents encode emotion-relevant structure that is not present in the base Qwen representation. The corresponding t-SNE projection provides a qualitative visualization: classes appear more separated in the z1 and [z1; z2] projections, while the no-adapter Qwen projection shows more inter-class overlap.

Figure 12: Visual Grounding. Visual grounding rate per response for Base vs. Ours on 200 held-out MMRole examples. Each point is one example’s share of generated sentences that mention a gold scene object; black markers show means. The metric is computed using a token-overlap classifier between generated sentence tokens and the gold scene_objects token bag, making it a proxy for grounding rather than grounding itself. Ours increases the rate by $\Delta=+0.103$ (paired Wilcoxon $p=1.3\times 10^{-4}$ , ***).

Figure 13: Probe signal accessibility. We compare the recursive mental-model latent against a compression-only VAE trained on the same Sotopia - $\pi$ contexts and responses with matched latent dimensionality. The VAE is optimized only for reconstruction, while the mental-model encoder is trained with ToM supervision and reward coupling. After freezing both encoders, we fit simple probes from $z$ and evaluate held-out recovery of social-label structure and reward-vector targets. Higher macro-F1/ $R^{2}$ for the mental latent shows that downstream social and reward signals are more accessible than in a generic compression code.

Explore similar work

May 29, 2026cs.MA

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

In cooperative multi-agent reinforcement learning (MARL), agents must coordinate with partners whose internal policies and intentions are not directly observable. While world models such as Dreamer have demonstrated strong generalization and sample efficiency in single-agent settings, their application to MARL remains limited by an inability to handle teammate-induced uncertainty. We propose a new perspective: treat teammates as structured, learnable components within the agent's world model. We introduce an architecture that factorizes the latent state of a Dreamer-style recurrent state-space model (RSSM) into environment and teammate components, and learns an auxiliary Theory-of-Mind (ToM) head to infer latent embeddings of partner behavior such as character, intent, and predicted actions from partial trajectories. These teammate latents condition the actor and critic, enabling the agent to imagine and adapt to diverse collaborators. We outline how this approach can support zero-shot and few-shot coordination in partially observable settings and propose a set of benchmarks and evaluation protocols to assess its impact. This work positions world models as not only predictors of environmental dynamics, but as simulators of social behavior, opening new directions for generalizable, human-compatible AI.

Tomas Leroy-Stone

Multi-Agent Coordination Theory of Mind Multi-Agent Collaboration Latent World Models Cooperative MARL

Ask PDF

May 10, 2026cs.CV

MOTOR-Bench: A Real-world Dataset and Multi-agent Framework for Zero-shot Human Mental State Understanding

Understanding human mental states from natural behavior is crucial for intelligent systems in the real world. However, most current research focuses on predicting isolated mental state labels, lacking structured annotations of complex interpersonal interactions. To support structured analysis, we introduce MOTOR-Bench, a carefully-designed benchmark with a real-world dataset MOTOR-dataset, containing 1,440 multimodal video clips in collaborative learning scenarios, reflecting key real-world data challenges including natural class imbalance, visual noise, and domain-specific language. Each sample is labeled by educational experts based on self-regulated learning theory. We further evaluate several state-of-the-art multimodal large language models and multi-agent systems in a zero-shot setting on our MOTOR-Bench. However, their performance on this task remains limited, suggesting that existing methods still struggle with structured reasoning from observable behavior to deeper mental states. To address this challenge, we propose a reasoning multi-agent framework, named MOTOR-MAS. It coordinates multiple agents through a structured agent coordination mechanism to infer explicit behaviors, internal cognitions, and psychological emotions. Experimental results show that our MOTOR-MAS outperforms the best single-model benchmark by 15.93 points in Macro-F1 scores for the three labels of behavior, cognition, and emotion, and outperforms the general multi-agent benchmark by 10.2 points in internal cognition prediction.

Xiaoyu Yuan, Niklas Heikkala, Tiina Törmänen +3

University of Oulu

Theory of Mind AI Agent Benchmarks Multi-Agent Reasoning Multimodal Agents

Ask PDF

May 11, 2026cs.AI

EnactToM: An Evolving Benchmark for Functional Theory of Mind in Embodied Agents

Theory of Mind (ToM), the ability to track others epistemic state, makes humans efficient collaborators. AI agents need the same capacity in multi agent settings, yet existing benchmarks mostly test literal ToM by asking direct belief questions. The ability act optimally on implicit beliefs in embodied environments, called functional ToM, remains largely untested. We introduce EnactToM, an evolving benchmark of 300 embodied multi-agent tasks set in a 3D household with partial observability, private information, and constrained communication. Each task is formally verified for solvability and required epistemic depth, and new tasks are generated increase difficulty as models improve. On the hard split, all seven evaluated frontier models score 0.0% Pass^3 on functional task completion, while averaging 45.0% on literal belief probes. Manual analysis traces 93% of sampled failures to epistemic coordination breakdowns such as withheld information, ignored partner constraints, and misallocated messages, providing a concrete target for future work.

Gurusha Juneja, Dylan Lu, Saaket Agashe +7

University of California, Santa Barbara · 2King’s College London · 3Cisco Research

Multi-Agent Coordination Theory of Mind AI Agent Benchmarks Embodied AI

Ask PDF