Prioritized Generative Replay

## Abstract

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.

## 1 Introduction

One central problem of online reinforcement learning (RL) lies in extracting signal from a continuous stream of experience. Not only does the non-i.i.d. form of this data induce learning instabilities, but agents may lose out on near-term experience that is not immediately useful, but becomes important much later. The standard solution to these problems is to use a replay buffer Lin (); Mnih et al. () as a form of memory. By storing a dataset of transitions to enable batch-wise sampling, the algorithm can decorrelate online observations and revisit past experience later in training.

However, the idea that an agent’s memory must identically reproduce past transitions, and at uniform frequency, is limiting. In general, the distribution of states an agent *visits* is different from the optimal distribution of states the agent should *train on*. Certain classes of transitions are more relevant to learning, i.e. data at critical decision boundaries or data that the agent has seen less frequently. Locating such long-tailed data is tricky, yet clearly important for efficient learning.
The challenge is thus to design a scalable memory system that replays *relevant* data, and at large quantities.

In this work, we propose a simple, plug-and-play formulation of an agent’s online memory that can be constructed using a generative model. We train a conditional generative model in an end-to-end manner to faithfully capture online transitions from the agent. This paradigm grants two key benefits, enabling 1) the densification of past experience, allowing us to create new training data that goes beyond the data distribution observed online, and 2) guidance via a family of “relevance functions“ $\mathcal{F}$ that push these generations towards more useful parts of an agent’s acquired experience.

An ideal relevance function $\mathcal{F}$ should be easy to compute, and naturally identify the most relevant experience. Intuitively, we should generate transitions at the “frontier” of the agent’s experience. These are regions of transition space that our generative model can accurately capture from the agent’s memory, but our policy has not yet completely mastered. We investigate a series of possible functions, concluding that *intrinsic curiosity* reliably approximates this distribution.

Figure: 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.
Refer to caption: x1.png

Our main contribution is this framework for a scalable, guidable generative replay, which we term “Prioritized Generative Replay” (pgr). We instantiate this framework by making use of strong diffusion model architectures. Experiments on both state-based and pixel-based RL tasks show pgr is consistently more sample-efficient than both model-free RL algorithms and generative approaches that do not use any guidance. In fact, by densifying the more relevant transitions, pgr is able to succeed in cases where unconditional generation struggles significantly. Moreover, we empirically demonstrate that pgr goes beyond simple prioritized experience replay; in particular, we show that conditioning on curiosity leads to more diverse and more learning-relevant generations.
Finally, we show how pgr improves with larger policy networks, and continues learning reliably with higher synthetic-to-real data ratios, setting up a promising recipe for data-efficient scaling.

## 2 Related Work

Model-based RL. Model-based reinforcement learning involves interacting with a predictive model of the environment, sometimes referred to as a world model (Ha & Schmidhuber, ), to learn a policy (Sutton, ). Two classes of approaches dominate: planning with the learned world model directly (Hafner et al., ; Schrittwieser et al., ; Ye et al., ; Chua et al., ; Ebert et al., ), or optimizing a policy by unrolling trajectories “in the imagination” of the world model (Janner et al., ; ; Hafner et al., ; Oh et al., ; Feinberg et al., ). This latter approach is most relevant to our problem setting. One distinction is we do not backprop or plan actions through a model directly, but rather wholly synthesize additional, high-relevance transitions. Closely related is PolyGRAD (Rigter et al., ) which generates entire on-policy trajectories. However, by generating *independent off-policy transitions* rather than attempting to approximate trajectories from the current policy, we avoid the issue of compounding errors in our model (Gu et al., ), and also enable easier plug-and-play with popular existing online RL methods.

Prioritized replay. Agents often benefit significantly more by learning from important experiences (Moore & Atkeson, ; Andre et al., ). One of the most well-known strategies which leverages this insight is prioritized experience replay (PER), which uses temporal difference (TD) error as a priority criterion to determine the relative importance between transitions (Schaul et al., ). Since then, a number of works have proposed many different criteria to determine transition importance, including state entropy (Ramicic & Bonarini, ), learnability (Sujit et al., ), or some combination of value/reward and TD-error (Cao et al., ; Gao et al., ). To alleviate the computational cost associated with evaluating the priority of all experiences in the replay buffer, different flavors of algorithmic improvements (Kapturowski et al., ; Schaul et al., ) or approximate parametric models of prior experience (Shin et al., ; Novati & Koumoutsakos, ) have been proposed. Our proposal to model the replay buffer as a parametric generative model is closer to this second class of works.
However, distinct from these methods which seek to relabel *existing* data, our method uses priority to generate entirely *new* data for more generalizable learning.

RL from synthetic data. The use of generative training data has a long history in reinforcement learning (Shin et al., ; Raghavan et al., ; Imre, ; Hafner et al., ). But only with the recent advent of powerful diffusion models have such methods achieved parity with methods trained on real data (Janner et al., ; Ajay et al., ; Zhu et al., ). The use of diffusion for RL was first introduced by Diffuser (Janner et al., ). The authors propose an offline diffusion model for generating trajectories of states and actions, and directly generate plans which reach goal states or achieve high rewards. Decision Diffuser (Ajay et al., ) generates state-only trajectories, and employs an inverse kinematics model to generate corresponding actions. More recently, works like (Ding et al., ; He et al., ) leverage diffusion models to augment datasets for offline RL, improving training stability. Most similar to our approach is synther (Lu et al., ), which augments the replay buffer online and can be seen as an unguided form of our framework. In contrast to these previous works, we posit that the strength of synthetic data models is that they can be guided towards novel transitions that are off-policy yet more relevant for learning.

