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

Yihao Qin; Yuanfei Wang; Hang Zhou; Peiran Liu; Hao Dong; Yiding Ji

cs.LGMar 4, 2026

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

Authors: Yihao Qin, Yuanfei Wang, Hang Zhou, Peiran Liu, Hao Dong, Yiding Ji

Organizations: The Hong Kong University of Science and Technology (Guangzhou) Guangzhou, China · Peking University Beijing, China

Abstract

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 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.

Figures & tables

Figure 1. An overview of IPD. The process begins by learning a world model with uncertainty measure and a quasi-optimal value function from the original offline suboptimal dataset. Suboptimal states are identified using the value function, and their corresponding trajectories are replaced with imaginary rollouts generated via Model Predictive Control, using the learned world model and value function. Each generated trajectory is evaluated for uncertainty before incorporated into the enhanced dataset. Finally, a Transformer-based sequential policy is trained on this MPC-enhanced dataset, with additional supervision from the value function via action gradients and a dynamic return-to-go. By distilling the imaginary planning, which streamlines both MPC and dynamic programming, into the Transformer based policy, IPD enables the generation of superior actions.

Gym Task	IQL	CQL	DT	QDT	DD	EDT	QT	Reinformer	IPD
walker-medium	78.3	83.0	74.0	67.1	82.5	72.8	87.6	80.5	89.5 $\pm$ 3.7
walker-medium-replay	73.9	77.2	79.4	58.2	68.9	74.8	94.2	72.9	96.2 $\pm$ 2.1
hopper-medium	66.3	69.4	67.6	66.5	79.3	64.5	78.0	63.5	81.6 $\pm$ 4.8
hopper-medium-replay	94.7	95.0	82.7	52.1	100.0	89.0	102.1	83.3	103.2 $\pm$ 3.7
halfcheetah-medium	47.4	49.2	42.6	39.3	49.1	42.5	49.1	42.9	51.2 $\pm$ 2.6
halfcheetah-medium-replay	44.2	45.5	36.6	35.6	39.3	37.8	48.9	39.0	49.9 $\pm$ 1.6

Table 1. Evaluating IPD and State-of-the-Art Q-Learning and Transformer-Based Sequential Policy Methods on D4RL Benchmarks: Performance Analysis Through Tenfold Episode Evaluations

Figure 2. Comparison between MPC and Greedy Q-Learning data augmentation in Walker2d-medium-replay task.

Figure 3. Performance Comparison between quasi-optimal value function and different setting of return-to-go in halfcheetah-medium-replay task.

Value Function	Value	World Model	Value
$\tau$	0.7	$\Sigma_{\text{reg}}$	$e^{-4.5}\cdot I$
$\beta$	3	$\gamma_{0}$	1.0
$\delta$	0.96	$T$	$10^{5}$
$\gamma$	0.99	$\alpha_{c}$	0.5
Dropout	0.01	$\alpha_{r}$	0.5

Table 2. Hyperparameters of Quasi-Optimal Value Function and World Model with Uncertainty Measure Learning.

Parameters	Value
$\gamma_{\text{mpc}}$	0.99
$H_{I}$	10
$H_{m}$	10
$H_{\text{con}}$	10
$N_{\text{mpc}}$	3

Table 3. Hyperparameters of Imaginary Data Augmentation.

Parameters	Value
Layers	4
Dropout	0.01
Embedding	256
Attention heads	4
Context length	20

Table 4. Hyperparameters of Sequential Policy

Offline RL

Offline reinforcement learning learns decision-making policies from fixed datasets without interacting with the environment during training.

See all 298 papers

RAD: Retrieval High-quality Demonstrations to Enhance Decision-making
Jul 21, 2025Lu Guo, Yixiang Shan, Zhengbang Zhu +5Offline RL Generalization in RL
School of Artificial Intelligence, Jilin University, Changchun, China · School of Computer Science, Shanghai Jiao Tong University, Shanghai, China
Q-learning Penalized Transformer for Safe Offline Reinforcement Learning
Sep 28, 2026Shengchao Hu, Peng Wang, Jifeng Hu +5Reinforcement Learning Transformer-Based RL
Shanghai Jiao Tong University, China · Shenzhen Campus of Sun Yat-sen University, China · Jilin University, China +1
Value-Rectified Distillation for Flow-based Offline Reinforcement Learning
Oct 21, 2026Ke Jiang, Wen Jiang, Yoshinobu Kawahara +1Offline RL Policy Distillation
College of Computer Science and Technology, Nanjing University of Aeronautics and Astronautics · MIIT Key Laboratory of Pattern Analysis and Machine Intelligence · Graduate School of Information Science and Technology, The University of Osaka +1

Jul 21, 2025J/K move · Enter open · S save

RAD: Retrieval High-quality Demonstrations to Enhance Decision-making

Lu Guo, Yixiang Shan, Zhengbang Zhu, Qifan Liang+4

School of Artificial Intelligence, Jilin University, Changchun, China · School of Computer Science, Shanghai Jiao Tong University, Shanghai, China

Offline RL Generalization in RL

Open paper AskPDF arXiv

Offline reinforcement learning (RL) learns policies from fixed datasets, thereby avoiding costly or unsafe environment interactions. However, its reliance on finite static datasets inherently restricts the ability to generalize beyond the training distribution. Prior solutions based on synthetic data augmentation often fail to generalize to unseen scenarios in the (augmented) dataset. To address these challenges, we propose Retrieval High-quAlity Demonstrations (RAD) for decision-making, which innovatively introduces a retrieval mechanism into offline RL. Specifically, RAD retrieves high-return and reachable states from the offline dataset as target states, and leverages a generative model to generate sub-trajectories conditioned on these targets for planning. Since the targets are high-return states, once the agent reaches such a target, it can continue to obtain high returns by following the associated high-return actions, thereby improving policy generalization. Extensive experiments confirm that RAD achieves competitive or superior performance compared to baselines across diverse benchmarks, validating its effectiveness. Our code is available at https://anonymous.4open.science/r/RAD_0925_1-690E.

Abstract

Offline RL

Explore similar work

RAD: Retrieval High-quality Demonstrations to Enhance Decision-making

Q-learning Penalized Transformer for Safe Offline Reinforcement Learning

Value-Rectified Distillation for Flow-based Offline Reinforcement Learning