Original: I MAGE A UGMENTATION I S A LL Y OU N EED : R EGULARIZING D EEP R EINFORCEMENT L EARNING FROM P IXELS
Revised: I MAGE A UGMENTATION I S A LL Y OU N EED : R EGULARIZING D EEP R EINFORCEMENT L EARNING FROM P IXELS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Existing model-free reinforcement learning (RL) approaches are effective when trained on states but struggle to learn directly from image observations.
Revised: Existing model- free reinforcement learning RL) approaches are effective when trained on states but struggle to learn directly from image observations.

Original: We propose an augmentation technique that can be applied to standard model-free
Revised: We propose an augmentation technique that can be applied to standard model-free

Original: RL algorithms, enabling robust learning directly from pixels without the need for auxiliary losses or pre-training.
Revised: RL algorithms, enabling robust learning directly from pixels without the need for auxiliary losses or pre-training.

Original: The approach leverages input perturbations commonly used in computer vision tasks to transform input examples, as well as regularizing the value function and policy.
Revised: The approach leverages input perturbations commonly used in computer vision tasks to transform input examples, as well as regularizing the value function and policy.

Original: Our approach reaches a new stateof-the-art performance on DeepMind control suite and Atari 100k benchmark, surpassing previous model-free (Haarnoja et al., 2018; van Hasselt et al., 2019a), model-based (Hafner et al., 2019; Lee et al., 2019; Hafner et al., 2018; Kaiser et al., 2019) and contrastive learning (Srinivas et al., 2020) approaches.
Revised: Our approach reaches a new stateof-the-art performance on DeepMind control suite and Atari 100k benchmark, surpassing previous model-free (Haarnoja et al., 2018; van Hasselt et al., 2019a) model-based (Hafner et al., 2019; Lee et al., 2019; Hafner et al., 2018; Kaiser et al., 2019) and contrastive learning (Srinivas et al., 2020) approaches.

Original: It also closes the gap between state-based and image-based RL training.
Revised: It also closes the gap between state-based and image-based RL training.

Original: Our method, which we dub
Revised: Our method which we dub

Original: DrQ : D ata- r egularized Q , can be combined with any model-free RL algorithm.
Revised: DrQ : D ata - r egularized Q , can be combined with any model-free RL algorithm.

Original: To the best of our knowledge, our approach is the ﬁrst effective data augmentation method for RL on these benchmarks.
Revised: To the best of our knowledge, our approach is the ﬁrst effective data augmentation method for RL on these benchmarks.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Sample-efﬁcient deep reinforcement learning (RL) algorithms capable of directly training from image pixels would open up many real-world applications in control and robotics.
Revised: Sample-efﬁcient deep reinforcement learning (RL) algorithms capable of directly training from image pixels would open up many real-world applications in control and robotics.

Original: However, simultaneously training a convolutional encoder alongside a policy network is challenging when given limited environment interaction, strong correlation between samples and a typically sparse reward signal.
Revised: However, simultaneously training a convolutional encoder alongside a policy network is challenging when given limited environment interaction, strong correlation between samples and a typically sparse reward signal.

Original: Limited supervision is a common problem across AI and two approaches are commonly taken: (i) training with an additional auxiliary losses, such as those based on self-supervised learning (SSL) and (ii) training with data augmentation.
Revised: Limited supervision is a common problem across AI and two approaches are commonly taken: (i) training with an additional auxiliary losses, such as those based on self-supervised learning (SSL) and (ii) training with data augmentation.

Original: A wide range of auxiliary loss functions have been proposed to augment supervised objectives, e.g. weight regularization, noise injection (Hinton et al., 2012), or various forms of auto-encoder (Kingma et al., 2014).
Revised: A wide range of auxiliary loss functions have been proposed to augment supervised objectives, e.g. weight regularization, noise injection (Hinton et al., 2012), or various forms of auto-encoder (Kingma et al., 2014).

Original: In RL, reconstruction losses (Jaderberg et al., 2017; Yarats et al., 2019) or SSL objectives (Dwibedi et al., 2018; Srinivas et al., 2020) are used.
Revised: In RL, reconstruction losses (Jaderberg et al., 2017; Yarats et al., 2019) or SSL objectives (Dwibedi et al., 2018; Srinivas et al., 2020) are used.

Original: However, these objectives are unrelated to the task at hand, thus have no guarantee of inducing an appropriate representation for the policy network.
Revised: However, these objectives are unrelated to the task at hand, thus have no guarantee of inducing an appropriate representation for the policy network.

Original: SSL losses are highly effective in the large data regime, e.g. in domains such as vision (Chen et al., 2020;
Revised: SSL losses are highly effective in the large data regime, e.g. in domains such as vision (Chen et al., 2020;

Original: He et al., 2019) and NLP (Collobert et al., 2011;
Revised: He et al., 2019) and NLP (Collobert et al., 2011;

Original: Devlin et al., 2018) where large (unlabeled) datasets are readily available.
Revised: Devlin et al., 2018) where large (unlabeled) datasets are readily available.

Original: However, in sample-efﬁcient RL, training data is more limited due to restricted interaction between the agent and the environment, limiting their effectiveness.
Revised: However, in sample-efﬁcient RL, training data is more limited due to restricted interaction between the agent and the environment, limiting their effectiveness.

Original: Data augmentation methods are widely used in vision and speech domains, where output-invariant perturbations can easily be applied to the labeled input examples.
Revised: Data augmentation methods are widely used in vision and speech domains, where output-invariant perturbations can easily be applied to the labeled input examples.

Original: Surprisingly, data augmentation has received little attention in the RL community.
Revised: Surprisingly, data augmentation has received little attention in the RL community.

Original: In this paper we propose augmentation approaches appropriate for sample-efﬁcient RL and comprehensively evaluate them.
Revised: In this paper we propose augmentation approaches appropriate for sample-efﬁcient RL and comprehensively evaluate them.

Original: The key idea of our approach is to use standard image transformations to perturb input observations, as well as regularizing the
Revised: The key idea of our approach is to use standard image transformations to perturb input observations, as well as regularizing the

Original: Q -function learned by the critic so that different transformations of the same input image have similar Q -function values.
Revised: Q -function learned by the critic so that different transformations of the same input image have similar Q -function values.

Original: No further modiﬁcations to standard actor-critic algorithms are required.
Revised: No further modiﬁcations to standard actor-critic algorithms are required.

Original: Our study is, to the best of our knowledge, the ﬁrst careful examination of image augmentation in sample-efﬁcient RL.
Revised: Our study is, to the best of our knowledge, the ﬁrst careful examination of image augmentation in sample-efﬁcient RL.

Original: The main contributions of the paper are as follows: (i) the ﬁrst to demonstrate that data augmentation greatly improves performance when training model-free RL algorithms from images; (ii) introducing a natural way to exploit MDP structure through two mechanisms for regularizing the value function, in a manner that is generally applicable to model-free RL and (iii) setting a new state-of-the-art performance on the standard DeepMind control suite (Tassa et al., 2018), closing the gap between learning from states, and Atari 100k (Kaiser et al., 2019) benchmarks.
Revised: The main contributions of the paper are as follows: (i) the ﬁrst to demonstrate that data augmentation greatly improves performance when training model-free RL algorithms from images; (ii) introducing a natural way to exploit MDP structure through two mechanisms for regularizing the value function, in a manner that is generally applicable to model-free RL and (iii) setting a new state-of-the-art performance on the standard DeepMind control suite (Tassa et al., 2018), closing the gap between learning from states, and Atari 100k (Kaiser et al., 2019) benchmarks.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Data Augmentation in Computer Vision Data augmentation via image transformations has been used to improve generalization since the inception of convolutional networks (Becker & Hinton, 1992; Simard et al., 2003; LeCun et al., 1989; Ciresan et al., 2011; Ciregan et al., 2012).
Revised: Data Augmentation in Computer Vision Data augmentation via image transformations has been used to improve generalization since the inception of convolutional networks (Becker & Hinton, 1992; Simard et al., 2003; LeCun et al., 1989; Ciresan et al., 2011; Ciregan et al., 2012).

Original: Following
Revised: Following

Original: AlexNet (Krizhevsky et al., 2012), they have become a standard part of training pipelines.
Revised: AlexNet (Krizhevsky et al., 2012), they have become a standard part of training pipelines.

Original: For object classiﬁcation tasks, the transformations are selected to avoid changing the semantic category, i.e. translations, scales, color shifts, etc.
Revised: For object classiﬁcation tasks, the transformations are selected to avoid changing the semantic category, i.e. translations, scales, color shifts, etc.

Original: While a similar set of transformations are potentially applicable to control tasks, the RL context does require modiﬁcations to be made to the underlying algorithm.
Revised: While a similar set of transformations are potentially applicable to control tasks, the RL context does require modiﬁcations to be made to the underlying algorithm.

Original: Data augmentation methods have also been used in the context of self-supervised learning.
Revised: Data augmentation methods have also been used in the context of self-supervised learning.

Original: Dosovitskiy et al.
Revised: Dosovitskiy et al.

Original: (2016) use per-exemplar perturbations in a unsupervised classiﬁcation framework.
Revised: (2016) use per-exemplar perturbations in a unsupervised classiﬁcation framework.

Original: More recently, several approaches (Chen et al., 2020; He et al., 2019; Misra & van der Maaten, 2019) have used invariance to imposed image transformations in contrastive learning schemes, producing state-of-the-art results on downstream recognition tasks.
Revised: More recently, several approaches (Chen et al., 2020; He et al., 2019; Misra & van der Maaten, 2019) have used invariance to imposed image transformations in contrastive learning schemes, producing state-of-the-art results on downstream recognition tasks.

