GammaZero: Learning To Guide POMDP Belief Space Search With Graph Representations

TL;DR

GammaZero tackles planning in POMDPs by converting belief states into action-centric graphs and learning value and policy estimates with a graph neural network trained offline on small, tractable problems. This graph-based belief representation, coupled with MCTS guided by learned priors and values, enables zero-shot generalization to problem instances 2–4× larger while preserving solution quality and reducing search effort. Compared to BetaZero, GammaZero matches performance on same-size problems and uniquely generalizes across scales, demonstrated in RockSample and LightDark domains where a model trained on small grids solves much larger instances with competitive returns. The work highlights structural invariance, local computation, and compositional reasoning as key properties that support transfer, offering a scalable approach to belief-space planning under uncertainty with practical implications for real-world POMDP applications.

Abstract

We introduce an action-centric graph representation framework for learning to guide planning in Partially Observable Markov Decision Processes (POMDPs). Unlike existing approaches that require domain-specific neural architectures and struggle with scalability, GammaZero leverages a unified graph-based belief representation that enables generalization across problem sizes within a domain. Our key insight is that belief states can be systematically transformed into action-centric graphs where structural patterns learned on small problems transfer to larger instances. We employ a graph neural network with a decoder architecture to learn value functions and policies from expert demonstrations on computationally tractable problems, then apply these learned heuristics to guide Monte Carlo tree search on larger problems. Experimental results on standard POMDP benchmarks demonstrate that GammaZero achieves comparable performance to BetaZero when trained and tested on the same-sized problems, while uniquely enabling zero-shot generalization to problems 2-4 times larger than those seen during training, maintaining solution quality with reduced search requirements.