## 3 Background

This section offers a primer on online reinforcement learning and diffusion models. Here, details on conditional generation via guidance are especially important to our method exposition in Section 4.

Reinforcement learning.
We model the environment as a fully-observable, infinite-horizon Markov Decision Process (MDP) (Sutton & Barto, ) defined by the tuple $\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},p_{0},\gamma)$. Here, $\mathcal{S},\mathcal{A}$ denote the state and action spaces, respectively. $\mathcal{P}(s^{\prime}\mid s,a)$ for $s,s^{\prime}\in\mathcal{S}$ and $a\in\mathcal{A}$ describes the transition dynamics, which are generally not known. $\mathcal{R}(s,a)$ is a reward function, $p_{0}$ is the initial distribution over states $s_{0}$, and $\gamma$ is the discount function. In online RL, a policy $\pi:\mathcal{S}\rightarrow\mathcal{A}$ interacts with the environment $\mathcal{M}$ and observes tuples $\tau=(s,a,s^{\prime},r)$ (i.e. transitions), which are stored in a replay buffer $\mathcal{D}$. The *action-value*, or $\mathcal{Q}$, function is given by:

$$ $Q_{\pi}(s,a)=\mathbb{E}_{a_{t}\sim\pi(\cdot\mid s_{t}),s_{t+1}\sim\mathcal{P}( \cdot\mid s_{t},a_{t})}\left[\sum\limits_{t=0}^{\infty}\gamma^{t}\mathcal{R}(s _{t},a_{t}\mid s_{0}=s,a_{0}=a)\right].$ (1) $$

The $\mathcal{Q}$-function describes the expected *return* after taking action $a$ in state $s$, under the policy $\pi$. Then our goal is to learn the optimal policy $\pi^{*}$ such that $\pi^{*}(a\mid s)=\operatorname*{arg\,max}_{\pi}Q_{\pi}(s,a)$.

Conditional diffusion models.
Diffusion models are a powerful class of generative models which employ an iterative denoising procedure for generation (Sohl-Dickstein et al., ; Ho et al., ). Diffusion first involves a *forward process* $q(x^{n+1}\mid x^{n})$ iteratively adding Gaussian noise to $x^{n}$ starting from some initial data point $x^{0}$. A *reverse process* $p_{\theta}(x^{n-1}\mid x^{n})$ then transforms random Gaussian noise into a sample from the original distribution. In particular, we learn a neural network $\epsilon_{\theta}$ which predicts the amount of noise $\epsilon\sim\mathcal{N}(0,\mathbf{I})$ injected for some particular forward step $x_{n}$.

For additional controllability, diffusion models naturally enable conditioning on some signal $y$, simply by formulating the forward and reverse processes as $q(x^{n+1}\mid x^{n},y)$ and $p_{\theta}(x^{n-1}\mid x^{n},y)$, respectively. Classifier-free guidance (CFG) is a common post-training technique which further promotes sample fidelity to the condition $y$ in exchange for more complete mode coverage (Ho & Salimans, ). To facilitate CFG at sampling-time, during training, we optimize $\epsilon_{\theta}$ with the following objective:

$$ $\mathbb{E}_{x^{0}\sim\mathcal{D},\epsilon\sim\mathcal{N}(0,\mathbf{I}),y,n\sim \textrm{Unif}(1,N),p\sim\textrm{Bernoulli}\left(p_{\textrm{uncond}}\right)} \lVert\epsilon_{\theta}\left(x^{n},n,(1-p)\cdot y+p\cdot\emptyset\right)\rVert ^{2}_{2},$ (2) $$

where $p_{\textrm{uncond}}$ is the probability of dropping condition $y$ in favor of a null condition $\emptyset$. During sampling, we take a convex combination of the conditional and unconditional predictions, i.e. $\omega\cdot\epsilon_{\theta}(x^{n},n,y)+(1-\omega)\cdot\epsilon_{\theta}(x^{n}
,n,\emptyset)$, where $\omega$ is a hyperparameter called the *guidance scale*.

## 4 Prioritized Generative Replay

In this section, we introduce Prioritized Generative Replay (pgr). At its core, pgr involves a parametric generative replay buffer that can be guided by a variety of relevance criteria. We first provide intuition and motivation for such a framework, and concretize how it can be instantiated. Next, we compare and contrast between various instantiations of relevance functions. We argue that we should use functions which promote generating higher-novelty states as they implicitly increase the diversity of relevant transitions our replay buffer, leading to more sample-efficient learning.

### 4.1 Motivation

Consider an agent with policy $\pi$ interacting online with some environment $\mathcal{M}$. The agent observes and stores transitions $\tau=(s,a,s^{\prime},r)$ in some finite replay buffer $\mathcal{D}$. Two main issues arise:

- 1. $|\mathcal{D}|$ must be sufficiently large and diverse to prevent overfitting (hard to guarantee online)
- 2. Transitions more relevant for updating $\pi$ might be rare (and thus heavily undersampled)

To address the first problem, we can *densify* the replay distribution $p_{\mathcal{D}}(\tau)$ by learning a generative world model using the buffer. This effectively trades in our finitely-sized, non-parametric replay buffer for an infinitely-sized, parametric one. By leveraging the generalization capabilities of, e.g., diffusion models, we can interpolate the replay distribution to more impoverished regions of data.

