Your Belief About the World Is a Graph. Now Your Planner Can Use It.

How encoding uncertainty as graph structure enables POMDP planners to generalize far beyond their training size. A deep dive into GammaZero.

Series context & running example. Within the Planning Under Uncertainty series, GammaZero is the learning strategy to Part 3 (HOO-POMDP)'s abstraction — and the partially-observable cousin of GABAR from the sibling Learning for Planning series. The running example throughout is warehouse-delivery — robot, packages, zones, the robot seeing only its current zone. GammaZero's experiments use RockSample, MultiObjectSearch, and other relational POMDPs, but the mental model stays the same warehouse: build a belief graph, score actions, act, observe, repeat.

The Setup: Planning When You Can't See Everything

Imagine you're a Mars rover. You can see the terrain directly in front of you, but you can't see what's behind the next ridge. You have a noisy sensor that gives you partial readings about whether a rock formation is scientifically valuable. You need to decide: do you spend time scanning that distant rock (which might be worthless), or do you move toward the exit to meet your deadline?

This is a Partially Observable Markov Decision Process (POMDP). Unlike fully observable problems where you know exactly what state you're in, here you maintain a belief—a probability distribution over possible states. Every action you take both changes the world and updates your belief about it.

The challenge is severe. Even for small problems, solving POMDPs exactly is computationally intractable. The gold-standard approach is Monte Carlo Tree Search (MCTS)—build a search tree by simulating possible futures. But MCTS struggles with long planning horizons. Without good heuristics to guide the search, it can't look far enough ahead to find rewarding action sequences that require extended information gathering.

Recent work like BetaZero has shown that neural networks can learn to guide MCTS effectively, replacing hand-crafted heuristics with learned value functions and policies. The approach works: train a network on expert demonstrations, then use it to prioritize actions and estimate values during search.

But there's a catch.

The Representation Bottleneck

BetaZero and similar approaches represent belief states as fixed-dimensional vectors—statistical summaries of the particle set. This creates a rigid architecture that must be redesigned for each problem size.

Consider a robot searching for objects on a grid:

A network trained on a 5x5 grid with 3 objects literally cannot process a 10x10 grid with 10 objects—the input dimensions don't match. You need to retrain from scratch, which requires running expensive expert planners on the larger problems to generate training data. But those expert planners are exactly what you're trying to avoid using.

This creates a circular dependency: you need the planner to train the network, but you need the network to make the planner tractable.

What if the representation itself could grow with the problem, while the patterns learned on small instances still applied?

The lineage: how GammaZero got here

GammaZero is the convergence point of three lines of work. To see why each design choice was forced rather than arbitrary, it helps to walk the lineage.

AlphaZero — learning to guide MCTS in fully observable games

AlphaZero (Silver et al., 2018) demonstrated that learned value and policy heads can replace the hand-crafted heuristics inside Monte Carlo tree search. The architecture is a network f_θ(s) = (V_θ(s), P_θ(a|s)) that takes a state and predicts both its value and a prior distribution over actions. During search, the PUCT selection rule weights the standard exploration term by the learned prior, so simulations concentrate on actions the policy judges promising. At leaf nodes, the value prediction replaces the rollout return directly — one forward pass instead of a noisy average over many random trajectories.

The training procedure is self-play. The network guides MCTS, the search produces an improved policy (the visit counts at the root), and the network is updated toward that improved policy and toward the observed returns. The improved policy then guides the next round of search. AlphaZero's success in Go, chess, and shogi without any hand-crafted heuristics established this as the dominant recipe for MCTS-driven game playing.

MuZero (Schrittwieser et al., 2020) later extended the template to settings where the transition model is not known, by jointly learning a latent dynamics model alongside the value and policy heads. This made AlphaZero-style learning applicable to Atari and other domains where ground-truth simulators aren't available.

Both AlphaZero and MuZero live in the fully observable setting. The state is observable; the network input is a fixed-shape board (or a fixed-shape frame). The recipe doesn't transfer to POMDPs out of the box because the belief space is high-dimensional and continuous — there is no fixed-shape input to feed the network.

BetaZero — AlphaZero for POMDPs, with fixed-dimensional beliefs

BetaZero (Moss et al., 2024) is the current state-of-the-art instantiation of the AlphaZero template for POMDPs and the most direct prior work for GammaZero. It trains neural networks that predict values and action probabilities from belief states, and uses those networks to guide MCTS at decision time. Training proceeds AlphaZero-style: run the planner on small instances, use the search to produce improved policies, update the network toward those policies and toward the observed returns. The benchmark results are strong — on LightDark, RockSample, and similar problems, BetaZero matches or exceeds classical online solvers while dramatically reducing per-decision cost.

