PriorZero: Unified LLM and World Model Planning

• PriorZero is a unified decision-making framework that integrates static LLM semantic priors with dynamic world models to overcome the prior–dynamics mismatch.
• It employs root-only injection in Monte Carlo Tree Search to fuse LLM scores with world-model policies, ensuring efficient multi-step planning.
• The decoupled rollout-training design enables stable learning and faster convergence, yielding superior performance on text-based adventure games and gridworld tasks.

PriorZero is a unified framework for decision making that integrates LLM-derived conceptual priors into world-model-based planning through a decoupled rollout-training design. It is formulated to address a “prior–dynamics mismatch” in which LLMs encode rich, static commonsense knowledge about actions and outcomes in text, whereas reinforcement-learning agents must learn environment-specific transition dynamics, assign long-term credit over sparse rewards, and plan multi-step trajectories under uncertain feedback. In PriorZero, semantic guidance from the LLM is injected only at the root node of Monte Carlo Tree Search, while the world model performs deep lookahead and is trained separately from LLM adaptation; the framework is evaluated on text-based adventure games in Jericho and instruction-following gridworld tasks in BabyAI (Xiong et al., 12 May 2026).

1. Prior–dynamics mismatch and problem setting

PriorZero is motivated by a structural mismatch between static LLM priors and dynamic environment models. LLMs encode rich, static commonsense knowledge about actions and outcomes in text but lack any notion of environment-specific transition dynamics. Reinforcement-learning agents, by contrast, must learn a model of state–action transitions, assign long-term credit over sparse rewards, and plan multi-step trajectories under uncertain feedback (Xiong et al., 12 May 2026).

The framework is positioned against three failure modes. If one simply uses the LLM as a fixed policy (“Naive Policy”), exploration becomes myopic: high language-model likelihood does not imply high environment return, so the agent may loop on locally plausible but ultimately unproductive actions. If one fine-tunes the LLM end-to-end with RL (“Naive RLFT”), gradients are extremely high-variance in long-horizon tasks with sparse signals; credit assignment fails, and the model may catastrophically unlearn its prior knowledge. If one uses the LLM as a frozen text encoder, representations remain non-adaptive and impose a ceiling on performance.

The central design requirement is therefore twofold: semantic guidance should be injected where it helps exploration most, and the LLM should adapt to environment feedback without destabilizing dynamics learning. This suggests a decomposition in which planning and representation learning remain anchored in a learned world model, while the LLM supplies structured action priors under controlled conditions.

2. Overall framework and two-phase organization

PriorZero integrates an LLM prior with a learnable latent world model under Monte Carlo Tree Search through a decoupled rollout–training design. The framework is organized into a rollout phase and a training phase. During rollout, at each real environment step, the historical context is encoded into a latent state representation $z_t$ via a Transformer world model and into a textual prompt $C_t$ for the LLM. The LLM scores admissible actions to produce $\pi_{\mathrm{LLM}}(a \mid C_t)$ through chain-of-thought reasoning. This prior is fused with the world-model policy $\pi_{\mathrm{WM}}$ only at the MCTS root node, and MCTS rollouts then proceed purely under the world model to compute an improved search policy $\pi_{\mathrm{MCTS}}(a \mid z_t)$. An action $a_t \sim \pi_{\mathrm{MCTS}}(a \mid z_t)$ is executed, and $$(o_t, a_t, r_t, \pi_{\mathrm{MCTS}}, \pi_{\mathrm{LLM}})$$ is stored in the replay buffer (Xiong et al., 12 May 2026).

During training, PriorZero alternates between world-model updates and LLM updates. World-model updates are performed more frequently, for example “10× more frequently,” and minimize a composite objective that combines dynamics prediction, reward prediction, value estimation over a discretized support, and policy distillation from MCTS. LLM updates use stored rollouts and low-variance $n$-step advantages derived from world-model value estimates to optimize a PPO-style clipped policy-gradient objective regularized by a KL term to a reference policy.

