Adaptive Replay Buffer for Offline-to-Online Reinforcement Learning

AI summary

Overview

Research area: Reinforcement Learning — specifically Offline-to-Online Reinforcement Learning (O2O RL) and replay buffer design.

Technical level: Intermediate. The paper assumes familiarity with Markov Decision Processes, off-policy RL, policy log-likelihoods, and the offline/online fine-tuning pipeline, but the core idea is conceptually simple.

Scope: The paper proposes the Adaptive Replay Buffer (ARB), a learning-free sampler that weights replayed transitions by how likely they are under the current policy ("on-policyness"), and evaluates it on D4RL benchmarks across three O2O RL algorithms.

What This Paper Is About

When an RL agent is pre-trained on a fixed offline dataset and then fine-tuned through online interaction, it must decide how to mix old offline data with newly collected online data — a "replay buffer dilemma." Keeping mostly offline data caps final performance, while switching too fast to online-only data causes catastrophic forgetting and performance collapse. The paper's goal is a simple, adaptive data-mixing rule that avoids both a hand-tuned fixed ratio and the extra machinery of learning an "onlineness" metric, so that offline data provides early stability while learning progressively shifts toward the online experiences most relevant to the current policy.

Key Contributions

1. A definition of "on-policyness." For a transition $(s,a,r,s')$, on-policyness is defined as the likelihood of that transition within an on-policy dataset, $\mathcal{O}(s,a,r,s';\pi_\theta) = d^{\pi_\theta}(s)\pi_\theta(a|s)$, where $d^{\pi_\theta}(s)$ is the stationary state distribution under the current policy. Because the state term is intractable, the authors approximate $d^{\pi_\theta}(s)$ as uniform and use only $\pi_\theta(a|s)$, yielding the practical metric $\tilde{\mathcal{O}}(s,a;\pi_\theta) = \pi_\theta(a|s)$.

2. A learning-free, trajectory-level prioritization scheme. Instead of assigning a sampling probability per transition (which has high variance and risks repeatedly sampling a few points), ARB computes on-policyness at the trajectory level as the geometric mean of the constituent transitions' values, and every transition in a trajectory inherits that trajectory's weight. The weight includes a temperature $\lambda$ controlling how strongly sampling follows on-policyness, and log-likelihoods are clamped between $\underline{p}$ and $\bar{p}$ and shifted by the maximum bounded log-probability.

3. A plug-in algorithm for O2O RL. ARB adds no extra model or learning procedure and integrates into existing O2O RL algorithms. It is presented as Algorithm 1 with a two-phase structure: offline pre-training, then online fine-tuning with periodic re-weighting of the buffer.

4. Extensive empirical validation on D4RL. ARB is integrated into Cal-QL, PEX (based on IQL), and FamO2O (based on IQL), and compared against Naive, Parallel, Top-N, and Balanced Experience Replay Buffer (BERB) baselines across three locomotion environments with four dataset levels and six Antmaze environments.

Main Findings

• ARB consistently beats all baselines across algorithms and environments. In Table 1, averaged across the three algorithms, ARB scores 32.0 vs. 18.8 (Naive), 22.6 (Parallel), 16.9 (Top-N), and 20.7 (BERB) on Locomotion random; 86.6 vs. 75.4/75.1/80.7/77.0 on Locomotion medium-replay; 86.7 vs. 70.3/82.3/79.7/80.3 on Locomotion medium; 100.3 vs. 97.5/97.7/97.1/98.6 on Locomotion medium-expert; and 88.2 vs. 78.6/84.8/75.6/81.7 on Antmaze.

• Gains are largest when offline data is low quality. On Locomotion random, FamO2O with ARB reaches 53.1 versus 30.5 (Naive), 40.9 (Parallel), 26.3 (Top-N), and 35.9 (BERB). The margin narrows on high-quality datasets such as Locomotion medium-expert, where ARB still leads but by a smaller amount (100.3 average vs. 98.6 for BERB).

• ARB is biased toward online data, and more so when offline data is poor. Figure 1 tracks the online-data ratio of minibatches over 1M environment steps for the four hopper datasets with FamO2O. ARB's curve is consistently above the grey line marking the overall online share of the dataset. When the offline dataset's average reward is low, ARB's online ratio rises sharply — behaviour not observed in other methods — which the authors present as direct evidence of adaptive prioritization.

