HiVAE: Hierarchical Latent Variables for Scalable Theory of Mind

Nigel Doering; Rahath Malladi; Arshia Sangwan; David Danks; Tauhidur Rahman

cs.LGFeb 18, 2026

HiVAE: Hierarchical Latent Variables for Scalable Theory of Mind

Authors: Nigel Doering, Rahath Malladi, Arshia Sangwan, David Danks, Tauhidur Rahman

Organizations: University of California San Diego, School of Computing, Information, and Data Sciences · New York University · University of Virginia

Abstract

Theory of mind (ToM) enables AI systems to infer agents' hidden goals and mental states, but existing approaches focus mainly on small human understandable gridworld spaces. We introduce HiVAE, a hierarchical variational architecture that scales ToM reasoning to realistic spatiotemporal domains. Inspired by the belief-desire-intention structure of human cognition, our three-level VAE hierarchy achieves substantial performance improvements on a 3,185-node campus navigation task. However, we identify a critical limitation: while our hierarchical structure improves prediction, learned latent representations lack explicit grounding to actual mental states. We propose self-supervised alignment strategies and present this work to solicit community feedback on grounding approaches.

Figures & tables

Figure 1: HiVAE first encodes the agent’s partial trajectory and the environment graph into a unified latent representation. This representation feeds a hierarchical mind-state module, sequentially inferring beliefs, desires, and intentions, which then drives the goal predictor to output a probability distribution over all possible goals.

Figure 2: Overall performance of models on goal prediction averaged across all trajectories. Lower is better.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Model	25%	50%	75%	95%
Bayesian ToM (BToM)	0.9555	0.9117	0.7400	0.4204
Extended BToM	0.8899	0.8292	0.6865	0.3927
GRU	1.5358	1.5853	1.4555	0.5399
LSTM	1.1886	1.1606	1.1171	0.4211
ToMNet	0.8583	0.7453	0.7215	0.7176
HiVAE (Ours)	0.1042	0.1036	0.1023	0.1023

Appendix

Table 1: Brier Score on Pedestrian Dataset at Different Path Completion Percentages - Experiment 1. Best results shown in bold.

Figure 3: Probability of False Goal on Pedestrian Dataset at Different Path Completion Percentages Leading up to the False Goal - Experiment 2. Lower is better.

Model	25%	50%	75%	95%
Bayesian ToM (BToM)	0.0644	0.0532	0.3755	0.8350
Extended BToM	0.0904	0.2814	0.5871	1.0222
GRU	-0.0184	0.0170	0.0290	0.6185
LSTM	0.0104	-0.0071	0.0373	0.2436
ToMNet	-0.0182	-0.0342	-0.0343	-0.0275
HiVAE (Ours)	0.0205	0.0309	0.0307	0.0307

Appendix

Table 2: Difference in Brier Score Between New and Original Pedestrian Dataset at Different Path Completion Percentages - Experiment 3. Closer to zero indicates same performance. Negative indicates model is performing better on new data.

Figure 4: Overall performance of models on goal prediction averaged across all trajectories for both original and new pedestrian datasets. Lower is better.

Theory of Mind

Theory of Mind studies how people and AI infer others’ beliefs, goals, and knowledge to predict behavior and coordinate interaction.

See all 92 papers

MindZero: Learning Online Mental Reasoning With Zero Annotations
May 29, 2026Shunchi Zhang, Jin Lu, Chuanyang Jin +3Self-Supervised Learning Theory of Mind
Johns Hopkins University · Peking University
Xiaomi OneVL: One-Step Latent Reasoning and Planning with Vision-Language Explanation
Apr 20, 2026Jinghui Lu, Jiayi Guan, Zhijian Huang +47Trajectory Generation VLMs for Autonomous Driving
Reasoning Without Inference Cost: Latent Semantic Scaffolding for Robot VLA Policies
Sep 4, 2026Andrew Ting Yan Li, Zhuo Li, Zhelin Yang +3Reasoning for VLA Models Vision-Language-Action Models
Collaborative and Versatile Robots (CLOVER) Laboratory, T-Stone Robotics Institute, The Chinese University of Hong Kong, Hong Kong

May 29, 2026J/K move · Enter open · S save

MindZero: Learning Online Mental Reasoning With Zero Annotations

Shunchi Zhang, Jin Lu, Chuanyang Jin, Yichao Zhou+2

Johns Hopkins University · Peking University

Self-Supervised Learning Theory of Mind Multimodal Large Language Models RL for VLMs

Open paper AskPDF arXiv

Effective real-world assistance requires AI agents with robust Theory of Mind (ToM): inferring human mental states from their behavior. Despite recent advances, several key challenges remain, including (1) online inference with robust uncertainty updates over multiple hypotheses; (2) efficient reasoning suitable for real-time assistance; and (3) the lack of ground-truth mental state annotations in real-world domains. We address these challenges by introducing MindZero, a self-supervised reinforcement learning framework that trains multimodal large language models (MLLMs) for efficient and robust online mental reasoning. During training, the model is rewarded for generating mental state hypotheses that maximize the likelihood of observed actions estimated by a planner, similar to model-based ToM reasoning. This method thus eliminates the need for explicit mental state annotations. After training, MindZero internalizes model-based reasoning into fast single-pass inference. We evaluate MindZero against baselines across challenging mental reasoning and AI assistance tasks in gridworld and household domains. We found that LLMs alone are insufficient; model-based methods improve accuracy but are slow, costly, and limited by backbone MLLM capacity. In contrast, MindZero enhances MLLMs' intrinsic ToM ability and significantly outperforms model-based methods in both accuracy and efficiency, showing that mental reasoning can be effectively learned as a self-supervised skill.

Abstract

Theory of Mind

Explore similar work

MindZero: Learning Online Mental Reasoning With Zero Annotations

Xiaomi OneVL: One-Step Latent Reasoning and Planning with Vision-Language Explanation

Reasoning Without Inference Cost: Latent Semantic Scaffolding for Robot VLA Policies