Диффузионная модель мира

Аннотация

Мы представляем Diffusion World Model (DWM) — условную диффузионную модель, способную одновременно предсказывать многократные будущие состояния и награды. В отличие от традиционных одношаговых моделей динамики, DWM обеспечивает долгосрочные прогнозы за один прямой проход, устраняя необходимость в рекурсивных запросах. Мы интегрируем DWM в оценку значений на основе моделей, где краткосрочный возврат моделируется с помощью будущих траекторий, сэмплированных из DWM. В контексте обучения с подкреплением на оффлайн-данных DWM можно рассматривать как консервативную регуляризацию значений через генеративное моделирование. Альтернативно, её можно рассматривать как источник данных, который позволяет проводить оффлайн Q-обучение с использованием синтетических данных. Наши эксперименты на наборе данных D4RL подтверждают устойчивость DWM к долгосрочному моделированию. По абсолютной производительности DWM значительно превосходит одношаговые модели динамики с приростом производительности на 44% и достигает результатов, соответствующих современным стандартам.

English

We introduce Diffusion World Model (DWM), a conditional diffusion model capable of predicting multistep future states and rewards concurrently. As opposed to traditional one-step dynamics models, DWM offers long-horizon predictions in a single forward pass, eliminating the need for recursive quires. We integrate DWM into model-based value estimation, where the short-term return is simulated by future trajectories sampled from DWM. In the context of offline reinforcement learning, DWM can be viewed as a conservative value regularization through generative modeling. Alternatively, it can be seen as a data source that enables offline Q-learning with synthetic data. Our experiments on the D4RL dataset confirm the robustness of DWM to long-horizon simulation. In terms of absolute performance, DWM significantly surpasses one-step dynamics models with a 44% performance gain, and achieves state-of-the-art performance.