• Trajectory-level prioritization outperforms transition-level. The ablation on antmaze-large-play-v2 with Cal-QL (Figure 2) shows the trajectory-based method achieving a significantly higher final normalized score. The authors note that transition-based sampling produces a higher online data ratio because the exponential of the geometric mean is less than the average of exponentials, and argue that aggregating weight over a sequence prevents over-fitting to a few noisy transitions.

• Temperature $\lambda$ governs the stability/adaptation trade-off. On antmaze-large-diverse-v2 with PEX (Figure 3), lower temperature values make the online data ratio increase faster, confirming that $\lambda$ controls how aggressively recent, on-policy data is prioritized; normalized scores follow a similar trend.

• Clipping bounds are robust except for very high upper bounds. Table 2 (antmaze-large-diverse-v2, PEX, four random seeds, defaults $\underline{p}=-12.0$, $\bar{p}=7.0$) reports 82.00 ± 2.94 at the defaults, with the lower bound generally robust (75.75 ± 2.16 at $-16.0$; 78.75 ± 1.71 at $-14.0$; 82.75 ± 3.40 at $-9.0$; 77.50 ± 2.52 at $-7.0$). Raising the upper bound degrades performance sharply: 70.25 ± 4.11 at $\bar{p}=9.0$ and 29.50 ± 34.41 at $\bar{p}=12.0$, which the authors attribute to over-sampling a small subset of data.

• Implementation details. ARB uses $\lambda = 5.0$ for Cal-QL and $\lambda = 0.5$ for PEX and FamO2O, with $d_{\text{weight}} = 1000$, $d_{\text{update}} = 1000$, and $N_{\text{update}} = 1000$. All three policies are Gaussian, and log-likelihoods are computed from the standard normal PDF and normalized by the number of action dimensions. Experiments ran on an AMD Ryzen Threadripper 2990WX 32-Core CPU with an NVIDIA RTX 3090 GPU, Python 3.9, CUDA 11.3, and PyTorch 1.11.

Methodology in Plain English

The approach rests on one intuition: online fine-tuning is off-policy learning, but the data being collected is on-policy. Learning from data that matches what the current policy would do lowers variance and improves final performance. So rather than deciding a fixed offline/online split, the buffer should simply prefer whatever the current policy is most likely to have produced.

Concretely:

1. Pre-train normally. Run the base O2O algorithm on the offline dataset for $N_{\text{pretrain}}$ iterations with uniform minibatch sampling, producing a warm-start policy and value function.

2. Load everything into one buffer. Fine-tuning begins by initializing an empty Adaptive Replay Buffer and loading the entire offline dataset into it. New transitions collected from the environment are appended.

3. Score each trajectory by on-policyness. For each transition, compute $\log \pi_\theta(a_t|s_t)$, clip it between $\underline{p}$ and $\bar{p}$, and subtract the maximum clipped log-probability in the dataset. Exponentiating gives a value in $(0,1]$. Because full stationary state distributions are intractable in continuous-state environments with unknown dynamics — and because over-weighting already-explored states could trap the agent in a local optimum — the state term is approximated as uniform, leaving only the action probability.

4. Aggregate at the trajectory level. Each trajectory's on-policyness is the geometric mean of its transitions' scores (with temperature $\lambda$), and every transition in it shares that weight. This smooths out the high variance of per-transition probabilities.

5. Sample, re-weight, repeat. Every $d_{\text{weight}}$ steps, the weights across the whole buffer are recomputed from the current policy. Every $d_{\text{update}}$ steps, minibatches are drawn by weighted sampling and the policy and value function are updated for $N_{\text{update}}$ steps. Offline data that drifts far from the current policy's distribution naturally loses weight over time.

Evaluation uses D4RL: three locomotion environments (hopper, walker2d, halfcheetah) at four dataset levels (random, medium-replay, medium, medium-expert) plus six Antmaze environments. All methods are pre-trained for 1M steps offline and fine-tuned for 1M online steps, with normalized scores averaged over four random seeds, measured from 100 sampled trajectories. Baselines are Naive (uniform sampling over all offline plus online data), Parallel (fixed 50/50 split), Top-N (retain the top 50,000 transitions by summed reward), and BERB.

Why This Matters

