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

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://github.com/LeahGL/RAD.

Many AI systems learn by trial and error, but in real-world settings such as robots, healthcare, or autonomous driving, trying unsafe actions can be costly or dangerous. Offline reinforcement learning avoids this by training agents only from previously collected data, but these fixed datasets often leave gaps: when an agent starts from an unfamiliar situation, it may not know how to reach a good outcome. We propose RAD, a method that helps the agent look up useful past experiences while it is making decisions. Instead of relying only on what was learned during training, RAD searches the reachable states that led to high rewards and uses them as intermediate goals. A planning model then helps the agent move toward these goals step by step. This is similar to asking for directions from an unfamiliar place to a promising landmark, rather than trying to solve the whole route at once. Experiments across navigation, locomotion, and manipulation tasks show that RAD often improves performance, especially when the agent faces unfamiliar starting states. This work may help offline decision-making systems make better use of limited data in settings where collecting new experience is expensive or risky.