Year: 2021

Abstract:
How to learn an effective reinforcement learning-based model for control tasks from high-level visual observations is a practical and challenging problem. A key to solve this problem is to learn low-dimensional state representations from observations, from which an effective policy can be learned. In order to boost the learning of state encoder, recent works proposed to capture behavioral similarities between states representations or apply data augmentation on visual observations. In this paper, we propose a novel meta-learner-based framework for representation learning regarding behavioral similarities in reinforcement learning. Specifically, our framework encodes the high-dimensional observations into two decomposed embeddings regarding reward and dynamics in a Markov Decision Process (MDP). A pair of meta-learners is developed, one of which quantifies the reward similarity and the other of which quantifies dynamics similarity over the correspondingly decomposed embeddings. The meta-learners are self-learned to update the state embeddings by approximating two disjoint terms in on-policy bisimulation metric. To incorporate the reward and dynamics terms, we further develop a strategy to adaptively balance their impacts based on different tasks or environments. We empirically demonstrate that our proposed framework outperforms state-of-theart baselines on several benchmarks, including conventional DM Control Suite, Distracting DM Control Suite and a self-driving task CARLA.

heading: 1 INTRODUCTION
text: Designing effective reinforcement learning algorithms to learn control from high-dimensional visual observations is crucial and has attracted more attentions in the reinforcement learning (RL) community (Yarats et al., 2021; Laskin et al., 2020a; Schwarzer et al., 2021; Lesort et al., 2018). To learn a policy efficiently from high-dimensional observations, prior approaches first learned an encoder to map high-dimensional observations, e.g., images, to low-dimensional representations, and subsequently trained a policy from low-dimensional representations to actions based on various RL algorithms. Therefore, how to learn low-dimensional representations which are able to provide semantic abstraction for high-dimensional observations plays a key role.
Early works on deep reinforcement learning train encoders based on various reconstruction losses (Lange & Riedmiller, 2010; Watter et al., 2015; Wahlström et al., 2015; Higgins et al., 2017), which aim to enforce that the learned low-dimensional representations can be decoded to reconstruct the original high-dimensional observations. Promising results have been achieved in some application domains, such as playing video games, simulated tasks, etc. However, the policy learned on representations of states trained with reconstruction losses may not generalize well to complex environments, which are even though semantically similar to the source environments. The reason is that a reconstruction loss is computed over all the pixels, which results in all the details of the high-dimensional images being preserved in low-dimensional representations. However, some of the observed details, e.g., complex background objects in an image, are task-irrelevant and highly environment-dependent. Encoding such details in representation learning makes the learned representation less effective for a downstream reinforcement learning task of interest and less generalizable to new environments. To address the aforementioned issue, some data augmentation techniques have been proposed to make the learned representations more robust (Yarats et al., 2021; Lee et al., 2020b;
Laskin et al., 2020b). However, these approaches rarely consider the properties of a Markov Decision Process (MDP), such as conditional transition probabilities between states given an action, in the representation learning procedure.
Recent research (Zhang et al., 2020; Agarwal et al., 2021; Castro, 2020; Ferns et al., 2004) has shown that the bisimulation metric and its variants are potentially effective to be exploited to learn a more generalizable reinforcement learning agent across semantically similar environments. The bisimulation metric measures the “behavioral similarity” between two states based on two terms: 1) reward difference that considers the difference in immediate task-specific reward signals between two states, and 2) dynamics difference that considers the similarity of the long-term behaviors between two states. A general idea of these approaches is to learn an encoder to map observations to a latent space such that the distance or similarity between two states in the latent space approximates their bisimulation metric. In this way, the learned representations (in the latent space) are task-relevant and invariant to task-irrelevant details. However, manually specifying a form of distance, e.g., the L1 norm as used in (Zhang et al., 2020), in the latent space may limit the approximation precision for the bisimulation metric and potentially lose parts of the state features that are useful for policy learning. Moreover, existing approaches rarely explore how to learn a linear combination of reward and dynamics differences in the bisimulation metric, which may vary in different tasks or environments.
We propose a novel framework for learning generalizable state representation in RL, called Adaptive Meta-learner of Behavioral Similarities (AMBS). In this framework we design a network with two encoders that map the high-dimensional observations to two decomposed representations regarding the two properties of MDP, rewards and dynamics. For the purpose of learning behavioral similarity in state representations, we introduce a pair of meta-learners that learn similarities in order to measure the reward and dynamics similarity between two states over the corresponding decomposed state representations, respectively. The meta-learners are self-learned by approximating the corresponding terms in the bisimulation metric, containing the reward difference and the dynamics difference. Then the meta-learners update the state representations according to their behavioral distance to the other state representations. Unlike previous approaches with a hand-crafted form of distance/similarity evaluating state encoding in L1 space, our learned similarities measure two state representations via a neural architecture, where more side information can be preserved for further policy learning. Our experiments also observe that a smaller loss can be obtained by using the learned similarity. This demonstrates that the proposed meta-learners can overcome the approximation precision issue and provide more stable gradients for robust learning of state representations. Moreover, we explore the impact between reward and dynamics terms in the bisimulation metric. We propose a learning-based adaptive strategy to balance the effect of reward and dynamics in different tasks or environments by introducing a learnable importance parameter, which is jointly learned with the state-action value function. Finally, we use a simple but effective data augmentation strategy to accelerate the RL procedure and help learning more robust state representations.
The main contributions of our work are 3-fold: 1) we propose a meta-learner-based framework to learn task-relevant and environment-details-invariant state representations; 2) we propose a network architecture that decomposes each state into two different types of representations for measuring the similarities in terms of reward and dynamics, respectively, and design a learnable adaptive strategy to balance them to estimate the bisimulation metric between states; 3) we verify our proposed framework on extensive experiments and demonstrate new state-of-the-art results on backgrounddistraction DeepMind control suite (Tassa et al., 2018; Zhang et al., 2018; Stone et al., 2021) and other visual-based RL tasks.

