BetaZero: Belief-State Planning for Long-Horizon POMDPs using Learned Approximations

TL;DR

BetaZero tackles long-horizon planning under partial observability by fusing online belief-state Monte Carlo tree search with offline neural networks that imitate policy and value functions. The method learns a two-head network (policy and value) and uses a parametric belief input, a Q-weighted visit-count policy, and progressive widening to manage stochastic belief transitions, enabling efficient planning with limited online budget. Empirical results show BetaZero outperforming several state-of-the-art POMDP solvers on benchmarks like LightDark, RockSample, and Mineral Exploration, demonstrating the value of offline experience to guide online decisions. The work advances scalable, heuristic-free planning for high-dimensional POMDPs and illustrates practical benefits for information-gathering and decision-making under uncertainty in complex domains.

Abstract

Real-world planning problems, including autonomous driving and sustainable energy applications like carbon storage and resource exploration, have recently been modeled as partially observable Markov decision processes (POMDPs) and solved using approximate methods. To solve high-dimensional POMDPs in practice, state-of-the-art methods use online planning with problem-specific heuristics to reduce planning horizons and make the problems tractable. Algorithms that learn approximations to replace heuristics have recently found success in large-scale fully observable domains. The key insight is the combination of online Monte Carlo tree search with offline neural network approximations of the optimal policy and value function. In this work, we bring this insight to partially observable domains and propose BetaZero, a belief-state planning algorithm for high-dimensional POMDPs. BetaZero learns offline approximations that replace heuristics to enable online decision making in long-horizon problems. We address several challenges inherent in large-scale partially observable domains; namely challenges of transitioning in stochastic environments, prioritizing action branching with a limited search budget, and representing beliefs as input to the network. To formalize the use of all limited search information, we train against a novel $Q$-weighted visit counts policy. We test BetaZero on various well-established POMDP benchmarks found in the literature and a real-world problem of critical mineral exploration. Experiments show that BetaZero outperforms state-of-the-art POMDP solvers on a variety of tasks.

Figures (23)

• Figure 1: The BetaZero POMDP policy iteration algorithm.
• Figure 2: The four stages of MCTS belief-state planning in BetaZero using the value $V_\theta$ and policy $P_\theta$ network heads (the policy evaluation step in \ref{['fig:betazero-alg']}).
• Figure 3: An illustrative example of when collecting policy data based purely on visit counts (left) or $Q$-values (middle) would perform worse than weighting the visit counts based on $Q$-values (right). This is useful when using a small MCTS budget with high exploration. Using both the $Q$-values and visit counts, we incorporate both what the tree search focused on and the values it found.
• Figure 4: POMDP space dimensions.
• Figure 5: LightDark$(10)$ value and policy plots over belief mean and std. High uncertainty (horizontal axis) makes the agent localize up near $y=10$, then moves down and stops at the origin.