Original: By contrast, our scheme addresses control tasks, utilizing different types of invariance.
Revised: By contrast, our scheme addresses control tasks, utilizing different types of invariance.

Original: Data Augmentation in RL In contrast to computer vision, data augmentation is rarely used in RL.
Revised: Data Augmentation in RL In contrast to computer vision, data augmentation is rarely used in RL.

Original: Certain approaches implicitly adopt it, for example Levine et al.
Revised: Certain approaches implicitly adopt it, for example Levine et al.

Original: ; Kalashnikov et al.
Revised: ; Kalashnikov et al.

Original: (2018) use image augmentation as part of the AlexNet training pipeline without analysing the beneﬁts occurring from it, thus being overlooked in subsequent work.
Revised: (2018) use image augmentation as part of the AlexNet training pipeline without analysing the beneﬁts occurring from it, thus being overlooked in subsequent work.

Original: HER (Andrychowicz et al., 2017) exploits information about the observation space by goal and reward relabeling, which can be viewed as a way to perform data augmentation.
Revised: HER (Andrychowicz et al., 2017) exploits information about the observation space by goal and reward relabeling, which can be viewed as a way to perform data augmentation.

Original: Other work uses data augmentation to improve generalization in domain transfer (Cobbe et al., 2018).
Revised: Other work uses data augmentation to improve generalization in domain transfer (Cobbe et al., 2018).

Original: However, the classical image transformations used in vision have not previously been shown to deﬁnitively help on standard RL benchmarks.
Revised: However, the classical image transformations used in vision have not previously been shown to deﬁnitively help on standard RL benchmarks.

Original: Concurrent with our work, RAD (Laskin et al., 2020) performs an exploration of different data augmentation approaches, but is limited to transformations of the image alone, without the additional augmentation of the Q-function used in our approach.
Revised: Concurrent with our work, RAD (Laskin et al., 2020) performs an exploration of different data augmentation approaches, but is limited to transformations of the image alone, without the additional augmentation of the Q-function used in our approach.

Original: Moreover, RAD can be regarded as a special case of our algorithm.
Revised: Moreover, RAD can be regarded as a special case of our algorithm.

Original: Multiple follow ups to our initial preprint appeared on ArXiv (Raileanu et al., 2020; Okada & Taniguchi, 2020), using similar techniques on other tasks, thus supporting the effectiveness and generality of data augmentation in RL.
Revised: Multiple follow ups to our initial preprint appeared on ArXiv (Raileanu et al., 2020; Okada & Taniguchi, 2020), using similar techniques on other tasks, thus supporting the effectiveness and generality of data augmentation in RL.

Original: Continuous Control from Pixels There are a variety of methods addressing the sample-efﬁciency of RL algorithms that directly learn from pixels.
Revised: Continuous Control from Pixels There are a variety of methods addressing the sample-efﬁciency of RL algorithms that directly learn from pixels.

Original: The most prominent approaches for this can be classiﬁed into two groups, model-based and model-free methods.
Revised: The most prominent approaches for this can be classiﬁed into two groups, model-based and model-free methods.

Original: The model-based methods attempt to learn the system dynamics in order to acquire a compact latent representation of high-dimensional observations to later perform policy search (Hafner et al., 2018;
Revised: The model-based methods attempt to learn the system dynamics in order to acquire a compact latent representation of high-dimensional observations to later perform policy search (Hafner et al., 2018;

Original: Lee et al., 2019; Hafner et al., 2019).
Revised: Lee et al., 2019; Hafner et al., 2019).

Original: In contrast, the model-free methods either learn the latent representation indirectly by optimizing the
Revised: In contrast, the model-free methods either learn the latent representation indirectly by optimizing the

Original: RL objective (Barth-Maron et al., 2018; Abdolmaleki et al., 2018) or by employing auxiliary losses that provide additional supervision (Yarats et al., 2019; Srinivas et al., 2020; Sermanet et al., 2018;
Revised: RL objective (Barth-Maron et al., 2018; Abdolmaleki et al., 2018) or by employing auxiliary losses that provide additional supervision (Yarats et al., 2019; Srinivas et al., 2020; Sermanet et al., 2018;

Original: Dwibedi et al., 2018).
Revised: Dwibedi et al., 2018).

Original: Our approach is complementary to these methods and can be combined with them to improve performance.
Revised: Our approach is complementary to these methods and can be combined with them to improve performance.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Reinforcement Learning from Images We formulate image-based control as an inﬁnite-horizon partially observable Markov decision process (POMDP) (Bellman, 1957; Kaelbling et al., 1998).
Revised: Reinforcement Learning from Images We formulate image-based control as an inﬁnite-horizon partially observable Markov decision process (POMDP) (Bellman, 1957; Kaelbling et al., 1998).

Original: An
Revised: An

Original: POMDP can be described as the tuple ( O , A , p, r, γ ) , where O is the high-dimensional observation space (image pixels), A is the action space, the transition dynamics p =
Revised: POMDP can be described as the tuple ( O , A , p, r, γ ) , where O is the high-dimensional observation space (image pixels), A is the action space, the transition dynamics p =

Original: Pr ( o (cid:48) t | o ≤ t , a t ) capture the probability distribution over the next observation o (cid:48)
Revised: Pr ( o (cid:48) t | o ≤ t , a t ) capture the probability distribution over the next observation o (cid:48)

Original: t given the history of previous observations o ≤ t and current action a t , r : O × A → R is the reward function that maps the current observation and action to a reward r t = r
Revised: t given the history of previous observations o ≤ t and current action a t , r : O × A → R is the reward function that maps the current observation and action to a reward r t = r

Original: ( o ≤ t , a t ) , and γ ∈ [0 , 1) is a discount factor.
Revised: ( o ≤ t , a t ) , and γ ∈ [0 , 1) is a discount factor.

Original: Per common practice (Mnih et al., 2013), throughout the paper the POMDP is converted into an MDP (Bellman, 1957) by stacking several consecutive image observations into a state s t = { o t , o t − 1 , o t − 2 , . . .
Revised: Per common practice (Mnih et al., 2013), throughout the paper the POMDP is converted into an MDP (Bellman, 1957) by stacking several consecutive image observations into a state s t = { o t , o t − 1 , o t − 2 , . . .

Original: For simplicity we redeﬁne the transition dynamics p =
Revised: For simplicity we redeﬁne the transition dynamics p =

Original: Pr ( s (cid:48)
Revised: Pr ( s (cid:48)

Original: t | s t , a t ) and the reward function r t = r ( s t , a t ) .
Revised: t | s t , a t ) and the reward function r t = r ( s t , a t ) .

Original: We then aim to ﬁnd a policy π ( a t | s t ) that maximizes the cumulative
Revised: We then aim to ﬁnd a policy π ( a t | s t ) that maximizes the cumulative

Original: discounted return E π [ (cid:80) ∞ t =1 γ t r t | a t ∼ π ( ·| s t ) , s (cid:48)
Revised: discounted return E π [ (cid:80) ∞ t =1 γ t r t | a t ∼ π ( ·| s t ) , s (cid:48)

Original: t ∼ p ( ·| s t , a t ) , s 1 ∼ p ( · )] .
Revised: t ∼ p ( ·| s t , a t ) , s 1 ∼ p ( · )] .

Original: Soft Actor-Critic The Soft Actor-Critic (SAC) (Haarnoja et al., 2018) learns a state-action value function Q θ , a stochastic policy π θ and a temperature α to ﬁnd an optimal policy for an MDP ( S , A , p, r, γ ) by optimizing a γ -discounted maximum-entropy objective (Ziebart et al., 2008).
Revised: The Soft Actor-Critic (SAC) (Haarnoja et al., 2018) learns a state-action value function Q θ , a stochastic policy π θ and a temperature α to ﬁnd an optimal policy for an MDP ( S , A , p, r, γ ) by optimizing a γ -discounted maximum-entropy objective (Ziebart et al., 2008).

Original: θ is used generically to denote the parameters updated through training in each part of the model.
Revised: θ is used generically to denote the parameters updated through training in each part of the model.

Original: Deep Q-learning DQN (Mnih et al., 2013) also learns a convolutional neural net to approximate
Revised: Deep Q-learning DQN (Mnih et al., 2013) also learns a convolutional neural net to approximate

Original: Q-function over states and actions.
Revised: Q-function over states and actions.

Original: The main difference is that DQN operates on discrete actions spaces, thus the policy can be directly inferred from Q-values.
Revised: The main difference is that DQN operates on discrete actions spaces, thus the policy can be directly inferred from Q-values.

Original: In practice, the standard version of DQN is frequently combined with a set of reﬁnements that improve performance and training stability, commonly known as Rainbow (van Hasselt et al., 2015).
Revised: In practice, the standard version of DQN is frequently combined with a set of reﬁnements that improve performance and training stability, commonly known as Rainbow (van Hasselt et al., 2015).

Original: For simplicity, the rest of the paper describes a generic actor-critic algorithm rather than DQN or SAC in particular.
Revised: For simplicity, the rest of the paper describes a generic actor-critic algorithm rather than DQN or SAC in particular.

Original: Further background on DQN and SAC can be found in Appendix A.
Revised: Further background on DQN and SAC can be found in Appendix A.

Original:  S AMPLE E FFICIENT R EINFORCEMENT L EARNING FROM P IXELS 4. O PTIMALITY I NVARIANT I MAGE T RANSFORMATIONS FOR Q FUNCTION
Revised:  S AMPLE E FFICIENT R EINFORCEMENT L EARNING FROM P IXELS 4. O PTIMALITY I NVARIANT I MAGE T RANSFORMATIONS FOR Q FUNCTION

Original: We ﬁrst introduce a general framework for regularizing the value function through transformations of the input state.
Revised: We ﬁrst introduce a general framework for regularizing the value function through transformations of the input state.

Original: For a given task, we deﬁne an optimality invariant state transformation f : S ×T → S as a mapping that preserves the Q -values
Revised: For a given task, we deﬁne an optimality invariant state transformation f : S ×T → S as a mapping that preserves the Q -values

Original: [Equation]
Revised: [Equation]

Original: where ν are the parameters of f ( · ) , drawn from the set of all possible parameters T .
Revised: where ν are the parameters of f ( · ) , drawn from the set of all possible parameters T .

Original: One example of such transformations are the random image translations successfully applied in the previous section.
Revised: One example of such transformations are the random image translations successfully applied in the previous section.

Original: For every state, the transformations allow the generation of several surrogate states with the same
Revised: For every state, the transformations allow the generation of several surrogate states with the same

Original: Q -values, thus providing a mechanism to reduce the variance of Q -function estimation.
Revised: Q -values, thus providing a mechanism to reduce the variance of Q -function estimation.

Original: In particular, for an arbitrary distribution of states µ ( · ) and policy π , instead of using a single sample s ∗ ∼ µ ( · ) , a ∗ ∼ π ( ·| s ∗ ) estimation of the following expectation
Revised: In particular, for an arbitrary distribution of states µ ( · ) and policy π , instead of using a single sample s ∗ ∼ µ ( · ) , a ∗ ∼ π ( ·| s ∗ ) estimation of the following expectation

Original: [Equation]
Revised: [Equation]

Original: we generate K samples via random transformations and obtain an estimate with lower variance 
Revised: we generate K samples via random transformations and obtain an estimate with lower variance E

Original: [Equation]
Revised: [Equation]

Original: Q
Revised: K

Revised: K

Original: [Equation]
Revised: [Equation]

Original: , ν
Revised: k

Original: [Equation]
Revised: [Equation]

Original:  ) , a k ) where ν k ∈ T and a k ∼ π ( ·| f ( s ∗ , ν k )) .
Revised: ν k ) , a k ) where ν k ∈ T and a k ∼ π ( ·| f ( s ∗ , ν k )) .

