PROWL: Prioritized Regret-Driven Optimization for World Model Learning

TL;DR

PROWL introduces an adversarial curriculum and prioritized failure replay to enhance the robustness of world models in rare, critical scenarios, improving their reliability for downstream planning tasks.

Contribution

It presents a novel KL-constrained adversarial training method with a prioritized failure buffer to actively discover and learn from rare model failures.

Findings

PROWL improves robustness of world models on out-of-distribution trajectories.

It reveals reward-hacking behaviors under weak constraints.

Effective training depends on balancing failure discovery with behavioral regularization.

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.