R Q DIA : R EGULARIZING Q-V ALUE D ISTRIBUTIONS
W ITH I MAGE A UGMENTATION
Anonymous authors
Paper under double-blind review
A BSTRACT
rQdia (pronounced “Arcadia”) regularizes Q-value distributions with augmented images in pixel-based deep reinforcement learning.
This simple idea, to equalize Q-values across statistical distributions of actions and states, affords image augmentation techniques like DrQ additional sample efﬁciency and better ﬁnal performance, while propelling discrete-action DER to nearly 1.5x the performance of DrQ and nearly 2x that of base DER.
Representation learning is a major hurdle in deep RL, which notoriously requires far too many environment interactions for real-world use cases.
rQdia decreases this data hunger and increases overall scores, bringing deep RL closer to real-world applicability.
I NTRODUCTION
Human perception is invariant to and remarkably robust against many perturbations, like discoloration, obfuscation, and low exposure.
On the other hand, artiﬁcial neural networks do not intrinsically carry these invariance properties, not without regularizers or hand-crafted inductive biases like convolution, kernel rotation, dilation, attention, and recurrence.
In deep reinforcement learning (RL) from pixels, an agent must learn to visually interpret a scene in order to decide on an action.
Thus, recent approaches in RL have turned to the self-supervision and data augmentation techniques found in computer vision.
Indeed, such contrastive learning auxiliary losses (Srinivas et al., 2020) or data augmentation regularizers (Yarats et al., 2021b) have afforded greater sample efﬁciency and ﬁnal scores in both the DeepMind Continuous Control Suite (Tassa et al., 2018) from pixels and Atari Arcade Learning Environments (Bellemare et al., 2013).
Nevertheless, pixel-based approaches continue to lag behind models that learn directly from state embeddings, not just in terms of sample efﬁciency, but also in longer-term asymptotic performance.
For example, the recent DrQ (Yarats et al., 2021b), an image augmentation-based regularizer added to SAC (Haarnoja et al., 2018), reports falling 14 .
5% short of its state embedding-based SAC counterpart on the Cheetah Run task.
Such discrepancies indicate that visual representations are not yet up to par with state embeddings, at least not for locomotive continuous control.
State embeddings have many properties that facilitate generalization, such as location invariance, and to a degree, invariance between morphological relations.
If a subset of the dimensions of a state embedding always indicates feet position, then the relation “one foot in front of the other” will be represented by those dimensions invariant to the placement of the robot’s arms, head, or other body parts.
Parametric visual encodings are not guaranteed to learn such invariances with respect to the robot’s morphology.
What other signals in deep RL can guide visual representation learning toward more invariant encodings?
We propose Q-value distributions, the expected cumulative discounted reward across stateaction pairs, offer such a signal.
For state s , Q-function Q , and actions a (0) , ..., a ( m ) ∼ D ( A ) from some statistical distribution D over action space A , we deﬁne Q-value distribution simply as:
[Equation]
This distribution over actions for a state is a measure of the current and future value of each action for that state.
We review states, actions, rewards, Q-values, and the MDP framework in the subsequent Section 2.1.
This distribution is amenable to many of the same invariances afforded by state embeddings, if not some additional ones.
For example, if the optimal action is “put one foot in front of the other,” then the Q-value distribution reﬂects the preferability of this action over other actions regardless of (invariant to) where the agent is located, how the agent is holding its hands, or any other potentially superﬂuous contextual information.
While recent work shows individual Q-values beneﬁt from regularizing across augmented images (Yarats et al., 2021b), we consider that Q-value distributions are also important, as they contain information not just about one action in isolation, but multiple actions in relation to one another.
B ACKGROUND 2. D EEP RL F ROM P IXELS
A Markov Decision Process (MDP) consists of an action
a ∈ IR
d a , state s ∈ IR
d s , and reward r ∈ IR .
“From pixels” assumes that state s is an image frame or multiple image frames stacked together.
The action is sampled from a policy at any t time step a t ∼ π ( s t ) .
Taking such actions yields a trajectory τ =
( s 0 , a 0 , s 1 , a 1 , ..., s T ) via the dynamics model s t +1 ∼ f
( s t , a t ) of the environment and its rewards r t +1 =
R ( s t , a t ) .
The goal is to maximize cumulative discounted reward R ( τ )
=
(cid:80) Tt =
r t γ t where γ is the temporal discount factor.
The optimal action for a state a ∗ ( s ) = argmax a Q ∗ ( s, a ) thus depends on the state-action value function, also known as the Qvalue, Q π ( s, a ) =
E τ ∼ π
[ R ( τ ) | s
= s, a
= a ] .
2. S OFT A CTOR -C RITIC & D R Q
Soft-Actor Critic (SAC) (Haarnoja et al., 2018) is an RL algorithm which learns a min-reduced ensemble of Q-value functions Q φ ( s, a ) = min i =1 , 2 Q φ [Equation]
( s [Equation] a
[Equation] optimized with one-step Bellman error, and a stochastic Gaussian policy π θ ( s, a ) optimized by maximizing Q φ ( s, a ) and entropy H ( s ) = − π θ ( s ) log ( π θ ( s )) , made differentiable via the reparameterization trick. To further encourage exploration and avoid premature policy collapse, entropy H ( s ) is also added as part of the agent’s reward.
In visual domains, Q φ and π θ are typically equipped with a shared convolutional neural net encoder.
DrQ (Yarats et al., 2021b) averages Q φ ( aug 1 ( s ) , a ) with Q φ ( aug 2 ( s ) , a ) to estimate Q π ( s, a ) on the LHS of the Bellman error, where aug 1 , aug 2 are random augmentations.
2. D ATA -E FFICIENT R AINBOW (DER)
Rainbow (Hessel et al., 2018) maps directly to Q-value estimates for discrete actions.
Compared to vanilla DQN (Mnih et al., 2013), several reﬁnements are used: Q-values are sampled from a multivariate Gaussian probabilistically (Dabney et al., 2018), noise is injected into network parameters (Plappert et al., 2017), double Q networks (Van Hasselt et al., 2016), dueling DQNs (Wang et al., 2016), n-step returns (Watkins, 1989), and mini-batches are sampled from a prioritized experience replay (Schaul et al., 2015). “Data-efﬁcient” refers to the Atari sample limit of 100k environmental interactions, a much more challenging setting for notoriously inefﬁcient RL.
R ELATED W ORK 3. D ATA -E FFICIENT RL
Data-efﬁcient RL is a paramount concern to the practicality of RL in real-world use cases. Image augmentation has proven an extremely effective regularizer for improving the sample efﬁciency of model-free off-policy RL algorithms (Yarats et al., 2019a; Srinivas et al., 2020; Sermanet et al., 2018;
Dwibedi et al., 2018), so much so that basic augmentation techniques sufﬁce to rival or surpass model-based RL algorithms (Hafner et al., 2019c; Lee et al., 2019b; Hafner et al., 2019b) in the sample efﬁciency metric.
Curiously, these methods have become progressively simpler. CURL (Srinivas et al., 2020) employed contrastive learning, using positive and negative samples extracted from the mini-batch, requiring a computation of quadratic complexity w.r.t.
the mini-batch size. RAD (Laskin et al., 2020) and DrQ (Yarats et al., 2021b) showed that simpler methods work just as well or better, either by augmenting images naively, or augmenting and averaging their Q-values, respectively. Even more recently (contemporaneously), the as-yet unpublished DrQv2 (Yarats et al., 2021a) found basic augmentation alone sufﬁces under a DDPG-based algorithm, with signiﬁcant efﬁciency improvements over prior methods despite the exceptional simplicity. While these recent methods have traded CURL’s mini-batch statistics for mere augmentation, rQdia marks the ﬁrst combination of the two that uses mini-batch statistics to enforce consistency across Q-value distributions, in a manner both simple and complementary to the above implements.
3. M INI -B ATCH R EGULARIZATION I N RL
Mini-batch-based regularization has not commonly been used in RL.
For example, Batch
Norm (Ioffe & Szegedy, 2015), a common regularizer in computer vision, is not as notably employed in RL.
This type of regularization remained unexploited in RL until CURL (Srinivas et al., 2020)
showed that contrastive learning via image augmentation (Chen et al., 2020; He et al., 2020; Misra & Maaten, 2020; Henaff, 2020) greatly improved RL data efﬁciency. CURL contrasts a state’s “positive” augmentation sample with the rest of the mini-batch’s “negative” augmentation samples. These negative samples can be thought of as sampled from a Uniform distribution over the agent’s experience replay, inspiring rQdia.
However, contrastive learning enforces a non-guaranteed ground truth, disassociating negative samples regardless of actual similarity. See Figure 2 for an example of this negative samples problem, where similar states are contrasted to have dissimilar encodings as an inadvertent byproduct of the uniform randomness.
rQdia bypasses this ﬂaw by only enforcing a guaranteed constraint: that the Q-value for any sampled action, regardless of statistical distribution, be consistent across the same states invariant to augmentation. This indeed should always be the case, thus yielding gains over CURL while remaining complementary to methods like DrQ and DrQv2.
3. I MAGE A UGMENTATION
Image augmentation is commonly used to counter over-ﬁtting in computer vision.
Techniques include color shift, afﬁnity translation, scale, etc. (Ciregan et al., 2012; Cires¸an et al., 2011; Simard et al., 2003; Krizhevsky et al., 2012; Chen et al., 2020).
In Yarats et al.
(2021b), the authors investigated several common image transformations and concluded that random shifts strike a good balance between simplicity and performance for the MuJoCo environments.
A variety of different augmentations are useful for different games in the Procgen benchmark (Raileanu et al., 2020), and Yarats et al.
(2021b)
used Intensity variation for the Atari environments.
These techniques have proven critical to MuJoCo from pixels.
M ETHODS
rQdia is the ﬁrst RL Q-value regularizer that does not necessarily depend on either the on-policy or the off-policy states and actions.
n, m such states and actions are instead sampled from arbitrary distributions, let’s call them D 1 and D 2 , of state space S and action space A , respectively:
[Equation]
Then, the following constraint is enforced for Q φ ( · ) , the neural network that models Q π ( · ) the
Q-value function:
[Equation]
where aug ( · ) is an arbitrary augmentation transform.
aug ( · ) and D 1 , D 2 could vary. For aug ( · ) , in line with Yarats et al.
(2021b), we use translation for MuJoCo and intensity jittering for Atari. Speciﬁcally, to translate, we pad by 4 pixels, then crop randomly inward by 4 pixels; to intensity jitter, each image is multiplied by some random noise s = 1 . 1 × clip ( r, − 2 , 2) , r ∼ N (0 , 1) . For D 1 , D 2 in MuJoCo continuous action spaces, we implement a simple approach analogous to
CURL’s negative-sample sampling, except sampling both states and actions rather than just states.
That is, states and actions are directly lifted from the mini-batch ( n = m are the batch size), in effect sampled from a Uniform distribution over the agent’s experience replay.
By using historical states and actions as opposed to random noise, we compute Q-value distributions over state-action pairs that could more feasibly be encountered in a deployed roll-out.
In discrete Atari, actions may be lifted from the action space directly, that is, a (0) , ..., a ( m − 1) =
A .
Thus, given mini-batch states s (0) , ..., s ( n − 1) and mini-batch (or action space) actions a (0) , ..., a ( m − 1) , the following auxiliary loss is proposed to constitute a basic implement of rQdia:
[Equation]
Then this auxiliary loss is simply added to the RL agent’s standard loss term.
Voila, rQdia (visualized “in a nutshell” in Figure 1).
This is applied in parallel for each s ( i ) , a ( j ) pair.
If mini-batches or actions spaces are very large, it is possible to convolve a smaller subset of n states and m actions.
In MuJoCo, we note that while D 1 , D 2 are treated as Uniform distributions over an agent’s history similar to CURL’s negative sampling, D 1 , D 2 could be more sophisticated. For example, the probability of sampling an action could be proportional to state similarities. Or, like MPO (Abdolmaleki et al., 2018), actions could be sampled directly from the policy itself.
Algorithm 1 provides pseudo-code for rQdia in continuous control algorithms like SAC-AE and
DrQ. All code for rQdia will be released open-source.
E XPERIMENTS
To measure the data-efﬁciency and overall performance of rQdia, we conducted experiments at 100k and 500k steps in the DeepMind Continuous Control Suite (MuJoCo) from pixels and 100k interaction steps in the Atari Arcade Learning Environments.
In the following, a ( j ) are all of the actions in the mini-batch for MuJoCo, a ( j ) are all of the actions in the action space for Atari, with s ( i ) as all of the states in the mini-batch and the same MSE-based loss as deﬁned in Equation 3 for both.
repeat if it’s time to update then for j in range(however many updates) do
Randomly sample a batch of transitions, B = { ( s, a, r, s (cid:48) , d ) } from D .
Compute targets for the Q functions:
[Equation]
Update Q-functions (minimize Bellman error) by one step of gradient descent using:
[Equation]
for i = 1 ,
Update rQdia by one step of gradient descent using:
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Note that s ( i ) and a ( j ) are not necessarily the historically paired state and action.
Update policy by one step of gradient ascent using:
[Equation]
where ˜ a θ ( s ) is a sample from π θ ( ·| s ) which is differentiable wrt θ via the reparametrization trick.
Update target networks with:
[Equation]
Most deep RL benchmarks report point estimates of aggregate performance such as mean and median scores across task runs, ignoring the statistical uncertainties that are a natural consequence of training with a ﬁnite number of random seeds.
In the recent analysis by Agarwal et al.
, the authors observe that viewing reported mean scores as random quantities that depend on a small number of sample runs exhibits substantial variability, and demonstrate that a lot of the reported improvements from previous works disproportionately beneﬁted from randomness in the experimental protocol. To account for the variability of results in RL, they propose a number of statistical best-practice protocols.
We followed these best recommended practices as closely as possible and report results in accord with their measured benchmark performances using their open-source library for RL statistical analysis, rliable ( https://github.com/google-research/rliable ).
Compared to previous works, we evaluate our method with this more thorough statistical analysis and prove rQdia with the recommended robust and efﬁcient aggregate metrics in Figures 3 and 4.
Sans stats, raw results show rQdia boosts DrQ’s 100k MuJoCo sample efﬁciency on 4/6 tasks by wide margins, despite a smaller batch size, also surpassing ground truth state embeddings on 4/6.
In the already-saturated 500k setting, rQdia boosts DrQ on 4/6 tasks, and additionally surpasses state embeddings on the Cheetah Run task. In Atari 100k, rQdia affords DER clear gains over DrQ and CURL, superseding mean baseline human-norm scores especially by near 200% .
[Figure]
We compare DrQ+rQdia with SAC+AE (Yarats et al., 2019b), SLAC (Lee et al., 2019a), Dreamer (Hafner et al., 2019a), CURL (Srinivas et al., 2020), RAD (Laskin et al.), DrQ (Yarats et al., 2021b)
[Figure] and PISAC (Lee et al., 2020). Due to our own computational limits, we used a batch size of 128, while other methods used 512. According to Yarats et al.
(2021a), bigger batch sizes yield bigger improvements. At 1/4 the batch size of other methods, rQdia still surpasses or matches benchmarked scores. (a) Although pure stochastic dominance is rarely observed (Dror et al., 2019), rQdia outperforms others at 100 steps, meaning we achieve higher scores per number of runs. At 500k steps, we still reach SOTA performance despite the hampered, more-efﬁcient batch size.
(b) If the lower bound of an algorithm’s interval is higher than another algorithm’s upper bound, there is a high conﬁdence that the algorithm is better.
rQdia not only outperforms the others but also has a relatively small interval, meaning the result is better and also more consistent across different runs.
(c) rQdia not only ranks ﬁrst, but has a high probability of being ranked ﬁrst on 5/6 tasks, indicating the holistic statistical performance exceeds baselines by a decisive margin.
In line with recent works, we evaluate 6 tasks in the DeepMind Continuous Control Suite at 100k and 500k steps to measure data efﬁciency and asymptotic performance respectively.
The reported improvements are based on higher mean scores per task, with large variability across random seeds.
When accounting for this variability, it turns out that many previous algorithms do not consistently rank above the algorithms they claim to improve on (Agarwal et al., 2021).
Therefore, the results and metrics aggregated by Agarwal et al.
and reported in Figure 3 are not necessarily consistent with the ones reported in previous papers.
We followed the protocol of Agarwal et al.
and instead stick to their benchmarked curves for our analysis.
Unlike prior works, our computational constraints required that we use a smaller batch size.
All of our rQdia results are reported with this smaller batch size.
Yarats et al.
(2021a) report that smaller batch sizes are disadvantaged.
Performance Proﬁle
The performance proﬁle shows the tail distribution of scores on combined runs across tasks, thus allowing us to compare different methods at a glance.
If one curve is strictly higher than another, it is said to “stochastically dominate” (Dror et al., 2019).
Figure 3a) indicates that at 100k, rQdia can achieve the same score as other algorithms with fewer runs and a higher score with equal runs.
At 500k, rQdia matches SOTA performance despite 1/4 the batch size.
[Table]
Interval Estimates We resampled with replacement independently for each task to construct an empirical bootstrap sample in which we computed 95% stratiﬁed bootstrap conﬁdence intervals (CIs). This process is repeated 50000 times to approximate the real sampling distributions. Normalized scores are computed by dividing by the maximum score ( = 1000 ). rQdia yields a high conﬁdence with a small interval, meaning its performance is both statistically better and more consistent.
Rank Comparisons show the probability that a given method is assigned rank i , averaged across all tasks.
The ranks are estimated using 200,000 stratiﬁed bootstrap re-samples.
rQdia ranks highest on 5/6 tasks with high probability.
5.1. R OBUSTNESS T O B ATCH S IZE
We report tabular results in Table 1. rQdia excels at sample efﬁciency, achieving SOTA results at the 100k benchmark, even rivaling the state embedding ground truth. At 500k, where results are already saturated, rQdia yields less pronounced improvements in terms of ﬁnal score, but notably attains SOTA scores despite a smaller batch size.
Yarats et al. (2021a) report that performance hinged on the whopping 512 batch size. Due to computational limits, we could not reproduce this, but achieved competitive results nevertheless at 1/4 the batch size, robustly using just 128. 5. A TARI A RCADE 100 K
To conduct a thorough statistical analysis of Atari 100k performance, we evaluated rQdia with the performance proﬁle described in Section 5.1.
We further scaled the x-axis such that the space between any τ 1 and τ 2 is proportional to the fraction of runs averaged across algorithms between τand τ 2 . This scaling shows the regions of the score distribution where most of the runs lie as opposed to comparing tail ends of the distribution. In addition to median and mean, we further report two additional aggregate metrics as recommended in (Dror et al., 2019), where all scores are computed with 95% stratiﬁed bootstrap conﬁdence intervals.
Our case study compares the performance of ﬁve recent deep RL algorithms, namely: (1) DER (van Hasselt et al., 2019), OTR(Kielak, 2020), SimPle (Kaiser et al., 2019), DrQ (Yarats et al., 2021b), and CURL (Srinivas et al., 2020). We also include SPR (Schwarzer et al., 2020), a slightly-apples-to-oranges-baseline which learns a selfsupervised environment dynamics model, to which rQdia is orthogonal to and amazingly approaches the performance of despite the substantial difference in simplicity.
Raw tabular results for Atari 100k are presented in Table 2, with benchmarks likewise pulled from rliable. [Table] Interquartile Mean (IQM) discards the bottom and top 25% of the runs and calculates the mean score of the remaining 50% runs
[Table]
which serves as an interpolation between mean and median across runs IQM is robust to outliers compared to mean and has considerably less bias than median.
Optimality Gap is the amount by which the algorithm fails to meet a minimum score (human score), which serves as a robust alternative to mean.
This metric assumes that a score of humanlevel performance is a desirable target beyond which improvements are not very important.
Probability Of Improvement is designed to measure how likely it is for X to outperform Y on a randomly selected task, which is computed by the Mann-Whitney U-statistic (Mann & Whitney, 1947). The interval reported estimates are based on stratiﬁed bootstrap with independent sampling with 2000 bootstrap re-samples.
D ISCUSSION
Limitations One limitation of rQdia is that it assumes a beneﬁt to a certain augmentation invariance.
Translation invariance for example might not be as useful in environments where most objects are held within a consistent axis.
Moreover, rQdia ensures consistency between Q-value distributions across such perturbations, which means that models in environments that do not require such visual invariances are needlessly expected to learn a more complex, more general Q function.
[Figure]
Notably, DER+rQdia, with simple image augmentation, rivals the SOTA results of the orthogonalpotentially-complementary SPR (Schwarzer et al., 2020), which learns a computationally intensive environment dynamics model.
We include SPR just for reference. On the other hand, such environments where these invariances are not useful or important may leverage rQdia to learn more invariant representations that could potentially better generalize to different, more complex environments.
Ethics rQdia is a simple regularizer that contributes to the generalization of deep reinforcement learning models.
While we hope deep RL continues to improve and its applications and abilities expand, we would be remiss not to note the destructive potential of the ﬁeld, ranging from autonomous weaponry to economic exploitation.
However, we are optimistic that RL can do much more good than bad for society.
Autonomous agents that can interact with the real world via RL-based, morestreamlined robotics opens the door for countless medical, social, and economic beneﬁts as well.
Reproducability We will release all code open-source; we have submitted code together with the paper; relevant code snippets are shared in Appendix C; continuous control pseudocode provided in Algorithm 1; hyperparams speciﬁed in Appendix B.
C ONCLUSION
We presented a simple regularizer for model-free reinforcement learning that may easily be integrated into existing reinforcement learning frameworks.
With the inclusion of this auxiliary loss, we attain strong performance compared to baseline models, including recent state of the arts.
By regularizing Q-value distributions, we further enforce the invariances afforded by image augmentation techniques such that Q-value distributions are preserved under these perturbations.
Consequently, we observe improvements in sample efﬁciency and ﬁnal reward in the DeepMind Continuous Control Suite and environments in the Atari Arcade Learning Environments.
A A RCHITECTURE
The SAC-AE base architecture we use for DrQ is the same as [38], consisting of a shared encoder and distinct policy and Q-function heads.
The CNN encoder is shared by the actor and critic; the critic consists of two ReLU-activated 3-layer MLP Q-networks; and the actor is a single ReLUactivated 3-layer MLP Gaussian policy head.
We modify the code provided by [30]: https: //github.com/
MishaLaskin/curl .
Atari environments were tested with a Rainbow architecture inspired by [35] and built on the variant implemented in tandem with CURL in [30].
We added the rQdia auxiliary loss to their code sans the CURL-related portion.
This code may be found here: https://github.com/ aravindsrinivas/curl_rainbow .
B T RAINING
All hyperparameters were preserved from the original implementations discussed above.
They are reviewed in Tables 3 and 4, except for the substitution of batch size since we used a batch size of 128 while [30, 37] used 512 , reportedly giving those models a decent advantage. [Table] [Table]
[Equation]
Code for continuous control and discrete Atari will be released on GitHub and is provided in the supplementary material.
The rQdia loss in Rainbow is a simple mean squared error between the anchor and augmentation’s respective Q-value distributions (Figure 5).
Continuous control involves a bit more handiwork, but is also simple to tweak into an existing RL library (Figure 6).
First, mini-batch actions and states have to be convolved in pairs with one another.
A scaling factor ∈ (0 , 1] can be modiﬁed for efﬁciency to determine how many such pairs should be used.
Then the double-critics predict a Q-value distribution for the convolved pairs, which is minimized w.r.t.
the augmentation.