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 No-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. Project code is available at: https://github.com/jkz-338/Robust-ToM-RL.git
Lay Summary
As AI systems become more common in everyday life, they need to better understand people: what someone knows, believes, wants, or intends. This ability is often called “Theory of Mind.” For example, if one person moves an object while another person is away, an AI should understand that the absent person may still believe the object is in its old location. This paper studies how to improve this kind of social reasoning in AI models. We first find an important problem: many existing tests for Theory of Mind can be solved by simple shortcuts rather than real understanding. For instance, a model may learn to pick an answer based on repeated patterns in the test, without actually reasoning about what each person knows or believes. This can make AI systems look more capable than they really are. To address this, we identify tests that are less affected by such shortcuts and study how different training methods perform on them. We find that training models to explicitly work through their reasoning, while rewarding correct final answers, leads to stronger and more reliable Theory of Mind behavior than standard training. These models perform better on harder questions, new settings, image-based scenarios, and stories where small details are changed. Overall, our work shows that improving social reasoning in AI requires both cleaner training data and training methods that encourage models to reason from the right evidence rather than exploit superficial patterns.