R OBUST R OBOTIC C ONTROL FROM P IXELS USING C ONTRASTIVE R ECURRENT S TATE -S PACE M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
Modeling the world can beneﬁt robot learning by providing a rich training signal for shaping an agent’s latent state space.
However, learning world models in unconstrained environments over high-dimensional observation spaces such as images is challenging.
One source of difﬁculty is the presence of irrelevant but hard-to-model background distractions, and unimportant visual details of taskrelevant entities.
We address this issue by learning a recurrent latent dynamics model which contrastively predicts the next observation.
This simple model leads to surprisingly robust robotic control even with simultaneous camera, background, and color distractions.
We outperform alternatives such as bisimulation methods which impose state-similarity measures derived from divergence in future reward or future optimal actions.
We obtain state-of-the-art results on the Distracting Control
Suite, a challenging benchmark for pixel-based robotic control.
I NTRODUCTION
For a robot, predicting the future conditioned on its actions is a rich source of information about itself and the world that it lives in.
The gradient-rich training signal from future prediction can be used to shape a robot’s internal latent representation of the world.
These models can be used to generate imagined interactions, facilitate model-predictive control, and dynamically attend to mispredicted events.
Video prediction models (Finn et al., 2016) have been shown to be effective for planning robot actions.
Model-based Reinforcement Learning (MBRL) methods such as Dreamer (Hafner et al., 2020) and SLAC (Lee et al., 2020) have been shown to reduce the sample complexity of learning control tasks from pixel-based observations.
These methods learn an action-conditioned model of the observations received by the robot.
This works well when observations are clean, backgrounds are stationary, and only task-relevant information is contained in the observations.
However, in a real-world unconstrained environment, observations are likely to include extraneous information that is irrelevant to the task (see Figure 1a).
Extraneous information can include independently moving objects, changes in lighting, colors, as well as changes due to camera motion.
High-frequency details, even in task-relevant entities, such as the texture of a door or the shininess of a door handle that a robot is trying to open are irrelevant.
Modeling the observations at a pixel-level requires spending capacity to explain all of these attributes, which is inefﬁcient.
The prevalence of extraneous information poses an immediate challenge to reconstruction based methods, e.g. Dreamer.
One promising approach is to impose a metric on the latent state space which groups together states that are indistinguishable with respect to future reward sequences (DBC, Zhang et al. or future action sequences (PSE, Agarwal et al. when running the optimal policy.
However, using rewards alone for grounding is sample inefﬁcient, especially in tasks with sparse rewards.
Similarly, actions may not provide enough training signal, especially if they are low-dimensional.
On the other hand, observations are usually high-dimensional and carry more bits of information.
One way to make use of this signal while avoiding pixel-level reconstruction is to maximize the mutual information between the observation and its learned representation (Hjelm et al., 2019).
This objective can be approximated using contrastive learning models where the task is to predict a learned representation which matches the encoding of the true future observation (positive pair), but does not match the encoding of other observations (negative pairs).
However, the [Figure] [Figure] performance of contrastive variants of Dreamer and DBC has been shown to be inferior, for example,
Fig. 8 in Hafner et al.
(2020) and Fig.
3 in Zhang et al. In this work, we show that contrastive learning can in fact lead to surprisingly strong robustness to severe distractions, provided that a recurrent state-space model is used.
We call our model
CoRe: Co ntrastive Re current State-Space Model.
The key intuition is that recurrent models such as GRUs and LSTMs maintain temporal smoothness in their hidden state because they propagate their state using gating and incremental modiﬁcations.
When used along with contrastive learning, this smoothness ensures the presence of informative hard negatives in the same mini-batch of training.
Using CoRe, we get state-of-the-art results on the challenging Distracting Control Suite benchmark (Stone et al., 2021) which includes background, camera, and color distractions applied simultaneously.
Video results from our model can be seen at supplementary/index.html.
M ODEL D ESCRIPTION
We formulate the robotics control problem from visual observations as a discrete-time partially observable Markov decision process (POMDP).
At any time step t , the robot agent has an internal state denoted by z t .
It receives an observation o t , takes action a t , and moves to state z t +1 , receiving a reward r t +1 .
We present our model with continuous-valued states, actions, and observations.
A straightforward extension can be made to discrete states.
The objective is to ﬁnd a policy that maximizes the expected sum of future rewards (cid:80) ∞ t =
γ t r t , where γ is a discount factor.
2. M ODEL C OMPONENTS
As shown in Figure 1b, our model consists of three components: a recurrent state space model (RSSM)
(Hafner et al., 2019) which encodes the learned dynamics, a behavior model based on
Soft-Actor-Critic (SAC) (Haarnoja et al., 2018) which is used to control the robot, and an observation encoder.
Observation encoder :
This component is a convolutional network f which takes an observation as input and outputs its encoding x t = f ( o t ) .
LayerNorm (Ba et al., 2016) is applied in the last layer of the encoder.
Recurrent State Space Model (RSSM) :
This component consists of the following modules,
Future predictor
[Equation]
Observation representation decoder
[Equation]
Correction Model
[Equation]
Inverse Dynamics model
[Equation]
Reward predictor [Equation] Each module is parameterized as a fully-connected MLP, except the future predictor which uses GRU cells (Cho et al., 2014).
The future predictor outputs a prior distribution ˆ z t over the next state, given the previous state and action.
This is decoded into y t which corresponds to an encoding of the observation that the agent expects to see next.
The correction model q is given access to the current observation encoding x t , along with the prior ˆ z t and outputs the posterior distribution z t , i.e., it corrects the prior belief given the new observation.
The posterior state z t is used to predict the reward, and also feeds into the actor and critic models that control the agent’s behavior.
The inverse dynamics model predicts the action that took the agent from state z t − 1 to z t .
The RSSM is similar to the one used in Dreamer and SLAC except that it includes an inverse dynamics model and that the observation representation is predicted from the prior latent state instead of the posterior latent state.
These differences lead to slightly better results and a different probabilistic interpretation of the underlying model which we discuss in Section 2.2.
Similar to Dreamer, the latent state z t consists of a deterministic and stochastic component (More details in Appendix A).
Behavior Model : This component consists of an actor π φ ( a t | z t ) and two critic networks Q θ i ( z t , a ) , i ∈ { 0 , 1 } which are standard for training using the SAC algorithm (Haarnoja et al., 2018).
2. R EPRESENTATION L EARNING WITH C O R E
Our goal is to represent the agent’s latent state by extracting task-relevant information from the observations, while ignoring the distractions.
We formulate this representation learning problem as dynamics-regularized mutual information maximization between the observations o t and the model’s prediction of the observation’s encoding y t .
Mutual information maximization (Hjelm et al., 2019) has been used extensively for representation learning.
In our case, it can be approximated using a contrastive learning objective deﬁned over a mini-batch of sequences, in a way similar to
Contrastive Predictive Coding (CPC) (van den Oord et al., 2018).
At time-step t in example i in the mini-batch, let x i,t and y i,t denote the real observation’s encoding and the predicted observation encoding, respectively.
The loss function is
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation] where λ is a learned inverse temperature parameter and i (cid:48) and t (cid:48) index over all sequences and timesteps in the mini-batch, respectively.
Intuitively, the loss function makes the predicted observation encoding match the corresponding real observation’s encoding, but not match other real observation encodings, and vice-versa.
In addition, the predicted and corrected latent state distributions should match, which can be formulated as minimizing a KL-divergence term
[Equation]
An additional source of training signal is the reward prediction loss, ( r t − ˆ r t ) 2 .
Furthermore, modeling the inverse dynamics, i.e. predicting a t from z t and z t +1 is a different way of representing the transition dynamics.
This does not provide any additional information because we already model forward dynamics.
However, we found that it helps in making faster training progress.
We incorporate this by adding an action prediction loss.
The combined loss is
[Equation]
where β, α r , α a are tunable hyperparameters and Θ is the combined set of model parameters which includes parameters of the observation encoder, the RSSM, and the inverse temperature λ .
Relationship with Action-Conditioned Sequential VAEs The above objective function bears resemblance to the ELBO from a Conditional Sequential VAE which is used as the underlying probabilistic model in Dreamer and SLAC.
The objective function there is the sum of an observation reconstruction term
||
o t − ˆ o t || 2 and L KL scaled by β , corresponding to a β -VAE formulation (Higgins et al., 2017).
This formulation relies on decoding the posterior latent state to compute ˆ o t .
Since the posterior latent state distribution q ( z t | ˆ z t , x t ) is conditioned on the true observation x t , observation reconstruction is an auto-encoding task.
Only the KL-term corresponds to the future prediction task.
On the other hand, in our case the prior latent state is decoded to a representation of the observation, which already constitutes a future prediction task.
The KL-term only provides additional regularization.
This makes the model less reliant on careful tuning of β (empirically validated in Appendix C.3).
In fact, setting it to zero also leads to reasonable performance in our experiments, though not the best.
2. B EHAVIOR L EARNING
Behavior is learned simultaneously with the state representation using SAC.
This involves learning an actor π φ that parameterizes a stochastic policy and two critics Q θ i , i ∈ { 1 , 2 } which use the latent state z t as the representation of the agent’s state.
The critic is trained by minimizing Bellman error,
[Equation]
where ¯ θ i represents an exponentially moving average of θ i , ‘sg’ denotes the stop-gradient operation, and α is the weight on the entropy regularization.
Note that J Q is a function of Θ as well.
In other words, the RSSM and observation encoder model are also trained using the gradients from the critic loss.
Appendix C.2 explores this design choice in more detail.
The parameters Θ are shared between the two critics.
The actor is updated using the following loss function
[Equation]
Overall training consists of three steps: collecting data using the current actor, using the data to update the representation model, and using the data to update the actor and critic as shown in Algorithm 1.
E XPERIMENTS
[Table]
We use the
Distracting
Control
Suite (DCS)
(Stone et al., 2021) and Robosuite (Zhu et al., 2020) to benchmark our model’s ability to withstand visual distractions.
DCS consists of six simulated robotic control tasks derived from the DeepMind Control Suite (Tassa et al., 2018).
There are three types of distraction: background, camera, and color.
They can either be ﬁxed for the duration of the episode (static) or changed smoothly during the episode (dynamic).
Three benchmarks are deﬁned: easy, medium, and hard, which have increasing levels of difﬁculty along all 3 distraction types.
We focus our experiments on two settings: dynamic-background-only and dynamicmedium.
Results on other settings follow the same trends (Appendix B).
In the dynamicbackground-only case, the backgrounds for training and validation environments are drawn from the 60 training and 30 validation videos in the DAVIS dataset (Pont-Tuset et al., 2017).
Our experiments show that (1) CoRe enables visual robotic control in the presence of severe distractions, performing better than strong baselines, (2) all three elements of CoRe - recurrence, contrastive learning, and dynamics modeling are important, (3) recurrent architectures are important to make contrastive learning work well, (4) CoRe can solve a challenging simulated robotic manipulation task, and (5) interpretable attention masks emerge that help visualize how CoRe works.
We report the mean validation reward over 10 random seeds and 100 episodes per seed along with the standard error.
Details of the network architectures and training hyper-parameters are provided in appendix A.
3. C OMPARISON OF ROBUSTNESS TO DISTRACTIONS
This set of experiments quantiﬁes the robustness of CoRe to visual distractions.
We compare to a model-free RL baseline SAC-RAD , a state-of-the-art method based on bisimilarity metrics PSE (Agarwal et al., 2021), Recon , a reconstruction variant of our model, where we replace the contrastive loss with an observation reconstruction loss (making it similar to SLAC), and CURL (Laskin et al., 2020a) which also uses contrastive learning but does not use recurrence or future prediction.
Figure 2 shows the comparison as a function of environment steps for 4 of the 6 environments.
We use the published code for PSE and CURL and only change the distraction settings.
In the no-distraction setting (Figure 2, top row), we can see that all methods perform well.
In the dynamic-background-only setting (Figure 2, middle row), CoRe and PSE both perform well overall
[Figure]
[Table]
[Figure]
[Table]
and better than the other methods.
CoRe does signiﬁcantly better than PSE on the ‘cheetah
SAC-RAD does not perform well, conﬁrming that model-free RL methods are insufﬁcient to solve these tasks.
Recon also does not work well, as expected, because it wastes capacity on modeling the background.
In the dynamic-medium setting (Figure 2, bottom row), we see that the performance of CoRe is far better than other baselines.
CoRe is able to solve most of the tasks in this challenging distraction setting even while other methods struggle to make progress.
In Table 1, we compare the performance of our proposed model with bisimulation-based methods (DBC and PSE) and CURL.
Results from DBC and PSE have been reported with different background videos making them hard to compare directly.
DBC results were reported using videos from the Car
Driving class in the Kinetics dataset (Kay et al., 2017), while PSE results were reported when training on 2 videos from DAVIS and testing on 30 validation videos.
To do a fair comparison, we train CoRe as well as DBC (using their published code) in the same setting as PSE.
This setting is denoted as
DAVIS (2 videos) in Table 1.
We can see that here, PSE outperforms DBC.
CoRe performs similar to
PSE.
However, when we use all 60 training videos, CoRe performs better.
3. A BLATION E XPERIMENTS
In this set of experiments, we compare CoRe to the variants shown in Table 2 in the dynamic-medium distraction setting. From Table 3, we can see that model-free methods, SAC and QT-Opt, trained
[Table]
[Table]
with data augmentation are not able to solve these tasks.
Adding recurrence (RSAC) and dynamics modeling (RSAC+Dyn) improves results.
However, trying to predict pixels hurts performance (Recon).
Contrastive prediction in a learned space (CoRe) works much better.
CURL and RCURL, which use a constrastive loss but do not model dynamics, perform worse that CoRe, showing that just having a contrastive loss is insufﬁcient.
More ablation experiments that study the relative effects of modeling rewards, forward and inverse dynamics are presented in appendix C.1.
3. I MPORTANCE OF HAVING A RECURRENT STATE SPACE MODEL
Contrastive models have been reported to underperform
DBC (Zhang et al., 2021).
However, in our experiments we found that contrastive models work quite well.
We hypothesize that this difference is because we use a recurrent state space whereas previous work used contrastive learning with a feed-forward encoder operating on stacked observation frames.
To validate this hypothesis, we compare our model to its non-recurrent version.
The models both do contrastive future prediction and use similar network architectures, but the non-recurrent one consumes three stacked frames.
The training mini-batches are sampled similarly using 32 sequences of 32 time-steps each.
Therefore, even the non-recurrent model has access to observations that are semantically close.
The results at 500K environment steps in the dynamic-medium setting on DCS are shown in Figure 3 (Top).
We can see that, in the absence of recurrence, contrastive learning does not perform well.
This indicates that just having nearby observations is not sufﬁcient to provide hard negatives and a recurrent network has help ﬁx this problem.
An intuitive explanation for this is the following.
A recurrent network (in our case, a GRU-RNN) has a smoothness bias built-in because at each time step, it carries forward previous state and only modiﬁes it slightly by gating and adding new information.
This is true to a large extent even during training, and not just at convergence.
Therefore, when CoRe is trained, it generates hard negatives throughout training in the form of nearby future predictions.
This is true even when the observations have distractions present which change the overall appearance of the observation dramatically.
On the other hand, starting from a random initialization, feed-forward networks are less likely to map observations that are semantically close but visually distinct to nearby points.
Therefore, hard negatives may not be found easily during training.
To conﬁrm this further, we train CoRe with a modiﬁed contrastive loss in which for each sample in the mini-batch, only the same temporal sequence is used to sample negatives.
As shown in Figure(bottom), this is not harmful (but actually beneﬁcial) to performance on all tasks, except reacher.
This means that for most tasks, CoRe doesn’t need other sequences in the mini-batch to ﬁnd hard negatives. This avoids a major concern in contrastive learning, which is having large mini-batches to ensure hard negatives can be found.
Essentially, recurrence provides an architectural bias for generating good negative samples locally.
Performance degrades on the reacher task because observations there contain a goal position for the arm.
Contrastive learning tries to remove this information because it is constant throughout the episode.
Therefore, the actor and critic may not get access to the goal information, causing the agent to fail.
This highlights a key limitation of contrastive learning – it discourages retaining constant information in the latent space.
For that type of information, it is important to have negative examples coming from other episodes.
3. R ESULTS ON R OBOSUITE
In this section, we present results of our method on Robosuite (Zhu et al., 2020), a MuJoCo-based simulator for robotic manipulation tasks.
We experiment with two kinds of arms: Panda and Jaco, and two distraction settings: (1) Color and lighting only, and (2) Color, lighting, and camera.
In this experiment, we use the static distraction setting to emulate a scenario where data collected within an episode has consistent lighting, color, and camera position, but these factors vary across episodes (depending on time of day, accidental movement of the camera relative to the robot setup, etc).
The action space in delta end-effector pose.
Control is run at 20Hz, with each episode lasting forsteps (25 seconds).
Architecture and training details are described in appendix A. Figure 4 compares CoRe with recurrent versions of SAC and CURL.
In the distraction-free case, we also compare to results from SAC operating on ground-truth object state instead of using image observations, as reported by Zhu et al.
This can be seen as an upper bound on performance because the agent gets perfect access to task-relevant state.
In the case where proprioception (the robot’s state) is provided separately (top row), the performance of CoRe is same as or better than that of CURL.
In the harder case where the robot’s state must also be obtained from the camera observation, CoRe signiﬁcantly outperforms other methods. This shows that CoRe is better at handling more complex vision tasks.
Added a pixel-level gating CNN to the encoder
Image
Sigmoid gating
Gated image
CNN CNN
3. V ISUALIZATIONS
In this section, we present qualitative visualizations to help understand what the model does to cope with distractions.
We make visualization easier by modifying the observation encoder to include a pixel-level gating network.
This is done using a four-layered CNN which takes the observation as input and outputs a full-resolution sigmoid mask.
This mask is multiplied into the observation image, and the result is fed to the encoder.
The gating network is trained along with the rest of the model.
We expected this gating to improve performance but found that it actually has no effect.
This lets us use it as a diagnostic tool that factors out the attention mask that the model is otherwise applying implicitly.
Figure 5 shows the attention masks inferred by the gating network.
We can see that the model learns to exclude distractions in the background, even in the dynamic-medium distraction setting for the ball in cup task where the task-relevant objects are very small.
For the door opening task, the model learns to attend to the bottom edge of the door handle, the top and right edges of the door, the edge of the table, and some parts of the robot’s base.
This includes all the information needed to ﬁnd the relative position of the door with respect to the agent.
Proprioception is provided separately.
Therefore, the agent doesn’t need to attend to the arm itself.
D ISCUSSION AND R ELATED WORK
Our work uses contrastive learning in the context of world modeling and pixel-based robotic control.
In this section, we review related work in these areas.
Pixel-based robotic control Controlling robots directly from raw environment observations is a general framework for robotic problem solving since it avoids the use of speciﬁc modules such as object detectors and instead relies on the data and reward function to learn a representation of the observations.
Model-free methods such as QT-Opt (Kalashnikov et al., 2018) and SAC (Haarnoja et al., 2018) have been applied to pixel-based control.
However, learning encoders for high-dimensional observations from reward only is sample inefﬁcient.
Simply adding observation autoencoding to
SAC as an auxiliary task has been shown to work well in Yarats et al.
Instead of using a generative model, CURL (Laskin et al., 2020a) uses contrastive learning to learn observation encoders yielding similar results.
Data augmentation techniques RAD (Laskin et al., 2020b) and DrQ (Yarats et al., 2021) dramatically improve sample efﬁciency.
However, these methods are still not sufﬁcient to solve visual tasks in the presence of distractions (Stone et al., 2021).
Other approaches shape the learned encoders to produce a latent space that is amenable to locally-linear control algorithms (Watter et al., 2015; Cui et al., 2020; Levine et al., 2020).
Improvements in the encoder architecture, for example, using attention Tang et al.
and efﬁcient transformer-based models (Choromanski et al., 2021) have been proposed.
Our proposed model can potentially beneﬁt from the use of these better architectures.
We use a simple convolutional encoder for our experiments and focus on the modeling aspect of the problem.
World Modeling World models aim to learn how the world changes as a result of the agent’s action instead of just learning the optimal policy.
Finn et al.
(2016) used video prediction parameterized as pixel-level ﬂow to predict the future video frames.
Dreamer (Hafner et al., 2020) and SLAC (Lee et al., 2020) model observations using a sequential action-conditioned VAE.
However, these methods are not designed to deal with distractions since they would waste capacity on reconstructing them along with the task-relevant portion of the observations.
In order to avoid observation reconstruction, Gelada et al.
(2019) proposed DeepMDP, a method that embeds the given high-dimensional RL problem into a latent space which preserves the essence of the problem (i.e. the reward and transition function).
While this method only preserves the current reward, Deep Bisimulation for Control (Zhang et al., 2021) uses the future sequence of rewards to extract a stronger learning signal.
In this case, the learned state space is trained to impose a bisimulation metric, under which similarity between states is determined by the similarity in future reward sequences produced from those states.
PSE (Agarwal et al., 2021) further extends this idea to use similarity in future action distributions.
While these models use dynamics to get more training signal, they do not make use of structure in the observation space itself.
In our proposed model, we seek to ﬁnd a middle ground between modeling observations and not modeling them at all.
This is done in a principled way using mutual information maximization between the observation and its prediction from the internal dynamics.
While previous work often found contrastive learning to be ineffective, we show that combining it with recurrent state space models makes it work.
Recently, a contrastive variant of Dreamer (Okada & Taniguchi, 2020) has been proposed which shares the same motivation.
Concurrent with our work, Nguyen et al.
(2021) explore a formulation similar to ours based on temporal predictive coding, but do not evaluate it on the difﬁcult camera and color distractions we do here.
Contrastive Learning : Learning to contrast positive and negative examples is a general representation learning tool since it can be applied to many scenarios by changing what positive and negative pairs mean.
They can be deﬁned based on identity, for example, in face veriﬁcation (Weinberger et al., 2006; Schroff et al., 2015), object identity for tracking (Wang & Gupta, 2015), closeness in temporal order in the case of time-contrastive models (Sermanet et al.
Hyvarinen & Morioka, 2016), word and sentence ordering in the case of language models such as Log-Bilinear models Mnih &
Hinton (2007), Word2Vec (Mikolov et al., 2013) and Skip-thoughts (Kiros et al., 2015).
Contrastive
Predictive Coding (CPC) (van den Oord et al., 2018) proposes an auto-regressive model to do future prediction at multiple time steps and applies it to various domains including image patches, audio, and RL.
In this case, positives are the corresponding predicted and observed representations, and negatives are all other pairs.
Our model uses a similar contrastive loss as CPC.
However, it also uses dynamics modeling in latent space to improve the representation quality.
Subsequent works have improved on CPC using momentum in the learned encoder (MoCo (He et al., 2020)), adding projection heads before computing contrastive loss (SimCLR (Chen et al., 2020)), avoiding negative examples altogether (BYOL (Grill et al., 2020) and SwAV (Caron et al., 2020)).
Our model can potentially be improved using these techniques but we found that the simple contrastive loss is already sufﬁcient to solve challenging visual control problems.
C ONCLUSION
We show that a contrastive objective can be a viable replacement for reconstructing observation pixels in solving visual robotics control problems using model-based RL.
It works even in the presence of severe distractions.
In addition, we show that having a recurrent architecture helps contrastive learning work well.
Even though our results are only in simulated environments, the surprisingly high degree of robustness indicates that this approach should transfer well to a real-world setting.
In future work, we plan to test this approach on real-world robots. A A RCHITECTURE AND T RAINING D ETAILS
In this section, we present the architecture and training hyper-parameters for the proposed CoRe model in more detail.
A. GRU-RNN WITH STOCHASTIC AND DETERMINISTIC STATES
CoRe uses a recurrent state-space model (RSSM) that is based on GRUs (Cho et al., 2014) and similar to the one used in Dreamer (Hafner et al., 2020), as shown in the ﬁgure below.
Latent states at any time-step consist of a (deterministic) recurrent hidden state and a stochastic state.
The prior latent state is ˆ z t =
[ h t , ˆ s t ] and the posterior latent state is z t
= [ h t , s t ] .
They share the same recurrent hidden state h t but differ in the stochastic component.
ˆ s t depends on h t , whereas s t depends on h t and the new observation features x t .
The feed-forward operation is h t
= GRUCell
( h t − 1 , GRU MLP ( s t − 1 , a t − 1 )) , ˆ µ , ˆ σ = PriorMLP ( h t ) , ˆ s t
= ˆ µ
+ ˆ σ ∗ N (0 , 1) , µ , σ = PosteriorMLP ( h t , ObsMLP ( x t )) , s t = µ + σ ∗ N (0 , 1) ,
Binder Networks | Progress 06/
[Table]
where MLP stands for a fully-connected multi-layered neural network.
The ﬁgure above illustrates this operation.
It can be seen as a more detailed view of the RSSM shown in Figure 1(b).
Note that the recurrent hidden state h t is propagated through time without sampling.
This is important because sampling can make it hard to retain information for long time-scales and destroys smoothness of the latent state.
However, information inserted into the RNN through s t is stochastic, making the
RSSM capable of generating multiple futures.
This prevents mode-averaging and allows the model to generate diverse trajectories when doing rollouts.
A. A RCHITECTURE
The model consists of a number of components: observation encoder CNN, RSSM (which includes
PriorMLP, PosteriorMLP, GRU MLP, GRUCell, ObsMLP, along with the reward, inverse dynamics, and observation representation decoder networks), and the actor and critic networks.
Table 4 describes the architecture of these networks.
CNN layers are denotes as [ num ﬁlters , kernel size , stride ] .
We borrow the observation encoder architecture from previous work (Laskin et al., 2020a; Yarats et al., 2019) and use it as-is to ensure an accurate comparison.
ELU non-linearity (Clevert et al., 2016) is used in all places except the observation encoder (which uses ReLU units).
For experiments with Robosuite, the observation encoder architecture is modiﬁed to accept 128 ×resolution images.
Two convolutional layers with 32 ﬁlters of kernel size 3 × 3 with a stride ofwere applied to reduce the spatial size from 128 × 128 to 32 × 32 .
This was followed by 3 layers of 3 × 3 convolutions with a stride of 1 .
A. T RAINING
The training is broken into iterations, where in each iteration one episode of data are collected, added to the replay buffer, and a number of gradient steps are taken using mini-batches sampled from the replay buffer.
Data Collection For the ﬁrst 1000 steps, data are collected by taking actions sampled from a uniform distribution in [ − 1 , 1] .
After that, data are collected by taking actions sampled from the learned actor policy, which requires rolling out the RSSM.
At each training iteration, an entire episode of data are collected, where the length of the episode is 1000 steps of the underlying MuJoCo simulator.
The actual number of steps is 1000 divided by the number of times the same action is repeated, which is standard based on the task (Table 6).
We rollout an entire episode because we use a recurrent model and it seemed natural to use the same model through time within an episode.
Typically, model-free
RL algorithms collect data one environment step at a time.
We did not explore per-step data collection
[Table]
[Table]
in this work, although that could work just as well.
The data from the Distracting Control Suite is rendered at a 320 × 240 resolution which is resized to 100 × 100.
Pixels are normalized to [0 , 1] by dividing by 255.
All actions are normalized to lie in [ − 1 , 1] and are modeled using tanh-squashed
Gaussian distributions.
For robosuite, visual observations were rendered from the front-view camera at a 256 × 256 resolution and a 128 × 128 crop near the center of the image was fed as input to the model.
Model updates Training is done with mini-batches of N = 32 sequences of length T = 32 each sampled from the replay buffer.
Data augmentation is done by taking random 84 × 84 crops.
The same crop position is used across the entire sequence.
Each mini-batch is used to do three updates sequentially corresponding to the three losses: J M , J Q and J π .
The number of updates is chosen to be 0.5 times the number of steps in an episode.
All training hyper-parameters are listed in Table 5.
Implementation The model is implemented using PyTorch (Paszke et al., 2019).
The environment is based on MuJoCo (Todorov et al., 2012).
All training was done on single Nvidia A100 GPUs.
Training time for 500K updates varies from 12-24 hr depending on the task.
Training times differ due to different values of action repeat.
Our implementation will be made public.
A. T RAINING CURVES FOR INDIVIDUAL LOSS TERMS
[Figure]
There are a number of loss terms which are linearly combined to create the total model loss J M .
In this section, we present training curves that show how these individual terms optimize with training.
Figure 6 shows these loss terms, along with the RL training losses ( J Q and J π ) and the episode reward.
We can see that the total loss J M goes down as expected, along with the individual terms.
Note that the reward loss is low in the beginning because the agent gets zero rewards which is easy for the model to predict.
As the agent starts receiving better rewards, the reward error increases but then eventually starts to come down.
The KL-term has a similar behavior.
At the very beginning, the prior and posterior latent state distributions match each other (making their KL divergence small) but both are bad at modeling the data.
As more training data is encountered, the KL rises sharply, then eventually reduces with more training.
The actor loss (which is dominated by the negative of the
Q-value of the actor’s chosen action) goes down as expected, indicating that the actor outputs actions that have high Q-values.
The critic loss appears to increase because the magnitude of the Q-values increases with training, causing the magnitude of the Bellman residual to go up as well.
B C OMPLETE R ESULTS ON D ISTRACTING C ONTROL S UITE
In this section, we report the performance of CoRe on all the distraction settings in the Distracting
Control Suite.
There are three difﬁculty levels: easy, medium, and hard 1 as shown in Figure 7.
The difﬁculty is set by increasing the scale of the camera pose and color change and the number of background distractions that are used (Table 7).
For each difﬁculty level, there are two settings: static , in which the distraction (a background image, or a particular choice of random camera pose) is ﬁxed throughout the episode, and dynamic in which the distraction changes smoothly.
In the case of background distractions, the video plays back-and-forth, ensuring no sharp cuts.
In the case of the camera distractions, the camera moves along a smooth trajectory.
Table 8 compares the performance of CoRe with model-free RL baselines such as SAC (Haarnoja et al., 2018) and QT-Opt (Kalashnikov et al., 2018) combined with data augmentation techniques
RAD (Laskin et al., 2020b) and DrQ (Yarats et al., 2021) as reported in (Stone et al., 2021).
We also compare with a recurrent SAC+RAD baseline (RSAC) which uses the same recurrent architecture as
CoRe, but does not contrastively predict the next observation, or model the dynamics and reward.
Comparisons are made at 500K environment steps, though we report our results at 1M environment
[Table]
[Table]
steps as well to show that our model continues to improve with more steps.
Performance is averaged over 10 random seeds, and 100 validation episodes at each checkpoint.
When no distractions are present (Table 8a), all methods perform well, though CURL and CoRe are slightly better than the rest in terms of mean scores.
On the easy benchmark (Table 8b), CoRe outperforms other methods in all tasks, except reacher.
As discussed in Section 3.3, this points to a key limitation of contrastive learning-based methods, which is that they tend to remove constant information (such as the goal location for reacher).
However, at 1M steps, the performance on reacher is much better, showing that the model is able to eventually solve the task.
On the medium benchmark (Table 8c) CoRe outperforms other methods across all tasks, showing that it can deal with the presence of more distractions.
The performance on the reacher tasks improves slightly over the easy setting, which shows that having more variation in the distractions actually helps training, whereas it hurts the baseline methods.
We also report results on the hard setting (Table 8d), which is documented in the DCS codebase.
Stone et al.
(2021) do not report model-free baselines for this setting, presumably because the baseline models fail to train reasonable policies at all.
However, CoRe is able to get off the ground and get reasonable performance even in this setting.
Training to 1M steps continues to improve results.
Figure 8 shows the progression of validation reward for CoRe over 1M environment steps for all tasks and distractions settings.
We can see that the hard-dynamic setting (green curves in the bottom two rows) is the hardest to ﬁt because the performance increases slowly.
However, in most cases the performance for that setting is still improving at 1M steps, and is likely to get better with more training.
Our model struggles on the reacher task in terms of performance reported at 500K and even 1M steps, but we can see from these plots that the model is likely to continue improving if trained beyond 1M steps in both static and dynamic settings.
In our experiments, we did not tune hyper-parameters speciﬁcally for each task, so it is possible that some tasks can beneﬁt from further tuning.
In particular, for the reacher task, a bigger batch-size can help since that is the only way to get access to diverse target positions.
C A DDITIONAL A BLATIONS C. I MPORTANCE OF INVERSE DYNAMICS AND REWARD PREDICTION
In addition to forward dynamics prediction, the proposed CoRe model includes reward prediction and inverse dynamics prediction as auxiliary tasks.
In this section, we present comparisons to ablated versions of our model that remove one or both of these tasks.
In Figure 9a we can see that removing both tasks (Fwd only) is signiﬁcantly worse than having both (Fwd + inv + reward).
Having any one of these alone is a big improvement on all tasks except reacher, where adding reward is much more important than inverse dynamics.
It is interesting to see that in the absence of reward, adding inverse dynamics prediction (Fwd + inv) improves over having forward dynamics only.
Asking the model to predict the action that takes the agent from one state to the next is a different way of expressing the model dynamics, compared to predicting the next state given the current state and action.
The fact that asking the model to do both simultaneously gives a boost in performance indicates that inverse dynamics prediction shapes the latent state in ways that are complementary to forward dynamics.
C. I MPORTANCE OF UPDATING Θ USING CRITIC LOSS
In our model we optimize the parameters Θ of the world model (observation encoder and RSSM) using the critic loss J Q .
This is similar to the choice made in SLAC (Lee et al., 2020) and DBC (Zhang
[Table]
et al., 2021).
However
[Table]
a reasonable alternative could be to optimize Θ only using J M and keep the critic training separate.
This would amount to separating the world model from the controller.
As shown in Figure 9b
[Table]
when excluding Θ from the critic vs including it
[Table]
exclusion performs comparably on two tasks (ball in cup, cartpole), better on one (ﬁnger) and worse on three (cheetah, reacher, and walker).
Overall, the inclusion regime works better.
Therefore, we chose to include Θ in the critic.
In future work, an important direction to explore is the exclusion regime, especially in the multi-task setting, because separating the controller from the world model enables separating general understanding of the world from task-speciﬁc control policies, which is key to generalization.
[Figure]
C. R OBUSTNESS TO β WHEN RECONSTRUCTING FROM PRIOR VS POSTERIOR
CoRe predicts the next observation’s feature from the prior latent state ˆ z t , making it different from sequential VAE-based models like SLAC (Lee et al., 2020) and Dreamer (Hafner et al., 2020) where the posterior latent state is used to reconstruct the observation.
We argued in the paper that doing so makes the model more robust to the choice of β , the weight applied to the KL-term during optimization.
In this experiment, we verify this argument by comparing CoRe with a variant where the posterior latent state is decoded to contrastively match the true observation’s representation.
We train both models with ﬁve values of β and 5 different seeds each.
In Figure 10 we can see that the posterior version (bottom row) has more variance in performance across different values of β .
The proposed CoRe model (top row) is more stable, and hence, easier to train.