PROWL: Prioritized Regret-Driven Optimization for World Model Learning

2026-05-11

AI Summary

This work addresses the fragility of existing action-conditioned video world models in handling rare yet critical state transitions, a limitation stemming from passive demonstration datasets that inadequately cover high-impact scenarios. To overcome this, the authors propose a KL-constrained adversarial curriculum learning framework that actively explores high-error trajectories in diffusion-based world models through an adversarial policy. Central to this approach is a Prioritized Adversarial Trajectory (PAT) replay buffer, which dynamically reweights samples based on prediction error, action fidelity, and learning progress to continuously focus optimization on model weaknesses. Evaluated on the MineRL benchmark, the method substantially enhances robustness to out-of-distribution trajectories, uncovers latent reward-hacking behaviors under weak behavioral constraints, and demonstrates the efficacy of selectively generating high-information data for world model training.

Abstract

Modern action-conditioned video world models achieve strong short-horizon visual realism, yet remain unreliable on rare, interaction-critical transitions that dominate downstream planning and policy performance. Because passive demonstration data systematically under-samples these high-impact regimes, improving robustness requires actively eliciting model failures rather than relying on their natural occurrence. We introduce a KL-constrained adversarial curriculum in which a policy is trained to expose high-error trajectories of a diffusion-based world model while remaining close to the behavior distribution. The world model is continuously fine-tuned on these adversarially discovered trajectories, yielding an adversarial training loop that converts rare failures into a stable, near-distribution training signal without drifting into out-of-distribution exploitation. To maintain pressure on unresolved weaknesses as the model improves, we propose a Prioritized Adversarial Trajectory (PAT) buffer that re-ranks trajectories based on prediction error, action fidelity, and learning progress, focusing training on unresolved failure modes rather than repeatedly revisiting solved cases. We implement our approach in the MineRL framework and evaluate it on held-out out-of-distribution trajectories; PROWL improves robustness over models trained on passive data alone, reveals reward-hacking behaviors under weak behavioral constraints, and demonstrates that effective adversarial world-model training critically depends on balancing exploratory failure discovery with explicit behavioral regularization. Our results suggest that scalable world models benefit not only from larger datasets, but also from selectively generating informative training data.

Problem

• world model learning
• rare transitions
• robustness
• adversarial training
• behavioral distribution

Innovation

• adversarial world model learning
• prioritized trajectory replay
• KL-constrained policy
• diffusion-based world model
• active failure discovery