EMPOWERMENT IN VISUAL MODEL-BASED RL
Year: 2022

Abstract:
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 modified 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 first. Furthermore, the same empowerment term also promotes faster exploration during the RL process, especially for sparse-reward tasks where the reward signal is insufficient 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 efficiency and episodic returns.

heading: 1 INTRODUCTION
text: 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-efficient 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
∗Work done during Homanga’s research internship at Google. hbharadh@cs.cmu.edu
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 significantly 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 efficiency at 100k interactions.

heading: 2 PROBLEM STATEMENT AND NOTATION
text: 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′|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[ ∑ t γ
t−1rt]. The value of a state V (Zt) is defined as the expected cumulative sum of discounted rewards starting at state Zt.
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., zt denotes the value of Zt). Since the underlying distributions are unknown, we evaluate all expectations through Monte-Carlo sampling with observed state-transition tuples (ot, at−1, ot−1, zt, zt−1, rt).

heading: 3 INFORMATION PRIORITIZATION FOR THE LATENT STATE-SPACE MODEL
text: 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 maxZ I(O, Z) subject to a prioritization of the empowerment objective maxZ I(At−1;Zt|Zt−1). For planning, we maximize the empowerment objective along with reward-based value with respect to the policy maxA I(At−1;Zt|Zt−1) + I(Rt;Zt). In the subsequent sections, we elaborate on our approach, InfoPower, and describe lower bounds to MI that yield a tractable algorithm.