The catch is the belief representation. BetaZero summarizes the particle-based belief into a fixed-dimensional feature vector: per-state probability estimates, attribute-wise marginals, or hand-engineered summary features. The dimensionality is chosen once per problem size. A network trained on RockSample with an 11×11 grid and 11 rocks cannot be applied to a 15×15 grid with 15 rocks, because the belief vector has a different length and the network's first-layer weights are tied to that length. The very problems where learning could help most — large instances where classical planners struggle — are exactly the ones that fixed-dimensional models trained on small instances cannot handle.

ConstrainedZero (Moss et al., 2025) extends BetaZero to chance-constrained POMDPs where the agent must satisfy safety constraints with high probability. It adds heads that predict constraint satisfaction and modifies the search to prune actions whose estimated constraint satisfaction falls below a threshold. LeTS-Drive (Cai et al., 2022) applies the learning-guided search template to autonomous driving, training the belief estimator and planning network jointly. Both extensions preserve BetaZero's fixed-dimensional belief encoder; both inherit its size-generalization limitation.

GABAR — the fully-observable cousin

If BetaZero is the obvious POMDP precursor, GABAR from the sibling series is the obvious representational precursor. GABAR demonstrated that a graph neural network over a relational planning state generalizes across instance sizes — the same trained network solves instances 8× larger than anything it trained on, without retraining. The graph topology grows with the problem; the network's weights are shared across nodes; size generalization comes for free.

GABAR is fully observable. The mapping from state to graph is immediate: each ground atom either holds or does not. There is no notion of probability over nodes. To carry the recipe into the partially observable setting, the graph construction itself has to change to encode uncertainty.

The convergence

GammaZero sits at the intersection. From AlphaZero/MuZero: the value+policy network template and the PUCT selection rule that exploits it. From BetaZero: the application to POMDPs and the particle-based belief substrate. From GABAR: the graph-based representation that handles variable-sized inputs through GNN message passing. The contribution — what no prior method had done — is the construction that encodes belief uncertainty in the graph topology itself, so the same network can guide MCTS on POMDPs of any size.

The Core Insight: Belief States Are Graphs

Here's the key idea behind GammaZero: a belief state—that probability distribution over possible worlds—has relational structure. Objects have attributes. Actions connect to objects. Uncertainty creates specific patterns of connectivity between what you know and what you don't.

Instead of flattening this structure into a fixed vector, we preserve it as a graph:

The graph has four types of nodes:

• Object nodes — entities in the environment (robot, rocks, locations)
• Attribute nodes — properties that hold with sufficient probability (e.g., "IsGood(Rock2) = 0.75")
• Action nodes — available actions connected to their parameter objects
• A global node — aggregates information across the entire graph

The critical innovation: attribute nodes are only created when belief probability exceeds a threshold tau. If you're 75% sure a rock is good, that belief becomes a node. If you're only 10% sure, it doesn't. This means the graph structure itself encodes what you believe—not just numerical features on nodes, but which nodes exist.

Adding more objects to the problem simply adds more nodes and edges to the graph. The same Graph Neural Network (GNN) processes graphs of any size using identical weights. Patterns learned on small graphs transfer directly to larger ones.

Why This Representation Works: The Three Properties

Not every graph representation would enable generalization. GammaZero's works because of three specific properties:

1. Structural Encoding of Uncertainty

Most approaches encode uncertainty as numbers—a 15-dimensional vector of belief statistics. GammaZero encodes it structurally. When a rock's quality is uncertain, both "IsGood(R1) = 0.4" and "IsBad(R1) = 0.6" exist as separate nodes. When you're certain, only one exists. The GNN learns to recognize these structural patterns:

• Two complementary attribute nodes with similar weights → high uncertainty → information gathering is valuable
• A single attribute node with high weight → low uncertainty → commitment is safe

These patterns are size-invariant. The same uncertainty signature around Rock 1 in a 5x5 grid means the same thing around Rock 15 in a 25x25 grid.

2. Action-Centric Connectivity

Each action node connects to exactly the objects it operates on. Check(R1) connects to R1. Sample(R3) connects to R3. MoveEast connects to the robot. This makes action evaluation a local computation—the GNN can assess action quality by examining the neighborhood of each action node.

The attribute-action edges are especially powerful: they directly encode which attributes are relevant to which actions. If IsGood(R3) has high belief and is connected to Sample(R3), the network learns this is a good sampling opportunity—regardless of how many other rocks exist in the problem.

