PaperlayerView exact revision on arXiv

Source-linked AI summary

ToM2Vec: Continuous Theory-of-Mind Tokens for Multi-Agent Coordination

Nithin Sivakumaran, Ali Payani, Hyunji Lee, Mohit Bansal, Elias Stengel-Eskin

arXiv:2610.29934v1cs.AIcs.CLcs.MA

TL;DR

ToM2Vec addresses the difficulty of representing multiple possible partner states in language by training agents to predict continuous latent summaries of a partner’s reasoning. Across five settings, it improves over baselines and transfers to unseen partner models, while analysis supports the tokens’ role in representing partner-specific reasoning.

  • Problem

    Text-based LLM agents can settle on a single description of a partner even when successful coordination requires reasoning over multiple possible partner states.

    Source: §1 · pp. 1–2

  • Method

    ToM2Vec trains an agent to predict continuous latent tokens summarizing a partner’s reasoning, then conditions its own subsequent reasoning and actions on those tokens.

    Source: §1 · p. 2

  • Results

    Across five settings, ToM2Vec improves over all baselines and generalizes to unseen partners; with GPT-5.4, it gains 17.3 points on average in Overcooked.

    Source: §4 · p. 7, §6 · p. 8

  • Takeaways & Limitations

    The results support latent-space partner modeling as a useful alternative to representing a partner’s reasoning in text.

    Source: §6 · p. 8

    Training requires access to partner chain-of-thought and shared representation spaces, and the experiments use English-only tasks.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Forming an accurate theory-of-mind (ToM) is key to successful multi-agent collaboration and requires reasoning under uncertainty to capture multiple possible partner states. This is hard to do in the text representation space native to large language model (LLM) agents, which more easily describe a single partner state. To address this challenge, we introduce ToM2Vec, which learns to form abstract theories-of-mind in a rich latent space via a fixed budget of continuous reasoning tokens. During training, supervision for these tokens comes from the partner's actual thoughts, which the ToM tokens are trained to abstract and match. At test-time, ToM tokens are generated by each agent and act as conditioners for that agent's thoughts and actions. Across several environments covering collaborative, competitive, and mixed-motive settings (Overcooked-AI, multiparty negotiation, Decrypto, Hanabi, and Collab Escape/Capture), we evaluate ToM2Vec's ability to support multi-agent interactions and collaboration, finding consistent improvements over the base model and a variety of competitive trained and prompted baselines. We show that latent reasoning is critical, with performance dropping when reasoning is expressed in text. Moreover, we show that ToM2Vec training generalizes to partner model families not seen at train-time, including a proprietary model (GPT-5.4). In our analysis, we provide evidence that the latent tokens encode a ToM, rather than task-specific information: ToM2Vec's tokens improve the agent's ability to directly answer questions about partner beliefs, desires, and next actions.

1 Introduction

Effective multi-agent coordination requires inferring partners’ goals, beliefs, and likely actions, but text can force an agent to commit to one partner hypothesis. ToM2Vec predicts continuous latent representations from the agent’s own perspective, improving results across tasks and supporting partner-state inferences.

1 Introduction

  • Text-based partner modeling can commit to one estimate even when multiple hypotheses are possible, whereas latent representations can preserve alternatives before a decision.The desired abstraction retains potentially useful future paths while lifting unnecessary details.Source: §1 · pp. 1–2
  • ToM2Vec trains an agent to predict a fixed number of continuous latent embeddings that abstract a partner’s thoughts, using the partner’s actual reasoning as supervision.At test time, the agent generates these tokens from its own perspective.Source: §1 · p. 2
  • Compared with untrained Qwen3-8B, TOM2VEC training improves average Overcooked score by 12.0 points and raises multi-party negotiation deal success from 62.0% to 71.0%.Source: §1 · p. 2
  • With an unseen GPT-5.4 partner, ToM2Vec improves Overcooked scores by 17.3 points on average and matches or exceeds the base agent with Gemma-4.Training used a Qwen3-8B partner; GPT-5.4 gains were positive on every layout.Source: §1 · p. 3
  • ToM2Vec tokens better support direct inferences about a partner’s beliefs, desires, and next actions, and context-corruption results indicate partner-specific information.Latents from the partner’s actual context assign lower perplexity to true reasoning than latents from altered layouts or adversarial intent.Source: §1 · p. 3

