WPT: World-to-Policy Transfer via Online World Model Distillation

TL;DR

WPT introduces an online world model distillation method that improves policy learning and planning efficiency by integrating world knowledge into lightweight policies, achieving state-of-the-art results in driving benchmarks.

Contribution

The paper presents a novel online distillation framework that transfers world model knowledge into a lightweight policy for real-time decision making.

Findings

Achieves a 0.11 collision rate in open-loop tests.

Attains a 79.23 driving score in closed-loop benchmarks.

Increases inference speed by up to 4.9 times.

Abstract

Recent years have witnessed remarkable progress in world models, which primarily aim to capture the spatio-temporal correlations between an agent's actions and the evolving environment. However, existing approaches often suffer from tight runtime coupling or depend on offline reward signals, resulting in substantial inference overhead or hindering end-to-end optimization. To overcome these limitations, we introduce WPT, a World-to-Policy Transfer training paradigm that enables online distillation under the guidance of an end-to-end world model. Specifically, we develop a trainable reward model that infuses world knowledge into a teacher policy by aligning candidate trajectories with the future dynamics predicted by the world model. Subsequently, we propose policy distillation and world reward distillation to transfer the teacher's reasoning ability into a lightweight student policy, enhancing planning performance while preserving real-time deployability. Extensive experiments on both open-loop and closed-loop benchmarks show that our WPT achieves state-of-the-art performance with a simple policy architecture: it attains a 0.11 collision rate (open-loop) and achieves a 79.23 driving score (closed-loop) surpassing both world-model-based and imitation-learning methods in accuracy and safety. Moreover, the student sustains up to 4.9x faster inference, while retaining most of the gains.