INFO RMATION P RIORITIZATION THROUGH E M POWER MENT IN V ISUAL M ODEL - BASED RL
Anonymous authors
Paper under double-blind review
A BSTRACT
Model-based reinforcement learning (RL) algorithms designed for handling complex visual observations typically learn some sort of latent state representation, either explicitly or implicitly.
Standard methods of this sort do not distinguish between functionally relevant aspects of the state and irrelevant distractors, instead aiming to represent all available information equally.
We propose a modiﬁed objective for model-based RL that, in combination with mutual information maximization, allows us to learn representations and dynamics for visual model-based
RL without reconstruction in a way that explicitly prioritizes functionally relevant factors.
The key principle behind our design is to integrate a term inspired by variational empowerment into a state-space learning model based on mutual information.
This term prioritizes information that is correlated with action, thus ensuring that functionally relevant factors are captured ﬁrst.
Furthermore, the same empowerment term also promotes faster exploration during the RL process, especially for sparse-reward tasks where the reward signal is insufﬁcient to drive exploration in the early stages of learning.
We evaluate the approach on a suite of vision-based robot control tasks with natural video backgrounds and show that the proposed prioritized information objective outperforms state-of-the-art model based RL approaches with higher sample efﬁciency and episodic returns
I NTRODUCTION
Model-based reinforcement learning (RL) provides a promising approach to accelerating skill learning: by acquiring a predictive model that represents how the world works, an agent can quickly derive effective strategies, either by planning or by simulating synthetic experience under the model.
However, in complex environments with high-dimensional observations (e.g., images), modeling the full observation space can present major challenges.
While large neural network models have made progress on this problem (Finn & Levine, 2017; Ha & Schmidhuber, 2018; Hafner et al., 2019a;
Watter et al., 2015; Babaeizadeh et al., 2017), sample-efﬁcient learning necessitates some mechanism to prioritize modeling latent representations from observations such that functionally-relevant factors for the task can be captured.
This needs to be done without wasting effort and capacity on irrelevant distractors, and without detailed reconstruction.
Several recent works have proposed contrastive objectives that maximize mutual information between observations and latent states (Hjelm et al., 2018;
Ma et al., 2020; Oord et al., 2018; Srinivas et al., 2020).
While such objectives avoid reconstruction, they still do not distinguish between relevant irrelevant factors of variation.
We thus pose the question: can we devise non-reconstructive representation learning methods that explicitly prioritize information that is most likely to be functionally relevant to the agent?
In this work, we derive a model-based RL algorithm from a combination of representation learning via mutual information maximization (Poole et al., 2019) and empowerment (Mohamed & Rezende, 2015).
The latter serves to drive both the representation and the policy toward exploring and representing functionally relevant factors of variation.
By integrating an empowerment-based term into a mutual information framework for learning state representations, we effectively prioritize information that is most likely to have functional relevance, which mitigates distractions due to irrelevant factors of variation in the observations.
By integrating this same term into policy learning, we further improve exploration, particularly in the early stages of learning in sparse-reward environments, where the reward signal provides comparatively little guidance.
Our main contribution is InfoPower , a model-based RL algorithm for high-dimensional systems with image observations that integrates empowerment into a mutual information based, nonreconstructive framework for learning state space models.
Our approach explicitly prioritizes information that is most likely to be functionally relevant, which signiﬁcantly improves performance in the presence of time-correlated distractors (e g., background videos) and also accelerates exploration in environments when the reward signal is weak.
We evaluate the proposed objectives on a suite of simulated robotic control tasks with explicit video distractors and demonstrate up to 20% better performance in terms of cumulative rewards at 1M environment interactions with 30% higher sample efﬁciency at 100k interactions.
P ROBLEM S TATEMENT AND N OTATION
A partially observed Markov decision process (POMDP) is a tuple ( S , A , T, R, O ) that consists of states s ∈ S , actions a ∈ A , rewards r ∈ R , observations o ∈
O , and a state-transition distribution T ( s (cid:48) | s, a ) .
In most practical settings, the agent interacting with the environment doesn’t have access to the actual states in S , but to some partial information in the form of observations O .
The underlying state-transition distribution T and reward distribution R are also unknown to the agent.
In this paper, we consider the observations o ∈ O to be high-dimensional images, and so, the agent should learn a compact representation space Z for the latent state-space model.
The problem statement is to learn effective representations from observations O when there are visual distractors present in the scene, and plan using the learned representations to maximize the cumulative sum of discounted rewards,
J = E
[ (cid:80) t γ t − 1 r t ] .
The value of a state V ( Z t ) is deﬁned as the expected cumulative sum of discounted rewards starting at state Z t .
We use q ( · ) to denote parameterized variational approximations to learned distributions.
We denote random variables with capital letters and use small letters to denote particular realizations (e.g., z t denotes the value of Z t ).
Since the underlying distributions are unknown, we evaluate all expectations through Monte-Carlo sampling with observed state-transition tuples ( o t , a t − 1 , o t − 1 , z t , z t − 1 , r t ) .
I NFORMATION P RIORITIZATION FOR THE L ATENT S TATE -S PACE M ODEL
Our goal is to learn a latent state-space model with a representation Z that prioritizes capturing functionally relevant parts of observations O , and devise a planning objective that explores with the learned representation.
To achieve this, our key insight is integration of empowerment in the visual model-based RL pipeline.
For representation learning we maximize MI max Z I ( O , Z ) subject to a prioritization of the empowerment objective max
Z I ( A t − 1 ; Z t | Z t − 1 ) .
For planning, we maximize the empowerment objective along with reward-based value with respect to the policy max
A I ( A t − 1 ; Z t | Z t − 1 )
+ I ( R t ; Z t ) .
In the subsequent sections, we elaborate on our approach, InfoPower , and describe lower bounds to MI that yield a tractable algorithm.
3. L EARNING C ONTROLLABLE F ACTORS AND P LANNING THROUGH E MPOWERMENT
Controllable representations are features of the observation that correspond to entities which the agent can inﬂuence through its actions.
For example, in quadrupedal locomotion, this could include the joint positions, velocities, motor torques, and the conﬁgurations of any object in the environment that the robot can interact with.
For robotic manipulation, it could include the joint actuators of the robot arm, and the conﬁgurations of objects in the scene that it can interact with.
Such representations are denoted by S + in Fig.
2, which we can formally deﬁne through conditional independence as the smallest subspace of S ,
S + ≤ S , such that I ( A t − 1 ; S t | S + t ) = 0 .
This conditional independence relation can be seen in Fig.
We explicitly prioritize the learning of such representations in the latent space by drawing inspiration from variational empowerment (Mohamed & Rezende, 2015).
The empowerment objective can be cast as maximizing a conditional information term
I ( A t − 1 ; Z t | Z t − 1 ) =
H ( A t − 1 | Z t − 1 ) − H ( A t − 1 | Z t , Z t − 1 ) .
The ﬁrst term H ( A t − 1 | Z t − 1 ) encourages the chosen actions to be as diverse as possible, while the second term −H ( A t − 1 | Z t , Z t − 1 ) encourages the representations Z t and Z t +1 to be such that the action A t for transition is predictable.
While prior approaches have used empowerment in the model-free setting to learn policies by exploration through intrinsic motivation (Mohamed & Rezende, 2015), we speciﬁcally use this objective in combination with MI maximization for prioritizing the learning of controllable representations from distracting images in the latent state-space model. Initialize dataset D with random episodes. Initialize model parameters φ, χ, ψ, η . Initialize dual variable λ .
while not converged do for update step c = 1 ..C do // Model learning Sample data { ( a t , o t , r t ) } k + L t = k ∼ D . Compute latents z t ∼ p φ ( z t | z t − 1 , a t − 1 , o t ) .
Calculate L based on section 3.4.
( φ, χ, ψ, η ) ← ( φ, χ, ψ, η ) + ∇ φ,χ,ψ,η L λ ← λ − ∇ λ L //
Behavior learning
Rollout latent plan, S ← S ∪ { z t , a t , r t }
V ( z t ) ≈ E π [ln q η ( r t | z t )+ln q ψ ( a t − 1 | z t , z t − 1 )]
Update policy π and value model end //
{+Environment interaction for time step t = 
+}
T − 1 do z t ∼ p φ ( z t | z t − 1 , a t − 1 , o t ) ; a t ∼ π ( a t | z t ) r t , o t +1 ← env.step( a t ) .
end
Add data D ← D ∪ { ( o t , a t , r t )
Tt =1 } . end We include the same empowerment objective in both representation learning and policy learning. For this, we augment the maximization of the latent value function that is standard for policy learning in visual model-based RL (Sutton, 1991), with max A I ( A t − 1 ; Z t | Z t − 1 ) .
This objectives complements value basedlearning and further improves exploration by seeking controllable states.
We empirically analyze the beneﬁts of this in sections 4.3 and 4.6.
In the Appendix A.1 we next describe two theorems regarding learning controllable representations.
We observe that the max (cid:80) t I ( A t − 1 ; Z t | Z t − 1 ) objective alone for learning latent representations Z , along with the planning objective provably recovers controllable parts of the observation
O , namely S + .
This result in Theorem 1 is important because in practice, we may not be able to represent every possible factor of variation in a complex environment.
In this situation, we would expect that when | Z | (cid:28) |O| , learning Z under the objective max (cid:80) t I ( A t − 1 ; Z t | Z t − 1 ) would encode S + .
We further show through Theorem 2 that the inverse information objective alone can be used to train a latent-state space model and a policy through an alternating optimization algorithm that converges to a local minimum of the objective max (cid:80) t
I ( A t − 1 ; Z t | Z t − 1 ) at a rate inversely proportional to the number of iterations.
In Section 4.3 we empirically show how this objective helps achieve higher sample efﬁciency compared to pure value-based policy learning.
3. M UTUAL I NFORMATION MAXIMIZATION FOR R EPRESENTATION L EARNING
For visual model-based RL, we need to learn a representation space Z , such that a forward dynamics model deﬁning the probability of the next state in terms of the current state and the current action can be learned.
The objective for this is (cid:80) t − I ( i t ; Z t | Z t − 1 , A t − 1 ) .
Here, i t denotes the dataset indices that determine the observations p ( o t | i t ) =
δ ( o t − o t (cid:48) ) .
In addition to the forward dynamics model, we need to learn a reward predictor by maximizing (cid:80) t
I ( R t ; Z t ) , such that the agent can plan ahead in the future by rolling forward latent states, without having to execute actions and observe rewards in the real environment.
Finally, we need to learn an encoder for encoding observations O to latents Z .
Most successful prior works have used reconstruction-loss as a natural objective for learning this encoder (Babaeizadeh et al., 2017; Hafner et al., 2019b;a).
A reconstruction-loss can be motivated by considering the objective I ( O , Z ) and computing its BA lower bound (Agakov, 2004).
I ( o t ; z t ) ≥ E p ( o t ,z t ) [log q φ (cid:48)
( o t | z t )]
+ H ( p ( o t )) .
The ﬁrst term here is the reconstruction objective, with q φ (cid:48)
( o t | z t ) being the decoder, and the second term can be ignored as it doesn’t depend on Z .
However, this reconstruction objective explicitly encourages encoding the information from every pixel in the latent space (such that reconstructing the image is possible) and hence is prone to not ignoring distractors.
In contrast, if we consider other lower bounds to I ( O , Z ) , we can obtain tractable objectives that do not involve reconstructing high-dimensional images.
We can obtain an NCE -based lower bound (Hjelm et al., 2018):
I ( o t ; z t ) ≥ E q φ ( z t | o t ) p ( o t ) [log f θ ( z t , o t ) − log (cid:80) t (cid:48)
(cid:54) =
t f θ ( z t , o t (cid:48) )] , where q φ ( z t | o t ) is the learned encoder, o t is the observation at timestep t (positive sample), and all observations in the replay buffer o t (cid:48) are negative samples.
f θ ( z t , o t (cid:48) )
=
exp ( z Tt W θ z t (cid:48) )
The lower-bound is a form of contrastive learning as it maximizes compatibility of z t with the corresponding observation o t while minimizing compatibility with all other observations across time and batch.
Although prior work has explored NCE-based bounds for contrastive learning in RL (Srinivas et al., 2020), to the best of our knowledge, prior work has not used this in conjunction with empowerment for prioritizing information in visual model-based RL.
Similarly, the Nguyen-Wainwright-Jordan (NWJ) bound (Nguyen et al., 2010), which to the best our knowledge has not been used by prior works in visual model-based RL, can be obtained as,
[Equation]
where f θ is a critic .
There exists an optimal critic function for which the bound is tightest and equality holds.
We refer to the InfoNCE and NWJ lower bound based objectives as contrastive learning , in order to distinguish them from a reconstruction-loss based objective, though both are bounds on mutual information.
We denote a lower bound to MI by I ( o t , z t ) .
We empirically ﬁnd the NWJ-bound to perform slightly better than the NCE-bound for our approach, explained in section 4.6.
3. O VERALL O BJECTIVE
We now motivate the overall objective, which consists of maximizing mutual information while prioritizing the learning of controllable representations through empowerment in a latent state-space model.
Based on the discussions in Sections 3.1 and 3.2, we deﬁne the overall objective for representation learning as
[Equation]
The objective is to maximize a MI term I ( O t ; Z t ) through contrastive learning such that a constraint on C t holds for prioritizing the encoding of forward-predictive, reward-predictive and controllable representations.
We deﬁne the overall planning objective as
[Equation]
The planning objective is to learn a policy as a function of the latent state Z such that the empowerment term and a reward-based value term are maximized over the horizon H .
We can perform the constrained optimization for representation learning through the method of
Lagrange Multipliers, by the primal and dual updates shown in Section A.2.
In order to analyze this objective, let |O| =
n and | Z | =
d .
Since, O corresponds to images and Z is a bottlenecked latent representation, d (cid:28) n .
[-Environment interaction for time step t = 
-]
I ( O , Z ) is maximized when Z contains all the information present in O such that Z is a sufﬁcient statistic of O .
However, in practice, this is not possible because | Z | (cid:28) |O| .
When c 0 is sufﬁciently large, and the constraint (cid:80) t C t ≥ c 0 is satisﬁed,
Z = [ S + , ˜ S − ] .
Hence, the objective max (cid:80) t I ( O t , Z t ) s.t.
(cid:80) t C t ≥ c 0 cannot encode anything else in Z , in particular it cannot encode distractors DS − .
To understand the importance of prioritization through C t ≥ c 0 , consider max (cid:80) t I ( O t , Z t ) without the constraint.
This objective would try to make Z a sufﬁcient statistic of O , but since | Z | (cid:28) |O| , there are no guarantees about which parts of O are getting encoded in Z .
This is because both distractors S − and non-distractors S + , D ˜ S − are equally important with respect to I ( O , Z ) .
Hence, the constraint helps in prioritizing the type of information to be encoded in Z .
3. P RACTICAL A LGORITHM AND I MPLEMENTATION D ETAILS
To arrive at a practical algorithm [Equation] we optimize the overall objective in section 3.3 through lower bounds on each of the MI terms.
For I ( O , Z ) we consider two variants, corresponding to the NCE and NWJ lower bounds described in Section 3.2.
We can obtain a variational lower bound on each of the terms in C t as follows, with detailed derivations in Appendix A.3: [Equation] Here [Equation]
q φ ( z t | o t ) is the observation encoder, q ψ ( a t − 1 | z t , z t − 1 ) is the inverse dynamics model, q η ( r t | z t ) is the reward decoder
[Equation]
and q χ ( z t | z t − 1 , a t − 1 )) is the forward dynamics model. The inverse model helps in learning representations such that the chosen action is predictable from successive latent states. The forward dynamics model helps in predicting the next latent state given the current latent state and the current action, without having the next observation. The reward decoder predicts the reward (a scalar) at a time-step given the corresponding latent state.
We use dense networks for all the models, with details in Appendix A.7.
We denote by C t , the lower bound to C t based on the sum of the terms above.
We construct a Lagrangian L = (cid:80) t I ( o t ; z t ) + λ ( C t - c 0 ) and optimize it by primal-dual gradient descent. An outline of this is shown in Algorithm 1.
Planning Objective.
For planning to choose actions at every time-step, we learn a policy π ( a | z ) through value estimates of task reward and the empowerment objective.
We learn value estimates with V ( z t )
≈ E π
[ln q η ( r t | z t ) + ln q ψ ( a t − 1 | z t , z t − 1 )] .
We estimate V ( z t ) similar to Equationof (Hafner et al., 2019a).
The empowerment term q ψ ( a t − 1 | z t , z t − 1 ) in policy learning incentivizes choosing actions a t − 1 such that they can be predicted from consecutive latent states z t − 1 , z t .
This biases the policy to explore controllable regions of the state-space.
The policy is trained to maximize the estimate of the value, while the value model is trained to ﬁt the estimate of the value that changes as the policy is updated.
Finally, we note that the difference in value function of the underlying MDP Q π ( o, a ) and the latent
MDP ˆ Q π ( z, a ) , where z ∼ q φ ( z | o ) is bounded, under some regularity assumptions.
We provide this result in Theorem 3 of the Appendix Section A.4.
The overall procedure for model learning, planning, and interacting with the environment is outlined in Algorithm 1.
E XPERIMENTS
Through experiments on robot control tasks, we aim to understand the following research questions:
How does InfoPower compare with the baselines in terms of episodic returns in environments with explicit background video distractors?
How sample efﬁcient is InfoPower when the reward signal is weak ( < 100 k env steps)?
How does InfoPower compare with baselines in terms of behavioral similarity of latent states?
4. S ETUP
We perform experiments with modiﬁed DeepMind Control Suite environments (Tassa et al., 2018), with natural video distractors from ILSVRC dataset (Russakovsky et al., 2015) in the background.
The agent receives only image observations at each time-step and does not receive the groundtruth simulator state.
This is a very challenging setting, because the agents must learn to ignore the distractors and abstract out representations necessary for control.
While natural videos that are unrelated to the task might be easy to ignore, realistic scenes might have other elements that resemble the controllable elements, but are not actually controllable (e.g., other cars in a driving scenario).
To emulate this, we also add distractors that resemble other potentially controllable robots, as shown for example in Fig.
3 (1st and 6th), but are not actually inﬂuenced by actions.
In addition to this setting, we also perform experiments with gray-scale distractors based on videos from the Kinetics dataset (Kay et al., 2017).
This setting is adapted exactly based on prior works (Fu et al., 2021; Nguyen et al., 2021; Zhang et al., 2020) and we compare directly to the published results.
We compare InfoPower with state-of-the art baselines that also learn world models for control:
Dreamer (Hafner et al., 2019a), C-Dreamer that is a contrastive version of Dreamer similar to Ma et al.
(2020), TIA (Fu et al., 2021) that learns explicit reconstructions for both the distracting background and agent separately, DBC (Zhang et al., 2020), and DeepMDP (Gelada et al., 2019).
We also compare variants of InfoPower with different MI bounds (NCE and NWJ), and ablations of it by removing the empowerment objective from policy learning.
4. A BEHAVIORAL SIMILARITY METRIC BASED ON GRAPH KERNELS
In order to measure how good the learned latent representations are at capturing the functionally relevant elements in the scene, we introduce a behavioral similarity metric with details in A.5.
Given the sets S z
= { z i } ni =1 and S z gt = { z i gt } ni =1 , we construct complete graphs by connecting every vertex with every other vertex through an edge.
The weight of an edge is the Euclidean distance between the respective pairs of vertices.
The label of each vertex is an unique integer and corresponding vertices in the two graphs are labelled similarly.
Let the resulting graphs be denoted as G z
= ( V z , E z ) and G z gt =
( V z gt , E z gt ) respectively.
Note that both these graphs are shortest path graphs, by construction.
Let, e i = { u i , v i } and e j = { u j , v j } .
We deﬁne a kernel ˆ k that measures similarity between the edges e i , e j and the labels on respective vertices. ˆ k ( e i , e j ) = k v ( l ( u i ) , l ( u j ))
k e ( l ( e i ) , l ( e j ))
k v ( l ( v i ) , l ( v j ))
+ k v ( l ( u i ) , l ( v j ))
k e ( l ( e i ) , l ( e j ))
k v ( l ( v i ) , l ( u j ))
Here, the function l ( · ) denotes the labels of vertices and the weights of edges.
k v is a Dirac delta kernel i.e. k e ( x, y )
= 1 − δ ( x, y )
(Note that δ ( x, y )
=
1 iff x = y , else δ ( x, y ) = 0 ) and k e is a Brownian ridge kernel, k e ( x, y )
= 1 c max(0 , c − | x − y | ) , where c is a large number.
We deﬁne the shortest path kernel to measure similarity between the two graphs, as, k
( G z , G z gt ) = 1 | E z | (cid:80) e i ∈ E z (cid:80)
e j ∈ E z gt ˆ
k
( e i , e j ) .
The value of k ( G z , G z gt ) is low when a large number of pairs of corresponding vertices in both the graphs have the same edge length. We expect methods that [Table] Walker
[Table]
Cheetah-Run [Table] Hopper-Stand [Table]
Ball-in-a-Cup
[Table]
[Table]
recover latent state representations which better reﬂect the true underlying simulator state (i
[Table]
e., the positions of the joints) to have higher values according to this metric.
4. S AMPLE EFFICIENCY ANALYSIS
In Fig.
[Table] we compare InfoPower and the baselines in terms of episodic returns.
This version of InfoPower corresponds to an NWJ bound on MI which we ﬁnd works slightly better than the
NCE bound variant analyzed in section 4.6.
It is evident that InfoPower achieves higher returns before 1M steps of training quickly compared to the baselines, indicating higher sample efﬁciency.
This suggests the effectiveness of the empowerment model in helping capture controllable representations early on during training, when the agent doesn’t take on actions that yield very high rewards.
4. R ESULTS ON DISTRACTORS FROM THE K INETICS DATASET
In this section, we evaluate InfoPower on the same distractors used in prior works, DBC and
TIA.
The distractor videos are from the Kinetics dataset (Kay et al., 2017)
Driving Car class, and converted to grayscale.
From Table 1 we see that InfoPower slightly outperforms the baselines. We believe the performance gap is not as signiﬁcant as Fig. 4 because the distractors being grayscale images, and not RGB might be easier to ignore by the baselines as well as InfoPower .
4. B EHAVIORAL SIMILARITY OF LATENT STATES
In this section, we analyze how similar are the learned latent representations with respect to the underlying simulator states.
The intuition for this comparison is that the proprioceptive features in the simulator state abstract out distractors in the image, and so we want the latent states to be behaviorally similar to the simulator states.
Quantitative results with the deﬁned metric.
In Table 2, we show results for behavioral similarity of latent states (Sim), based on the metric in section 4.2.
We see that the value of Sim for
InfoPower is around 20% higher than the most competitive baseline, indicating high behavioral similarity of the latent states with respect to the corresponding ground-truth simulator states.
Qualitative visualizations with t-sne.
Fig.
5 shows a t-SNE plot of latent states z ∼ q φ ( z | o ) for
InfoPower and the baselines with visualizations of 3 nearest neighbors for two randomly chosen latent states.
We see that the state of the agent is similar is each group for InfoPower , although the background scenes are signiﬁcantly different.
However, for the baselines, the nearest neighbor states are signiﬁcantly different in terms of the pose of the agent, indicating that the latent representations encode signiﬁcant background information.
4. A BLATION S TUDIES
In Fig.
6, we compare different ablations of InfoPower .
Keeping everything else the same, and changing only the MI lower bound to NCE, we see that the performance is almost similar or slightly worse.
However, when we remove the empowerment objective from policy optimization (the versions with ‘-Policy’ in the plot), we see that performance drops.
The drop is signiﬁcant in the regions < 200 k environment interactions, particularly in the sparse reward environments - cartpole balance and ball in a cup, indicating the necessity of the empowerment objective in exploration for learning controllable representations, when the reward signal is weak.
R ELATED W ORKS
Visual model-based RL.
Recent developments in video prediction and contrastive learning have enabled learning of world-models from images (Watter et al., 2015; Babaeizadeh et al., 2017; Finn & Levine, 2017; Hafner et al., 2019a; Ha & Schmidhuber, 2018; Hafner et al., 2019b; Xie et al., 2020).
All of these approaches learn latent representations through reconstruction objectives that are amenable for planning.
Other approaches have used similar reconstruction based objectives for control, but not for MBRL (Lee et al., 2019; Gregor et al., 2019).
MI for representation learning.
Mutual Information measures the dependence between two random variables.
The task of learning latent representations Z from images O for downstream applications, has been very successful with MI objectives of the form max f 1 ,f 2 I ( f 1 ( O ) , f 2 ( Z ))
(Hjelm & Bachman, 2020; Tian et al., 2020; Oord et al., 2018; Tschannen et al., 2019; Nachum & Yang, 2021).
Since calculating MI exactly is intractable optimizing MI based objectives, it is important to construct appropriate MI estimators that lower-bound the true MI objective (Hjelm et al., 2018;
Nguyen et al., 2010; Belghazi et al., 2018; Agakov, 2004).
The choice of the estimator is crucial, as shown by recent works (Poole et al., 2019), and different estimators yield very different behav- iors of the algorithm.
We incorporate MI maximization through the NCE (Hjelm et al., 2018) and
NWJ (Nguyen et al., 2010) lower bounds, such that typical reconstruction objectives for representation learning which do not perform well with visual distractors, can be avoided.
Inverse models and empowerment.
Prior approaches have used inverse dynamics models for regularization in representation learning (Agrawal et al., 2016; Zhang et al., 2018) and as bonuses for improving policy gradient updates in RL (Shelhamer et al., 2016; Pathak et al., 2017).
The importance of information theoretic approaches for learning representations that maximize predictive power has been discussed in prior work (Still, 2009) and more recently in Lee et al.
Empowerment (Mohamed & Rezende, 2015) has been used as exploration bonuses for policy learning (Leibfried et al., 2019; Klyubin et al., 2008), and for learning skills in RL (Gregor et al., 2016;
Sharma et al., 2019; Eysenbach et al., 2018).
In contrast to prior work, we incorporate empowerment both for state space representation learning and policy learning, in a visual model-based RL framework, with the aim of prioritizing the most functionally relevant information in the scene.
RL with environment distractors.
Some recent RL frameworks have studied the problem of abstracting out only the task relevant information from the environment when there are explicit distractors (Hansen & Wang, 2020; Zhang et al., 2020; Fu et al., 2021; Ma et al., 2020).
Zhang et al.
(2020) constrain the latent states by enforcing a bisimulation metric, without a reconstruction objective.
Fu et al.
(2021) model both the relevant and irrelevant aspects of the environment separately, and differ from our approach that prioritizes learning only the relevant aspects.
(Shu et al., 2020) was among the ﬁrst approaches to use contrastive representations for control.
More recently, (Nguyen et al., 2021) used temporal predictive coding and (Ma et al., 2020) used InfoNCE based contrastive loss for learning representations while ignoring distractors in visual MBRL.
Incorporating a lot of data augmentations for improving robustness with respect to environment variations is an orthogonal line of work (Laskin et al., 2020;
Hansen & Wang, 2020; Raileanu et al., 2020; Srinivas et al., 2020), complementary to our approach.
C ONCLUSION
In this paper we derived an approach for visual model-based RL such that an agent can learn a latent state-space model and a policy by explicitly prioritizing the encoding of functionally relevant factors.
Our prioritized information objective integrates a term inspired by variational empowerment into an objective performing MI without reconstruction.
We evaluate our approach on a suite of vision-based robot control tasks with two sets of challenging video distractor backgrounds.
In comparison to state-of-the-art visual model-based RL methods, we observe higher sample efﬁciency and episodic returns across different environments and distractor settings.
A A
PPENDIX A. T HEORETICAL RESULTS ON CONTROLLABLE REPRESENTATIONS
Let Φ( · ) denote the encoder such that Z =
Φ( O ) .
O is the observation seen by the agent and S is the underlying state.
S + is deﬁned as that part of underlying state S which is directly inﬂuenced by actions A i.e. S + ≤ S s.t.
I ( A t − 1 ; S t | S + t ) = 0 .
We next describe two theorems regarding learning controllable representations, with proofs in the
Appendix.
We observe that the max (cid:80) t I ( A t − 1 ; Z t | Z t − 1 ) objective alone for learning latent representations Z , along with the planning objective provably recovers controllable parts of the observation O , namely S + .
Theorem 1.
The objective max (cid:80) t I ( A t − 1 ; Z t | Z t − 1 ) provably recovers controllable parts S + of the observation O .
S
+ is deﬁned as that part of underlying state S which is directly inﬂuenced by actions A i.e. S + ⊂ S s.t.
I ( S t ; A t − 1 | S + t ) = 0 .
Proof.
Based on Fig.
2 and Fig.
7, we have the following max
I ( A t − 1 ; Z t ) ≤ max I ( A t − 1 ; O t )
Data-Processing Inequality ≤
I ( A t − 1 ; S t )
Data-Processing Inequality ≤
I ( A t − 1 ; S t , S + t ) ≤ I ( A t − 1 ; S t | S + t ) + I
( A t − 1 ; S + t )
Chain Rule of MI ≤ I ( A t − 1 ; S + t )
Conditional
independence
So, the maximum above can be obtained when the encoder Φ is an identity function over S + and a zero function over [ ˜ S − , S − ] .
So, Z t = S + t .
Hence, given Z t − 1 , I ( A t − 1 ; Z t | Z t − 1 ) is maximized when Φ is an identity function over S + .
So, max (cid:80) t I ( A t − 1 ; Z t | Z t − 1 ) learns the encoder Φ such that controllable representations are recovered.
This result is important because in practice, we may not be able to represent every possible factor of variation in a complex environment.
In this situation, we would expect that when | Z | (cid:28) |O| , learning Z under the objective max (cid:80) t I ( A t − 1 ; Z t | Z t − 1 ) would encode S + .
We next show that the inverse information objective alone can be used to train a latent-state space model and a policy through an alternating optimization algorithm that converges to a local minimum of the objective max (cid:80) t
I ( A t − 1 ; Z t | Z t − 1 ) at a rate inversely proportional to the number of iterations.
Theorem 2.
[Equation]
can be optimized through an alternating optimization scheme that has a convergence rate of O (1 /N ) to a local minima of the objective, where N is the number of iterations.
Proof.
Let z ∼ q φ ( z | o ) denote a latent state sampled from the encoder distribution.
Learning a world model such that reward prediction and inverse information are maximized can be summarized as:
[Equation]
The policy π ( a t | z t ) is learned by
[Equation]
Model-based RL in this case involves alternating optimization between policy learning and learning the world model.
In the case where there is no reward signal from the environment, i.e. p ( r t | z t ) , we have the following objective
[Equation]
The optimal q ∗ ψ give a ﬁxed π and φ can be derived as in page 335 of Cover & Thomas (2006).
[Equation]
The optimal π ∗ given a ﬁxed ψ and φ can be derived as
[Equation]
We note that these alternating iterations correspond to an instantiation of the Blahut-Arimoto algorithm (Blahut, 1972) for determining the information theoretic capacity of a channel.
Based on (Nakagawa et al., 2021)
we see that this iterative optimization procedure between q ∗ ψ and π ∗ converges, and the worst-case rate of convergence is O (1 /N ) where N is the number of iterations.
This result is useful because it shows that even in the absence of rewards from the environment ( r t = 0 ∀ t ) , when planning to minimize regret is not possible, the inverse information objective can be used to train a policy that explores to seek out controllable parts of the state-space.
When rewards from the environment are present, we can train the policy with this objective and the value estimates obtained from the cumulative rewards during planning.
In Section 4.3 we empirically show how this objective helps achieve higher sample efﬁciency compared to pure value-based policy learning.
A. P RIMAL D UAL UPDATES
[Equation] [Equation] A. F ORWARD , I NVERSE , AND R EWARD M ODELS
[Equation]
We can obtain a variational lower bound on the empowerment
I ( a t ; z t + k | z t ) =
[Equation]
p
[Equation]
Here, q ( a t | z t + k , z t ) is the inverse dynamics model
[Equation]
and H ( p ( a t | z t )) is the entropy of the policy.
We can obtain a variational lower bound on the reward model as
[Equation]
Here, q ( r t | z t ) is the reward decoder.
Since the reward is a scalar, reconstructing it is computationally simpler compared to reconstructing high dimensional observations as in the BA bound for the observation model.
We show that the difference in value function of the underlying MDP Q π ( o, a ) and the latent MDP ˆ Q π ( z, a ) , where z ∼ q φ ( z | o ) is bounded, under some regularity assumptions.
Similar to the assumption in (Gelada et al., 2019), let the policy π have bounded semi-norm value functions under the total variation distance D TV i.e. | ˆ V π ( z ) | D TV ≤ K and | E z 1 ∼ P V ( z 1 )
− E z 2 ∼ Q V ( z 2 ) | ≤ KD TV ( P, Q ) .
Deﬁne the forward dynamics learning objective as min χ L T = min
χ D KL ( p ( o (cid:48) | o, a ) ||
q χ ( z (cid:48) | z, a )) and the reward model objective as min η L R = min
η D KL ( p ( r | o ) ||
q η ( r | z ))
Theorem 3.
The difference between the Q-function in the latent MDP and that in the original MDP is bounded by the loss in reward predictor and forward dynamics model.
[Equation]
Proof.
[Equation]
So
[Equation]
we have shown that
[Equation]
We can see that, as L R , L T → 0 , the Q-function in the representation space of the MDP, ˆ Q , becomes increasingly closer to that of the original MDP, Q .
Since this result holds for all Q-functions, it also holds for the optimal Q ∗ and ˆ Q ∗ .
This result extends sufﬁciency results in prior works Rakelly et al.
(2021) to a setting with stochastic encoders and KL-divergence losses on forward dynamics and reward, as opposed to Wasserstein metrics (Gelada et al., 2019), and bisimulation metrics (Zhang et al., 2020).
A. B EHAVIORAL S IMILARITY METRIC DETAILS
Given the sets S z
=
{ z i } ni =1 and S z gt =
{ z i gt } ni =1 , we construct complete graphs by connecting every vertex with every other vertex through an edge.
The weight of an edge is the Euclidean distance between the respective pairs of vertices.
The label of each vertex is an unique integer and corresponding vertices in the two graphs are labelled similarly.
Let the resulting graphs be denoted as G z
= ( V z , E z ) and G z gt =
( V z gt , E z gt ) respectively.
Note that both these graphs are shortest path graphs, by construction.
Since the Euclidean distances in both the graphs cannot be directly compared, we scale the distances such that the shortest edge weight among all edge weights are equal in both the graphs.
Scaling every edge weight by the same number doesn’t change the structure of the graph.
Now, we deﬁne the shortest path kernel to measure similarity between the two graphs, as follows
[Equation]
Let, e i = { u i , v i } and e j = { u j , v j } .
Here, the kernel ˆ k measures similarity between the edges e i , e j and the labels on respective vertices.
Let the function l ( · ) denote the labels of vertices and the weights of edges.
[Equation]
i ) , l ( e j )) k v ( l ( v i ) , l ( u j )) Here, k v is a Dirac delta kernel i.e. k e ( x, y )
= 1 − δ ( x, y )
(Note that δ ( x, y )
=
1 iff x = y , else δ ( x, y ) = 0 ) and k e is a Brownian ridge kernel, k e ( x, y )
= 1 c max(0 , c − | x − y | ) , where c is a large number.
So, in summary, the value of ˆ k ( e i , e j ) is low when the edge lengths between corresponding pairs of vertices in both the graphs are close.
Hence, the value of k ( G z , G z gt ) is low when a large number of pairs of corresponding vertices in both the graphs have the same edge length.
We expect methods that recover latent state representations which better reﬂect the true underlying simulator state (i.e., the positions of the joints) to have higher values according to this metric.
A. E XPERIMENTS A.6. D ESCRIPTION OF THE DISTRACTORS
ILSVRC dataset (Russakovsky et al., 2015)
The distractors used for the main experiments are natural video backgrounds with RGB images from the ILSVRC dataset.
We use 200 videos during training, and reserve 50 videos for testing.
An illustration of this for different environments are shown in Fig.
These images are snapshots of what is received by the algorithm (64x64x3 images) and no information about proprioceptive features is provided.
While natural videos that are unrelated to the task might be easy to ignore, realistic scenes might have other elements that resemble the controllable elements, but are not actually controllable (e.g., other cars in a driving scenario).
To emulate this, we also add distractors that resemble other potentially controllable robots, as shown for example in Fig.
8 (top left and bottom right frames), but are not actually inﬂuenced by the actions of the algorithm.
These are particularly challenging because the background video has a prerecorded motion of the same (or similar) agent that is being controlled.
We include such challenging agent-behind-agent background videos with a frequency of 1 in very 3 videos.
For the experiments in Table 2, examples of different levels of distractors are shown in Fig.
The three different levels have distractor windows of size 32x32, 40x40, and 64x64 respectively.
Kinetics dataset (Kay et al., 2017)
For comparing directly with results in prior works DBC (Zhang et al., 2020), and TIA (Fu et al., 2021), we evaluate on the setting where the random videos are grayscale images from the Kinetics dataset driving car class.
The settings are same as in the prior works and we compare directly to the results in the respective papers.
A. T RAINING AND N ETWORK DETAILS
The agent observations have shape 64 x 64 x 3, action dimensions range from 1 to 12, rewards per time-step are scalars between 0 and 1.
All episodes have randomized initial states, and the initial distractor video background is chosen randomly.
We implement our approach with TensorFlow 2 and use a single Nvidia V100 GPU and 10 CPU cores for each training run.
The training time for 1e6 environment steps is 4 hours on average.
In comparison, the average training time for 1e6 steps for Dreamer is 3 hours, for CDreamer is 4 hours, for TIA is 4 hours, for DBC is 4.5 hours, and for DeepMDP is 5 hours.
The encoder consists of 4 convolutional layers with kernel size 4 and channel numbers 32, 65, 128, 256.
The forward model consists of deterministic and stochastic parts, as in standard in visual modelbased RL (Hafner et al., 2019b;a).
The stochastic states have size 30 and deterministic state has size 200.
The compatibility function f θ ( z t , o t (cid:48) ) =
exp ( z Tt W θ z t (cid:48) ) for contrastive learning is learned with a 200 x 200 W θ matrix.
Here, z t (cid:48) is the encoding of o t (cid:48) .
All other models are implemented as three dense layers of size 300 with ELU activations.
We use ADAM optimizer, with learning rate of 6e-4 for the latent-state space model, and 8e-5 for the value function and policy optimization.
The hyper-parameter c 0 for the prioritized information constraint is set to 1000 , after doing a grid-search over the range [100 , 10000] and observing similar performance for 1000 and 10000 .
100 training steps are followed by 1 episode of interacting with the environment with the currently trained policy, for observing real transitions.
The dataset is initialized with 7 episodes collected with random actions in the environment.
We kept the hyper-parameters of
InfoPower same across all environments and distractor types.
For fair comparisons, we kept all the common hyper-parameter values same as Dreamer (Hafner et al., 2019a).
This was the protocol followed in prior works TIA (Fu et al., 2021) and Ma et al.
For the baseline TIA, we tuned the environment-speciﬁc parameters λ R adv and λ O s as mentioned in Table 2 of (Fu et al., 2021).
For λ R adv , we performed a gridsearch over the range 10 k − 50 k and for λ O s we performed a gridsearch over the range 1 − 3 to obtain the parameters for best performance, which we used for the plots in Fig. 4 and Fig.
Respectively for walker-walk, cheetah-run, ﬁnger-spin, quadruped-walk, hopper-stand, ball-in-a-cup-catch, cartpolebalance-sparse, quadruped-run, the values of λ R adv are 30k,30k,20k,40k,30k,40k,20k,30k.
The values of λ O s are 2,2,1.5,2.5,2.5,2,2,2.
For baseline DBC (Zhang et al., 2020), we kept all the parameters same as in Table 2 of the paper, because DBC does not have any environment-speciﬁc parameters and the same values were used for all environments in the DBC paper.
The DeepMDP agent and its hyperparameters are adapted from the implementation provided in the github repo of DBC (Zhang et al., 2020).
A. C ONTRASTIVE L EARNING
Since the distracting backgrounds are changing videos of a certain duration, for contrastive learning, we contrast both across time and across batches.
Contrasting across time helps model invariance against temporally contiguous distractors (for example in the same video) while contrasting across batches helps in learning invariance across different videos.
Concretely, we sample batches from the dataset ( i ) { ( a t , o t , r t ) } Ht =1 , where i denotes a batch index, and for each observation o t in batch i , all observations o t (cid:48) both in batch i and in other batches j (cid:54)
=
i are considered negative samples.
[Table] A. R ESULT ON VARYING DISTRACTOR LEVELS
We consider different levels of distractors by varying the size of the window where distractors in the background are active.
Fig.
9 in the Appendix illustrates this visually.
Table 2 shows results for the different approaches on varying distractor levels.
The size of the frame for distractors at levels
L1, L2, and L3 respectively are 32x32, 40x40, and 64x64.
We observe that the baselines perform worse as the distractor window size increases.
While, for InfoPower , the performance decrease with increasing level is minimal.
This indicates the effectiveness of InfoPower in ﬁltering out background distractors from the observations while learning latent representations.
Note that the distractors in Figs.
4 and 6 are of level 3 i.e. same size as the observations 64x64.
A. S ETTING WITH ONLY SHIFTED AGENT BEHIND AGENT DISTRACTORS
In this section we evaluate InfoPower with the baselines only for the challenging setting of agent-behindagent distrators where the background contains the same agent being controlled (a walker or a cheetah).
An illustration of this is shown in Fig.
From the results in Table 3, we see that InfoPower signiﬁcantly outperforms the competitive baselines both at 500k and 1M environment interactions.
This suggests the utility of
InfoPower in separating out informative elements of the observation from uninformative ones even in challenging settings.
[Table]
Cheetah-Run-shifted 500k 1M 500k 1M
DBC 35 ±104 ±41 ±78 ±
CDreamer 108 ±164 ±73 ±110 ±
TIA 79 ±152 ±65 ±115 ±
InfoPower 168 ±278 ±130 ±207 ±
A. C OMPARISON OF BASELINES WITH EMPOWERMENT IN POLICY LEARNING
In this section, we compare against modiﬁed versions of the baselines, such that we include empowerment in the policy learning of the baselines.
This is the only modiﬁcation we make to the baselines, and keep everything else unchanged.
We described in section 3.3 that inclusion of empowerment in the objective for visual model-based RL (by modifying both the representation learning objective and the policy learning objective) is an important contribution of our paper.
The aim of this experiment is to disentangle the beneﬁts of empowerment in representation learning, for InfoPower with that in policy learning.
We have plotted the results in Fig.
We can see that the performance of the baselines is slightly better than Fig. 4 because of the inclusion of empowerment in the policy learning objective which as we motivated previously in section 3.4 helps with exploring controllable regions of the state-space.
Note that by adding empowerment, we have effectively modiﬁed the baselines.
A. C OMPARISON ON ENVIRONMENTS WITHOUT DISTRACTORS
We compare the performance of InfoPower with a state-of-the-art visual model-based RL method,
Dreamer in the default DM Control environments without distractors.
We see from Table 4 that
InfoPower performs slightly better or similar to Dreamer in all the environments consistently.
This result complements our main paper result on environments with distractors and conﬁrms that the beneﬁt of InfoPower in challenging distracting environments does not come at a cost of performance in simpler non-distracting environments. [Table] Walker Walk
[Table]
Cheetah Run
[Table] [Table] [Table] [Table]
Finger Spin
[Table]
[Table]
[Table]
[Table]
Quadruped Walk
[Table]
[Table]
[Table]
[Table]
Hopper Stand
[Table]
[Table]
[Table]
[Table]
Ball in a Cup
[Table]
[Table]
[Table]
[Table]
Cart-Pole Sparse
[Table]
[Table] [Table] [Table] Quadruped Run
[Table]
[Table]
[Table]
[Table]
A. A DDITIONAL ABLATIONS
In this section
[Table]
we consider additional ablations, namely a version of of InfoPower that does not have the empowerment objective, a version of InfoPower that does not have the MI maximization objective.
We see from Table 5 that removing either of these terms drops performance
[Table] [Table]