heading: 2 RELATED WORK
text: Representation learning for reinforcement learning from pixels Various existing deep RL methods have been proposed to address the sample-efficiency and generalization problem of RL algorithms learning from pixel observation. In end-to-end deep RL, neural networks learn representations implicitly by optimizing some RL objective (Mnih et al., 2015; Espeholt et al., 2018). Watter et al. (2015) and Wahlström et al. (2015) proposed to learn a encoder with training the dynamics model jointly to produce observation representations. Lee et al. (2020a), Hafner et al. (2019), Hafner et al. (2020) and Zhang et al. (2019) aimed to learn an environment dynamics model with reconstruction loss to compact pixels to latent representation. Gelada et al. (2019) proposed to learn representations
by predicting the dynamics model along with the reward, and analyzed its theoretical connection to the bisimulation metric. Agarwal et al. (2021) learns representations by state distances based on policy distribution. Kemertas & Aumentado-Armstrong (2021) modified the learning of bisimulation metric with intrinsic rewards and inverse dynamics regularization. Castro et al. (2021) replaced the Wasserstein distance in bisimulation metric with a parametrized metric.
Data Augmentation in reinforcement learning Laskin et al. (2020b) and Ye et al. (2020b) explored the various data augmentation techniques from computer vision, which are limited to the transformation on input images. Cobbe et al. (2019) applied data augmentation on domain transfer. Ye et al. (2020a) studied data augmentation on game environments for zero-shot generalization. RandConv (Lee et al., 2020b) proposes a randomized convolutional neural network to generate randomized observations in order to perform data augmentation. A recent work DrQ (Yarats et al., 2021) suggests to perform random crop on image observations and provides regularized formulation for updating Q-function. Raileanu et al. (2021) proposed to use data augmentation for regularizing on-policy actor-critic RL objectives. Pitis et al. (2020) proposed a counterfactual data augmentation technique by swapping observed trajectory pairs.

heading: 3 PRELIMINARIES
text: We define an environment as a Markov Decision Process (MDP) described by a tuple M = (S,A,P,R, γ), where S is the high-dimensional state space (e.g., images), A is the action space, P(s′|s,a) is the transition dynamics model that captures the probability of transitioning to next state s′ ∈ S given current state s ∈ S and action a ∈ A, R is the reward function yielding a reward signal r = R(s,a) ∈ R, and γ ∈ [0, 1) is the discounting factor. The goal of reinforcement learning is to learn a policy π(a|s) that maximizes the expected cumulative rewards: maxπ E[ ∑ t γ
tr(st,at)|at ∼ π(·|st), s′t ∼ P(s′|s,a)]. In the scope of this paper, we do not consider partial observability, and use use stacked consecutive frames as the fully observed states.
Soft Actor-Critic (SAC) is a widely used off-policy model-free reinforcement learning algorithm (Haarnoja et al., 2018). SAC aims to maximize the entropy objective (Ziebart, 2010) which is the reinforcement learning objective augmented with an entropy term J(π) =∑ t E [r(st,at) + αH (π(·|st))]. In order to maximize the objective, SAC learns a state-value function Qθ(s,a), a stochastic policy πψ(a|s) and the temperature α, where θ and ψ are the parameters of Qθ and πψ , respectively. The Q-function is trained by minimizing the squared soft Bellman residual
JQ(θ) = E(st,at)∼D [ 1
2
( Qθ(st,at)− ( r(st,at) + γEst+1∼P(·|st,at)[Vθ̄(st+1)] ))2] , (1)
where D is the dataset or replay buffer storing the transitions, and Vθ̄ is the value function parameterized by θ̄. The parameters ψ of policy is learned by maximizing the following objective
Jπ(ψ) = Est∼D [ Eat∼πψ [α log(πψ(a|s))−Qθ(st,at)] ] . (2)
Bisimulation Metric defines a pseudometric d : S × S → R,1 where d quantifies the behavioral similarity of two discrete states (Ferns et al., 2004). An extension of bisimulation metric to both continuous and discrete state spaces has also been developed (Ferns et al., 2011). An variant of the bisimulation metric proposed in Castro (2020) defines a metric with respect to a policy π, which is known as the on-policy bisimulation metric. It removes the requirement of matching actions in the dynamics model but focuses on the policy π, which is able to better capture behavioral similarity for a specific task. Because of this property, in this paper, we focus on the π-bisimulation metric (Castro, 2020):
Fπ(d)(si, sj) = (1− c)|Rπsi −R π sj |+ cW1(d)(P π si ,P π sj ). (3) where Rπs := ∑ a π(a|s)R(s,a), Pπs := ∑
a π(a|s)P(·|s,a), mapping Fπ : met → met, met denotes the set of all pseudometrics on state space S and W1(d)(·, ·) denotes the 1-Wasserstein distance given the pseudometric d. Then Fπ has a least fixed point denoted by d∗. Deep Bisimulation
1If the pseudometric d of two states is 0, then the two states belong to an equivalence relation.
for Control (DBC) (Zhang et al., 2020) proposed to learn latent representations such that the L1 distances in the latent space equal to the π-bisimulation metric in the state space:
argmin ϕ
( ∥ϕ(si)− ϕ(sj)∥1 − |Rπsi −R π sj | − γW2(d)(P π si ,P π sj ) )2 ,
where ϕ is the state encoder, and W2(d)(·, ·) denotes the 2-Wasserstein distance. DBC is combined with the reinforcement learning algorithm SAC where ϕ(s) is the input for SAC.

