PaperlayerView exact revision on arXiv

Source-linked AI summary

Mental-Models for Multi-Agent Systems

Hanan Gani, Lulu Shao, Manmohan Chandraker

arXiv:2610.12453v1cs.MA

TL;DR

Multi-agent agents must infer what partners know, want, and may do under partial observability, but current systems typically lack reusable explicit partner-state representations. This paper learns a recursive latent mental model and belief-conditioned reward to train standalone policies, reporting improvements across language-only and multimodal benchmarks.

  • Problem

    Under partial observability, robust multi-agent decision-making requires inferring partners’ knowledge, intentions, and likely actions, while current LLM agents typically lack an explicit partner-state representation reusable across tasks.

    Source: §1 · p. 2

  • Method

    The framework learns a recursive latent belief state and a belief-conditioned reward model that scores candidate actions, then uses the reward signal to train a standalone policy.

    Source: §1 · pp. 2–3, §3 · p. 4

  • Results

    Across SOTOPIA, negotiation transfer, BIGTOM/TOMI/FANTOM, and MMROLE, the framework improves open-source agents; on SOTOPIA, Qwen2.5-7B achieves an average relative gain of 11.5% over the base model across two partner settings.

    Source: §1 · p. 3, §6 · p. 10

  • Takeaways & Limitations

    The results indicate that augmenting datasets with mental-state annotations, utility signals, rationales, and contrastive negatives provides useful supervision for reusable partner representations.

    Source: §6 · p. 10

    Training environments are primarily dyadic; zero-shot FANTOM transfer offers initial multi-party evidence, while end-to-end multi-party training with participant-specific latent updates remains future work.

    Source: §6 · p. 10

Abstract

from arXiv · show

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

1 Introduction

The introduction argues that multi-agent agents need explicit models of counterpart mental states under partial observability, yet current LLM systems usually represent partner dynamics only implicitly. It presents a latent, belief-conditioned framework for guiding policy learning and reports consistent gains across language-only and multimodal benchmarks.

1 Introduction

  • Current LLM agents typically lack explicit models of partners’ hidden goals, beliefs, and strategies, and may misread intent or struggle to adapt when interactions unfold unexpectedly.The introduction notes that LLM agents lag humans on Theory-of-Mind-style inference, interaction norms, and goal-driven multi-agent scenarios.Source: §1 · p. 2
  • Prior approaches shape behavior through tuning, prompts, memory, or self-play, but represent partner dynamics implicitly rather than as a reusable belief state with clear semantics and quantified uncertainty.The paper identifies interpretability, cross-task reuse, and integration with different policy learners as challenges for implicit representations.Source: §1 · p. 2
  • The framework learns a compact latent belief state from interaction history, uses it to score actions relative to inferred partner states, and trains a policy that acts standalone at inference time.The latent covers hidden goals, beliefs, knowledge, strategies, and higher-order expectations.Source: §1 · pp. 2–3
  • Across language-only and multimodal benchmarks, explicit mental-state modeling consistently improves downstream performance over base agents.On SOTOPIA, Qwen2.5-7B achieves an average relative gain of 11.5% over the base model across two partner settings.Source: §1 · p. 3

2 Related Work

Related work develops multi-agent systems through external orchestration, interactive-agent environments and training pipelines, and Theory-of-Mind inference and evaluation. This work instead learns a latent partner representation, couples it to utility prediction, and uses the resulting signal to train policies across language-only and multimodal benchmarks.

2 Related Work

  • Unlike workflow-oriented systems that emphasize role decomposition, protocols, and tools, this work learns latent partner beliefs under partial observability to shape downstream control.CAMEL, MetaGPT, ChatDev, and AutoGen are cited as examples of systems based on external orchestration.Source: §2 · p. 3
  • Interactive environments such as SOTOPIA and MMROLE provide testbeds and training recipes, but generally do not learn an amortized partner model reused as a decision variable across benchmarks.The proposed approach couples a latent partner representation from interaction history with reward modeling, then distills its signal into a policy.Source: §2 · p. 3
  • Prior ToM work includes latent agent modeling, benchmark development, explicit belief tracking, and multimodal inference; this paper instead trains policies using an amortized latent state coupled to utility prediction.The framework is evaluated across both language-only and multimodal benchmarks.Source: §2 · pp. 3–4