Original: This suggests two distinct ways to regularize Q -function.
Revised: This suggests two distinct ways to regularize Q -function.

Original: First, we use the data augmentation to compute the target values for every transition tuple ( s i , a i , r i , s (cid:48) i ) as
Revised: First, we use the data augmentation to compute the target values for every transition tuple ( s i , a i , r i , s (cid:48) i ) as

Original: [Equation]
Revised: [Equation]

Original: k ) [Equation] a (cid:48) i,k ∼ π ( ·| f ( s (cid:48)
Revised: where

Original: [Equation]
Revised: [Equation]

Original: , ν (cid:48)
Revised: where ν (cid:48)

Original: [Equation]

Original: k [Equation] where ν (cid:48) i,k ∈ T corresponds to a transformation parameter of s (cid:48)
Revised: i,k ∈ T corresponds to a transformation parameter of s (cid:48)

Original: i .
Revised: i .

Original: Then the Q-function is updated using these targets through an SGD update using learning rate λ θ
Revised: Then the Q-function is updated using these targets through an SGD update using learning rate λ θ

Original: [Equation]
Revised: [Equation]

Original: In tandem, we note that the same target from Equation (1) can be used for different augmentations of s i , resulting in the second regularization approach
Revised: In tandem, we note that the same target from Equation (1) can be used for different augmentations of s i , resulting in the second regularization approach

Original: [Equation]
Revised: [Equation]

Original: When both regularization methods are used, ν i,m and ν (cid:48)
Revised: When both regularization methods are used, ν i,m and ν (cid:48)

Original: i,k are drawn independently.
Revised: i,k are drawn independently.

Original: 4. P RACTICAL I NSTANTIATION OF O PTIMALITY I NVARIANT I MAGE T RANSFORMATION
Revised: 4. P RACTICAL I NSTANTIATION OF O PTIMALITY I NVARIANT I MAGE T RANSFORMATION

Original: A range of successful image augmentation techniques have been developed in computer vision (Ciregan et al., 2012; Ciresan et al., 2011; Simard et al., 2003; Krizhevsky et al., 2012; Chen et al., 2020).
Revised: A range of successful image augmentation techniques have been developed in computer vision (Ciregan et al., 2012; Ciresan et al., 2011; Simard et al., 2003; Krizhevsky et al., 2012; Chen et al., 2020).

Original: These apply transformations to the input image for which the task labels are invariant, e.g. for object recognition tasks, image ﬂips and rotations do not alter the semantic label.
Revised: These apply transformations to the input image for which the task labels are invariant, e.g. for object recognition tasks, image ﬂips and rotations do not alter the semantic label.

Original: However, tasks in
Revised: However, tasks in

Original: RL differ signiﬁcantly from those in vision and in many cases the reward would not be preserved by these transformations.
Revised: RL differ signiﬁcantly from those in vision and in many cases the reward would not be preserved by these transformations.

Original: We examine image transformations from Chen et al.
Revised: We examine image transformations from Chen et al.

Original: (random shifts, random cutouts, horizontal/vertical ﬂips, rotations and intensity shifts) in Appendix E and conclude that random shifts strike a good balance between simplicity and performance, we therefore limit our choice of transformation function f ( · ) to random shifts.
Revised: (random shifts, random cutouts, horizontal/vertical ﬂips, rotations and intensity shifts) in Appendix E and conclude that random shifts strike a good balance between simplicity and performance, we therefore limit our choice of transformation function f ( · ) to random shifts.

Original: We apply shifts to the images sampled from the replay buffer.
Revised: We apply shifts to the images sampled from the replay buffer.

Original: For example, images from the
Revised: For example, images from the

Original: DeepMind control suite used in our experiments are 84 × 84 .
Revised: DeepMind control suite used in our experiments are 84 × 84 .

Original: We pad each side by 4 pixels (by repeating boundary pixels) and then select a random 84 × 84 crop, yielding the original image shifted by ± 4 pixels.
Revised: We pad each side by 4 pixels (by repeating boundary pixels) and then select a random 84 × 84 crop, yielding the original image shifted by ± 4 pixels.

Original: This procedure is repeated every time an image is sampled from the replay buffer.
Revised: This procedure is repeated every time an image is sampled from the replay buffer.

Original: 4. O UR APPROACH : D ATA - REGULARIZED Q (D R Q)
Revised: 4. O UR APPROACH : D ATA - REGULARIZED Q (D R Q)

Original: Our approach, DrQ , is the union of the three separate regularization mechanisms introduced above:
Revised: Our approach, DrQ , is the union of the three separate regularization mechanisms introduced above:

Original: 1. transformations of the input image (Section 4.2).
Revised: 1. transformations of the input image (Section 4.2).

Original: averaging the Q target over K image transformations (Equation (1)).
Revised: averaging the Q target over K image transformations (Equation (1)).

Original: averaging the Q function itself over M image transformations (Equation (3)).
Revised: averaging the Q function itself over M image transformations (Equation (3)).

Original: Algorithm 1 details how they are incorporated into a generic pixel-based off-policy actor-critic algorithm.
Revised: Algorithm 1 details how they are incorporated into a generic pixel-based off-policy actor-critic algorithm.

Original: Note that if [K=1,M=1] then DrQ reverts to image transformations alone, this makes applying DrQ to any model-free RL algorithm straightforward.
Revised: Note that if [K=1,M=1] then DrQ reverts to image transformations alone, this makes applying DrQ to any model-free RL algorithm straightforward.

Original: For the experiments in this paper, we pair DrQ with SAC (Haarnoja et al., 2018) and DQN (Mnih et al., 2013), popular model-free algorithms for control in continuous and discrete action spaces respectively.
Revised: For the experiments in this paper, we pair DrQ with SAC (Haarnoja et al., 2018) and DQN (Mnih et al., 2013), popular model-free algorithms for control in continuous and discrete action spaces respectively.

Original: We select image shifts as the class of image transformations f , with ν ± 4 , as explained in Section 4.2.
Revised: We select image shifts as the class of image transformations f , with ν ± 4 , as explained in Section 4.2.

Original:  E XPERIMENTS 5. A BLATION E XPERIMENT
Revised:  E XPERIMENTS 5. A BLATION E XPERIMENT

Original: Figure 1 shows the effect of image shift augmentation applied to three tasks from the DeepMind control suite (Tassa et al., 2018).
Revised: Figure 1 shows the effect of image shift augmentation applied to three tasks from the DeepMind control suite (Tassa et al., 2018).

Original: Figure 1a shows unmodiﬁed SAC (Haarnoja et al., 2018) parameterized with different image encoders, taken from: NatureDQN (Mnih et al., 2013), Dreamer (Hafner et al., 2019), Impala (Espeholt et al., 2018),
Revised: Figure 1a shows unmodiﬁed SAC (Haarnoja et al., 2018) parameterized with different image encoders, taken from: NatureDQN (Mnih et al., 2013), Dreamer (Hafner et al., 2019), Impala (Espeholt et al., 2018),

Original: SAC-AE (Yarats et al., 2019), and D4PG (Barth-Maron et al., 2018).
Revised: SAC-AE (Yarats et al., 2019), and D4PG (Barth-Maron et al., 2018).

Original: The encoders vary signiﬁcantly in their architecture and capacity, with parametercounts ranging from 220k to 2.4M. None of these train satisfactorily, with performance decreasing for the larger capacity models. Figure 1b shows SAC with the application of our random shifts transformation of the input images (i.e. just Section 4.2, not Q augmentation also). The results for all encoder architectures improve dramatically, suggesting that our method is general and can assist many different encoder architectures. To the best of our knowledge, this is the ﬁrst successful demonstration of applying image augmentation on the standard benchmarks for continuous control.
Revised: The encoders vary signiﬁcantly in their architecture and capacity, with parameter for each timestep t = 1 ..

Original: Furthermore, Figure 2 shows the full DrQ , with both image shifts and Q augmentation (Section 4.1), as well as ablated versions. Q augmentation provides additional consistent gain over image shift augmentation alone (full results are in Appendix F). 5. D EEP M IND C ONTROL S UITE E XPERIMENTS In this section we evaluate our algorithm ( DrQ ) on the two commonly used benchmarks based on the DeepMind control suite (Tassa et al., 2018), namely the PlaNet (Hafner et al., 2018) and
Revised: T do

Original: Dreamer (Hafner et al., 2019) setups.

Original: Throughout these experiments all hyper-parameters of the for each timestep t = 1 ..T do

Original: [Equation]
Revised: [Equation]

Original:  t )
Revised: s t )