The textual block diagram given for the method is:

$$[\mathrm{Environment}] \to [\mathrm{World\ Model}\ (z_t \leftarrow h_\theta(H_t))]$$

with MCTS receiving inputs from $\pi_{\mathrm{WM}}$ and $C_t$0 through root fusion, after which the sampled action is executed in the environment and the resulting rollout is stored in the replay buffer. Alternation between world-model training and LLM training then repeats.

A plausible implication is that PriorZero treats the LLM neither as a standalone policy nor merely as a static feature extractor, but as a structured source of action priors whose influence is mediated by search and by a separately trained latent dynamics model.

3. Root-prior injection in Monte Carlo Tree Search

The defining algorithmic mechanism in PriorZero is root-prior injection. At timestep $C_t$1, the world-model latent state is formed as $C_t$2, and a $C_t$3-tempered LLM prior $C_t$4 is obtained. The fused root prior is

$C_t$5

For non-root nodes in MCTS, the prior is set to the world-model policy alone:

$C_t$6

Node selection uses PUCT / pUCT:

$C_t$7

where at the root $C_t$8 and elsewhere $C_t$9. After $\pi_{\mathrm{LLM}}(a \mid C_t)$0 rollouts under the world model, the search policy is extracted as

$\pi_{\mathrm{LLM}}(a \mid C_t)$1

The rationale for root-only injection is explicit. The root state $\pi_{\mathrm{LLM}}(a \mid C_t)$2 aligns with a real textual context $\pi_{\mathrm{LLM}}(a \mid C_t)$3, whereas deeper nodes are imagined states with no faithful text mapping. Root-only fusion therefore preserves the world model’s multi-step lookahead free of spurious textual bias in rollout dynamics. It focuses search on semantically plausible actions while letting the world model adjudicate final value (Xiong et al., 12 May 2026).

This design distinguishes PriorZero from schemes that inject language priors throughout the search tree. The reported ablations indicate that always injecting at all nodes leads to degraded dynamics modeling because of a granularity mismatch, whereas no injection causes catastrophic collapse in Detective and Zork1. Root-only injection is reported as the best trade-off.

4. Decoupled optimization of world model and LLM

The world model in PriorZero is trained by minimizing

$\pi_{\mathrm{LLM}}(a \mid C_t)$4

The four components are, respectively, dynamics loss, reward loss, value loss over a discretized support, and policy distillation from the search policy. This objective jointly refines the model’s dynamics, policy, and value predictions (Xiong et al., 12 May 2026).

LLM fine-tuning is based on a low-variance $\pi_{\mathrm{LLM}}(a \mid C_t)$5-step advantage:

$\pi_{\mathrm{LLM}}(a \mid C_t)$6

with

$\pi_{\mathrm{LLM}}(a \mid C_t)$7

The LLM parameters $\pi_{\mathrm{LLM}}(a \mid C_t)$8 are updated through a PPO-style objective

$\pi_{\mathrm{LLM}}(a \mid C_t)$9

where

$\pi_{\mathrm{WM}}$0

The paper further states that advantages are normalized and can optionally be blended with a small format-compliance reward. Alternating optimization proceeds by interleaving rollout collection, multiple world-model updates, and multiple LLM updates from replayed trajectories.

The decoupling is significant because it separates noisy long-horizon return optimization from dynamics learning. This suggests that value estimates produced by the world model serve as an intermediate credit-assignment mechanism for LLM adaptation, rather than propagating sparse end-task reward directly through the full LLM parameter space.

5. Empirical performance on Jericho and BabyAI

Experiments compare PriorZero with UniZero, described here as a baseline with no LLM prior, across four Jericho text-based games and 18 BabyAI levels. On Jericho, the reported comparison is as follows (Xiong et al., 12 May 2026):

Game	PriorZero Early Steps to 80% Return	UniZero Steps
Detective	5K steps	8K
Acorncourt	10K	25K
Zork1	30K (asymptotic +15% higher score)	50K
OmniQuest	12K	35K