3 Methodology

The methodology learns a recursive mental model and a mental-model-aware utility, then trains a policy to prefer actions that utility scores highly. At deployment, only the trained policy is used.

3.1 Amortized Mental-model Learning

  • The encoder maps interaction history to a compact latent counterpart state, recursively separating first-order state from higher-order state conditioned on that inference.The Gaussian posterior represents uncertainty over plausible mental states; latent factors are decomposed into belief, intent, and thought-process subspaces.Source: §3.1 · pp. 4–5

3.2 Mental-Model-Aware Utility

  • A utility head scores candidate actions by expected outcomes across plausible counterpart states, while joint training makes the latent decision-relevant and the utility model use it.The scoring objective may also use hard-negative ranking when contrastive responses are available.Source: §3.2 · p. 5

3.3 Policy Learning under Mental-conditioned Reward

  • The policy is trained to favor sampled actions with higher mental-model-aware reward relative to other plausible choices, then deployed without the mental model or reward model.The update uses mental-reward-guided GRPO, and the reference term keeps the update near the base policy.Source: §3.3 · p. 6

4 Experimentation

The experiments evaluate the framework across language-only and multimodal benchmarks, including transfer to new partner types, domains, and ToM task formats. Results report improvements over comparison agents and training variants across these settings.

4 Experimentation

  • Experiments use LoRA adapters on 7B-scale backbones, with separate language and multimodal backbones and benchmark-specific training and evaluation sources.The mental-state/reward model is trained before policy optimization with supervised warmup when used, followed by mental-reward-guided GRPO using a frozen scorer.Source: §4 · p. 6

4.1 Main Experimentation Results

  • 3.330→3.541 (LLaMA), 3.511→3.784 (Mistral), and 3.290→3.856 (Qwen) on SOTOPIA-ALL with a fixed GPT-4o-mini partner.On SOTOPIA-HARD, average scores also improve for all three agents.Source: Table 1 · p. 7, §4.1 · p. 7, Table 2 · p. 7
  • Figure 3 shows the Qwen2.5 7B agent outperforming stronger closed-source baselines, while explicit mental-model training improves over SFT-only and GRPO-only variants.The open-source variant comparisons are averaged across LLaMA, Mistral, and Qwen.Source: §4 · p. 7, §4.1 · p. 7
  • 6.588→7.245 (LLaMA), 7.109→7.552 (Mistral), and 7.087→7.633 (Qwen) in partner-averaged SOTOPIA settings, with improvements across all agents.The averages cover GPT-4o-mini and self-partner settings.Source: Table 3 · p. 7, §4.1 · p. 7
  • SOTOPIA-trained agents retain high deal rates and reduce walkaways on 100 CRAIGSLIST-BARGAIN examples without task-specific fine-tuning.The cross-domain transfer result is reported in Table 4.Source: Table 4 · p. 7, §4.1 · p. 7
  • On BIGTOM, the mental-model variant is more consistent across forward belief, forward action, and backward belief; on TOMI, the transferred policy improves Average and Joint over the base model.Table 5 reports BIGTOM true-belief, false-belief, and paired-branch accuracy; Table 7 reports TOMI generalization.Source: §4.1 · pp. 7–8, Table 5 · p. 8, Table 7
  • MMROLE mental-model training improves overall scores over MMRole-Agent in both in-domain and out-of-domain splits, with the largest stated gain on Out: 1.019 versus 0.981.Reported improvement dimensions include Coherence, Response Accuracy, Personality Consistency, Knowledge Consistency, and Tone Consistency.Source: Table 6 · p. 8, §4.1 · p. 8

5 Analysis and Ablations

