- The paper presents ADWM, an autoregressive diffusion framework that factorizes policy-guided trajectories for effective off-policy evaluation of LLM agents.
- It integrates a diffusion prior, action-posterior guidance, and a continuation factor to mitigate compounding errors in long-horizon simulations.
- Experiments across diverse benchmarks show ADWM achieving a mean rank correlation of 0.82, outperforming classical OPE methods significantly.
Autoregressive Diffusion World Models for Off-Policy Evaluation of LLM Agents
Problem Setting and Limitations of Existing OPE Approaches
As the application space for LLM agents expands to complex, multi-turn, high-stakes domains, reliable off-policy evaluation (OPE) becomes necessary to anticipate agent behavior without costly and potentially unsafe online rollouts. Existing approaches to OPE suffer from fundamental limitations: importance sampling methods are hampered by exponentially increasing variance with trajectory length, while direct model-based methods using learned simulators frequently encounter catastrophic compounding error—small errors in rollout predictions are amplified across long horizons due to the autoregressive nature of world simulations. Moreover, classical diffusion-based OPE methods, successful in continuous control, are fundamentally at odds with agentic LLM settings where actions are discrete and causally depend on step-wise observations. This circular dependency in LLM agent evaluation is unresolved by prior work and precludes their application to evaluation of text-generating policies.
Method: Autoregressive Diffusion World Models (ADWM)
The core technical contribution is ADWM, a framework that leverages a latent diffusion world model tailored for sequential, agentic LLM settings. The key insight is an exact, autoregressive factorization of the policy-guided trajectory law into single-step conditionals. Each conditional admits a principled decomposition as the product of (i) a diffusion prior over environment transitions, (ii) an action-posterior enforcing stepwise action consistency, and (iii) a continuation factor aligning generated environments with the long-horizon behavior of the evaluation policy.
Guided diffusion is used to simulate the environment in latent space, with the evaluation policy actively steering the denoising process at each step via gradients of its log-likelihood over actions conditioned on world model state projections (via a learned ψ-adapter). The combination of autoregressive generation with policy-conditioned guidance allows the evaluation LLM to participate causally in rollouts without retraining the world model, resolving the circular dependency that defeats both classical model-based and diffusion-based OPE approaches for LLMs.
The architecture consists of:
- a text-to-latent encoder trained end-to-end for action-aware, environment-grounded representation,
- a structured diffusion model operating on latent states with classifier-free guidance for integrating action information,
- a ψ-adapter that projects latents to soft tokens in the evaluation LLM’s input embedding space, and
- terminal/reward prediction heads for value estimation.
Theoretical Insights
A critical theoretical result is the exact, stepwise factorization (Theorem 1), which (i) enables the decomposition of the target distribution into three additive scores corresponding precisely to the prior, action-posterior, and continuation guidance, and (ii) demonstrates that the required guidance can be delivered through gradients with respect to the latent, with no need for reweighting or explicit policy-density ratios. Notably, the continuation factor, intractable in general, is efficiently approximated via Monte Carlo world model rollouts from denoised latents.
The effect of each component can be modulated, interpolating smoothly between different guidance regimes, and ADWM overall subsumes behavior-regularization, action-consistency, and policy-alignment as controllable factors in the generative model.
Experimental Evaluation
Experiments span a diverse suite of LLM-agent benchmarks that exercise environments (HotpotQA, ScienceWorld, ALFWorld, WebShop) and policy families (RLHF-style, fine-tuned iterates, task transfer) under a variety of reward structures (dense F1, sparse/binary, shaped partial, continuous), configured to enforce significant train-eval behavior policy mismatch. For each, granular ablations and baselines are compared.
The primary evaluation is Spearman rank correlation between estimated and ground-truth policy ordering across ε-greedy mixes. ADWM achieves strictly positive correlation (minρ=0.67, mean $0.82$) in all settings, outperforming all classical OPE baselines by at least +0.34 mean correlation. Importance Sampling and Doubly Robust estimators collapse under LLM policy mismatch, producing high-variance or degenerate solutions, while FQE and DM are inconsistent, yielding inverted or policy-invariant estimates in multiple settings. Notably, even when supplying behavior log-probabilities unavailable in practical settings, classical methods are unreliable.
Ablation studies establish the indispensability of all three ADWM guidance components: action-posterior guidance is critical for sparse reward, continuation guidance for long-horizon, goal-directed rollouts, and the ψ-adapter for linguistic environments. The combination is essential for robustness; omitting any yields substantial performance degradation or even perfect negative correlation in select environments.
The practical implications of latent capacity and diversity of offline behavior data are also systematically analyzed, showing that both are prerequisites for positive transfer and ranking fidelity.
Practical and Theoretical Implications
ADWM enables policy evaluation of LLM agents from purely offline data in settings with severe distribution shift, inaccessible behavior policy densities, or costly environment interactions, as commonly encountered in practical deployments. Crucially, it achieves this without trajectory-level importance reweighting, diagnosis of behavior support, or retraining. The approach is fundamentally robust to long-horizon settings that defeat standard estimators.
From a theoretical standpoint, the factorization and corresponding guidance mechanism unify and extend behavior regularized and guided world modeling, providing a template for integrating black-box, non-differentiable, or modular policies into generative environment model sampling. This reinforces the view that generative modeling of environment dynamics, when aligned with policy structure and causality, can overcome the limitations of value-based or density-ratio-based OPE.
Future Directions
Potential avenues for future research include:
- Extension to cross-family or cross-architecture evaluation, relaxing the requirement for a ψ-adapter tied to the evaluation LLM,
- Scaling world model capacity and adaptation to reinforcement learning in high-dimensional, compositional action spaces,
- Generalization to offline RL algorithms that exploit such policy-guided, non-compounding world models for safe policy improvement,
- Exploration of richer, multi-agent or partially observable environments where latent space factorization and guidance may play a critical role.
Further, bridging the gap to closed-source or proprietary APIs via adapter distillation, and extending the fidelity of offline evaluation beyond reward estimation to safety, fairness, and robustness assessments, offers substantial practical impact.
Conclusion
ADWM establishes a new state-of-the-art methodology for OPE of LLM agents by introducing an exact, stepwise, policy-guided diffusion framework that is robust to the fundamental pathologies of long-horizon, agentic policy evaluation. Its empirical dominance over classical baselines and its principled, modular design recommend it as the practical standard for offline evaluation and analysis of next-generation LLM-based interactive agents (2606.05558).