Original: [Equation]
Revised: [Equation]

Original: (cid:46) Add a transition to the replay buffer
Revised: (cid:46) Add a transition to the replay buffer

Original: U PDATE C RITIC ( D )
Revised: U PDATE C RITIC ( D )

Original: U PDATE A CTOR ( D ) (cid:46) Data augmentation is applied to the samples for actor training as well end for procedure
Revised: U PDATE A CTOR ( D ) (cid:46) Data augmentation is applied to the samples for actor training as well end for procedure

Original: U PDATE C RITIC
Revised: U PDATE C RITIC

Original: ( D )
Revised: ( D )

Original: [Equation]
Revised: [Equation]

Original: (cid:46)
Revised: (cid:46)

Original: Sample a mini batch from the replay buffer(cid:110) ν
Revised: Sample a mini batch from the replay buffer

Original: [Equation]
Revised: [Equation]

Original: (cid:46)  for each i = 1 ..
Revised: (cid:46) Uniformly sample target for each i = 1 ..

Original: N do
Revised: N do

Original: [Equation]
Revised: [Equation]

Original: (cid:48) i , ν (cid:48)
Revised: i , ν (cid:48)

Original: i,k ) , a (cid:48) i,k )
Revised: i,k ) , a (cid:48) i,k )

Original: y
Revised: y

Original: i ← r ( s i , a i ) + γ ˆ Q i end for { ν i,m |
Revised: i ← r ( s i , a i ) + γ ˆ Q i end for { ν i,m |

Original: ν i,m ∼ U ( T ) , i = 1 ..N, m = 1 ..M } (cid:46) Uniformly sample Q augmentations
Revised: ν i,m ∼ U ( T ) , i = 1 ..N, m = 1 ..M } (cid:46) Uniformly sample Q augmentations

Original: J Q ( θ ) =
Revised: J Q ( θ ) =

Original: 1 N (cid:80) Ni =

Revised: 1 N (cid:80) Ni =

Original: ( Q θ ( s i , a i ) − y i ) 2 or J Q ( θ ) = 1 NM (cid:80) N,Mi,m =1 ( Q θ ( f ( s i , ν i,m ) , a i ) − y i )
Revised: ( Q θ ( s i , a i ) − y i ) 2 or J Q ( θ ) = 1 NM (cid:80) N,Mi,m =1 ( Q θ ( f ( s i , ν i,m ) , a i ) − y i )

Original: θ ← θ − λ θ ∇ θ J Q ( θ ) (cid:46) Update the critic θ (cid:48)
Revised: θ ← θ − λ θ ∇ θ J Q ( θ ) (cid:46) Update the critic θ (cid:48)

Original: ← (1 − τ ) θ (cid:48)
Revised: ← (1 − τ ) θ (cid:48)

Original: + τθ (cid:46) Update the critic target end procedure 
Revised: + τθ (cid:46) Update the critic target end procedure counts ranging from 220k to 2.4M. None of these train satisfactorily, with performance decreasing for the larger capacity models.

Original: [Figure]
Revised: Figure 1b shows SAC with the application of our random shifts transformation of the input images (i.e. just Section 4.2, not Q augmentation also).

Original:  of improvement
Revised: The results for all encoder architectures improve dramatically, suggesting that our method is general and can assist many different encoder architectures. To the best of our knowledge, this is the ﬁrst successful demonstration of applying image augmentation on the standard benchmarks for continuous control. Furthermore, Figure 2 shows the full DrQ , with both image shifts and Q augmentation (Section 4.1), as well as ablated versions. Q augmentation provides additional consistent gain over image shift augmentation alone (full results are in Appendix F).

Original: [Figure]
Revised: In this section we evaluate our algorithm ( DrQ ) on the two commonly used benchmarks based on the DeepMind control suite (Tassa et al., 2018), namely the PlaNet (Hafner et al., 2018) and Dreamer (Hafner et al., 2019) setups. Throughout these experiments all hyper-parameters of the algorithm are kept ﬁxed: the actor and critic neural networks are trained using the Adam optimizer (Kingma & Ba, 2014) with default parameters and a mini-batch size of 512 1 .

Original:  SAC, the soft target update rate τ is 0 .
Revised: For SAC, the soft target update rate τ is 0 .

Original: 01 , initial temperature is 0 .
Revised: 01 , initial temperature is 0 .

Original: 1 , and target network and the actor updates are made every 2 critic updates (as in Yarats et al.
Revised: 1 , and target network and the actor updates are made every 2 critic updates (as in Yarats et al.

Original: We use the image encoder architecture from SAC-AE (Yarats et al., 2019) and follow their training procedure.
Revised: We use the image encoder architecture from SAC-AE (Yarats et al., 2019) and follow their training procedure.

Original: The full set of parameters can be found in Appendix B. Following Henderson et al.
Revised: The full set of parameters can be found in Appendix B. Following Henderson et al.

Original: (2018), the models are trained usingdifferent seeds; for every seed the mean episode returns are computed every 10000 environment steps, averaging over 10 episodes.
Revised: (2018), the models are trained usingdifferent seeds; for every seed the mean episode returns are computed every 10000 environment steps, averaging over 10 episodes.

Original: All ﬁgures plot the mean performance over the 10 seeds, together with ± 1 standard deviation shading.
Revised: All ﬁgures plot the mean performance over the 10 seeds, together with ± 1 standard deviation shading.

Original: We compare our DrQ approach to leading model-free and model-based  approaches: PlaNet (Hafner et al., 2018), SAC-AE (Yarats et al., 2019), SLAC (Lee et al., 2019),
Revised: We compare our DrQ approach to leading model-free and model-based [Figure] [Figure] approaches: PlaNet (Hafner et al., 2018), SAC-AE (Yarats et al., 2019), SLAC (Lee et al., 2019),

Original: CURL (Srinivas et al., 2020) and Dreamer (Hafner et al., 2019).
Revised: CURL (Srinivas et al., 2020) and Dreamer (Hafner et al., 2019).

Original: The comparisons use the results provided by the authors of the corresponding papers.
Revised: The comparisons use the results provided by the authors of the corresponding papers.

Original: PlaNet Benchmark (Hafner et al., 2018) consists of six challenging control tasks from (Tassa et al., 2018) with different traits.
Revised: PlaNet Benchmark (Hafner et al., 2018) consists of six challenging control tasks from (Tassa et al., 2018) with different traits.

Original: The benchmark speciﬁes a different action-repeat hyper-parameter for each of the six tasks 2 .
Revised: The benchmark speciﬁes a different action-repeat hyper-parameter for each of the six tasks 2 .

Original: Following common practice (Hafner et al., 2018; Lee et al., 2019; Yarats et al., 2019; Mnih et al., 2013), we report the performance using true environment steps, thus are invariant to  the action-repeat hyper-parameter.
Revised: Following common practice (Hafner et al., 2018; Lee et al., 2019; Yarats et al., [Figure] 2,M=2]) outperforms the other methods and demonstrates the state-of-the-art performance. Furthermore, on several tasks DrQ is able to match the upper-bound performance of SAC trained directly on internal state, rather than images.

Revised: Finally, our algorithm not only shows improved sample-efﬁciency relative to other approaches, but is also faster in terms of wall clock time.

Revised: 2019; Mnih et al., 2013), we report the performance using true environment steps, thus are invariant to the action-repeat hyper-parameter.

Original: Aside from action-repeat, all other hyper-parameters of our algorithm are ﬁxed across the six tasks, using the values previously detailed.
Revised: Aside from action-repeat, all other hyper-parameters of our algorithm are ﬁxed across the six tasks, using the values previously detailed.

Original: Figure 3 compares DrQ [K=2,M=2] to PlaNet (Hafner et al., 2018), SAC-AE (Yarats et al., 2019),
Revised: Figure 3 compares DrQ [K=2,M=2] to PlaNet (Hafner et al., 2018), SAC-AE (Yarats et al., 2019),