heading: 4 ADAPTIVE META-LEARNER OF BEHAVIORAL SIMILARITIES
text: In this section, we propose a framework named Adaptive Meta-learner of Behavioral Similarities (AMBS) to learn generalizable states representation regarding π-bisimulation metric. The whole learning procedure of AMBS is demonstrated in Figure 1. Observe that the π-bisimulation metric is composed of two terms: |Rπsi −R π sj |, which computes the difference of reward signals between states, and W2(d)(Pπsi ,P π sj ), which computes the difference of the outputs of dynamics model between states. We propose a network architecture for AMBS, where encoders transform the highdimensional visual observation to two decomposed encodings regarding reward and dynamics. We develop two meta-learners, one of which quantifies the reward difference on reward representations and the other one captures dynamics distance (Section 4.1). Each meta-learner is self-learned to update the corresponding state representations, by learning to approximate a difference term in the π-bisimulation metric, respectively. Rather than enforcing the L1 distance between embedded states equal to the π-bisimulation metric, our meta-learners are able to use more flexible form of similarities, i.e., a well-designed non-linear neural network, with each similarity evaluates the reward difference or dynamics difference between states beyond the original Euclidean space. Moreover, we propose a strategy for learning to combine the outputs of the two learned similarities in a specific environment (Section 4.2). We introduce a learnable weight for the combination and such weight is adaptively learned together with the policy learning procedure (in this work we use SAC as the base reinforcement learning algorithm). In addition, we also use a data augmentation technique to make the learned representations more robust (Section 4.3). The whole learning procedure is trained in an end-to-end manner.

heading: 4.1 META-LEARNERS FOR DECOMPOSED REPRESENTATIONS
text: We design a network architecture with two encoders, ϕr and ϕd, to encode decomposed features from visual observations, as shown in the left side of Figure 1. Each encoder maps a high-dimensional observation to a low-dimensional representation, i.e., ϕr : S → Zr capturing reward-relevant features and ϕd : S → Zd capturing dynamics-relevant features. We design two meta-learners, fr and fd, where fr learns to measure reward difference and fd aims to measure dynamics difference between two state representations. Specifically, each meta-learner takes two embedded states as input and outputs the corresponding similarity. The procedure is summarized as follows.
si, sj → ϕ∗(si), ϕ∗(sj) → f∗(ϕ∗(si), ϕ∗(sj)), where ∗ ∈ {r, d}.
Note that the aggregation of the outputs of the two encoders ϕr and ϕd of each observation (i.e., state) will also be fed into SAC to learn a policy (the right side of Figure 1). Each self-learned meta-learner learns to capture similarity by approximating a distance term in the π-bisimulation
metric respectively and provides stable gradients for updating ϕr or ϕd according to the distance to the other state representations. The details of network architecture can be found in Appendix B.2.
As discussed in Section 1, restricting the form of distance to be the L1/L2 norm in the latent space may limit the approximation precision. As shown in Figure 2, L1 and L2 norm for measuring distance of two latent representation ϕ(si) and ϕ(sj) lead to large regression losses during RL training, which destabilize the representation learning and consequently decrease final performance. Besides, such hand-crafted distances make semantically similar states encoding representation close to each other in terms of the L1 norm in the Euclidean space, which however may lose part of useful information.
To overcome aforementioned limitations, we propose to exploit meta-learner fr and fd to learn similarities. By learning with a regression target, the f∗ learned similarity is easier to converge to the target and leads to faster descent tendency of regression loss (as shown in Figure 2 where y-axis is log-scale). Such property provides more stable gradients for state representation comparing to the L1/L2 norm, and therefore the meta-learner f∗ is able to guide the process of updating the state encoder. Besides, f∗ is a non-linear transformation that evaluates the state similarity in a more flexible space instead of the original Euclidean space, consequently it is able to preserve more task-relevant information in state representation that is required for further SAC policy learning.
As the goal of meta-learners is to approximate different types of similarities in the π-bisimulation metric, we design the loss functions for the meta-learners using the mean squared error:
ℓ(fr, ϕr) = ( fr(ϕr(si), ϕr(sj))− |Rπsi −R π sj | )2 , (4)
ℓ(fd, ϕd) = ( fd(ϕd(si), ϕd(sj))−W2(Pπsi ,P π sj ) )2 . (5)
To make the resultant optimization problems compatible with SGD updates with transitions sampled from replay buffer, we replace the reward metric term |Rπsi −R π sj | by the distance of rewards in the sampled transi-
tions in (4), and use a learned parametric dynamics model P̂ to approximate the true dynamics model Pπs in (5). We also use 2-Wasserstein distance W2 in (5) (as in DBC) because it has close form for Gaussian distribution. With denoting by (si,ai, ri, s′i) and (sj ,aj , rj , s ′ j) two transitions sampled from the replay buffer, the loss functions (4) and (5) can be revised as follows,
ℓ(fr, ϕr) = (fr(ϕr(si), ϕr(sj))− |ri − rj |)2 , (6) ℓ(fd, ϕd) = ( fd(ϕd(si), ϕd(sj))−W2(P̂(·|ϕd(si),ai), P̂(·|ϕd(sj),aj)) )2 , (7)
where P̂(·|ϕd(s∗),a∗) is a learned probabilistic dynamics model, which is implemented as a neural network that takes the dynamics representation ϕd(s∗) and action a∗ as input, and predicts the distribution of dynamics representation of next state s′∗. The details of P̂ can be found Appendix A.

