Structured Machine Theory of Mind from Agent Trajectories

AI summary

Overview

Research area: Machine Theory of Mind (MToM) — the intersection of causal representation learning, inverse planning, and neural sequence modeling of human behavior, applied to longitudinal trajectory data.

Technical level: Intermediate. The paper assumes familiarity with transformer encoders, multi-task losses (KL divergence, binary cross-entropy), and the basics of structural causal models and interventions, but the central architectural idea is conveyed clearly enough for a reader with general deep-learning background.

Scope: The paper proposes a framework (SMToM) and a concrete model (BDIBottleneck) that learns explicit, separately supervised belief and desire representations from pedestrian navigation trajectories, routes them through a goal-prediction bottleneck, and evaluates whether editing those representations produces coherent counterfactual changes in predicted destinations.

What This Paper Is About

Predictive models trained on large behavioral trajectory datasets can forecast where people go, but they do not represent the beliefs and desires that cause the behavior — so they cannot answer "what would this agent do if it believed a location was closed?" The paper argues that answering such questions requires explicitly recovering mental states, and it builds a model that attributes belief and desire vectors from observed paths and then feeds only those vectors (plus a current-trajectory embedding) into a destination predictor, making mental-state edits a direct architectural operation. Because real-world data contains no ground-truth mental-state labels, the authors instantiate and test the framework in a controlled simulated pedestrian navigation domain where the true beliefs and preferences are known by construction.

Key Contributions

1. The SMToM framework. A formulation of Machine Theory of Mind that combines the explicit, intervenable representation of model-based Bayesian inverse planning with the data-driven flexibility of neural predictors, organized around the Belief-Desire-Intention (BDI) causal structure.

2. The goal bottleneck architecture (BDIBottleneck). A model in which desire and belief heads produce explicit probability vectors, and the goal head consumes only [π̂ ‖ b̂ ‖ z_τ] — the predicted mental-state channels plus a current-trajectory embedding — so that counterfactual intervention on beliefs and desires requires no re-encoding of context.

3. A supervised, ground-truth-instrumented testbed. A controlled navigation simulator with M = 24 points of interest grouped into C = 6 desire categories, binary belief vectors over POIs, eight training belief configurations (B₀–B₇) and two held-out configurations (BH0, BH1) reserved for belief-generalization evaluation.

4. A counterfactual evaluation protocol. Local channel interventions (editing one mental-state channel and re-running the frozen goal head) plus a model-planned diagnostic that compares the model's counterfactual goal against the goal the simulator's agent would actually pursue under the same edited mental state.

Main Findings

• Goal inference is competitive with a matched auxiliary-supervision baseline but not uniformly better. BDIBottleneck reaches 55.56% top-1 accuracy on the test split and 56.88% on held-out new agents at full reveal, and 36.32% (test) and 36.11% (new agent) at 50% reveal. The authors explicitly frame this as calibrated: auxiliary mental-state supervision itself accounts for part of the improvement, while the bottleneck's distinctive benefit is the intervenable interface rather than guaranteed higher raw accuracy.

• Stronger early disambiguation than trajectory-only and unauxiliary context baselines. BDIBottleneck numerically leads GoalPredictor and ContextGoalPredictor at all path fractions on both standard splits, with the advantage growing at longer reveals; McNemar's exact tests with Holm–Bonferroni correction show significant gains over GoalPredictor at most path fractions on both splits, and over ContextGoalPredictor at 80–100% reveal and on held-out agents at 70–100%.

• Best learned model in seven of eight held-out-belief columns. At full reveal under held-out belief configurations, BDIBottleneck reaches 51.87% (held-in agents) and 55.90% (held-out agents), versus 43.61% and 44.86% for the unbottlenecked AuxContext. The sole exception is the 25% held-out-agent column, where AuxContext is higher by 0.21 percentage points.

