Original: O N THE D ATA -E FFICIENCY WITH C ONTRASTIVE I MAGE T RANSFORMATION IN R EINFORCEMENT L EARNING
Revised: O N THE D ATA -E FFICIENCY WITH C ONTRASTIVE I MAGE T RANSFORMATION IN R EINFORCEMENT L EARNING

Original: A BSTRACT
Revised: A BSTRACT

Original: Data-efﬁciency has always been an essential issue in pixel-based reinforcement learning (RL).
Revised: Data-efficiency has always been an essential issue in pixel-based reinforcement learning (RL).

Original: As the agent not only learns decision-making but also meaningful representations from images.
Revised: As the agent not only learns decision-making but also meaningful representations from images.

Original: The line of reinforcement learning with data augmentation shows signiﬁcant improvements in sample-efﬁciency.
Revised: The line of reinforcement learning with data augmentation shows significant improvements in sample-efficiency.

Original: However, it is challenging to guarantee the optimality invariant transformation, that is, the augmented data are readily recognized as a completely different state by the agent.
Revised: However, it is challenging to guarantee the optimality invariant transformation, that is, the augmented data are readily recognized as a completely different state by the agent.

Original: In the end, we propose a contrastive invariant transformation (CoIT), a simple yet promising learnable data augmentation combined with standard model-free algorithms to improve sample-efﬁciency.
Revised: In the end, we propose a contrastive invariant transformation (CoIT), a simple yet promising learnable data augmentation combined with standard model-free algorithms to improve sample-efficiency.

Original: Concretely, the differentiable CoIT leverages original samples with augmented samples and hastens the state encoder for a contrastive invariant embedding.
Revised: Concretely, the differentiable CoIT leverages original samples with augmented samples and hastens the state encoder for a contrastive invariant embedding.

Original: We evaluate our approach on DeepMind Control
Revised: We evaluate our approach on DeepMind Control

Original: Suite and Atari100K. Empirical results verify advances using CoIT, enabling it to outperform the new state-of-the-art on various tasks.
Revised: Suite and Atari100K. Empirical results verify advances using CoIT, enabling it to outperform the new state-of-the-art on various tasks.

Original: Source code is available at https://github.com/mooricAnna/CoIT .
Revised: Source code is available at https://github.com/mooricAnna/CoIT .

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Improving data-efﬁciency to accomplish sequential decisions has always been a crucial problem in pixel-based reinforcement learning.
Revised: Improving data-efficiency to accomplish sequential decisions has always been a crucial problem in pixel-based reinforcement learning.

Original: As the agent has to learn an optimal policy with a meaningful information abstraction from observations parallel.
Revised: As the agent has to learn an optimal policy with a meaningful information abstraction from observations parallel.

Original: Unlike supervised representation learning with strong supervised high-dimensional signals, the training process in RL is fragile.
Revised: Unlike supervised representation learning with strong supervised high-dimensional signals, the training process in RL is fragile.

Original: It could be harmful to the training process and cause performance degradation consequently using inappropriate manners.
Revised: It could be harmful to the training process and cause performance degradation consequently using inappropriate manners.

Original: Hence, it is an urgent request to seek subtle representation learning methods for visual RL.
Revised: Hence, it is an urgent request to seek subtle representation learning methods for visual RL.

Original: Previous works have been proposed in the literature to demonstrate that introducing auxiliary loss functions such as pixel reconstruction (Yarats et al., 2019) and contrastive learning (Laskin et al., 2020b) alleviates this issue.
Revised: Previous works have been proposed in the literature to demonstrate that introducing auxiliary loss functions such as pixel reconstruction (Yarats et al., 2019) and contrastive learning (Laskin et al., 2020b) alleviates this issue.

Original: In particular, data augmentations have already proven beneﬁcial to dataefﬁciency.
Revised: In particular, data augmentations have already proven beneficial to dataefficiency.

Original: RAD (Laskin et al., 2020a) performs an extension of experiments and widely analyzes the impact of various techniques in data augmentation.
Revised: RAD (Laskin et al., 2020a) performs an extension of experiments and widely analyzes the impact of various techniques in data augmentation.

Original: DrQ (Yarats et al., 2020) and DrQ-v2 (Yarats et al., 2021) make use of appropriate image augmentation with great success.
Revised: DrQ (Yarats et al., 2020) and DrQ-v2 (Yarats et al., 2021) make use of appropriate image augmentation with great success.

Original: Also, previous works have carried out the potential of data augmentation in terms of generalization (Hansen et al., 2021;
Revised: Also, previous works have carried out the potential of data augmentation in terms of generalization (Hansen et al., 2021;

Original: Raileanu et al., 2020; Zhang & Guo, 2021; Hansen & Wang, 2021;
Revised: Raileanu et al., 2020; Zhang & Guo, 2021; Hansen & Wang, 2021;

Original: Fan et al., 2021).
Revised: Fan et al., 2021).

Original: Despite the mentioned efforts, it is pretty hard to guarantee that the augmented representations are sufﬁciently diverse yet semantically consistent.
Revised: Despite the mentioned efforts, it is pretty hard to guarantee that the augmented representations are sufficiently diverse yet semantically consistent.

Original: To this end, we explore the underlying condition for representation learning in RL.
Revised: To this end, we explore the underlying condition for representation learning in RL.

Original: It is rational to hypothesize that there is an optimal transformation enabling an encoder to abstract informative latent space.
Revised: It is rational to hypothesize that there is an optimal transformation enabling an encoder to abstract informative latent space.

Original: This line of works belongs to the regime of state abstraction (Du et al., 2019;
Revised: This line of works belongs to the regime of state abstraction (Du et al., 2019;

Original: Zhang et al., 2020b; Tomar et al., 2021;
Revised: Zhang et al., 2020b; Tomar et al., 2021;

Original: Wang et al., 2022), which derives from grouping similar world-states for descriptions of the environment (Dietterich, 2000;
Revised: Wang et al., 2022), which derives from grouping similar world-states for descriptions of the environment (Dietterich, 2000;

Original: Andre & Russell, 2002; Castro & Precup, 2010).
Revised: Andre & Russell, 2002; Castro & Precup, 2010).

Original: Inspired by spatial transformer networks (STN)
Revised: Inspired by spatial transformer networks (STN)

Original: (Jaderberg et al., 2015), a data augmentation model in the vision domain, we consider that merging the parameterized transformation with visual RL could be beneﬁcial.
Revised: (Jaderberg et al., 2015), a data augmentation model in the vision domain, we consider that merging the parameterized transformation with visual RL could be beneficial.

Original: The designed transformation not only discovers the optimality of state abstraction but also produces diverse virtual samples for the agent.
Revised: The designed transformation not only discovers the optimality of state abstraction but also produces diverse virtual samples for the agent.

Original: To do so, we enforce a learnable data augmentation that updates its parameters along with the
Revised: To do so, we enforce a learnable data augmentation that updates its parameters along with the

Original: RL objective.
Revised: RL objective.

Original: To understand parameterized augmentation and its relation to representation learning in RL, we focus on fundamental data manipulation by generating augmented data from a learnable Gaussian distribution.
Revised: To understand parameterized augmentation and its relation to representation learning in RL, we focus on fundamental data manipulation by generating augmented data from a learnable Gaussian distribution.

Original: To be clear, we present the image transformation to control the margin of the augmentation under an RL training-friendly data distribution.
Revised: To be clear, we present the image transformation to control the margin of the augmentation under an RL training-friendly data distribution.

Original: Since changed data distributions meanwhile being controlled by learning algorithms would be helpful in high-dimensional cases (Balestriero et al., 2021).
Revised: Since changed data distributions meanwhile being controlled by learning algorithms would be helpful in high-dimensional cases (Balestriero et al., 2021).

