Prioritized Generative Replay

TL;DR

The paper tackles sample efficiency in online reinforcement learning by addressing limitations of uniform replay through a parametric memory built from conditional diffusion models. It introduces Prioritized Generative Replay (pgr), which densifies past experience and guides generation via relevance functions, with curiosity emerging as a robust default cue to promote diverse, informative transitions. Empirical results across state-based and pixel-based tasks show that curiosity-guided pgr delivers consistent gains in data efficiency and performance, and scales effectively with larger policies and higher synthetic data budgets. The approach offers a flexible, scalable path to leveraging synthetic data in online RL while mitigating overfitting to generated samples and enabling higher update-to-data ratios.

Abstract

Sample-efficient online reinforcement learning often uses replay buffers to store experience for reuse when updating the value function. However, uniform replay is inefficient, since certain classes of transitions can be more relevant to learning. While prioritization of more useful samples is helpful, this strategy can also lead to overfitting, as useful samples are likely to be more rare. In this work, we instead propose a prioritized, parametric version of an agent's memory, using generative models to capture online experience. This paradigm enables (1) densification of past experience, with new generations that benefit from the generative model's generalization capacity and (2) guidance via a family of "relevance functions" that push these generations towards more useful parts of an agent's acquired history. We show this recipe can be instantiated using conditional diffusion models and simple relevance functions such as curiosity- or value-based metrics. Our approach consistently improves performance and sample efficiency in both state- and pixel-based domains. We expose the mechanisms underlying these gains, showing how guidance promotes diversity in our generated transitions and reduces overfitting. We also showcase how our approach can train policies with even higher update-to-data ratios than before, opening up avenues to better scale online RL agents.

Figures (10)

• Figure 1: We model an agent's online memory using a conditional diffusion model. By conditioning on measures of data relevance, we can generate samples more useful for policy learning.
• Figure 2: pgr improves performance by densifying subspaces of data where transitions more relevant for learning reside. We project 10K generations for both our pgr and the unconditional baseline synther to the same tSNE plot. A: At epoch 1, the distribution of data generated by pgr and synther are similar. B: At the inflection point of performance near epoch 130, pgr generates a distinct sub-portion of the data space from synther (i.e. red and blue dots are largely separate.) C: At the end of learning, pgr still densely covers a distinct subspace of the synther transitions.
• Figure 3: Comparison to baselines that use (a) prioritized experience replay (PER) and (b) exploration reward bonuses.redq using PER, with priority determined by curiosity \ref{['eqn:curiosity_f']} or TD-error \ref{['eqn:td_f']}, still underperform their pgr counterparts. pgr also remains superior after directly adding an exploration bonus in the form of curiosity to either synther or redq.
• Figure 4: Sample efficiency on DMC (a) state-based and (b) pixel-based tasks. We show mean and standard deviation over five seeds. Curiosity-pgr consistently demonstrates the best sample efficiency. synther, which uses unconditional generation to augment replay, underperforms model-free algorithms like sac and redq on harder sparse-reward tasks, like finger-turn-hard.
• Figure 5: pgr does not outperform baselines due to improved generation quality. We compute mean-squared error (MSE) of dynamics over 10K generated transitions for synther and curiosity-pgr across 3 OpenAI gym environments. Top: Average MSE. Bottom: Histograms of MSE values.