PriorZero — Injecting LLM Priors into MuZero-Style World Models at the MCTS Root
Paper. Junyu Xiong, Yuan Pu, Jia Tang, Yazhe Niu. PriorZero: Bridging Language Priors and World Models for Decision Making. arXiv preprint, 2026. [arXiv]
0. The Picture in One Paragraph
MuZero-family agents (and their scalable descendant, UniZero) learn a latent world model and plan over it with MCTS — powerful, but they learn everything about the environment’s dynamics and semantics from scratch through interaction. LLMs, in contrast, carry rich static world knowledge but no grounding in a specific environment’s actual transition dynamics — injecting their priors naively risks a prior-dynamics mismatch that corrupts the agent’s own learned lookahead. PriorZero resolves this with a deliberately narrow injection point: an LLM policy prior (computed from prompt log-probabilities over admissible actions, using a chain-of-thought prefix) is blended into the action distribution only at the MCTS root node. Every deeper simulation step — dynamics rollout, value backup, expansion — remains purely world-model-driven. The LLM cheaply biases which actions get explored first without ever touching the model’s own learned notion of how the environment actually behaves. On sparse-reward Jericho text-adventures and the compositional- generalization BabyAI suite, this gives PriorZero a clear edge over plain UniZero — including reaching 0.96 on a compositional level (SynthLoc) where UniZero alone learns almost nothing.
1. The Problem — Prior-Dynamics Mismatch
An LLM’s world knowledge is static and general; a MuZero-style world model’s knowledge is learned, task-specific, and grounded in actual observed transitions. If an LLM prior is injected deeply into a planning process — say, at every internal MCTS node — it risks steering the search toward actions that sound plausible in general but are wrong for the actual (possibly idiosyncratic) dynamics of the specific environment, corrupting the value estimates the world model would otherwise have produced through honest simulation.
2. The PriorZero Method
Three gradient-decoupled components:
- UniZero world model + latent MCTS — the base scalable MuZero-family planner, unchanged in its own dynamics/value/policy learning.
- LLM prior module — scores admissible action strings via prompt log-probabilities, using a chain-of-thought prefix, normalized with a softmax (temperature 1) into a smooth prior distribution over the actions available at the current state.
- Alternating trainer — the world model continually updates on interaction data (dynamics, policy, value losses) while its own value estimates, in alternation, provide a fine-grained credit-assignment signal used to fine-tune the LLM.
Root-prior injection
The core design decision: the LLM’s policy prior is blended into the action distribution only at the root node of each MCTS search. Internal nodes, dynamics rollouts, and value backups are left entirely to the world model. This means the LLM’s contribution is limited to biasing which branches get explored/expanded first from the current real state — a cheap, low-risk way to inject semantic knowledge without letting a potentially wrong static prior propagate through deep, compounding simulation steps.
3. Results
Jericho (text-adventure interactive fiction — Detective, Acorncourt, Zork1, Omniquest):
- On the sparse-reward games (Acorncourt, Omniquest), PriorZero reaches high-return regions substantially earlier than UniZero.
- On Detective, UniZero has an early advantage but PriorZero overtakes it as training progresses.
- On Zork1 (a large action space, >50 admissible actions/step, long horizon), PriorZero reaches a higher asymptotic return than UniZero.
BabyAI (18 instruction-following gridworld levels):
- PriorZero’s average asymptotic score across levels is 0.82 versus UniZero’s 0.79.
- On the compositional-generalization level SynthLoc, where UniZero learns almost nothing (near-zero), PriorZero reaches 0.96 — the clearest single result in the paper for where language priors help most: generalizing to novel action/object compositions the world model has not directly experienced.
LLM capacity ablation
Swapping the default Qwen2.5-3B-Instruct prior source for Qwen2.5-7B-Instruct converges faster and reaches a higher asymptotic score on Zork1 — the quality of the injected prior scales with the underlying LLM’s capacity, as one would hope.
4. Ablations
- Alternating vs. non-alternating fine-tuning schedule — confirms the alternating (decoupled rollout/training) design is preferable to jointly optimizing world model and LLM prior in lockstep.
- CoT-prefixed prior extraction with weighted token loss vs. plain action scoring — the CoT-prefixed variant is the one used in the main results.
- LLM capacity (3B vs. 7B) — larger LLM priors help, as above.
- Necessity of MCTS lookahead depth — confirms that PriorZero’s gains depend on retaining genuine world-model search depth, not just on the LLM prior alone; a shallow-search variant loses much of the advantage.
5. Relationship to Prior Work
PriorZero is built directly on UniZero, itself part of the MuZero/EfficientZero lineage and released alongside the LightZero benchmark codebase. Its contribution relative to that lineage is narrow and specific: rather than proposing a new world model or a new MCTS variant, it identifies the root node as the right, minimally-invasive injection point for external LLM priors — a design choice that is easy to describe but whose value is only demonstrated empirically through the sparse-reward and compositional generalization results above.
6. Why It Matters
MCTS-based planners are excellent at exploiting a well-learned world model but can be slow to discover good actions in sparse-reward or combinatorially large action spaces, precisely because undirected search has to stumble onto the right branch before the world model can even evaluate it well. PriorZero’s root-only injection is a minimal-risk way to use an LLM’s general knowledge to narrow that initial search — a design pattern that generalizes beyond text adventures to any MCTS-based agent operating in an action space large enough that undirected exploration is the actual bottleneck, not world-model accuracy.
7. Limitations Worth Knowing
- Evaluated only on text/language-grounded environments (Jericho, BabyAI) — no results on classic MCTS domains like Atari, board games, or continuous-control robotics, where action spaces are not naturally describable as admissible strings for an LLM to score.
- Inference-cost overhead of the added LLM scoring step at the MCTS root (extra LLM forward passes per search) is not reported in the sources reviewed here.
- Root-only injection may leave value on the table in settings where useful semantic guidance would help deeper in the tree, not just at the root — the paper doesn’t explore that trade-off space.
- Generalization beyond Jericho/BabyAI-scale action spaces is untested; both benchmarks are relatively small/discrete compared to many real-world decision-making settings.
8. The Takeaway for a First Reader
If you remember three things:
- Naively injecting LLM priors deep into MCTS risks corrupting the world model’s own learned lookahead — a prior-dynamics mismatch problem PriorZero avoids by design.
- Root-prior injection blends the LLM’s policy prior into the action distribution only at the MCTS root node, leaving all deeper simulation, value backup, and dynamics rollout purely world-model-driven.
- The gains are largest exactly where undirected search struggles most: sparse-reward text-adventures and compositional generalization (0.96 vs. near-zero on BabyAI’s SynthLoc level) — suggesting the technique targets exploration bottlenecks specifically, not general sample efficiency.
References
- Xiong, J., Pu, Y., Tang, J., & Niu, Y. (2026). PriorZero: Bridging Language Priors and World Models for Decision Making. arXiv:2605.12289.
- Related on this site: BG-MCTS — budget-aware MCTS for LLM reasoning, a complementary use of search budget rather than search guidance; Magellan — another MCTS-plus-LLM-guidance system, applied to creative idea generation rather than sequential decision-making.
Enjoy Reading This Article?
Here are some more articles you might like to read next: