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.
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.
Chain-of-Thought (CoT) reasoning has become a powerful driver of trajectory prediction in VLA-based autonomous driving, yet its autoregressive nature imposes a latency cost that is prohibitive for real-time deployment. Latent CoT methods attempt to close this gap by compressing reasoning into continuous hidden states, but consistently fall short of their explicit counterparts. We suggest that this is due to purely linguistic latent representations compressing a symbolic abstraction of the world, rather than the causal dynamics that actually govern driving. Thus, we present OneVL (One-step latent reasoning and planning with Vision-Language explanations), a unified VLA and World Model framework that routes reasoning through compact latent tokens supervised by dual auxiliary decoders. Alongside a language decoder that reconstructs text CoT, we introduce a visual world model decoder that predicts future-frame tokens, forcing the latent space to internalize the causal dynamics of road geometry, agent motion, and environmental change. A three-stage training pipeline progressively aligns these latents with trajectory, language, and visual objectives, ensuring stable joint optimization. In inference, the auxiliary decoders are discarded, and all latent tokens are prefilled in a single parallel pass, matching the speed of answer-only prediction. Across four benchmarks, OneVL becomes the first latent CoT method to surpass explicit CoT, delivering superior accuracy at answer-only latency. These results show that with world model supervision, latent CoT produces more generalizable representations than verbose token-by-token reasoning. Code has been open-sourced to the community. Project Page: https://xiaomi-embodied-intelligence.github.io/OneVL
Vision-language-action (VLA) models are trained by imitation and capture what action to take but not why; adding causal reasoning improves manipulation, but current methods pay for it at inference time - generating reasoning tokens or rolling out predicted future states at every step, a cost that compounds over long horizons. We ask whether this benefit can instead be captured during training and discarded before deployment. We introduce Latent Semantic Scaffolding (LSS), an auxiliary loss applied during human-demonstration pretraining that aligns a VLA's action-token representations to text embeddings of physical-reasoning rationales through a small projection head. The head is dropped at inference, leaving the unmodified base policy with zero added cost. Our central finding concerns alignment granularity: aligning each action token to the rationale of its own manipulation phase (Dense LSS) rather than to a single pooled episode-level embedding (Pooled LSS) yields representations that transfer markedly better to held-out tasks. Dense LSS attains both the best in-distribution success and the best transfer to tasks unseen during alignment, whereas pooled alignment over-specializes to the training task. A representational probe shows Dense LSS induces roughly twice the per-phase separability in the backbone, supporting that phase-local alignment is the operative mechanism.
Andrew Ting Yan Li, Zhuo Li, Zhelin Yang +3
Collaborative and Versatile Robots (CLOVER) Laboratory, T-Stone Robotics Institute, The Chinese University of Hong Kong, Hong Kong