However, uniformly sampling from this parametric buffer, in expectation, implies simply replaying past transitions at the same frequency at which they were observed. In more challenging environments, there may only a small fraction of data which is most relevant to updating the current policy $\pi$. Thus, to address the second problem and enable sample-efficient learning, we need some mechanism to not only *densify* $p_{\mathcal{D}}(\tau)$, but actually *guide* it towards the more immediately useful set of transitions.

Our key insight is to frame this problem via the lens of conditional generation. Specifically, we seek to model $p_{\mathcal{D}}(\tau\mid c)$, for some condition $c$.
By choosing the right $c$, we “marginalize out” parts of the unconditional distribution $p_{\mathcal{D}}(\tau)$ which are less relevant under $c$ (see Fig. 2.) More concretely, we propose to learn a *relevance function* $\mathcal{F}(\tau)=c$ jointly with the policy $\pi$.
Intuitively, this relevance function measures the “priority” $c$ of $\tau$.
Overall, we can achieve both more *complete* as well as more *relevant* coverage of $p_{D}(\tau)$.
Clearly, the choice of $\mathcal{F}$ in our framework is critical. We now explore a number of instantiations for $\mathcal{F}$,
and perform an analysis on their strengths and weaknesses.

Figure: 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.
Refer to caption: x2.png

### 4.2 Relevance Functions

We begin by outlining two desiderata for our relevance functions. First, given our online setting, these functions should incur minimal computation cost. This removes conditioning on strong priors which demand extensive pretraining (e.g. prior works in offline RL (Du & Narasimhan, ; Schwarzer et al., )). Second, we want a guidance signal that will not overfit easily, thereby collapsing generation to stale or low-quality transitions even as the agent’s experience evolves over time.

Return.
Our first proposal is one of the most natural: we consider a relevance function based on episodic return.
Concretely, we use the value estimate of the learned $Q$-function and current policy $\pi$:

$$ $\mathcal{F}(s,a,s^{\prime},r)=Q(s,\pi(s)).$ (3) $$

Return-as-relevance sensibly pushes generations to be more on-policy, since $\pi$ by construction seeks out high $Q$-estimates states. Also, many online RL algorithms already learn a $Q$-function, and so we readily satisfy the first condition (Mnih et al., ). However, the second condition is not adequately addressed by this choice of $\mathcal{F}$. In particular, the diversity of high-return transitions might in practice be quite low, making overfitting to the conditional generations more likely.

Temporal difference (TD) error.
Another possible choice for our relevance function is TD-error, first proposed for replay prioritization by Schaul et al. ().
Our relevance function in this case is given by the difference between the current Q-value estimate and the bootstrapped next-step estimate:

$$ $\mathcal{F}(s,a,s^{\prime},r)=r+\gamma Q_{\textrm{target}}\left(s^{\prime}, \operatorname*{arg\,max}\limits_{a^{\prime}}Q(s^{\prime},a^{\prime})\right)-Q( s,a),$ (4) $$

One immediate shortcoming is that in practice we have no guarantees on the $Q$-estimates for rarer, out-of-distribution transitions. In fact, overly greedy prioritization of high TD-error transitions can lead to low-quality, myopic $Q$-estimates (Schaul et al., ).

Curiosity.
A glaring problem with both return-based and TD error-based relevance functions is their reliance on high-quality $Q$-functions. Estimation errors can thus lead to $\mathcal{F}$ providing a poor conditioning signal. Moreover, online RL agents tend to overfit $Q$-functions to early experience (Nikishin et al., ), which will in turn lead to a rapidly overfitted $\mathcal{F}$ under these two choices.

We argue that a natural way to reduce overfitting is via some relevance function which promotes generation *diversity*.
As prior work has shown, an effective way to decrease overfitting to early, noisy signal in online RL is to leverage diverse experience (Zhang et al., ). To achieve this diversity, we model $\mathcal{F}$ after exploration objectives which promote engaging with “higher-novelty” transitions that are more rarely seen (Strehl & Littman, ). Moreover, by learning a separate function entirely, we decorrelate our relevance function from the $Q$-function, making overfitting less likely.

We thus turn to prior work on intrinsic motivation (Schmidhuber, ; Oudeyer & Kaplan, ) to operationalize these insights.
Concretely, we take inspiration from the intrinsic curiosity module (Pathak et al., ) to parameterize $\mathcal{F}$. Given a feature encoder $h$, we learn a forward dynamics model $g$ which models the environment transition function $\mathcal{P}(s^{\prime}\mid s,a)$, in the latent space of $h$. Then $\mathcal{F}$ is given by the *error* of this forward dynamics model:

$$ $\mathcal{F}(s,a,s^{\prime},r)=\frac{1}{2}||g(h(s),a)-h(s^{\prime})||^{2}.$ (5) $$

Figure: Algorithm 1 Overview of our outer loop + inner loop framework.

### 4.3 PGR Framework Summary