3. Global Context Through Aggregation

The global node connects to all object nodes and maintains a holistic representation of the belief state. This provides "shortcut" information propagation: even in a large graph where two nodes might be far apart, they both communicate through the global node at every message-passing round.

This is crucial for decisions that require global context—like deciding whether to gather more information or head for the exit. The local neighborhood of "MoveEast" shows the robot's position, but the global node aggregates how many high-value rocks remain unsampled across the entire grid.

The Full System: From Graphs to Decisions

GammaZero operates in two phases:

Offline: Learn from Small Problems

1. Run optimal planners on small POMDP instances (5x5 grids, 3-5 objects)
2. Collect belief states encountered during planning, along with optimal actions and values
3. Convert each belief to a graph using the construction above
4. Train a GNN to predict both values V(G) and action distributions P(a|G)

Training takes 2-4 hours on a single GPU. The training data is essentially free—small problems are solved in milliseconds by existing planners.

Online: Guide Search on Large Problems

During deployment on larger problems (15x15 grid, 15 rocks), the learned GNN enhances MCTS in three ways:

The GNN processes the larger graph using the exact same weights trained on small instances. No retraining, no architecture changes, no new expert data needed.

Results: Generalization That Actually Works

The experiments answer two questions: Can GammaZero match existing methods on same-sized problems? And can it generalize to larger ones?

Same-Size Performance

When trained and tested on identical problem sizes, GammaZero matches or exceeds BetaZero across all domains:

This validates that the graph representation doesn't sacrifice performance—it captures everything the fixed-dimensional approach captures, and more.

Zero-Shot Generalization

The unique capability: GammaZero can generalize to problems significantly larger than training instances. BetaZero cannot—its architecture physically won't accept larger inputs.

Key observations:

• Graceful degradation: Performance drops gradually as problem size increases, rather than catastrophically failing
• Competitive without retraining: On RockSample(20,20), GammaZero (trained only on small instances) achieves 87% of AdaOPS's performance (which was specifically configured for that problem size)
• Classical planners fail: At RockSample(25,25) and MOS(8,6), DESPOT and POMCPOW simply timeout, while GammaZero still produces reasonable solutions
• Raw policy suffices at extreme scales: On RS(25,25), the raw learned policy (without MCTS) achieves the best result, suggesting the learned structural patterns are robust even when search becomes too expensive

The Ablation Story: What Makes It Work?

The ablation results reveal a clear hierarchy of what matters most. Let me walk through each component:

Full MCTS + Value + Policy: The Gold Standard

The complete system combines three learned components with online search. Each plays a distinct role:

The policy network alone achieves 60-80% of full performance—demonstrating that the graph representation genuinely captures action quality. The value network performs slightly worse in isolation, which makes sense: estimating absolute values is harder than ranking relative action quality (echoing findings from fully observable planning).

But the combination through MCTS always wins. Search provides error correction: even when the policy's top choice is wrong, MCTS explores alternatives and uses the value network to course-correct.

Comparing Approaches: A Mental Model

To understand where GammaZero sits in the landscape, here's how it compares to existing approaches:

GammaZero is essentially the POMDP analog of GABAR (our prior work on fully observable planning), but with two critical additions: belief-weighted attributes that encode uncertainty, and integration with MCTS for online error correction.

The Deeper Lessons: Invariants for Other Research

Beyond the specific results, GammaZero demonstrates principles that generalize broadly:

1. Structure Your Input to Match Your Problem's Structure

POMDPs have relational structure: objects have attributes, actions operate on objects, beliefs assign probabilities to propositions. The graph representation preserves all of this. A flat vector discards it.

The principle applies beyond POMDPs: if your problem has entities, relations, and properties, your representation should make these first-class. Don't flatten structure into features—you're forcing the network to reconstruct what you already know.

2. Encode Uncertainty Structurally, Not Just Numerically

Instead of a feature vector [0.4, 0.6, 0.75, ...] encoding belief probabilities, GammaZero creates/removes graph nodes based on belief. The topology changes as beliefs change. This means:

• A network can learn structural patterns ("two competing attribute nodes = high uncertainty")
• These patterns are scale-invariant (same local structure, regardless of graph size)
• The graph is naturally sparse (only plausible hypotheses get nodes)

This applies to any domain with uncertainty: medical diagnosis (create symptom nodes only for plausible conditions), autonomous driving (instantiate obstacle hypotheses only above detection threshold), portfolio optimization (represent asset categories only when position is significant).