2 Method

The method trains an agent to anticipate a partner from its own partial context by predicting a fixed budget of latent tokens. Partner reasoning supplies training targets; at inference, the agent conditions its subsequent reasoning and actions on its predicted tokens.

2.1 Problem Setup

  • The task is formulated as a decentralized partially observable Markov decision process, with agents acting from role-permitted observations and histories.The shared-reward assumption applies only to cooperative tasks.Source: §2.1 · p. 3

2.2 Modeling Partners via Latent Reasoning

  • From its own context, the main agent predicts K continuous hidden-space vectors intended to summarize the partner’s likely reasoning.LatentK uses K autoregressive hidden-space steps rather than decoding each step to vocabulary tokens.Source: §2.2 · p. 3

2.3 Self-Supervised Training via Self-Play

  • In self-play, the partner’s reasoning trace is segmented into K latent targets, and the main agent is trained to match them with a cosine-similarity objective.The targets include start and end states and representative hidden states from evenly divided reasoning chunks; K=10 in practice.Source: §2.3 · pp. 3–4

2.4 Inference

  • At inference, the agent predicts ToM tokens from its own context and includes them in the context for its subsequent reasoning and actions.The partner’s private context is unavailable at test time; the predicted tokens serve as an abstract theory of its likely reasoning.Source: §2.4 · p. 4

3 Experiments and Results

The experiments evaluate TOM2VEC in cooperative, mixed-motive, and competitive multi-agent settings, using domain-specific metrics and baselines. TOM2VEC consistently improves over the untrained base model across the tested domains.

3.1 Experimental Setup

  • The setup covers Overcooked-AI, negotiation, Decrypto, Hanabi, and Collab Capture/Escape, with Qwen3-8B trained in self-play and unseen model-family partners used for generality tests.Every agent receives only its role-permitted observation and cannot see partners’ private contexts; GPT-5.4 and Gemma-4 are unseen during training.Source: §3.1 · pp. 4–5

3.2 Metrics

  • The evaluation reports mean episode score, negotiation deal and utility scores, Decrypto role-specific win rates, Hanabi game score, and Capture/Escape success rates and turn counts.Source: §3.2 · p. 5

3.3 Baselines

  • Baselines include an untrained base model, prompted and trained language ToM, Self-Consistency, SimToM, Adaptive ToM, and an oracle with partner-private-thought access.The oracle is an upper bound; Language ToM verbalizes a summary of the partner’s reasoning rather than predicting latent tokens.Source: §3.3 · p. 5

3.4 Main Results

  • Across five domains, TOM2VEC improves over the Qwen3-8B base model, with consistent results indicating a useful partner model across varied tasks and configurations.The evaluation spans cooperative, mixed-motive, and competitive settings.Source: §3.4 · pp. 5–6
  • TOM2VEC raises Overcooked mean episode score by 12.0 points on average over the untrained model, while negotiation deal success rises from 62.0% to 71.0%.It outperforms all baselines on Overcooked; negotiation deal success and own score are significantly higher than all baselines.Source: §3.4 · pp. 5–6, Table 1 · p. 5, Table 2 · p. 6
  • In Decrypto, TOM2VEC improves decoder and interceptor win rates by 4.4 and 6.1 percentage points, respectively; in Hanabi, its score is 11.3 out of 25.The Hanabi score is 2.4 points above the untrained baseline and 0.7 points above Language ToM.Source: §3.4 · p. 6, Table 3 · p. 6, Table 4 · p. 6
  • In Collab Capture, TOM2VEC raises capture rate from 0.67 to 1.00; in Collab Escape, it raises escape rate from 0.76 to 0.92.Capture turn count falls from 4.76 to 4.32; Escape ties the highest success rate while keeping turn count consistent. In Collab Escape, Language ToM drops 10 percentage points while TOM2VEC improves by 16 percentage points.Source: §3.4 · pp. 6–7, Table 5 · p. 7
  • Language ToM recovers only a small part of TOM2VEC’s Overcooked gains, while prompting methods produce marginal gains and can decrease performance on Cramped Room.Oracle access to the partner’s private context raises mean score by 23.0 points over the base model; prompted ToM averages a 1.7-point gain.Source: §3.4 · p. 6

