Learned Replay Policy in RL

• Learned replay policy is a data-driven mechanism that selects past transitions to boost learning performance in RL and continual learning settings.
• It employs policy-gradient and meta-learning methods to dynamically assign sampling weights across scenarios like single-agent, multi-agent, and continual learning.
• Empirical studies show these policies improve sample efficiency and stability, outperforming traditional uniform or rule-based replay approaches.

A learned replay policy is a data-driven mechanism for selecting and prioritizing transitions in experience replay buffers, where the selection or scheduling is optimized through policy-gradient or meta-learning methods targeting the downstream performance metrics of reinforcement learning (RL) or continual learning systems. By parameterizing the replay process as a policy, these methods transcend hand-crafted or rule-based schemes, allowing selective reuse of past experiences to maximize sample efficiency, stability, and asymptotic performance in high-dimensional, nonstationary, or multi-agent settings.

1. Mathematical Formulation of Learned Replay Policies

A learned replay policy may be parameterized explicitly as a discrete or continuous probability distribution over the replay buffer, or implicitly as a set of adaptive sampling weights derived from an optimization principle. Formalizations differ by context:

• Single-agent RL (Experience Replay Optimization, ERO): Define a replay buffer $\mathcal{D}$ of $N$ past transitions $e_1,\dots,e_N$, and a replay policy $\pi_r^\psi(e_i\,|\,\mathcal{D}) = \lambda_i$, with $\lambda_i\in(0,1)$ computed as $\lambda_i=\phi(f_{e_i};\psi)$. The feature vector $f_{e_i}$ summarizes aspects such as reward, TD-error, and transition age (Zha et al., 2019).
• Multi-agent RL (MAC-PO): Given a Dec-POMDP $G=\langle S, U^n, P, r, Z, O, n, \gamma\rangle$, define weights $w_k(s,u)\geq0$ for joint state-action transitions. MAC-PO frames the optimal selection problem as regret minimization:

$$\min_{w_k\geq0} \ \eta(\pi^*)-\eta(\pi_k) \quad \text{s.t. } Q_k = \arg\min_Q E_{(s,u)\sim\mu}[w_k(s,u) (Q(s,u)-\mathcal{B}^*Q_{k-1}(s,u))^2], \quad \sum_{s,u} w_k(s,u)\mu(s,u)=1$$

where $\pi^*$ is the nominal optimal policy, and $Q^*$ is the corresponding optimal Q-function (Mei et al., 2023).

• Continual Learning Scheduling: Replay decisions are modeled as an MDP. At each time $t$, the state $s_t$ summarizes validation performance over seen tasks, and the action $a_t$ selects a proportion vector $p_t$ allocating the replay budget $M$ across history. The RL-trained scheduling policy $\pi_\theta(a|s)$ maximizes cumulative retained accuracy (Klasson et al., 2022).

2. Learning Objectives and Policy Update Mechanisms

Learned replay policies are directly optimized to improve the downstream RL or continual learning agent’s performance, often through meta-learning or regret minimization:

• Meta-Objective (ERO): The replay-policy parameters $\psi$ are updated via a REINFORCE-style gradient to maximize the agent’s expected improvement after replay-based updates:

$J_r(\psi) = \mathbb{E}_{I\sim\pi_r^\psi}[\, r^r\,]$

where $r^r = R(\theta^+) - R(\theta_\text{old})$ is the difference in agent performance before and after a replay batch, and $I\sim\text{Bernoulli}(\lambda)$ denotes the selected transitions (Zha et al., 2019).

• Regret-minimization (MAC-PO): The closed-form optimal sampling weights $w_k(s,u)$ are derived via Lagrangian relaxation and KKT conditions, minimizing regret upper-bounded by a Jensen-relaxed term involving $|Q_k-Q^*|$, discounted state visitation, Bellman error, and a joint-action coupling unique to MARL:

$$w_k(s,u) = \frac{1}{Z^*}\left( E_k(s,u) + \epsilon_k(s,u) \right)$$

with $E_k(s,u)$ incorporating $d^{\pi_k}(s,u)/\mu(s,u)$, Bellman error, exponential penalty for Q error, and a term $$1+\sum_{i=1}^{n}\prod_{j\neq i}\pi_k^j-n\prod_{i=1}^n\pi_k^i$$ (Mei et al., 2023).

• Replay Scheduling via RL: The scheduling policy $\pi_\theta(a|s)$ is trained via DQN or A2C to maximize cumulative validation (or test) accuracy over an episode. Dense intermediate rewards or end-to-end average accuracy serve as optimization signals (Klasson et al., 2022).

3. Algorithmic Structures and Pseudocode

The following summarization identifies key algorithmic loops for implementing learned replay policies.

• Experience Replay Optimization (ERO):
  • For each episode:
  • 1. Interact and collect transitions into $\mathcal{D}$.
  • 2. At episode end, compute agent return $R_\text{new}$.
  • 3. If prior return exists, update $\psi$ using $\nabla_\psi J_r$ sampled via episode-level REINFORCE.
  • 4. Use $\pi_r^\psi$ to sample replayed mini-batches for standard agent updates.
  • 5. Periodically update agent networks and the replay-policy as described above (Zha et al., 2019).