Original: Here we raise our idea:  we parameterize the data augmentation by sampling from a dynamic distribution to obtain a friendly along
Revised: Here we raise our idea: Can we parameterize the data augmentation by sampling from a dynamic distribution to obtain a training along

Original:  In light of this challenge, we present a contrastive invariant transformation (CoIT), a novel contrastive learning to ameliorate the data-efﬁciency for visual RL.
Revised: RL In light of this challenge, we present a contrastive invariant transformation (CoIT), a novel contrastive learning to ameliorate the data-efficiency for visual RL.

Original: CoIT integrates a learnable transformation for model-free methods with minimal modiﬁcation to the architecture and training pipeline.
Revised: CoIT integrates a learnable transformation for model-free methods with minimal modification to the architecture and training pipeline.

Original: Specifically, we parameterize the mean and variance of a Gaussian distribution for transforming data and update the parameters together with RL by using constraints to urge faster algorithm convergence empirically.
Revised: Specifically, we parameterize the mean and variance of a Gaussian distribution for transforming data and update the parameters together with RL by using constraints to urge faster algorithm convergence empirically.

Original: As the learning goes on, the agent approximates the TRANSFORM distribution that is optimal for the task at hand to solve the task.
Revised: As the learning goes on, the agent approximates the TRANSFORM distribution that is optimal for the task at hand to solve the task.

Original: In addition, we evaluate
Revised: In addition, we evaluate

Original: CoIT on DeepMind Control Suite and Atari100K, and experimental results demonstrate that the learnable transformation outperforms the current SOTA methods.
Revised: CoIT on DeepMind Control Suite and Atari100K, and experimental results demonstrate that the learnable transformation outperforms the current SOTA methods.

Original: Besides, our method does not claim any custom architecture choices and is essential for reproducing end-to-end training.
Revised: Besides, our method does not claim any custom architecture choices and is essential for reproducing end-to-end training.

Original: Based on these results, we demonstrate that a learnable transformation improves dataefﬁciency effectively for visual RL.
Revised: Based on these results, we demonstrate that a learnable transformation improves dataefficiency effectively for visual RL.

Original: Key Contributions: (i) We present CoIT, a simple yet effective framework with a learnable image transformation that integrates invariant representations with model-free RL to improve data-efﬁciency.
Revised: Key Contributions: (i) We present CoIT, a simple yet effective framework with a learnable image transformation that integrates invariant representations with model-free RL to improve data-efficiency.

Original: (ii)
Revised: (ii)

Original: We propose a theoretical analysis of how our method can approximate a stationary distribution over the transformed data by the optimal invariant metric, thus learning better representations.
Revised: We propose a theoretical analysis of how our method can approximate a stationary distribution over the transformed data by the optimal invariant metric, thus learning better representations.

Original: (iii)
Revised: (iii)

Original: We evaluate CoIT on popular benchmarks and show that our method outperforms previous state-of-the-art methods on data-efﬁciency and stability.
Revised: We evaluate CoIT on popular benchmarks and show that our method outperforms previous state-of-the-art methods on data-efficiency and stability.

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

Original: Several concurrent methods have been proposed for improving data-efﬁciency whose common ingredients containing data augmentation and self-supervised learning are listed.
Revised: Several concurrent methods have been proposed for improving data-efficiency whose common ingredients containing data augmentation and self-supervised learning are listed.

Original: Data augmentation in RL.
Revised: Data augmentation in RL.

Original: Like the success of data augmentation in computer vision (Zhong et al., 2020; DeVries & Taylor, 2017; Yun et al., 2019; Zhang et al., 2017), these methods have played a key role in improving the data-efﬁciency of visual RL problems Mnih et al.
Revised: Like the success of data augmentation in computer vision (Zhong et al., 2020; DeVries & Taylor, 2017; Yun et al., 2019; Zhang et al., 2017), these methods have played a key role in improving the data-efficiency of visual RL problems Mnih et al.

Original: ; Yarats et al.
Revised: ; Yarats et al.

Original: ; Hafner et al.
Revised: ; Hafner et al.

Original: ; Lee et al.
Revised: ; Lee et al.

Original: RAD (Laskin et al., 2020a) conducted mounts of experiments and ﬁnds out that different data augmentations lead to entirely different results.
Revised: RAD (Laskin et al., 2020a) conducted mounts of experiments and finds out that different data augmentations lead to entirely different results.

Original: It provides a broader perspective for the follow-up study of data augmentation in RL.
Revised: It provides a broader perspective for the follow-up study of data augmentation in RL.

Original: DrQ (Yarats et al., 2020) proposed an effective augmentation method called random shift and introduced a regularization term for Q-learning.
Revised: DrQ (Yarats et al., 2020) proposed an effective augmentation method called random shift and introduced a regularization term for Q-learning.

Original: Based on DrQ, the DrQ-v2 (Yarats et al., 2021) conducted minimal changes and demonstrated that merely a simple augmentation method could match the state-of-the-art model-based algorithm on data-efﬁciency and performance.
Revised: Based on DrQ, the DrQ-v2 (Yarats et al., 2021) conducted minimal changes and demonstrated that merely a simple augmentation method could match the state-of-the-art model-based algorithm on data-efficiency and performance.

Original: Self-supervised learning in RL.
Revised: Self-supervised learning in RL.

Original: Motivated by the breakthrough in self-supervised learning (Chen et al., 2020;
Revised: Motivated by the breakthrough in self-supervised learning (Chen et al., 2020;

Original: He et al., 2020; Caron et al., 2020; Grill et al., 2020), it is natural to combine these methods with visual RL to learn rich representations.
Revised: He et al., 2020; Caron et al., 2020; Grill et al., 2020), it is natural to combine these methods with visual RL to learn rich representations.

Original: CURL (Laskin et al., 2020b) introduces a framework similar to SimCLR (che) into visual RL.
Revised: CURL (Laskin et al., 2020b) introduces a framework similar to SimCLR (che) into visual RL.

Original: CoBERL
Revised: CoBERL

Original: (Banino et al., 2021) also tried to minimize the consistency between positive samples by semantic-preserving data augmentation.
Revised: (Banino et al., 2021) also tried to minimize the consistency between positive samples by semantic-preserving data augmentation.

Original: Besides, STDIM (Mazoure et al., 2020) and PI-SAC (Lee et al., 2020) maximize the temporal mutual information (MI) between the nearby states.
Revised: Besides, STDIM (Mazoure et al., 2020) and PI-SAC (Lee et al., 2020) maximize the temporal mutual information (MI) between the nearby states.

Original: SPR (Schwarzer et al., 2020) and PlayVirtual (Yu et al., 2021) follow their idea, but they utilize the dynamics model to predict nearby states in latent space.
Revised: SPR (Schwarzer et al., 2020) and PlayVirtual (Yu et al., 2021) follow their idea, but they utilize the dynamics model to predict nearby states in latent space.

Original: DBC (Zhang et al., 2020b) and PSM (Agarwal et al., 2021) focus on learning task-relevant information.
Revised: DBC (Zhang et al., 2020b) and PSM (Agarwal et al., 2021) focus on learning task-relevant information.

Original: They utilize signals in the environment to achieve an invariant representation learning and thereby generalize the agent to unseen environments.
Revised: They utilize signals in the environment to achieve an invariant representation learning and thereby generalize the agent to unseen environments.

Original:  P RELIMINARIES 3. R EINFORCEMENT L EARNING FROM O BSERVATIONS
Revised:  P RELIMINARIES 3. R EINFORCEMENT L EARNING FROM O BSERVATIONS