Original: CURL (Srinivas et al., 2020), SLAC (Lee et al., 2019), and an upper bound performance provided by
Revised: CURL (Srinivas et al., 2020), SLAC (Lee et al., 2019), and an upper bound performance provided by

Original: SAC (Haarnoja et al., 2018) that directly learns from internal states.
Revised: SAC (Haarnoja et al., 2018) that directly learns from internal states.

Original: We use the version of SLAC that performs one gradient update per an environment step to ensure a fair comparison to other approaches.
Revised: We use the version of SLAC that performs one gradient update per an environment step to ensure a fair comparison to other approaches.

Original: DrQ achieves state-of-the-art performance on this benchmark on all the tasks, despite being much simpler than other methods.
Revised: DrQ achieves state-of-the-art performance on this benchmark on all the tasks, despite being much simpler than other methods.

Original: Furthermore, since DrQ does not learn a model (Hafner et al., 2018; Lee et al., 2019) or any auxiliary tasks (Srinivas et al., 2020), the wall clock time also compares favorably to the other methods.
Revised: Furthermore, since DrQ does not learn a model (Hafner et al., 2018; Lee et al., 2019) or any auxiliary tasks (Srinivas et al., 2020), the wall clock time also compares favorably to the other methods.

Original: In Table 1 we also compare performance given at a ﬁxed number of environment interactions (e.g.
Revised: In Table 1 we also compare performance given at a ﬁxed number of environment interactions (e.g.

Original: 100 k and 500 k).
Revised: 100 k and 500 k).

Original: Furthermore, in Appendix G we demonstrate that DrQ is robust to signiﬁcant changes in hyper-parameter settings.
Revised: Furthermore, in Appendix G we demonstrate that DrQ is robust to signiﬁcant changes in hyper-parameter settings.

Original: Dreamer Benchmark is a more extensive testbed that was introduced in Dreamer (Hafner et al., 2019), featuring a diverse set of tasks from the DeepMind control suite.
Revised: Dreamer Benchmark is a more extensive testbed that was introduced in Dreamer (Hafner et al., 2019), featuring a diverse set of tasks from the DeepMind control suite.

Original: Tasks involving sparse reward were excluded (e.g. Acrobot and Quadruped) since they require modiﬁcation of SAC to incorporate multi-step returns (Barth-Maron et al., 2018), which is beyond the scope of this work.
Revised: Tasks involving sparse reward were excluded (e.g. Acrobot and Quadruped) since they require modiﬁcation of SAC to incorporate multi-step returns (Barth-Maron et al., 2018), which is beyond the scope of this work.

Original: We evaluate on the remaining 15 tasks, ﬁxing the action-repeat hyper-parameter to 2 as in Hafner et al.
Revised: We evaluate on the remaining 15 tasks, ﬁxing the action-repeat hyper-parameter to 2 as in Hafner et al.

Original: We compare DrQ [K=2,M=2] to Dreamer (Hafner et al., 2019) and the upper-bound performance of
Revised: We compare DrQ [K=2,M=2] to Dreamer (Hafner et al., 2019) and the upper-bound performance of

Original: SAC (Haarnoja et al., 2018) from states 3 .
Revised: SAC (Haarnoja et al., 2018) from states 3 .

Original: Again, we keep all the hyper-parameters of our algorithm ﬁxed across all the tasks.
Revised: Again, we keep all the hyper-parameters of our algorithm ﬁxed across all the tasks.

Original: In Figure 4, DrQ demonstrates the state-of-the-art results by collectively outperforming Dreamer
Revised: In Figure 4, DrQ demonstrates the state-of-the-art results by collectively outperforming Dreamer

Original: (Hafner et al., 2019), although Dreamer is superior on 3 of the 15 tasks [Figure] =2]) outperforms the other methods and demonstrates the state-of-the-art performance.
Revised: (Hafner et al., 2019), although Dreamer is superior on 3 of the 15 tasks (Walker Run, Cartpole Swingup Sparse and Pendulum Swingup).

Original: Furthermore, on several tasks DrQ is able to match the upper-bound performance of SAC  trained directly on internal state, rather than images.
Revised: On many tasks DrQ approaches the upper-bound performance of SAC (Haarnoja et al., 2018) trained directly on states.

Original: Finally, our algorithm not only shows improved sample-efﬁciency relative to other approaches, but is also faster in terms of wall clock time.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Figure]

Original: (Walker Run, Cartpole Swingup Sparse and Pendulum Swingup). On many tasks DrQ approaches the upper-bound performance of SAC (Haarnoja et al., 2018) trained directly on states.
Revised: p i s o d e R e t u r n [Figure]

Original: 5. A TARI 100 K E XPERIMENTS
Revised: 5. A TARI 100 K E XPERIMENTS

Original: We evaluate DrQ [K=1,M=1] on the Atari 100k benchmark (Kaiser et al., 2019) – a sampleconstrained evaluation for discrete control algorithms.
Revised: We evaluate DrQ [K=1,M=1] on the Atari 100k benchmark (Kaiser et al., 2019) – a sampleconstrained evaluation for discrete control algorithms.

Original: The underlying RL approach to which DrQ is applied is a DQN, combined with double Q-learning (van Hasselt et al., 2015), n-step returns (Mnih et al., 2016), and dueling critic architecture (Wang et al., 2015).
Revised: The underlying RL approach to which DrQ is applied is a DQN, combined with double Q-learning (van Hasselt et al., 2015), n-step returns (Mnih et al., 2016), and dueling critic architecture (Wang et al., 2015).

Original: As per common practice (Kaiser et al., 2019; van Hasselt et al., 2019a), we evaluate our agent for 125k environment steps at the end of training and average its performance over 5 random seeds.
Revised: As per common practice (Kaiser et al., 2019; van Hasselt et al., 2019a), we evaluate our agent for 125k environment steps at the end of training and average its performance over 5 random seeds.

