Mutual Information Maximization for Robust Plannable Representations

1 INTRODUCTION . A fundamental challenge in applying reinforcement learning ( RL ) to real robotics is the need to define a suitable state space representation . Designing perception systems manually makes it difficult to apply the same reinforcement learning algorithm to a wide range of tasks . Real-world environments are unstructured , cultured , and changing . A suitable representation has to capture features relevant for the task , while ignoring the other elements in the scene . How do we impose inductive biases in the learned representation for real-world robotics ? One could , in principle , apply model-free reinforcement algorithms from raw , low-level observations ; however , these tend to be slow , sensitive to hyperparameters , and sample inefficient ( Lee et al. , 2019 ) . Current work has focused on learning latent spaces for model-based RL constraining them to be feature points ( Finn et al. , 2016 ) , or learning them with reconstruction objectives ( Hafner et al. , 2018 ; Watter et al. , 2015 ; Ha & Schmidhuber , 2018 ) . While feature points objectives provide a good inductive bias , they are too rigid for unstructured tasks and can not incorporate semantic knowledge of the environment . Learned representations with reconstruction objectives ( Kingma & Welling , 2013 ) force the model to capture all the diversity of the environment in the learned representation in order to accurately reconstruct the observations . As a result , in environments with variability , the relevant features for the task either represent a small component of the latent space or are ignored , both of which prevents deriving a controller from them . This work tackles the problem of representation learning from an information theoretic point of view : learning a latent space that maximizes the mutual information between the latent and future observations ( van den Oord et al. , 2018 ; Hénaff et al. , 2019 ) . The mutual information objective induces an inductive bias towards the shared information with future observations , which allows the latent space to discard low-level information and local noise . Hence , by foregoing the reliance on reconstruction objectives , we obtain a latent space that is robust to the presence of variations on the scene if they are unrelated to the task . The main contribution of our work is a representational learning approach , MIRO ( Mutual Information Maximization for RObust Plannable Representations ) , that jointly optimizes for the latent representation and model , and results in a latent space that is robust against disturbances , noisy observations , and achieves performance comparable to state-of-the-art algorithms . We show that previous approaches are brittle even to minor changes in the environment . Finally , we experimentally investigate these effects on the latent variables in our approach and previous work , pinpointing that our learned space is invariant to such changes in the environment . The experimental evaluation is carried out in standard DeepMind Control Suite ( Tassa et al. , 2018 ) environments . 2 RELATED WORK . Model-based Reinforcement Learning from Image Observations . Recent work has shown that model-based reinforcement learning RL is able to achieve the same performance as model-free methods while being substantially more sample efficient ( Wang et al. , 2019 ) . However , such achievements have been limited to envitonments where a compact state space representation is available , and extending them to raw sensory high-dimensional input spaces poses a major challenge . In such cases , learning a compact and accurate latent space is crucial to relieve the dynamics model from directly modeling the raw sensor input space . Recent work on high-dimensional observations on model-based RL can be categorized in two main classes : 1 ) video prediction models ( Jayaraman et al. , 2018 ; Ebert et al. , 2018 ; Kaiser et al. , 2019 ) , and 2 ) latent space learning with reconstruction objectives ( Hafner et al. , 2018 ; Lee et al. , 2019 ; Watter et al. , 2015 ; Wahlström et al. , 2015 ; Zhang et al. , 2018 ; Ha & Schmidhuber , 2018 ) . Both video prediction models and latent space with reconstruction objectives share the commonality that the latent space is learned using loss on the raw pixel observations . As a result , the latent space model needs to incorporate all information to reconstruct every detail on the observations , which is redundant as the task is usually represented by a small fraction of the scene in real world environments . Representation Learning . Variational autoencoders ( VAE ) ( Kingma & Welling , 2013 ; Higgins et al . ; van den Oord et al. , 2017 ) learn an embedding by optimizing the variational lower bound of the likelihood of data . The learned latent follows a prior distribution , usually Gaussian , which makes sampling and distance meaningful in the latent space . Our method uses a similar architecture to the extension of VAE to sequential data ( Chung et al. , 2015 ; Lee et al. , 2019 ; Hafner et al. , 2018 ; Fraccaro et al. , 2016 ) . However , instead of a reconstruction objective , we propose to learn a latent space by maximizing the mutual information ( MI ) between the latent and future observations . The omission of a reconstruction objective relieves the dynamics model from encoding unpredictable variation in the scene , which proves to be more robust when the observation is visually noisy . Our method inspired by the InfoMax principle ( Linsker , 1988 ) , learns a latent space by maximizing the mutual information ( MI ) between observation and latent representation . Despite the difficulty of estimating MI directly , recent work has shown that representations learned by maximizing a noise contrastive estimation lower bound of MI are able to achieve state-of-the-art performance in downstream benchmark tasks in various domains , including image classification , video captioning and natural language processing ( van den Oord et al. , 2018 ; Hénaff et al. , 2019 ; Hjelm et al. , 2018 ; Sun et al. , 2019 ; Tian et al. , 2019 ; Bachman et al. , 2019 ) . The maximizing the MI has been applied to the context of reinforcement learning , but not in the context of learning robust representations , but to either speed up learning ( van den Oord et al. , 2018 ) , to achieve hierarchical represenations ( Nachum et al. , 2018 ) , or model the believe of state space ( Gregor et al. , 2019 ; Guo et al. , 2018 ) , which use a similar architecture to our method . ( Anand et al. , 2019 ) use an MI objective to learn representations for Atari games . 3 BACKGROUND . This work tackles the problem of learning a latent space that is suitable for planning from highdimensional observations in POMDPs . In this work , the observations constitute single images , which are not enough to completely capture the dynamics of the system . For instance , a single image does not contain any information of the velocity or an object might be occluded . Partially Observable Markov Decision Process . A discrete-time finite partially observable Markov decision process ( POMDP ) M is defined by the tuple ( S , A , T , R , O , O , γ , ρ0 , H ) . Here , S is the set of states , A the action space , T ( st+1|st , at ) the transition distribution ( p ( st+1|st , at ) ) , R ( st is the probability of obtaining the reward rt at the state st ( p ( rt|st ) ) , O is the observation space , O ( ot|st ) ( p ( ot|st ) ) , γ the discount factor , ρ0 : S → R+ represents the initial state distribution , and H is the horizon of the process . We define the return as the sum of rewards rt along a trajectory τ : = ( s0 , a0 , ... , sH−1 , aH−1 , sH ) . The goal of reinforcement learning is to find a controller π : S ×A → R+ that maximizes the expected return , i.e . : maxπ J ( π ) = Eat∼π st∼p [ ∑H t=1 γ trt ] . Mutual Information . The mutual information between two random variables X and Y , denoted by I ( X ; Y ) is a reparametrization-invariant measure of dependency . Specifically , it characterizes the Kullback-Leibler divergence between the joint distribution ( X , Y ) and the product of the marginals X and Y : I ( X ; Y ) = Ep ( x , y ) [ log p ( x|y ) p ( x ) ] = Ep ( x , y ) [ log p ( x , y ) p ( x ) p ( y ) ] = DKL ( ( X , Y ) ‖X ⊗ Y ) . Estimating and optimizing the mutual information objective poses a challenging problem . In this work , we use a multi-sample unnormalized lower bound based on noise contrastive estimation , INCE van den Oord et al . ( 2018 ) . 4 ROBUST PLANNABLE REPRESENTATIONS . Enabling complex real robotics tasks requires extending current model-based methods to low-level high-dimensional observations . However , in order to do so , we need to specify which space they should operate on . Real-world environments are unstructured , cluttered , and present distractors . Our approach , MIRO , is able to learn latent representations that capture the relevant aspects of the dynamics and the task by framing the representational learning problem in information theoretic terms : maximizing the mutual information between the latent space and the future observations . This objective advocates for representing just the relevant aspects of the dynamics , removing the burden of reconstructing the entire pixel observation . We first motivate the use of the mutual information objective as a representational learning objective for control , then we derive our objective that entangles dynamics and reward learning with the representational objective , and finally we instantiate this objective in a concrete model-based algorithm . 4.1 ROBUSTNESS OF LEARNED LATENT SPACES . Representations that ignore the presence of elements unrelated to the task at hand would allow us to apply successfully off-the-shelf reinforcement learning algorithms in real-world environments . In contrast , current representational learning approaches for control are based on reconstruction objectives . Here we show that these methods are not suitable for real-world tasks , since even in the presence of simple distractors , they completely fail to capture the important aspects of the task . Concretely , several prior work learns representations by reconstructing the observation using a variational auto-encoder ( VAE ) ( Kingma & Welling , 2013 ) : max θ , φ Eqθ ( z|x ) [ log pφ ( x|z ) ] −DKL ( qθ ( z|x ) ‖p ( z ) ) here p ( z ) ∼ N ( 0 , 1 ) , and θ , φ are the parameters of the encoder and decoder , respectively . This objective captures the variation of the input in the latent space while the commonalities are stored in the weights of the auto-encoder . As a result , an unstructured and changing environment will be encoded in the latent space even when this variation is irrelevant to the control problem . The mutual information objective presents itself as an alternative to learn representations . We argue that this objective is more suitable for learning robust representations . In the POMDP setup , we propose to maximize the mutual information between past observations and future observations , i.e. , maxθ I ( o1 : t ; ot+h|at : t+h ) , where θ are the parameters of the encoder ( no decoder is needed ) . Then , given limited capacity of the encoder , the optimal representation should not capture the details on the scene that do not depend on the actions . Including more information in latent space always increases the mutual information , hence limiting the capacity or adding an information bottleneck is crucial for obtaining a robust representation . We estimate the mutual information using InfoNCE lower bound ( van den Oord et al. , 2018 ) : max θ E [ fθ ( o1 : t , ot+h ) ∑ h′ fθ ( o1 : t , ot+h′ ) ] Here , the expectation is on sets of observtions { oi1 , ... , oh , ... , oiN+1 } with N negative examples and one positive example , oh , and fθ represent a score function . The mutual information objective allows us to encode an inductive bias towards the important aspects of the task . The inductive bias arises from the fact that we are optimizing a discriminative model , instead of generative such as the VAE . Figure 1 shows the effect of distractors on the learned representation in a simple point-mass environment . Both representations are learned using the same data , raw pixel images of the point mass with the distractors . Then , we regress on top of the learned representations to the position of the point-mass . The mutual information objective allows to maintain the accuracy even when more distractors are present . However , the autoencoder accuracy is increasingly harmed by the appearance of unrelated objects to the task . 4.2 LEARNING LATENT SPACES FOR CONTROL In this section , we develop the MIRO objective that entangles the representation , dynamics model , and reward predictor . Consequently , the objective not only learns a representation that is robust , but also emphasizes the important aspects of the dynamics and the task . Figure 2 shows the POMDP set-up , where the only observed variables are the observations ot and rewards rt , and we are trying to model the underlying latent space st and the underlying process that generates the observed variables . We model the latent space st and reward rt to be Gaussian with diagonal covariance given the previous state and action , i.e. , st+1|st , at ∼ N ( µt+1 , Σt+1 ) and rt|st , at ∼ N ( νt , Λt ) . In order to learn the underlying process of the POMDP and successfully use this representation for control , we learn four functions : an encoder , filter , transition dynamics , and reward predictor . In order to plan , we need to process the current observation ( encoder ) , update the latent space with the current observation ( filter ) , predict the next states ( transition dynamics ) and predict the future rewards ( reward predictor ) . In the following , we describe in detail the parametrization , and functionality of each of these components : Encoder . The encoder parametrizes a map from a high-dimensional observation to a lower dimensional manifold , eθ : O → Z ⊆ Rm . We represent the encoded observation eθ ( ot ) with the notation zt ∈ Z . Contrary to prior work in learning latent spaces for planning , we do not make any assumption on the underlying distribution of zt or ot ( such as they are Gaussian ( Kingma & Welling , 2013 ) ) . Filter . This function , gζ : Z ×S → S , filters the belief of the prior ŝt state variable with the current encoded observation zt resulting in the filtered variable st . Essentially , it updates the latent to be coherent with the most recent observation . Dynamics model . The dynamics model approximates the transition function on the learned latent space , fφ : S ×A → S . The dynamics parametrizes the mean µt and diagonal covariance Σt of the underlying transition distribution of the latent process . We denote by ŝt+1 the prior latent variable resulting from sampling from the dynamics model distribution . Reward predictor . The reward predictor regresses onto the observed rewards from the latent variable . Hence , it is a mapping from latent states to rewards , rψ : S → R. As said , we assume that the rewards follow a Gaussian distribution with unit variance . We denote by Θ = { θ , ζ , φ , ψ } the parameters of our function approximators . These are learned altogether with the following constrained optimization objective : max Θ I ( st ; ot+h|at : t+h ) s.t . : DKL ( gζ ( zt+1 , ŝt+1 ) ‖st+1 ) ≤ DKL ( rt‖rψ ( st , at ) ) ≤ As seen in the previous section , the mutual information maximization term allows us to obtain latent spaces that just contain information about the relevant components of the scene . The enforcement of the constraints has two purposes : 1 ) ensures small prediction errors when planning , and 2 ) guides the state space to capture the components for predicting rewards and future states . In practice , the previous objective is intractable to optimize . The exact maximization , or even evaluation , of the mutual information term is intractable , as well as the set of non-linear inequalities . Algorithm 1 Model-Based Reinforcement Learning with MIRO 1 : Initialize encoder eθ , filter gζ , dynamics model fφ , reward predictor rψ and dataset D ← ∅ . 2 : repeat 3 : T ← ∅ 4 : for k = 1 . . .K do 5 : Sample initial state s0 ∼ ρ0 6 : for t = 0 . . . H do 7 : Filter the state ŝt with ot using gζ to get st 8 : Get action at from planner using st , fφ and rψ 9 : Take action at in the environment and obtain ot+1 , rt 10 : T ← T ∪ { st , at , rt , ot } 11 : Estimate next state ŝt+1 ∼ fφ ( st , at ) 12 : end for 13 : D ← D ∪ T 14 : end for 15 : Optimize Equation 1 with mini-batches from dataset D 16 : until Until desired performance is achieved 17 : return Optimal model parameters Θ Instead , we optimize a tractable lower bound of this objective . First , we replace the mutual information term with the noise constrative estimator lower bound INCE . Second , we formulate its Lagrangian and treat the dual parameters as hyperparameters . Altogether , it gives an optimization objective that can be easily optimized with stochastic gradient descent max Θ INCE ( st ; ot+h|at : t+h ) − λ1DKL ( gζ ( zt , ŝt+1 ) ‖st+1 ) − λ2DKL ( rt‖rψ ( st , at ) ) ( 1 ) We optimize this objective by using Monte-Carlo estimates and backpropagating through the stochastic nodes with the path-wise derivative , also known as reparametrization trick ( Kingma & Welling , 2013 ) . As in van den Oord et al . ( 2018 ) , we use a log-bilinear model for the score function , which leads to the following formulation of the InfoNCE objective . Be xt = [ st ; at : t+h ] : INCE ( st ; ot+h|at : t+h ) = E [ log ( exp ( x > t Whzt+h ) exp ( x > t Whzt+h ) + ∑N j=1 exp ( x > t Whzj ) ) ] ( 2 ) Even though more complicated models can be used , as well as models that encode the future sequences of actions ; This simple score model tends to work well in practice .

Summary: The authors propose a model-based reinforcement algorithm in which the model is a sequential latent variable model and the actions planned with a cross-entropy method (CEM) planner. The model is learnt by maximizing a lower bound on the mutual information between the latent states and their successor observations (instead of the classical sequential ELBO). The authors argue that the latter objective function yield robustness to distraction in visual scenes. The algorithm, named MIRO, is experimented on 4 simulated environments.