Paper deep dive
Belief Embedding Tree Search
Christopher Solinas, Radovan Haluška, Martin Schmid, Michael Buro, Nathan R. Sturtevant, Michael Bowling
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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.
Tags
Links
- Source: https://arxiv.org/abs/2610.15046v1
- Canonical: https://arxiv.org/abs/2610.15046v1
Trouble viewing inline? Open PDF directly →
Full Text
63,444 characters extracted from source content.
Expand or collapse full text
Belief Embedding Tree Search Christopher Solinas Affiliation: University of Alberta Email: solinas@ualberta.ca Radovan Haluška Affiliation: Charles University, Prague Email: mburo@ualberta.ca Martin Schmid Affiliation: Charles University, Prague Affiliation: EquiLibre Technologies, Inc. Email: nathanst@ualberta.ca Michael Buro Affiliation: University of Alberta Email: mbowling@ualberta.ca Nathan R. Sturtevant Affiliation: University of Alberta Affiliation: Alberta Machine Intelligence Institute (Amii) Email: haluskr@kam.mff.cuni.cz Michael Bowling Affiliation: University of Alberta Affiliation: Alberta Machine Intelligence Institute (Amii) Email: schmidm@kam.mff.cuni.cz Abstract 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 approximation quality, and prior methods with learned networks rely on hand-chosen summaries of the sampled states to compute them. A belief representation should be computable from the parent node’s representation using the POMDP’s dynamics and it should let the search sample from it, but it should also be learned so that it provides the information needed to condition policies and value functions. 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 Agents planning in a Partially Observable Markov Decision Process (POMDP) must act on a belief, the posterior distribution over hidden states given their history of actions and observations. The action an agent takes should depend on the belief rather than on any single state, so policies map beliefs to actions and value functions map beliefs to expected returns. Computing them exactly is intractable in general, even when the state, action, and observation spaces are finite (Papadimitriou and Tsitsiklis, 1987). Monte Carlo planners instead search forward from the agent’s current history by simulating the POMDP’s dynamics and observation model (Silver and Veness, 2010; Somani et al., 2013). These simulations draw states from a representation of the belief at every node of the search tree and stop at leaves that need value estimates. Monte Carlo planners typically represent each node’s belief with a set of sampled states and differ in how they maintain that set and how they estimate leaf values. At each node POMCP and POMCPOW keep the states that simulations carried into it and estimate leaf values with rollouts of a default policy (Silver and Veness, 2010; Sunberg and Kochenderfer, 2018). DESPOT and AdaOPS search a fixed set of sampled scenarios and estimate leaf values with upper and lower bounds (Somani et al., 2013; Wu et al., 2021). The accuracy of these leaf values depends on the rollout policies and bounds, which are either designed for each domain or loose. Maintaining a set of sampled states also carries the weaknesses of particle filters. The number of states needed for an accurate approximation grows rapidly with the size of the state space (Thrun, 2002; Snyder et al., 2008), and repeated resampling also collapses the set onto a few distinct states, a failure called impoverishment (Sokota et al., 2022). Policy and value networks trained on experience generated by the planner can replace rollouts and bounds (Silver et al., 2018). The networks of BetaZero (Moss et al., 2024) condition on a hand-chosen summary of the particle set, the mean and standard deviation of its features. This summary discards much of what distinguishes one belief from another, such as its support and the dependence between its variables. BetaZero also keeps a particle set at every node of its tree, so it carries the same weaknesses of particle filters. A learned belief representation that does not collapse under repeated updates would address both, provided updating it and sampling from it remain tractable. Solinas et al. (2025) proposed Neural Bayesian Filtering (NBF), which represents beliefs as generative distribution embeddings (GDEs) (Fishman et al., 2025) and approximates their Bayesian update after actions and observations with a forward pass through the GDE’s generator and encoder. At each step, NBF generates particles from the previous belief’s embedding, propagates and reweights them under the POMDP’s dynamics and likelihood, and embeds the weighted set as the next belief. Every update therefore draws a fresh particle set, so impoverishment from resampling does not accumulate across steps. NBF’s conditional generator is trained to sample from the belief the embedding represents, so the search can generate the states it needs. If the embedding is accurate it is also a sufficient statistic of the history (Fishman et al., 2025), so policies and value functions can condition on it in place of the belief. We introduce Belief Embedding Tree Search (BETS), a Monte Carlo planner that searches over a tree of embeddings. Every node stores one embedding, expanding a leaf is one NBF update, and the policy and value networks that guide the search condition on the embedding, the same representation of the belief that the generator samples from. Hidden states in many discrete POMDPs are configurations of coupled variables, so BETS adds a variable-coupling stream to the GDE that pools a joint feature of every pair of variables over the particle set. The belief model and both networks train together on episodes collected by running BETS, and the belief model’s target is the hidden state the environment realizes at each step (Solinas et al., 2025). We evaluate BETS against BetaZero, POMCPOW, DESPOT, and AdaOPS on six configurations of RockSample, LightDark, and PocMan (Smith and Simmons, 2004; Platt et al., 2010; Silver and Veness, 2010). Each baseline falls far below the best return in at least one setting, while BETS is competitive or leading on all six. On RockSample(20,20) and PocMan, the configurations with the highest-dimensional beliefs, BETS outperforms every baseline by a large margin, even when it acts with its policy network alone. On PocMan, value networks predict the return more accurately from embeddings than from summary statistics of a particle filter with 16× as many particles. BETS also achieves a higher return with this embedding than with a pooled embedding twice as large. 2 Background BETS plans in POMDPs whose dynamics and observation model are given as input to the problem. A POMDP is a tuple G=(X,A,Y,T,,R,γ,b0)G=(X,A,Y,T,Y,R,γ,b_0) of a set of states X, a finite set of actions A, a set of observations Y, a transition function T:X×A→ΔXT:X× A→ X, an observation function :X×A×X→ΔYY:X× A× X→ Y, a reward function R:X×A×X→ℝR:X× A× X , a discount factor γ, and an initial state distribution b0∈ΔXb_0∈ X. An episode starts from x(0)∼b0x^(0) b_0, and at each step t the agent takes an action a(t)a^(t), the environment moves to x(t+1)∼T(⋅∣x(t),a(t))x^(t+1) T(· x^(t),a^(t)), and the agent receives the observation y(t+1)∼(⋅∣x(t),a(t),x(t+1))y^(t+1) (· x^(t),a^(t),x^(t+1)) and the reward R(x(t),a(t),x(t+1))R(x^(t),a^(t),x^(t+1)). The agent observes its own actions, so what it knows at time t is the history ht=def(a(0),y(1),…,a(t−1),y(t))h_t def=(a^(0),y^(1),…,a^(t-1),y^(t)), and it acts according to a policy π:H→ΔAπ:H→ A, a mapping from the set H of histories to distributions over actions. The objective is a policy that maximizes the expected discounted return G,π[∑t=0tmax−1γtR(x(t),a(t),x(t+1))]E_G,π[ _t=0^t_max-1γ^tR(x^(t),a^(t),x^(t+1))] over episodes of tmaxt_max steps, and the action value Qπ(ht,a)=defG,π[∑t′=ttmax−1γt′−tR(x(t′),a(t′),x(t′+1))∣ht,a(t)=a]Q^π(h_t,a) def=E_G,π[ _t =t^t_max-1γ^t -t\,R(x^(t ),a^(t ),x^(t +1)) h_t,a^(t)=a] is the expected return from taking a after hth_t and following π thereafter. The initial state is random, transitions can be stochastic, and distinct states can emit the same observation, so histories do not identify states; they induce a distribution over X. This distribution is the belief btb_t, which depends on hth_t. Given a(t)a^(t) and y(t+1)y^(t+1), Bayes’ rule gives the posterior belief, bt+1(x′)∝∑x∈Xbt(x)T(x′∣x,a(t))(y(t+1)∣x,a(t),x′).b_t+1(x ) _x∈ Xb_t(x)\,T(x x,a^(t))\,Y(y^(t+1) x,a^(t),x ). (1) The update, with an integral in place of the sum when X is continuous, is exact but intractable when the state space is large or continuous. The belief is a sufficient statistic of the history (Kaelbling et al., 1998), so an optimal policy and its action values depend on hth_t only through btb_t. Planning in a POMDP is therefore planning in a belief MDP, whose states are beliefs, whose transitions apply Eq. 1 to the next observation, and whose values are expected returns under btb_t. Its state space is continuous even when X is finite, and the planners of Section 2.2 search with approximations of btb_t. 2.1 Approximating Beliefs Particle filters (Gordon et al., 1993; Doucet and Johansen, 2009) approximate btb_t by a set of weighted states, or particles, xi,wii=1n\x_i,w_i\_i=1^n. Given a(t)a^(t) and y(t+1)y^(t+1), the update simulates a successor xi′∼T(⋅∣xi,a(t))x_i T(· x_i,a^(t)) for every particle and multiplies its weight by the likelihood (y(t+1)∣xi,a(t),xi′)Y(y^(t+1) x_i,a^(t),x_i ). When many weights become small, the filter resamples by duplicating high-weight particles and discarding the rest. Resampling never introduces a state the set did not already contain, so repeated resampling concentrates the set on a few distinct states, a failure called impoverishment. The approximation of btb_t then collapses, and the value estimates a planner computes from it inherit the error (Poupart and Boutilier, 2000). Regularization (Musso et al., 2001), MCMC (Gilks and Berzuini, 2001), and locally optimal proposals (Doucet and Johansen, 2009) slow the collapse, but every new particle still descends from the surviving set, so loss compounds over the episode. Preventing weight degeneracy takes exponentially many particles in the dimension of the state (Snyder et al., 2008). Neural Bayesian Filtering (NBF) (Solinas et al., 2025) approximates the Bayes update of Eq. 1 in the latent space of a generative distribution embedding (GDE) (Fishman et al., 2025). A GDE consists of an encoder ℰϕ:⋃n≥1Xn×ℝn→ℝdembE_φ: _n≥ 1X^n×R^n ^d_emb that maps a weighted set of states to a vector θ, and a conditional generator ψ(θ)G_ψ(θ) that maps the vector back to sample states. The encoder is distributionally invariant (Fishman et al., 2025): it depends on a weighted set only through the set’s weighted empirical measure, so reordering the set or duplicating every sample leaves its output unchanged. The NBF update θ′←ϕ,ψ(θ,a,y)θ _φ,ψ(θ,a,y) generates n particles from ψ(θ)G_ψ(θ), propagates each through T, weights it by Y, and embeds the weighted set with ℰϕE_φ. Each update regenerates its particle set from the embedding, so the collapse that resampling causes in a particle filter does not accumulate across steps. If the particle set loses most or all of the belief’s support in a single update, the approximation can collapse irrecoverably. In practice, this happens only in extreme cases, such as a sharp likelihood on a belief far from those BETS visits during training. 2.2 Planning with Partial Observability Full-width planners expand every branch with exact belief updates (Ross et al., 2008), which limits them to small state and observation spaces. Monte Carlo planners (Kocsis and Szepesvári, 2006; Browne et al., 2012) scale better by simulating states drawn from an approximate belief through the environment dynamics. POMCP (Silver and Veness, 2010) runs UCT over histories with a particle set at each node, and POMCPOW (Sunberg and Kochenderfer, 2018) adds progressive widening for large observation spaces. DESPOT (Somani et al., 2013) and AdaOPS (Wu et al., 2021) search a sparse tree over a fixed set of sampled scenarios and prune it with upper and lower bounds on the value. These planners have no learned components, so their search is only as informed as their rollout policies and bounds, but they need no training phase. BetaZero (Moss et al., 2024) approximates each node’s belief with a particle filter, and uses learned policy and value networks to guide search. These networks condition on the mean and standard deviation of features of the particle set and are trained on targets from BetaZero’s own searches (Silver et al., 2018). Summary statistics of particle features can discard the structure a decision depends on. Consider two cells, each of which may hold a hidden object, so the state is x∈0,12x∈\0,1\^2 with xj=1x_j=1 when cell j is occupied. The belief bexclb_excl is uniform over 01,10\01,10\, where exactly one cell is occupied, and bcob_co is uniform over 00,11\00,11\, where both or neither are. Both beliefs give each cell probability 12 12 of holding the object, so their per-cell means and standard deviations are identical, yet their supports are disjoint. An action rewarded only when both cells are occupied has zero expected reward under bexclb_excl and positive expected reward under bcob_co. A network that conditions on per-variable statistics must assign the two beliefs the same policy and value. A generator conditioned on per-variable statistics fails in the same way, since it must produce one distribution for both beliefs and so places at least half of its mass on configurations that one of them rules out. A planner that simulates those configurations estimates values on states the belief rules out. The next section builds a planner whose networks and generator condition on an embedding that represents this dependence explicitly. 3 Belief Embedding Tree Search Belief Embedding Tree Search (BETS) is a decision-time planner for POMDPs that searches over a tree of embeddings (Figure 1). At each step, BETS roots the tree at the root’s embedding and returns a policy over actions that determines the agent’s next action. Every node in the tree stores one embedding θ, and expanding a leaf is exactly one NBF update. A policy network πω:ℝdemb×ℝdξ→ΔA _ω:R^d_emb×R^d_ξ→ A and a value network Vη:ℝdemb×ℝdξ→ℝV_η:R^d_emb×R^d_ξ guide the search. Both networks condition on the embedding and a small vector ξ of features of the observed part of the state, such as the fraction of the horizon elapsed (Appendix A). Conditioning on θ rather than on per-variable statistics lets the networks assign different policies and values to beliefs that share those statistics. a1a_1a2a_2a3a_3NBF updateleaf value VηV_ηRoot: paired rounds. one x~∼ψ(θt) x _ψ( _t) per round, every a once, shared random numbersInterior: PUCT. prior ρ=(1−m)u+mπωρ=(1-m)\,u+m\, _ωa1a_1a2a_2a3a_3root Q(θt,a)Q( _t,a)a1a_1a2a_2a3a_3π¯∝ρeQ/τ π ρ\,e^Q/τafter KKsimulationsprior ρ, temperature τ :a∼π¯a π (train), argmaxaπ¯ _a π (eval)policy target (Sec. 4)value bootstrap:∑aπ¯(a)Q(θt,a) _a π(a)\,Q( _t,a) Figure 1: BETS. Each round starts every root action from the same sampled state and interior actions are selected via PUCT. Expansion costs one NBF update. 3.1 Search At decision time, BETS follows the standard Monte Carlo tree search procedure of selection, expansion, evaluation, and backup (Kocsis and Szepesvári, 2006; Browne et al., 2012). It runs K simulations of this procedure through a tree rooted at the belief embedding θt _t. Root. BETS groups the K simulations into K/|A|K/|A| rounds, and each round selects every root action exactly once. At the start of a round, BETS samples a state x~∼ψ(θt) x _ψ( _t), and every action in the round starts from x~ x with the same random numbers, so the actions’ returns are paired (Veness et al., 2011). The decision depends only on the differences between root action values, and pairing reduces the variance that the sampled state adds to those differences. When UCT (Kocsis and Szepesvári, 2006) or PUCT (Rosin, 2011; Silver et al., 2017) selects the root action, each action is evaluated on its own sampled states, and that state sampling variance can influence the relative value estimates of actions for finite K. Selecting every action once per round leaves root visit counts uniform, so the decision and the policy target are computed from the root action values instead (Section 3.2). Interior. Below the root, the traversal applies the selected action to x~ x, samples the next state and observation from T and Y, and descends to the child node for that action and observation if one exists. At an interior node with embedding θ, the traversal selects the action maximizing PUCT, argmaxa∈AQ¯(θ,a)+cρ(a∣θ,ξ)N(θ)1+N(θ,a), *arg\,max_a∈ A\; Q(θ,a)+c\,ρ(a θ,ξ)\, N(θ)1+N(θ,a), where N(θ)N(θ) and N(θ,a)N(θ,a) are the node and edge visit counts, Q¯ Q is the action value estimate normalized to [−1,1][-1,1] within the node (Schrittwieser et al., 2020), and c is the exploration constant. The prior ρ mixes the policy network with the uniform distribution u over legal actions, ρ(a∣θ,ξ)=def(1−m)u(a)+mπω(a∣θ,ξ),ρ(a θ,ξ) def=(1-m)\,u(a)+m\, _ω(a θ,ξ), so no action’s prior reaches zero and BETS keeps exploring regardless of policy network bias. Expansion, evaluation, and backup. When the current node has no child for the action a and observation y of the traversal’s last step, BETS expands by adding a new leaf. If the current node lies d levels below the root with embedding θt+d _t+d, the leaf’s embedding is computed with an NBF update, θt+d+1←ϕ,ψ(θt+d,a,y) _t+d+1 _φ,ψ( _t+d,a,y). The value network evaluates the leaf, Vη(θt+d+1,ξt+d+1)V_η( _t+d+1, _t+d+1), with no rollout. The value backed up through the traversal is the discounted sum of the rewards generated below it, plus the leaf value. The visit counts along the path are then incremented. Each simulation applies the NBF update, the policy network, and the value network once, and every node stores one embedding in place of a particle set. 3.2 Action Selection After K simulations, the root node holds a paired estimate Q(θt,a)Q( _t,a) for every action. BETS turns these estimates into a distribution over actions with one step of KL-regularized policy improvement (Abdolmaleki et al., 2018), π¯=defargmaxπ′∈ΔA∑aπ′(a)Q(θt,a)−τKL(π′∥ρ(⋅∣θt,ξt)), π def= *arg\,max_π ∈ A\; _aπ (a)\,Q( _t,a)-τ\,KL\! (π \,\|\,ρ(· _t, _t) ), where the temperature τ sets how far the search can move the policy away from the prior. The maximizer has the closed form π¯(a)∝ρ(a∣θt,ξt)exp(Q(θt,a)/τ) π(a) ρ(a _t, _t) (Q( _t,a)/τ), which is one mirror descent step from the prior with step size 1/τ1/τ (Beck and Teboulle, 2003). AlphaZero’s visit-count distribution approximates the solution of this problem with the KL taken in the other direction (Grill et al., 2020), and Danihelka et al. (2022) use this closed form. BETS evaluates this closed form exactly. During training, the agent acts by sampling a∼π¯a π; at evaluation, it takes argmaxaπ¯(a) _a π(a). The same π¯ π is the policy network’s training target, and its expected root value ∑aπ¯(a)Q(θt,a) _a π(a)\,Q( _t,a) is the bootstrap for the value network’s target (Section 4.2). After BETS acts, one more NBF update computes the next root embedding from the action taken and the observation received. 3.3 Variable-Coupling Embeddings Standard GDE encoders compute a feature vector for each sample and pool the per-sample features over the set (Fishman et al., 2025). Each per-sample feature depends on every variable of its sample, so in principle the pooled features can capture how the variables of a sample depend on one another. The two beliefs in the example of Section 2.2 give each cell the same probability of holding the object, so the dependence between the two cells is all that separates them. In practice, the encoder has to discover the dependence inside generic per-sample features, and if the dependence does not survive the pooling, the two beliefs receive the same embedding. We extend GDEs by giving the encoder and the generator an explicit representation of the co-occurrence between pairs of variables. A discrete state is a vector x∈0,…,q−1Dx∈\0,…,q-1\^D of D variables. The pooled stream is a set transformer (Lee et al., 2019) whose per-sample features are averaged over the set to give the pooled embedding μ. A second stream, which we call the variable-coupling stream, averages a joint feature of each pair of variables over the same samples. From these pairwise averages it forms a coupling summary for each variable that describes how the variable relates to the others (Figure 2). The D coupling summaries are combined with μ to form θ. The policy and value networks take all of θ as input, so they condition on every coupling summary as well as on μ. Both streams average over samples, so θ remains distributionally invariant, and both are trained jointly with the generator. samples ×variablesvariable summariesvariable × variable co-occurrenceper-variablecouplingspooledembeddingμ∥\|θ Figure 2: The variable-coupling stream. Pooling over samples gives one summary per variable and a variable-by-variable co-occurrence matrix. Attention over the variable summaries, biased by the co-occurrence matrix, produces coupling summaries, which are concatenated with μ to form θ. BETS pairs the variable-coupling stream with an autoregressive generator that treats each variable as a token and factorizes the belief over them, pψ(x∣θ)=∏j=1Dpψ(xj∣x<j,θ)p_ψ(x θ)= _j=1^Dp_ψ(x_j x_<j,θ), where pψ(⋅∣θ)p_ψ(· θ) is the density of ψ(θ)G_ψ(θ). The generator conditions factor j on the preceding variables x<jx_<j, the shared pooled embedding μ, and the coupling summary of variable j (Appendix C.1). Each factor is a normalized categorical, so the product is an exact joint likelihood computable in one forward pass. 4 Training BETS learns three models: the belief model, which is the GDE encoder and generator that NBF uses in its update, and the policy and value networks that guide the search. The beliefs BETS visits depend on its policy, so the beliefs the model must represent change as the networks train and the search policy improves. All three models train together on episodes that BETS collects. Training alternates between gradient steps and collection rounds. In each collection round, a set of collectors runs BETS with the parameters from the most recent gradient step, and each collector takes a fixed number of environment steps. The consecutive transitions that one collector records in one round form a fragment, and an episode that ends inside a fragment is reset, so the fragment continues with the next episode. Before the first gradient step, a fixed number of collection rounds run with the initial model parameters. Collection rounds repeat at a fixed interval of gradient steps, and every gradient step draws one batch for each of the three losses from the fragments of the most recent rounds. The policy and value networks are trained on the embeddings θt _t stored during collection, so their losses never reach the encoder and the GDE is trained on the belief loss alone. The exploration constant c, the mix weight m, and the temperature τ are annealed from exploratory to greedy values over the first part of training (Appendix B). 4.1 Belief Model Each training example for the GDE is a tuple (θt,a(t),y(t+1),x(t+1))( _t,a^(t),y^(t+1),x^(t+1)) stored from a collected episode. For each example, the NBF update generates particles from ψ(θt)G_ψ( _t), propagates them under a(t)a^(t), and weights them by the likelihood of y(t+1)y^(t+1), which yields the weighted set St+1S_t+1 and its embedding θt+1=ℰϕ(St+1) _t+1=E_φ(S_t+1). The update runs at training time, so the particle sets the belief model embeds are regenerated from the generator being trained rather than loaded from storage. At the start of an episode, the particles are drawn from b0b_0 instead. The target is the hidden state x(t+1)x^(t+1) that the environment realized after a(t)a^(t). Though the agent never observes x(t+1)x^(t+1), it is the ground truth state for simulation, so it is available to compute the loss at training time. The belief model minimizes the realized-state objective of Solinas et al. (2025), ℒ(ϕ,ψ)=def[ℓψ(x(t+1),ℰϕ(St+1))],L(φ,ψ) def=E [ _ψ (x^(t+1);\,E_φ(S_t+1) ) ], (2) where the expectation is over the collected transitions and over the particle sets the NBF update generates from them. ℓψ(x,θ)=−logpψ(x∣θ) _ψ(x;θ)=- p_ψ(x θ) is the GDE’s negative log-likelihood of the single target x given the embedding θ=ℰϕ(St+1)θ=E_φ(S_t+1).11 1 On LightDark, ℓψ _ψ is instead the conditional flow-matching regression at x(t+1)x^(t+1) (Appendix C.2). Gradients pass through the embedding of the set and through the generator, but not through the sampling that produced the set. BETS’s actions depend on the hidden state only through the history, so each x(t+1)x^(t+1) is a sample from the true posterior bt+1b_t+1 given the history ht+1h_t+1. The particle set St+1S_t+1 is computed from that history and the filter’s own random draws, and the set posterior p(⋅∣St+1)p(· S_t+1) is the distribution of x(t+1)x^(t+1) given the set alone. It equals bt+1b_t+1 when the set retains all the information that the history carries about the state, and in general it is the belief with the lowest expected log loss among all beliefs computed from the set. By the decomposition of Solinas et al. (2025, Proposition 4.1), minimizing ℒ(ϕ,ψ)L(φ,ψ) for fixed sets is equivalent to minimizing [KL(p(⋅∣St+1)∥pψ(⋅∣θt+1))]E [KL (p(· S_t+1)\,\|\,p_ψ(· _t+1) ) ], the expected KL divergence from the set posterior to the GDE’s belief pψ(⋅∣θt+1)p_ψ(· _t+1). 4.2 Policy and Value Targets The policy and value networks are trained on targets derived from the search (Anthony et al., 2017; Silver et al., 2017). At step t, the root gives the improved policy π¯t π_t and its value v¯t=def∑aπ¯t(a)Q(θt,a) v_t def= _a π_t(a)\,Q( _t,a), and the transition (θt,ξt,a(t),rt,π¯t,v¯t)( _t, _t,a^(t),r_t, π_t, v_t) is stored, where rtr_t is the reward of step t. The policy target is π¯t π_t, and the policy network minimizes the cross-entropy H(π¯t,πω(⋅∣θt,ξt))H( π_t, _ω(· _t, _t)). The value target is the λ-return bootstrapped from the root values v¯ v (Sutton and Barto, 1998), ztλ=defrt+γ[(1−λ)v¯t+1+λzt+1λ],z^λ_t def=r_t+γ [(1-λ)\, v_t+1+λ\,z^λ_t+1 ], computed backward over each fragment, with zt+1λ=defv¯t+1z^λ_t+1 def= v_t+1 past the end of a fragment and v¯t+1=zt+1λ=0 v_t+1=z^λ_t+1=0 at terminal states. The value network is trained as a classifier over a fixed grid of bins. Each λ-return is squashed by f(u)=sign(u)(|u|+1−1)+ϵuf(u)=sign(u) ( |u|+1-1 )+ε u (Pohlen et al., 2018), and f(ztλ)f(z^λ_t) is encoded as a two-hot distribution over the bins (Schrittwieser et al., 2020). Training minimizes the cross-entropy between the value network’s prediction and this two-hot target. During search, a leaf’s value is f−1f^-1 applied to the mean of the predicted distribution over the bins. 5 Experiments We evaluate BETS on six configurations of three POMDP domains: RockSample(15,15) and RockSample(20,20) (Smith and Simmons, 2004), LightDark(5) and LightDark(10) (Platt et al., 2010), and PocMan on a medium maze (M) and the large maze (L) of Silver and Veness (2010) (Appendix A). The domains differ in belief structure. In RockSample, the agent’s position is known, and the belief is a product of 15 or 20 independent Bernoulli distributions, one per rock, over 2152^15 or 2202^20 joint configurations. In LightDark, the agent moves along a line, and the belief is a one-dimensional density over its position that stays broad in the dark and sharpens near the light. In PocMan, it is a joint distribution over the positions and headings of three or four ghosts, which move through the same maze, chasing or running away from the agent when it is nearby. Baselines. We compare against BetaZero (Moss et al., 2024) and three planners without learned components: AdaOPS (Wu et al., 2021), DESPOT (Somani et al., 2013), and POMCPOW (Sunberg and Kochenderfer, 2018). At evaluation, all four use the same number of simulations per decision and particles as BETS, and otherwise use the settings of Moss et al. (2024) on RockSample and LightDark (Appendix B). Protocol. BETS is evaluated in two ways on the same model: with search, acting by argmaxaπ¯(a) _a π(a) with the search settings training ends with, and with the policy network alone, acting by argmaxaπω(a∣θ,ξ) _a _ω(a θ,ξ). We tuned the training hyperparameters of BETS on every domain and those of BetaZero on PocMan, where Moss et al. (2024) give no settings. BETS and BetaZero are trained at three seeds, and each trained model is evaluated on 300 episodes at each of three evaluation seeds. AdaOPS, DESPOT, and POMCPOW have no learned components, so each runs three evaluation seeds with 300 episodes each. Table 1: Mean discounted return with the SE across training seeds (BETS, BetaZero) or evaluation seeds (AdaOPS, DESPOT, POMCPOW). AdaOPS has no bounds implementation for PocMan. BETS Domain search policy BetaZero AdaOPS DESPOT POMCPOW RockSample(15,15) 20.78 ± 0.09 20.65 ± 0.10 19.79 ± 0.18 21.45 ± 0.10 19.45 ± 0.10 11.45 ± 0.08 RockSample(20,20) 14.97 ± 0.13 13.91 ± 0.16 11.30 ± 0.15 10.77 ± 0.20† 0.28 ± 0.02† 4.33 ± 0.18† LightDark(5) 3.61 ± 0.10 3.77 ± 0.12 3.76 ± 0.15 3.08 ± 0.03† 2.29 ± 0.09 2.82 ± 0.25 LightDark(10) 13.01 ± 0.41 13.88 ± 0.28 14.43 ± 0.33 3.18 ± 0.40† 0.33 ± 0.05 0.98 ± 0.13 PocMan(M) 78.08 ± 0.54 72.42 ± 2.05 40.87 ± 3.26 – 54.20 ± 0.34 36.34 ± 1.93 PocMan(L) 78.95 ± 0.23 77.47 ± 0.46 69.87 ± 1.64 – 60.49 ± 0.71 39.91 ± 1.41 † Without domain heuristics: exit-policy lower and fixed upper bounds for AdaOPS and DESPOT on RockSample(20,20), fixed bounds for AdaOPS on LightDark, and no rollout heuristic for POMCPOW. 5.1 Results Table 1 shows that BETS earns the highest return on RockSample(20,20) and on both PocMan mazes, where baseline search heuristics or belief approximations fail. It trails AdaOPS on RockSample(15,15), where AdaOPS’s heuristic bounds are feasible, and BetaZero on LightDark(10). On RockSample(20,20) computing the heuristic upper bounds of AdaOPS and DESPOT is infeasible, so both run on a fixed upper bound, which reduces discounted return. BetaZero also learns to guide its search, so we attribute most of the difference from BetaZero to what its networks condition on: the mean and standard deviation of each state variable over its particle set. Either the particle set or its summary statistics can make a poor input to the policy and value networks. On RockSample(20,20) the belief is a product of Bernoulli distributions, so the per-rock means are a sufficient statistic for it. We attribute the gap there to BetaZero’s particle set, which covers a small fraction of the 2202^20 configurations. Each resampling step of the bootstrap filter duplicates the surviving particles, and with deterministic transitions the duplicates stay identical, so the per-rock means rest on ever fewer distinct configurations that carry disproportionate weight, whereas BETS regenerates its particles from the embedding at every update. A resampler that redraws each rock from its per-rock mean would restore distinct configurations, but it relies on the factorization, which is a domain-specific property. On PocMan we attribute the gap to the summary statistics, which lose two kinds of structure. A ghost that may be in any of several corridors has a multimodal belief, and the mean and standard deviation of its position cannot distinguish those corridors from the region between them. These statistics also discard the dependence between state variables, such as ghost positions and headings, which is what the variable-coupling stream was designed to model. BetaZero matches BETS on LightDark(5) and leads on LightDark(10). With a one-dimensional state the particle set approximates the belief well. Far from the light, the Gaussian likelihood spreads the weight over many particles even under an inaccurate proposal, and near it the belief itself narrows almost to a point, so its mean and standard deviation are close to a sufficient statistic for it. The policy network alone matches search on LightDark(5) and exceeds it on LightDark(10). We hypothesize that, because every sampled observation is new, the tree stays shallow and a single search compares actions mostly through the value network’s estimates at a few sampled leaves, whose noise the policy network averages out by training on π¯ π from many searches. On RockSample(20,20) and both PocMan mazes, the policy network alone earns more return than every baseline does with search. Effect of Variable Coupling on PocMan. VC is the BETS configuration of Table 1, and Pooled removes its variable-coupling stream and doubles the pooled embedding to 128 dimensions, which more than doubles the model’s capacity. Everything else is held fixed, and both are trained with three seeds on PocMan(M). Table 2(a) shows that Pooled earns less return than VC with both search and the policy network, so the added capacity does not replace the variable-coupling stream. Pooled has to discover the dependence between ghost positions and headings inside generic per-sample features. With search, Pooled also clears the maze in fewer than half as many episodes, and removing the stream costs more discounted return. Search scores its leaves with the value network and simulates states drawn from the belief, so the larger difference under search suggests that the coupling summaries improve both the value estimates and the simulated states. Table 2: Variable coupling and value prediction on PocMan(M). VC Pooled Discounted return (policy) 72.42 ± 2.05 65.39 ± 0.72 Discounted return (search) 78.08 ± 0.54 68.20 ± 0.48 Maze clears (search) 14.1% 6.7% Episodes Collected 5000 10000 25000 50000 VC 0.256 0.280 0.315 0.341 Pooled 0.245 0.276 0.311 0.330 Moments 0.218 0.250 0.285 0.311 (a) Mean discounted return over three training seeds. (b) Held-out R2R^2 of value networks. Value Prediction from Belief Representations. We collect 50000 PocMan(M) episodes with the policy network from one training seed, while a bootstrap particle filter with the true model and 8192 particles tracks the belief on the same actions and observations. Value networks with the same architecture as those of Tables 1 and 2(a) are trained by squared error to predict the discounted return from each step, conditioned on either the mean and variance of the filter’s particles (Moments) or the model’s embedding (VC / Pooled). We report the held-out R2R^2 under five-fold cross-validation over episodes, where each network trains on episodes from three folds and early stops with the fourth. Table 2(b) shows that, for both models and at every training-set size, value networks conditioned on the embedding predict the return significantly more accurately than those conditioned on the moments (paired t-test over held-out episodes). The moment estimates use 16 times as many particles, so a shortage of particles does not explain the gap. This supports the attribution of BetaZero’s gap on PocMan to conditioning on particle summary statistics, and conditioning on the combination of the pooled embedding and coupling summaries leads to better value estimates than conditioning on the pooled embedding alone. 6 Limitations Paired simulations at the root spend the budget evenly among all actions, so each action’s share shrinks as the action space grows. Sequential halving discards half of the remaining root actions after each phase and assigns each unvisited action an estimate of the root’s value, so π¯ π remains defined over every root action (Danihelka et al., 2022). Pairing the actions within each phase would keep the variance reduction of common sampled states and extend BETS to large action spaces. BETS must be trained on each domain before it can search effectively, while AdaOPS, DESPOT, and POMCPOW have no training phase. Against BetaZero, training improves returns where a particle set and its summary statistics represent the belief poorly, as on RockSample(20,20) and PocMan, but not on LightDark, where it adds computational overhead with no benefit over simpler belief models. Each simulation applies the NBF update, the policy network, and the value network once, so search takes roughly K times as long per decision as the policy network alone (Appendix D). The policy network alone earns more return than every baseline does with search on RockSample(20,20) and both PocMan mazes (Table 1), so BETS can act without search where decisions must be fast. Training also requires a simulator that exposes the realized hidden state, which might not always be available in practice. 7 Conclusion This paper introduced Belief Embedding Tree Search (BETS), a Monte Carlo planner that represents the belief at search tree nodes with a generative distribution embedding. The embedding lets the search sample states, and it gives the policy and value networks a learned compact representation of the belief to condition on. We show that these function approximators benefit from conditioning on embeddings, rather than summaries of particle sets. BETS is competitive or leading on every configuration, and even without search it outperforms every baseline by a large margin on RockSample(20,20) and PocMan, the configurations with the highest-dimensional beliefs. Our experiments also suggest that adding per-variable coupling summaries to a GDE improves the model’s representation on PocMan. BETS shows that a learned belief embedding can take the place of a particle set inside a planner, and the same substitution is open to any algorithm that learns or plans on beliefs. Embedding and updating beliefs is not specific to tree search, so the same ideas extend to other learning and planning algorithms for partially observable environments. In games, player policies would enter the NBF update through the same particle propagation as the dynamics, so an embedding could represent public beliefs needed for many algorithms (Moravčík et al., 2017; Brown et al., 2020). AI Use Statement Large language models assisted the authors with writing and coding. They were used to help the authors refine writing and help implement and test the experiment pipeline. The authors revised and verified all LLM work carefully. Reproducibility Statement All source code for BETS is provided as supplementary material and will be made publicly available on GitHub upon publication. Appendix B gives the hyperparameter settings of BETS and of every baseline. References Abdolmaleki et al. (2018) A. Abdolmaleki, J. T. Springenberg, Y. Tassa, R. Munos, N. Heess, and M. Riedmiller Maximum a posteriori policy optimisation. In International Conference on Learning Representations, Cited by: §3.2. Anthony et al. (2017) T. Anthony, Z. Tian, and D. Barber Thinking fast and slow with deep learning and tree search. In Advances in Neural Information Processing Systems, Vol. 30. Cited by: §4.2. Beck and Teboulle (2003) A. Beck and M. Teboulle Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters 31 (3), p. 167–175. Cited by: §3.2. Brown et al. (2020) N. Brown, A. Bakhtin, A. Lerer, and Q. Gong Combining deep reinforcement learning and search for imperfect-information games. Advances in Neural Information Processing Systems 33, p. 17057–17069. Cited by: §7. Browne et al. (2012) C. B. Browne, E. Powley, D. Whitehouse, S. M. Lucas, P. I. Cowling, P. Rohlfshagen, S. Tavener, D. Perez, S. Samothrakis, and S. Colton A survey of Monte Carlo tree search methods. IEEE Transactions on Computational Intelligence and AI in Games 4 (1), p. 1–43. Cited by: §2.2, §3.1. Danihelka et al. (2022) I. Danihelka, A. Guez, J. Schrittwieser, and D. Silver Policy improvement by planning with Gumbel. In International Conference on Learning Representations, Cited by: §3.2, §6. Doucet and Johansen (2009) A. Doucet and A. M. Johansen A tutorial on particle filtering and smoothing: fifteen years later. In Handbook of Nonlinear Filtering, p. 656–704. Cited by: §2.1. Fishman et al. (2025) N. Fishman, G. Gowri, P. Yin, J. Gootenberg, and O. Abudayyeh Generative distribution embeddings. arXiv preprint arXiv:2505.18150. Cited by: §1, §2.1, §3.3. Gilks and Berzuini (2001) W. R. Gilks and C. Berzuini Following a moving target: Monte Carlo inference for dynamic Bayesian models. Journal of the Royal Statistical Society: Series B 63 (1), p. 127–146. Cited by: §2.1. Gordon et al. (1993) N. J. Gordon, D. J. Salmond, and A. F. Smith Novel approach to nonlinear/non-Gaussian Bayesian state estimation. IEE Proceedings F (Radar and Signal Processing) 140 (2), p. 107–113. Cited by: §2.1. Grill et al. (2020) J. Grill, F. Altché, Y. Tang, T. Hubert, M. Valko, I. Antonoglou, and R. Munos Monte-Carlo tree search as regularized policy optimization. In Proceedings of the 37th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 119, p. 3769–3778. Cited by: §3.2. Kaelbling et al. (1998) L. P. Kaelbling, M. L. Littman, and A. R. Cassandra Planning and acting in partially observable stochastic domains. Artificial Intelligence 101 (1–2), p. 99–134. Cited by: §2. Kocsis and Szepesvári (2006) L. Kocsis and C. Szepesvári Bandit based Monte-Carlo planning. In European Conference on Machine Learning, p. 282–293. Cited by: §2.2, §3.1, §3.1. Lee et al. (2019) J. Lee, Y. Lee, J. Kim, A. R. Kosiorek, S. Choi, and Y. W. Teh Set transformer: a framework for attention-based permutation-invariant neural networks. In International Conference on Machine Learning, Cited by: §3.3. Lipman et al. (2023) Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le Flow matching for generative modeling. In International Conference on Learning Representations, Cited by: §C.2. Moravčík et al. (2017) M. Moravčík, M. Schmid, N. Burch, V. Lisỳ, D. Morrill, N. Bard, T. Davis, K. Waugh, M. Johanson, and M. Bowling DeepStack: expert-level artificial intelligence in heads-up no-limit poker. Science 356 (6337), p. 508–513. Cited by: §7. Moss et al. (2024) R. J. Moss, A. Corso, J. Caers, and M. J. Kochenderfer BetaZero: belief-state planning for long-horizon POMDPs using learned approximations. Reinforcement Learning Journal 1, p. 158–181. Cited by: Appendix A, Appendix B, §1, §2.2, §5, §5. Musso et al. (2001) C. Musso, N. Oudjane, and F. Le Gland Improving regularised particle filters. In Sequential Monte Carlo Methods in Practice, A. Doucet, N. de Freitas, and N. Gordon (Eds.), p. 247–271. Cited by: §2.1. Papadimitriou and Tsitsiklis (1987) C. H. Papadimitriou and J. N. Tsitsiklis The complexity of Markov decision processes. Mathematics of Operations Research 12 (3), p. 441–450. Cited by: §1. Platt et al. (2010) R. Platt, R. Tedrake, L. Kaelbling, and T. Lozano-Pérez Belief space planning assuming maximum likelihood observations. In Robotics: Science and Systems, Cited by: Appendix A, §1, §5. Pohlen et al. (2018) T. Pohlen, B. Piot, T. Hester, M. G. Azar, D. Horgan, D. Budden, G. Barth-Maron, H. van Hasselt, J. Quan, M. Večerík, M. Hessel, R. Munos, and O. Pietquin Observe and look further: achieving consistent performance on Atari. arXiv preprint arXiv:1805.11593. Cited by: §4.2. Poupart and Boutilier (2000) P. Poupart and C. Boutilier Value-directed belief state approximation for POMDPs. In Conference on Uncertainty in Artificial Intelligence, p. 497–506. Cited by: §2.1. Rosin (2011) C. D. Rosin Multi-armed bandits with episode context. Annals of Mathematics and Artificial Intelligence 61 (3), p. 203–230. Cited by: §3.1. Ross et al. (2008) S. Ross, J. Pineau, S. Paquet, and B. Chaib-Draa Online planning algorithms for POMDPs. Journal of Artificial Intelligence Research 32, p. 663–704. Cited by: §2.2. Schrittwieser et al. (2020) J. Schrittwieser, I. Antonoglou, T. Hubert, K. Simonyan, L. Sifre, S. Schmitt, A. Guez, E. Lockhart, D. Hassabis, T. Graepel, et al. Mastering Atari, Go, chess and shogi by planning with a learned model. Nature 588 (7839), p. 604–609. Cited by: §3.1, §4.2. Silver et al. (2018) D. Silver, T. Hubert, J. Schrittwieser, I. Antonoglou, M. Lai, A. Guez, M. Lanctot, L. Sifre, D. Kumaran, T. Graepel, T. Lillicrap, K. Simonyan, and D. Hassabis A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play. Science 362 (6419), p. 1140–1144. Cited by: §1, §2.2. Silver et al. (2017) D. Silver, J. Schrittwieser, K. Simonyan, I. Antonoglou, A. Huang, A. Guez, T. Hubert, L. Baker, M. Lai, A. Bolton, et al. Mastering the game of Go without human knowledge. Nature 550 (7676), p. 354–359. Cited by: §3.1, §4.2. Silver and Veness (2010) D. Silver and J. Veness Monte-Carlo planning in large POMDPs. Advances in Neural Information Processing Systems 23. Cited by: Appendix A, Appendix B, §1, §1, §1, §2.2, §5. Smith and Simmons (2004) T. Smith and R. Simmons Heuristic search value iteration for POMDPs. In Proceedings of the 20th Conference on Uncertainty in Artificial Intelligence, p. 520–527. Cited by: Appendix A, §1, §5. Snyder et al. (2008) C. Snyder, T. Bengtsson, P. Bickel, and J. Anderson Obstacles to high-dimensional particle filtering. Monthly Weather Review 136 (12), p. 4629–4640. Cited by: §1, §2.1. Sokota et al. (2022) S. Sokota, H. Hu, D. J. Wu, J. Z. Kolter, J. N. Foerster, and N. Brown A fine-tuning approach to belief state modeling. In International Conference on Learning Representations, Cited by: §1. Solinas et al. (2025) C. Solinas, R. Haluska, D. Sychrovsky, F. Timbers, N. Bard, M. Buro, M. Schmid, N. R. Sturtevant, and M. Bowling Neural Bayesian filtering. arXiv preprint arXiv:2510.03614. Cited by: §1, §1, §2.1, §4.1, §4.1. Somani et al. (2013) A. Somani, N. Ye, D. Hsu, and W. S. Lee DESPOT: online POMDP planning with regularization. Advances in Neural Information Processing Systems 26. Cited by: Table 5, §1, §1, §2.2, §5. Sunberg and Kochenderfer (2018) Z. Sunberg and M. Kochenderfer Online algorithms for POMDPs with continuous state, action, and observation spaces. In Proceedings of the International Conference on Automated Planning and Scheduling, Vol. 28, p. 259–263. Cited by: Table 5, §1, §2.2, §5. Sutton and Barto (1998) R. S. Sutton and A. G. Barto Reinforcement learning: an introduction. MIT Press. Cited by: §4.2. Thrun (2002) S. Thrun Particle filters in robotics.. In Conference on Uncertainty in Artificial Intelligence, p. 511–518. Cited by: §1. Tong et al. (2023) A. Tong, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, K. Fatras, G. Wolf, and Y. Bengio Conditional flow matching: simulation-free dynamic optimal transport. arXiv preprint arXiv:2302.00482. Cited by: §C.2. Veness et al. (2011) J. Veness, M. Lanctot, and M. Bowling Variance reduction in Monte-Carlo tree search. Advances in Neural Information Processing Systems 24. Cited by: §3.1. Wu et al. (2021) C. Wu, G. Yang, Z. Zhang, Y. Yu, D. Li, W. Liu, and J. Hao Adaptive online packing-guided search for POMDPs. Advances in Neural Information Processing Systems 34, p. 28419–28430. Cited by: Table 5, §1, §2.2, §5. Appendix A Domains RockSample. RockSample (Smith and Simmons, 2004) places a rover on a square grid containing k rocks at known positions, each independently good or bad. The rover moves deterministically in the four cardinal directions, samples the rock at its current cell, or checks any rock with a noisy long-range sensor. A check returns the rock’s quality corrupted by noise whose accuracy decays exponentially with the distance to the rock. Sampling a good rock earns a positive reward and turns the rock bad, sampling a bad rock is penalized, and the episode ends with a positive reward when the rover exits the east edge of the grid. The rover’s position is fully observable and transitions are deterministic. We use the two configurations of Moss et al. (2024): RockSample(15,15) is a 15×1515× 15 grid with 15 rocks and 20 actions, and RockSample(20,20) is a 20×2020× 20 grid with 20 rocks and 25 actions. Both use a fixed rock layout, sensor efficiency 20, discount 0.95, and horizon 100. LightDark. LightDark (Platt et al., 2010) is a continuous one-dimensional localization domain. The agent’s position is a real-valued state that is never observed directly. Each step yields an observation of the position corrupted by Gaussian noise whose standard deviation grows with the distance from a light. The actions move the agent one unit left or right or stop, and stopping ends the episode with a positive reward if the agent’s true position is within distance 1 of the origin and a negative reward otherwise. Both configurations draw the agent’s start position from (2,32)N(2,3^2), and use a discount factor 0.9 and horizon 100. LightDark(5) places the light at position 5. The standard deviation of the noise is |x−5|/2+0.01|x-5|/ 2+0.01 at position x, and stopping earns +10+10 or −10-10. LightDark(10) places the light at position 10 with noise standard deviation |x−10|+10−4|x-10|+10^-4, and stopping earns +100+100 or −100-100. PocMan. PocMan (Silver and Veness, 2010) is a partially observable adaptation of the arcade game Pac-Man. The agent navigates a maze, collecting food pellets while avoiding ghosts that chase it when nearby, and eating a power pill briefly lets the agent eat ghosts for a bonus. The observation is 11 bits: sight of a ghost along each of the four cardinal directions, contact with walls in each direction, hearing a ghost within Manhattan distance 2, and smelling a pellet and smelling a power pill in the eight surrounding cells. Rewards are −1-1 per step, +10+10 for a pellet or a power pill, +25+25 for eating a ghost, −100-100 for being caught, and +1000+1000 for clearing the maze; the discount is 0.95. PocMan(M) is a 9×99× 9 maze with 59 reachable cells, three ghosts, three power pills, and horizon 100. PocMan(L) is the 17×1917× 19 maze of Silver and Veness (2010) with 189 reachable cells, four ghosts, four power pills, and horizon 250. In both mazes, the agent, the ghosts, and the power pills start on the same cells in every episode. The hidden state is the ghosts’ positions and headings together with the pellet field, in which each reachable cell holds a pellet independently with probability 0.5 at the start of an episode. The simulator forbids a ghost from reversing its direction, so its heading constrains its next move. The pellet field is not modelled by the generator and an observed pellet mask is given to all filters. Policy and Value Network Inputs. The policy and value networks condition on the embedding θ together with a vector ξ of features of the observed part of the state. In every domain, ξ contains the fraction of the horizon elapsed. On PocMan it also contains the fraction of pellets remaining, the power-pill timer, and the observed pellet mask over the reachable cells. Appendix B Hyperparameters BETS. Search settings are shared between training and evaluation, and an annealed setting moves linearly from its first to its second value across the anneal window and keeps the second value afterwards. An action node keeps at most 16 observation children, and once it is full, a traversal whose observation matches none of them descends to an existing child sampled in proportion to its visit count plus one. On LightDark every sampled observation is new, so observation progressive widening lowers this limit to ⌈kN(θ,a)α⌉ k\,N(θ,a)^α (Table 3). Table 3: BETS settings per configuration. A range is annealed over the anneal window. RockSample LightDark PocMan (15,15) (20,20) (5) (10) (M), (L) Search Simulations per decision K 600 600 99 99 100 Particles per node n 256 256 256 256 512 Exploration constant c 20 to 3 Prior mix m 0 to 0.75 Temperature τ 2 to 0.5 2 to 0.5 2 to 0.5 20 to 5 10 to 2.5 Observation widening (k,α)(k,α) – – (4, 0.1) (4, 0.1) – Training Gradient steps 150k 300k 75k 150k 60k Anneal window (gradient steps) 10k to 40k 15k to 60k 7.5k to 30k 15k to 60k 4k to 16k Collectors × fragment length 84 × 100 84 × 100 84 × 100 84 × 100 84 × 25 Collection interval (gradient steps) 400 400 400 400 100 Initial collection rounds 8 8 5 5 5 Value target λ 1 1 1 1 0.95 Optimizer Adam, learning rate 10−310^-3, batch 32, gradient clip 1.0 Networks Pooled embedding μ dimension 32 32 32 32 64 VC stream (width × layers) 16 × 3 16 × 3 – – 32 × 4 Set transformer width 128 128 128 128 256 Set transformer blocks, inducing points 3, 16 Attention heads 4 Generator (layers × width) AR, 3 × 48 AR, 3 × 48 FM, 3 × 64 FM, 3 × 64 AR, 3 × 48 Value bins on [fmin,fmax][f_ ,f_ ] 66, ±8.1± 8.1 66, ±8.1± 8.1 66, ±3± 3 66, ±9.2± 9.2 128, [−14,38][-14,38] Policy and value networks MLP, 3 hidden layers of 128 BetaZero. BetaZero runs with the same number of simulations per decision and particles as BETS in every domain, and otherwise with its default settings on RockSample and LightDark. On PocMan, BetaZero tracks the belief with a bootstrap particle filter. Its networks condition on the mean and standard deviation of the player and ghost positions and the mean pellet and power-pill grids over its particles. During policy iteration, BetaZero samples the root action from the visit counts weighted by the root action values (zq=zn=1z_q=z_n=1) at temperature τ, where τ=0τ=0 takes the argmax, and at evaluation it takes the argmax under the weights (zq,zn)(z_q,z_n). A dash for widening means it is disabled, so every action is expanded at once and each node keeps only the first belief it is reached with (Table 4). Table 4: BetaZero settings per configuration. A cell with two values gives the value during policy iteration, then the value at evaluation. RockSample LightDark PocMan (15,15) (20,20) (5) (10) (M), (L) Policy iteration Iterations × episodes 50 × 500 50 × 500 30 × 500 30 × 500 40 × 250 Epochs per iteration 10 10 50 50 10 Optimizer RMSProp RMSProp Adam Adam RMSProp Learning rate 10−310^-3 10−310^-3 10−410^-4 10−410^-4 10−410^-4 Value loss MSE MSE MSE MSE MAE Dropout 0.5 0.5 0.2 0.2 0.5 Hidden layers × width 4 × 128 4 × 128 4 × 64 4 × 64 4 × 256 Batch, L2L_2 penalty, samples 1024, 10−510^-5, 10510^5 from the latest iteration Search Simulations per decision 100 / 600 100 / 600 100 / 99 100 / 99 100 Particles 1000 / 256 1000 / 256 500 / 256 500 / 256 512 Exploration constant c 50 50 1 1 50 Maximum depth 15 4 10 10 10 Action widening (ka,αa)(k_a, _a) – / (5, 0.9) – (2, 0.25) (2, 0.25) – Belief widening (kb,αb)(k_b, _b) – / (1, 0) – / (1, 0) (2, 0.1) (2, 0.1) – Temperature τ 1 / 0 1.5 / 0 0 0 1 / 0 Evaluation weights (zq,zn)(z_q,z_n) (0.4, 0.9) (0.5, 0.8) (1, 1) (1, 1) (0, 1) Bootstrapped Q0Q_0 no / yes no / yes no no no Online planners. AdaOPS, DESPOT, and POMCPOW run with the same number of simulations per decision and particles as BETS in every domain, and otherwise with the settings and heuristics from Moss et al. (2024) on RockSample and LightDark. They use the preferred-actions rollout policy of Silver and Veness (2010) on PocMan. The exit policy moves the rover east to the exit without sampling any rock, and on LightDark r is the reward for a correct stop (Table 5). Table 5: AdaOPS, DESPOT, and POMCPOW settings per domain. RockSample LightDark PocMan AdaOPS (Wu et al., 2021) Lower, upper bound exit policy, QMDP† −r-r, r – mminm_ , δ 100, 0.1 10, 1.0 – Maximum depth 90 90 – DESPOT (Somani et al., 2013) Lower, upper bound exit policy, QMDP† random policy, r preferred actions, fixed Scenarios K, regularization λ 100, 0 30, 0.1 30, 0.01 Maximum depth 90 POMCPOW (Sunberg and Kochenderfer, 2018) Leaf value exit policy‡ MDP value preferred-actions rollout UCB constant 10 10 2 2 Action widening (ka,αa)(k_a, _a) – (10, 0.5) – Observation widening (ko,αo)(k_o, _o) (1, 1) (4, 0.03) (1, 1) Maximum depth 90 20 100 † Fixed upper bound Vmax=138.3V_ =138.3 on RockSample(20,20). ‡ Random rollouts on RockSample(20,20). Appendix C Generator Architectures C.1 Autoregressive Transformer The generator of Section 3.3 is a causal transformer that keeps one token per variable of the sample, and token j emits the logits of the factor pψ(xj∣x<j,θ)p_ψ(x_j x_<j,θ). Each token’s input is the sum of three terms: a learned positional embedding for token j, an embedding of the preceding value xj−1x_j-1 with a learned start token in place of x0x_0 at j=1j=1, and a conditioning vector computed from θ. VCjVC_j∥\|μ jjx<jx_<jcausal blockslogits for xjx_j Figure 3: Per-variable conditioning of the autoregressive generator. Token j combines the coupling summary VCjVC_j of variable j with the pooled embedding μ shared by every token. The causal blocks attend to the preceding values x<jx_<j, and the output is the logits of the factor pψ(xj∣x<j,θ)p_ψ(x_j x_<j,θ). Pooled conditions the generator on one vector. A single linear map sends θ to the token dimension and the result is added to every token identically. The conditioning is the same at every token. VC instead splits θ into the pooled embedding μ and the D per-variable coupling summaries, and applies one linear map, shared by every token, to the concatenation [μ;VCj][μ\,;\,VC_j] at token j (Figure 3). Token j and summary j index the same variable, so factor j sees the shared pooled statistics and the summary for the variable it predicts. C.2 Flow-Matching LightDark’s hidden state is a real-valued position, so its generator is a conditional flow-matching model (Lipman et al., 2023; Tong et al., 2023). A velocity field, an MLP with three hidden layers of width 64 that takes the current sample, the flow time, and θ as input, transports noise to a sample of the belief in 50 integration steps. Its belief loss is the conditional flow-matching regression at the realized state, and its embedding uses the pooled stream. Appendix D Wall-Clock Time We measure the wall clock time used by BETS, averaged over 100 episodes. Standard errors are below 3% of the mean. Each simulation applies the NBF update, the policy network, and the value network once, so search takes roughly K times as long per decision as the policy network alone (Table 6). The policy network alone, which earns more return than every baseline on RockSample(20,20) and both PocMan mazes (Table 1), takes under 25 ms per decision on every configuration. Table 6: Milliseconds per decision of BETS on one NVIDIA GeForce RTX 5080. search Domain n K=50K=50 K=100K=100 K=200K=200 policy RockSample(15,15) 256 140 229 460 2.8 512 248 410 812 4.5 RockSample(20,20) 256 176 349 688 4.0 512 323 641 1289 6.6 LightDark(5) 256 237 454 901 5.6 512 452 866 1722 9.7 LightDark(10) 256 238 455 902 5.6 512 452 867 1724 9.7 PocMan(M) 256 253 517 1042 5.6 512 425 879 1729 9.2 PocMan(L) 256 894 1876 3722 19.6 512 1107 2320 4692 24.1 • Each K is rounded up to a multiple of |A||A|, so that every round of the search selects every root action once.