M UTUAL I NFORMATION M AXIMIZATION FOR R OBUST P LANNABLE R EPRESENTATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
Extending the capabilities of robotics to real-world complex, unstructured environments requires the capability of developing better perception systems while maintaining low sample complexity.
When dealing with high-dimensional state spaces, current methods are either model-free, or model-based with reconstruction based objectives.
The sample inefﬁciency of the former constitutes a major barrier for applying them to the real-world.
While the latter present low sample complexity, they learn latent spaces that need to reconstruct every single detail of the scene.
Real-world environments are unstructured and cluttered with objects.
Capturing all the variability on the latent representation harms its applicability to downstream tasks.
In this work, we present mutual information maximization for robust plannable representations (MIRO), an information theoretic representational learning objective for model-based reinforcement learning.
Our objective optimizes for a latent space that maximizes the mutual information with future observations and emphasizes the relevant aspects of the dynamics, which allows to capture all the information needed for planning.
We show that our approach learns a latent representation that in cluttered scenes focuses on the task relevant features, ignoring the irrelevant aspects.
At the same time, state-of-the-art methods with reconstruction objectives are unable to learn in such environments.
I NTRODUCTION
A fundamental challenge in applying reinforcement learning (RL) to real robotics is the need to deﬁne a suitable state space representation.
Designing perception systems manually makes it difﬁcult 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 inefﬁcient (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 cannot 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 ( M utual I nformation
Maximization for RO bust 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.
R ELATED W ORK
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 efﬁcient (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 difﬁculty 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 classiﬁcation, 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.
B ACKGROUND
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 ﬁnite partially observable Markov decision process (POMDP)
M is deﬁned by the tuple ( S , A , T, R, O , O, γ, ρ 0 , H ) .
Here, S is the set of states, A the action space, T ( s t +1 | s t , a t ) the transition distribution ( p ( s t +1 | s t , a t ) ), R ( s t is the probability of obtaining the reward r t at the state s t ( p ( r t | s t ) ), O is the observation space,
O ( o t | s t ) ( p ( o t | s t ) ), γ the discount factor, ρ 0 : S → R + represents the initial state distribution, and
H is the horizon of the process.
We deﬁne the return as the sum of rewards r t along a trajectory τ :=
( s 0 , a 0 , ..., s H − 1 , a H − 1 , s H ) .
The goal of reinforcement learning is to ﬁnd a controller π :
S ×
A → R + that maximizes the expected return, i.e.: max π J ( π ) =
E a t ∼ π s t ∼ p
[ (cid:80)
Ht
=1 γ t r t ] .
Mutual Information.
The mutual information between two random variables X and Y , denoted by
I ( X ; Y ) is a reparametrization-invariant measure of dependency.
Speciﬁcally, it characterizes the
Kullback-Leibler divergence between the joint distribution ( X, Y ) and the product of the marginals
X and Y : I ( X ; Y ) =
E p ( x,y ) (cid:104) log p ( x | y ) p ( x ) (cid:105) =
E p ( x,y ) (cid:104) log p ( x,y ) p ( x ) p ( y ) (cid:105) =
D KL (( X, Y ) (cid:107) 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,
I NCE van den Oord et al.
R OBUST P LANNABLE R EPRESENTATIONS
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 ﬁrst 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 ﬁnally we instantiate this objective in a concrete model-based algorithm.
4. R OBUSTNESS OF L EARNED L ATENT S PACES
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):
[Equation]
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 ( o 1: t ; o t + h | a t : 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):
[Equation]
Here, the expectation is on sets of observtions { o i 1 , ..., o h , ..., o i N +1 } with N negative examples and one positive example, o h , 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. L EARNING L ATENT S PACES FOR C ONTROL
[Equation] o [Equation] o
[Equation]
o
[Equation] o [Equation] 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 o t and rewards r t , and we are trying to model the underlying latent space s t and the underlying process that generates the observed variables.
We model the latent space s t and reward r t to be Gaussian with diagonal covariance given the previous state and action, i.e., s t +1 | s t , a t ∼ N ( µ t +1 , Σ t +1 ) and r t | s t , a t ∼ 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, ﬁlter, 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 (ﬁlter), 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 ⊆ R m .
We represent the encoded observation e θ ( o t ) with the notation z t ∈ Z .
Contrary to prior work in learning latent spaces for planning, we do not make any assumption on the underlying distribution of z t or o t (such as they are Gaussian (Kingma & Welling, 2013)).
Filter.
This function, g ζ : Z × S → S , ﬁlters the belief of the prior ˆ s t state variable with the current encoded observation z t resulting in the ﬁltered variable s t .
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:
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.
Instead, we optimize a tractable lower bound of this objective.
First, we replace the mutual information term with the noise constrative estimator lower bound I NCE .
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
[Equation]
1: Initialize ﬁlter
T ← ∅ 4: for
K do 5:
Sample initial state s 0 ∼ ρ6: for t = 0 . . .
H do 7:
Filter the state ˆ s t with o t using g ζ to get s t 8:
Get action a t from planner using s t , f φ and r ψ 9:
Take action a t in the environment and obtain o t +1 , r t 10:
T ← T ∪ { s t , a t , r t , o t } 11:
Estimate next state ˆ s t +1 ∼ f φ ( s t , a t ) 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 Θ
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 x t = [ s t ; a t : t + h ] :
[Equation]
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.
R OBUST L ATENT S PACE M ODEL -B ASED R EINFORCEMENT L EARNING
In this section, we present our architecture, implementation, and instantiation into a model-based method.
The presented algorithm uses planning, but our representational learning method for control is agnostic to the derived controller.
Methods such as Lee et al.
(2019) could be used with MIRO.
The pseudocode for the model-based approach is shown in Alg. 1, and is typical in model-based
RL algorithms, it iterates between three steps: 1) data collection, 2) model learning, and 3) policy improvement or planning.
Data Collection.
At each iteration, we ﬁrst collect data by planning with the current dynamics model.
On-policy data collection relieves the model from learning the entire space, and instead it just focuses on the regions that the agents visits.
It also allows to overcome the insufﬁcient coverage of the initial data distribution.
The data collected is stored in a replay buffer that is used for training the model.
Model Learning.
The architecture of the model is shown in Figure 3.
It is composed by an convolutional encoder, a feed forward ﬁlter, and a dynamics model that has the form of a recurrent neural network.
The reward predictor, not shown in Figure 3, is modelled as a feed forward neural network that has as input the latent state variable.
The model is trained using all the data collected so far on the MIRO objective (Equation 1).
Planning.
In this work, we use model-predictive control (MPC) with cross-entropy method (CEM) for action selection (Botev et al., 2013).
The CEM algorithm selects the action sequence that maximizes the expected return under the learned dynamics.
Speciﬁcally, it is a population based procedure that iteratively reﬁts a Gaussian distribution, starting from a unit Gaussian, on the best sequences of actions.
The MPC component prevents overﬁtting to the model deﬁciencies by selecting the best CEM action and replanning at each step.
Note that our represenational learning objective is orthogonal to the planning method used.
R ESULTS
In this section, we empirically corroborate the claims in the previous sections.
Speciﬁcally, the experiments are designed to address the following questions:
Is our approach able to maintain its performance in front of distractors in the scene?
How does our method compare with state-of-the-art reconstruction objectives?
Does our latent variable remain unchanged in the presence of visual noise?
To answer the posed questions, we evaluate our framework, in four continuous control benchmark tasks MuJoCo simulator: cartpole-balance, reacher, ﬁnger-spin and half-cheetah (Todorov et al., 2012; Tassa et al., 2018).
We choose PlaNet (Hafner et al., 2018) as the state-of-the-art reconstruction objective baseline for plannable representations.
6. R OBUST C ONTINUOUS C ONTROL C OMPARISON
To test the robustness of MIRO and PlaNet in visually noisy environments, we add distractors to each of the four environments, as shown in Fig. 4.
[Figure]
The results, depicted in Figure 5, show that in all four environments the performance of MIRO is not undermined by the presence of objects unrelated to the task.
Interestingly, in the half-cheetah environment, the performance even improves with visual noise in the background.
We argue that the presence of visual noise forces the embedding to focus even more on information relevant to the task, and thus makes the embedding more suitable for planning.
In comparison, the performance of PlaNet struggles in face of distractors: the agent is unable to learn in the reacher environment, and struggle to learn in the other environment obtaining sub-optimal performance.
When no distractors are present in the scene, Figure 5, shows that our approach achieves comparable or superior performance than its reconstruction-based counterpart. 6. R EALISTIC D ISTRACTORS
In the previous section, the distractors added to the scene is placed at a random position at each time step.
While this resembles particular types of random noise, the movement of majority of objects in real life do not conform to such behavior. To simulate a more realistic visual noise, we add the same geometric objects in 4 to the scene but move the i -th distractor according to the following linear rule: [Equation] where l is size of each step, ˜u i is a unit vector denoting the direction of the movement and it is resampled with probability p = 0 .
2 at each timestep.
We term the resulting distractors “linear distractors" and term the distractors Figure 4 “random distractors".
We compared two ways of obtaining a dynamics model on tasks with linear distractors: 1.
Transfer :
Train with random distractors and test on linear distractors; 2.
From scratch :
Train with linear distractors and test on linear distractors.
The performance of both MIRO and PlaNet on two schemes are presented in table.
The results showcase that MIRO, when presented with realistic distractors, is still able to perform the task.
Since MIRO is maximizing an MI objective, the latent space is more robust to the distractors when it is exposed to higher entropy distractors at training time.
PlaNet (Transfer)
MIRO (Transfer)
PlaNet (From scratch)
MIRO (From scratch)
Cartpole Balance 615 ±984 ±583 ±780 ±
Reacher 139 ±673 ±274 ±551 ±
{+Finger 
+}
±784 ±582 ±648 ±
Half Cheetah 308 ±615 ±211 ±576 ±
[Figure] [Figure] Here, we provide insight into robustness of latent learned by MIRO.
An ideal latent space should only encode information essential to the dynamics of the task.
Thus, given observations o clean and o dis , the later being the same as the former but with added distractors in the background, the desired behaviour of our encoder would be to output the same latent varible; i.e., z clean ≈ z dis , being z clean the encoded o clean ( e θ ( o clean ) ) and z dis the encoded o dis
( e θ ( o dis ) ).
In order to quantify the discrepancy on the latent, we measure its normalized difference
[Equation]
[Figure] [Equation] Figure 6 shows the average normalized difference on the Cartpole Balance environment when the latent space is learned using the MIRO objective and a sequential VAE objective (Hafner et al., 2018; Lee et al., 2019; Ha &
Schmidhuber, 2018).
The results manifest that the embedding learned by MIRO excludes irrelevant information over the course of training.
In contrast, with sequential VAE, the normalized error increases and lands at higher ﬁnal value, which indicates that the model fails to understand the essential elements in the task despite being trained jointly with the dynamics.
C ONCLUSION
In this paper we present MIRO, an information theoretic representational learning objective that tackles the problem of learning robust and generalizable representations for control.
Our approach maximizes the mutual information between the learned latent space and future observations, emphasizing the features that are relevant for the dynamics and the reward predictions.
We compare our approach in standard benchmark control environments with and without visual disturbances.
Our method remains unaffected by those disturbances while state-of-the-art methods fail to complete the tasks.
When disturbances are removed it attains comparable or superior performance to prior work.
Furthermore, we experimentally investigate the effect of such disturbances in the latent variables, showing that MIRO’s latents are invariant to visual noises.
The development of this method has been motivated for the applicability of general reinforcement learning to real-world and unstructured environments.
Consequently, an enticing direction for future work would be to analyze MIRO’s behaviour in real-robotic agents outside structured environments  A PPENDIX 8. H YPERPARAMETERS
Environment: Dimension of the input observation is 32 × 32 × 3 .
Horizon of all the tasks are
H = 1000 with the same action repeat parameters as in (Hafner et al., 2018).
Network: The dimension of the observation encoding space Z is 1024, dimension of the latent space
S is 30.
Horizon of the CEM planner is 12.
The horizon h of the Mutual Information Objective is 3.
Number of negative examples for NCE objective is N = 10 .
Loss Function:
In equation 1, the coefﬁcient for dynamics model loss is set to λ 1 = 0 .
002 , the coefﬁcient for reward prediction loss is set to λ 2 = 0 .
A BLATION S TUDIES
In the following we present the ablation studies.
8.2. M UTUAL I NFORMATION O BJECTIVE
To study the effect of the mutual information objective in the loss function, we ablated the MI objective in 1 and compared its performance with our full model (MIRO).
As shown in 7, using solely dynamics model loss and reward prediction loss, planning fails on all four environments.
This shows that the MI objective is a central part of the algorithm.
[Figure]
[Figure]
M ODEL A RCHITECTURE
In this section we study the effect two architecture choices of the recurrent neural network used in the model.
We compare the performance when the latent state is stochastic (as used in MIRO) and when the latent state is deterministic as shown in Figure 9.
With deterministic latent space, the performance hardly improves.
This could be due to the fact that deterministic latents are more prone to overﬁt.
The observation is consistent with (Hafner et al., 2018).
[Figure]
8.2. M UTUAL I NFORMATION C ONDITIONED ON A CTIONS
In this section we compare two different ways of incorporating actions in the mutual information objective.
When predicting z t + k from s t , we can incorporate actions a t , · · · , a t + k − 1 in two ways: 1.
stacking the actions and concatenating them with s t when predicting z t + k ; 2. execute an open-loop rollout from s t given a t , · · · , a t + k − 1 and get ˆ s t + k , and predict z t
+ k from ˆ s t + k .
For the ﬁrst variant, we also examined the effect of different horizon h .
[Figure]