• BToM (Informed) behaves as an approximate ceiling that the neural models approach at intermediate reveals. The informed Bayesian inverse planner reaches 34.8% at 50% reveal and 67.7% at 90% reveal on test (40.4% and 71.3% on new agents), and 93.1% at full reveal. BDIBottleneck marginally exceeds it in the 40–60% interval (31.74% vs. 30.21% at 40%; 36.32% vs. 34.79% at 50%; 40.62% vs. 38.12% at 60%) but falls far short at 100% reveal, because the Boltzmann likelihood pins the unique destination under a near-shortest-path route.

• Visit-frequency priors are far weaker. VisitFrequency is constant at 12.3% on test and 15.4% on new agents, independent of trajectory length; BDIBottleneck surpasses it from the earliest fractions on held-in agents and from 20% onwards on held-out new agents.

• A pooling artifact explains the 90%-to-100% accuracy dip. Several models decline slightly from 90% to 100% reveal because at full reveal the goal node is included in the input but its contribution is diluted by mean pooling over preceding steps, with no explicit termination signal.

• Desire attribution converges tightly. The desire head reaches a validation KL divergence of 0.22 against the ground-truth preference distribution, down from 0.69 at epoch 1.

• Belief attribution is harder and reported honestly. Overall accuracy on the 24-entry binary belief vector is 76.1%, but this is dominated by majority open-class entries; false-belief F1 peaks at 40.4% in early training and converges to 33.3% at the selected checkpoint, against 0% for an all-open predictor. The decline reflects the multi-task objective shifting gradient toward goal accuracy.

• Desire counterfactuals behave as intended. Substituting an agent's inferred preferences with a different activity type coherently shifts predicted destinations toward locations relevant to that activity type, evaluated as a 6×6 delta matrix over all category prototypes at path fractions 0.25, 0.50, 0.75 and 1.00.

• Belief counterfactuals work but in a complex way. Marking a location as unavailable in the belief state reduces its predicted probability as a destination, measured by alternative gain and goal-probability change with one-sided random-sign permutation tests, exact sign tests, and bootstrap confidence intervals.

Methodology in Plain English

The authors build a simulated town. There are 24 places to go, organized into 6 categories of 4 places each. Each simulated person has a stable preference distribution over categories and over places within categories, plus a binary belief about whether each place is open. The environment is physically all-open; a "false belief" is injected by flipping entries in the belief vector to zero. Each episode, the agent samples a category from its preferences, samples a destination among the believed-open places in that category (renormalizing if a preferred place is believed shut), and then walks a noisily near-shortest path to it. This makes the generator a clean causal graph: preferences and beliefs both point to the goal, and the goal points to the trajectory.

The model inverts that graph with three separate transformer encoders — one over 10 randomly chosen past trajectories from the same agent (capturing stable desire), one over the 5 most recent trajectories (capturing recent behavior), and one over the current partial path. A desire head maps the first embedding to a 24-dimensional probability vector via softmax. A belief head maps the concatenation of the desire and recent embeddings to 24 independent probabilities via sigmoid, on the reasoning that a contrast between long-run habit and recent avoidance is the cue a human would use to infer that an agent thinks somewhere is closed. The critical design decision is the goal head: it sees only the predicted desire vector, the predicted belief vector, and the current-trajectory embedding — never the raw context embeddings. Because of this, a researcher can overwrite the desire or belief vector by hand and re-run the goal head to get a counterfactual destination distribution, without re-encoding anything.

Training is multi-task: cross-entropy for the goal, KL divergence for desire, and class-weighted binary cross-entropy for belief, with belief loss upweighted (λ_b = 5.0) to counter false-belief sparsity. Only the current trajectory is randomly truncated during training; context streams stay intact. The authors compare against five baselines: a visit-frequency prior, an informed Bayesian inverse planner (BToM), a trajectory-only transformer, a three-stream context transformer with no mental-state supervision, and — most importantly — an "AuxContext" control that has identical auxiliary losses and identical encoders but an unbottlenecked goal head. That last comparison isolates how much accuracy is due to supervision versus how much is lost by forcing prediction through a low-dimensional interpretable interface.

Why This Matters