The analyses test whether recursive mental-model latents encode reward-relevant, structured, and nested partner states, and examine how supervision and posterior choices affect learning. Latent traversals further probe whether learned directions correspond to belief and strategy variables.

Significance of BIT structure

  • Figure 4b–c shows that BIT slices are separated yet related and recover ToM-relevant MMROLE factors more accurately than a flat latent of equal dimensionality.The slices are interpreted as complementary rather than redundant components.Source: §5 · pp. 8–9, Figure 4 · p. 9

Is mental-model simply a compression code?

  • Figure 4a shows that scrambling mental-model z sharply disrupts joint reward prediction, while scrambling the flat compression VAE latent has only a weak effect.This supports a tighter coupling between the mental-model latent and reward-relevant decision structure than for the reconstructive flat latent.Source: §5 · p. 8, Figure 4 · p. 9

Recursive z1 →z2 structure

  • Shuffling z2 reduces recoverable second-order mental-state information more than reward information, supporting nested ToM structure in the z1 →z2 link.The shuffled-pair control preserves each block’s marginal distribution while breaking its recursive alignment.Source: §5 · p. 9, Figure 4 · p. 9

Sensitivity to mental-state supervision

  • 0.350 overall-score points over the mental-free baseline with 25% annotations; performance rises with annotation coverage to 3.856.The 25% subset is contained in the 50% subset, and the architecture and evaluation setup are held fixed.Source: §5 · p. 9, Table 8 · p. 9

Posterior-family ablation

  • A full-covariance posterior provides no measurable benefit over the diagonal posterior on held-out BIGTOM reward ranking, margin, or belief recovery.The 0.2-point accuracy difference is within the ±0.65-point sampling uncertainty at n=5,976; cross-coordinate dependence remains available elsewhere in the model.Source: §5 · p. 9, Table 9 · p. 9

Functional Latent Traversals

  • Figure 5 shows that belief traversal shifts the branch classifier and frozen reward head, while offer/proposal traversal raises strategy signal and local label enrichment.The offer/proposal direction has greater second-order ToM block mass, suggesting these strategies rely more on modeling possible partner responses.Source: §5 · p. 10, Figure 5 · p. 10

6 Conclusion

The paper concludes that its framework couples recursive Theory-of-Mind latents with belief-conditioned reward learning, then trains policies that act independently at test time. The framework improves open-source agents across social-interaction and Theory-of-Mind benchmarks, while multi-party training and broader evaluations remain future work.

Conclusion

  • The framework couples an amortized recursive Theory-of-Mind latent state to a belief-conditioned reward model for policy training.The resulting policies act independently at test time while retaining benefits of explicit partner-state modeling.Source: §6 · p. 10
  • Across SOTOPIA, negotiation transfer, BIGTOM/TOMI/FANTOM, and MMROLE, the framework improves open-source agents on social-interaction and Theory-of-Mind benchmarks.Dataset augmentation with mental-state annotations, utility signals, rationales, and contrastive negatives also provides useful supervision for reusable partner representations.Source: §6 · p. 10
  • Training environments are primarily dyadic; end-to-end multi-party training with participant-specific latent updates remains future work.The paper also identifies larger policies, more diverse partners, richer online updates, and human-in-the-loop evaluation as future directions.Source: §6 · p. 10

A Appendix

The appendix is organized into methodology and training details, dataset-specific instantiations, additional results and ablations, and an impact statement.

Appendix

A.1.1 Dataset-Specific Methodology and Notation

Dataset examples use a tuple with observable context and optional supervision fields, allowing one objective to support several interaction formats. Missing fields are handled by masking their corresponding auxiliary losses.

A.1.1 Dataset-Specific Methodology and Notation

  • Each dataset example is represented as a tuple beginning with observable context and an optional hard negative.Source: §A.1.1 · p. 14
  • Optional tuple fields include dataset-specific utility vectors, dimension masks, first- and second-order mental-state targets, and rationale or reward-explanation text.The dimension mask indicates which utility dimensions a contrastive negative is intended to violate.Source: §A.1.1 · p. 14
  • Masking losses for missing fields lets the same objective support open-ended dialogue, structured belief QA, and multimodal role-play.Source: §A.1.1 · p. 14

