English

HiVAE: Hierarchical Latent Variables for Scalable Theory of Mind

Machine Learning 2026-02-20 v1 Artificial Intelligence

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.

Keywords

Cite

@article{arxiv.2602.16826,
  title  = {HiVAE: Hierarchical Latent Variables for Scalable Theory of Mind},
  author = {Nigel Doering and Rahath Malladi and Arshia Sangwan and David Danks and Tauhidur Rahman},
  journal= {arXiv preprint arXiv:2602.16826},
  year   = {2026}
}

Comments

Accepted at the Workshop on Theory of Mind for AI (ToM4AI) at the 40th AAAI Conference on Artificial Intelligence (AAAI-26), Singapore, 2026