4 Discussion and Analysis

TOM2VEC transfers to unseen partner model families and improves partner-state inference. Context-corruption tests also indicate that its latent tokens encode partner-specific reasoning.

4 Discussion and Analysis

  • With GPT-5.4, TOM2VEC improves over the base agent by 17.3 points on average, and it matches or exceeds the base agent with Gemma-4.Training used only a Qwen3-8B partner; the GPT-5.4 gain was positive on every layout.Source: §4 · p. 7, Table 7 · p. 7
  • TOM2VEC achieves the highest accuracy on belief, desire, and next-action probes, beating the next-best baseline by 5.6, 2.0, and 2.7 points, respectively.Belief questions test second-order inference: what the partner thinks the agent is trying to do.Source: §4 · pp. 7–8, Table 7 · p. 7
  • Perplexity is 1.412 with normal latents, versus 1.702 for altered-layout latents and 1.613 for adversarial-intent latents.The correct CoT receives the lowest perplexity in 98.6% of cases with true partner reasoning; context corruption degrades prediction.Source: §4 · p. 8, Table 8 · p. 8

5 Related Work

Related work covers ToM research for LLM agents and latent reasoning, while positioning TOM2VEC as a latent partner-modeling approach trained through self-play.

5 Related Work

  • Prior work on LLM ToM studies model capabilities and uses prompting, structured decomposition, symbolic belief tracking, or training to improve them.The paper’s review focuses on LLM agents rather than ToM work in non-LLM agents.Source: §5 · p. 8
  • Unlike prior prompting, symbolic-tracking, and ToM-training approaches, TOM2VEC learns a latent partner representation supervised directly during self-play.Source: §5 · p. 8
  • Latent-reasoning methods explore continuous or soft tokens and alternative paths, but TOM2VEC applies latent reasoning to modeling other agents under partial observability.The authors describe this as the first exploration of the paradigm for modeling other agents.Source: §5 · p. 8

6 Conclusion

The conclusion presents TOM2VEC as a self-play method for modeling partners with continuous thought tokens. Across five cooperative and competitive settings, it improves over baselines, generalizes to unseen partners, and supports more accurate partner-state inferences.

6 Conclusion

  • TOM2VEC trains an agent via self-play to model its partner using continuous thought tokens.Source: §6 · p. 8
  • Across five cooperative and competitive settings, TOM2VEC improves over all baselines and generalizes to partners unseen during training.The analysis also finds more accurate partner-state inferences, which the authors interpret as stronger ToM understanding.Source: §6 · p. 8
  • The results suggest that modeling another agent in latent space is a useful alternative to modeling it in text.Source: §6 · p. 8

Limitations

TOM2VEC’s training and evaluation have requirements that constrain its use.

Limitations

  • TOM2VEC requires partner CoT during training and a shared representation space, so training uses the same base model for agent and partner.Training against one partner can transfer to other test-time partners, but does not remove the training-time CoT requirement; experiments also use English-only tasks.Source: Limitations · p. 9

A Model Details

The appendix specifies the generation settings used for Qwen3-8B, Gemma-4-E4B, and GPT-5.4, and the optimizer and batch setup used to train ToM2Vec.

A Model Details

  • Qwen3-8B and Gemma-4-E4B use temperature 1.0 and top_p 0.95, with top_k set to 20 and 64, respectively.Qwen3-8B also uses repetition_penalty 1.0 and presence_penalty 1.5; inference uses vLLM with GPU memory utilization at 0.90.Source: A Model Details · p. 12
  • GPT-5.4 is tested with reasoning effort set to high and verbosity set to medium.The Qwen3-8B configuration is initialized from HuggingFace's default generation configuration.Source: A Model Details · p. 12
  • ToM2Vec is trained with AdamW at a learning rate of 10^-6, batch size 4, and one gradient-accumulation step.Source: A Model Details · p. 12