A.1.2 Architecture Overview

A shared causal language-model or vision-language backbone separately encodes task context and candidate actions, whose pooled representations feed latent-posterior and reward heads.

A.1.2 Architecture Overview

  • The backbone separately encodes context and candidate actions, then sends their pooled states to the latent-posterior and reward heads.The implementation uses 128-dimensional recursive latent levels, partitioned into belief, intent, and thought subspaces as an interface for mental-state decoders.Source: §A.1.2 · p. 14

A.1.3 Implementation Details

The implementation uses LoRA-adapted 7B-scale backbones, with separate language-only and multimodal model choices. Training draws on ToM-augmented data, while several listed benchmarks are reserved for evaluation or transfer checks.

A.1.3 Implementation Details

  • Open-source models use LoRA adapters on 7B-scale backbones, with Qwen2.5-7B-Instruct for language-only work and Qwen2.5-VL-Chat for multimodal work.LLaMA2-7B-Chat and Mistral-7B-Instruct-v0.3 provide additional SOTOPIA comparisons.Source: §A.1.3 · p. 14
  • Training uses ToM-augmented SOTOPIA-π episodes, synthetic BIGTOM-style scenarios, and MMROLE turns; TOMI, official BIGTOM and SOTOPIA, CRAIGSLIST-BARGAIN, and FANTOM are used for evaluation or transfer checks.SOTOPIA training episodes do not overlap with the original SOTOPIA evaluation episodes.Source: §A.1.3 · p. 14
  • Across benchmarks, training may include supervised warmup or mental-prefix SFT before GRPO, and typically uses 128-dimensional latent states with belief, intent, and thought subspaces.The reported setup uses bf16, AdamW, warmup scheduling, KL annealing, delayed second-order latent updates, and NVIDIA RTX A6000 48GB GPUs.Source: §A.1.3 · p. 14

A.1.4 Shared Mental–Reward Objective

The shared objective couples recursive mental-state inference with action-utility prediction, so latent partner representations are supervised by mental-state targets and tied to downstream decision quality. The same loss template accommodates datasets with different supervision through masking unavailable targets.

A.1.4 Shared Mental–Reward Objective

  • The shared objective learns a latent representation jointly constrained by mental-state supervision and action-utility supervision, rather than fitting either in isolation.Dataset-specific context, response, utility, mental-state targets, and negatives use a shared loss template.Source: §A.1.4 · p. 15
  • A recursively factorized two-level latent hierarchy represents first-order counterpart inferences and higher-order inferences conditioned on the first level.Belief, intent, and thought subspaces are optional implementation structure, not requirements of the methodology.Source: §A.1.4 · p. 15
  • The action-dependent utility head predicts how utility changes with a candidate action under the inferred mental state.The shared backbone encodes context and candidate actions into representations used by the model.Source: §A.1.4 · p. 15
  • When textual mental-state supervision is available, decoders reconstruct first- and second-order descriptions, while an explanation-conditioned head predicts the same utility target from rationale text.The explanation-conditioned prediction tests whether the rationale is consistent with the outcome; supervised terms are masked when targets are unavailable.Source: §A.1.4 · p. 15
  • Preference training compares observed or preferred actions with hard negatives only on dimensions the negative was constructed to violate.When a dataset supplies scalar preferences, the mask is all ones; for generated negatives, masking avoids penalizing unviolated dimensions.Source: §A.1.4 · p. 16
  • The objective’s KL terms regularize the first-order posterior toward a unit Gaussian and the second-order posterior toward a learned prior conditioned on the first latent level.The conditional prior recursively ties the second latent level to the first rather than imposing unconditional independence.Source: §A.1.4 · p. 16
  • Removing mental-state reconstruction or rationale conditioning reduces reward-regression correlation by 29% and 26%, respectively; removing hard-negative training lowers hard-negative preference accuracy from 100% to 81%.The hard-negative removal reduces regression correlation by 7%, indicating a distinct discriminative signal.Source: §A.1.4 · p. 16
  • The resulting latent state is amortized from history, grounded in available mental-state supervision, and tied to decision quality through coupled utility prediction.Unavailable annotations mask their auxiliary terms, allowing the objective to span datasets with different supervision granularity.Source: §A.1.4 · p. 16