Finally, we provide a concrete overview of our framework in Algorithm 1. In the outer loop, the agent interacts with the environment, receiving a stream of real data and building a replay buffer $\mathcal{D}_{\textrm{real}}$, as in regular online RL. In the event that we are using a curiosity-based relevance function, we also perform an appropriate gradient update for $\mathcal{F}$ using samples from $\mathcal{D}_{\textrm{real}}$, via the loss function given by Eq. 5. Then periodically in the inner loop, we learn a conditional generative model $G$ of $\mathcal{D}_{\textrm{real}}$ and generatively densify these transitions to obtain our synthetic replay buffer $\mathcal{D}_{\textrm{syn}}$. Concretely, we take $G$ to be a conditional diffusion model. To leverage the conditioning signal given by $\mathcal{F}$, we use CFG and a “prompting” strategy inspired by Peebles et al. (). We choose some ratio $k$ of the transitions in the real replay buffer $\mathcal{D}_{\textrm{real}}$ with the highest values for $\mathcal{F}(s,a,s^{\prime},r)$, and sample their conditioning values randomly to pass to $G$. Implementation-wise, we keep both $\mathcal{D}_{\textrm{real}}$ and $\mathcal{D}_{\textrm{syn}}$ at 1M transitions, and randomly sample synthetic and real data mixed according to some ratio $r$ to train our policy $\pi$. Here, any off-policy RL algorithm can be used to learn $\pi$. For fair comparison to prior art, we instantiate our framework with both sac (Haarnoja et al., ) and redq (Chen et al., ).

## 5 Experiments

In this section, we answer three main questions. (1) First, to what extent does our pgr improve performance and sample efficiency in online RL settings? (2) Second, what are the underlying mechanisms by which pgr brings about performance gains? (3) Third, what kind of scaling behavior does pgr exhibit, if indeed conditionally-generated synthetic data under pgr is better?

Environment and tasks.
Our results span a range of state-based and pixel-based tasks in the DeepMind Control Suite (DMC) (Tunyasuvunakool et al., ) and OpenAI Gym (Brockman, ) environments.
In particular, our benchmark follows exactly the online RL evaluation suite of prior work in generative RL by Lu et al. (), facilitating direct comparison. In all tasks, we allow 100K environment interactions, a standard choice in online RL (Li et al., ; Kostrikov et al., ).

Model details, training and evaluation.
For state-based tasks,
in additional to the sac and redq model-free baselines, we also include two strong model-based online RL baselines in mbpo (Janner et al., ) and dreamer-v3 (Hafner et al., ). We also compare against synther (Lu et al., ), a recent work which learns an unconditional generative replay model, allowing sac to be trained with an update-to-data (UTD) ratio of 20. To facilitate conditional sampling, during training we randomly discard the scalar given by our relevance function $\mathcal{F}$ with probability 0.25.

For pixel-based tasks, our policy is based on drq-v2 (Yarats et al., ) as in Lu et al. (). To maintain the same approach and architecture for our generative model, we follow Lu et al. (); Esser et al. () and generate data in the latent space of the policy’s CNN visual encoder. That is, given a visual encoder $f_{\theta}$, and a transition $(s,a,s^{\prime},r)$ for pixel observations $s,s^{\prime}\in\mathbb{R}^{3\times h\times w}$ of height $h$ and width $w$, we learn to (conditionally) generate transitions $(f_{\theta}(s),a,f_{\theta}(s^{\prime}),r)$.

**Table 1: Average returns on state and pixel-based DMC after 100K environment steps (5 seeds, 1 std. dev. err.). * is a harder environment with sparser rewards, and so we present results over 300K timesteps.**
|  | DMC-100k (Online) | Pixel-DMC-100k (Online) |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Environment | Quadruped-<br>Walk | Cheetah-<br>Run | Reacher-<br>Hard | Finger-Turn-<br>Hard* | Walker-<br>Walk | Cheetah-<br>Run |
| mbpo | 505.91 $\pm$ 252.55 | 450.47 $\pm$ 132.09 | 777.24 $\pm$ 98.59 | 631.19 $\pm$ 98.77 | - | - |
| dreamer-v3 | 389.63 $\pm$ 168.47 | 362.01 $\pm$ 30.69 | 807.58 $\pm$ 156.38 | 745.27 $\pm$ 90.30 | 353.40 $\pm$ 114.12 | 298.13 $\pm$ 86.37 |
| sac | 178.31 $\pm$ 36.85 | 346.61 $\pm$ 61.94 | 654.23 $\pm$ 211.84 | 591.11 $\pm$ 41.44 | - | - |
| redq | 496.75 $\pm$ 151.00 | 606.86 $\pm$ 99.77 | 733.54 $\pm$ 79.66 | 520.53 $\pm$ 114.88 | - | - |
| redq + curiosity | 687.14 $\pm$ 93.12 | 682.64 $\pm$ 52.89 | 725.70 $\pm$ 87.78 | 777.66 $\pm$ 116.96 | - | - |
| drq-v2 | - | - | - | - | 514.11 $\pm$ 81.42 | 489.30 $\pm$ 69.26 |
| synther | 727.01 $\pm$ 86.66 | 729.35 $\pm$ 49.59 | 838.60 $\pm$ 131.15 | 554.01 $\pm$ 220.77 | 468.53 $\pm$ 28.65 | 465.09 $\pm$ 28.27 |
| pgr (Reward) | 510.39 $\pm$ 121.11 | 660.87 $\pm$ 87.54 | 715.43 $\pm$ 97.56 | 540.85 $\pm$ 73.29 | - | - |
| pgr (Return) | 737.62 $\pm$ 20.13 | 779.42 $\pm$ 30.00 | 893.65 $\pm$ 55.71 | 805.42 $\pm$ 92.07 | - | - |
| pgr (TD Error) | 802.18 $\pm$ 116.52 | 704.17 $\pm$ 96.49 | 917.61 $\pm$ 37.32 | 839.26 $\pm$ 49.90 | - | - |
| pgr (Curiosity) | 927.98 $\pm$ 25.18 | 817.36 $\pm$ 35.93 | 915.21 $\pm$ 48.24 | 885.98 $\pm$ 67.29 | 570.99 $\pm$ 41.44 | 529.70 $\pm$ 27.76 |

