From Shortcuts to Reasoning: Robust Post-Training of Theory of Mind with Reinforcement Learning

Abstract: Theory of Mind (ToM) is a must-acquire skill for modern foundation model systems to operate effectively and safely in the real world. Recent works have explored honing ToM via post-training; however, we show that such progress is confounded by a pervasive "shortcut" issue: tasks can reach up to 99% accuracy by simply exploiting spurious causal correlations, leading to a false sense of ToM. Motivated by this, we first develop a framework to systematically examine ToM datasets for shortcuts and provide guidance for future development. We find that questions reducible to pure state tracking, such as "belief," are especially shortcut-prone compared to mind questions, such as "intention," where reasoning beyond tracking is required. Using four shortcut-free datasets across three ToM contexts, we then comprehensively study whether Reinforcement Fine-Tuning with verifiable rewards and explicit reasoning chains, called Thinking-RFT, elevates ToM beyond Supervised Fine-Tuning, or SFT. Our key findings are as follows. First, Thinking-RFT effectively improves ToM in all scenarios, with a 6% improvement over SFT, particularly in complex higher-order reasoning, with a 10% improvement over SFT, and multimodal cases, with a 7% improvement over SFT. It also generalizes notably better to unseen domains and higher-order queries while being more robust to counterfactuals. Second, ToM benefits specifically from the joint effect of reasoning and RL: Thinking-RFT outperforms Non-Thinking-RFT by 7% on average. Third, RFT works by learning to ground its reasoning on anchor cues, such as keywords and state changes, that correspond to causal factors. We believe our study is useful for developing effective and robust ToM post-training datasets and advancing critical ToM capabilities.

• The paper introduces a novel RL-based post-training method that filters dataset shortcuts to improve robust Theory of Mind capabilities.
• The approach employs a two-stage shortcut probing pipeline to diagnose and eliminate spurious correlations in widely-used ToM benchmarks.
• Empirical results show that Thinking-RFT outperforms SFT with gains up to 10% in higher-order reasoning tasks and superior generalization.

Reinforcement Learning for Robust Theory of Mind: From Shortcuts to Reasoning

Introduction

"From Shortcuts to Reasoning: Robust Post-Training of Theory of Mind with Reinforcement Learning" (2606.09092) addresses a critical challenge in the development of Theory of Mind (ToM) capabilities within foundation models. While previous ToM post-training endeavors often report impressive gains, the paper finds that many such improvements are confounded by dataset shortcuts—spurious correlations or procedural cues that allow models to bypass genuine mental-state reasoning. This work systematically audits these datasets, proposes a rigorous framework to filter shortcut-prone examples, and evaluates various post-training methods, emphasizing reinforcement learning with verifiable rewards and explicit reasoning traces (Thinking-RFT). The results provide strong evidence for the efficacy and robustness of RL-based ToM post-training when shortcut issues are neutralized.

Auditing and Diagnosing Theory of Mind Datasets

The paper formalizes a two-stage shortcut probing pipeline: (1) LLM/Agent-guided rules used to discover deterministic procedural or causal heuristics, and (2) a mutual-information-based lexical shortcut metric capturing surface correlations. Eight widely-used ToM benchmarks were audited, and four exhibited severe shortcut issues, e.g., Hi-ToM’s higher-order queries were “solved” by predicting the object location after an agent leaves, regardless of agent beliefs. These shortcut solutions result in artificially inflated accuracy—sometimes up to 99%—even for strong models like GPT-4o in zero-shot mode, with incoherent or logically flawed reasoning traces. Key recommendations include restricting training to shortcut-free datasets and explicitly designing ToM tasks that require reasoning beyond state tracking (e.g., intention, attitude).

Impact of Shortcuts on Post-Training and Evaluation

Training on shortcut-prone datasets produces models that generalize poorly, invert the ranking among post-training strategies, and exhibit diminished reasoning trace quality. This effect is robust across SFT, RFT, and even scaling (3B vs 7B). For example, in controlled experiments on ExploreToM, both SFT and RL-based methods reach >95% in-domain accuracy but fail to improve out-of-domain or higher-order queries, with RL chains-of-thought collapsing to shortcut-driven answers. This demonstrates that genuine ToM requires more than pattern extraction from spurious correlations.

Reinforcement Learning with Explicit Reasoning: Thinking-RFT

The authors implement rule-based reinforcement learning fine-tuning (RFT) using the Group Relative Preference Optimization (GRPO) algorithm, with rewards composed of accuracy, format compliance (e.g., correct reasoning chain structure), and ToM-specific role attribution. The instruction prompts elicit chains-of-thought and clear answer delineation, and reward models check both answer correctness and reasoning trace validity. Experiments show that RFT is especially effective in shortcut-free datasets across multiple modalities (narrative, conversational, multimodal)—yielding average gains of +6% over SFT and +10% in higher-order reasoning tasks. No-Thinking-RFT, which omits explicit reasoning chains, lags behind Thinking-RFT by 7% on average, underlining the necessity of reasoning for robust ToM post-training.

Numerical Results and Robustness

The empirical results highlight several strong claims:

• RFT achieves state-of-the-art accuracy, often >10% over SFT and >30% over zero-shot, especially in second-order ToM queries.
• RFT lifts performance on both mind-state (intention/desire) and tracking tasks, with higher gains in tasks requiring recursive reasoning (+10% in second-order).
• Gains persist across both 3B and 7B model scales, as well as in vision-language settings (+10.55% multimodal).
• Generalization to unseen environments and higher-order queries is markedly stronger for RFT, with only ∼2% drop compared to 7% for SFT.
• Counterfactual robustness experiments confirm that RFT reasoning chains are anchored on true causal factors (CFC-LB = 0.79–0.84), whereas SFT is brittle to minor changes.

Analysis: Reasoning Trace Quality and Attention Visualization

Quantitative LLM-judge metrics demonstrate that Thinking-RFT traces are more logical and faithful (9.4/10), with only a marginal tradeoff in efficiency. Attention visualization reveals that RFT-trained models reliably attend to human-identified causal anchor sentences when generating answers—covering the decisive narrative events in 89% of cases—while SFT and No-Thinking-RFT display more diffuse and irrelevant attention. Error analysis shows that most failures are due to agent-of-mind/target role confusion, especially in higher-order queries, further motivating ToM-specific reward design.

Implications and Future Directions

Theoretical implications include a refined understanding that explicit reasoning is required for ToM acquisition, contradicting analyses that attribute SFT or RL gains to model capacity or dataset size alone. Practically, the results validate GRPO rule-based RL as a scalable, annotation-efficient post-training mechanism for ToM, with immediate application in scenarios requiring socially and causally coherent foundation models. Future directions involve developing higher-order shortcut-free ToM datasets, extending RL post-training to functional ToM tasks (beyond QA), and exploring richer reward functions for even finer-grained mental-state reasoning.

Conclusion

This paper establishes that reinforcement learning with explicit reasoning traces (Thinking-RFT), when applied to shortcut-free ToM benchmarks, reliably elicits robust, generalized, and causally grounded Theory of Mind abilities in foundation models. Dataset shortcut auditing is essential—otherwise, apparent gains are illusory. The work not only provides reproducible, efficient post-training protocols but also advances the design principles for future ToM datasets and post-training strategies, fostering more human-aligned AI systems.