IPD: Boosting Sequential Policy with Imaginary Planning Distillation in Offline Reinforcement Learning

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Qin, Yihao, Wang, Yuanfei, Zhou, Hang, Liu, Peiran, Dong, Hao, Ji, Yiding
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917314454093824
author Qin, Yihao
Wang, Yuanfei
Zhou, Hang
Liu, Peiran
Dong, Hao
Ji, Yiding
author_facet Qin, Yihao
Wang, Yuanfei
Zhou, Hang
Liu, Peiran
Dong, Hao
Ji, Yiding
contents Decision transformer based sequential policies have emerged as a powerful paradigm in offline reinforcement learning (RL), yet their efficacy remains constrained by the quality of static datasets and inherent architectural limitations. Specifically, these models often struggle to effectively integrate suboptimal experiences and fail to explicitly plan for an optimal policy. To bridge this gap, we propose \textbf{Imaginary Planning Distillation (IPD)}, a novel framework that seamlessly incorporates offline planning into data generation, supervised training, and online inference. Our framework first learns a world model equipped with uncertainty measures and a quasi-optimal value function from the offline data. These components are utilized to identify suboptimal trajectories and augment them with reliable, imagined optimal rollouts generated via Model Predictive Control (MPC). A Transformer-based sequential policy is then trained on this enriched dataset, complemented by a value-guided objective that promotes the distillation of the optimal policy. By replacing the conventional, manually-tuned return-to-go with the learned quasi-optimal value function, IPD improves both decision-making stability and performance during inference. Empirical evaluations on the D4RL benchmark demonstrate that IPD significantly outperforms several state-of-the-art value-based and transformer-based offline RL methods across diverse tasks.
format Preprint
id arxiv_https___arxiv_org_abs_2603_04289
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle IPD: Boosting Sequential Policy with Imaginary Planning Distillation in Offline Reinforcement Learning
Qin, Yihao
Wang, Yuanfei
Zhou, Hang
Liu, Peiran
Dong, Hao
Ji, Yiding
Machine Learning
Artificial Intelligence
Decision transformer based sequential policies have emerged as a powerful paradigm in offline reinforcement learning (RL), yet their efficacy remains constrained by the quality of static datasets and inherent architectural limitations. Specifically, these models often struggle to effectively integrate suboptimal experiences and fail to explicitly plan for an optimal policy. To bridge this gap, we propose \textbf{Imaginary Planning Distillation (IPD)}, a novel framework that seamlessly incorporates offline planning into data generation, supervised training, and online inference. Our framework first learns a world model equipped with uncertainty measures and a quasi-optimal value function from the offline data. These components are utilized to identify suboptimal trajectories and augment them with reliable, imagined optimal rollouts generated via Model Predictive Control (MPC). A Transformer-based sequential policy is then trained on this enriched dataset, complemented by a value-guided objective that promotes the distillation of the optimal policy. By replacing the conventional, manually-tuned return-to-go with the learned quasi-optimal value function, IPD improves both decision-making stability and performance during inference. Empirical evaluations on the D4RL benchmark demonstrate that IPD significantly outperforms several state-of-the-art value-based and transformer-based offline RL methods across diverse tasks.
title IPD: Boosting Sequential Policy with Imaginary Planning Distillation in Offline Reinforcement Learning
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2603.04289