heading: 3.1 LEARNING CONTROLLABLE FACTORS AND PLANNING THROUGH EMPOWERMENT
text: Controllable representations are features of the observation that correspond to entities which the agent can influence through its actions. For example, in quadrupedal locomotion, this could include the joint positions, velocities, motor torques, and the configurations 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 configurations of objects in the scene that it can interact with. Such representations are denoted by S+ in Fig. 2, which we can formally define through conditional independence as the smallest subspace of S, S+ ≤ S, such that I(At−1;St|S+t ) = 0. This conditional independence relation can be seen in Fig. 2. We explicitly priori-
tize 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(At−1;Zt|Zt−1) = H(At−1|Zt−1)−H(At−1|Zt, Zt−1). The first term H(At−1|Zt−1) encourages the chosen actions to be as diverse as possible, while the second term−H(At−1|Zt, Zt−1) encourages the representations Zt and Zt+1 to be such that the action At 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 specifically use this objective in combination with MI maximization for prioritizing the learning of controllable representations from distracting images in the latent state-space model.
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 maxA I(At−1;Zt|Zt−1). This objectives complements value based-learning and further improves exploration by seeking controllable states. We empirically analyze the benefits of this in sections 4.3 and 4.5.
In the Appendix A.1 we next describe two theorems regarding learning controllable representations. We observe that the max ∑ t I(At−1;Zt|Zt−1) objective alone for learning latent representations Z, along with the planning objective provably recovers controllable parts of the observationO, 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| |O|, learning Z under the objective max ∑ t I(At−1;Zt|Zt−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 ∑ t I(At−1;Zt|Zt−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 efficiency compared to pure value-based policy learning.

heading: 3.2 MUTUAL INFORMATION MAXIMIZATION FOR REPRESENTATION LEARNING
text: Algorithm 1: Information Prioritization in Visual Model-based RL (InfoPower)
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 {(at, ot, rt)}k+Lt=k ∼ D. Compute latents zt ∼ pφ(zt|zt−1, at−1, ot). Calculate L based on section 3.4. (φ, χ, ψ, η)← (φ, χ, ψ, η) +∇φ,χ,ψ,ηL λ← λ−∇λL // Behavior learning Rollout latent plan, S ← S ∪ {zt, at, rt} V (zt) ≈ Eπ[ln qη(rt|zt)+ ln qψ(at−1|zt, zt−1)] Update policy π and value model
end // Environment interaction for time step t = 0..T − 1 do
zt ∼ pφ(zt|zt−1, at−1, ot); at ∼ π(at|zt) rt, ot+1 ← env.step(at).
end Add data D ← D ∪ {(ot, at, rt)Tt=1}.
end
For visual model-based RL, we need to learn a representation space Z, such that a forward dynamics model defining the probability of the next state in terms of the current state and the current action can be learned. The objective for this is∑ t−I(it;Zt|Zt−1, At−1). Here, it denotes the dataset indices that determine the observations p(ot|it) = δ(ot − ot′). In addition to the forward dynamics model, we need to learn a reward predictor by maximizing ∑ t I(Rt;Zt), 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 reconstructionloss 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(ot; zt) ≥ Ep(ot,zt)[log qφ′(ot|zt)] +H(p(ot)). The first term here is the reconstruction objective, with qφ′(ot|zt) 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(ot; zt) ≥ Eqφ(zt|ot)p(ot)[log fθ(zt, ot)− log ∑ t′ 6=t fθ(zt, ot′)], where qφ(zt|ot) is the learned encoder, ot is the observation at timestep t (positive sample), and all observations in the replay buffer ot′ are negative samples. fθ(zt, ot′) = exp (zTt Wθzt′) The lower-bound is a form of contrastive learning as it maximizes compatibility of zt with the corresponding observation ot 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,
I(ot; zt) ≥ Eqφ(zt|ot)p(ot)[fθ(zt, ot)]− e −1Eqφ(zt|ot)p(ot)e fθ(zt,ot),
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(ot, zt). We empirically find the NWJ-bound to perform slightly better than the NCE-bound for our approach, explained in section 4.5.

heading: 3.3 OVERALL OBJECTIVE
text: 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 define the overall objective for representation learning as
max Z0:H−1 H−1∑ t=0 I(Ot;Zt) s.t. H−1∑ t=0 Ct︷ ︸︸ ︷ (−I(it;Zt|Zt−1, At−1) + I(At−1;Zt|Zt−1) + I(Rt;Zt)) ≥ c0.
The objective is to maximize a MI term I(Ot;Zt) through contrastive learning such that a constraint on Ct holds for prioritizing the encoding of forward-predictive, reward-predictive and controllable representations. We define the overall planning objective as
max A0:H−1 H−1∑ t=0 I(At−1;Zt|Zt−1) + V (Zt) ;At = π(Zt) ;V (Zt) ≈ ∑ t Rt.
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 n. I(O, Z) is maximized when Z contains all the information present in O such that Z is a sufficient statistic of O. However, in practice, this is not possible because |Z| |O|. When c0 is sufficiently large, and the constraint ∑ t Ct ≥ c0 is satisfied, Z = [S+, S̃−]. Hence, the objective
max ∑ t I(Ot, Zt) s.t. ∑ t Ct ≥ c0 cannot encode anything else in Z, in particular it cannot encode distractors DS−.
To understand the importance of prioritization through Ct ≥ c0, consider max ∑ t I(Ot, Zt) without the constraint. This objective would try to make Z a sufficient statistic of O, but since |Z| |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+, DS̃− are equally important with respect to I(O, Z). Hence, the constraint helps in prioritizing the type of information to be encoded in Z.

heading: 3.4 PRACTICAL ALGORITHM AND IMPLEMENTATION DETAILS
text: To arrive at a practical algorithm, 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 Ct as follows, with detailed derivations in Appendix A.3:
−I(it; zt|at−1) ≥ − ∑ t E[DKL(p(zt|zt−1, at−1, ot)||qχ(zt|zt−1, at−1))]
I(rt; zt) ≥ Ep(rt|ot)[log qη(rt|zt)] +H(p(rt)) I(at−1; zt|zt−1) ≥ Ep(ot|zt−1,at−1)qφ(zt|ot)[log qψ(at−1|zt, zt−1)] + E[H(π(at−1|zt−1))]
Here, qφ(zt|ot) is the observation encoder, qψ(at−1|zt, zt−1) is the inverse dynamics model, qη(rt|zt) is the reward decoder, and qχ(zt|zt−1, at−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 Ct, the lower bound to Ct based on the sum of the terms above. We construct a Lagrangian L = ∑ t I(ot; zt) + λ (Ct - c0) 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 (zt) ≈ Eπ[ln qη(rt|zt) + ln qψ(at−1|zt, zt−1)]. We estimate V (zt) similar to Equation 6
of (Hafner et al., 2019a). The empowerment term qψ(at−1|zt, zt−1) in policy learning incentivizes choosing actions at−1 such that they can be predicted from consecutive latent states zt−1, zt. 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 fit 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.

heading: 4 EXPERIMENTS
text: Through experiments on robot control tasks, we aim to understand the following research questions:
1. How does InfoPower compare with the baselines in terms of episodic returns in environments with explicit background video distractors? 2. How sample efficient is InfoPower when the reward signal is weak (< 100k env steps)? 3. How does InfoPower compare with baselines in terms of behavioral similarity of latent states?
Please refer to the website for a summary and qualitative visualization results https://sites.google.com/view/information-empowerment

heading: 4.1 SETUP
text: We perform experiments with modified 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 influenced 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; Zhang et al., 2020) and we compare directly to the published results.
We compare InfoPowerwith 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.

heading: 4.2 A BEHAVIORAL SIMILARITY METRIC BASED ON GRAPH KERNELS
text: 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 Sz = {zi}ni=1 and Szgt = {zigt}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 Gz = (Vz, Ez) and Gzgt = (Vzgt , Ezgt) respectively. Note that both these graphs are shortest path graphs, by construction. Let, ei = {ui, vi} and ej = {uj , vj}. We define a kernel k̂ that measures similarity between the edges ei, ej and the labels on respective vertices. k̂(ei, ej) = kv(l(ui), l(uj))ke(l(ei), l(ej))kv(l(vi), l(vj)) + kv(l(ui), l(vj))ke(l(ei), l(ej))kv(l(vi), l(uj)) Here, the function l(·) denotes the labels of vertices and the weights of edges. kv is a Dirac delta kernel i.e. ke(x, y) = 1 − δ(x, y) (Note that δ(x, y) = 1 iff x = y, else δ(x, y) = 0) and ke is a Brownian ridge kernel, ke(x, y) = 1c max(0, c − |x − y|), where c is a large number. We define the shortest path kernel to measure similarity between the two graphs, as, k(Gz, Gzgt) = 1 |Ez| ∑ ei∈Ez ∑ ej∈Ezgt
k̂(ei, ej). The value of k(Gz, Gzgt) 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 reflect the true underlying simulator state (i.e., the positions of the joints) to have higher values according to this metric.

heading: 4.3 SAMPLE EFFICIENCY ANALYSIS
text: In Fig. 4, we compare InfoPower and the baselines in terms of episodic returns. This version of InfoPower corresponds to an NWJ bound on MI which we find works slightly better than the NCE bound variant analyzed in section 4.5. It is evident that InfoPower achieves higher returns before 1M steps of training quickly compared to the baselines, indicating higher sample efficiency. 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.

heading: 4.4 BEHAVIORAL SIMILARITY OF LATENT STATES
text: 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 defined metric. In Table 1, 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 significantly different. However, for the baselines, the nearest neighbor states
are significantly different in terms of the pose of the agent, indicating that the latent representations encode significant background information.

heading: 4.5 ABLATION STUDIES
text: 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 significant in the regions < 200k 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.

heading: 5 RELATED WORKS
text: 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 maxf1,f2 I(f1(O), f2(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 behaviors 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. (2020). 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 first 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.

heading: 6 CONCLUSION
text: 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 efficiency and episodic returns across different environments and distractor settings.

heading: A APPENDIX

heading: A.1 THEORETICAL RESULTS ON CONTROLLABLE REPRESENTATIONS
text: Let Φ(·) denote the encoder such that Z = Φ(O). O is the observation seen by the agent and S is the underlying state. S+ is defined as that part of underlying state S which is directly influenced by actions A i.e. S+ ≤ S s.t. I(At−1;St|S+t ) = 0. We next describe two theorems regarding learning controllable representations, with proofs in the Appendix. We observe that the max ∑ t I(At−1;Zt|Zt−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 ∑ t I(At−1;Zt|Zt−1) provably recovers controllable parts S+ of the observation O. S+ is defined as that part of underlying state S which is directly influenced by actions A i.e. S+ ⊂ S s.t. I(St;At−1|S+t ) = 0.
Proof. Based on Fig. 2 and Fig. 7, we have the following max I(At−1;Zt) ≤ max I(At−1;Ot) Data-Processing Inequality
≤ I(At−1;St) Data-Processing Inequality ≤ I(At−1;St, S+t ) ≤ I(At−1;St|S+t ) + I(At−1;S+t ) Chain Rule of MI ≤ I(At−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, Zt = S+t . Hence, given Zt−1, I(At−1;Zt|Zt−1) is maximized when Φ is an identity function over S+. So, max ∑ t I(At−1;Zt|Zt−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| |O|, learning Z under the objective max ∑ t I(At−1;Zt|Zt−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 ∑ t I(At−1;Zt|Zt−1) at a rate inversely proportional to the number of iterations.
Theorem 2. maxπ,ψ ∑ t I(At−1;Zt|Zt−1) = ∑ t Eπ(at−1|zt−1)p(zt|zt−1,at−1,ot) log qψ(at−1|zt,zt−1) π(at−1|zt−1) can be optimized through an alternating optimization scheme that has a convergence rate ofO(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:
max φ,ψ,η ∑ t Eπ(at|zt),qφ(zt|ot) [ Ep(rt|zt,at) log qη(rt|zt, at) + Ep(zt+1|zt,at) log qψ(at|zt+1, zt) π(at|zt) ] The policy π(at|zt) is learned by
max π ∑ t Eπ(at|zt),qφ(zt|ot) [ Ep(rt|zt,at) log qη(rt|zt, at) + Ep(zt+1|zt,at) log qψ(at|zt+1, zt) π(at|zt) ] 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(rt|zt), we have the following objective
max π max φ,ψ ∑ t Eπ(at|zt),qφ(zt|st)p(zt+1|zt,at) log qψ(at|zt+1, zt) π(at|zt)
The optimal q∗ψ give a fixed π and φ can be derived as in page 335 of Cover & Thomas (2006).
q∗ψ(at|zt+1, zt) = qφ(zt|st)p(zt+1|zt, at)π(at|zt)∫ a qφ(zt|st)p(zt+1|zt, at)π(at|zt)
The optimal π∗ given a fixed ψ and φ can be derived as
π∗(at|zt) = exp (Eqφ(zt|st)p(zt+1|zt,at)[log qψ(at|zt+1, zt)])∫ a
exp (Eqφ(zt|st)p(zt+1|zt,at)q[log qψ(at|zt+1, zt)]) 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 (rt = 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 efficiency compared to pure value-based policy learning.

heading: A.2 PRIMAL DUAL UPDATES
text: Z0:H−1 = Z0:H−1+ ∂ ∑H−1 t=0 I(Ot;Zt) + λ (−I(it;Zt|Zt−1, At−1) + I(At−1;Zt|Zt−1) + I(Rt;Zt)− c0)
∂Z0:H−1
λ = λ− −I(it;Zt|Zt−1, At−1)︸ ︷︷ ︸ fwd dynamics model + I(At−1;Zt|Zt−1)︸ ︷︷ ︸ inv dynamics model + I(Rt;Zt)︸ ︷︷ ︸ rew model −c0 

heading: A.3 FORWARD, INVERSE, AND REWARD MODELS
text: −I(it;Zt|Zt−1, At−1) = − ∫ p(zt, zt−1, at−1, it, ot) log
p(zt|zt−1, at−1, ot) p(zt|zt−1, at−1)
= − ∫ p(zt, zt−1, at−1, ot) ( log
p(zt|zt−1, at−1, ot) qχ(zt|zt−1, at−1) + log qχ(zt|zt−1, at−1) p(zt|zt−1, at−1) ) ≥ − ∫ p(zt, zt−1, at−1, ot) log
p(zt|zt−1, at−1, ot) qχ(zt|zt−1, at−1)
= −DKL(p(zt|zt−1, at−1, ot)||qχ(zt|zt−1, at−1))
We can obtain a variational lower bound on the empowerment I(at; zt+k|zt) = ∫ p(at, zt+k, zt) log
p(at|zt+k, zt) p(at|zt)
= ∫ p(at, zt+k, zt) ( log
q(at|zt+k, zt) p(at|zt) + log p(at|zt+k, zt) q(at|zt+k, zt) ) ≥ ∫ p(at, zt+k, zt) log
q(at|zt+k, zt) p(at|zt)
= ∫ p(at, zt+k, zt) log q(at|zt+k, zt)− ∫ p(zt)p(zt+k|at, zt)p(at|zt) log p(at|zt)
= Ep(at,zt+k,zt)[log q(at|zt+k, zt)] + Ep(zt)p(zt+k|at,zt)[H(p(at|zt))] Here, q(at|zt+k, zt) is the inverse dynamics model, andH(p(at|zt)) is the entropy of the policy. We can obtain a variational lower bound on the reward model as
I(rt; zt) = ∫ p(rt, zt) log
p(rt|zt) p(rt)
= ∫ p(rt, zt) ( log
q(rt|zt) p(rt) + log p(rt|zt) q(rt|zt) ) = ∫ p(rt, zt) log
q(rt|zt) p(rt) +KL[p(rt, zt)||q(rt, zt)]
≥ ∫ p(rt, zt) log
q(rt|zt) p(rt)
= ∫ p(rt, zt) log q(rt|zt)− ∫ p(rt, zt) log p(rt)
= Ep(rt,zt)[log q(rt|zt)] + Constant
Here, q(rt|zt) 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.
A.4 VALUE DIFFERENCE RESULT
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 DTV i.e. |V̂ π(z)|DTV ≤ K and |Ez1∼PV (z1)− Ez2∼QV (z2)| ≤ KDTV(P,Q). Define the forward dynamics learning objective as minχ LT = minχDKL(p(o′|o, a)||qχ(z′|z, a)) and the reward model objective as minη LR = minηDKL(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.
E(o,a)∼dπ(o,a),z∼qφ(z|o)[Q π(o, a)− Q̂π(z, a)] ≤
√ LR + γK √ LT
1− γ

heading: Proof.
text: E(o,a)∼dπ(o,a),z∼qφ(z|o)[Q π(o, a)− Q̂π(z, a)]
≤ E(o,a)∼dπ(o,a),z∼qφ(z|o)(|r(o, a)− r̂(z, a)|+ γ|Eo′∼p(o′|o,a)V (o ′)− Ez′∼qχ(z′|z,a)V̂ (z ′)|) ≤ √ DKL(p(r|o)||qη(r|z)) + E(o,a)∼dπ(o,a),z∼qφ(z|o)(γ|Eo′∼p(o′|o,a),z′∼qφ(z′|o′)[V (o ′)− V̂ (z′)]|
+ γ|Eo′∼p(o′|o,a),z′∼qχ(z′|z,a)[V (o ′)− V̂ (z′)]|) ≤ √ LR + E(o,a)∼dπ(o,a),z∼qφ(z|o)(γ|Eo′∼p(o′|o,a),z′∼qφ(z′|o′)[V (o
′)− V̂ (z′)]|) + E(o,a)∼dπ(o,a),z∼qφ(z|o)(γKDTV(p(o ′|o, a)||qχ(z′|z, a)))
≤ √ LR + γE(o,a)∼dπ(o,a),z∼qφ(z|o)Eo′∼p(o′|o,a),z′∼qφ(z′|o′)[V (o ′)− V̂ (z′)])
+ γK √ DKL(p(o′|o, a)||qχ(z′|z, a))
≤ √ LR + γE(o,a)∼dπ(o,a),z∼qφ(z|o)[V (o)− V̂ (z)]) + γK √ LT
≤ √ LR + γE(o,a)∼dπ(o,a),z∼qφ(z|o)[Q(o, a)− Q̂(z, a)]) + γK √ LT
So, we have shown that
E(o,a)∼dπ(o,a),z∼qφ(z|o)[Q π(o, a)− Q̂π(z, a)] ≤
√ LR + γK √ LT
1− γ
We can see that, as LR, LT → 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 sufficiency 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).

heading: A.5 BEHAVIORAL SIMILARITY METRIC DETAILS
text: Given the sets Sz = {zi}ni=1 and Szgt = {zigt}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 Gz = (Vz, Ez) and Gzgt = (Vzgt , Ezgt) 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 define the shortest path kernel to measure similarity between the two graphs, as follows
k(Gz, Gzgt) = 1 |Ez| ∑ ei∈Ez ∑ ej∈Ezgt k̂(ei, ej)
Let, ei = {ui, vi} and ej = {uj , vj}. Here, the kernel k̂ measures similarity between the edges ei, ej and the labels on respective vertices. Let the function l(·) denote the labels of vertices and the weights of edges.
k̂(ei, ej) = kv(l(ui), l(uj))ke(l(ei), l(ej))kv(l(vi), l(vj))+kv(l(ui), l(vj))ke(l(ei), l(ej))kv(l(vi), l(uj))
Here, kv is a Dirac delta kernel i.e. ke(x, y) = 1 − δ(x, y) (Note that δ(x, y) = 1 iff x = y, else δ(x, y) = 0) and ke is a Brownian ridge kernel, ke(x, y) = 1c max(0, c− |x− y|), where c is a large number.
So, in summary, the value of k̂(ei, ej) is low when the edge lengths between corresponding pairs of vertices in both the graphs are close. Hence, the value of k(Gz, Gzgt) 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 reflect the true underlying simulator state (i.e., the positions of the joints) to have higher values according to this metric.

heading: A.6 DESCRIPTION OF THE DISTRACTORS
text: 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. 8. 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 influenced 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 1, examples of different levels of distractors are shown in Fig. 9. 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.

heading: A.7 TRAINING AND NETWORK DETAILS
text: 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θ(zt, ot′) = exp (zTt Wθzt′) for contrastive learning is learned with a 200 x 200 Wθ matrix. Here, zt′ is the encoding of ot′ . 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 c0 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.
(2020). For the baseline TIA, we tuned the environment-specific parameters λRadv and λOs as mentioned in Table 2 of (Fu et al., 2021). For λRadv , we performed a gridsearch over the range 10k − 50k and for λOs 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. 11. Respectively for walker-walk, cheetah-run, finger-spin, quadruped-walk, hopper-stand, ball-in-a-cup-catch, cartpolebalance-sparse, quadruped-run, the values of λRadv are 30k,30k,20k,40k,30k,40k,20k,30k. The values of λOs 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-specific 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).

heading: A.8 CONTRASTIVE LEARNING
text: 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){(at, ot, rt)}Ht=1, where i denotes a batch index, and for each observation ot in batch i, all observations ot′ both in batch i and in other batches j 6= i are considered negative samples.

heading: A.9 RESULT ON VARYING DISTRACTOR LEVELS
text: 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 1 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 filtering 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.

heading: A.10 RESULTS ON DISTRACTORS FROM THE KINETICS DATASET
text: 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 2 we see that InfoPower slightly outperforms the baselines. We believe the performance gap is not as significant as Fig. 4 because the distractors being grayscale images, and not RGB might be easier to ignore by the baselines as well as InfoPower.

heading: A.11 SETTING WITH ONLY SHIFTED AGENT BEHIND AGENT DISTRACTORS
text: 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. 10. From the results in Table 3, we see that InfoPower significantly out-
19
performs 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.

heading: A.12 COMPARISON OF BASELINES WITH EMPOWERMENT IN POLICY LEARNING
text: In this section, we compare against modified versions of the baselines, such that we include empowerment in the policy learning of the baselines. This is the only modification 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 benefits of empowerment in representation learning, for InfoPower with that in policy learning.
We have plotted the results in Fig. 11. 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 modified the baselines.

heading: A.13 COMPARISON ON ENVIRONMENTS WITHOUT DISTRACTORS
text: We compare the performance of InfoPowerwith 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 confirms that the benefit of InfoPower in challenging distracting environments does not come at a cost of performance in simpler non-distracting environments.

heading: A.14 ADDITIONAL ABLATIONS
text: In this section, 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.

Review summaries:
This paper tackles the problem of prioritizing functionally relevant information from complex observations for model-based RL. To this end, previous work has proposed to replace the reconstruction loss with contrastive loss. Building on that, this paper introduces an additional empowerment objective that is used for both representation learning and policy learning. Experiments show that the proposed model outperforms baselines on a set of DeepMind Control tasks with custom background distractions, which include other visually similar but uncontrollable agents. It is also shown that the similarity between learned states can match well to the similarity between groundtruth simulator states, according to the proposed metric.

The paper proposes a new non-reconstruction method for model-based RL from high-dimensional observations. The main idea is to make use of an information empowerment objective that prioritizes encoding parts of the environment that are influenceable by the actions. This allows the model to focus on functionally relevant information and filter out distractors. The same empowerment term can also be used to promote faster exploration when the reward signal is sparse. The proposed method outperformed the existing baselines in difficult Deepmind control tasks with natural video backgrounds.

The paper introduces a method to learn representations for model-based RL. The key idea is to find a representation that maximises information between the action at the previous step and the current latent representation, given the representation at the previous step, thereby maximising the information that the representation encodes about effects of an action. The work contributes a method and evaluation on tasks where high-dimensional visual inputs are used to control an agent in an environment with complex backgrounds as distractors.