B Dataset Details

The appendix describes training and evaluation data across Overcooked-AI, negotiation, Decrypto, Hanabi, and Collab Escape/Capture, with task-specific sampling and splits.

B Dataset Details

  • Overcooked-AI training uses 4,000 self-play states collected with a random agent, while baselines run 400 steps across three layouts and 30 seeds.The layouts are Cramped Room, Asymmetric Advantage, and Coordination Ring.Source: B Dataset Details · p. 12
  • Negotiation data comprise 500 GPT-5.4-generated settings, with discussion lasting 24 turns and an 80/20 training-validation split over turn samples.Source: B Dataset Details · p. 12
  • The Sally-Anne table reports overseer accuracy for predicting where the first player would look and a separate measure of object-state understanding.Source: Table 9 · p. 12
  • Decrypto training uses 1,000 randomly sampled keyword sets for both decoder and interceptor roles, with another 1,000 sets for evaluation.The sets are formed from the original paper's 640-word list.Source: B Dataset Details · p. 12
  • Hanabi training uses 4,000 randomly generated game states, and evaluation uses 100 randomly initialized games.Source: B Dataset Details · p. 12
  • Collab Escape and Capture each use 4,000 random-agent self-play states for training, with evaluation on 12 and 6 seeds, respectively.Source: B Dataset Details · p. 12

C Additional Results

Additional analyses report ToM2Vec's results on a false-belief test, evidence-based belief revision, statistical comparisons, and inference costs. The analyses find improved partner-state understanding and movement-sensitive shifts in the likelihood of partner thoughts.

C Additional Results

  • ToM2Vec scores 0.88 on the Sally-Anne test, improving 0.14 over the untrained agent, tying Self-Consistency, and outperforming other baselines.State-understanding scores differ by at most 0.08 across methods, which the authors take to indicate that partner-state gains reflect ToM rather than better state understanding.Source: C Additional Results · p. 13
  • As the partner moves toward a dispenser, perplexity falls for the corresponding dish or onion thoughts, indicating evidence-sensitive belief revision.The analysis samples five partner chains of thought per intended direction and scores them as the partner moves east or west.Source: C Additional Results · p. 13, Figure 3 · p. 13
  • ToM2Vec significantly outperforms baselines in 16 of 18 Overcooked comparisons, in seven of 10 Decrypto comparisons, and over all Language ToM baselines in Hanabi.It beats every baseline on the more complex Asymmetric Advantages and Coordination Ring layouts; none of five negotiation comparisons is statistically significant.Source: C Additional Results · p. 14, Table 10 · p. 14, Table 11 · p. 14
  • Emitting 10 latent tokens adds 1,150.1 ms per query on average, while inference peak memory remains 16.3 GB for both untrained and ToM2Vec models.The authors describe the latency overhead as a fixed constant and the 83.96 GB peak training cost as one-time, with no test-time deployment effect.Source: C Additional Results · p. 14, Table 14 · p. 14

D Prompts

The appendix supplies prompts for Overcooked and notes that the original dataset prompts are used for negotiation and Decrypto. The Overcooked prompt specifies game rules, map interpretation, available actions, and questions for probing partner beliefs and desires.

D Prompts

  • The appendix uses the original dataset prompts for negotiation and Decrypto, while the Overcooked prompt emphasizes avoiding duplicated work and coordinating shared-kitchen tasks.Source: D Prompts · pp. 14–15
  • The Overcooked prompt explains the soup-delivery loop, including ingredient collection, pot preparation, cooking, plating, and delivery against active orders.Source: D Prompts · p. 15
  • The prompt defines map symbols and player-cell notation, then describes how directional actions move or turn agents and how interaction requires facing an object.STAY preserves position and facing; INTERACT acts on the cell in front of the agent.Source: D Prompts · p. 15
  • Overcooked partner probes ask what the partner believes the main agent is doing and what the partner is trying to do, using multiple-choice action categories.Source: D Prompts · pp. 16–17

Related papers

Loading 2610.29934v1…