heading: 4.2 BALANCING IMPACT OF REWARD AND DYNAMICS
text: As the π-bisimulation metric (3) is defined as a linear combination of the reward and the dynamics difference, an intuitive approach to combine the two terms is to adapt the difference values of combination weight in difference environments and tasks. Rather than considering the combination weight as a hyper-parameter, which is tuned in advance, we introduce a learnable parameter c ∈ (0, 1) to balance the impact of the reward and the dynamics.
The impact c should be automatically determined when learning in a specific task or environment. We integrate the learning of c into the update of the Q-function in SAC such that the value of c is learned from the state-action value which is highly relevant to a specific task and environment. To be specific, we concatenates the low-dimensional representations ϕr(s) and ϕd(s) weighted by 1− c and c, where 1− c and c are output of softmax function to ensure c ∈ (0, 1). The loss for Q-function is revised as the Q-function takes input of ϕr(s), ϕd(s), weight c and action a:
JQ(θ, c, ϕr, ϕd) = E(s,a,s′)∼D [ 1
2
( Qθ(ϕr(s), ϕd(s), c,a)− ( r(s,a) + γV (s′) ))2] , (8)
where Qθ is the Q-function parameterized by θ and V (·) is the target value function.
Moreover, to balance the learning of two meta-learners fr and fd, we jointly minimize the regression losses of fr and fd ((6) and (7), respectively) with combination weight c. The loss function is modified as follows,
ℓ(Θ) = (1− c)ℓ(fr, ϕr) + cℓ(fd, ϕd), where Θ = {fr, fd, ϕr, ϕd}. (9)
Note that parameter c is already learned along with Q-function, we stop gradient of cwhen minimizing the loss (9). This is because if we perform gradient descent w.r.t. c, it may only capture which loss (ℓ(fr, ϕr) or ℓ(fd, ϕd)) is larger/smaller and fail to learn the quantities of their importance. We also stop gradient for the dynamics model P̂ since P̂ only predicts one-step dynamics.

heading: 4.3 OVERALL OBJECTIVE WITH DATA AUGMENTATION
text: In this section, we propose to integrate a data augmentation strategy into our proposed framework to learn more robust state representations. We follow the notations of an existing work on data augmentation for reinforcement learning, DrQ (Yarats et al., 2021), and define a state transformation h : S × T → S that maps a state to a data-augmented state, where T is the space of parameters of h. In practice, random crop is an efficient transformation for data augmentation. Therefore, we use random crop as the transformation h in this work. Then T contains all possible crop positions and v ∈ T is one crop position drawn from T . We apply DrQ to regularize the objective of the Q-function (see Appendix for the full objective of Q-function).
Besides, by using augmentation transformation h, we rewrite the loss of representation learning and similarity approximation in (9) as follows,
ℓ(Θ) =(1− c) ( fr ( ϕr(s (1) i ), ϕr(s (1) j ) ) − |ri − rj | )2 +c ( fd ( ϕd(s (1) i ), ϕd(s (1) j ) ) −W2 ( P̂(·|ϕd(s(1)i ),ai), P̂(·|ϕd(s (1) j ),aj)
))2 +(1− c) ( fr ( ϕr(s (2) j ), ϕr(s (2) i ) ) − |ri − rj |
)2 +c ( fd ( ϕd(s (2) j ), ϕd(s (2) i ) ) −W2 ( P̂(·|ϕd(s(1)i ),ai), P̂(·|ϕd(s (1) j ),aj) ))2 ,
(10)
Algorithm 1 AMBS+SAC 1: Input: Replay Buffer D, initialized Θ = {fr, ϕr, fd, ϕd} , Q
networkQθ , actor piψ , target Q networkQθ̄ , 2: Sample a batch with sizeB: {(si, ai, ri, s′i)}Bb=1 ∼ D. 3: Shuffle batch {(si, ai, ri, s′i)}Bb=1 to {(sj , aj , rj , s′j)}Bb=1. 4: Update Q network by (8) with DrQ augmentation. 5: update actor network by (11). 6: update alpha α. 7: update encoder Θ by (10) 8: update dynamics model P̂ .
where s(k)∗ = h(s∗,v (k) ∗ ) is the transformed state with parameters v(k)∗ . Specifically, s (1) i , s (1) j , s (2) i and s (2) j are transformed states with parameters of v(1)i , v (1) j , v (2) i and v (2) j respectively, which are all drawn from T independently. The first two terms in (10) are similar to the two terms in (9), while the observation s∗ in (9) is transformed to s (1) ∗ . For the last two terms in (10), the observation is transformed by a different parameter v(2)∗ except for the last term, where the parameter v (1) ∗ does not change. The reason why the parameter v(1)∗ is used in last term is that we expect the meta-learner fd is able to make a consensus prediction on s(1)∗ and s (2) ∗ as they are transformed from a same state s∗. Another major difference between first two terms and last two terms is the order of the subscripts i and j: i is followed by j in first two terms while j is followed by i in last two terms. The reason behind this is that we aim to make fr and fd to be symmetric.
The loss of the actor of SAC comprises of encoder is Jπ(ψ) = E(s,a)∼D [ α log(πψ(a|ϕ̂(s)))−Qθ(ϕ(s),a) ] , (11)
where ϕ̂ denotes the convolutional layers of ϕ and Qθ(ϕ(s),a) is a convention form of Q-function with c-weighted state representations input. We also stop gradients for ϕ̂ and ϕ (ϕr and ϕd) regarding the actor loss. In other words, loss (11) only calculates the gradients w.r.t. ψ.
The overall learning algorithm is summarized in Algorithm 1. Note that due to limit of space, we do not present all the details here. More details can be found in Appendix A.
The architecture of AMBS is illustrated in Figure 1. We share the convnet for encoders ϕr and ϕd. The shared convnet consists of four convolutional layers. The final output of convnet is fed into two branches: 1) one is a fully-connected layer to form reward representation ϕr(s), and 2) the other one is also a fully-connected layer but form dynamic representation ϕd(s). The meta-learners fr and fd are both MLPs with two layers. More implementation details of AMBS are provided at Appendix B.

