Explicit Latent Belief State
Infer hidden goals, beliefs, knowledge, strategies, and higher-order expectations from observed history.
Mental Models for Multi-Agent Systems
An amortized latent belief state inferred from interaction history provides a stable training signal for multi-agent policy learning under partial observability.
University of California, San Diego
NeurIPS 2026
01 · Introduction
Mental-Model-Enabled Agents under Partial Observability
Multi-agent interaction is partially observable. A dialogue reveals actions and language, but not the other agent’s goals, knowledge, constraints, or expectations. Policies trained only on surface behavior can therefore choose plausible responses that are strategically or socially wrong.
Abstract
We introduce mental-model-enabled agents for multi-agent systems under partial observability. The framework learns a compact latent belief state from interaction history, jointly trains it with a mental-model-conditioned utility, and uses the resulting signal for policy learning. The final policy remains standalone at inference time. Across language-only and multimodal benchmarks, explicit mental-state modeling improves downstream multi-agent performance.
Recursive Mental-State Architecture
Represent first- and second-order belief, intent, and thought subspaces.
Coupled Mental-State/Reward Learning
Make partner representations directly useful for evaluating candidate actions.
02 · Methodology
Methodology
We jointly learn an amortized recursive mental-model and a mental-model-conditioned utility, then use the learned signal to train a standalone policy.
Amortized Mental-Model and Reward Learning
Jointly learn recursive mental states and a utility model over candidate actions.
Policy Learning under Mental-Conditioned Reward
Score sampled actions with the learned mental reward and optimize the policy.
Standalone Policy Deployment
Deploy only the trained policy in the multi-agent setting.
Train + evaluateSOTOPIA · BigToM · MMRole
Zero-shot transferCraigslist-Bargain · ToMi · FANToM
03 · Experimentation
Main Experimentation Results
We test open-ended social interaction, controlled belief tracking, and multimodal role-playing with the same mental-model learning principle.
Qwen2.5-7B, averaged across two partner settings.
Average improvement across the six reported TB∧FB settings.
Overall score versus 0.981 for MMRole-Agent.
Zero-shot BigToM transfer with the official scorer.
SOTOPIA-All · GPT-4o-mini partner
Results on SOTOPIA
Mental-model guidance improves the overall score for LLaMA, Mistral, and Qwen, with the largest gain on Qwen2.5-7B.
| Backbone | Base | Ours | Gain |
|---|---|---|---|
| LLaMA2-7B | 3.330 | 3.541 | +0.211 |
| Mistral-7B | 3.511 | 3.784 | +0.273 |
| Qwen2.5-7B | 3.290 | 3.856 | +0.566 |
Paired true-belief and false-belief accuracy
Results on BigToM
TB∧FB gives credit only when both branches of the same scenario are correct. Our model improves paired consistency across forward belief, forward action, and backward belief tasks.
| Task | Model | TB∧FB accuracy | |
|---|---|---|---|
| Without initial belief | With initial belief | ||
| Forward belief | Base | 77.0 | 48.0 |
| SFT | 84.5 | 92.0 | |
| Ours | 96.0 | 97.5 | |
| Forward action | Base | 63.5 | 67.5 |
| SFT | 76.5 | 73.0 | |
| Ours | 78.5 | 80.0 | |
| Backward belief | Base | 77.5 | 52.5 |
| SFT | 85.0 | 91.0 | |
| Ours | 95.5 | 98.5 | |
Official normalized multimodal role-play metrics
Results on MMRole
The strongest gains occur on coherence, response accuracy, personality consistency, knowledge consistency, and tone consistency—dimensions that require grounded character-state reasoning rather than fluent captioning alone.
| Split | MMRole-Agent | Ours | Gain |
|---|---|---|---|
| Overall | 0.994 | 1.023 | +0.029 |
| In-domain | 0.999 | 1.024 | +0.025 |
| Out-of-domain | 0.981 | 1.019 | +0.038 |
04 · Generalization
Cross-Domain Transfer and OOD Generalization
Transfer datasets are evaluation-only: no target labels are used for policy training.
Craigslist-Bargain
Cross-domain negotiation retains high deal rates while reducing walkaways.
- Qwen deal rate
- 91% to 94%
- Qwen walkaway
- 7% to 5%
ToMi
A different synthetic format tests whether nested belief structure is reusable.
- Average
- 87.92 to 89.16
- Joint
- 55.16 to 60.46
FANToM
Five thousand probes test transfer from dyadic training to multi-party reasoning.
- Second order
- 49.7 to 52.2
- Belief distance
- 35.2 to 38.0
05 · Analysis
Analysis and Ablations
SOTOPIA · Qwen2.5-7B
Sensitivity to Mental-State Supervision
Even 25% supervision improves over the mental-free baseline; performance continues to rise with additional annotation coverage.
MMRole · BigToM
Mental-State Annotation and Human Verification
98.04%
MMRole training turns retained after six automated checks.
4.45 / 5
Mean faithfulness in a blinded human audit of 100 annotations.
Posterior-Family Ablation
Diagonal Gaussian96.6reward accuracy
Full covariance96.8reward accuracy
The more expensive full-covariance posterior provides no measurable practical gain on the held-out BigToM split.
06
Citation
The repository contains separate pipelines for every benchmark reported in the paper, plus reproducible transfer evaluators and pinned upstream sources.
Open the code repository@inproceedings{gani2026mentalmodels,
title = {Mental Models for Multi-Agent Systems},
author = {Gani, Hanan and Shao, Lulu and Chandraker, Manmohan},
booktitle = {Advances in Neural Information Processing Systems},
year = {2026}
}