Impact on research. The paper reframes the O2O replay-buffer question: instead of choosing how much offline versus online data to use, ARB asks which specific experiences are worth replaying, and answers with a quantity already computed by the policy network. This suggests that expensive auxiliary machinery for measuring "onlineness" — as in BERB's learned metric — may be unnecessary. It also provides a clean ablation baseline (trajectory- vs. transition-level weighting) for future work on prioritized sampling in the O2O setting.

Real-world applications:

• Robotics and locomotion control, where policies are pre-trained on logged demonstrations or prior deployments and then refined on the physical robot — domains like the legged locomotion and maze navigation tasks in the D4RL suite.
• Industrial process control, where historical operating logs are abundant but safe exploration is limited, and the controller must improve with each new run.
• Healthcare treatment optimization, where retrospective patient records form the offline dataset and subsequent clinical interactions provide the online signal.
• Autonomous driving and recommendation/advertising systems, where large offline logs coexist with a continuous stream of online feedback and the policy must adapt without discarding prior knowledge.

Industry relevance. The main selling point for practitioners is operational: ARB requires no additional learned model, no extra training loop, and no per-environment tuning of a mixing ratio. The default clipping bounds $\underline{p}=-12.0$ and $\bar{p}=7.0$ are reported as robust across the tests shown, and the only hyperparameter needing per-algorithm choice is the temperature ($\lambda = 5.0$ for Cal-QL, $0.5$ for PEX and FamO2O). Public code is released at https://github.com/song970407/ARB.

Future Directions

• Automatic temperature selection. $\lambda$ was manually set per base algorithm and controls the entire stability/adaptation trade-off; the paper shows its effect but does not report a procedure for choosing it automatically.
• Restoring the state-distribution term. ARB approximates the stationary state distribution $d^{\pi_\theta}(s)$ as uniform. The authors flag this as a deliberate simplification to avoid intractable estimation and to prevent over-prioritizing already-explored states. Whether a tractable approximation of the full on-policyness expression would help remains open.
• Extending beyond Gaussian policies and low-dimensional states. Log-likelihood is computed from the Gaussian PDF and normalized by action dimension; the paper reports results on D4RL locomotion and Antmaze. Applicability to image-based observations, discrete action spaces, or non-Gaussian policy classes is not established in the content provided.
• Combination with other O2O techniques. The paper positions ARB as a plug-in and demonstrates it with Cal-QL, PEX, and FamO2O. Whether it composes additively with value-calibration, policy-constraint, or generative-sampling methods (for example Energy-guided Diffusion Sampling) is not reported.
• Complexity and compute reporting. The checklist states that time, space, and sample complexity are analyzed, but the provided content does not include those numbers; concrete wall-clock overhead relative to BERB's learned metric is therefore not reported here.

Target Audience

This paper is most useful to RL researchers and engineers working on offline RL, offline-to-online fine-tuning, or experience replay design — particularly those who already have a pre-trained policy and are struggling with the performance dip that occurs when online fine-tuning begins. It is also relevant to practitioners who need a low-overhead improvement to an existing O2O pipeline rather than a new architecture, since ARB is explicitly designed as a learning-free plug-in. Readers unfamiliar with log-likelihoods, replay buffers, and the offline/online split will find the paper manageable at an intermediate level but should first be comfortable with standard off-policy RL concepts.

Authors’ abstract

Offline-to-Online Reinforcement Learning (O2O RL) faces a critical dilemma in balancing the use of a fixed offline dataset with newly collected online experiences. Standard methods, often relying on a fixed data-mixing ratio, struggle to manage the trade-off between early learning stability and asymptotic performance. To overcome this, we introduce the Adaptive Replay Buffer (ARB), a novel approach that dynamically prioritizes data sampling based on a lightweight metric we call 'on-policyness'. Unlike prior methods that rely on complex learning procedures or fixed ratios, ARB is designed to be learning-free and simple to implement, seamlessly integrating into existing O2O RL algorithms. It assesses how closely collected trajectories align with the current policy's behavior and assigns a proportional sampling weight to each transition within that trajectory. This strategy effectively leverages offline data for initial stability while progressively focusing learning on the most relevant, high-rewarding online experiences. Our extensive experiments on D4RL benchmarks demonstrate that ARB consistently mitigates early performance degradation and significantly improves the final performance of various O2O RL algorithms, highlighting the importance of an adaptive, behavior-aware replay buffer design. Our code is publicly available at https://github.com/song970407/ARB.