L EARNING G ENERALIZABLE R EPRESENTATIONS FOR R EINFORCEMENT L EARNING VIA A DAPTIVE M ETALEARNER OF B EHAVIORAL S IMILARITIES
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
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; Wahlstr¨om 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 Lnorm 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 L 1 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.
R ELATED W ORK
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 Wahlstr¨om 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.
, 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.
P RELIMINARIES
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 [ (cid:80) t γ t r ( s t , a t ) | a t ∼ π ( ·| s t ) , 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 ( π )
= (cid:80) t E [ r ( s t , a t ) + α H ( π ( ·| s t ))] .
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
[Equation]
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
[Equation]
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):
[Equation]
where R π s := (cid:80) a π ( a | s ) R ( s , a ) , P π s := (cid:80) a π ( a | s ) P ( ·| s , a ) , mapping F π : met → met , met denotes the set of all pseudometrics on state space S and W 1 ( d )( · , · ) denotes the 1-Wasserstein distance given the pseudometric d .
Then F π has a least fixed point denoted by d ∗ .
Deep Bisimulation [Figure] for Control (DBC) (Zhang et al., 2020) proposed to learn latent representations such that the Ldistances in the latent space equal to the π -bisimulation metric in the state space:
[Equation]
where ϕ is the state encoder, and W 2 ( d )( · , · ) denotes the 2-Wasserstein distance.
DBC is combined with the reinforcement learning algorithm SAC where ϕ ( s ) is the input for SAC.
A DAPTIVE M ETA - LEARNER OF B EHAVIORAL S IMILARITIES
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 π s i −
R π s j | , which computes the difference of reward signals between states, and W 2 ( d )( P π s i , P π s j ) , 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 L 1 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.
4. M ETA - LEARNERS FOR D ECOMPOSED R EPRESENTATIONS
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 → Z r capturing reward-relevant features and ϕ d : S → Z d capturing dynamics-relevant features.
We design two meta-learners, f r and f d , where f r learns to measure reward difference and f d 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.
[Equation]
where ∗ ∈ { r, d }
[Equation] 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 L 1 / L 2 norm in the latent space may limit the approximation precision.
As shown in Figure 2, L 1 and L 2 norm for measuring distance of two latent representation ϕ ( s i ) and ϕ ( s j ) 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 L 1 norm in the Euclidean space, which however may lose part of useful information.
To overcome aforementioned limitations, we propose to exploit meta-learner f r and f d 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 L 1 / L 2 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:
[Equation]
To make the resultant optimization problems compatible with SGD updates with transitions sampled from replay buffer, we replace the reward metric term
|R π s i −
R π s j | by the distance of rewards in the sampled transitions 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 W 2 in (5) (as in DBC) because it has close form for
Gaussian distribution. With denoting by ( s i , a i , r i , s ′ i ) and ( s j , a j , r j , s ′ j ) two transitions sampled from the replay buffer, the loss functions (4) and (5) can be revised as follows,
[Equation]
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.
4. B ALANCING I MPACT OF R EWARD AND D YNAMICS
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 :
[Equation]
where Q θ is the Q-function parameterized by θ and V ( · ) is the target value function.
Moreover, to balance the learning of two meta-learners f r and f d , we jointly minimize the regression losses of f r and f d ((6) and (7), respectively) with combination weight c .
The loss function is modified as follows,
[Equation]
Note that parameter c is already learned along with Q-function, we stop gradient of c when minimizing the loss (9).
This is because if we perform gradient descent w.r.t.
c , it may only capture which loss ( ℓ ( f r , ϕ r ) or ℓ ( f d , ϕ 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.
4. O VERALL O BJECTIVE WITH D ATA A UGMENTATION
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,
[Equation]
c (cid:16) f d (cid:16) ϕ
d ( s (1) i ) , ϕ d ( s (1) j ) (cid:17) − W 2 (cid:16) ˆ P ( ·| ϕ d ( s (1) i ) , a i ) , ˆ P ( ·| ϕ d ( s (1) j ) , a j ) (cid:17)(cid:17)+(1 − c ) (cid:16) f r (cid:16) ϕ r ( s (2) j ) , ϕ r ( s (2) i ) (cid:17) − | r i − r j | (cid:17)+ c (cid:16) f d (cid:16) ϕ d ( s (2) j ) , ϕ d ( s (2) i ) (cid:17) − W 2 (cid:16) ˆ P ( ·| ϕ d ( s (1) i ) , a i ) , ˆ P ( ·| ϕ d ( s (1) j ) , a j ) (cid:17)(cid:17) 2 , (10)
[Table]
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 f d 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 f r and f d to be symmetric.
The loss of the actor of SAC comprises of encoder is
[Equation]
(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 f r and f d are both MLPs with two layers.
More implementation details of AMBS are provided at Appendix B.
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 L 1 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.
5. DMC SUITE WITH BACKGROUND DISTRACTION
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.
[Figure]
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.
5. A BLATION S TUDY
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 f r and f d by L 1 distance, as DBC use L 1 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 f r and f d and decomposing states into reward and dynamics representations help to improve RL performance than DBC that uses L 1 distance to approximate the whole π -bisimulation metric.
5. T RANSFER OVER R EWARD F UNCTIONS
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.
5. G ENERALIZATION OVER B ACKGROUND V IDEO
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.
[Table]
[Figure]
5. A UTONOMOUS D RIVING T ASK 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.
Figureshows 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.
C ONCLUSION
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.
A O BJECTIVES OF A DAPTIVE M ETA - LEARNER OF B EHAVIORAL S IMILARITIES (AMBS)
We denote two encoders ϕ r and ϕ d , where each encoder maps a high-dimensional observation to a low-dimensional representation.
The encoder ϕ r : S → Z r captures reward-relevant features and the encoder ϕ d : S → Z d captures dynamics-relevant features, where Z r and Z d are representation spaces.
We design two meta-learners f r :
Z r × Z r → R and f d : Z d × Z d → 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,
[Equation]
where Θ = { f r , f d , ϕ 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 Z d for predicting dynamics representation of next state s ′∗ .
W 2 denotes the 2-Wasserstein distance which has close-form for Gaussian distributions: W 2 ( N ( µ i , Σ i ) , N ( µ j , Σ j )) 2 = ∥ µ
i − µ j ∥ 22 + ∥ Σ 12 i − Σ 12 j ∥ 2 F , where µ ∗ ∈ R |Z
d | are the mean vectors of Gaussian distribution,
Σ ∗ ∈ R |Z
d | 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 ∈ R 2 .
It can be formulated as
[Equation]
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
[Equation]
where s ( l ) =
h ( s , v ( l ) ) , v ( l ) are all drawn from T independently, and { Q θ k } 2 k =1 are the double Q networks.
The target value function with DrQ augmentation is
[Equation]
where s ′ ( l ) =
h ( s ′ , v ′ (
[Equation]
) ) , 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,
[Equation]
The loss of α of SAC is,
[Equation]
where ¯ H ∈ R is the target entropy hyper-parameter which in practice is ¯ H = −|A| .
The loss for updating of dynamic model ˆ P is
[Equation]
where µ ( ˆ P ( · )) and σ ( ˆ P ( · )) are mean and standard deviation of ˆ P output Gaussian distribution, respectively.
Algorithm 2 shows the algorithm at each learning step.
B I MPLEMENTATION D ETAILS B. P IXELS P ROCESSING
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.
B. N ETWORK A RCHITECTURE
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 withdimensions 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 f r and f d 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.
B. H YPERPARAMETERS
The hyperparameters used in our algorithm are listed in Table 2.
B. A CTION R EPEAT FOR DMC
We use different action repeat hyper-parameters for each task in DeepMind control suite, which are listed in Table 3. [Table] [Table]
B. S OURCE C ODES
We will open the source codes to public once this paper is accepted.
C A NALYSIS C. R EGRESSION L OSSES OF A PPROXIMATING B ISIMULATION M ETRIC
We compare the regression losses of approximating the bisimulation metric over RL environment steps among DBC and our AMBS.
DBC uses L 1 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. V ISUALIZATIONS OF C OMBINATION W EIGHT 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.
C. N ORMS OF S TATE E MBEDDINGS
We record the L1 norms of reward representation ϕ r ( s ) and dynamic representation ϕ d ( s ) during training. We record 1 n r || ϕ r ( s ) || 1 for reward representation where n r is the number of dimensions of ϕ r ( s ) , and 1 n d || ϕ d
( s )
|| 1 for dynamic representation where n r 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.
D A DDITIONAL E XPERIMENTAL R ESULTS D. A DDITIONAL R ESULT OF DMC S UITE WITH AND W / O B ACKGROUND D ISTRACTION
Figure 14 shows the training curves on original background setting. Figure 15 shows
D. A DDITIONAL R ESULT OF G ENERALIZATION OVER B ACKGROUND V IDEO
Table 4 shows the scores on DAVIS video background setting comparing to PSEs (Agarwal et al., 2021). Methods
[Table]
D. A DDITIONAL R ESULT ON CARLA
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 16 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.
E P ROOFS E. V ALUE - FUNCTION B OUND
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 π .
[Equation]
Definition E.1 ( c -weighted on policy bisimulation metric) .
Given a policy π , and c ∈ (0 , 1) , the on-policy bisimulation metric exists,
[Equation]
Theorem E.2 (Value difference bound) .
For any c ∈ [ γ, 1) , given two state s i and s j , (1 − c
[Equation]
proof.
Lemma 6 of Kemertas & Aumentado-Armstrong (2021).
Theorem E.3 (Generalized value difference bound) .
For any c ∈ (0 , 1) , given two states s i and s j ,
[Equation]
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
[Equation]
Then,
[Equation]
Replace γ ′ by c , finally we have
[Equation]