**Table 2: Results on state-based OpenAI gym tasks. We report average return after 100K environment steps. Results are over 3 seeds, with 1 std. dev. err.**
|  | Walker2d-v2 | HalfCheetah-v2 | Hopper-v2 |
| --- | --- | --- | --- |
| mbpo | 3781.34 $\pm$ 912.44 | 8612.49 $\pm$ 407.53 | 3007.83 $\pm$ 511.57 |
| dreamer-v3 | 4104.67 $\pm$ 349.74 | 7126.84 $\pm$ 539.22 | 3083.41 $\pm$ 138.90 |
| sac | 879.98 $\pm$ 217.52 | 5065.61 $\pm$ 467.73 | 2033.39 $\pm$ 793.96 |
| redq | 3819.17 $\pm$ 906.34 | 6330.85 $\pm$ 433.47 | 3275.66 $\pm$ 171.90 |
| synther | 4829.32 $\pm$ 191.16 | 8165.35 $\pm$ 1534.24 | 3395.21 $\pm$ 117.50 |
| pgr (Curiosity) | 5682.33 $\pm$ 370.04 | 9234.61 $\pm$ 658.77 | 4101.79 $\pm$ 244.05 |

In the interest of fair comparison, our diffusion and policy architectures mirror synther exactly. Thus, training FLOPs, parameter count and generation time as presented in Table 3 are all directly comparable to synther. Our only additional parameters lie in a lightweight curiosity head, which is updated for only $5\%$ of all policy gradient steps. pgr thus incurs minimal additional overhead.

### 5.1 Results

As shown in Table 1 and Table 3, all variants of pgr successfully solve the tasks, with curiosity guidance consistently outperforming both
model-free (sac, redq and drq-v2), and model-based (mbpo and dreamer-v3) algorithms, as well as the unconditional generation baseline (synther).

Comparison to prioritized experience replay (PER).
We also compare pgr to PER baselines that use different measures of priority. In particular, we train redq with a prioritized replay buffer, using
the classic TD-error (Schaul et al., ), or a priority function based on Eq. 5. The former is comparable to our pgr approach using Eq. 4 as relevance, and the latter to using Eq. 5. As shown in  Fig. 3a, pgr demonstrates superior performance in both cases. This emphasizes the importance of *densifying* the replay distribution with generations, and not simply reweighting past experience.

Figure: 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 Eq. 5 or TD-error Eq. 4, 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.
Refer to caption: x3.png

Figure: 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.
Refer to caption: x4.png

Comparison to exploration bonuses.
We examine how baselines improve when given a bonus in the form of an intrinsic curiosity reward (*c.f*. Eq. 5). In Fig. 3b, we see that curiosity-pgr continues to outperform synther or redq when either is augmented this way. This suggests pgr goes beyond just improving exploration.
We propose our gains are a result of generating higher novelty transitions from the replay buffer with higher diversity. This additional diversity thereby reduces overfitting of the $Q$-function to synthetic transitions. We provide further evidence for this idea in Section 5.2.

### 5.2 Sample Efficiency

In this section, we show pgr achieves its strong performance in a sample-efficient manner, and reveal the mechanisms underlying this sample-efficiency.
As seen in Fig. 4a, for tasks with state-based inputs, synther often attains its best performance around $100K$ environment steps. In contrast, curiosity-pgr is able to match this performance in both the cheetah-run and quadruped-walk tasks after only ${\sim}50K$ steps.
Especially noteworthy is the finger-turn-hard task, where synther actually underperforms compared to the vanilla model-free redq baseline, while our curiosity-pgr continues outperforming. We argue conditioning is particularly valuable in sparse reward tasks, where efficient exploration of the environment presents a significant challenge.
These observations hold for pixel-based tasks. As we see in Fig. 4b,
while synther is eventually overtaken by drq-v2 in both environments, our curiosity-pgr continues to consistently improve over drq-v2.

|  | Walker2d-v2 | HalfCheetah-v2 | Hopper-v2 |
| --- | --- | --- | --- |
| synther | 4.228 | 0.814 | 0.178 |
| pgr | 6.683 (+58.1%) | 0.699 (-14.1%) | 0.171 (-3.9%) |

Figure: 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.
Refer to caption: x5.png

Is the solution to simply replace unconditional with conditional generation?
Conditional generation is well-known to improve sample quality in the generative modeling literature (Ho & Salimans, ; Chen et al., ). Thus, one sensible assumption is that pgr exhibits stronger performance simply because our generated samples are better in quality. We show this is not the case.

Specifically, we borrow the methodology of Lu et al. () and measure faithfulness of generated transitions to environment dynamics. Given a generated transition $(s,a,s^{\prime},r)$, we roll out the action $a$ given the current state $s$ in the environment simulator to obtain the ground truth next state and reward. We then measure the mean-squared error (MSE) on these two values, comparing against generated next state $s^{\prime}$ and reward $r$, respectively. This analysis is performed at epoch 50 (halfway through online policy learning), over 10K generated transitions, and across 3 different environments.

As seen in Fig. 5, synther and pgr are highly similar in terms of generation quality. Thus, our motivations for using conditional generation have nothing to do with traditional arguments in generative modeling surrounding enhanced generation quality. Generations under pgr are better because what matters is not simply quality, but generating the *right classes/kinds of transitions*. More generally, our core contribution is to elegantly cast widely-used prioritized replay in online RL through this lens of conditional generation.