A.1.5 Policy Optimization Details

After learning the mental-state and reward model, the method freezes it and trains a policy using its belief-aware reward signal. Candidate-level scores rank responses, while context-only scores are reserved for diagnostics or optional prompt filtering.

A.1.5 Policy Optimization Details

  • A scalarization operator maps vector-valued utility to a scalar control signal for policy optimization.The response-dependent term ranks candidate actions; the context-only latent score is used for diagnostics and optional prompt filtering, not candidate-level GRPO advantages.Source: §A.1.5 · pp. 16–17
  • For each policy input, the method samples candidate actions and normalizes their response-level advantages within the group.It clips the centered response-level advantage.Source: §A.1.5 · p. 17
  • Policy updates use a clipped policy-gradient objective regularized by a reference model, with the reward supplied by the frozen mental-state and reward model.Prompt-level scores may filter or stratify prompts but are not shifted into within-group advantages.Source: §A.1.5 · p. 17
  • The policy is trained under the frozen mental-state and reward model, although its latent state is not fed explicitly to the policy at deployment.The policy may be initialized with supervised fine-tuning on observed actions.Source: §A.1.5 · pp. 16–17

A.1.6 Deployment and Evaluation

This section describes deployment of the learned policy without the reward model in benchmark-native environments and reports appendix evaluations of visual-grounding proxies and mental-model latent structure. It also summarizes impact considerations and checklist disclosures on asset documentation, human-subject research, and method-specific LLM use.

A.2.1 BIGTOM

  • BIGTOM augmentation expanded 200 scenarios to 5,000 synthetic scenarios and 9,964 training rows, adding validated second-order beliefs and rationale supervision while holding out official BIGTOM and TOMI for evaluation.The source scenarios contain belief, desire, and action questions across aware and unaware branches; each synthetic scenario is flattened into one example per perceptual condition.Source: §A.2.1 · pp. 17–18

A.2.2 SOTOPIA

  • SOTOPIA augmentation produced 1,647 annotated episodes with seven-dimension reward reasoning and first- and second-order mental-state targets, used to train while official evaluation episodes were held out.Training contrasts observed utterances with fluent but socially or strategically poor responses, and evaluation uses the official evaluator across the seven dimensions and their mean.Source: §A.2.2 · pp. 18–19

A.2.3 MMROLE

  • MMROLE’s visual ToM augmentation annotated image-grounded percepts and first- and second-order beliefs, then quality filtering retained 15,607 turns for training.The multimodal model jointly encodes images and textual context, reconstructs visual mental states, and contrasts in-character responses with ToM violations; evaluation uses the official protocol.Source: §A.2.3 · pp. 19–20

A.2.4 Mental-State Annotation and Human Verification

  • A blinded audit of 100 MMROLE and SOTOPIA annotations received a mean faithfulness rating of 4.45/5 (89.0% rescaled); 90 scored at least 4 and 56 scored 5.The audit was a quality spot-check, not an inter-rater agreement study; lower ratings mainly reflected overconfident, unsupported, underspecified, or contradictory inferences.Source: §A.2.4 · p. 21, Table 11 · p. 21