Original: Visual RL control is formulated as an inﬁnite-horizon Markov Decision Process (MDP) (Bellman, 1957), as the observations can not fully describe the underlying state.
Revised: Visual RL control is formulated as an infinite-horizon Markov Decision Process (MDP) (Bellman, 1957), as the observations can not fully describe the underlying state.

Original: To address this problem, we stack multiple consecutive frames together to represent the current underlying state s (Mnih et al., 2013).
Revised: To address this problem, we stack multiple consecutive frames together to represent the current underlying state s (Mnih et al., 2013).

Original: In this mind, the MDP M is a 5-tuple h O , S , A , r,  i .
Revised: In this mind, the MDP M is a 5-tuple ⟨O , S , A , r, γ ⟩ .

Original: Here, the observation space O generally consists of multiple-stack frames.
Revised: Here, the observation space O generally consists of multiple-stack frames.

Original: The state space S is either observable or unobservable (Silver et al., 2017; Zhang et al., 2020a).
Revised: The state space S is either observable or unobservable (Silver et al., 2017; Zhang et al., 2020a).

Original: The agent uses observations O to sample actions from the action space
Revised: The agent uses observations O to sample actions from the action space

Original: A .
Revised: A .

Original: Every time the agent interacts with the environment, it obtains a reward r .
Revised: Every time the agent interacts with the environment, it obtains a reward r .

Original: The end goal is to train an agent to maximize the cumulative reward R .
Revised: The end goal is to train an agent to maximize the cumulative reward R .

Original: The policy evaluation used as estimating the performance of the policy ⇡  is normally deﬁned by rewards in inﬁnite-horizon tasks,
Revised: The policy evaluation used as estimating the performance of the policy π ϕ is normally defined by rewards in infinite-horizon tasks,

Original: [Equation]
Revised: [Equation]

Original: 1) where  2 [0 , 1) is the discount factor and r t denotes the reward at time t .
Revised: where γ ∈ [0 , 1) is the discount factor and r t denotes the reward at time t .

Original: 3. Q L EARNING
Revised: 3. Q L EARNING

Original: The state-action value function Q ✓ is trained by minimizing the Bellman error to estimate the cumulative reward at the current state:
Revised: The state-action value function Q θ is trained by minimizing the Bellman error to estimate the cumulative reward at the current state:

Original: [Equation]
Revised: [Equation]

Original: where e is a transition from the replay buffer D .
Revised: where e is a transition from the replay buffer D .

Original: And ¯ ✓ denotes an exponential moving average of ✓ .
Revised: And ¯ θ denotes an exponential moving average of θ .

Original: For the continuous control tasks, we utilize an actor-critic algorithm called Deep Deterministic Policy
Revised: For the continuous control tasks, we utilize an actor-critic algorithm called Deep Deterministic Policy

Original: Gradient (DDPG) (Silver et al., 2014; Lillicrap et al., 2015) which consists of the aforementioned state-action value function Q ✓ and a deterministic policy ⇡  .
Revised: Gradient (DDPG) (Silver et al., 2014; Lillicrap et al., 2015) which consists of the aforementioned state-action value function Q θ and a deterministic policy π ϕ .

Original: The policy ⇡  aims at maximizing
Revised: The policy π ϕ aims at maximizing

Original: J  ( D ) =
Revised: J ϕ ( D ) =