• MAC-PO:
  • For each environment step:
  • 1. Collect new joint-state/action/reward/next transitions.
  • 2. Periodically sample a minibatch.
  • 3. For each batch element, compute Bellman error, joint policies, and $Q_k-Q^*$ estimates.
  • 4. Assign sampling weights $w_i$ as per the closed-form, normalize, and perform weighted Bellman updates.
  • 5. Update target networks as needed (Mei et al., 2023).
• Continual Learning Scheduling:
  • Either MCTS or parametric RL explores/optimizes the replay allocation schedule across task epochs.
  • MCTS: Roll out full schedules, backpropagate test accuracy, return best schedule.
  • RL: Train $\pi_\theta$ via DQN/A2C on environment distributions, deploy as a fast schedule policy at test time (Klasson et al., 2022).

4. Distinctive Theoretical and Practical Properties

Learned replay policies differ from classical or rule-based replay by offering:

• Direct optimization for downstream task performance: Parameters are tuned to maximize agent improvement, e.g., empirical future return or minimal regret relative to an ideal policy (Zha et al., 2019, Mei et al., 2023).
• Task- and environment-adaptivity: By conditioning on observed statistics (e.g., TD-error, age, validation accuracy), the learned policy exploits domain structure, shifting priorities dynamically (Zha et al., 2019, Klasson et al., 2022).
• Multi-agent coordination: In MARL, replay weights encode not only individual agent priorities but joint action couplings to emphasize “rare but critical” transitions—a property not present in single-agent settings (Mei et al., 2023).
• Time-aware scheduling: In continual learning, the scheduling policy learns when each task’s exemplars should be replayed, exhibiting generalization across task permutations and dataset variations (Klasson et al., 2022).

5. Empirical Results and Comparative Findings

Empirical validations across several domains support the effectiveness of learned replay policies:

Setting	Replay Policy Method	Main Empirical Gains
MuJoCo RL (DDPG)	ERO ($\pi_r^\psi$)	Faster learning, higher final returns on 6/8 tasks, modest overhead vs. uniform replay; outperforms PER-rule approaches (Zha et al., 2019)
Multi-Agent RL (Predator-Prey, SMAC)	MAC-PO	Higher final win rates and faster convergence vs. uniform, PER, DisCor, ReMERN, PSER; ablations confirm necessity of each weight term (Mei et al., 2023)
Continual Learning (CLIFAR/MNIST)	Scheduling via MCTS, DQN, A2C	Learned schedules outperform uniform/heuristics by up to 3–4 points (ACC), 2–6 points (BWT); RL-learned schedule generalizes to new task orders/datasets (Klasson et al., 2022)

Significant findings include:

• In the MuJoCo continuous control domain, ERO learns to down-weight high-TD-error and stale transitions in favor of recent, near-on-policy, low-variance samples, unlike PER-based methods (Zha et al., 2019).
• In MAC-PO, omitting the joint-probability, Bellman-error, or value-enhancement factors each leads to substantial decreases in MARL performance—by 10–18% depending on the term (Mei et al., 2023).
• Scheduling the timing (not just content) of replay is critical for mitigating forgetting and optimizing average accuracy in continual learning (Klasson et al., 2022).

6. Implementation Considerations and Theoretical Insights

Key implementation details and theoretical properties include:

• Replay-policy architectures: Lightweight MLPs (2 layers, width 64) suffice for experience selection (Zha et al., 2019); joint action weighting in MARL admits closed-form expressions (Mei et al., 2023).
• Update frequency and compute overhead: ERO increases wall-clock time by ≈10–15%; most computational cost remains dominated by environment or simulator runtime (Zha et al., 2019, Klasson et al., 2022).
• Sample efficiency and stability: Learned policies demonstrate more stable and effective agent learning than uniform or rule-based alternatives, with meta-objective gradients (episode-level differences) empirically yielding robust improvements (Zha et al., 2019, Mei et al., 2023).
• Generalization: RL-trained scheduling policies ($\pi_\theta$) generalize across new task orders and, partially, to new datasets, offering scalable alternatives to re-running tree search in every environment (Klasson et al., 2022).
• Convergence: Two-loop REINFORCE optimizations used in ERO can be viewed as stochastic meta-gradients, empirically validating stability when reward estimates are smoothed over adequate windows (Zha et al., 2019). Closed-form optimal weights in MAC-PO guarantee monotonic regret reduction under reasonable conditions (Mei et al., 2023).

7. Research Directions and Open Problems

Current learned replay policy methods highlight several ongoing directions:

• Exploring more expressive or partially observed replay-policy models, especially in the presence of heavy-tailed, nonstationary, or sparse-reward environments.
• Extending closed-form optimal weighting (as in MAC-PO) to multi-step, hierarchical, or partially observable settings.
• Improving meta-gradient estimation in large-scale RL to further stabilize and accelerate learned replay policies.
• Adapting scheduling strategies for practical continual learning deployments requiring low runtime and memory cost but high sample reuse and resilience to catastrophic forgetting.
• Quantifying generalization regimes in which RL-learned scheduling policies transfer across datasets with different class distributions and task structures.

The learned replay policy paradigm thus provides a principled foundation for optimizing sample reuse, discovery, and curriculum in episodic reinforcement and continual learning architectures (Zha et al., 2019, Mei et al., 2023, Klasson et al., 2022).