SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shivakanth Sujit
Mila, Quebec AI Institute, ÉTS Montréal shivakanth.sujit@gmail.com
Somjit Nath
Mila, Quebec AI Institute, ÉTS Montréal somjitnath@gmail.com
Pedro H.M. Braga
Mila, Quebec AI Institute, ÉTS Montréal, Universidade Federal de Pernambuco pedromagalhaes.hb@gmail.com
Mila, Quebec AI Institute, ÉTS Montréal, CIFAR AI Chair samira.ebrahimi.kahou@gmail.com
Most reinforcement learning algorithms take advantage of an experience replay buffer to repeatedly train on samples the agent has observed in the past. Not all samples carry the same amount of significance and simply assigning equal importance to each of the samples is a naive strategy. In this paper, we propose a method to prioritize samples based on how much we can learn from a sample. We define the learn-ability of a sample as the steady decrease of the training loss associated with this sample over time. We develop an algorithm to prioritize samples with high learn-ability, while assigning lower priority to those that are hard-to-learn, typically caused by noise or stochasticity. We empirically show that across multiple domains our method is more robust than random sampling and also better than just prioritizing with respect to the training loss, i.e. the temporal difference loss, which is used in prioritized experience replay. The code to reproduce our experiments can be found here.
Deep reinforcement learning has shown great promise in recent years, particularly with its ability to solve difficult games such as Go [Silver et al., 2016], chess [Silver et al., 2018], and Atari [Mnih et al., 2015]. However, online Reinforcement Learning (RL) suffers from sample inefficiency because updates to network parameters take place at every time-step with the data being discarded immediately. One of the landmarks in the space of online RL learning has been Deep Q Networks (DQN) [Mnih et al., 2015], where the agent learns to achieve human-level performance in Atari 2600 games. A key feature of that algorithm was the use of batched data for online learning. Observed transitions are stored in a buffer called the experience replay [Lin, 2004], from which one randomly samples batches of transitions for updating the RL agent.
Instead of randomly sampling from the experience replay, we propose to sample based on the learnability of the samples. We consider a sample to be learnable if there is a potential for reducing the agent's loss with respect to that sample. We term the amount by which we can reduce the loss of a sample to be its reducible loss (ReLo). This is different from vanilla prioritization in Schaul et al.
[2016] which just assigns high priority to samples with high loss, which can potentially lead to repeated sampling of data points which can not be learned from due to noise.
In our paper, we first briefly describe the current methods for prioritization while sampling from the buffer, followed by an intuition for reducible loss in reinforcement learning. We demonstrate the performance of our approach empirically on the DeepMind Control Suite [Tassa et al., 2018], OpenAI Gym, MinAtar [Young and Tian, 2019] and Arcade Learning Environment [Bellemare et al., 2013] benchmarks. These experiments show how prioritizing based on the reducible loss is a more robust approach compared to just the loss term [Schaul et al., 2016] used in Hessel et al. [2017] and that it can be integrated without adding any additional computational complexity.
In Reinforcement Learning (RL), an agent is tasked with maximizing the expected total reward it receives from an environment via interaction with it. This problem is formulated using a Markov Decision Process (MDP) [Bellman, 1957] that is described by  $< S, \mathcal{A}, \mathcal{R}, \mathcal{P} >$ , where  $S, \mathcal{A}, \mathcal{R}$  and  $\mathcal{P}$  represent the state space, the action space, the reward function, and the transition function of the environment, respectively. The objective of RL is to learn an optimal policy  $\pi^{*}$ , which is a mapping from states to actions that maximizes the expected discounted sum of rewards it receives from the environment, that is
$$
\pi^ {*} = \operatorname * {a r g m a x} _ {\pi} \mathbb {E} _ {\pi} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r _ {t} | S _ {t} = s, A _ {t} = a \right], \tag {1}
$$
where  $\gamma \in [0,1]$  is the discount factor. Action value methods obtain a policy by learning the action value  $(Q^{\pi}(s_t,a_t))$  of a policy which is the expected return by taking action  $a_{t}$  in state  $s_t$  and then following the policy  $\pi$  to choose further actions. This is done using the Bellman equation, which defines a recursive relationship in terms of the  $Q$  value function, as follows
$$
Q ^ {\pi} \left(s _ {t}, a _ {t}\right) = r _ {t} + \gamma \max  _ {a} Q ^ {\pi} \left(s _ {t + 1}, a\right) \tag {2}
$$
The difference between the left and right sides of Eq. 2 is called the temporal difference error (TD error), and  $Q$  value methods minimize the TD error of the learned  $Q$  function  $Q^{\theta}$  (implemented as a neural network) using stochastic gradient descent. That is, the loss for the  $Q$  network is
$$
L _ {\theta} = \left(Q ^ {\theta} \left(s _ {t}, a _ {t}\right) - \left(r _ {t} + \gamma \max  _ {a} Q ^ {\theta} \left(s _ {t + 1}, a\right)\right)\right) ^ {2}. \tag {3}
$$
We can then use the  $Q$  value to implicitly represent a policy by choosing actions with high  $Q$  values. While this is easy in discrete control tasks which have a small action space, it can be difficult in continuous action spaces because finding the action that maximizes the  $Q$  value can be an optimization problem in itself. This can be computationally expensive to do at every instant, so recent methods alleviate this problem through an actor network  $\mu_{\theta}$  that learns the action that produces the maximum  $Q$  value through stochastic gradient ascent, that is
$$
\mu_ {\theta} = \underset {\theta} {\operatorname {a r g m a x}} Q ^ {\theta} \left(s _ {t}, \mu_ {\theta} \left(s _ {t}\right)\right). \tag {4}
$$
The loss for the  $Q$  network in Eq. 3 is then modified so that the argmax is evaluated using the actor network,
$$
L _ {\theta} = \left(Q ^ {\theta} \left(s _ {t}, a _ {t}\right) - \left(r _ {t} + \gamma Q ^ {\theta} \left(s _ {t + 1}, \mu_ {\theta} \left(s _ {t}\right)\right)\right)\right) ^ {2} \tag {5}
$$
2.1 Experience Replay
Online RL algorithms perform updates immediately after observing a transition. However, these not only make learning inefficient but can cause issues in training due to high correlation between recent transitions. To eliminate this problem, Lin [2004] introduced experience replay, which stores the observed transitions and provides an interface to sample batches of transitions. This has been successfully used in DQN [Mnih et al., 2015] to play Atari 2600 games. Since Eqs. 3 and 5 do not require that the states and actions are generated from the current policy, algorithms trained this way are called off-policy RL algorithms. During training, data is collected from the environment and stored in a replay buffer from which mini-batches are sampled to be trained on.
A naive method of sampling is to uniformly sample all data in the buffer, however, this is inefficient because not all data is necessarily equally important. Schaul et al. [2016] proposes Prioritized Experience Replay (PER), that samples points with probabilities proportional to their TD error – which has been shown to have a positive effect on performance by efficiently replaying samples that the model has not yet learned, i.e., data points with high TD error. Each transition in the replay buffer is assigned a priority  $p_i$ , and the transitions are sampled based on this priority. To ensure that data points, even with low TD error, are sampled sometimes by the agent, instead of greedy sampling based on TD error, the replay buffer in PER stochastically samples points with probability  $P_i$ .
$$
P _ {i} = \frac {p _ {i} ^ {\alpha}}{\sum_ {j} p _ {j} ^ {\alpha}} \tag {6}
$$
where  $\alpha \in [0,1)$  is a hyper-parameter introduced to smoothen out very high TD errors. Setting  $\alpha$  to 0 makes it equivalent to uniform sampling. Since sampling points non-uniformly changes the expected gradient of a mini-batch, PER corrects for this by using importance sampling (IS) weights  $w$
$$
w _ {i} = \left(\frac {p _ {\text {u n i f o r m}}}{P _ {i}}\right) ^ {\beta} \tag {7}
$$
where  $\beta \in [0,1]$  controls the amount by which the change in gradient should be corrected and  $p_{\mathrm{uniform}} = \frac{1}{N}$  where  $N$  is the number of samples in the replay buffer. The loss attributed to each sample is weighed by the corresponding  $w_{i}$  before the gradient is computed. In practice,  $\beta$  is either set to 0.5 or linearly annealed from 0.4 to 1 during training.
2.2 Target Networks
In Eqs. 3 and 5, the target action value depends not only on the rewards but also on the value of the next state, which is not known. So, the value of the next state is approximated by feeding the next state to the same network used for generating the current  $Q$  values. As mentioned in DQN [Mnih et al., 2015], this leads to a very unstable target for learning due to the frequent updates of the  $Q$  network. To alleviate this issue, Mnih et al. [2015] introduce target networks, where the target  $Q$  value is obtained from a lagging copy of the  $Q$  network used to generate the current  $Q$  value. This prevents the target from changing rapidly and makes learning much more stable. So Eq. 3 can be suitably modified to
$$
L _ {\theta} = \left(Q ^ {\theta} \left(s _ {t}, a _ {t}\right) - \left(r _ {t} + \gamma \max  _ {a} Q ^ {\bar {\theta}} \left(s _ {t + 1}, a\right)\right)\right) ^ {2} \tag {8}
$$
respectively, where  $\bar{\theta}$  are the parameters of the target network, which are updated at a low frequency. Mnih et al. [2015] copies the entire training network  $\theta$  to the target network, whereas Haarnoja et al. [2018] performs a soft update, where the new target network parameters are an exponential moving average (with a parameter  $\tau$ ) of the old target network parameters and the online network parameters.
2.3 Reducible Loss
The work of Mindermann et al. [2022] proposes prioritized training for supervised learning tasks based on focusing on data points that reduce the model's generalization loss the most. Prioritized training keeps a held-out subset of the training data to train a small capacity model,  $\theta_{ho}$  at the beginning of training. During training, this hold-out model is used to provide a measure of whether a data point could be learned without training on it. Given training data  $(x_i,y_i)\in \mathcal{D}$ , the loss of the hold-out model's prediction,  $\hat{y}_{iho}$  on a data point  $x_{i}$  could be considered an estimate of the remaining loss after training on datapoints other than  $(x_{i},y_{i})$ , termed the irreducible loss. This estimate becomes more accurate as one increases the size of the held-out dataset. The difference between the losses of the main model,  $\theta$ , and the hold-out model on the actual training data is called the reducible loss,  $L_{r}$  which is used for prioritizing training data in mini-batch sampling.  $L_{r}$  can be thought of as a measure of information gain by also training on data point  $(x,y)$ .
$$
L _ {r} = \operatorname {L o s s} (\hat {y} \mid x, \theta) - \operatorname {L o s s} (\hat {y} \mid x, \theta_ {h o}) \tag {9}
$$
2.4 Prioritization Schemes
Alternate prioritization strategies have been proposed for improvements in sample efficiency. Sinha et al. [2020] proposes an approach that re-weights experiences based on their likelihood under the stationary distribution of the current policy in order to ensure small approximation errors on the value function of recurring seen states. Lahire et al. [2021] introduces the Large Batch Experience Replay (LaBER) to overcome the issue of the outdated priorities of PER and its hyperparameter sensitivity by employing an importance sampling view for estimating gradients. LaBER first samples a large batch from the replay buffer then computes the gradient norms and finally down-samples it to a mini-batch of smaller size according to a priority. Kumar et al. [2020] presents Distribution Correction (DisCor), a form of corrective feedback to make learning dynamics more steady. DisCor computes the optimal distribution and performs a weighted Bellman update to re-weight data distribution in the replay buffer. Inspired by DisCor, Regret Minimization Experience Replay (ReMERN) [Liu et al., 2021] estimates the suboptimality of the Q value with an error network. Yet, Hong et al. [2022] uses Topological Experience Replay (TER) to organize the experience of agents into a graph that tracks the dependency between Q-value of states.
3 Reducible Loss for Reinforcement Learning
While PER helps the agent to prioritize points that the model has not yet learned based on high TD error, we argue that there are some drawbacks. Data points could have high TD error because they are noisy or not learnable by the model. It might not be the case that a data point with high TD error is also a sample that the model can actually learn or get a useful signal from. In supervised learning, a known failure condition of loss-based prioritization schemes is when there are noisy points which can have high loss but are not useful for repeated training [Hu et al., 2021]. Instead of prioritization based on the TD error, we propose that the agent should focus on samples that have higher reducible TD error. This means that instead of the TD error, we should use a measure of how much the TD error can be potentially decreased, as the priority  $p_i$  term in Eq. 6. We contend that this is better because it means that the algorithm can avoid repeatedly sampling points that the agent has been unable to learn from and can focus on minimizing error on points that are learnable, thereby improving sample efficiency. Motivated by prioritized training in supervised learning, we propose a scheme of prioritization tailored to the RL problem.
In the context of supervised learning, learn-ability and reducible loss for a sample are well-defined as one has access to the true target. However, in RL, the true target is approximated by a bootstrapped target (TD methods) or the return obtained from that state (Monte Carlo). For policy evaluation with a fixed policy, the true value function can be obtained, however, since we are interested in control, it will be computationally intensive to capture the true value function with every change in policy as in policy iteration. Thus, to determine the learn-ability of a sample, we need access to how the targets of the sample behave and how it changes across time. Since the concepts of a hold-out dataset or model are ill-defined in the paradigm of RL, we replace them with a moving estimate of the targets. Unlike in supervised learning, where we draw i.i.d. batches from a fixed training set, the training data in RL are not i.i.d since they are generated by a changing policy. So the holdout model would need to be updated from time to time. Thus, in  $Q$  learning-based RL methods, a good proxy for the hold-out model is the target network used in the Bellman update in Eq. 8. Since the target network is only periodically updated with the online model parameters, it retains the performance of the agent on older data which are trained with outdated policies. Schaul et al. [2022] demonstrates how the policies keep changing with more training even when the agent receives close to optimal rewards. Thus, the target network can be easily used as an approximation of the hold out model that was not trained on the new samples. Therefore, we define the Reducible Loss (ReLo) for RL as the difference between the loss of the data point with respect to the online network (with parameters  $\theta$ ) and with respect to the target network (with parameters  $\bar{\theta}$ ). So the Reducible Loss (ReLo) can be computed as
$$
\operatorname {R e L o} = L _ {\theta} - L _ {\bar {\theta}} \tag {10}
$$
There are similarities between ReLo as prioritization scheme in the sampling behavior of low priority points when compared to PER. Data points that were not important under PER, i.e. they have low  $L_{\theta}$ , will also remain unimportant in ReLo. This is because if  $L_{\theta}$  is low, then as per Eq. 10, ReLo will also be low. This ensures that we retain the desirable behavior of PER, which is to not repeatedly sample points that have already been learned.
However, there is a difference in sampling points that have high TD error. PER would assign high priority to data points with high TD error, regardless of whether or not those data points are noisy or unlearnable. For example, a data point can have a high TD error which continues to remain high even after being sampled several times due to the inherent noise of the transition itself, but it would continue to have high priority with PER. Thus, PER would continue to sample it, leading to inefficient learning. But, its priority should be reduced since there might be other data points that are worth sampling more because they have useful information which would enable faster learning. The ReLo of such a point would be low because both  $L_{\theta}$  and  $L_{\bar{\theta}}$  would be high. In case a data point is forgotten, then the  $L_{\theta}$  would be higher than  $L_{\bar{\theta}}$ , and the ReLo would ensure that these points are revisited. Thus ReLo can also help to overcome forgetting.
3.1 Implementation
The probability of sampling a data point is related to the priority through Eq. 6 and requires the priority to be non-negative. Since  $Q$  value methods use the mean-squared error (MSE) loss, the priority is guaranteed to be non-negative. However, ReLo computes the difference between the MSE losses and it does not have the same property. This value can go to zero when the target network is updated with the main network using hard updates. However, it quickly becomes non-zero after one update. So, we should create a mapping  $f_{map}$  for the ReLo error that is monotonically increasing and non-negative for all values. In practice, we found that clipping the negative values to zero, followed by adding a small  $\epsilon$  to ensure samples had some minimum probability, worked well. That is,  $p_i = \max(\text{ReLo}, 0) + \epsilon$ . This is not the only way we can map the negative values and we have studied one other mapping in the supplementary material. Note also that during initial training when the agent sees most of the points for the first time, it would assign low priority to all depending on the value of ReLo and the sampling would be close to random. However, the priority will never go to zero because of  $\epsilon$ , so we will always have a non-zero probability of drawing the sample. Once it is sampled and the loss goes down, then the priority would again increase as per ReLo and it would prioritize such points until the ReLo becomes low again. Thus, it kind of achieves the best of both worlds, where initially with less information we sample uniformly, but once ReLo increases we sample points more frequently.
ReLo is not computationally expensive since it does not require any additional training. It only involves one additional forward pass of the states through the target network. This is because the Bellman backup (i.e., the right hand side of Eq. 2) is the same for  $L_{\theta}$  and  $L_{\bar{\theta}}$ . The only additional term that needs to be computed for ReLo is  $Q_{tgt}(s_t, a_t)$  to compute  $L_{\bar{\theta}}$ .
Algorithm 1 Computing ReLo for prioritization
Given off-policy algorithm  $A$  with loss function  $L^{alg}$ , online  $Q$  network parameters  $\theta$ , target  $Q$  network parameters  $\bar{\theta}$ , replay buffer  $B$ , max priority  $p_{max}$ , ReLo mapping  $f_{map}$ , epsilon priority  $\epsilon$ , training timesteps  $T$ , gradient steps per timestep  $T_{grad}$ , batch size  $b$ .  
for t in 1, 2, 3, ...  $T$  do  
Get current state  $s_t$  from the environment  
Compute action  $a_t$  from the agent  
Store the transition  $< s_t, a_t, r_t, s_{t+1}>$  in the replay buffer  $B$  with priority  $p_{max}$ .  
for steps in 1, 2, 3, ...  $T_{grad}$  do  
Sample minibatch of size  $b$  from replay buffer  
Compute the loss  $L_{\theta}^{alg}$  and update the agent parameters  $\theta$   
Compute  $L_{\theta tgt}^{alg}$  and calculate ReLo as per Eq. 10  
Update priorities of samples in mini-batch with the new ReLo values as  $f_{map}(\mathrm{ReLo}_i) + \epsilon$   
end for  
Update target network following the original RL algorithm  $A$   
end for
In our implementation, we saw a negligible change in the computational time between PER and ReLo. ReLo also does not introduce any additional hyper-parameters that need to be tuned and works well with the default hyper-parameters of  $\alpha$  and  $\beta$  in PER. An important point to note is that ReLo does not necessarily depend on the exact loss formulation given in Eq. 8 and can be used with the loss function  $L_{\theta}^{alg}$  of any off-policy  $Q$  value learning algorithm. In order to use ReLo, we only have to additionally
(a) DeepMind Control Suite
(c) MinAtar
(b) OpenAI Gym
(d) Atari
Figure 1: Metrics aggregated across each benchmark based on proposed metrics from Agarwal et al. [2021]. 10 seeds per environment-algorithm pair.
compute  $L^{alg}$  with respect to the target network parameters  $\bar{\theta}$ . If the loss is just the mean square error, then ReLo can be simplified and can be represented by the difference between  $Q_{\theta}$  and  $Q_{\bar{\theta}}$ . But other extensions to off policy Q learning methods modify this objective, for example distributional learning [Bellemare et al., 2017] minimizes the KL divergence and the difference between two KL divergences can not be simplified the same way. To make ReLo a general technique that can be used across these methods, we define it in terms of  $L_{\theta}$  and  $L_{\bar{\theta}}$ . Our experiments also show that ReLo is robust to the target network update mechanism, whether it is a hard copy of online parameters at a fixed frequency (as in DQN [Mnih et al., 2015], and Rainbow [Hessel et al., 2017]) or if the target network is an exponential moving average of the online parameters (as in Soft Actor Critic [Haarnoja et al., 2018]).
4.1 GridWorld Experiments
The major goal of these Gridworld studies is to draw attention to two PER issues: subpar performance when faced with stochasticity and forgetting when faced with many tasks. We also highlight how ReLo can potentially solve both issues.
4.1.1 Pitfalls of TD Error Prioritization
To illustrate the potential downsides of using PER with TD error prioritization, we created a  $7 \times 7$  empty gridworld with a single goal. The agent always spawns at the top left corner and it has to reach the goal at the bottom right for a reward of  $+2$ . The reward is zero everywhere else except, near the center of the grid, where there is another state with a reward sampled uniformly from  $[-0.5, 0.5]^1$ . The presence of this critical point would mean that during the initial phase of training, PER would keep prioritizing transitions to that state due to the unpredictable reward and the corresponding high TD error, and fail to learn about the true goal. Given a sufficient number of samples, all algorithms would ultimately converge to the optimal policy. The ReLo criterion, however, does not prioritize unlearnable points and prioritizes other points where it can receive a larger reward. Fig. 2 shows that PER samples this unlearnable point very often and with limited experience is unable to solve the task. In contrast, both uniform experience replay and ReLo do not over-sample this transition and thus can perform better in this scenario. This is further affirmed by the number of updates each algorithm makes in the neighborhood of the unlearnable point as visualized in Fig. 2. Prioritization using ReLo enables the agent to sample more from the transitions leading to the main reward which explains its better performance.
<sup>1</sup>A detailed description of the environment is given in the supplementary material.