3. Train Cheap, Deploy Expensive

GammaZero trains on problems that cost milliseconds to solve optimally. It deploys on problems that would cost hours. The training data is essentially free—just run a planner on small instances.

This works because the patterns are compositional: how to evaluate a single rock's check-vs-sample tradeoff on a 5x5 grid is the same tradeoff on a 25x25 grid. The GNN learns these local patterns, and they compose into global policies.

The broader lesson: if your problem has compositional structure, you can often train on trivial instances and deploy on hard ones. The key is choosing a representation with the right inductive biases—in this case, GNNs that process local neighborhoods identically regardless of global size.

4. Combine Learning with Search

The policy network alone achieves 60-80% of full performance. Adding MCTS closes the gap to 100%. This isn't surprising—learned policies make mistakes, and lookahead search catches them.

But the reverse is also true: MCTS alone (without learned guidance) performs poorly on long-horizon POMDPs because it can't search deeply enough. The learned heuristics focus search where it matters.

This learning + search combination is a general recipe: learning provides fast approximate guidance, search provides deliberate error correction. Neither alone is sufficient for hard problems.

5. The Global Node Trick

A fixed-depth GNN (3 layers in GammaZero's case) has a limited receptive field. In a large graph, distant nodes can't communicate through local message passing alone. The global node solves this by acting as a communication hub—every node reads from it and writes to it every round.

This is the graph equivalent of the [CLS] token in transformers. If you're using GNNs on variable-size graphs, a global aggregation node is essentially free and dramatically improves scalability.

What This Means for Different Audiences

For POMDP/Planning Researchers

GammaZero shows that graph-based generalization—previously limited to fully observable domains—can work in belief space. The belief-threshold mechanism (only instantiate nodes above tau) is a clean, principled way to handle the continuous nature of beliefs within a discrete graph framework.

Practical implication: you can solve "toy" versions of your POMDP domain (which are cheap), train GammaZero on those, and get reasonable policies for the full-scale version without any additional expert computation.

For ML Researchers

The representation design is the key contribution from an ML perspective. The insight that belief uncertainty can be encoded topologically (through node existence) rather than just numerically (through features) is powerful. It converts a continuous estimation problem into a discrete structure recognition problem—and GNNs are good at structure.

The graceful degradation results are also notable. Most ML models fail catastrophically on out-of-distribution inputs. GammaZero degrades gracefully—performance drops proportionally to the scale increase, not suddenly. This suggests the learned features are genuinely capturing transferable patterns rather than memorizing training distributions.

For Robotics/Applications Researchers

The practical value is clear: train once on simulation-easy instances, deploy on the real (hard) problem. For domains like robotic search, autonomous exploration, or sensor placement, this means you can develop planning heuristics without needing to solve the full-scale problem even once during development.

The approach also handles multiple POMDP domains with the same architecture—only the graph construction rules change. RockSample, MultiObjectSearch, and Rearrangement all use the same GNN weights with different domain-specific node/edge definitions.

Limitations and Open Questions

GammaZero has clear limitations worth noting:

• Discrete actions and attributes: The current graph construction assumes discrete attribute values. Continuous state spaces would require discretization or a different node construction strategy.
• Threshold sensitivity: The tau threshold for node creation is a hyperparameter. Too low creates cluttered graphs; too high misses important beliefs. Domain-adaptive thresholds could help.
• Training data quality: Performance depends on having good expert demonstrations. For domains where even small instances are hard to solve optimally, training data generation becomes a bottleneck.
• Performance at extreme scales: While degradation is graceful, performance does drop at 6x scale. Hierarchical graph representations or curriculum-based training might push the boundary further.

Summary

GammaZero demonstrates that representing belief states as uncertainty-aware graphs enables POMDP planners to generalize across problem scales—something no prior learning-based POMDP approach could do.

The recipe is concrete:

1. Convert particle beliefs to graphs with belief-weighted attribute nodes
2. Train a GNN on small, cheaply-solved problem instances
3. Use the learned policy and value function to guide MCTS on arbitrarily larger problems

The result: competitive performance on same-sized problems, graceful generalization to 2-6x larger instances, and the ability to solve problems where classical planners simply timeout.

The broader insight: when your problem has relational structure and uncertainty, encode both in your representation—structurally, not just numerically. Let the architecture match the problem. And train cheap, deploy expensive.

Paper: "GammaZero: Learning to Guide Belief-Space Search for Long-Horizon POMDPs with Generalizable Graph Representations"

arXiv: 2510.14035