Robust Robotic Control from Pixels using Contrastive Recurrent State-Space Models

Modeling the world can benefit 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 difficulty 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 . 1 INTRODUCTION . 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 inefficient . 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 . ( 2021 ) ) or future action sequences ( PSE , Agarwal et al . ( 2021 ) ) when running the optimal policy . However , using rewards alone for grounding is sample inefficient , 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 Under review as a conference paper at ICLR 2022CoRe | Visualization Binder Networks | Progress 06/07 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 . ( 2021 ) . 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 : Contrastive Recurrent 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 modifications . 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 . 2 MODEL DESCRIPTION . 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 zt . It receives an observation ot , takes action at , and moves to state zt+1 , receiving a reward rt+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 find a policy that maximizes the expected sum of future rewards ∑∞ t=0 γ trt , where γ is a discount factor . 2.1 MODEL COMPONENTS . 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 xt = f ( ot ) . 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 p ( ẑt|zt−1 , at−1 ) , Observation representation decoder yt = g ( ẑt ) , Correction Model q ( zt|ẑt , xt ) , Inverse Dynamics model ât−1 = a ( zt , zt−1 ) , Reward predictor r̂t = r ( zt ) . 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 yt 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 xt , along with the prior ẑt and outputs the posterior distribution zt , i.e. , it corrects the prior belief given the new observation . The posterior state zt 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 zt−1 to zt . 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 zt consists of a deterministic and stochastic component ( More details in Appendix A ) . Behavior Model : This component consists of an actor πφ ( at|zt ) and two critic networks Qθi ( zt , a ) , i ∈ { 1 , 2 } which are standard for training using the SAC algorithm ( Haarnoja et al. , 2018 ) . 2.2 REPRESENTATION LEARNING WITH CORE . 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 ot and the model ’ s prediction of the observation ’ s encoding yt . 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 defined 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 xi , t and yi , t denote the real observation ’ s encoding and the predicted observation encoding , respectively . The loss function is Lc = − ∑ i , t log ( exp ( λx > i , tyi , t ) ∑ i′ , t′ exp ( λx > i , tyi′ , t′ ) ) − ∑ i , t log ( exp ( λx > i , tyi , t ) ∑ i′ , t′ exp ( λx > i′ , t′yi , t ) ) , ( 1 ) where λ is a learned inverse temperature parameter and i′ and t′ 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 LKL = KL ( q ( zt|ẑt , xt ) ||p ( ẑt|zt−1 , at−1 ) ) . ( 2 ) An additional source of training signal is the reward prediction loss , ( rt−r̂t ) 2 . Furthermore , modeling the inverse dynamics , i.e . predicting at from zt and zt+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 JM ( Θ ) = Lc + βLKL + αr ( rt − r̂t ) 2 + αa||at − ât||2 , 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 ||ot − ôt||2 and LKL 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 ( zt|ẑt , xt ) is conditioned on the true observation xt , 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 .

Summary: The paper proposes a recurrent state-space model that learns robust representations for robotic control. The proposed method builds on top of prior works on world-models which learn a latent dynamics model of the agent which can be used for planning and action selection. Different from prior work such as Dreamer and SLAC which rely on pixel-based observation reconstruction, this paper highlights that a simpler contrastive loss for the next-observation prediction achieves better results **if** a recurrent state-space model is used for the latent space. Results are presented on the Distracting Control Suite benchmark and show strong improvements over prior approaches.