heading: 5 EXPERIMENTS
text: The major goal of our approach is to learn generalizable representation for RL control from pixels. We evaluate our approach on two environments: 1) the DeepMind Control (DMC) suite (Tassa et al., 2018) with background distraction; and 2) autonomous driving task on CARLA (Dosovitskiy et al., 2017) simulator. Serval baselines methods are to be compared with our approach: 1) Deep Bisimulation Control (DBC) (Zhang et al., 2020) which is a recent research on RL representation learning by using L1 distances to approximate bisimulation metric; 2) PSEs (Agarwal et al., 2021) which proposes a new state-pair metric called policy similarity metric for representation learning and learn state embedding by incorporating a contrastive learning method SimCLR (Chen et al., 2020); 3) DrQ (Yarats et al., 2021) which is a state-of-the-art data augmentation method on deep RL; 4) Stochastic Latent Actor-Critic (SLAC) (Lee et al., 2020a) which is a state-of-the-art approach for partial observability on controlling by learning a sequential latent model with reconstruction loss; 5) Soft Actor-Critic (SAC) (Haarnoja et al., 2018) which is a commonly used off-policy deep RL method that above baselines and our approach build upon.

heading: 5.1 DMC SUITE WITH BACKGROUND DISTRACTION
text: DeepMind Control (DMC) suite (Tassa et al., 2018) is an environment that provides high dimensional pixel observation for RL tasks. DMC suite with background distraction (Zhang et al., 2018; Stone et al., 2021) replaces the background with natural videos which are task-irrelevant to the RL control tasks. We perform the comparison experiments on two settings, original background and nature video background. For each setting we evaluate our approach and baselines methods on four environments, cartpole-swingup, finger-spin, cheetah-run and walker-walk. Results of additional environments are shown in Appendix D.
Original Background. Figure 3a shows the DMC observation of original background which is static and clean. The experiment result is shown in Figure 4. Our method AMBS has comparable learning efficiency and converge scores comparing to state-of-the-art pixel observation RL method DrQ and SLAC. However, DBC performs worse in original background setting.
Natural Video Background. Figure 3b shows the DMC observation with natural video background. The background video is sampled from Kinetics (Kay et al., 2017) dataset under label “driving car”. By following the experiment configuration of DBC, we sample 1000 continuous frames as background video for training, and sample another 1000 continuous frames for evaluation. The experiment result
is shown in Figure 5. Our method AMBS outperforms other baseline methods on both efficiency and scores. It converges to the scores that are similar to the scores on original background setting within 1 million steps. DrQ can only converges to the same scores on cartpole-swingup and finger-spin but the learning is slightly slower than AMBS. The sequential latent model method SLAC performs bad on this setting. This experiment demonstrates that AMBS is robust to the background taskirrelevant information and the AMBS objectives improve the performance when comparing to other RL representation learning methods, e.g., DBC.

heading: 5.2 ABLATION STUDY
text: Our approach AMBS consists of several major components: meta-learner, balancing impact of reward and dynamics, and data augmentation on representation learning. To evaluate the importance of each component, we eliminate or replace each component by baseline methods. We perform the experiments on cheetah-run and walker-walk environments under natural video background setting. Figure 6 shows the performances of AMBS, variants of AMBS and baseline methods. AMBS w/o data aug is constructed by removing the data augmentation in our method. AMBS w/o f is that we replace the meta-learners fr and fd by L1 distance, as DBC use L1 distance to measure distance of representation. AMBS w/o c removes the component of balancing impact of reward and dynamics but encodes the observation into a single embedding. It uses single meta-learner to approximate the sum of reward and dynamics distance. DBC + Drq is DBC with data augmentation applied on representation learning and SAC. AMBS c = 0.5 is our method with fixing weight c to 0.5. AMBS w/ (1, γ) replace the weight c in AMBS by freezing reward term weight to 1 and dynamic weight to γ. Such weight setting (1, γ) is used in DBC. Figure 6 demonstrates that our full model performs better than any its variant. Comparing to DBC that does not utilize data augmentation, our variant AMBS w/o data aug still outperforms it. This comparison shows that using meta-learners fr and fd and decomposing states into reward and dynamics representations help to improve RL performance than DBC that uses L1 distance to approximate the whole π-bisimulation metric.

heading: 5.3 TRANSFER OVER REWARD FUNCTIONS
text: The environments walker-walk, walker-run and walker-stand share the same dynamics but have different reward functions according to the moving speed. To evaluate the transfer ability of our method over reward functions, we transfer the learned model from walker-walk to walker-run and walker-stand. We train agents on walker-run and walker-stand with frozen convolutional layers of encoders that are well trained in walker-walk. These experiments are done under natural video setting and are compared with baseline DBC. The results are shown in Figure 7. It shows that the transferring encoder converges faster than training from scratch. Besides our approach has a better transfer ability than DBC.

