Source-linked AI summary
Belief Embedding Tree Search
Christopher Solinas, Radovan Haluška, Martin Schmid, Michael Buro, Nathan R. Sturtevant, Michael Bowling
TL;DR
POMDP planning requires approximating beliefs and values, but particle sets and hand-chosen summaries can lose decision-relevant information. BETS plans over generative distribution embeddings updated through the dynamics and trains its belief, policy, and value models together. It is competitive or leading across six configurations, with large margins over baselines on the highest-dimensional beliefs, even without search.
Problem
Particle-based belief approximations can suffer impoverishment, while learned planners may summarize particles in ways that discard belief structure needed by policies and values.
Source: §1 · p. 2Method
BETS searches over generative distribution embeddings, updates them through POMDP dynamics, and trains the belief model alongside policy and value networks using realized hidden states as targets.
Source: §1 · p. 2Results
Takeaways & Limitations
BETS shows that a learned belief embedding can replace a particle set inside a planner, and the same approach can extend to other learning and planning algorithms for partially observable environments.
Source: §7 · p. 10BETS requires domain-specific training and a simulator, and its paired root-action sampling divides the budget across actions as the action space grows.
Source: §6 · p. 9
Abstract
from arXiv · showhide
Beliefs in a Partially Observable Markov Decision Process (POMDP) are posterior distributions over the hidden state given the agent's history of actions and observations. They summarize everything the history says about the state, so planning in a POMDP is planning in a Markov decision process (MDP) over beliefs. Belief MDPs have continuous state spaces even for finite POMDPs, and exact planning in them is intractable in general. Planners must therefore approximate beliefs and estimate values. Monte Carlo planners do both inside a search tree, with an approximate belief at every node and a value estimate at every leaf. Approximations usually consist of sets of sampled states, with leaf values coming from rollouts, hand-designed bounds, or learned networks. Policies and values computed from these approximations are limited by their quality, and prior methods with learned networks rely on hand-chosen summaries of the sampled states. A belief representation should be computable from the parent node's representation through the POMDP's dynamics, let the search sample from it, and be learned so that it provides the information that policies and value functions need. We present Belief Embedding Tree Search (BETS), a Monte Carlo planner that meets these requirements by representing each node's belief with a generative distribution embedding (GDE) while updating it through the POMDP dynamics. BETS trains its belief model alongside its policy and value networks through experience, with the realized hidden state of the POMDP as the belief model's training target. We evaluate BETS on six configurations of RockSample, LightDark, and PocMan against BetaZero and three planners without learned components. BETS is competitive or leading on all six, while every baseline falls far behind on at least one. Where beliefs are highest-dimensional, it outperforms them by a large margin, even without search.
1 INTRODUCTION
POMDP planners must approximate beliefs and values, but particle-based beliefs can degrade and learned networks may rely on summaries that discard decision-relevant structure. BETS searches over learned belief embeddings and is competitive or leading across six evaluated configurations, with large gains on the highest-dimensional beliefs.
1 INTRODUCTION
- Particle sets need rapidly growing numbers of states for accurate approximations and can collapse onto a few distinct states through repeated resampling.Leaf-value accuracy also depends on rollouts or bounds that are domain-designed or loose.Source: §1 · pp. 1–2
- BetaZero's policy and value networks use particle-set means and standard deviations, which discard belief support and variable dependencies.BetaZero also retains a particle set at every tree node, with the same particle-filter weaknesses.Source: §1 · p. 2
- Neural Bayesian Filtering represents beliefs as generative distribution embeddings and regenerates particles at each update, preventing resampling impoverishment from accumulating across steps.Its conditional generator samples from the represented belief, allowing search to generate states as needed.Source: §1 · p. 2
- BETS stores a generative distribution embedding at each search node, updates it through the POMDP dynamics, and uses it for both state sampling and policy and value conditioning.The belief model and both networks train together on BETS-collected episodes, using the realized hidden state as the belief model's training target.Source: §1 · p. 2
- BETS is competitive or leading across six configurations, and on RockSample(20,20) and PocMan it outperforms every baseline by a large margin, even without search.PocMan value networks also predict returns more accurately from embeddings than from particle-filter summary statistics with 16× as many particles.Source: §1 · p. 2
2 BACKGROUND
A POMDP's history induces a belief over hidden states, and planning can be expressed as an MDP over those beliefs. Because exact updates and particle approximations have limitations, planners use learned or sampled representations to guide search.
2 BACKGROUND
- A belief is a posterior over hidden states given history, and it is a sufficient statistic for policies and action values in the belief MDP.The belief update is exact but intractable for large or continuous state spaces.Source: §2 · pp. 2–3
2.1 APPROXIMATING BELIEFS
- Particle filters update weighted states using transition dynamics and observation likelihoods, but repeated resampling can collapse support and require exponentially many particles as state dimension grows.Neural Bayesian Filtering instead regenerates particles from a generative distribution embedding at each update, so resampling collapse does not accumulate across steps.Source: §2.1 · p. 3
2.2 PLANNING WITH PARTIAL OBSERVABILITY
- POMCP and POMCPOW use node particle sets, while DESPOT and AdaOPS search sampled scenarios with value bounds; BetaZero adds learned networks conditioned on particle-feature means and standard deviations.Two beliefs with identical per-variable statistics can have disjoint support and different action values, so such summaries can give networks indistinguishable inputs for decision-relevant beliefs.Source: §2.2 · pp. 3–4
- BETS uses paired root-action rounds with a shared sampled state, selects interior actions with PUCT, and charges one NBF update for expansion.The figure caption describes the root sampling and expansion cost; the planning notes specify shared random numbers and a mixed prior over actions.Source: Figure 1 · p. 4, §2.2 · p. 4
3 BELIEF EMBEDDING TREE SEARCH
BETS plans over belief embeddings, using them as the shared representation for tree updates and policy and value networks. Its search combines NBF updates with a variable-coupling representation and turns root action estimates into an improved policy.
3 BELIEF EMBEDDING TREE SEARCH
- Each tree node stores a belief embedding; expanding a leaf applies one NBF update, and policy and value networks guide search using that representation.Leaf values come from the value network rather than a rollout.Source: §3 · p. 4, §3.1 · p. 5
3.1 SEARCH
- Root actions are evaluated in paired rounds, while interior traversal uses PUCT and a policy prior mixed with uniform mass to preserve exploration.Every root action is selected once per round, pairing actions on the same sampled state and random numbers.Source: §3.1 · pp. 4–5
3.2 ACTION SELECTION
- After search, BETS converts root action values into a KL-regularized policy improvement, which also supplies training and value-bootstrap targets.The temperature controls how far the improved policy moves from its prior.Source: §3.2 · p. 5
3.3 VARIABLE-COUPLING EMBEDDINGS
- The variable-coupling stream pools pairwise variable features into per-variable summaries and combines them with the pooled embedding for θ.Attention biased by the variable-by-variable co-occurrence matrix produces the coupling summaries; both streams average over samples.Source: §3.3 · p. 5, Figure 2 · p. 6
4 TRAINING
BETS jointly trains its belief, policy, and value models on collected experience. The belief model learns from realized hidden states, while the policy and value networks use targets derived from search.
4 TRAINING
- BETS jointly trains its belief model and policy and value networks on episodes collected by the planner, alternating collection rounds with gradient steps.The policy and value losses use stored embeddings and do not update the encoder; the GDE is trained on its belief loss alone.Source: §4 · p. 6
4.1 BELIEF MODEL
- The GDE minimizes realized-state negative log-likelihood, using the environment’s next hidden state as its target even though the agent does not observe it.Training regenerates NBF particle sets from the current generator; each realized state is a sample from the true posterior, and the set posterior has the lowest expected log loss among beliefs computed from that set.Source: §4.1 · pp. 6–7
4.2 POLICY AND VALUE TARGETS
- The policy target is the search-improved policy, and the value target is a λ-return bootstrapped from root values, trained as a two-hot distribution over fixed bins.Terminal states use a zero bootstrap; leaf evaluation applies the inverse transform to the predicted bin mean.Source: §4.2 · p. 7
5 EXPERIMENTS
Across six configurations of RockSample, LightDark, and PocMan, BETS is competitive or leading, with its strongest results in high-dimensional settings. The experiments also compare belief representations and test the variable-coupling stream.
5 EXPERIMENTS
- The six test configurations cover two RockSample sizes, two LightDark settings, and medium and large PocMan mazes.Beliefs range from independent Bernoulli variables in RockSample to a one-dimensional position density in LightDark and joint ghost position-heading distributions in PocMan.Source: §5 · p. 7
- BETS is compared with BetaZero, AdaOPS, DESPOT, and POMCPOW using equal simulation and particle counts at evaluation.BETS and BetaZero use three training seeds and 300 episodes at each of three evaluation seeds; the other planners use three evaluation seeds with 300 episodes each.Source: §5 · p. 7
5.1 RESULTS
- BETS earns the highest return on RockSample(20,20) and both PocMan mazes, while trailing AdaOPS on RockSample(15,15) and BetaZero on LightDark(10).At RockSample(20,20), AdaOPS and DESPOT use a fixed upper bound because computing their heuristic bounds is infeasible, reducing discounted return.Source: Table 1 · p. 8, §5.1 · p. 8
- On RockSample(20,20), authors attribute BetaZero’s gap to particle impoverishment across 2^20 configurations, despite per-rock means being sufficient statistics.With deterministic transitions, repeated resampling duplicates surviving particles; BETS regenerates particles from its embedding at each update.Source: §5.1 · p. 8
- On PocMan, moment summaries blur multimodal ghost locations and discard dependence between positions and headings, which the coupling stream models.The authors attribute BetaZero’s gap on PocMan to these losses of structure.Source: §5.1 · p. 8
- BetaZero matches BETS on LightDark(5) and leads on LightDark(10); policy-only BETS matches search on LightDark(5) and exceeds it on LightDark(10).The paper notes that the one-dimensional belief is well approximated by particles in this domain.Source: §5.1 · p. 8
- On PocMan(M), Pooled earns less return than VC with both search and policy-only, despite more than doubling model capacity.With search, Pooled clears the maze in fewer than half as many episodes, and removing the coupling stream costs more discounted return.Source: §5.1 · p. 8
- For both models and every training-set size, embeddings yield significantly more accurate PocMan(M) value predictions than moments, even though moments use 16 times as many particles.The value-prediction comparison reports held-out R2; combining pooled and coupling summaries also improves estimates over the pooled embedding alone.Source: §5.1 · p. 9, Table 2 · p. 9
6 LIMITATIONS
BETS has practical constraints in large action spaces and requires domain-specific training. Its training also depends on access to a simulator that exposes realized hidden states.
6 LIMITATIONS
- Root simulations divide the budget evenly across actions, so each action receives a smaller share as the action space grows.The paper suggests sequential halving as a way to extend BETS to large action spaces.Source: §6 · p. 9
- BETS must be trained on each domain and requires a simulator exposing realized hidden states, which may not always be available.Search takes roughly K times as long per decision as the policy network alone; training can add overhead without benefit over simpler belief models on LightDark.Source: §6 · p. 9
7 CONCLUSION
BETS represents tree-node beliefs with generative distribution embeddings that support sampling and condition the policy and value networks. Experiments show broad performance, while the authors identify possible extensions beyond tree search.
7 CONCLUSION
- BETS is competitive or leading on every tested configuration, and without search it beats every baseline by a large margin on the highest-dimensional RockSample(20,20) and PocMan settings.The conclusion also reports that experiments suggest per-variable coupling summaries improve the GDE’s representation on PocMan.Source: §7 · p. 9
- Embedding beliefs instead of retaining particle sets is presented as applicable to other algorithms that learn or plan on beliefs.The authors also describe extending the approach to public beliefs in games by incorporating player policies into the NBF update.Source: §7 · p. 10
AI USE STATEMENT
- Large language models assisted with writing and coding, and the authors revised and verified that work.Source: AI USE STATEMENT · p. 10
REPRODUCIBILITY STATEMENT
- BETS source code is supplementary material and is planned for public release on GitHub upon publication; Appendix B lists BETS and baseline hyperparameters.Source: REPRODUCIBILITY STATEMENT · p. 10
A DOMAINS
The evaluation uses two RockSample configurations, two LightDark configurations, and two PocMan mazes. Their hidden-state structure and observations differ, with PocMan networks also receiving features of the observed state.
RockSample
- RockSample uses a rover on a grid with independently good or bad rocks, noisy distance-dependent checks, and deterministic movement; its two settings have 15 or 20 rocks and horizons of 100.The rover position is fully observable, and the settings are RockSample(15,15) and RockSample(20,20).Source: A DOMAINS · p. 12
LightDark
- LightDark tests continuous one-dimensional localization with unobserved position and distance-dependent Gaussian observation noise, placing the light at 5 or 10 across its two configurations.Both start positions are drawn from N(2, 32), and both use a 100-step horizon.Source: A DOMAINS · p. 12
PocMan
- PocMan tests 9×9 and 17×19 mazes with hidden ghost positions and headings and a random pellet field; the generator does not model the pellet field, and filters receive the observed pellet mask.The mazes have 59 and 189 reachable cells, respectively; horizons are 100 and 250.Source: A DOMAINS · p. 13
Policy and Value Network Inputs
- Policy and value networks use the embedding plus elapsed-horizon features, with PocMan also supplying remaining-pellet fraction, power-pill timer, and observed pellet mask.Source: A DOMAINS · p. 13
B HYPERPARAMETERS
The appendix describes configuration-specific settings for BETS, BetaZero, and three online planners. These planners use shared simulation and particle counts relative to BETS, with domain-specific settings and rollout policies.
BETS
- BETS settings are shared between training and evaluation, with annealed values changing linearly across the anneal window; LightDark uses observation progressive widening because every sampled observation is new.Action nodes otherwise keep at most 16 observation children.Source: B HYPERPARAMETERS · p. 13
BetaZero
- BetaZero matches BETS’s simulations and particles, uses default RockSample and LightDark settings, and uses a bootstrap particle filter on PocMan.Its policy and value networks condition on summary statistics of player and ghost positions and the pellet and power-pill grids.Source: B HYPERPARAMETERS · p. 13
Online planners
- AdaOPS, DESPOT, and POMCPOW match BETS’s simulations and particles, use cited settings on RockSample and LightDark, and use the preferred-actions rollout policy on PocMan.The described exit policy moves the rover east without sampling rocks.Source: B HYPERPARAMETERS · p. 13
C.1 AUTOREGRESSIVE TRANSFORMER
The autoregressive generator conditions each variable’s output on preceding sample values and on the belief embedding. Its conditioning can include both pooled and variable-specific coupling information; LightDark instead uses a flow-matching generator.
C.1 AUTOREGRESSIVE TRANSFORMER
- The autoregressive generator uses one token per variable, and token j emits logits for pψ(xj | x<j, θ) from the preceding value and a conditioning vector.With variable coupling, each token receives the shared pooled embedding and the coupling summary for its corresponding variable.Source: §C.1 · pp. 13–14, Figure 3 · p. 15
- The appendix identifies configuration-specific settings tables for BETS, BetaZero, and AdaOPS, DESPOT, and POMCPOW.The BETS table notes that ranges are annealed over an anneal window; the supplied passages do not report the table values.Source: Table 3 · p. 14, Table 4 · p. 14, Table 5 · p. 15
C.2 FLOW-MATCHING
- LightDark uses conditional flow matching: a three-hidden-layer, width-64 MLP transports noise into a belief sample in 50 integration steps, with training regression at the realized state.Its embedding uses the pooled stream.Source: §C.2 · p. 15
D WALL-CLOCK TIME
The policy network alone takes under 25 ms per decision on every configuration, while search takes roughly K times as long; timing is averaged over 100 episodes.
D WALL-CLOCK TIME
- Under 25 ms per decision on every configuration, the policy network alone also earns more return than every baseline on RockSample(20,20) and both PocMan mazes.Source: D WALL-CLOCK TIME · p. 15
- Search takes roughly K times as long per decision as policy-only inference because each simulation runs the NBF update, policy network, and value network once.Source: D WALL-CLOCK TIME · p. 15
- Timing averages 100 episodes, with standard errors below 3% of the mean; Table 6 reports milliseconds per decision on an NVIDIA GeForce RTX 5080.Source: D WALL-CLOCK TIME · p. 15, Table 6 · p. 16
- K is rounded up to a multiple of |A|, ensuring every search round selects every root action once.Source: D WALL-CLOCK TIME · p. 16