Original: Figure 5 shows the median humannormalized episode returns performance (as in Mnih et al.
Revised: Figure 5 shows the median humannormalized episode returns performance (as in Mnih et al.

Original: of the underlying model, which we refer to as Efﬁcient DQN, in pink.
Revised: of the underlying model, which we refer to as Efﬁcient DQN, in pink.

Original: When DrQ is added there is a signiﬁcant increase in performance (cyan), surpassing OTRainbow (Kielak, 2020) and Data Efﬁcient Rainbow (van Hasselt et al., 2019a).
Revised: When DrQ is added there is a signiﬁcant increase in performance (cyan), surpassing OTRainbow (Kielak, 2020) and Data Efﬁcient Rainbow (van Hasselt et al., 2019a).

Original: DrQ is also superior to CURL (Srinivas et al., 2020) that uses an auxiliary loss built on top of a hybrid between OTRainbow and Efﬁcient rainbow.
Revised: DrQ is also superior to CURL (Srinivas et al., 2020) that uses an auxiliary loss built on top of a hybrid between OTRainbow and Efﬁcient rainbow.

Original: DrQ combined with Efﬁcient DQN thus achieves 
Revised: DrQ combined with Efﬁcient DQN thus achieves state-of-the-art performance, despite being signiﬁcantly simpler than the other approaches. The experimental setup and full results are detailed in Appendix C and Appendix D respectively.

Original: [Figure]
Revised: [Figure]

Original: t u r n
Revised: hu

Original: [Figure]
Revised: [Figure]

Original: state-of-the-art performance, despite being signiﬁcantly simpler than the other approaches. The experimental setup and full results are detailed in Appendix C and Appendix D respectively.
Revised: n

Original: [Figure]
Revised: [Figure]

Revised: r

Revised: [Figure]

Revised: a li z e d r e t u r n s

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We have introduced a regularization technique, based on image shifts and Q-function augmentation, that signiﬁcantly improves the performance of model-free RL algorithms trained directly from images.
Revised: We have introduced a regularization technique, based on image shifts and Q-function augmentation, that signiﬁcantly improves the performance of model-free RL algorithms trained directly from images.

Original: In contrast to the concurrent work of Laskin et al.
Revised: In contrast to the concurrent work of Laskin et al.

Original: (2020), which is a special case of DrQ , our method exploits the MDP structure of the problem, demonstrating gains over image augmentations alone.
Revised: (2020), which is a special case of DrQ , our method exploits the MDP structure of the problem, demonstrating gains over image augmentations alone.

Original: Our method is easy to implement and adds a negligible computational burden.
Revised: Our method is easy to implement and adds a negligible computational burden.

Original: We compared our method to state-of-the-art approaches on the DeepMind control suite, outperforming them on the majority of tasks and closing the gap with state-based training.
Revised: We compared our method to state-of-the-art approaches on the DeepMind control suite, outperforming them on the majority of tasks and closing the gap with state-based training.

Original: On the Atari 100k benchmark DrQ outperforms other SOTA methods in the median metric.
Revised: On the Atari 100k benchmark DrQ outperforms other SOTA methods in the median metric.

Original: To the best of our knowledge, this is the ﬁrst convincing demonstration of the utility of data augmentation on these standard benchmarks.
Revised: To the best of our knowledge, this is the ﬁrst convincing demonstration of the utility of data augmentation on these standard benchmarks.

Original: Furthermore, we demonstrate the method to be robust to the choice of hyper-parameters.
Revised: Furthermore, we demonstrate the method to be robust to the choice of hyper-parameters.

Original: A PPENDIX A E XTENDED B ACKGROUND
Revised: A PPENDIX A E XTENDED B ACKGROUND

Original: Reinforcement Learning from Images We formulate image-based control as an inﬁnite-horizon partially observable Markov decision process (POMDP) (Bellman, 1957; Kaelbling et al., 1998).
Revised: Reinforcement Learning from Images We formulate image-based control as an inﬁnite-horizon partially observable Markov decision process (POMDP) (Bellman, 1957; Kaelbling et al., 1998).

Original: An
Revised: An

Original: POMDP can be described as the tuple ( O , A , p, r, γ ) , where O is the high-dimensional observation space (image pixels), A is the action space, the transition dynamics p =
Revised: POMDP can be described as the tuple ( O , A , p, r, γ ) , where O is the high-dimensional observation space (image pixels), A is the action space, the transition dynamics p =

Original: Pr ( o (cid:48) t | o ≤ t , a t ) capture the probability distribution over the next observation o (cid:48)
Revised: Pr ( o (cid:48) t | o ≤ t , a t ) capture the probability distribution over the next observation o (cid:48)

Original: t given the history of previous observations o ≤ t and current action a t , r : O × A → R is the reward function that maps the current observation and action to a reward r t = r
Revised: t given the history of previous observations o ≤ t and current action a t , r : O × A → R is the reward function that maps the current observation and action to a reward r t = r

Original: ( o ≤ t , a t ) , and γ ∈ [0 , 1) is a discount factor.
Revised: ( o ≤ t , a t ) , and γ ∈ [0 , 1) is a discount factor.

Original: Per common practice (Mnih et al., 2013), throughout the paper the POMDP is converted into an MDP (Bellman, 1957) by stacking several consecutive image observations into a state s t = { o t , o t − 1 , o t − 2 , . . .
Revised: Per common practice (Mnih et al., 2013), throughout the paper the POMDP is converted into an MDP (Bellman, 1957) by stacking several consecutive image observations into a state s t = { o t , o t − 1 , o t − 2 , . . .

Original: For simplicity we redeﬁne the transition dynamics p =
Revised: For simplicity we redeﬁne the transition dynamics p =

Original: Pr ( s (cid:48)
Revised: Pr ( s (cid:48)

Original: t | s t , a t ) and the reward function r t = r ( s t , a t ) .
Revised: t | s t , a t ) and the reward function r t = r ( s t , a t ) .

Original: We then aim to ﬁnd a policy π ( a t | s t ) that maximizes the cumulative
Revised: We then aim to ﬁnd a policy π ( a t | s t ) that maximizes the cumulative

Original: discounted return E π [ (cid:80) ∞ t =1 γ t r t | a t ∼ π ( ·| s t ) , s (cid:48)
Revised: discounted return E π [ (cid:80) ∞ t =1 γ t r t | a t ∼ π ( ·| s t ) , s (cid:48)

Original: t ∼ p ( ·| s t , a t ) , s 1 ∼ p ( · )] .
Revised: t ∼ p ( ·| s t , a t ) , s 1 ∼ p ( · )] .

Original: Soft Actor-Critic
Revised: Soft Actor-Critic

Original: The Soft Actor-Critic (SAC) (Haarnoja et al., 2018) learns a state-action value function Q θ , a stochastic policy π θ and a temperature α to ﬁnd an optimal policy for an MDP ( S , A , p, r, γ ) by optimizing a γ -discounted maximum-entropy objective (Ziebart et al., 2008).
Revised: The Soft Actor-Critic (SAC) (Haarnoja et al., 2018) learns a state-action value function Q θ , a stochastic policy π θ and a temperature α to ﬁnd an optimal policy for an MDP ( S , A , p, r, γ ) by optimizing a γ -discounted maximum-entropy objective (Ziebart et al., 2008).

Original: θ is used generically to denote the parameters updated through training in each part of the model.
Revised: θ is used generically to denote the parameters updated through training in each part of the model.

Original: The actor policy π θ ( a t | s t ) is a parametric tanh -Gaussian that given s t samples a t = tanh( µ θ ( s t )
Revised: The actor policy π θ ( a t | s t ) is a parametric tanh -Gaussian that given s t samples a t = tanh( µ θ ( s t )

Original: + σ θ ( s t ) (cid:15) ) , where (cid:15) ∼ N (0 , 1) and µ θ and σ θ are parametric mean and standard deviation.
Revised: + σ θ ( s t ) (cid:15) ) , where (cid:15) ∼ N (0 , 1) and µ θ and σ θ are parametric mean and standard deviation.

Original: The policy evaluation step learns the critic Q θ ( s t , a t ) network by optimizing a single-step of the soft
Revised: The policy evaluation step learns the critic Q θ ( s t , a t ) network by optimizing a single-step of the soft

Original: Bellman residual
Revised: Bellman residual

Original: [Equation]
Revised: [Equation]

Original: where D is a replay buffer of transitions, θ (cid:48) is an exponential moving average of the weights as done in (Lillicrap et al., 2015).
Revised: where D is a replay buffer of transitions, θ (cid:48) is an exponential moving average of the weights as done in (Lillicrap et al., 2015).

Original: SAC uses clipped double-Q learning (van Hasselt et al., 2015; Fujimoto et al., 2018), which we omit from our notation for simplicity but employ in practice.
Revised: SAC uses clipped double-Q learning (van Hasselt et al., 2015; Fujimoto et al., 2018), which we omit from our notation for simplicity but employ in practice.

Original: The policy improvement step then ﬁts the actor policy π θ ( a t | s t ) network by optimizing the objective
Revised: The policy improvement step then ﬁts the actor policy π θ ( a t | s t ) network by optimizing the objective

Original: [Equation]
Revised: [Equation]

Original: Finally, the temperature α is learned with the loss
Revised: Finally, the temperature α is learned with the loss

Original: [Equation]
Revised: [Equation]

Original: where ¯ H ∈ R is the target entropy hyper-parameter that the policy tries to match, which in practice is usually set to ¯ H = −|A| .
Revised: where ¯ H ∈ R is the target entropy hyper-parameter that the policy tries to match, which in practice is usually set to ¯ H = −|A| .

Original: Deep Q-learning DQN (Mnih et al., 2013) also learns a convolutional neural net to approximate
Revised: Deep Q-learning DQN (Mnih et al., 2013) also learns a convolutional neural net to approximate

Original: Q-function over states and actions.
Revised: Q-function over states and actions.

Original: The main difference is that DQN operates on discrete actions spaces, thus the policy can be directly inferred from Q-values.
Revised: The main difference is that DQN operates on discrete actions spaces, thus the policy can be directly inferred from Q-values.

Original: The parameters of DQN are updated by optimizing the squared residual error
Revised: The parameters of DQN are updated by optimizing the squared residual error

Original: [Equation]
Revised: [Equation]

Original: In practice, the standard version of DQN is frequently combined with a set of tricks that improve performance and training stability, wildly known as Rainbow (van Hasselt et al., 2015).
Revised: In practice, the standard version of DQN is frequently combined with a set of tricks that improve performance and training stability, wildly known as Rainbow (van Hasselt et al., 2015).

Original: B T HE D EEP M IND C ONTROL S UITE E XPERIMENTS S ETUP
Revised: B T HE D EEP M IND C ONTROL S UITE E XPERIMENTS S ETUP

Original: Our PyTorch SAC (Haarnoja et al., 2018) implementation is based off of Yarats & Kostrikov (2020).
Revised: Our PyTorch SAC (Haarnoja et al., 2018) implementation is based off of Yarats & Kostrikov (2020).

Original: B. A CTOR AND C RITIC N ETWORKS
Revised: B. A CTOR AND C RITIC N ETWORKS

Original: We employ clipped double Q-learning (van Hasselt et al., 2015; Fujimoto et al., 2018) for the critic, where each Q -function is parametrized as a 3-layer MLP with ReLU activations after each layer except of the last.
Revised: We employ clipped double Q-learning (van Hasselt et al., 2015; Fujimoto et al., 2018) for the critic, where each Q -function is parametrized as a 3-layer MLP with ReLU activations after each layer except of the last.

Original: The actor is also a 3-layer MLP with ReLU s that outputs mean and covariance for the diagonal Gaussian that represents the policy.
Revised: The actor is also a 3-layer MLP with ReLU s that outputs mean and covariance for the diagonal Gaussian that represents the policy.

Original: The hidden dimension is set to 1024 for both the critic and actor.
Revised: The hidden dimension is set to 1024 for both the critic and actor.

Original: B. E NCODER N ETWORK
Revised: B. E NCODER N ETWORK

Original: We employ an encoder architecture from Yarats et al.
Revised: We employ an encoder architecture from Yarats et al.

Original: This encoder consists of four convolutional layers with 3 × 3 kernels and 32 channels.
Revised: This encoder consists of four convolutional layers with 3 × 3 kernels and 32 channels.

Original: The ReLU activation is applied after each conv layer.
Revised: The ReLU activation is applied after each conv layer.

Original: We use stride to 1 everywhere, except of the ﬁrst conv layer, which has stride 2 .
Revised: We use stride to 1 everywhere, except of the ﬁrst conv layer, which has stride 2 .

Original: The output of the convnet is feed into a single fully-connected layer normalized by LayerNorm (Ba et al., 2016).
Revised: The output of the convnet is feed into a single fully-connected layer normalized by LayerNorm (Ba et al., 2016).

Original: Finally, we apply tanh nonlinearity to the 50 dimensional output of the fully-connected layer.
Revised: Finally, we apply tanh nonlinearity to the 50 dimensional output of the fully-connected layer.

Original: We initialize the weight matrix of fully-connected and convolutional layers with the orthogonal initialization (Saxe et al., 2013) and set the bias to be zero.
Revised: We initialize the weight matrix of fully-connected and convolutional layers with the orthogonal initialization (Saxe et al., 2013) and set the bias to be zero.

Original: The actor and critic networks both have separate encoders, although we share the weights of the conv layers between them.
Revised: The actor and critic networks both have separate encoders, although we share the weights of the conv layers between them.

Original: Furthermore, only the critic optimizer is allowed to update these weights (e.g. we stop the gradients from the actor before they propagate to the shared conv layers)
Revised: Furthermore, only the critic optimizer is allowed to update these weights (e.g. we stop the gradients from the actor before they propagate to the shared conv layers)

Original: . B. T RAINING AND E VALUATION S ETUP
Revised: . B. T RAINING AND E VALUATION S ETUP

Original: Our agent ﬁrst collects 1000 seed observations using a random policy.
Revised: Our agent ﬁrst collects 1000 seed observations using a random policy.

Original: The further training observations are collected by sampling actions from the current policy.
Revised: The further training observations are collected by sampling actions from the current policy.

Original: We perform one training update every time we receive a new observation.
Revised: We perform one training update every time we receive a new observation.

Original: In cases where we use action repeat, the number of training observations is only a fraction of the environment steps (e.g. a 1000 steps episode at action repeatwill only results into 250 training observations).
Revised: In cases where we use action repeat, the number of training observations is only a fraction of the environment steps (e.g. a 1000 steps episode at action repeatwill only results into 250 training observations).

Original: We evaluate our agent every 10000 true environment steps by computing the average episode return over 10 evaluation episodes.
Revised: We evaluate our agent every 10000 true environment steps by computing the average episode return over 10 evaluation episodes.

Original: During evaluation we take the mean policy action instead of sampling.
Revised: During evaluation we take the mean policy action instead of sampling.

Original: B. P LA N ET AND D REAMER B ENCHMARKS
Revised: B. P LA N ET AND D REAMER B ENCHMARKS

Original: We consider two evaluation setups that were introduced in PlaNet (Hafner et al., 2018) and
Revised: We consider two evaluation setups that were introduced in PlaNet (Hafner et al., 2018) and

Original: Dreamer (Hafner et al., 2019), both using tasks from the DeepMind control suite (Tassa et al., 2018).
Revised: Dreamer (Hafner et al., 2019), both using tasks from the DeepMind control suite (Tassa et al., 2018).

Original: The PlaNet benchmark consists of six tasks of various traits.
Revised: The PlaNet benchmark consists of six tasks of various traits.

Original: Importantly, the benchmark proposed to use a different action repeat hyper-parameter for each task, which we summarize in Table 2.
Revised: Importantly, the benchmark proposed to use a different action repeat hyper-parameter for each task, which we summarize in Table 2.

Original: The Dreamer benchmark considers an extended set of tasks, which makes it more difﬁcult that
Revised: The Dreamer benchmark considers an extended set of tasks, which makes it more difﬁcult that

Original: the
Revised: the

Original: PlaNet setup.
Revised: PlaNet setup.

Original: Additionally, this benchmark requires to use the same set hyper-parameters for each task, including action repeat (set to 2 ), which further increases the difﬁculty.
Revised: Additionally, this benchmark requires to use the same set hyper-parameters for each task, including action repeat (set to 2 ), which further increases the difﬁculty.

Original: [Table]
Revised: [Table]

Original: B. P IXELS P REPROCESSING
Revised: B. P IXELS P REPROCESSING

Original: We construct an observational input as an 3 -stack of consecutive frames (Mnih et al., 2013), where each frame is a RGB rendering of size 84 × 84 from the 0 th camera.
Revised: We construct an observational input as an 3 -stack of consecutive frames (Mnih et al., 2013), where each frame is a RGB rendering of size 84 × 84 from the 0 th camera.

Original: We then divide each pixel by 255 to scale it down to [0 , 1] range.
Revised: We then divide each pixel by 255 to scale it down to [0 , 1] range.

Original: B. O THER H YPER P ARAMETERS
Revised: B. O THER H YPER P ARAMETERS

Original: Due to computational constraints for all the continuous control ablation experiments in the main paper and appendix we use a minibatch size of
Revised: Due to computational constraints for all the continuous control ablation experiments in the main paper and appendix we use a minibatch size of

Original: [Table]
Revised: [Table]

Original: while for the main results we use minibatch of size 512 .
Revised: while for the main results we use minibatch of size 512 .

Original: In Table 3 we provide a comprehensive overview of all the other hyper-parameters
Revised: In Table 3 we provide a comprehensive overview of all the other hyper-parameters

Original: [Table] [Table]
Revised: [Table] [Table]

Original: C T HE A TARI 100 K E XPERIMENTS S ETUP
Revised: C T HE A TARI 100 K E XPERIMENTS S ETUP

Original: For ease of reproducibility in Table 4 we report the hyper-parameter settings used in the Atari 100k experiments.
Revised: For ease of reproducibility in Table 4 we report the hyper-parameter settings used in the Atari 100k experiments.

Original: We largely reuse the hyper-parameters from OTRainbow (Kielak, 2020), but adapt them for DQN (Mnih et al., 2013).
Revised: We largely reuse the hyper-parameters from OTRainbow (Kielak, 2020) [Table] but adapt them for DQN (Mnih et al., 2013).

Original: Per common practise, we average performance of our agent overrandom seeds.
Revised: Per common practise, we average performance of our agent overrandom seeds.

Original: The evaluation is done for 125k environment steps at the end of training for 100k environment steps.
Revised: The evaluation is done for 125k environment steps at the end of training for 100k environment steps.

Original: [Table]
Revised: [Table]

Original: D F ULL A TARI 100 K R ESULTS
Revised: D F ULL A TARI 100 K R ESULTS

Original: Besides reporting in Figure 5 median human-normalized episode returns over the 26 Atari games used in (Kaiser et al., 2019), we also provide the mean episode return for each individual game in Table
Revised: Besides reporting in Figure 5 median human-normalized episode returns over the 26 Atari games used in (Kaiser et al., 2019), we also provide the mean episode return for each individual game in Table

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: 20072.2 DemonAttack 163.
Revised: 20072.2 DemonAttack 163.

Original: 1086.0 Freeway 0.
Revised: 1086.0 Freeway 0.

Original: 20.0 Frostbite 60.
Revised: 20.0 Frostbite 60.

Original: 889.9 Gopher 431.
Revised: 889.9 Gopher 431.

Original: 678.0 Hero 487.
Revised: 678.0 Hero 487.

Original: 4083.7 Jamesbond 47.
Revised: 4083.7 Jamesbond 47.

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: 330.3 Kangaroo 0.
Revised: 330.3 Kangaroo 0.

Original: 1282.6 Krull 1468.
Revised: 1282.6 Krull 1468.

Original: 4163.0 KungFuMaster 0.
Revised: 4163.0 KungFuMaster 0.

Original: 7649.0 MsPacman 67.
Revised: 7649.0 MsPacman 67.

Original: 1015.9 Pong -20.
Revised: 1015.9 Pong -20.

Original: -17.1 PrivateEye 0.
Revised: -17.1 PrivateEye 0.

Original: Qbert 123.
Revised: Qbert 123.

Original: 769.1 RoadRunner 1588.
Revised: 769.1 RoadRunner 1588.

Original: [Table]
Revised: [Table]

Original: E I MAGE A UGMENTATIONS A BLATION
Revised: E I MAGE A UGMENTATIONS A BLATION

Original: Following (Chen et al., 2020), we evaluate popular image augmentation techniques, namely random shifts, cutouts, vertical and horizontal ﬂips, random rotations and imagewise intensity jittering.
Revised: Following (Chen et al., 2020), we evaluate popular image augmentation techniques, namely random shifts, cutouts, vertical and horizontal ﬂips, random rotations and imagewise intensity jittering.

Original: Below, we provide a comprehensive overview of each augmentation.
Revised: Below, we provide a comprehensive overview of each augmentation.

Original: Furthermore, we examine effectiveness of these techniques in Figure 6.
Revised: Furthermore, we examine effectiveness of these techniques in Figure 6.

Original: Random Shift
Revised: Random Shift

Original: We bring our attention to random shifts that are commonly used to regularize neural networks trained on small images (Becker & Hinton, 1992; Simard et al., 2003; LeCun et al., 1989;
Revised: We bring our attention to random shifts that are commonly used to regularize neural networks trained on small images (Becker & Hinton, 1992; Simard et al., 2003; LeCun et al., 1989;

Original: Ciresan et al., 2011; Ciregan et al., 2012).
Revised: Ciresan et al., 2011; Ciregan et al., 2012).

Original: In our implementation of this method images of size 84 × 84 are padded each side by 4 pixels (by repeating boundary pixels) and then randomly cropped back to the original 84 × 84 size.
Revised: In our implementation of this method images of size 84 × 84 are padded each side by 4 pixels (by repeating boundary pixels) and then randomly cropped back to the original 84 × 84 size.

Original: Cutout Cutouts introduced in DeVries & Taylor (2017) represent a generalization of Dropout (Hinton et al., 2012).
Revised: Cutout Cutouts introduced in DeVries & Taylor (2017) represent a generalization of Dropout (Hinton et al., 2012).

Original: Instead of masking individual pixels cutouts mask square regions.
Revised: Instead of masking individual pixels cutouts mask square regions.

Original: Since image pixels can be highly correlated, this technique is proven to improve training of neural networks.
Revised: Since image pixels can be highly correlated, this technique is proven to improve training of neural networks.

Original: Horizontal/Vertical Flip
Revised: Horizontal/Vertical Flip

Original: This technique simply ﬂips an image either horizontally or vertically with probability 0 . 1 .
Revised: This technique simply ﬂips an image either horizontally or vertically with probability 0 . 1 .

Original: Rotate Here, an image is rotated by r degrees, where r is uniformly sampled from [ − 5 , − 5] .
Revised: Rotate Here, an image is rotated by r degrees, where r is uniformly sampled from [ − 5 , − 5] .

Original: Intensity Each N × C × 84 × 84 image tensor is multiplied by a single scalar s , which is computed as s = µ + σ · clip( r, − 2 , 2) , where r ∼ N (0 , 1) .
Revised: Intensity Each N × C × 84 × 84 image tensor is multiplied by a single scalar s , which is computed as s = µ + σ · clip( r, − 2 , 2) , where r ∼ N (0 , 1) .

Original: For our experiments we use µ = 1 . 0 and σ = 0 . 
Revised: For our experiments we use µ = 1 . 0 and σ = 0 . [Figure] i

Original: [Figure]
Revised: [Figure]

Original: R e t u r 
Revised: R e t u r n

Original: [Figure]
Revised: [Figure]

Original:  i s o d e
Revised: p i s o d e

Original: R e t u r n
Revised: R e t u r n

Original: [Figure]
Revised: [Figure]

Original: import torch import torch.nn as nn import kornia.augmentation as aug random_shift = nn.Sequential(nn.
Revised: import torch import torch.nn as nn import kornia.augmentation as aug random_shift = nn.Sequential(nn.

Original: ReplicationPad2d(4),aug.
Revised: ReplicationPad2d(4),aug.

Original: RandomCrop((84, 84))) cutout =
Revised: RandomCrop((84, 84))) cutout =

Original: aug.RandomErasing(p=0.5) h_flip =
Revised: aug.RandomErasing(p=0.5) h_flip =

Original: aug.RandomHorizontalFlip(p=0.1) v_flip
Revised: aug.RandomHorizontalFlip(p=0.1) v_flip

Original: = aug.RandomVerticalFlip(p=0.1)
Revised: = aug.RandomVerticalFlip(p=0.1)

Original: rotate = aug.
Revised: rotate = aug.

Original: RandomRotation(degrees=5.0) intensity = Intensity(scale=0.1) class Intensity(nn.
Revised: RandomRotation(degrees=5.0) intensity = Intensity(scale=0.1) class Intensity(nn.

Original: Module): def __init__(self, scale): super().__init__() self.scale=
Revised: Module): def __init__(self, scale): super().__init__() self.scale=

Original: scale def forward(self, x): r = torch.randn((x.size(0), 1, 1, 1), device=x.device) noise = 1.0 + (self.scale * r.clamp(-2.0, 2.0)) return x * noise
Revised: scale def forward(self, x): r = torch.randn((x.size(0), 1, 1, 1), device=x.device) noise = 1.0 + (self.scale * r.clamp(-2.0, 2.0)) return x * noise

Original: F K AND M H YPER - PARAMETERS A BLATION
Revised: F K AND M H YPER - PARAMETERS A BLATION

Original: We further ablate the K,M hyper-parameters from Algorithm 1 to understand their effect on performance.
Revised: We further ablate the K,M hyper-parameters from Algorithm 1 to understand their effect on performance.

Original: In Figure 7 we observe that increase values of K,M improves the agent’s performance.
Revised: In Figure 7 we observe that increase values of K,M improves the agent’s performance.

Original: We choose to use the [K=2,M=2] parametrization as it strikes a good balance between performance and computational demands
Revised: We choose to use the [K=2,M=2] parametrization as it strikes a good balance between performance and computational demands

Original: [Figure]
Revised: [Figure]

Original: G R OBUSTNESS I NVESTIGATION
Revised: G R OBUSTNESS I NVESTIGATION

Original: To demonstrate the robustness of our approach (Henderson et al., 2018), we perform a comprehensive study on the effect different hyper-parameter choices have on performance.
Revised: To demonstrate the robustness of our approach (Henderson et al., 2018), we perform a comprehensive study on the effect different hyper-parameter choices have on performance.

Original: A review of prior work (Hafner et al., 2018; 2019; Lee et al., 2019; Srinivas et al., 2020) shows consistent values for discount γ = 0 .
Revised: A review of prior work (Hafner et al., 2018; 2019; Lee et al., 2019; Srinivas et al., 2020) shows consistent values for discount γ = 0 .

Original: 99 and target update rate τ = 0 .
Revised: 99 and target update rate τ = 0 .

Original: 01 parameters, but variability on network architectures, mini-batch sizes, learning rates.
Revised: 01 parameters, but variability on network architectures, mini-batch sizes, learning rates.

Original: Since our method is based on SAC (Haarnoja et al., 2018), we also check whether the initial value of the temperature is important, as it plays a crucial role in the initial phase of exploration.
Revised: Since our method is based on SAC (Haarnoja et al., 2018), we also check whether the initial value of the temperature is important, as it plays a crucial role in the initial phase of exploration.

Original: We omit search over network architectures since Figure 1b shows our method to be robust to the exact choice.
Revised: We omit search over network architectures since Figure 1b shows our method to be robust to the exact choice.

Original: We thus focus on three hyper-parameters: mini-batch size, learning rate, and initial temperature.
Revised: We thus focus on three hyper-parameters: mini-batch size, learning rate, and initial temperature.

Original: Due to computational demands, experiments are restricted to a subset of tasks from Tassa et al.
Revised: Due to computational demands, experiments are restricted to a subset of tasks from Tassa et al.

Original: Walker Walk, Cartpole Swingup, and Finger Spin.
Revised: Walker Walk, Cartpole Swingup, and Finger Spin.

Original: These were selected to be diverse, requiring different behaviors including locomotion and goal reaching.
Revised: These were selected to be diverse, requiring different behaviors including locomotion and goal reaching.

Original: A grid search is performed over minibatch sizes { 128 , 256 , 512 } , learning rates { 0 . 0001 , 0 . 0005 , 0 .
Revised: A grid search is performed over minibatch sizes { 128 , 256 , 512 } , learning rates { 0 . 0001 , 0 . 0005 , 0 .

Original: 005 } , and initial temperatures { 0 . 005 , 0 . 01 , 0 .
Revised: 005 } , and initial temperatures { 0 . 005 , 0 . 01 , 0 .

Original: We follow the experimental setup from Appendix B, except that only 3 seeds are used due to the computation limitations, but since variance is low the results are representative.
Revised: We follow the experimental setup from Appendix B, except that only 3 seeds are used due to the computation limitations, but since variance is low the results are representative.

Original: [Figure]
Revised: [Figure]

Original: Figure 8 shows performance curves for each conﬁguration as well as a heat map over the mean performance of the ﬁnal evaluation episodes, similar to Mnih et al.
Revised: Figure 8 shows performance curves for each conﬁguration as well as a heat map over the mean performance of the ﬁnal evaluation episodes, similar to Mnih et al.

Original: Our method demonstrates good stability and is largely invariant to the studied hyper-parameters.
Revised: Our method demonstrates good stability and is largely invariant to the studied hyper-parameters.

Original: We emphasize that for simplicity the experiments in Section 5 use the default learning rate of Adam (Kingma & Ba, 2014) (0.001), even though it is not always optimal.
Revised: We emphasize that for simplicity the experiments in Section 5 use the default learning rate of Adam (Kingma & Ba, 2014) (0.001), even though it is not always optimal.

Original: H I MPROVED D ATA -E FFICIENT R EINFORCEMENT L EARNING FROM P IXELS
Revised: H I MPROVED D ATA -E FFICIENT R EINFORCEMENT L EARNING FROM P IXELS

Original: Our method allows to generate many various transformations from a training observation due to the data augmentation strategy.
Revised: Our method allows to generate many various transformations from a training observation due to the data augmentation strategy.

Original: Thus, we further investigate whether performing more training updates per an environment step can lead to even better sample-efﬁciency.
Revised: Thus, we further investigate whether performing more training updates per an environment step can lead to even better sample-efﬁciency.

Original: Following van Hasselt et al.
Revised: Following van Hasselt et al.

Original: (2019b)
Revised: (2019b)

Original: we compare a single update with a mini-batch of 512 transitions with 4 updates withdifferent mini-batches of size 128 samples each.
Revised: we compare a single update with a mini-batch of 512 transitions with 4 updates withdifferent mini-batches of size 128 samples each.

Original: Performing more updates per an environment step leads to even worse over-ﬁtting on some tasks without data augmentation (see Figure 9a), while our method DrQ , that takes advantage of data augmentation, demonstrates improved sample-efﬁciency (see Figure 9b).
Revised: Performing more updates per an environment step leads to even worse over-ﬁtting on some tasks without data augmentation (see Figure 9a), while our method DrQ , that takes advantage of data augmentation, demonstrates improved sample-efﬁciency (see Figure 9b).

Original: [Figure]
Revised: [Figure]

Original: i [Figure] d 
Revised: i s o d e

Original: [Figure]

Original: R [Figure] t u r 
Revised: R e t u r n

Original: [Figure]
Revised: [Figure]