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

TL;DR

RAD introduces a retrieval and generative modeling approach to improve offline reinforcement learning by dynamically retrieving high-quality states and planning towards them, enhancing generalization and decision-making in complex environments.

Contribution

The paper presents a novel method combining retrieval and diffusion-based generative modeling to improve trajectory stitching and generalization in offline RL.

Findings

RAD outperforms baselines across multiple benchmarks.

Retrieval-guided generation enhances decision-making in OOD states.

The approach improves long-horizon planning in offline RL.

Abstract

Offline reinforcement learning (RL) enables agents to learn policies from fixed datasets, avoiding costly or unsafe environment interactions. However, its effectiveness is often limited by dataset sparsity and the lack of transition overlap between suboptimal and expert trajectories, which makes long-horizon planning particularly challenging. Prior solutions based on synthetic data augmentation or trajectory stitching often fail to generalize to novel states and rely on heuristic stitching points. To address these challenges, we propose Retrieval High-quAlity Demonstrations (RAD) for decision-making, which combines non-parametric retrieval with diffusion-based generative modeling. RAD dynamically retrieves high-return states from the offline dataset as target states based on state similarity and return estimation, and plans toward them using a condition-guided diffusion model. Such retrieval-guided generation enables flexible trajectory stitching and improves generalization when encountered with underrepresented or out-of-distribution states. Extensive experiments confirm that RAD achieves competitive or superior performance compared to baselines across diverse benchmarks, validating its effectiveness.