Moreover, naïve choices for conditioning fail entirely. For example, training on a larger quantity of higher reward transitions should intuitively yield a high-reward policy. However, as we show in  Fig. 4 and Table 1, naïvely conditioning on high reward (Reward-pgr) actually results in worse performance than any of the other pgr variants we propose, and in fact does worse than unconditional generation (synther). This is because greedily densifying high reward transitions does not actually provide the policy with any data on how to navigate to those high reward environment regions. Thus, it is non-trivial to even know what the correct conditioning signal should be to obtain strong performance.

Figure: Figure 6: Curiosity-pgr reduces overfitting and improves diversity of replay data. (a) Dormant ratio (Xu et al., ) (DR) of policy networks for different approaches. DR is consistently lower for pgr, indicating a minimally overfitting policy. (b) Curiosity $\mathcal{F}$-values throughout training for unconditional baseline and pgr. We show that the distribution of states accessed by the curiosity-pgr policy is significantly shifted towards higher novelty environment transitions.
Refer to caption: x6.png

What is the mechanism underlying our performance gains?
We now validate our argument in Section 4.2 for conditioning on relevance functions which mitigate overfitting. We quantify the “dormant ratio” (DR)  (Sokar et al., ) over the course of learning on the quadruped-walk task. DR is the fraction of inactive neurons in the policy network (i.e. activations below some threshold). Prior work has shown this metric effectively quantifies overfitting in value-based RL, where higher DR correlates with policies that execute unmeaningful actions (Sokar et al., ; Xu et al., ).

As we see in Fig. 6a, the redq baseline exhibits high and increasing DR over training, indicating aggressive overfitting. This reflects the underperformance of redq on quadruped-walk. Crucially, our curiosity-pgr displays a low and stable DR, which increases only marginally (after task performance has saturated, as seen in Fig. 4.) Moreover, our DR remains consistently below that of the unconditional synther baseline. This further suggests our pgr generates higher-relevance transitions with more diversity, and better addresses the issue of overfitting $Q$-functions to the synthetic data.

How does curiosity lead to pgr’s reduction in overfitting?
To show that conditioning on curiosity contributes to mitigating overfitting, we characterize the signal we obtain from $\mathcal{F}$ over time.
Specifically, we examine the curiosity-pgr variant on the quadruped-walk task, measuring the distribution of $\mathcal{F}(s,a,s^{\prime},r)$ using Eq. 5 over $10K$ *real* transitions. We perform this evaluation every $10K$ timesteps, which is the frequency at which we retrain our generative model. For comparison, we also evaluate this trained $\mathcal{F}$ on $10K$ real transitions encountered by the unconditional synther baseline.

As we see in Fig. 6b, as early as $10K$ timesteps, immediately after the policy has been trained on the first batch of synthetic data, the distribution of curiosity values is more left-skewed for the conditional model. This distribution becomes increasingly long-tailed over training, suggesting a growing diversity in the observed states as training progresses. This is contrasted by the increased biasedness of the unconditional distribution towards low-curiosity transitions. Thus, the agent trained on synthetic data from pgr is learning to move towards environment states that are more “novel” over time, improving diversity in both the real replay buffer $\mathcal{D}_{\textrm{real}}$ as well as synthetic replay buffer $\mathcal{D}_{\textrm{syn}}$. As further confirmation, we see that $100K$ timesteps, the environment is relatively well-explored, and curiosity values from $\mathcal{F}$ diminish, which correlates with the task reward saturation that we observe. We conclude that this is ultimately the mechanism underlying the improved sample efficiency of pgr.

### 5.3 Scaling Properties

Finally, we perform a series of experiments examining the scaling behavior of our pgr, in comparison to the unconditional synther baseline. For baselines, we reduce the capacity of the diffusion model in both pgr and synther such that the resultant policies attain the same performance as the model-free redq baseline on the quadruped-walk task in DMC-100K (*c.f*. dashed lines in Fig. 7). We then introduce a sequence of experiments — first increasing the size of the policy network, holding synthetic data constant, then increasing both the amount of synthetic data as well as how aggressively we rely on this data for training. Results show that pgr leverages conditionally generated synthetic data in a more scalable fashion than synther is able to use its unconditionally generated data.

Figure: Figure 7: Scaling behavior on DMC-100K quadruped-walk. pgr can effectively combine (a) larger networks with (b) higher ratios of synthetic data, (c) allowing us to train with a much higher UTD of 40. In comparison, the unconditional synther does not scale as well, and combining (a) and (b) actually underperforms using either independently. Runs shown are averaged over 3 seeds.
Refer to caption: x7.png

Network size.
We employ the experimental protocol of synther (Lu et al., ): for both networks, we increase the number of hidden layers from 2 to 3, and their widths from 256 to 512. This results in ${\sim}6$x more parameters, so we also increase batch size from 256 to 1024 to maintain per-parameter throughput. We see in Fig. 7a that this improves both pgr and synther, commensurately with redq trained with real data. This affirms that the synthetic data is a reasonable stand-in for real data.

Amount of generative data.
Next, we analyze the behavior of pgr and synther when varying the *fraction* of generative data used per batch. This better answers the full extent to which synthetic data can replace real data. Recall our baseline models use a batch size of 256 and a synthetic data ratio $r$ of 0.5 (i.e. every batch has 128 real and 128 generated transitions). We now double the batch size to 512 and then to 1024, each time scaling $r$ to 0.75 and 0.875, respectively. This keeps the number of real transitions per-batch fixed at 128, increasing only the number of generated transitions used.