The paper's contribution is less about beating a benchmark than about establishing that a neural model can be given an editable mental-state interface without collapsing on prediction. Most high-accuracy behavior models are opaque; most interpretable Theory-of-Mind models require hand-specified dynamics. SMToM is an attempt to occupy the middle. For research, it reframes the question from "how accurately can we predict the next move" to "can we intervene on the representation and get a causally sensible answer," and it demonstrates a rigorous protocol — matched baselines, held-out belief configurations, permutation tests — for evaluating that claim.

Real-world applications:

• Urban and transit planning: asking how ridership patterns would shift if a station, route, or venue were closed, rather than merely extrapolating from historical flow.
• Assistive and service robotics: a robot that infers a person believes a corridor is blocked can plan around that belief rather than around physical reality.
• Agent-based simulation and policy stress-testing: populating simulations with agents whose preferences can be edited to test interventions before deployment.
• Human-AI interaction design: systems that maintain an explicit, inspectable model of what a user wants and what they think is true, which can be corrected directly instead of retrained.

Industry relevance: mobility, logistics, retail site selection, and any domain with dense longitudinal trace data where stakeholders want to ask counterfactual "what if" questions rather than just forecast. The bottleneck pattern — forcing a predictor to pass through a small, semantically named channel — is also a transferable architectural recipe for interpretability-constrained prediction generally.

Future Directions

• Extending beyond predefined mental-state spaces. The framework fixes 6 categories and 24 binary beliefs; the authors explicitly leave open-ended or data-discovered mental states (e.g., EM-style or Chinese-restaurant-process category discovery) to future work.
• Dynamic environments and online belief updating. The current model assumes a static world, fixed beliefs, and a single goal per episode, so false beliefs live entirely in the belief vector rather than in world dynamics.
• Moving from supervised identification to proxy-based validation. The authors are explicit that they make no claim of unsupervised identifiability — they validate recovery because ground truth exists in simulation. Extending to real data without mental-state labels will require proxy-based validation, which they name as the main path to external validity.
• Improving belief attribution specifically. False-belief F1 converging to 33.3% (from a 40.4% peak) is the clearest weakness, and the paper attributes part of the decline to the multi-task objective reallocating gradient toward goal accuracy — suggesting loss balancing or belief-specific training schedules as an open problem.

Target Audience

Researchers working on Machine Theory of Mind, inverse reinforcement learning, or causal representation learning who want a concrete, controlled testbed for intervenable mental-state representations. It is also valuable for practitioners building behavioral forecasting systems on mobility or interaction logs who need models that support "what if" queries rather than only prediction, and for methodologists interested in how to evaluate counterfactual claims in learned representations when ground truth is unavailable in the wild — the matched AuxContext baseline and the held-out belief-configuration splits are the paper's most reusable methodological contributions.

Authors’ abstract

Predictive models of human behavior trained on large-scale trajectory data optimize for statistical accuracy without representing the mental states that causally generate behavior. Such models support prediction but not principled intervention: they cannot answer how an agent's behavior would change if its beliefs or preferences were different. We introduce Structured Machine Theory of Mind (SMToM), a framework that attributes explicit, independently supervised belief and desire representations from observed trajectories within a Belief-Desire-Intention (BDI) causal structure. The central architectural element is a goal head that consumes only the predicted mental-state channels and a current-trajectory embedding; counterfactual intervention on beliefs and desires is a direct operation. We instantiate SMToM on a controlled pedestrian navigation domain where ground-truth mental states are known by construction, enabling evaluation of attribution accuracy, local channel interventions, and trajectory-level counterfactual diagnostics. The resulting model, BDIBottleneck, improves over trajectory-only and unsupervised-context neural baselines. It is also competitive with a matched auxiliary-supervised context baseline while providing an editable mental-state interface that the unbottlenecked baseline lacks. Counterfactual experiments on desires confirm that substituting an agent's inferred preferences with a different activity type coherently shifts predicted destinations toward relevant locations. Counterfactual experiments on beliefs confirm that marking a location as unavailable in the agent's belief state reduces its predicted probability as a destination, though in a complex way. These results demonstrate, in a controlled navigation setting, that explicit BDI-structured supervision is a viable foundation for interpretable and counterfactual analysis of longitudinal trajectory data.