A.3 Additional Results and Ablations

  • Zero-shot on 5,000 FANTOM probes, the policy improved over Base on all four metrics and scored 52.2 on second-order ToM, versus 49.7 for Base and 41.4 for SFT.It was trained only on BIGTOM-style data, with no FANTOM training, fine-tuning, or adaptation.Source: §A.3 · p. 21, Table 12 · p. 22
  • On BIGTOM, the model had the strongest open-source paired accuracy in forward and backward belief and improved paired forward-action scores over the base model in both settings.The latent geometry also separates aware and not-aware branches within nearby task/scenario regions, with branch F1 = 0.989; belief is recoverable but not the largest visual axis of variation.Source: §A.3 · pp. 21–23, Table 13 · p. 22, Figure 6 · p. 22
  • On SOTOPIA, the second-order latent broadened gains, BIT preference loss was 9.6×10−5 vs. 1.0×10−2 for flat supervision, and transferred reward scored 3.557 vs. 3.541.BIT’s median preference margin was +25.06 versus +12.76 for flat supervision; the second-order latent showed broader gains on knowledge, goal completion, financial reasoning, and relationship dimensions.Source: §A.3 · pp. 23–24, Figure 7 · p. 23, Figure 8 · p. 23
  • The transferred Qwen mental-reward model preserves overall SOTOPIA performance when used with LLaMA2-7B, supporting generalization across agents.The figure compares same-architecture mental-reward modeling with the transferred Qwen mental-reward model.Source: Figure 9 · p. 24
  • Role-swap interventions show weak but measurable role-addressability only for BIT: swapping its role-specific latent block pulls the matching role prediction more than non-matching roles.Flat and Shuffled variants remain diffuse, and the absolute effects are small.Source: §A.3 · p. 24, Figure 10 · p. 24
  • 0.662 5-NN accuracy (silhouette +0.046) for [z1; z2] on the 5-way emotion task compares with 0.529 for Qwen without the adapter and 0.200 chance.The probe uses a filtered, balanced DailyDialog subset; t-SNE also shows more class separation for z1 and [z1; z2] than for no-adapter Qwen.Source: §A.3 · p. 24, Figure 11 · p. 25
  • ∆= +0.103 mean visual grounding rate for Ours over Base on 200 held-out MMROLE examples (paired Wilcoxon p = 1.3 × 10−4).The response-level metric uses token overlap with gold scene objects, so it is a proxy for grounding rather than grounding itself.Source: Figure 12 · p. 25, §A.3 · p. 25
  • Linear probes recover social-label structure and reward-vector dimensions more effectively from mental-model latents than from a matched compression-only VAE.The encoders share latent dimensionality and training contexts; the VAE uses reconstruction alone, while the mental-model encoder uses ToM supervision and reward coupling.Source: §A.3 · p. 26, Figure 13 · p. 26

A.4 Impact Statement

  • The impact statement identifies potential benefits for interpretable, controllable collaboration and negotiation, while warning that inferred states could enable manipulation, privacy invasion, or profiling.It treats inferred variables as task-conditioned constructs rather than ground-truth psychology and recommends consent-aware collection, privacy filtering, bias audits, uncertainty reporting, and restrictions on sensitive decisions; real-world use requires future study of oversight and safety constraints.Source: §A.4 · p. 26

NeurIPS Paper Checklist

  • The checklist reports repeated SOTOPIA evaluations with nearly identical aggregate scores but no formal run-to-run error bars or confidence intervals for all main trained-policy comparisons.It marks reproducibility as Yes and cites sections 4, A.2, and A.1.3; the MMROLE grounding analysis includes a paired Wilcoxon test.Source: NeurIPS Paper Checklist · pp. 28–29
  • The paper refers readers to Sections 4, A.1.3, and A.2 for its documentation of new assets.Source: NeurIPS Paper Checklist · p. 32
  • The paper reports that it involves neither crowdsourcing nor research with human subjects, and that it includes no experiments with humans.Source: NeurIPS Paper Checklist · p. 32
  • The authors answer “No” to whether LLMs are an important, original, or non-standard method component, stating that LLMs were not used for method-specific ideas.Source: NeurIPS Paper Checklist · p. 33

Related papers

Loading 2610.12453v1…