heading: 5.4 GENERALIZATION OVER BACKGROUND VIDEO
text: We evaluate the generalization ability of our method when background video is changed. We follow the experiment setting of PSE (Agarwal et al., 2021) for fair comparison. We pick 2 videos, ‘bear’ and ‘bmx-bumps’, from training set of DAVIS 2017 (Pont-Tuset et al., 2018) to form the training set of background videos. We randomly sample one video and randomly sample a frame at each episode start, then play the video forward and backward until episode end. The RL agent is evaluated on validation environment where the background video is sampled from validation set of DAVIS 2017. Those validation videos are unseen during the training. The evaluation scores at 500K environment interaction steps is shown in Table 1. The comparison method DrQ+PSEs is data augmented by DrQ. Our method outperforms DrQ+PSEs on all of four tasks. It demonstrates that our AMBS is generalizable on unseen background video.
5.5 AUTONOMOUS DRIVING TASK ON CARLA
CARLA is an autonomous driving simulator that provides a 3D simulation for realistic on-road scenario. In real world cases or in realistic simulation, the learning of RL agents may suffer from complex background that contains task-irrelevant details. To argue that our AMBS can address this issue, we perform experiments on CARLA simulator. We follow the setting of DBC to create an autonomous driving task on map “Town04” for controlling a car to drive as far as possible in 1000 frames. The environment contains a highway with 20 moving vehicles. The reward functions prompts the driving distance and penalizes
collisions. The observation shape is 84 × 420 pixels which consists of five 84 × 84 cameras. Each camera has 60 degrees view, together they produce a 300 degrees first-person view, as shown in Figure 8. The weather, clouds and brightness may change slowly during the simulation. Figure 9 shows the training curves of AMBS and other comparison methods. Our AMBS perform the best over other methods. It is empirical evidence that AMBS has high potential to handle the generalization over task-irrelevant details, e.g., weather and clouds, in real world scenarios.

heading: 6 CONCLUSION
text: This paper proposes a novel framework with meta-learners for learning task-relevant and environmentdetail-invariant state representation in deep RL. In this framework, we propose state encoders that decompose the observation into rewards and dynamic representations for measuring behavioral similarities by self-learned meta-learners, and design a learnable strategy to balance two types of representations. Our framework archives the new state-of-the-art results in various environments and shows its transfer ability on RL tasks. Experimental results show that our framework is able to effectively learn generalizable state representation.

heading: A OBJECTIVES OF ADAPTIVE META-LEARNER OF BEHAVIORAL SIMILARITIES (AMBS)
text: We denote two encoders ϕr and ϕd, where each encoder maps a high-dimensional observation to a low-dimensional representation. The encoder ϕr : S → Zr captures reward-relevant features and the encoder ϕd : S → Zd captures dynamics-relevant features, where Zr and Zd are representation spaces. We design two meta-learners fr : Zr × Zr → R and fd : Zd × Zd → R to output reward
difference and dynamics difference, respectively. We follow DrQ (Yarats et al., 2021) to define the state transformation h : S × T → S that maps a state to a data-augmented state, where T is the space of parameters of h. In practice, we use random crop as transformation h.
With augmentation transformation h, the loss of representation learning is, ℓ(Θ) =(1− c) ( fr(ϕr(s (1) i ), ϕr(s (1) j ))− |ri − rj | )2 + c ( fd(ϕd(s (1) i ), ϕd(s (1) j ))−W2(P̂(·|ϕd(s (1) i ),ai), P̂(·|ϕd(s (1) j ),aj))
)2 + (1− c) ( fr(ϕr(s (2) j ), ϕr(s (2) i ))− |ri − rj |
)2 + c ( fd(ϕd(s (2) j ), ϕd(s (2) i ))−W2(P̂(·|ϕd(s (1) i ),ai), P̂(·|ϕd(s (1) j ),aj)) )2 ,
(12)
where Θ = {fr, fd, ϕr, ϕd}, s(l)∗ = h(s∗,v(l)∗ ) is the transformed state with parameters v(l)∗ . Specifically, s(1)i , s (1) j , s (2) i and s (2) j are transformed states with parameters of v (1) i , v (1) j , v (2) i and v (2) j respectively, which are all drawn from T independently. P̂(·|ϕd(s∗),a∗) is a learned probabilistic dynamics model, which outputs a Gaussian distribution over Zd for predicting dynamics representation of next state s′∗. W2 denotes the 2-Wasserstein distance which has close-form for Gaussian distributions: W2(N (µi,Σi),N (µj ,Σj))2 = ∥µi − µj∥22 + ∥Σ 1 2 i −Σ 1 2 j ∥2F , where µ∗ ∈ R|Zd| are the mean vectors of Gaussian distribution, Σ∗ ∈ R|Zd| are the diagonals of covariance matrices, and ∥ · ∥F is Frobenius norm. Note that we stop gradients for c as mentioned in Section 4.2. The adaptive weight c and (1 − c) are the softmax output of corresponding learnable parameter ηc ∈ R2. It can be formulated as
[c : (1− c)] = softmax(ηc). (13) where the operation : denotes concatenation of two vectors/scalars,
We integrate the learning of c into the update of Q-function of SAC. The objective of the Q-function with DrQ augmentation and clipped double trick is
JQ(θk, ηc) = E(s,a,s′)∼D [ 1
2
( Qθk([(1− c)ϕr(s(1)) : cϕd(s(1))],a)− (r(s,a) + γV (s′)) )2 + 1
2
( Qθk([(1− c)ϕr(s(2)) : cϕd(s(2))],a)− (r(s,a) + γV (s′)) )2] , (14)
where s(l) = h(s,v(l)), v(l) are all drawn from T independently, and {Qθk}2k=1 are the double Q networks. The target value function with DrQ augmentation is
V (s′) = 1
2 ( min k=1,2 Qθ̄k ( [(1− c̄)ϕ̄r(s′(1)) : c̄ϕ̄d(s′(1))],a′(1) ) + α log πψ(a ′(1)|s′(1))
+ min k=1,2 Qθ̄k
( [(1− c̄)ϕ̄r(s′(2)) : c̄ϕ̄d(s′(2))],a′(2) ) + α log πψ(a ′(2)|s′(2)) ) , (15)
where s′(l) = h(s′,v′(l)), a′(l) ∼ πψ(a′(l)|s′(l)), θ̄k is the set of parameters of the target Q network, ϕ̄r and ϕ̄d are target encoders specifically for the target Q network and c̄ is the adaptive weight for target encoders. The parameters θ̄k, ϕ̄r, ϕ̄d and η̄c are softly updated.
The loss of actor of SAC is, Jπ(ψ) = E(s)∼D [ Ea∼πψ [α log(πψ(a|ϕ̂(s(1))))− min
k=1,2 Qθk(ϕ(s
(1)),a)] ] . (16)
The loss of α of SAC is, J(α) = E(s)∼D [ Ea∼πψ [α log(πψ(a|ϕ̂(s(1))))− αH̄] ] , (17)
where H̄ ∈ R is the target entropy hyper-parameter which in practice is H̄ = −|A|.
The loss for updating of dynamic model P̂ is
ℓ(P̂) = E(s,a,s′)∼D (ϕd(s′(1))− µ(P̂(·|ϕd(s(1)),a)) 2σ(P̂(·|ϕd(s(1)),a)) )2 , (18) where µ(P̂(·)) and σ(P̂(·)) are mean and standard deviation of P̂ output Gaussian distribution, respectively.
Algorithm 2 shows the algorithm at each learning step.
Algorithm 2 AMBS+SAC 1: Input: Replay Buffer D, initialized Θ = {fr, ϕr, fd, ϕd} , Q network Qθk , actor piψ, target Q
network Qθ̄k , 2: Sample a batch with size B: {(si,ai, ri, s′i)}Bb=1 ∼ D. 3: Shuffle batch {(si,ai, ri, s′i)}Bb=1 to {(sj ,aj , rj , s′j)}Bb=1. 4: Update Q network by equation 14 with DrQ augmentation. 5: Update actor network by equation 16 . 6: Update alpha α by equation 17 . 7: Update encoder Θ by equation 12 8: Update dynamics model P̂ by equation 18. 9: Softly update target Q network: θ̄k = τQθk + (1− τQ)θ̄k .
10: Softly update target encoder: ϕ̄ = τϕϕ+ (1− τϕ)ϕ̄ . 11: Softly update target c̄: η̄c = τϕηc + (1− τϕ)η̄c .
B IMPLEMENTATION DETAILS