Original: E D [ Q ✓ ( s t , ⇡  ( s t )] .
Revised: E D [ Q θ ( s t , π ϕ ( s t )] .

Original: Various effective improvements have also been lead to DDPG.
Revised: Various effective improvements have also been lead to DDPG.

Original: The Qlearning process incorporates n -step returns (Watkins, 1989; Peng & Williams, 1994).
Revised: The Qlearning process incorporates n -step returns (Watkins, 1989; Peng & Williams, 1994).

Original: The scheduled exploration noise is produced by a linear decay ˜  ( t ) for the variance ˜  2 which provides different levels of exploration at different training steps: ˜  ( t ) =
Revised: The scheduled exploration noise is produced by a linear decay ˜ σ ( t ) for the variance ˜ σ 2 which provides different levels of exploration at different training steps: ˜ σ ( t ) =

Original: ˜  init + (1  min( t/T, 1))(˜  ﬁnal  ˜  init ) .
Revised: ˜ σ init + (1 − min( t/T, 1))(˜ σ final − ˜ σ init ) .

Original: The initial and ﬁnal value for standard deviation are deﬁned by ˜  init and ˜  ﬁnal , and the decay horizon T is related to the total training steps of the environment.
Revised: The initial and final value for standard deviation are defined by ˜ σ init and ˜ σ final , and the decay horizon T is related to the total training steps of the environment.

Original: For the discrete control, we use the data-efﬁcient Rainbow DQN (Van Hasselt et al., 2019) which applied multiple improvements on top of the original Nature DQN (Mnih et al., 2015).
Revised: For the discrete control, we use the data-efficient Rainbow DQN (Van Hasselt et al., 2019) which applied multiple improvements on top of the original Nature DQN (Mnih et al., 2015).

Original: 3. S TATE A BSTRACTION
Revised: 3. S TATE A BSTRACTION

Original: While visual RL has achieved many successes in simulated tasks, it remains challenging to learn robust representations from real vision, where images reveal detailed scenes of a complex and unstructured world (Zhang et al., 2020b; Agarwal et al., 2021;
Revised: While visual RL has achieved many successes in simulated tasks, it remains challenging to learn robust representations from real vision, where images reveal detailed scenes of a complex and unstructured world (Zhang et al., 2020b; Agarwal et al., 2021;

Original: Wang et al., 2022).
Revised: Wang et al., 2022).

Original: Therefore, abstracting meaningful elements from the visual scene to present the underlying state is signiﬁcantly important for visual RL.
Revised: Therefore, abstracting meaningful elements from the visual scene to present the underlying state is significantly important for visual RL.

Original: We follow the Block Markov Decision Process (BMDP) (Du et al., 2019), which refers to episodic learning tasks via an unobservable latent space S and an observable context space X .
Revised: We follow the Block Markov Decision Process (BMDP) (Du et al., 2019), which refers to episodic learning tasks via an unobservable latent space S and an observable context space X .

Original: The environment generates a context by x ⇠ p ( ·| s ) .
Revised: The environment generates a context by x ∼ p ( ·| s ) .

Original: They present a fundamental assumption as: each observation x uniquely determines its generating state s .
Revised: They present a fundamental assumption as: each observation x uniquely determines its generating state s .

Original: Similarly, in model-free RL without modeling dynamics, the manipulated context x can be conditioned on a certain probability given an environment transition e which is x ⇠ p ( ·| e ) .
Revised: Similarly, in model-free RL without modeling dynamics, the manipulated context x can be conditioned on a certain probability given an environment transition e which is x ∼ p ( ·| e ) .

Original:  T HE C O IT 4. L EARNABLE I NVARIANT T RANSFORMATION
Revised:  T HE C O IT 4. L EARNABLE I NVARIANT T RANSFORMATION

Original: Following the motivation of smoothing training experiences to stabilize the target Q network (Mnih et al., 2013; Lillicrap et al., 2015), the transformed x 0 is required to satisfy x 0 ⇠ p ( ·| e ) , where environment transition e is ideally in the replay distribution D .
Revised: Following the motivation of smoothing training experiences to stabilize the target Q network (Mnih et al., 2013; Lillicrap et al., 2015), the transformed x ′ is required to satisfy x ′ ∼ p ( ·| e ) , where environment transition e is ideally in the replay distribution D .

Original: Note that e is a random variables.
Revised: Note that e is a random variables.

Original: Formally, we are ready to introduce the optimal invariant metric to reach the stationary distribution D over the augmented context x 0 , through the deﬁnition,
Revised: Formally, we are ready to introduce the optimal invariant metric to reach the stationary distribution D over the augmented context x ′ , through the definition,

Original: Deﬁnition 4.1.
Revised: Definition 4.1.

Original: (Optimal Invariant Metric).
Revised: (Optimal Invariant Metric).

Original: Given a transition distribution D for tuples in the replay buffer, suppose the block structure assumption holds, the shift between transitions x and its context x 0 can be measured by a conditional divergence:
Revised: Given a transition distribution D for tuples in the replay buffer, suppose the block structure assumption holds, the shift between transitions x and its context x ′ can be measured by a conditional divergence:

Original: [Equation]
Revised: [Equation]

Original: where d KL ( ·||· ) is the Kullback-Leibler (KL) divergence.
Revised: where d KL ( ·||· ) is the Kullback-Leibler (KL) divergence.

Original: It indicates the expected distance between x and x 0 conditioning on e .
Revised: It indicates the expected distance between x and x ′ conditioning on e .

Original: One may argue that the dynamics cannot be assumed as a ﬁxed distribution when it comes to new observations, especially after data manipulation.
Revised: One may argue that the dynamics cannot be assumed as a fixed distribution when it comes to new observations, especially after data manipulation.

Original: Nevertheless, it generally claims that the experiences are uniformly sampled in a replay memory (Mnih et al., 2013).
Revised: Nevertheless, it generally claims that the experiences are uniformly sampled in a replay memory (Mnih et al., 2013).

Original: Also, the fact that the given conditions of the transition are consistent for the observed data as well as the transformed data, makes the above deﬁnition reasonable.
Revised: Also, the fact that the given conditions of the transition are consistent for the observed data as well as the transformed data, makes the above definition reasonable.

Original: Next, we will show why the conditional divergence deﬁned in Eq.(3) is an optimal invariant metric from theoretical perspectives.
Revised: Next, we will show why the conditional divergence defined in Eq.(3) is an optimal invariant metric from theoretical perspectives.

Original: We employ the Bayes’ rule on the conditionally distribution p ( x | e ) =
Revised: We employ the Bayes’ rule on the conditionally distribution p ( x | e ) =

Original: p ( e | x ) p ( x ) /p
Revised: p ( e | x ) p ( x ) /p

Original: ( e ) , 8 x 
Revised: ( e ) , ∀ x ∈

Original: O , e 2 D .
Revised: O , e ∈ D .

Original: Then the transition operator p ( e | x ) can be further divided as p ( e | s ) p ( s | x ) for any x 2 O , if e and x are conditional independent given s 1 . Eq.(3) is rewritten as:
Revised: Then the transition operator p ( e | x ) can be further divided as p ( e | s ) p ( s | x ) for any x ∈ O , if e and x are conditional independent given s 1 . Eq.(3) is rewritten as,

Original: [Equation]
Revised: [Equation]

Original: Therefore, minimizing the conditional divergence leads to encoding the observation x and the transformed context x 0 into an invariant latent state space S .
Revised: Therefore, minimizing the conditional divergence leads to encoding the observation x and the transformed context x ′ into an invariant latent state space S .

Original: As a consequence, the learnable pixel transformation is an optimality invariant combining a qualiﬁed encoder.
Revised: As a consequence, the learnable pixel transformation is an optimality invariant combining a qualified encoder.

Original: Now we have the observation encoder g : O ! S mapping from the observed state O to the latent state S by a non-trivial function g such that g ( x ) =
Revised: Now we have the observation encoder g : O → S mapping from the observed state O to the latent state S by a non-trivial function g such that g ( x ) =

Original: p ( s | x ) , 8 x .
Revised: p ( s | x ) , ∀ x .

Original: Traditionally, there is only one encoder dubbed feature backbone in RL models.
Revised: Traditionally, there is only one encoder dubbed feature backbone in RL models.

Original: Since the pixel transformation could drift away, e.g., supported by different components with those supporting x (Du et al., 2019).
Revised: Since the pixel transformation could drift away, e.g., supported by different components with those supporting x (Du et al., 2019).

Original: To enforce the invariant hidden states, another state encoder g 0 that can map x 0 to s should exist, which is g 0 ( x 0 ) =
Revised: To enforce the invariant hidden states, another state encoder g ′ that can map x ′ to s should exist, which is g ′ ( x ′ ) =

Original: p ( s | ⌫ ( x , · )) , 8 x 0 , ⌫ ( x , · ) is the transformation.
Revised: p ( s | ν ( x , · )) , ∀ x ′ , ν ( x , · ) is the transformation.

Original: So far, the goal of learning the optimal transformed data and encoders boils down to minimizing the distance between representations g ( x ) and g 0 ( x 0 ) .
Revised: So far, the goal of learning the optimal transformed data and encoders boils down to minimizing the distance between representations g ( x ) and g ′ ( x ′ ) .

Original: To tackle the issue, we ﬁrst provide two deﬁnitions to introduce measurements as follows,
Revised: To tackle the issue, we first provide two definitions to introduce measurements as follows,

Original: Deﬁnition 4.2.
Revised: Definition 4.2.

Original: ( ✏ -Approximation).
Revised: ( ϵ -Approximation).

Original: Given a distance metric d : O ⇥ S !
Revised: Given a distance metric d : O×S → R + satisfies d

Original: R + satisﬁes d ( s , s ) = 0 , 8 s , and let g, g 0 :O !
Revised: ( s , s ) = 0 , ∀ s , and let g, g ′ :

Original:  S be two functions.
Revised: O → S be two functions.

Original: Let ✏  0 , given a distribution ˆ D on O , then g and g 0 are ✏ - approximate w.r.t.
Revised: Let ϵ ≥ 0 , given a distribution ˆ D on O , then g and g ′ are ϵ - approximate w.r.t.

Original: ( d, [Equation]
Revised: ( d, D ) if E x ∼ ˆ D [ d ( g ( x ) , g ′ ( x ))] ≤ ϵ .

Original: Deﬁnition 4.3.
Revised: Definition 4.3.

Original: (  -Similarity).
Revised: ( β -Similarity).

Original: Given a distance metric d : O ⇥ S !
Revised: Given a distance metric d : O × S → R + satisfies

Original: [Equation]
Revised: d ( s , s ) = 0 , ∀ s .

Original: There exists g : O ! S .
Revised: There exists g : O → S .

Original: Let   0 , given distributions ˆ D and ˆ D 0 on O , then x and x 0 are  - similar if 
Revised: Let β ≥ 0 , given distributions ˆ D and ˆ D ′ on O , then x and x ′ are β - similar if E x ∼ ˆ D , x ′ ∼ ˆ D ′

Original: [Equation]
Revised: [ d ( g ( x ) , g ( x ′ ))] ≤ β .

Original: Without loss of generality, the distance between the encoded states g ( x ) and g 0 ( x 0 ) can be expressed as the following triangular inequality.
Revised: Without loss of generality, the distance between the encoded states g ( x ) and g ′ ( x ′ ) can be expressed as the following triangular inequality.

Original: To obtain a metric, Kullback-Leibler divergence is rewritten in a form of the square root of Jensen-Shannon divergence.
Revised: To obtain a metric, Kullback-Leibler divergence is rewritten in a form of the square root of Jensen-Shannon divergence.

Original: Therefore, we have,
Revised: Therefore, we have,

Original: [Equation]
Revised: [Equation]

Original: From the view of invariant learning, minimizing the right side of the inequality can upper bound our problem.
Revised: From the view of invariant learning, minimizing the right side of the inequality can upper bound our problem.

Original: The ﬁrst term on the right side is the so-called ✏ -approximation to measure the functional similarity after state abstraction, whereas the second term exists based on the procedure of data augmentation.
Revised: The first term on the right side is the so-called ϵ -approximation to measure the functional similarity after state abstraction, whereas the second term exists based on the procedure of data augmentation.

Original: Thus, we learn the encoders and shifted data simultaneously through the upper bound.
Revised: Thus, we learn the encoders and shifted data simultaneously through the upper bound.

Original: 4. O PTIMAL S TATE A BSTRACTION
Revised: 4. O PTIMAL S TATE A BSTRACTION

Original: To restrict the functional similarity of Eq.(5) from the perspective of learning a good encoding function with consistent semantics, the approaches formulating the main and momentum feature learning are utilized, motivated by contrastive learning (He et al., 2020).
Revised: To restrict the functional similarity of Eq.(5) from the perspective of learning a good encoding function with consistent semantics, the approaches formulating the main and momentum feature learning are utilized, motivated by contrastive learning (He et al., 2020).

Original: In particular, we enforce the encoding functions with exactly the same architecture, and use ¯ ⇠ t =(1  ⌧ m )
Revised: In particular, we enforce the encoding functions with exactly the same architecture, and use ¯ ξ t =

Original: ¯ ⇠ t  1 + ⌧ m ⇠ t at timestep t to update the parameters of momentum function g ¯ ⇠ with g ⇠ , where ⌧ m 2 [0 , 1] is the updating rate.
Revised: (1 − τ m )¯ ξ t − 1 + τ m ξ t at timestep t to update the parameters of momentum function g ¯ ξ with g ξ , where τ m ∈ [0 , 1] is the updating rate.

Original: Furthermore, we design a projection that is f : S ! Y using a ReLU network (Petersen & Voigtlaender, 2018) to upper bound the divergence by minimizing the distance in the projected space Y .
Revised: Furthermore, we design a projection that is f : S → Y using a ReLU network (Petersen & Voigtlaender, 2018) to upper bound the divergence by minimizing the distance in the projected space Y .

Original: Previously, the projection has been proposed by Chen et al.
Revised: Previously, the projection has been proposed by Chen et al.

Original: (2020), while the theoretical guarantees of the underlying mechanism with momentum updating for model-free RL are explained in this work.
Revised: (2020), while the theoretical guarantees of the underlying mechanism with momentum updating for model-free RL are explained in this work.

Original: Suppose the Markov chain O g ! S f ! Y holds.
Revised: Suppose the Markov chain O g −→ S f −→ Y holds.

Original: For two functions g and f in the compatible ranges, we use f  g to denote the function composition f
Revised: For two functions g and f in the compatible ranges, we use f ◦ g to denote the function composition f

Original: ( g ( · )) .
Revised: ( g ( · )) .

Original: Before showing the proposed data transformed method, we introduce technical lemmas to take advantage of the designable projection function by leveraging the convexity.
Revised: Before showing the proposed data transformed method, we introduce technical lemmas to take advantage of the designable projection function by leveraging the convexity.

Original: The momentum updating paradigm is capable of turning into momentum feature updating through a convex function or an equivalence of the convex function.
Revised: The momentum updating paradigm is capable of turning into momentum feature updating through a convex function or an equivalence of the convex function.

Original: Lemma 4.1.
Revised: Lemma 4.1.

Original: Assume that h : R |S| ! R |Y| can be written as h ( ⇠ ) =
Revised: Assume that h : R |S| → R |Y| can be written as h ( ξ ) =

Original: f ( < ⇠ , s > ) , for some s 2 R |S| , and f : R |Y| ! R |Y| with parameter ⇠ .
Revised: f ( < ξ, s > ) , for some s ∈ R |S| , and f : R |Y| → R |Y| with parameter ξ .

Original: Then, convexity of f implies the convexity of h .
Revised: Then, convexity of f implies the convexity of h .

Original: Lemma 4.2.
Revised: Lemma 4.2.

Original: Given the dynamical updating: ¯ ⇠ t = (1  ⌧ m ) ¯ ⇠ t  1 + ⌧ m ⇠ t .
Revised: Given the dynamical updating: ¯ ξ t = (1 − τ m )¯ ξ t − 1 + τ m ξ t .

Original: By Lemma 4.1, f ⇠ = f ¯ ⇠ holds after convergence.
Revised: By Lemma 4.1, f ξ = f ¯ ξ holds after convergence.

Original: As a result, the problem of min E x 
Revised: As a result, the problem of min E x [ ∥

Original: [ k f ⇠  g ⇠ ( x )  f ¯ ⇠  g ¯ ⇠ ( x ) k ] is equivalent to the problem of min E x [ k g ⇠ ( x )  g ¯ ⇠ ( x ) k ] .
Revised: f ξ ◦ g ξ ( x ) − f ¯ ξ ◦ g ¯ ξ ( x ) ∥ ] is equivalent to the problem of min E x [ ∥ g ξ ( x ) − g ¯ ξ ( x ) ∥ ] .

Original: To meet the requirement of a small upper bound, we state a theorem that provides some insights into why it is necessary to learn optimal transformed data together with the encoders.
Revised: To meet the requirement of a small upper bound, we state a theorem that provides some insights into why it is necessary to learn optimal transformed data together with the encoders.

Original: Theorem 4.1.
Revised: Theorem 4.1.

Original: (CoIT)
Revised: (CoIT)

Original: Suppose that Lipschitzness holds for functions g ⇠ , g ¯ ⇠ , f ⇠ and f ¯ ⇠ , respectively.
Revised: Suppose that Lipschitzness holds for functions g ξ , g ¯ ξ , f ξ and f ¯ ξ , respectively.

Original: The updating dynamics is: ¯ ⇠ t = (1  ⌧ m ) ¯ ⇠ t  1 + ⌧ m ⇠ t , ⌧ m 2 [0 , 1] .
Revised: The updating dynamics is: ¯ ξ t = (1 − τ m )¯ ξ t − 1 + τ m ξ t , τ m ∈ [0 , 1] .

Original: For any input x ⇠ ˆ D and transformed x 0 obtained via the transform operator ⌫ ( x , · ) , optimizing the conditional divergence in
Revised: For any input x ∼ ˆ D and transformed x ′ obtained via the transform operator ν ( x , · ) , optimizing the conditional divergence in

Original: Deﬁnition 4.1 means to minimize the upper bound as follows,
Revised: Definition 4.1 means to minimize the upper bound as follows,

Original: [Equation]
Revised: [Equation]

Original:  where ⇢ = L g ( CL f + k ⇠ f k ) , C = 1+ ⌧ 1  ⌧ , ⌧ = 1  ⌧ m are constants.
Revised: ) where ρ = L g ( CL f + ∥ ξ f ∥ ) , C = 1+ τ 1 − τ , τ = 1 − τ m are constants.

Original: L f and L g are Lipschitz constants of the functions f ( s ) and g ( x ) , respectively.
Revised: L f and L g are Lipschitz constants of the functions f ( s ) and g ( x ) , respectively.

Original: The upper bound of the right side measures the margin of augmentation between original and transformed data.
Revised: The upper bound of the right side measures the margin of augmentation between original and transformed data.

Original: The left side measures the distribution changes in the latent space.
Revised: The left side measures the distribution changes in the latent space.

Original: Since both transformed data x 0 and encoder g are updating, incorporating augmentation directly cannot well meet the basic stationary environment.
Revised: Since both transformed data x ′ and encoder g are updating, incorporating augmentation directly cannot well meet the basic stationary environment.

Original: Hereby, the theorem suggests us that the automatic transformation is used to bound the representation learning so that the abstracted states enhance the stationary distribution of tuples e and facilitate efﬁcient training.
Revised: Hereby, the theorem suggests us that the automatic transformation is used to bound the representation learning so that the abstracted states enhance the stationary distribution of tuples e and facilitate efficient training.

Original: Proofs are given in Appendix A.
Revised: Proofs are given in Appendix A.

Original: The empirical comparisons of the projection network are presented in Appendix D.
Revised: The empirical comparisons of the projection network are presented in Appendix D.

Original: 4. P ARAMETERIZABLE O BSERVATION
Revised: 4. P ARAMETERIZABLE O BSERVATION

Original: From the Theorem 4.1, we know the relation between the parameterized augmentation and the learnable latent state.
Revised: From the Theorem 4.1, we know the relation between the parameterized augmentation and the learnable latent state.

Original: That is, image transformation needs to be constrained by the distance between an observation and its associated augmentation.
Revised: That is, image transformation needs to be constrained by the distance between an observation and its associated augmentation.

Original: The optimal embedding can be obtained by minimizing of Eq.(5).
Revised: The optimal embedding can be obtained by minimizing of Eq.(5).

Original: Particularly, we parameterize the transformed data x 0 as ⌫ ( x , G ) , where G is a Gaussian distribution dynamically changed along with the RL training.
Revised: Particularly, we parameterize the transformed data x ′ as ν ( x , G ) , where G is a Gaussian distribution dynamically changed along with the RL training.

Original: In Algorithm 1, it deﬁnes the MDP M with Gaussian random variables G 0 ⇠ G |O| for initialization.
Revised: In Algorithm 1, it defines the MDP M with Gaussian random variables G 0 ∼ G |O| for initialization.

Original: The TRANSFORM operator is fulﬁlled by the aforementioned pixel transformation ⌫ which is a shift  x 0 = ⌫ ( x , G ) is normalized according to both x and learned distribution G , and thereby contributes to the cumulative reward maximization in an interactive way.
Revised: The TRANSFORM operator is fulfilled by the aforementioned pixel transformation ν which is a shift following Gaussian distribution G t ( µ t , σ t ) on the top of data interpolation. The transformation x ′ = ν ( x , G ) is normalized according to both x and learned distribution G , and thereby contributes to the cumulative reward maximization in an interactive way.

Original: We use the scope to the observation sampling for an optimal state abstraction.
Revised: We use the scope to the observation sampling for an optimal state abstraction.

Original: It can be regarded as data sampling from the replay buffer to reach a training-friendly distribution, dubbed Contrastive Invariant 
Revised: It can be regarded as data sampling from the replay buffer to reach a training-friendly distribution, dubbed Contrastive Invariant Transformation .

Original: 4. S TABILIZING R EWARD F UNCTION
Revised: 4. S TABILIZING R EWARD F UNCTION

Original: It identiﬁes one of the key impacts that CoIT in the RL training procedure as parametrizing the underlying invariant optimization to smooth the distribution D in the replay buffer.
Revised: It identifies one of the key impacts that CoIT in the RL training procedure as parametrizing the underlying invariant optimization to smooth the distribution D in the replay buffer.

Original: To further stabilize the reward function, we propose a mixed CoIT that samples multiple transformed data from the learned distribution
Revised: To further stabilize the reward function, we propose a mixed CoIT that samples multiple transformed data from the learned distribution

Original: G ( µ,  ) , and then mix up the learned observation x 0 .
Revised: G ( µ, σ ) , and then mix up the learned observation x ′ .

Original: Similarly, we provide the invariant learning guarantee by optimizing the right side transformation in Theorem 4.2.
Revised: Similarly, we provide the invariant learning guarantee by optimizing the right side transformation in Theorem 4.2.

Original: Theorem 4.2.
Revised: Theorem 4.2.

Original: (Mixed CoIT) Suppose that
Revised: (Mixed CoIT) Suppose that

Original: Lipschitzness holds for functions g ⇠ , g ¯ ⇠ , f ⇠ and f ¯ ⇠ , respectively.
Revised: Lipschitzness holds for functions g ξ , g ¯ ξ , f ξ and f ¯ ξ , respectively.

Original: The updating dynamics is: ¯ ⇠ t = (1  ⌧ m ) ¯ ⇠ t  1 + ⌧ m ⇠ t .
Revised: The updating dynamics is: ¯ ξ t = (1 − τ m )¯ ξ t − 1 + τ m ξ t .

Original: For any input x ⇠ ˆ D and transformed x 0 , the divergence with mixed transformed observation can be bound by,
Revised: For any input x ∼ ˆ D and transformed x ′ , the divergence with mixed transformed observation can be bound by,

Original: [Equation]
Revised: [Equation]

Original: where ⇢ = L g ( CL f + k ⇠ f k ) , C = 1+ ⌧ 1  ⌧ , ⌧ = 1  ⌧ m .
Revised: where ρ = L g ( CL f + ∥ ξ f ∥ ) , C = 1+ τ 1 − τ , τ = 1 − τ m .

Original: L f and L g are Lipschitz constants of the functions f ( s ) and g ( x ) respectively.
Revised: L f and L g are Lipschitz constants of the functions f ( s ) and g ( x ) respectively.

Original: The proof of Theorem 4.2 is straightforward based on Theorem 4.1, and the details are presented in the supplement.
Revised: The proof of Theorem 4.2 is straightforward based on Theorem 4.1, and the details are presented in the supplement.

Original: 4. L EARNING C ONTRASTIVE I NVARIANT T RANSFORMATION
Revised: 4. L EARNING C ONTRASTIVE I NVARIANT T RANSFORMATION

Original: With theoretical analysis of invariant transformations, we presented a new framework with normalization variants to ensure above discussed learning guarantees by optimizing parameters.
Revised: With theoretical analysis of invariant transformations, we presented a new framework with normalization variants to ensure above discussed learning guarantees by optimizing parameters.

Original: We initialize a distribution G t ( µ,  ) and use the TRANSFORM operator to produce different views of x t (Algorithm 1).
Revised: We initialize a distribution G t ( µ, σ ) and use the TRANSFORM operator to produce different views of x t (Algorithm 1).

Original: The transformed data x 0 t is viewed as the positive pair of x t .
Revised: The transformed data x ′ t is viewed as the positive pair of x t .

Original: We also utilize a similarity metric d to learn contrastive invariant transformation for the encoder g ⇠ ( · ) .
Revised: We also utilize a similarity metric d to learn contrastive invariant transformation for the encoder g ξ ( · ) .

Original: We ﬁrst apply the bilinear interpolation to x t and sample shift terms from G t to produce multiple positive samples x 1 t , x 2 t , .., x n t and mix them together as x 0 t following Theorem 4.2.
Revised: We first apply the bilinear interpolation to x t and sample shift terms from G t to produce multiple positive samples x 1 t , x 2 t , .., x nt and mix them together as x ′ t following Theorem 4.2.

Original: To prevent the distribution of transformed data from shifting too far away from the replay buffer D , we borrow a similar idea from Yin et al.
Revised: To prevent the distribution of transformed data from shifting too far away from the replay buffer D , we borrow a similar idea from Yin et al.

Original: (2020) to regularize and smooth the distribution shift between the transformed and overall data.
Revised: (2020) to regularize and smooth the distribution shift between the transformed and overall data.

Original: We use the statistical data stored in the BatchNorm layers to approximate the distribution of the overall data.
Revised: We use the statistical data stored in the BatchNorm layers to approximate the distribution of the overall data.

Original: In this way, the distribution shift between the transformed and overall data can be estimated by the following formulation,
Revised: In this way, the distribution shift between the transformed and overall data can be estimated by the following formulation,

Original: [Equation]
Revised: [Equation]

Original: where ˜ µ ( x 0 t ) and ˜  2 ( x 0 t ) are the mean and variance of the transformed data and ! represents parameter collection { µ t ,  t } of the Gaussian distribution G t ( µ,  ) .
Revised: where ˜ µ ( x ′ t ) and ˜ σ 2 ( x ′ t ) are the mean and variance of the transformed data and ω represents parameter collection { µ t , σ t } of the Gaussian distribution G t ( µ, σ ) .

Original: The expectation terms E (˜ µ l ( x ) |O ) and E (˜  2 l ( x ) |O ) respectively denote the estimation of the batch-wise mean and variance for the feature map corresponding to the l -th convolution layer, and O is the given observations.
Revised: The expectation terms E (˜ µ l ( x ) |O ) and E (˜ σ 2 l ( x ) |O ) respectively denote the estimation of the batch-wise mean and variance for the feature map corresponding to the l -th convolution layer, and O is the given observations.

Original: Second, we utilize the similarity metric d proposed by Chen et al.
Revised: Second, we utilize the similarity metric d proposed by Chen et al.

Original: (2020) for learning the encoder g ⇠ ( · ) which maps high-dimensional observation to embeddings to meet the invariant transformation in Eq.5.
Revised: (2020) for learning the encoder g ξ ( · ) which maps high-dimensional observation to embeddings to meet the invariant transformation in Eq.5.

Original: Given a positive observation pair ( x t , x 0 t ) , the loss is given by
Revised: Given a positive observation pair ( x t , x ′ t ) , the loss is given by

Original: [Equation]
Revised: [Equation]

Original: Here ¯ ⇠ denotes the momentum version of parameters ⇠ and f ⇠ ( · ) is a non-linear projection of the representations embedded by g ⇠ ( · ) .
Revised: Here ¯ ξ denotes the momentum version of parameters ξ and f ξ ( · ) is a non-linear projection of the representations embedded by g ξ ( · ) .

Original: D indicates the tuples stored in the replay buffer.
Revised: D indicates the tuples stored in the replay buffer.

Original: Next, we update the critic network Q ✓ with transformed data x 0 t and x 0 t + n to minimize the TD error for n -steps returns.
Revised: Next, we update the critic network Q θ with transformed data x ′ t and x ′ t + n to minimize the TD error for n -steps returns.

Original: This is regarded as a regularized Q learning by Yarats et al.
Revised: This is regarded as a regularized Q learning by Yarats et al.

Original: (2020) where the regularized representation learning is beneﬁcial for optimal action taking (Zhang et al., 2020a).
Revised: (2020) where the regularized representation learning is beneficial for optimal action taking (Zhang et al., 2020a).

Original: [Equation]
Revised: [Equation]

Original: Eventually, we give the uniﬁed objective function as the full version of the CoIT,
Revised: Eventually, we give the unified objective function as the full version of the CoIT,

Original: [Equation]
Revised: [Equation]

Original: where ↵ and  are hyper-parameters and the overall architecture is presented in Figure 2.
Revised: where α and λ are hyper-parameters and the overall architecture is presented in Figure 2.

Original: We replace the vanilla Q-learning by J ✓ , ⇠ , ! ( D ) and the entire algorithm is presented in Algorithm 2 in Appendix
Revised: We replace the vanilla Q-learning by J θ,ξ,ω ( D ) and the entire algorithm is presented in Algorithm 2 in Appendix

Original: B. Then, we evaluate CoIT on popular benchmarks to demonstrate the beneﬁts of our method.
Revised: B. Then, we evaluate CoIT on popular benchmarks to demonstrate the benefits of our method.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we benchmark our method on the DeepMind control suite and Atari100K.
Revised: In this section, we benchmark our method on the DeepMind control suite and Atari100K.

Original: We compare
Revised: We compare

Original: CoIT with prior model-free methods ﬁrst, then we present ablation studies to show the details of our method.
Revised: CoIT with prior model-free methods first, then we present ablation studies to show the details of our method.

Original: Implementation details can be found in Appendix C
Revised: Implementation details can be found in Appendix C.

Original: 5. E NVIRONMENTS
Revised: 5. E NVIRONMENTS

Original: DMControl.
Revised: DMControl.

Original: DeepMind control suite (Tassa et al., 2018) is a widely used benchmark with several robot control tasks.
Revised: DeepMind control suite (Tassa et al., 2018) is a widely used benchmark with several robot control tasks.

Original: Each episode is set to be 1 , 000 frames and we use the total experienced frames to measure the data-efﬁciency.
Revised: Each episode is set to be 1 , 000 frames and we use the total experienced frames to measure the data-efficiency.

Original: The per-frame reward is in the unit interval [0 , 1] , so each episode contains a total reward of no more than 1 , 000 .
Revised: The per-frame reward is in the unit interval [0 , 1] , so each episode contains a total reward of no more than 1 , 000 .

Original: Considering the different difﬁculties depending on tasks, we refer to setting more episodes with hard tasks for better evaluation.
Revised: Considering the different difficulties depending on tasks, we refer to setting more episodes with hard tasks for better evaluation.

Original: Atari100K. There have been a number of prior papers that have benchmarked data-efﬁciency on the Atari 2600 Games for discrete control.
Revised: Atari100K. There have been a number of prior papers that have benchmarked data-efficiency on the Atari 2600 Games for discrete control.

Original: Van Hasselt et al.
Revised: Van Hasselt et al.

Original: and Kielak (2019) propose the data-efﬁcient version of Rainbow DQN (Hessel et al., 2018) compared with human performance (Kaiser et al., 2019) within 100K time steps (400K frames, frame skip of 4).
Revised: and Kielak (2019) propose the data-efficient version of Rainbow DQN (Hessel et al., 2018) compared with human performance (Kaiser et al., 2019) within 100K time steps (400K frames, frame skip of 4).

Original: This sample-constrained evaluation is the so-called Atari100K and we benchmark CoIT on all 26 Atari Games.
Revised: This sample-constrained evaluation is the so-called Atari100K and we benchmark CoIT on all 26 Atari Games.

Original: 5. B ASELINES
Revised: 5. B ASELINES

Original: DMControl.
Revised: DMControl.

Original: For continuous control we present several baselines, including methods of using data augmentation and contrastive learning to improve data-efﬁciency: (i) DrQ-v2 (Yarats et al., 2021), (ii)
Revised: For continuous control we present several baselines, including methods of using data augmentation and contrastive learning to improve data-efficiency: (i) DrQ-v2 (Yarats et al., 2021), (ii)

Original: Atari100K.
Revised: Atari100K.

Original: To benchmark the data-efﬁciency of CoIT for discrete control tasks, we compare our method to (i) DrQ (Yarats et al., 2020), (ii) CURL (Laskin et al., 2020b),
Revised: To benchmark the data-efficiency of CoIT for discrete control tasks, we compare our method to (i) DrQ (Yarats et al., 2020), (ii) CURL (Laskin et al., 2020b),

Original: (iii) SPR (Schwarzer et al., 2020), (iv) Random Agent and Human Performance (Kaiser et al., 2019), All the algorithms are evaluated within 100K time steps for interaction.
Revised: (iii) SPR (Schwarzer et al., 2020), (iv) Random Agent, and Human Performance (Kaiser et al., 2019). All the algorithms are evaluated within 100K time steps for interaction.

Original: We average CoIT’s performance over 10 random seeds and report the best score for each game following prior works.
Revised: We average CoIT’s performance over 10 random seeds and report the best score for each game following prior works.

Original: 5. M AIN R ESULTS
Revised: 5. M AIN R ESULTS

Original: DMControl.
Revised: DMControl.

Original: We choose 8 complex tasks from the DMC for evaluation and present the results in Figure 3 and Table 2 in Appendix C. We also report the percentage (%) of score solved in the
Revised: We choose 8 complex tasks from the DMControl for evaluation and present the results in Figure 3 and Table 2 in Appendix C. We also report the percentage (%) of score solved in the

Original: DMControl for baselines and CoIT in 500K and 100K steps in Figure 1.
Revised: DMControl for baselines and CoIT in 500K and 100K steps in Figure 1.

Original: Below are key ﬁndings: (i)
Revised: Below are key findings: (i)

Original: CoIT outperforms vanilla DDPG and SAC in a wide range.
Revised: CoIT outperforms vanilla DDPG and SAC in a wide range.

Original: (ii) We also compare CoIT with DrQ-v2, a remarkable method for continuous control, to better demonstrate our method’s data-efﬁciency.
Revised: (ii) We also compare CoIT with DrQ-v2, a remarkable method for continuous control, to better demonstrate our method’s data-efficiency.

Original: (iii)
Revised: (iii)

Original: From general trends of the learning curves, CoIT improves or keeps the data-efﬁciency in a more stable manner which is not trivial on DM control tasks.
Revised: From general trends of the learning curves, CoIT improves or keeps the data-efficiency in a more stable manner which is not trivial on DMControl tasks.

Original: Atari100K. We present results for Atari100K in Table 1 and below are key ﬁndings: (i) CoIT achieves top-performance on 10 of 26 games while still being competitive in the rest.
Revised: Atari100K. We present results for Atari100K in Table 1 and below are key findings: (i) CoIT achieves top-performance on 10 of 26 games while still being competitive in the rest.

Original: (ii) CoIT surpasses superhuman performance on 6 games  .
Revised: (ii) CoIT surpasses superhuman performance on 6 games on the basis of Rainbow DQN .

Original: (iii)
Revised: (iii)

Original: We also report the mean and std of the scores achieved by CoIT in 10 of 26 games which are top-performance.
Revised: We also report the mean and std of the scores achieved by CoIT in 10 of 26 games which are top-performance.

Original: We present the results of two versions of CoIT (mixed & no-mixed) in Figure 7 in the appendix.
Revised: We present the results of two versions of CoIT (mixed & no-mixed) in Figure 6 in the appendix.

Original: From the histogram, we ﬁnd that CoIT has much better stability, which is similar to the observation in the DMControl.
Revised: From the histogram, we find that CoIT has much better stability, which is similar to the observation in the DMControl.

Original: 5. A BLATION S TUDIES
Revised: 5. A BLATION S TUDIES

Original: We ﬁrst visualize the TRANSFORM operator to demonstrate that there exists an invariant transformation for each task.
Revised: We first visualize the TRANSFORM operator to demonstrate that there exists an invariant transformation for each task.

Original: We initialize the Gaussian distribution G t ( µ,  ) based on the range of pixel shifts and plot the curves of the mean µ and std  during training in the DMC in Figure 4.
Revised: We initialize the Gaussian distribution G t ( µ, σ ) based on the range of pixel shifts and plot the curves of the mean µ and std σ during training in the DMControl in Figure 4.

Original: [Table]
Revised: [Table]

Original: According to the curves below
Revised: According to the curves below

Original: [Table]
Revised: [Table]

Original: we ﬁnd that the mean and std converge to an interval as the training goes on.
Revised: we find that the mean and std converge to an interval as the training goes on.

Original: These results demonstrate that the gaussian distribution proposed in CoIT could automatically ﬁnd a TRANSFORM to smooth the distribution shift between the different views of the same observation, therefore being beneﬁcial to the representation learning.
Revised: These results demonstrate that the gaussian distribution proposed in CoIT could automatically find a TRANSFORM to smooth the distribution shift between the different views of the same observation, therefore being beneficial to the representation learning.

Original: Then, we study the effects of different components in Eq.(11).
Revised: Then, we study the effects of different components in Eq.(11).

Original: This object function is composed of two parts: K ! ( x 0 t ) for regularization and L ⇠ , ¯ ⇠ , ! ( D ) for similarity metric.
Revised: This object function is composed of two parts: K ω ( x ′ t ) for regularization and L ξ, ¯ ξ,ω ( D ) for similarity metric.

Original: On this basis, we divide CoIT into 4 versions: (i) Critic .
Revised: On this basis, we divide CoIT into 4 versions: (i) Critic .

Original: Transformation is only updated with the critic.
Revised: Transformation is only updated with the critic.

Original: (ii) X-stats & Critic .
Revised: (ii) X-stats & Critic .

Original: Transformation is updated by critic and K ! ( x 0 t ) together.
Revised: Transformation is updated by critic and K ω ( x ′ t ) together.

Original: (iii) H-dist & Critic .
Revised: (iii) H-dist & Critic .

Original: Transformation is updated by critic and L ⇠ , ¯ ⇠ , ! ( D ) together.
Revised: Transformation is updated by critic and L ξ, ¯ ξ,ω ( D ) together.

Original: (iv) Uniﬁed Objective .
Revised: (iv) Unified Objective .

Original: We evaluate all of these versions on 4 representative tasks from the DMC and present the results in Figure 5 and in the appendix.
Revised: We evaluate all of these versions on 8 representative tasks from the DMControl and present the results in Figure 9 in the appendix.

Original: Compared Critic with other variants, we demonstrate that both of the components are beneﬁcial to the performance.
Revised: Compared Critic with other variants, we demonstrate that both of the components are beneficial to the performance.

Original: Though Critic is data-efﬁcient on most tasks, it may fall into trivial solutions.
Revised: Though Critic is data-efficient on most tasks, it may fall into trivial solutions.

Original: To solve this issue, we utilize the regularization in Eq.(8) with the similarity metric in Eq.(9) to meet the invariant transformation.
Revised: To solve this issue, we utilize the regularization in Eq.(8) with the similarity metric in Eq.(9) to meet the invariant transformation.

Original: Thus the Uniﬁed Objective ’s performance leads ahead of all tasks.
Revised: Thus the Unified Objective ’s performance leads ahead of all tasks.

Original: See
Revised: See

Original: Appendix C for extra ablation studies.
Revised: Appendix C for extra ablation studies.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: A novel pixel transformation CoIT under model-free RL algorithms that signiﬁcantly improves the data-efﬁciency and stability for visual tasks is introduced in this work.
Revised: A novel pixel transformation CoIT under model-free RL algorithms that significantly improves the data-efficiency and stability for visual tasks is introduced in this work.

Original: We theoretically analyze how the learnable transformation constrains the distribution of transformed data, and dissect its beneﬁts to representation learning.
Revised: We theoretically analyze how the learnable transformation constrains the distribution of transformed data, and dissect its benefits to representation learning.

Original: CoIT is no need for any additional modiﬁcations to the backbone RL algorithm and is easy to implement.
Revised: CoIT is no need for any additional modifications to the backbone RL algorithm and is easy to implement.

Original: We compare CoIT to SOTA methods on popular benchmarks and certify that it gains promising performance with advanced stability.
Revised: We compare CoIT to SOTA methods on popular benchmarks and certify that it gains promising performance with advanced stability.

Original: Hopefully, contrastive invariant transformation can lead to a new branch for representation learning in RL.
Revised: Hopefully, contrastive invariant transformation can lead to a new branch for representation learning in RL.

Original: A CKNOWLEDGMENT
Revised: A CKNOWLEDGMENT

Original: This work was supported in part by the National Key Research and Development Program of
Revised: This work was supported in part by the National Key Research and Development Program of

Original: China No. 2020AAA0103400, National Key Research and Development Program of China No.
Revised: China No. 2020AAA0103400, National Key Research and Development Program of China No.

Original: 2021ZD0201504, National Natural Science Foundation of China under Grant 62273347, and CCFTencent Open Research Fund RAGR20220104.
Revised: 2021ZD0201504, National Natural Science Foundation of China under Grant 62273347, and CCFTencent Open Research Fund RAGR20220104.

Original: We thank anonymous reviewers for their discussions and feedback on the paper
Revised: We thank anonymous reviewers for their discussions and feedback on the paper