We hypothesize that our generated transitions conditioned on curiosity are more effective for learning than synther’s unconditional synthetic transitions. Indeed, we observe in Fig. 7b that at $r=0.75$, pgr continues to enjoy sample-efficient learning, whereas synther fails to improve as significantly. Surprisingly, we find that both pgr and synther fail catastrophically when the synthetic data ratio is pushed to 0.875. While useful as a mechanism to augment past online experience, further work is needed to completely supplant real environment interactions with synthetic ones.

Merging insights.
Finally, we combine the above two insights, and show that this allows us to push the UTD ratio of pgr to new heights. For both synther and pgr, we again use larger MLPs to parameterize the policies, and a synthetic data ratio of 0.75 and batch size of 512. Finally, we double the UTD from 20 to 40, and the size of the synthetic data buffer $\mathcal{D}_{\textrm{syn}}$, from 1M to 2M transitions. The idea here is to preserve average diversity of the sampled synthetic transitions. We show in Fig. 7c
pgr clearly leverages both components — larger networks and more generative data — in a complementary fashion to scale more effectively without additional real interactions. In contrast, synther actually degrades in performance compared with using either component independently.

## 6 Conclusion

In this work, we propose Prioritized Generative Replay (pgr): a parametric, prioritized formulation of an online agent’s memory based on an end-to-end learned conditional generative model. pgr conditions on a relevance function $\mathcal{F}$ to guide generation towards more learning-informative transitions, improving sample efficiency in both state- and pixel-based tasks.
We show that it is not conditional generation itself, but rather conditioning on the correct $\mathcal{F}$ that is critical.
We identify curiosity as the right choice for $\mathcal{F}$, and provide a decisive answer for why: curiosity-pgr improves the diversity of generative replay, thereby reducing overfitting to synthetic data. pgr also showcases a promising formula for scalable training using synthetic data,
opening up new directions in generative RL.

#### Acknowledgments

The authors would like to thank Qiyang Li for discussion on experimental design for Section 5.2. The authors would also like to thank Qianqian Wang, Yifei Zhou, Yossi Gandelsman and Alex Pan for helpful edits of prior drafts. Amy Lu and Chung Min Kim also provided comments on Fig. 1. RW is supported in part by the NSERC PGS-D Fellowship (no. 587282). KF is supported in part by an National Science Foundation Fellowship for KF, under grant No. DGE 2146752. Any opinions, findings, and conclusions or recommendations expressed in this
material are those of the author(s) and do not necessarily reflect the views of the NSF. PA holds
concurrent appointments as a Professor at UC Berkeley and as an Amazon Scholar. This paper describes work performed at UC Berkeley and is not associated with Amazon.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: I'll rephrase my above concerns as questions. 

1. How is this method novel with respect to prior work that uses intrinsic rewards on rollouts from a learned dynamics model? It seems like a very similar approach to acquiring data that scores well under a given guidance function F, where F can be ICM or another intrinsic reward. 

2. How does this method handle noisy-tvs?
soundness: 3
strengths: * This work proposes a scalable method for training model-free or model-based agents in a variety of domains. I believe the formulation is simple enough to be integrated into and improve other approaches. 

* I also found the presentation clear and easy to read. 

* I found the scaling experiments to be very compelling, I'm a little concerned about the general thrust of driving up the syn-real data ratio as high as possible, since we do need to ground the generations in real experience. But I still think insights here are valuable.
summary: The work proposes a form of sample based experience replays that leverages a generative model to provide and augment samples drawn from the replay buffer. To avoid overfitting, a set of guidance functions are used to steer the generative process toward diverse and useful samples. The generative replay mechanism is a diffusion model that is conditioned on some auxiliary information. The authors propose a few different versions of this conditioning such as intrinsic curiosity, TD error, or Q values. The idea is that using these scores, the generative model can be steered towards generating high quality samples. Given such a replay mechanism, this work evaluates model free and model-based RL agents trained via this generative replay on gym and dmc.The results show improvement on both pixel based and state based tasks. There are also ablations with larger policy networks and higher generative data rations, which show further improvements.

-------------------------------------------------
I thank the authors for a substantive rebuttal that addressed my and (as far as I can tell) other concerns. I therefore raise my score to an 8.
weaknesses: I have two points of contention with this work. 
1. From a paradigm perspective, I don't understand how this is different from prior work in model-based RL that apples intrinsic rewards to a learned dynamics model [1] or world-model [2]. These methods also utilize a generative model as a copy of the environment, then train the agent in simulation to acquire interesting data (under the intrinsic reward). It seems that this method does the same, except that instances, rather than full trajectories are generated. I do see how this is different than just applying an intrinsic bonus during training, since here the synthetic data has a chance to be more diverse. 

2. I thank the authors for providing numerous experiments, but I am not at all convinced that this method is robust to the choice of guiding function F. ICM is known to be susceptible to the noisy TV problem, where difficult-to-model environmental factors score arbitrarily high under ICM. The chosen tasks are too simple perceptually to see this problem. This in and of itself is not a problem, but it means that we need to search for another F that works for our task which is hard in practice. In the meantime, there are other intrinsic rewards that do not suffer from this pathology [3]. 



[1] Shyam, Pranav, Wojciech Jaśkowski, and Faustino Gomez. "Model-based active exploration." International conference on machine learning. PMLR, 2019.

[2] Hafner, Danijar, et al. "Dream to control: Learning behaviors by latent imagination." arXiv preprint arXiv:1912.01603 (2019).