heading: B.1 PIXELS PROCESSING
text: We stack 3 consecutive frames as observation, where each frame is 84×84 RGB images in DeepMind control suite. Each pixel is divided by 255 to down scale to [0, 1]. We consider the stacked frames as fully observed state.

heading: B.2 NETWORK ARCHITECTURE
text: We share the convnet for encoder ϕr and encoder ϕd. The shared convnet consists of four convolutional layers with 3× 3 kernels and 32 output channels. We set stride to 1 everywhere, except for the first convolutional layer, which has stride 2. We use ReLU activation for each convolutional layer output. The final output of convnet is fed into two branches: 1) one is a fully-connected layer with 50 dimensions to form reward representation ϕr(s), and 2) the other one is also a fully-connected layer with 50 dimensions but form dynamic representation ϕd(s). The meta-learners fr and fd are both MLPs with two layers with 50 hidden dimensions.
The critic network takes ϕr(s), ϕd(s) and action as input, and feed them into three stacked fullyconnected layers with 1024 hidden dimensions. The actor network takes the shared convnet output as input, and feed it into four fully-connected layers where the first layer has 100 hidden dimensions and other layers have 1024 hidden dimensions. The dynamic model are MLPs with two layers with 512 hidden dimensions. ReLU activations are used in every hidden layers.

heading: B.3 HYPERPARAMETERS
text: The hyperparameters used in our algorithm are listed in Table 2.

heading: B.4 ACTION REPEAT FOR DMC
text: We use different action repeat hyper-parameters for each task in DeepMind control suite, which are listed in Table 3.

heading: B.5 SOURCE CODES
text: We will open the source codes to public once this paper is accepted.

heading: C ANALYSIS

heading: C.1 REGRESSION LOSSES OF APPROXIMATING BISIMULATION METRIC
text: We compare the regression losses of approximating the bisimulation metric over RL environment steps among DBC and our AMBS. DBC uses L1 norm to calculate the distances between two states embeddings. AMBS performs meta-learners on state representations to measure the distance with learned similarities. Figure 10 demonstrates that AMBS has smaller regression loss and also faster descent tendency compared to DBC. Meta-learners in AMBS are able to provide more stable gradients to update the state representations.
C.2 VISUALIZATIONS OF COMBINATION WEIGHT C
Figure 11 shows the values of combination weight c in AMBS trained on DMC with origin background and Figure 11 is on DMC with video background setting. The values of c at 1M steps vary in different tasks. The common trend is that they increase at the beginning of training. In original background setting c changes slowly after the beginning. In natural video background setting , c has a fast drop after the beginning. At the beginning, agent learns a large weight c for dynamics feature that may boost the learning of RL agent. Then it drops for natural video background setting because it learns that the video background is distracting the agent. Lower weight is learned for the dynamic features.

heading: C.3 NORMS OF STATE EMBEDDINGS
text: We record the L1 norms of reward representation ϕr(s) and dynamic representation ϕd(s) during training. We record 1nr ||ϕr(s)||1 for reward representation where nr is the number of dimensions of ϕr(s), and 1nd ||ϕd(s)||1 for dynamic representation where nr is the number of dimensions of ϕd(s). Figure 13 shows the values average on each sampled training batch. Although the L1 norm of ϕd(s) decreases during training, it converges around 0.4-0.5. The L1 norm of ϕr(s) also decreases in the training procedure. It can verify that dynamic representation ϕd(s) will not converge to all zeros.

