S HIFT N ORM : O N D ATA E FFICIENCY IN R EINFORCEMENT L EARNING WITH S HIFT N ORMALIZATION
A BSTRACT
We propose ShiftNorm, a simple yet promising data augmentation that can be applied to standard model-free algorithms to improve data-efficiency in highdimensional observation-based reinforcement learning (RL).
Concretely, the differentiable ShiftNorm leverages original samples with reparameterized virtual samples, and hasten the image encoder to generate invariant representations.
Our approach demonstrates certify substantial advances enabling it to outperform the new state-of-the -art on 8 of 9 tasks on the DeepMind Control
Suite at 500k steps.
I NTRODUCTION
Improving agent’s performance and data-efficiency have always been a key problem in visual reinforcement learning (RL).
Unlike the representation learning in supervised learning, as the model has a strongly supervised singal and various methods could be applied to learn representations useful to the task, there are not enough supervised signals in visual RL, thus the trainging process is quite fragile.
The network need to learn useful representations for performance improvement, while inappropriate methods may do damage to the training process, causing performance degradation.
In this case, we urgently need representation learning methods suitable for reinforcement learning.
Previous works have demonstrated that data augmentation could better this situation, both for sampleefficiency (Laskin et al., 2020a;
Yarats et al., 2020; 2021) and generalization (Hansen et al., 2021;
Raileanu et al., 2020; Zhang & Guo, 2021; Hansen & Wang, 2021;
Fan et al., 2021).
This method without further modifications to the backbone of RL algorithms is being noticed by others.
Inspired by Spatial Transformer Networks (STN) (Jaderberg et al., 2015), we hold the point that the observation transformation process can be parameterized.
We focus on transforming handcrafted perturbations into an optimizeable process and propose ShiftNorm to improve the data-efficiency at pixel-based tasks.
We hypothesize that there is a suitable transformation of the observation for the agent, thus the learned transformation can enable the encoder to abstract more useful semantic representations from high-dimensional obsevations, and control algorithms based on these representations should be more sample-efficiency.
Following this way, we propose ShiftNorm to improve the data-efficiency at pixel-based tasks.
Through this process, the pixel shift in different degrees and directions will occur in the whole image and the shift controls the level of perturbation of the image.
Here we raise our ideas:
Can we reparameterize this shift procedure by sampling it from a dynamic distribution to cope with the assumption of stationary environment in
model -free RL training?
To this aim, we parameterize the mean and variance of the distribution and update them with the
RL process.
Two constraints are also proposed to ensure that the automatic augmentation will not ruin the convergence of the RL algorithm.
As the iterations increase, the agent will find the shift distribution suitable for the task at hand.
Key Contributions: We introduce ShiftNorm, a reparameterized data shift method to integrate invariant representations with model-free RL methods.
We evaluate our algorithm on 9 control tasks from the DM control suite and show that the optimizable transformation is effective to improve data-efficiency for visual RL.
T HE S HIFT N ORM 2. P ROBLEM F ORMULATION
We formulate the visual RL as an infinite-horizon Markov Decision Process (MDP) (Bellman, 1957).
The MDP M can be described as a 5-tuple ⟨O , S , A , r, γ ⟩ .
Here O consists of a stack of images (Mnih et al., 2013).
The state space S is either observable or unobservable (Silver et al., 2017; Zhang et al., 2020).
A is the action space for the agent.
The goal is to maximize the cumulative rewards
R = (cid:80) t γ t r t , where γ ∈ [0 , 1) is the discount factor and r t denotes the reward at time t .
2. L EARNABLE I NVARIANT T RANSFORMATION
We first introduce the optimal invariant metric to reach the stationary distribution D over the augmented context x ′ , where x ′ is required to satisfy x ′ ∼ q ( ·| e ) .
Environment transition e belongs to the replay distribution D .
Below are the definition:
Definition 2.1.
(Optimal Invariant Metric).
Given a transition distribution D , suppose the block structure assumption holds, the shift between x and x ′ can be measured by a conditional divergence:
[Equation]
Following the Bayes’ rule on the conditional distribution, Eq.(1) can be rewritten as:
[Equation]
where d KL ( · , · ) is the Kullback–Leibler (KL) divergence.
Therefore, minimizing the conditional divergence leads to encoding x and x ′ into an invariant latent state space S .
Now we define a non-trivial function g : O → S mapping from the observed state O to the latent state S such that g ( x ) =
p ( s | x ) , ∀ x .
Since the pixel transformation is reparameterized as ν ( x , G ) and could drift away, it is natural to find another state encoder g ′ ( x ′ )
=
p ( s | ν ( x , G )) , ∀ x ′ , and this encoding function g ′ should be different with the original observation encoder g .
So far, the learning goal boils down to minimizing the distance between g ( x ) and g ′ ( x ′ ) ,
Definition 2.2.
(State ϵ -Approximation).
Given a distance metric d : O × S → R + satisfies
d ( s , s ) = 0 , ∀ s , and let g, g ′ :
O → S be two functions.
Let ϵ ≥ 0 , given a distribution ˆ D on O , then g and g ′ are ϵ - approximate w.r.t.
( d, D )
if
[Equation]
Define the distance d ( · , · ) as l p -norm.
If g ′ satisfies L p -Lipschitzness
according to Assumption B.1, the distance between the s and s ′ can be expressed as the following triangular inequality,
[Equation]
Minimizing the right side of the inequality is able to upper bound our problem.
To restrict the functional similarity, we combine the momentum updating (He et al., 2020) with a projection f : S → Y (Chen et al., 2020a) and minimize the distance in the projected space Y for model-free RL.
Suppose the Markov chain O g −→ S f −→ Y holds.
We use f ◦ g to denote the function composition f
( g ( · )) .
Also, there exists the assumption of Lipschitzness for the projection f .
Before present data shifted method, we leverage the convexity in momentum updating paradigm.
Lemma 2.1.
Assume that h : R |S| → R |Y| can be written as h ( ξ ) =
f ( < ξ, s > ) , for some s ∈ R |S| , and f : R |Y| → R |Y| with parameter ξ .
Then, convexity of f implies the convexity of h .
Lemma 2.2.
Given the dynamical updating: ¯ ξ t = (1 − τ m )¯ ξ t − 1 + τ m ξ t .
By Lemma 2.1, f ξ = f ¯ ξ holds after convergence.
As a result, the problem of min E x [ ∥
f ξ ◦ g ξ ( x ) − f ¯ ξ ◦ g ¯ ξ ( x ) ∥ ] is equivalent to the problem of min E x [ ∥ g ξ ( x ) − g ¯ ξ ( x ) ∥ ] .
We then provide helpful insight for learning an optimal shifted data together with the encoders.
Theorem 2.1.
(Shift Normalization.)
Suppose Assumption B.1 hold for functions g ξ , g ¯ ξ , f ξ , and f ¯ ξ , respectively.
The updating dynamics is: ¯ ξ t = (1 − τ m )¯ ξ t − 1 + τ m ξ t , τ m ∈ [0 , 1] .
For any input x ∼ ˆ D and shifted x ′ obtained via ν ( x , G ) , optimizing the conditional divergence in Definition 2. means to minimize the upper bound as follows, where C = 1+ τ 1 − τ , τ = 1 − τ m is a constant.
[Equation]
The above theorem suggests two analysis results.
First, minimizing the representations of the projected layer
∈
Y is useful for encoding the optimal invariant state.
Second, the pixel transformation needs to be regularized with the observation such that the divergence in upper bounded.
Next, we will place emphasis on the illustration of this shift via a reparameterization method.
2. R EPARAMETERIZABLE O BSERVATION
For the bound of Theorem2.1, the Algorithm 1 in appendix is a procedure defined on the MDP M with Gaussian random variables G 0 ∼ G |O| for initialization.
The TRANSFORM is fulfilled by aforementioned pixel transformation ν .
To parameterize the augmentation, we add a shift subject to
G t ( µ t , σ t ) on the bilinear interpolation.
Therefore, the output will have different levels of pixel shift.
2. S TABILIZING R EWARD F UNCTION
Reparametrizing the underlying invariant optimization plays a key role to smooth the distribution D in the replay buffer.
We sample multiple augmentation data following G ( µ, σ ) and then mixup the learned hidden state s ′ for further stabilization.
Theorem 2.2.
(Mixed Shift Normalization.)
Suppose Assumption B.1 hold for functions g ξ , g ¯ ξ , f ξ , and f ¯ ξ , respectively.
For any input x ∼ ˆ D and shifted x ′ ∼ G , the divergence with mixed augmented states can be bound by, where
C = 1+ τ 1 − τ , τ = 1 − τ m .
[Equation]
The proof of Theorem 2.2 is straightforward based on Jensen’s inequality and the Theorem 2.1.
2. A UTOMATIC M ANIPULATED O BSERVATION
With theoretical analysis on invariant transformations, we presented a new framework with normalization variants to ensure our discussed learning guarantees.
Critic oriented.
We switch the shift into a parameterized Gaussian distribution.
The mean and variance will participate in optimizing the objective function of the critic network, where s ( x ′ t ) and s ( x ′ t
+
) are states embedded by encoder and ω represents the transformation.
[Equation]
Hidden representation oriented.
For similarity measurement, we use the loss proposed in BYOL.
ξ and ¯ ξ represent the parameters of online encoder g ξ ◦ f ξ and momentum encoder g ¯ ξ ◦ f ¯ ξ respectively. [Equation]
[Equation]
Batch statistics oriented.
The distribution shift between the transformed samples and the overall data can be measured by the following formulation, where ˜ µ ( x ′ t ) and ˜ σ 2 ( x ′ t ) is the mean and variance corresponding to the l -th convolution layer, X is the overall observations.
[Equation]
Architectural overview.
We summarize the objective function of the transformation below.
α and λ as hyperparameters represent the magnitude of the constraints.
[Equation]
E XPERIMENTS
We compare ShiftNorm with prior model-free methods on 9 visual tasks from the DM control suite. and results are presented in Figure 2.
Below are key findings: (i) Compared to vanilla DDPG and
SAC, ShiftNorm gains outstanding results and outperform in a wide range.
(ii) When compared with DrQ-v2 which has already performed remarkable for continuous control, we improve the sampleefficiency on multiple tasks.
(iii) We also find that ShiftNorm has better stability during training while still keep convincing performance. More details and ablation studies can be found in appendix.
C ONCLUSION
We introduce a simple automatic transformation for model-free RL algorithm for visual continuous control tasks.
We proposed an auxiliary loss to improve the performance and training stability at the same time.
Our method achieves convincing performance without specific architectural selection compared to SOTA approachs on DeepMind control suite.
We hope that our method can promote the progress of automatic augmentation and representation learning in RL.
A M ISSING P ROOFS .
Lemma A.1.
Assume that h : R |S| → R |Y| can be written as h ( ξ ) =
f ( < ξ, s > ) , for some s ∈ R |S| , and f : R |Y| → R |Y| with parameter ξ .
Then, convexity of f implies the convexity of h .
Proof.
Let ξ 1 , ξ 2 ∈ R |S| and τ ∈ [0 , 1] .
We have
[Equation]
where the last inequality follows from the convexity of f .
Lemma A.2.
Given the updating dynamics: ¯ ξ t = (1 − τ m )¯ ξ t − 1 + τ m ξ t .
By Lemma 2.1, f ξ = f ¯ ξ holds after convergence.
As a result, the problem of min E x [ ∥
f ξ ◦ g ξ ( x ′ ) − f ¯ ξ ◦ g ¯ ξ ( x ) ∥ ] is equivalent to the problem of min E x
[ ∥ g ξ ( x ′ ) − g ¯ ξ ( x ) ∥ ] .
Proof.
Through Lemma A.1, we know the convexity of a designed function f can give rise to the convexity of the function h with the parameters as the input.
Therefore, we design our projections f ξ and f ¯ ξ as f ( < ξ, s > ) and f ( < ¯ ξ, s > ) respectively.
For instance, ReLU and MLP can be adopted here.
Using the dynamic: ¯ ξ t = (1 − τ m )¯ ξ t − 1 + τ m ξ t together with h ( ξ ) mentioned in Lemma A.1, we obtain the divergence of hidden representations f ξ ◦ g ξ ( x ′ ) and f ¯ ξ ◦ g ¯ ξ ( x )] ,
[Equation]
where ∥ ξ f ∥ is the parameter of the projection ∥ f ξ ∥ .
The first equality is determined by the approximation of convergence analysis, which is f ξ = f ¯ ξ .
We use Cauchy–Schwarz inequality here.
Note that a premise in this lemma is that the momentum updating reached convergence, which means f ξ = f ¯ ξ .
Minimizing the right side bound equals optimizing the problem of the left side in Eq.(12).
When the norm of ∥ ξ f ∥ is fixed, the proof completes.
Theorem A.1.
(Shift Normalization.)
Suppose Assumption B.1 hold for functions g ξ , g ¯ ξ , f ξ , and f ¯ ξ , respectively.
The updating dynamics is: ¯ ξ t = (1 − τ m )¯ ξ t − 1 + τ m ξ t , τ m ∈ [0 , 1] .
For any input x ∼ ˆ D and shifted x ′ , optimizing the conditional divergence in Definition 2.1 means to minimize the upper bound as follows,
[Equation]
13) where C = 1+ τ 1 − τ , τ = 1 − τ m is a constant.
Proof.
According to the following triangular inequality,
[Equation]
The distance d ( · , · ) is set as l p -norm for simplicity.
By incorporating Lemma A.2 into the left side of
Eq.(14), it leads to a divergence with projections where a new triangular inequality holds,
[Equation]
Set s ′ =
g ξ ( x ) and s = g ¯ ξ ( x ) .
Now we use Lemma A.1 and the updating dynamics for the designed projection f ξ and f ¯ ξ , τ m ∈ [0 , 1] , we can obtain,
[Equation]
[Equation]
By designing a projection that satisfies h (¯ ξ t ) =
f t ¯ ξ ( s ) and h ( ξ t ) =
f tξ ( s ) , we have
[Equation]
The goal is to minimize ϵ -approximation on latent distance d ( f ¯ ξ ( s ) , f ξ ( s ′ )) such that the left side of
Eq.(15) is minimized.
Particularly, given l p -norm as distance d ( · , · ) at the timestep t
[Equation]
and a ReLU network as function f , it leads to,
[Equation]
where τ = 1 − τ m .
L f -Lipschitzness assumption is employed.
Suppose the updating has achieved convergence [Equation] Eq.(18) turns to the following inequality,
[Equation]
On the other hand, the second term on the right side of Eq.(15) can be rewritten as,
[Equation]
Altogether, we substitute Eq.(19) and Eq.(20) in Eq.(15), we finally obtain,
[Equation]
The proof is finished.
Theorem A.2.
(Mixed Shift Normalization.)
Suppose Assumption B.1 hold for functions g ξ , g ¯ ξ , f ξ , and f ¯ ξ , respectively.
The updating dynamics is: ¯ ξ t = (1 − τ m )¯ ξ t − 1 + τ m ξ t .
For any input x ∼ ˆ D and shifted x ′ ∼ G , the divergence with mixed augmented states can be bound by,
[Equation]
where C = 1+ τ 1 − τ , τ = 1 − τ m .
Proof.
Based on Theorem A.
[Equation]
we can view the mixup E x ′ [ x ′ ] as a sort of augmented data
[Equation]
and thereby we have,
[Equation]
Since E x [ ∥ x − E x ′ [ x ′ ] ∥ ] = E x [ ∥ E x ′ [ x − x ′ ] ∥ ] ≤ E x E x ′ [ ∥ x − x ′ ∥ ] , we can finally obtain,
E
[Equation]
(24) which completes the proof.
B T ECHNICAL T OOLS
Assumption B.1.
(Lipschitzness).
Let the encoding function g ( x ) : R |O| → R |S| is L g -Lipschitz, we have that ∀ x , x ′
[Equation]
C R ELATED W ORK
Data augmentation.
Training agents from high-dimensional images data has always been highly concerned (Mnih et al., 2013; Yarats et al., 2019; Hafner et al., 2019; Lee et al., 2019).
Previous work has shown the great potential of solving the visual RL tasks by data augmentation.
Many results show that even the simplest augmentation method can greatly improve the agent’s sample efficiency and asymptomatic performance (Laskin et al., 2020a;
Yarats et al., 2020; 2021) without further modifications to the backbone RL algorithm.
Self-supervised learning.
Recent years have witnessed the self-supervised learning methods achieving huge success in representation learning (Chen et al., 2020a;
He et al., 2020; Caron et al., 2020; Grill et al., 2020; Chen et al., 2020b).
For SSL combined with RL, CURL (Laskin et al., 2020b), ReLIC (Mitrovic et al., 2020) and CoBERL (Banino et al., 2021) focus on compute the consistent between positive samples.
On the other hand, Mazoure et al.
(2020) maximized the mutual information between the nearby states, while PSM (Agarwal et al., 2021) measures behavioral similarity between states through optimal policies similarity with future state transition probability.
D B ENCHMARKS .
We follow the settings in DrQ-v2 (Yarats et al., 2021) and classify these tasks into easy and challenging and provide a summary for each task in Table 1.
[Table]
E P SEUDOCODE
[Table]
[Table]
t + n ← g ξ ( x ′ t ) , g ξ ( x ′ t + n ) 20: s ¯ ξ ← g ¯ ξ ( x t ) 21:
Measure similarity by L ξ, ¯ ξ,ω 22: a t ← π ϕ ( s t + n ) + ˜ ϵ and ˜ ϵ ∼ G (0 , ˜ σ t ) 23:
Compute J θ,ω ( D ) for Q θ and TRANSFORM updating 24: µ ← µ − δ aug ∇ µ ( J θ,ω ( D )) 25: σ ← σ − δ aug ∇ σ ( J θ,ω ( D )) 26: ξ ← ξ − δ ∇ ξ J θ,ω ( D ) 27: θ ← θ − δ ∇ θ J θ,ω ( D ) 28: ˆ θ ← (1 − τ ) ˆ θ + τθ 29: ¯ ξ ← (1 − τ m )¯ ξ + τ m ξ 30: end procedure 31: procedure UpdateActor( D , ˜ σ ) 32: { x t } Ni =1 ∼ D 33: s t ← g ξ ( TRANSFORM ( x t , G t ))
34: a t ← π ϕ ( s t ) + ˜ ϵ and ˜ ϵ ∼ clip ( G (0 , ˜ σ )) 35:
Update the actor using the sampled policy gradient 36: ∇ ϕ J ≈ 1 N (cid:80) i ∇ a Q ( s , a ) | s = s t , a = a t ∇ ϕ π ( s ) | s t 37:
ϕ ← ϕ − δ ∇ ϕ J 38: end procedure
F E XPERIMENTS F. S ETUP
Environments.
The DeepMind control suite Tassa et al.
(2018) is a popular benchmark that has been widely used in prior algorithms.
This benchmark is built on MuJoCo Todorov et al.
and contained several robot control tasks, which provides different difficulties.
Following previous works, we set our observations as stacks of 3 consecutive images.
The size of the
RGB images is 84 × 84, and we concatenate 3 of them at the dimension of the color channel to ensure that the dynamic and temporal information is fitted into the agent.
Details.
We select nine tasks from DM Control with different difficulties (as mentioned by Drq-v2) to test the performance of ShiftNorm both on sample-efficiency and asymptotical performance .
Since each episode is set to a total of 1000 frames in all tasks of DM Control, it is quite reasonable to set the total number of frames experienced during training as x-Axis, so that we can refer to the number of episodes we have gone through to evaluate the tested algorithms.
We will also use the settings in
DM Control to calculate the reward: a per-frame reward is in the unit interval [0 , 1] , so each episode will get an episode reward of no more than 1000.
For a fair comparison, we refer to the settings of the training episodes for different tasks in DrQ-v2, e.g., more episodes will be trained for hard tasks to facilitate asymptotical performance comparison.
Baselines.
We present several baselines, including methods of using data augmentation to improve performance and sample-efficiency to benchmark performance for continuous control on DM Control suite: (i) DrQ-v2 Yarats et al.
(2021) where the authors change the backbone RL algorithm from SAC
Haarnoja et al.
to a better designed DDPG.
In addition, they optimized the details in DrQ
Yarats et al.
(2020), (ii) Pixel SAC and (iii)
Pixel DDPG: Vanilla SAC and DDPG operating purely from images.
Evaluation.
To facilitate fair performance comparison all algorithms will be evaluated with the same periodicity of 20000 environment steps and we average over 10 episodes return for each evaluation query.
we also use environment steps to measure sample complexity for a well-defined comparison with action repeat of 2.
Hyperparameters.
We set the parameters as consistent as possible with the baselines.
To prevent premature convergence, the mean and standard deviation in automatic shift augmentation, the learning rate is 2e-6 and the momentum tau is set to 0.0001.
Both policy and Q-function networks are trained using Adam optimizer, and the batch size is the same as DrQ-v2 of 256.
SAC and DDPG’s parameters are following prior algorithms.
The detailed description of all tasks will be shown in table 1.
F. A CTOR AND C RITIC N ETWORKS
The clipped double Q-learning Van Hasselt et al.
; Fujimoto et al.
(2018) is applied for the critic, where each Q-function is parametrized as a 3-layer MLP with ReLU activations after each layer except of the last.
The actor is also a 3-layer MLP with ReLU s that outputs mean for the action.
The hidden dimension is set to 1024 for both the critic and actor.
F. E NCODER N ETWORKS
The architecture of encoder is based on Yarats et al.
(2019), which has four convolutional layers with 3 × 3 kernels and 32 channels.
The ReLU activation is applied after each conv layer.
We also use BatchNorm Ioffe & Szegedy (2015) after each activations rather than LayerNorm Ba et al.
(2016) after a single fully-connected layer.
The stride for the first conv layer is 2 while 1 for the rest.
BatchNorm is also applied to normalize the fully-connected layer where the output of the convent is feed into.
Finally, the use of tanh nonlinearity and the initialization of weight are consistent with the prior work.
The actor and critic share the same encoder, although the encoder only uses the gradients from the critic for updating.
F. R ESULTS
The curves of 9 complex tasks are in Table
[Table]
F. A BLATION S TUDIES
In this section, we present an ablation study to discuss the effects of different terms in shift loss.
As mentioned in previous sections, the shift loss function is composed of two parts: R ω ( x ′ t ) for normalization with observation statistics, and L ξ, ¯ ξ,ω for similarity measure.
On this basis, we can divide ShiftNorm into 4 versions: (i) with critic .
Transformation is only updated with critic.
(ii) with xstats & critic .
Transformation is updated by critic and R ω ( x ′ t ) together.
(iii) with h-dist & critic .
Transformation is updated by critic and L ξ, ¯ ξ,ω together.
(iv) with all .
Transformation will be updated by all components.
We evaluate all of these versions on 5 tasks from the DeepMind control suite and present the results in Figure 3.
Compared with critic with other methods, we demonstrate that both of the components in shift loss has improved the performance.
Though the version of with critic has the ability to solve most of the tasks and improve the data-efficiency, it may suffer from weak supervision signals and sink into suboptimal cases.
However, once the constraints of shiftnorm are added, the performance has been improved in a degree compared with the version of with critic .
If both terms are used, the performance can lead ahead of all tasks, rather than task-specific.
We can also figure out from the shades that the curve with complete constraints achieves better stability (i.e. smaller variance), which shows that representation learning may be closely related to policy learning.
This result demonstrates that even if the automatic transformation can enable agents to learn features in a more appropriate way, such a process is often uncontrollable.
Introducing simple distribution constraints and similarity learning can greatly alleviate this problem, and will not do damage to the sample-efficiency or asymptotic performance.
O
YPER
[Table] [Table]