[3] Savinov, Nikolay, et al. "Episodic curiosity through reachability." arXiv preprint arXiv:1810.02274 (2018).

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: What exploration method does the agent use?
Could the exploration method be improved instead of the sample generation to improve diversity of samples?
Would a combination of both a better exploration and this method be the optimal and a possible solution?
soundness: 4
strengths: One of the strength of this paper are the clear and concise language as well as good structured presentation of the proposed method.
It is quite logical to improve on the already existing prioritized experience replay method and implement it in the generative domain. The method is explained well and should be quite easily reproducable.
Overall the research could be a valuable contribution to the reinforcement learning community.
summary: This paper proposes a conditional diffusion model as a synthetic replay buffer to combat early overfitting in online reinforcement learning and to diversify experiences for the learning model. This is achieved with a relevance function that selects samples that are rare but important for learning based on the temporal difference error, the value of a learned Q-function and a modified intrinsic curiosity module.
weaknesses: A topic i feel like missed somewhat are the different ways to approach generative replay such as mentions of other generative models (e.g. variational auto encoders, gaussian mixture models) and why they were not used.
One thing i found rather off putting and this is very nitpicky is that the Tables 1, 2 and 3 are a bit crammed and slightly off from each other.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Is the method compatible with different kinds of exploration bonuses? If so, how do you think they would compare?
2. How do you think the method would do when simply having diverse samples does not imply usefulness? An example is the noisy tv problem.
3. How sensitive is the algo towards the frequency of the inner loop in Algo 1?
4. Can multiple relevance functions be combined?
soundness: 3
strengths: 1. The paper is well written and provides a clear explanation of their method.
2. The research problem addressed in the paper is well laid out and is an important one to improve the performance of RL methods.
summary: The paper proposes to use conditional diffusion models to improve the experience replay for an RL learning agent. The method proposed improves performance by improving the diversity of the samples in the experience replay buffer and reducing overfitting.
weaknesses: 1. While the method shows improved performance, it is a bit simple as it combines existing elements in diffusion models and RL to propose the solution.
2. It would be useful to compare the effect of different kinds of exploration bonuses.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1) How robust is PGR to errors in the learned dynamics model? Are there ways to mitigate the impact of inaccurate dynamics predictions on the curiosity-based relevance function?
2) Could PGR be extended to offline RL settings? If so, what modifications would be necessary?
3) How does PGR's performance compare against PER baselines which use approximate parametric models of prior experience?
4) Are there any other relevance functions thats been tried out? As thats core to the working of PGR.
soundness: 4
strengths: 1) PGR offers a fresh perspective on replay buffers by combining generative modeling with guided replay. Framing this problem as a conditional generation problem with diffusion models is novel.
2) Diffusion model typically uses one single set of HPs requires no additional tuning I'd assume. This works well for PGR
3) Empirical results on various benchmarks demonstrate that PGR consistently outperforms existing model-free and model-based RL algorithms, as well as a generative replay baseline without guidance. Also has been shown to work in both state-based and pixel-based environments. 
4) PGR is shown to scale well with larger policy networks and higher synthetic-to-real data ratios (important ablation that I wanted to see), potentially enabling more data-efficient training of large-scale RL agents. Really important result for scaling to many real use cases.
5) The authors also provide insights into why PGR works, particularly highlighting the role of curiosity in promoting diversity and reducing overfitting.
summary: This paper introduces a framework called Prioritized Generative Replay (PGR), a novel approach to enhance sample efficiency in online reinforcement learning (RL). Traditionally, replay buffers store experienced transitions and replay them uniformly or with prioritization based on metrics like TD-error. However, the authors point out that uniform replay can be inefficient, and prioritization can lead to overfitting. PGR addresses these issues by using a conditional generative model to create a parametric replay buffer. 

The paper claims that this allows for two key advantages:
1) Densification: The generative model can create new, plausible transitions beyond those directly experienced, enriching the training data, especially in sparsely explored regions.
2) Guidance: By conditioning the generative model on "relevance functions," the generated transitions can be steered towards areas more critical for learning, such as states with high novelty or uncertainty.

The authors also explore various relevance functions, including return, TD-error, and curiosity. They find that curiosity, based on the prediction error of a learned dynamics model, performs best. This is attributed to its ability to promote diversity in the generated transitions, thus reducing overfitting. They also show that their approach consistently improves performance and sample efficiency in both state- and pixel based domains.
weaknesses: 1) The curiosity-based relevance function relies on a learned dynamics model, which might be challenging to train accurately in complex environments.
2) Increasing Synthetic Data ratio does not benefit PGR and the unconditional baseline (SynthER) equally. PGR scales better at r=0.75 than SYNTHER but neither benefits from 0.875. We would think the trend would be consistent? whats the intution behind this? Also this figure 7 could be improved with the variation in r being shown
3) (Minor) writing issues throughout the paper with some missing words etc. Please re-read the paper and make the necessary changes.

Meta-review: This paper uses generative models for samples for experience replay in RL. The paper uses conditional diffusion models, and explores different relevance functions, showing empirical improvements across experiments in online RL. 

Reviewers agree that this paper is well-written, and the research problem is important and well-addressed. The method is novel and the results are strong, as well as the insights as to why the method works. 

Reviewers had numerous concerns about the relevance function (eg robustness), which the authors addressed convincingly in their rebuttal (and reviewers have increased scores as a result). All reviewers agree that this work is above the acceptance threshold, with most agreeing it is a clear accept. Overall, this is a strong paper which will be a good addition to ICLR.

Decision: Accept (Oral)