BETS: Cracking POMDP Planning with Belief Embedding Tree Search

In the field of AI planning, the Partially Observable Markov Decision Process (POMDP) has always been a tough nut to crack. When an agent cannot fully observe the environment state, it must rely on the history of actions and observations to infer the hidden state — and this inference result is what we call a "belief." How to efficiently search in continuous belief space has become the core bottleneck limiting the practical application of POMDPs.

Why POMDP Planning Is So Difficult

A belief in a POMDP is the posterior distribution over hidden states given the agent's action and observation history. It condenses all information about the state contained in the history, so POMDP planning is essentially planning in a Markov Decision Process (MDP) over belief space. The problem is that even when the original POMDP's state space is finite, its corresponding belief MDP has a continuous state space, making exact planning infeasible in general.

This forces planners to do two things: approximate beliefs and estimate value functions. Monte Carlo planners accomplish both simultaneously within a search tree — each node maintains an approximate belief, and each leaf node provides a value estimate. The traditional approach uses a set of sampled states to approximate beliefs, while leaf values come from rollouts, hand-crafted bounds, or learned networks.

The effectiveness of these approximation methods is limited by their quality. Especially when using learned networks, prior methods rely on manually chosen summaries of sampled states, which both limits expressiveness and increases design burden.

Belief MDPs and Continuous State Spaces

A POMDP can be formally transformed into an MDP defined over the belief simplex, called a belief MDP. Its states are beliefs (i.e., posterior distributions over hidden states), its actions are the same as the original POMDP, but the transition function is determined by the Bayes update rule. The key difficulty is that even if the original POMDP has only finitely many states, its belief space is a continuous simplex. For example, the belief over N hidden states is an N-dimensional probability vector lying within an (N-1)-dimensional simplex. Solving for the optimal value function over this continuous space generally requires infinite-dimensional representations, and exact dynamic programming is infeasible. Therefore, practical algorithms must apply finite parameterization or discretization to belief space while simultaneously approximating the value function, which constitutes the fundamental challenge of POMDP planning.

Monte Carlo Tree Search and Rollouts

Monte Carlo Tree Search (MCTS) is an online planning method that estimates action values by repeatedly sampling within a decision tree. In a POMDP, each search node corresponds to a belief, and edges correspond to actions and observations. Since exhaustive enumeration is impossible, MCTS starts from the current belief, selects actions according to some strategy (e.g., UCT), samples observations, and gradually expands the tree. Leaf node values are typically estimated by rollouts: starting from that node, a default policy simulates several steps until termination or a depth limit, and the average return is taken as the value estimate. The quality and computational cost of rollouts directly affect planning performance. BETS maintains learnable belief embeddings in the search tree precisely to replace the coarse belief approximations based on sampled state sets in traditional MCTS, thereby improving the expressiveness of node representations.

BETS's Core Idea: Generative Distribution Embeddings

The research team proposed Belief Embedding Tree Search (BETS), a Monte Carlo planner that satisfies three key requirements: belief representations can be computed from parent nodes via POMDP dynamics, the search can sample from them, and they can be learned to provide the information needed by policy and value functions.

BETS represents each node's belief using a Generative Distribution Embedding (GDE), updated through POMDP dynamics. Unlike prior methods that rely on hand-crafted summaries, GDE makes the belief representation itself a trainable parametric model.

The training approach is also quite clever: BETS simultaneously trains the belief model, policy network, and value network as experience accumulates, where the belief model's training objective is the actual hidden state realized by the POMDP. This means the belief model directly learns the ability to "recover the true state from observation history," rather than fitting some indirect proxy objective.

Generative Distribution Embeddings and Co-training

A Generative Distribution Embedding (GDE) is a method for representing probability distributions using parametric generative models. Rather than directly storing a set of sampled states, a GDE learns a mapping from a low-dimensional latent variable to the state space, while using an encoder to map observation history to a latent variable distribution, thereby enabling the generation of arbitrarily many candidate states and computation of densities. In BETS, each tree node's belief is represented as an instance of a GDE, and can be computed from the parent node's belief through POMDP dynamics (prediction-update). During training, BETS jointly optimizes three components: the belief model (GDE), the policy network, and the value network. The belief model's loss minimizes the discrepancy between the generated distribution and the true hidden state, commonly using negative log-likelihood or maximum mean discrepancy. Since the belief representation is differentiable, gradients from the policy and value networks can flow back to the belief model, enabling end-to-end co-adaptation.

Experimental Performance: Comprehensive Testing Across Six Configurations

The study evaluated BETS on six configurations across three classic POMDP environments: RockSample, LightDark, and PocMan. Baselines included BetaZero and three planners without learning components.

Results show that BETS is either competitive or leading on all six configurations, while every baseline falls far behind on at least one configuration. Most notably, in the configuration with the highest belief dimensionality, BETS substantially outperforms all compared methods, maintaining its advantage even without search.

The significance of this result deserves careful consideration. High-dimensional belief spaces are precisely where traditional approximation methods are most prone to failure — the coverage capability of sampled state sets declines exponentially with dimensionality, and hand-crafted summaries struggle even more to capture key information. BETS bypasses this dilemma through learnable embedding representations, keeping belief representations informative even in high-dimensional settings.

Implications for the POMDP Planning Paradigm

The value of BETS goes beyond performance numbers. It proposes an idea: belief representations should not be a pre-designed fixed form, but rather learned in coordination with policy and value functions, and naturally propagated through the search tree via dynamics.

This aligns with the current AI trend of integrating "end-to-end learning + structured search." Pure end-to-end methods lack the precision of search, while pure traditional planners are constrained by hand-crafted approximations. BETS finds a scalable balance point between the two.

Of course, BETS still has limitations. It relies on the accessibility of POMDP dynamics to update belief embeddings, and how it applies in scenarios where dynamics are unknown or difficult to model remains to be explored. Additionally, the belief model's training objective requires the environment to provide true hidden states as supervision signals, which may not be available in some practical applications.

Regardless, BETS provides a compelling new baseline for POMDP planning. When belief spaces become high-dimensional and complex, learnable embedding representations may be a more reliable approximation tool than sampled sets.

One-line Summary

BETS uses learnable generative distribution embeddings to represent POMDP beliefs, co-training policy and value functions within the search tree, leading by a wide margin in high-dimensional settings.

Paragraph Summary

POMDP planning is notoriously difficult due to continuous, high-dimensional belief spaces. Traditional Monte Carlo planners rely on sampled state sets or hand-crafted summaries to approximate beliefs, limiting expressiveness. This paper proposes Belief Embedding Tree Search (BETS), which uses Generative Distribution Embeddings (GDEs) to represent each search node's belief, updated from parent nodes via POMDP dynamics. BETS simultaneously trains the belief model, policy network, and value network as experience accumulates, with the belief model directly supervised by true hidden states to learn state recovery from observation history. Across six configurations in RockSample, LightDark, and PocMan, BETS matches or leads all baselines, substantially outperforming them in high-dimensional belief configurations even without search. This offers a scalable new direction for POMDP planning, though it still depends on accessible dynamics and true state supervision.