The key metrics summarized for Jericho are that PriorZero reaches high-return regions 2–3× faster, achieves asymptotic returns with a +10–20% improvement on sparse, long-horizon tasks, and exhibits root-node entropy that is initially lower and later higher than UniZero, corresponding respectively to more focused search and better exploration spread.

On BabyAI, the aggregated return over 18 levels is reported as 0.82 for PriorZero versus 0.79 for UniZero. On compositional levels, specifically SynthLoc, UniZero stalls near 0.0 while PriorZero reaches 0.96. A standalone fine-tuned LLM without a world model often fails on multi-step tasks, which the paper presents as evidence for the necessity of planning.

These results are presented as improvements in both exploration efficiency and asymptotic performance. A plausible interpretation is that root-level semantic priors are most useful in shaping the search frontier early in training, while the world model becomes increasingly responsible for long-horizon evaluation as learning progresses.

6. Ablations, stability analyses, and capacity effects

The ablation studies target the three principal components of the method: root-prior injection, alternating optimization, and LLM-mediated reasoning signals. For injection frequency, the paper reports that no injection leads to catastrophic collapse in Detective and Zork1, always injecting at all nodes degrades dynamics modeling, and root-only injection provides the best trade-off (Xiong et al., 12 May 2026).

For LLM fine-tuning stability, a frozen LLM prior yields early gains but plateaus at suboptimal performance. Naive RLFT from scratch leads to training collapse due to uninitialized value estimates. By contrast, an alternating schedule with initial world-model warm-up produces stable improvement of the LLM’s standalone policy over training.

The chain-of-thought loss weight is also examined. The reported outcomes are: $\pi_{\mathrm{WM}}$1 produces numerical instability (NaN), $\pi_{\mathrm{WM}}$2 yields volatile learning dynamics and poorer final return, and $\pi_{\mathrm{WM}}$3 is the default and gives stable and best performance.

Model-capacity analysis compares Qwen2.5-7B and Qwen2.5-3B. The larger LLM accelerates convergence and raises asymptotic return on Zork1, while on the simpler Detective task the 7B model exhibits higher variance but similar peak performance. The role of MCTS is isolated by disabling lookahead, which causes performance collapse; model-free LLM priors alone cannot learn multi-step credit assignment.

Collectively, these ablations support a narrow claim: the gains attributed to PriorZero do not arise from language priors alone, nor from world-model planning alone, but from the specific combination of root-only fusion, MCTS lookahead, and decoupled alternating training.

7. Limitations, scope, and relation to broader research directions

The stated limitations are threefold. First, the benchmarks are limited to discrete-action, text-based and grid-world domains, and extension to continuous or visual tasks is identified as needed. Second, the method relies on API calls to a large LLM, so compute and latency overhead remain significant. Third, the fusion weight $\pi_{\mathrm{WM}}$4 is fixed and may not adapt optimally across states or tasks (Xiong et al., 12 May 2026).

The future extensions proposed in the paper include adaptive root-prior weighting based on world-model uncertainty or entropy, hybrid modalities that integrate visual encoders with LLM priors for embodied control, hierarchical injection combining LLM guidance at abstract levels with low-level control, and meta-learning schedules for balancing world-model and LLM updates.

A common misconception is that PriorZero reduces decision making to prompt-conditioned action scoring. The reported failures of a standalone fine-tuned LLM and of MCTS-disabled variants indicate otherwise: the framework depends on explicit latent dynamics learning and search. Another potential misconception is that stronger language priors should be injected everywhere in the search tree. The root-only design and its ablations argue against this, on the grounds that imagined states deeper in the tree lack faithful textual grounding.

Within the scope tested, PriorZero is best understood as a method for bridging static language priors and environment-specific world models. Its core claim is not that LLMs can replace planning, but that LLM-derived semantic priors can be made operational for long-horizon decision making when coupled to root-only MCTS fusion and decoupled alternating training.