Source-linked AI summary
Uncertainty-Driven Replay Memory for Reinforcement Learning
Sheeraja Rajakrishnan, Alexander G. Ororbia, Travis Desell, Daniel E. Krutz
TL;DR
RL agents can make unreliable, overconfident predictions, motivating replay strategies that use uncertainty to guide learning. The paper introduces UDRM, which biases replay memory toward uncertain transitions, and reports better or competitive performance against uncertainty-aware baselines across tested environments. Its recalibration intervals are manually selected, and their effects across environments remain for further study.
Problem
Existing RL replay buffers commonly use recency, temporal difference error, or transition features, while uncertainty-aware replay aims to address limited information about the action space.
Method
UDRM populates an experience replay buffer using uncertainty estimates from the RL model and biases storage toward transitions about which the agent is uncertain.
Results
UDRM achieves better or competitive performance compared with other uncertainty-aware baselines, including higher scores in four of five MinAtar games.
Takeaways & Limitations
UDRM shifts replay toward uncertain experiences during training, supporting exploration and generalization while later allowing more certain transitions to support exploitation.
Takeaways & Limitations
The α and β recalibration intervals are manually selected, and their effects across different RL environments remain for further study.
Abstract
from arXiv · showhide
Uncertainty estimation provides promising capabilities for reinforcement learning (RL) agents. Notably, estimating uncertainty can reduce the training time and enable agents to obtain greater rewards over time by exploiting information related to whether an action would facilitate exploration of portions of an environment that are well-known versus those that are relatively unknown. In this work, we propose a novel formulation of the experience replay buffer commonly used in RL that we call uncertainty-driven replay memory (UDRM), which entails an update scheme for internally stored memories based on uncertainty estimates obtained by an RL model during training. In contrast to existing forms of RL, which typically use temporal difference error or the distribution of transitions to update the replay memory buffer and train RL controllers, our scheme biases the memory buffer to store more uncertain transitions that will improve an RL agent's generalization throughout training. Experimental results demonstrate that our proposed uncertainty-aware replay buffer enables an RL agent to obtain higher rewards during training compared to other existing uncertainty-aware RL frameworks.
1 Introduction
The paper introduces uncertainty-driven replay memory, which populates the RL replay buffer using an agent’s prediction uncertainty. It reports higher training scores and broader applicability across varied environments.
- Uncertainty-aware replay memory is proposed to improve decisions when the agent has limited information about its action space.
- UDRM populates an RL memory buffer according to the agent’s prediction uncertainty rather than conventional replay-buffer criteria.
- UDRM obtains higher scores during training and higher overall average cumulative return than UADQN and CEQR-DQN baselines.
- The model is evaluated on MinAtar, Classic Control, and Toy Text Gym environments to demonstrate applicability across varied environments.
2 Related Work
Prior work estimates uncertainty for exploration, filters or prioritizes replay transitions, and modifies replay-buffer sampling. UDRM instead adds uncertain transitions without blocking other transitions and avoids extra prioritization metrics.
- UADQN estimates aleatoric and epistemic uncertainty with quantile regression, while CEQR-DQN uses evidential deep learning to estimate uncertainty for exploration.
- UBER filters transitions before replay-buffer insertion, whereas Wasserstein-based replay stores a transition when its uncertainty exceeds a threshold.
- UDRM injects additional uncertain transitions without preventing other transitions from entering the buffer, enabling sampling of both uncertain and certain experiences.
- As training increases certainty, UDRM stops receiving duplicate transitions and shifts toward exploitation near the end of training.
- Existing replay-buffer research has examined capacity, retention criteria, learned sampling contexts, TD-error prioritization, and sequence prioritization.
- Unlike several replay-buffer variants, UDRM directly samples transitions using epistemic uncertainty without calculating additional weights or prioritization metrics.
3 Methodology
UDRM updates replay memory using epistemic uncertainty, favoring uncertain, rewarding transitions during exploration and progressively tightening selection toward exploitation. The method combines evidential uncertainty estimation with adaptive threshold recalibration and evaluates the resulting buffer across multiple environment families.
- Uncertainty-driven replay memory: UDRM updates replay memory by prioritizing transitions with high epistemic uncertainty rather than relying on TD error or transition position.This biases sampling toward uncertain actions and supports exploration during training.
- Uncertainty estimation: Evidential deep learning estimates prediction, aleatoric, and epistemic uncertainty in a single pass using a normal-inverse-gamma distribution.The DQN uses evidential losses and uncertainty estimates to guide replay-memory updates.
- Adaptive threshold: The uncertainty threshold ξ = αe^−βt decays over time, with α and β recalibrated from stored uncertainty values to balance exploration and exploitation.The initial α uses the 99th percentile of pre-training uncertainties, while β is adjusted using a moving-average comparison with the threshold.
- Evaluation: UDRM variants outperform or match competing models across MinAtar, classic-control, and toy-text evaluations, with higher scores in four of five MinAtar games.The experiments compare UDRM with UADQN, CEQR-DQN, and a prioritized-replay variant.
- Buffer update scheme: During initial data collection, every transition enters the replay memory; afterward, uncertain transitions with positive rewards can be stored twice.Reward-positive re-injection increases the probability of sampling uncertain actions.
- Limitations: The α and β recalibration intervals are manually selected, and their effects across environments remain an open direction for further study.The authors report no drastic performance change from the tested β interval but identify environment-specific tuning as future work.
4 Experimental Results
UDRM is evaluated against uncertainty-aware baselines across MinAtar, Classic Control, and Toy Text environments. It achieves stronger performance in most MinAtar games and remains competitive in smaller environments.
- Ten uniquely seeded trials ran for 2.5 million time steps across five MinAtar games.UDRM used α recalibration intervals of 50 and 500.
- UDRM achieves higher maximum scores than UADQN and CEQR-DQN in four of five MinAtar games.The gains occur in Asterix, Breakout, Seaquest, and SpaceInvaders; Freeway is competitive.
- A smaller CartPole pre-training phase appears to enable greater rewards.The ablation tested a limited set of pre-training step counts.
- UDRM learns as well as CEQR-DQN in Classic Control and Toy Text environments.UADQN performs better in three of four smaller environments, while the PER variant performs worst.
5 Conclusions
UDRM biases replay memory toward transitions with higher estimated uncertainty, aiming to improve exploration and generalization. Across tasks, it achieves better or competitive performance than other uncertainty-aware baselines, although Freeway favors simpler models.
- UDRM populates replay memory using model-estimated uncertainty, biasing stored transitions toward uncertain actions.This distribution is intended to support better decisions, exploration, and generalization.
- UDRM achieves better or competitive performance than other uncertainty-aware baseline models across several tasks.
- Freeway is a sparse-reward environment in which simpler models perform better.