Technical Summary: PriorZero

Problem Statement

The integration of Large Language Models (LLMs) into Reinforcement Learning (RL) agents offers a promising avenue for leveraging rich world knowledge to solve long-horizon, sparse-reward decision-making tasks. However, existing approaches are hindered by a fundamental prior-dynamics mismatch.

Current paradigms face specific structural bottlenecks:

1. Naive Policy (P1): Using LLMs as zero-shot policies leads to myopic decisions because linguistic likelihood does not equate to environment return, creating a "knowing-doing gap."
2. Naive RL Fine-Tuning (P2): End-to-end fine-tuning of LLMs via RL suffers from high-variance credit assignment and optimization instability, particularly in sparse-reward environments.
3. Dynamics Modeling (P3): Using LLMs to predict state transitions fails due to a granularity mismatch; LLMs are pretrained on token-level text generation, whereas latent dynamics models require state-action level transition prediction.
4. Text Encoders (P4): Using frozen LLM features as static representations limits task-specific adaptation, imposing a performance ceiling.

The core challenge is bridging the gap between the static, semantic priors of LLMs and the dynamic, environment-specific transition models required for effective planning.

Methodology: PriorZero Framework

PriorZero proposes a unified framework that decouples the LLM from the world model, integrating them through a decoupled rollout-training design. The system places a learnable world model between the LLM and the environment, allowing the LLM to provide semantic guidance while the world model handles transition modeling, value estimation, and multi-step planning.

1. Root-Prior Injection (Rollout Phase)

During the rollout phase, PriorZero utilizes a Latent World Model (based on UniZero) to perform Monte Carlo Tree Search (MCTS) in a latent space.

• Mechanism: A novel Root-Prior Injection mechanism fuses the LLM's semantic action prior exclusively at the root node of the MCTS.
• Implementation: The LLM generates a situation analysis (Chain-of-Thought) and an action distribution ($\pi_{LLM}$) based on the current textual context. This is fused with the world model's policy ($\pi_{WM}$) at the root:
  $$P_{root}(a | z_t) = (1 - \alpha) \pi_{WM}(a | z_t) + \alpha \pi_{LLM}(a | C_t)$$
• Constraint: Non-root nodes in the search tree rely entirely on the world model's latent dynamics. This prevents the "granularity mismatch" of Paradigm P3, as the LLM is never forced to predict state transitions for imagined latent states.
• Benefit: This guides the search toward semantically promising actions, preventing the MCTS from falling into meaningless trajectories (e.g., reward loops) while preserving the world model's deep lookahead capability.

2. Alternating Reinforcement Fine-Tuning (Training Phase)

To address the credit assignment failure of Paradigm P2, PriorZero employs an Alternating Fine-Tuning strategy that decouples world-model learning from LLM adaptation.

• Schedule: The training follows an asymmetric schedule (e.g., 10 world-model updates for every 1 LLM update).
• World Model Learning: The world model is updated on interaction data to refine its dynamics, value, and policy predictions.
• LLM Adaptation: Instead of using high-variance on-policy returns, the LLM is fine-tuned using low-variance advantage signals derived from the world model's value estimates.
  • The advantage is calculated using an $n$-step bootstrapped TD target: $A(z_t, a_t) = Q_n(z_t, a_t) - v_\theta(z_t)$.
  • The LLM is updated via a clipped PPO objective, with a specific weighting scheme that applies full weight to action tokens and reduced weight ($w_{cot}=0.1$) to Chain-of-Thought reasoning tokens to maintain stability.
• Benefit: This creates a closed-loop where the world model provides stable, fine-grained supervision for the LLM, allowing the LLM's priors to co-evolve with the environment dynamics without the instability of end-to-end optimization.

Key Contributions

1. Analysis of Paradigms: The paper identifies the prior-dynamics mismatch as the central obstacle in transferring LLM knowledge to long-horizon RL, categorizing existing approaches (P1–P4) and their specific failure modes.
2. PriorZero Framework: A novel architecture that synergizes LLM priors with world-model planning via two core mechanisms:
   • Root-Prior Injection: Injects semantic guidance at the MCTS root to facilitate exploration without contaminating latent dynamics.
   • Alternating Reinforcement Fine-Tuning: Uses world-model value estimates to construct stable, low-variance credit assignment signals for LLM adaptation.
3. Empirical Validation: Comprehensive experiments on Jericho (text-based adventure games) and BabyAI (grid-world instruction following) demonstrate that PriorZero consistently improves exploration efficiency and asymptotic performance compared to baselines.

Experimental Results

• Jericho Benchmarks: PriorZero outperforms UniZero (a strong world-model baseline without LLM priors) across four distinct games (Detective, Acorncourt, Zork1, Omniquest).
  • In Zork1 (long-horizon, sparse rewards), PriorZero achieves higher asymptotic returns and faster convergence.
  • In Detective, PriorZero surpasses UniZero after initial gains, demonstrating the value of adapting priors to task-specific dynamics.
• BabyAI Benchmarks: On 18 levels of the BabyAI grid-world, PriorZero reaches higher return plateaus (0.82 vs. 0.79 for UniZero) and shows significant gains in compositional tasks (e.g., SynthLoc), where the LLM prior helps decompose goals into sub-tasks.
• Ablation Studies:
  • Alternating Schedule: Crucial for stability; non-alternating variants suffer from "knowledge corruption" due to noisy value signals.
  • Root-Only Injection: Essential for preventing the LLM from overriding the world model's dynamics during lookahead.
  • CoT Loss Weight: A balanced weight ($w_{cot}=0.1$) is required; removing CoT causes numerical instability, while full weight leads to volatile learning.
  • MCTS: Disabling MCTS causes performance collapse, confirming the necessity of planning-improved targets.

Significance and Claims

The paper claims that PriorZero establishes a promising framework for LLM-empowered decision-making by resolving the structural bottlenecks of prior transfer. Its significance lies in:

• Decoupling: Successfully separating the LLM's role (providing semantic priors) from the world model's role (learning dynamics and planning), thereby avoiding the granularity mismatch and credit assignment failures of previous methods.
• Stability: Enabling stable fine-tuning of LLMs in long-horizon, sparse-reward environments by leveraging the world model as a low-variance critic.
• Scalability: Demonstrating that the framework scales with LLM capacity (improving performance with 7B models in complex tasks) while maintaining robustness through the alternating training schedule.

The authors position PriorZero not as a replacement for world models or LLMs, but as a synergistic integration where the LLM focuses the search and the world model refines the LLM's priors through continuous interaction.