heading: C.4 SHARING ENCODERS BETWEEN ACTOR AND CRITIC
text: In AMBS, we share only the convolutional layers ϕ̂ of encoders between actor and critic network. We compare to sharing the whole encoders ϕr and ϕd in this experiment. Figure 14 shows the training curves on walker-walk with natural video background. Sharing full encoders ϕr and ϕd is worse than sharing CNN ϕ̂ only.

heading: D ADDITIONAL EXPERIMENTAL RESULTS

heading: D.1 ADDITIONAL RESULT OF DMC SUITE WITH AND W/O BACKGROUND DISTRACTION
text: Figure 15 shows the training curves on original background setting. Figure 16 shows the training curves on Kinetics (Kay et al., 2017) video background setting.

heading: D.2 ADDITIONAL RESULT OF GENERALIZATION OVER BACKGROUND VIDEO
text: Table 4 shows the scores on DAVIS video background setting comparing to PSEs (Agarwal et al., 2021).

heading: D.3 ADDITIONAL RESULT ON CARLA
text: The CARLA experiment is to learn an autonomous driving RL agent to avoid collisions with other moving vehicles and drive as far as possible on a highway within 1000 frames. Those vehicles are randomly generated, including vehicles’ types, colors and initial positions. We raise the number of generated vehicles from 20 to 40 so that the road becomes more crowded for RL agent. Besides, we also extend the training steps to 2e5 for AMBS and DBC. Figure 17 shows the training curve on CARLA. AMBS and DBC run on the precious setting of CARLA, while AMBS40 and DBC40 run on the highway with 40 generated vehicles. When increasing the number of vehicles from 20 to 40, DBC becomes more difficult to learn but AMBS just slightly slower. DBC still doesn’t achieve high scores after 1e5 steps.

heading: E PROOFS
text: E.1 VALUE-FUNCTION BOUND
Theorem E.1 (Value function difference bound for different discounting factors (Petrik & Scherrer, 2009; Kemertas & Aumentado-Armstrong, 2021)). Consider two identical MDPs except different discounting factor γ1 and γ2, where γ1 ≤ γ2 and reward r ∈ [0, 1]. Let V πγ denotes the value function for MDP with discounting factor γ given policy π.∣∣V πγ1(s)− V πγ2(s)∣∣ ≤ γ2 − γ1(1− γ1)(1− γ2) (19) Definition E.1 (c-weighted on policy bisimulation metric). Given a policy π, and c ∈ (0, 1), the on-policy bisimulation metric exists,
d(si, sj) = (1− c)|Rπsi −R π sj |+ cW1(d)(P π si ,P π sj ). (20)
Theorem E.2 (Value difference bound). For any c ∈ [γ, 1), given two state si and sj ,
(1− c) |V π(si)− V π(sj)| ≤ d(si, sj) (21)
proof. Lemma 6 of Kemertas & Aumentado-Armstrong (2021). Theorem E.3 (Generalized value difference bound). For any c ∈ (0, 1), given two states si and sj ,
(1− c)|V π(si)− V π(sj)| ≤ d(si, sj) + 2(1− c)(γ −min(c, γ))
(1− γ)(1− c) (22)
proof. We follow the proof of Theorem 1 in Kemertas & Aumentado-Armstrong (2021). Suppose another MDP with discounting factor γ′ = c exists. From Theorem E.2 we have
(1− c)|V πγ′(si)− V πγ′(sj)| ≤ d(si, sj). (23)
Then,
(1− c)|V π(si)− V π(sj)| =(1− c)|V π(si)− V πγ′(si) + V πγ′(si)− V π(sj) + V πγ′(sj)− V πγ′(sj)| ≤(1− c)(|V πγ′(si)− V πγ′(sj)|+ |V π(si)− V πγ′(si)|+ |V π(sj)− V πγ′(sj)|) ≤d(si, sj) + (1− c)|V π(si)− V πγ′(si)|+ |V π(sj)− V πγ′(sj)) ≤d(si, sj) + 2(1− c)(γ −min(γ′, γ))
(1− γ)(1− γ′)
(24)
Replace γ′ by c, finally we have
(1− c)|V π(si)− V π(sj)| ≤ d(si, sj) + 2(1− c)(γ −min(c, γ))
(1− γ)(1− c) . (25)

Review summaries:
This work considers the problem of learning representations of high-dimensional pixel observations for RL. High-dimensional pixel observations often include many task-irrelevant details and ideally, an effective representation of such observations should only encode the task-relevant details. To implement this intuition, this work proposes to learn a two-part representation based off of a bisimulation metric. Two states that have similar first parts of the representation should exhibit similar reward structure, and two states that have similar second parts of the representation should exhibit similar dynamics structure. This work evaluates this learned representation on several continuous control tasks with distracting backgrounds and on CARLA, and shows that the proposed learned representation improves over baselines.

The paper studies the problem of learning invariant visual representations for reinforcement learning. Specifically, it builds on prior work Deep Bisimulation for Control (DBC) (Zhang et al), which learns an image representation predictive of rewards and dynamics and uses it with SAC. This paper makes a few modifications on top of DBC, specifically (1) instead of a single state embedding, it learns two separate ones (one for rewards and one for dynamics), (2) instead of using L1 distance in the embedding space to measure distance it instead uses a learned MLP which predicts distance, (3) using gradients from SAC it learns to dynamically adjust the balance between the dynamics weight and reward weight, and (4) adds image augmentation to the training. Experiments indicate on Distracting Control that on 2/4 tasks it matches DrQ and DBC and on the other 2 tasks outperforms them. Ablations suggest that all components are important, particularly (1) and (2).