Year: 2021

Abstract:
Deep Reinforcement Learning (deep RL) has been successfully applied to solve various decision-making problems in recent years. However, the observations in many real-world tasks are often high dimensional and include much task-irrelevant information, limiting the applications of RL algorithms. To tackle this problem, we propose LCER, a representation learning method that aims to provide RL algorithms with compact and sufficient descriptions of the original observations. Specifically, LCER trains representations to retain the controllable elements of the environment, which can reflect the action-related environment dynamics and thus are likely to be task-relevant. We demonstrate the strength of LCER on the DMControl Suite, proving that it can achieve state-of-the-art performance. To the best of our knowledge, LCER is the first representation learning algorithm that enables the pixel-based SAC to outperform state-based SAC on the DMControl 100K benchmark, showing that the obtained representations can match the oracle descriptions (i.e. the physical states) of the environment.

heading: 1 INTRODUCTION
text: Deep Reinforcement Learning (deep RL) has proven its ability to solve difficult sequential decisionmaking problems such as Dota2 (Berner et al., 2019) and StarCraft (Vinyals et al., 2019). However, it is still a challenge to apply deep RL to many real-world tasks, because the observations in these tasks are often high dimensional (e.g. pixels) and include much task-irrelevant noise. To address these issues, representation learning is introduced in deep RL to provide information-dense descriptions of the original observations, and thus reduce the difficulty of solving RL problems.
Many representations learning algorithms are proposed with the design goal to make the representations be predictive of properties of future states. Such a goal is typically achieved by imposing the representations to reflect the dynamics of the environment (Mazoure et al.; Lee et al., 2020b; Schwarzer et al., 2020; Gelada et al., 2019). However, the dynamics of the environment are not always task-relevant(e.g. the noisy and varying background in control tasks). This can be problematic because those task-irrelevant dynamics will entice the representations into encoding task-irrelevant information(e.g. the noisy background). Since the capacity of representations is limited, the obtained representations can miss important task-relevant information (Gelada et al., 2019) or encode too much unimportant task-irrelevant information (Zhang et al., 2020). Therefore, it is essential to learn representations with more task-relevant environmental dynamics.
We envision that action-related dynamics are likely to be more task-relevant because actions are the only feedback from the agent to the environment. For example, in control tasks and robotics particularly, the agent can complete its task only if it can realize which part of the environment can be influenced by its actions. For simplicity, we use controllable elements to denote those environment elements that can be influenced by the agent’s actions and thus reflect the action-related dynamics. We argue that these controllable elements should be explicitly retained for better representations.
Following the intuition above, we develop an efficient and flexible framework named Learning Controllable Elements oriented Representations (LCER) for reinforcement learning to capture the controllable elements of the environment. Specifically, we formally define a metric to measure the amount of controllable elements that representations capture, and then propose a surrogate objective that is derived from this metric and can make the representations encode controllable elements as
well as sufficient information. By emphasizing the importance of the controllable elements, LCER can capture more task-relevant information automatically.
We evaluate our method LCER on the DMControl Suite (Tassa et al., 2018), which has been widely used in recent years to benchmark sample efficiency (Kostrikov et al., 2020b; Hafner et al., 2019; Laskin et al.; Lee et al., 2020b). Our experiments show that LCER combined with SAC (Haarnoja et al., 2018) can outperform other representation learning methods on a majority of tasks, and even achieve better performance than the state-based SAC on the DMControl 100K benchmark. We also evaluate LCER on the Distracting Control Suite (Stone et al., 2021), showing LCER’s ability to filter out distracting factors. The implementation of our work is available at https://anonymous. 4open.science/r/LCER-FF40

heading: 2 RELATED WORK
text: In this section, we review the common approaches of representation learning in RL briefly.
Predictive representations in RL. Predictive representations are typically trained to predict the properties of future states. One commonly used approach is to predict the representation of future states given the current state’s representation and actions. Such an approach is widely used in modelbased RL algorithms. The representations are prone to collapse if without further constraints (Gelada et al., 2019), thus a reconstruction loss is usually applied to ensure non-trivial representations. However, this approach tries to encode everything, unable to filter out task-irrelevant information. SPR (Schwarzer et al., 2020) introduces a target encoder to avoid collapse and reconstruction, but its effect lacks theoretical guarantee, and it does not emphasize the controllable elements as LCER does. There are some methods that train predictive representations from information-theoretic perspectives, such as (van den Oord et al., 2018; Hjelm et al., 2019; Anand et al., 2019; Lee et al., 2020b; Mazoure et al.). These methods do not give a bias explicitly towards what kind of environment’s dynamics should the representations reflect, while our method LCER focuses more on the actionrelated dynamics which is more likely to be task-relevant. Our experiments in Section 5.2 shows that encoding predictive information does not ensure representations of high quality in complex environments, and it is crucial to retain controllable elements explicitly.
There are some prior works that also emphasize controllable elements. Typically, (Zhang et al., 2018; pok, 2016; Badia et al., 2020) learn an inverse model that predicts the action given successive states. However, these approaches typically rely on a reconstructive loss, which will significantly harm their performance in complex environments. There are also some works that emphasize the importance of actions in representation learning by introducing various action-related objectives such as policy selectivity(Bengio et al., 2017) and contingency awareness(Bellemare et al., 2012). However, their methods are often built on heuristics and do not give a clear definition of controllable elements. Compared with these works, LCER firstly introduces a novel measure of controllable elements (Eq.(3)) and then develops an efficient method to optimize it without the aid of reconstruction loss.
Bisimulation-based Representations in RL. Bisimulation-based representations group those states into a cluster which are indistinguishable w.r.t. future return, given any action sequence tested. The indistinguishability of two states is typically measured by the bisimulation metrics, which are firstly introduced in (Ferns & doina Precup, 2014). Recently, (Zhang et al., 2020) proposes DBC, a representation learning algorithm based on bisimulation that presents impressive results in distracting environments. However, these methods rely on a well-defined reward function, which is probably unavailable in real-world tasks. Besides, they often fail to solve simple tasks efficiently such as ones from DMControl Suite (Zhang et al., 2020).
Prior-knowledge-based Representations in RL There are other methods that use prior knowledge to constraint the representations (Jonschkowski & Brock, 2015; Jonschkowski et al., 2017; Thomas et al., 2018). For example, (Jonschkowski & Brock, 2015) forces the representations to satisfy multiple priors such as proportionality and repeatability, which generally describe how the optimal representations should look like. However, such priors need to be specially designed for each environment, and much domain knowledge is required to form good representations. These drawbacks significantly limit the applications of these methods.

heading: 3 PRELIMINARIES

heading: 3.1 NOTATION
text: We assume the underlying environment is a Markov decision process (MDP), described by the tuple M = (S,A, P,R, γ), where S is the state space, A the action space, P : St × At × St+1 → [0, 1] the transition probability function which determines the distribution of next state given current state and action, and γ ∈ [0, 1] a discount factor. Given the current state s ∈ S, an agent choose its action a ∈ A according to a policy function a ∼ π(·|s). This action will update the system state to a new state s′ according to the transition function P , and then a reward r = R(s, a, s′) ∈ R is given to the agent. The goal of the agent is to maximize the expected cumulative rewards by learning a policy π.
We denote the state at time step t by st. We use the upper letter of s (i.e. S) to refer to the random variable of state if there is no ambiguity. We use st:t+k to present the state sequence from t to t+ k, i.e. st:t+k := (st, st+1, ..., st+k). These notations can be extended to other variables such as actions and rewards. We use φ : S → Z to denote the embedding function of LCER, which maps s ∈ S to a latent space Z.

heading: 3.2 MUTUAL INFORMATION
text: The mutual information(MI) of two random variables X,Y is defined as:
I(X;Y ) = DKL(PXY ‖PXPY ), (1)
where DKL is the Kullback–Leibler divergence that measures the difference of two distributions.
The conditional mutual information(CMI) I(X;Y |Z) is defined as the expected value of I(X;Y ) given the value of Z:
I(X;Y |Z) = Ez∼ZI(X|z;Y |z). (2)

heading: 4 LEARNING CONTROLLABLE ELEMENTS ORIENTED REPRESENTATIONS
text: Given an embedding function: φ : S → Z, we argue that the number of controllable elements that φ contains can be measured by the conditional mutual information(CMI):
Ik(φ) := I(φ(St+k);At:t+k−1|φ(St)) (3)
Intuitively, Ik measures the information that φ(St+k) and At:t+k−1 shares given φ(St). Note that the non-controllable elements captured by φ do not contribute to Ik because they are not related to the actions. Training representations to maximize Eq.(3) is an efficient way to make the representations capture controllable elements. However, representations obtained in this way may be too compressive1 to provide sufficient information for solving tasks in some cases, because not all task-relevant elements are controllable. For example, in navigation tasks, the position of the goal is task-relevant but not controllable by the agent. Therefore, we need to utilize Eq.(3) in a more delicate manner. Specifically, in Section 4.1 we propose a surrogate objective derived from Eq.(3) which additionally provides a mechanism to control the degree of compression. In Section 4.2, we describe how to optimize this objective, leading to our final algorithm LCER.

heading: 4.1 THE OPTIMIZATION OBJECTIVE FOR LCER
text: By the chain rule of mutual information, we can break the Ik defined in Eq.(3) into the subtraction of two terms: Ik(φ) = I([φ(St), At:t+k−1];φ(St+k)) − I(φ(St);φ(St+k)). The first term tries to capture predictive information, while the second term serves as a compression term that aims to filter out action-irrelevant information. To make φ capture the controllable elements and maintain sufficient information at the same time, we introduce a hyper-parameter β to control the degree of compression:
Ik(φ;β) := I([φ(St), At:t+k−1];φ(St+k))− βI(φ(St);φ(St+k)). (4) 1By saying “compressive”, we mean that these representations only encode controllable elements and filter
out any other information.
By choosing different β when training φ to maximize Eq.(4), we can control the degree of compression, making it possible to adapt to different environments by choosing different β.
However, directly maximizing Ik(φ;β) is intractable, because maximizing the first term and minimizing the second term in Eq.(4) at the same time have an antagonistic effect to each other, thus it is difficult to tune in practice. In fact, we find this will often lead to unstable training process in our early experiments.
To tackle this problem, we substitute I(φ(St);φ(St+k)) with Ik(φ;β = 0)− Ik(φ;β = 1): Ik(φ;β) = I([φ(St), At:t+k−1];φ(St+k))− β(Ik(φ;β = 0)− Ik(φ;β = 1))
= I([φ(St), At:t+k−1];φ(St+k))︸ ︷︷ ︸ J1k ·(1− β) + I(φ(St+k);At:t+k−1|φ(St))︸ ︷︷ ︸ J2k ·β. (5)
In Eq.(5), we maximize two terms (J1k and J 2 k ) at the same time, and their coefficients are 1 − β and β respectively. The first term J1k is closely related to the predictive information encoded by φ. The second term J2k is related to the controllable elements, which will play an important role in distracting environments as shown in Section 5.3.
In practice, we sum over different choices of k, leading to the final objective to maximize:
IK(φ;β) := K∑ k=1 Ik(φ;β). (6)

heading: 4.2 HOW TO OPTIMIZE THE OBJECTIVE FOR LCER
text: To maximize IK(φ, β), we propose to maximize its corresponding lower bound, which is common practice in mutual information estimation. In this section, we first provides the lower bound for J1k and J2k respectively that are used in this paper, and finally put them together to get our final loss for training φ.
Optimizing J1k . In order to maximize J1k (φ), we propose to maximize the InfoNCE (van den Oord et al., 2018) lower bound of J1k (φ):
Ĵ1k (φ) = − 1
B B∑ i=1 log exp(f(φ(sit), a i t:t+k−1;φ(s i t+k)))∑B j=1 exp(f(φ(s i t), a i t:t+k−1;φ(s j t+k))) , (7)
where f is any scalar function, B is the batch size, and sit, a i t:t+k−1, s i t+k is a transition sequence from sit to s i t+k. There are other alternatives to InfoNCE such as NWJ (Poole et al., 2019) and MINE (Belghazi et al., 2018), however InfoNCE can often lead to better representations as shown in (Tschannen et al., 2020). In order to efficiently optimize Ĵ1k for all 1 ≤ k ≤ K, we recursively define ψit+k to encode the information of (φ(s i t), a i t:t+k−1):
ψit+k =
{ φ(sit), if k = 0
h(ψit+k−1, a i t+k−1), otherwise
(8)
and set f to be the bi-linear inner product of ψit+k and φ(s i t+k):
f(φ(sit), a i t:t+k−1;φ(s i t+k)) := (ψ i t+k) TWφ(sit+k). (9)
Note that the φit+k is shared across all Ĵ 1 k , thus we do not need to re-compute it. The h in Eq.(8) is quite similar to a “prediction” network, but is trained together with φ instead of prediction errors.
Optimizing J2k . In order to maximize J2k (φ), we turn to maximizing a Jensen-Shannon CMI estimator Eq.(10), because it is proven to be more stable in our experiments, which is also suggested in prior works (Hjelm et al., 2019; Sanchez et al., 2020). We introduce a statistics network Tk : S ×At:t+k−1 × S → R and maximize:
Ĵ2k (φ) := − 1
B B∑ i=1 [log(1 + e−Tk(φ i t,a i t:t+k−1,φ i t+k)) + log(1 + eTk(φ i t,a j(i) t:t+k−1,φ i t+k)], (10)
where j(1), .., j(B) are obtained by randomly shuffling 1, .., B. The derivation of Eq.(10) is provided in Appendix B. In practice, we choose Tk(φt, at:t+k−1, φt+k) = T (ψt+k, φt+k), where ψt+k is a latent variable defined in Eq.(8) and T : Z × Z → R is a network shared across different Tk.
Final loss for training φ. Inspired by (He et al., 2020), we propose to utilize a target encoder φtarg to stabilize the training process, whose weights are the exponential moving average(EMA2) of φ. Specifically, the embeddings of sit+k (k ≥ 1) are calculated by φtarg. All pieces together, we obtain our loss for φ to minimize:
LK(φ;h,W, T ) = − K∑ k=1 Ĵ1k (φ) · (1− β) + Ĵ2k (φ) · β
= 1
B B∑ i=1 K∑ k=1 [(1− β) log exp((ψit+k) TWφtarg(s i t+k))∑B j=1 exp((ψ i t+k) TWφtarg(s j t+k))
+ β(log(1 + exp (−T (ψit+k, φtarg(sit+k))) + log(1 + exp (T (ψ̂it+k, φtarg(sit+k)))], (11)
where B is the batch size, h, ψit+k are defined in Eq.(8), W is defined in Eq.(9), ψ̂ i t+k is a latent defined in the same way with ψit+k but it encodes φ(s i t) and a j(i) t:t+k−1(instead of a i t:t+k−1), T is a scalar function, and φtarg is the target encoder. The parameters of φ, h,W, T are training together by minimizing LK . In practice, Eq.(11) can be used as an auxiliary loss for any RL algorithms such as SAC (Haarnoja et al., 2018).

heading: 5 EXPERIMENTS
text: In this section, we first compare LCER with other methods in Section 5.1 on the PlaNet benchmark. In Section 5.2, we then demonstrate LCER’s strength to filter out distracting factors on the Distracting Control Suite, and show that it is not enough to obtain good representations just by encoding predictive information. In Section 5.3, we investigate the effect of K and β, which are the most important parameter in LCER. Finally, we investigate how well the representations obtained by LCER can be generalized to other tasks from the same domain in Section 5.4.
In our experiments, LCER is equipped with SAC (Haarnoja et al., 2018). Most of the parameters in LCER are the same with PI-SAC (Lee et al., 2020b), including the batch size, learning rate, action repeat, and the network architecture of φ, see Appendix A for more details. Following (Laskin et al.; Lee et al., 2020a;b), we report the performance using the true environment steps, which is invariant to the choice of action repeat. Unless specially specified, We set K = 5, β = 0.1 and run every experiment using 10 random seeds.
2EMA is a widely used trick in RL. For example, in SAC, the weights of target Q are the EMA of Q.

heading: 5.1 EVALUATION ON SAMPLE EFFICIENCY
text: In this section, We evaluate LCER on the PlaNet benchmark, which consists of six challenging control tasks from the DMControl Suite (Tassa et al., 2018). The PlaNet benchmark is first introduced in (Hafner et al., 2019) and later widely used to benchmark sample efficiency in (Kostrikov et al., 2020b; Laskin et al.; Lee et al., 2020b;a).
We equip LCER with SAC (Haarnoja et al., 2018) and compare LCER to 3 leading representation learning algorithms3: CURL (Laskin et al.),SLAC (Lee et al., 2020a),and PI-SAC (Lee et al., 2020b). CURL introduces contrastive learning into RL and actually trains representations to encode information that is invariant under data augmentation. SLAC efficiently combines both representation learning and control into a single objective via Bayesian inference, but it relies on a reconstruction loss. PI-SAC tries to encode predictive information, while LCER focuses on encoding controllable elements. All these methods are also built on SAC (Haarnoja et al., 2018). Following (Laskin et al.; Kostrikov et al., 2020b), we use the version of PI-SAC that performs one gradient update per environment step to ensure a fair comparison. We also compare LCER with a model-based algorithm Dreamer (Hafner et al., 2020). To further evaluate the quality of the representations obtained by LCER, we also consider the state-based SAC (which called StateSAC) that directly learns from internal physical states provided by DMControl Suite. The performance of StateSAC can be seen as an upper bound since it utilizes the oracle representations of the environments. All hyperparameters of StateSAC are the same with LCER except that it uses states as input instead of pixels.
In Figure 1, we report the learning curves of LCER, PI-SAC (Lee et al., 2020b), SLAC (Lee et al., 2020a), and Dreamer (Hafner et al., 2020). LCER consistently achieves better or comparable performance than other methods. In Table 1, We report the performance of each algorithm at 100k environment steps. LCER is the only algorithm that enables pixel-based SAC to outperform statebased SAC w.r.t the average returns, showing that the obtained representations can match the oracle descriptions(i.e. the physical states) of the environment. We also provide additional aggregate metrics on DMControl 100k benchmark, see Section C for details.
In practice, pixel-based observations are often equipped with data augmentation to serve as a regularizer (Kostrikov et al., 2020b; Lee et al., 2020b; Laskin et al.; 2020). In LCER, we also randomly shift the observations by [−4, 4] pixels. However, a proper data augmentation method is not always available for general forms of observations. Therefore, We further test our method LCER in the
3The performance data of SLAC, Dreamer, CURL is provided by the authors of the corresponding papers.
situation where no data augmentation is available. As shown in Figure 2, LCER can still improve the performance to a great extent without the help of data augmentation. Note that some sampleefficient RL algorithms such as DrQ (Kostrikov et al., 2020b) and RAD (Laskin et al., 2020) that rely on data augmentation can not be extended to such a situation naturally.
Our experiments on the PlaNet benchmark show that our proposed method LCER can obtain representations that help solve the tasks, and thus improve the sample efficiency.

heading: 5.2 ROBUSTNESS TO DISTRACTIONS
text: In this section, we test LCER on the Distracting Control Suite (Stone et al., 2021), a challenging benchmark that is implemented by introducing several kinds of visual distractions to the DMControl Suite, such as random changes to camera pose, object color, and background (see Figure 3). These distractions widely exist in the real world and can reduce the performance of RL algorithms to a great extent (Stone et al., 2021). We aim to show LCER’s ability to filter out those distractions and that it is not enough to just encode predictive information. Note that some representation learning methods such as SLAC (Lee et al., 2020a) rely on a reconstruction loss, which will lead to significant performance degradation as shown in (Zhang et al., 2020).
In Figure 3, we run experiments on three control tasks from the Distracting Control Suite: D.cheetahrun, D.walker-walk, D.finger-spin4. In all tasks, we set the difficulty to be ‘easy’ and use a dynamically varying background (refer to (Stone et al., 2021) for more details). We choose β = 0.4 for D.cheetah-run and D.finger-spin, and β = 1 for D.walker-walk5. We additionally compare LCER with DBC (Zhang et al., 2020) and CVRL (Ma et al., 2020), which are specially designed to deal with distractions. As shown in Figure 3, proper representation learning methods are crucial in these tasks, because it is difficult for the agent to filter out task-irrelevant information by itself. LCER achieves much better performance than PI-SAC, which indicates that it can not lead to high-quality representations just by encoding predictive information, instead we have to attach more importance to the controllable elements. We argue that this is because the distractions are often not controllable by the agent, and thus can be filtered out by LCER. Although these distractions do have a negative effect on LCER’s performance, LCER suffer less than other methods, showing LCER’s robustness to these distractions.
5.3 THE EFFECT OF K AND β
In this section, we try to specify the effect of K and β, which are the most important hyperparameters in LCER.
To investigate the effect of K, we set β = 0 and report the performance for K = 0, 1, 2, 5 in Figure 4(a). Note that K = 0 is a special case in LCER that does not take into account any actions, and exactly recovers the work of CURL (Laskin et al.). We add K = 0 to emphasize the importance of actions. As shown in Figure 4(a), larger K leads to better performance, especially when K changes from 0 to 1. We think this is because large K can provide more supervision signals thanks to our special design of Eq.(11). To investigated the effect of β, we report the performance for β = 0.1, 0.2, 0.4, 1 in Figure 4(b). As shown in Figure 4(b), larger β does lead to better performance
4We use “D.” to distinguish these tasks from the ones in the DMControl Suite. 5We sweep over {0.1,0.2,0.4,1} to get the optimal β. The learning curves for these β’s are provided in
Figure 4(b).
within a certain range, however when β goes beyond this range (e.g.β = 1), it may hinder the performance in some tasks. The result in Figure 4(b) indicates that the optimal β should be chosen according to the given environment. By introducing a hyper-parameter β, LCER is able to strike a balance on the proportion of controllable elements that φ captures, allowing for searching for better performance.

heading: 5.4 THE QUALITY OF REPRESENTATIONS
text: In this section, we try to measure the quality of φ by generalizing φ to unseen tasks. Specifically, we train φ on a source task, froze φ, and then apply it to a target task within the same domain. If LCER can capture the controllable elements efficiently, then the representations obtained on source task should also fit in with the target task, because tasks within the same domain share the same environment dynamics (but with different reward functions). In Figure 5, we report 3 learning curves on target tasks6: SAC, LCER(tsf.) and LCER(self tsf.). SAC is trained directly on target tasks. LCER(tsf.) and LCER(self tsf.) is trained with frozen encoders which are obtained via training on source tasks and target tasks respectively for 300k environment steps. Note that LCER(self tsf.) can be seen as a ground truth baseline where an LCER agent learns representations directly on the target task and is then made to ”transfer” to the same task. On those source tasks where the environment dynamics are well explored (such as cartpole-swingup and walker-walk), LCER(tsf.) can perform as well as LCER(self tsf.) on target tasks. However, when solving target tasks needs to exploit more environment dynamics which are not explored in source tasks (such as hopper-stand7), LCER(tsf.) shows inferior performance. We can conclude that the transfer ability of LCER is mainly influenced by the exploration on the environment dynamics instead of task-specific features. If the agent can explore the full dynamics (such as in walker-walk and cartpole-swingup), LCER(tsf.) can learn as effective representations as LCER(self tsf) does even on source tasks.

heading: 6 CONCLUSION
text: In this paper, we introduce the LCER, an efficient framework to train representations that can reflect the action-related dynamics of the environment. By focusing on the controllable elements of the environment, LCER can encode task-relevant information into representations, improving the performance of RL algorithms. Our experiments on DMControl 100K benchmark show that LCER can enable the pixel-based SAC to outperform state-based SAC, demonstrating improvement over prior works. We also evaluate LCER on the Distracting Control Suite, showing LCER’s robustness to distractions.
6including cartpole-swingup → cartpole-balance, walker-walk → walker-run, and hopper-stand → hopperhop.
7This because in hopper-hop the agent needs to move forward, however in hopper-stand it does not.

heading: B ESTIMATION OF CONDITIONAL MUTUAL INFORMATION
text: In this section, we describe how we optimize IJSD defined in Section 4.2. In the followings, the φ(St+k), At:t+k−1, φ(St) are denoted by X,Y, Z respectively.
We first borrow some theoretical results from (Nowozin et al., 2016):
IJSD(X;Y |Z) = Ez∼ZDJS(PXY |z‖PX|zPY |z) = DJSD(PXY Z‖PXY PY |Z) ≥ sup
V Ex,y,z∼XY ZV (x, y, z)− Ex,z∼XZ,y∼Y |zf∗JSD(V (x, y, z)).
(12)
Here, fJSD is the generator function of Jensen-Shannon divergence: fJSD(u) = −(u+1) log 1+u2 + u log u, and f∗JSD is its corresponding convex conjugate function: f ∗ JSD(t) = supu{ut − fJSD(u)} = − log(2 − exp(t)). We choose V (x, y, z) = log(2) − log(1 + exp(−Tk(x, y, z))), then:
IJSD(X;Y |Z) ≥ sup V Ex,y,z∼XY ZV (x, y, z)− Ex,z∼XZ,y∼Y |zf∗JSD(V (x, y, z))
= sup Tk
−Ex,y,z∼XY Z log(1 + e−Tk(x,y,z))− Ex,z∼XZ,y∼Y |z log(1 + eTk(x,y,z)) + 2 log 2
≈ sup Tk −Ex,y,z∼XY Z log(1 + e−Tk(x,y,z))− Ex,z∼XZ,y∼Y log(1 + eTk(x,y,z)) + 2 log 2.
(13) In Eq.(13), we suppose p(y|z) ≈ p(y)(i.e. At:t+k−1 is independent from φ(St)) for simplicity. Such an assumption is appropriate in our case, because (St, At:t+k−1, St+k) is sampled from a large replay buffer(= 105) whose action policy is quite non-stationary due to policy updates.
To estimate the RHS of Eq.(13), we first randomly sample B transitions {sit, ait:t+k−1, sit+k}Bi=1 from the replay buffer. Assume (j(1), .., j(B)) is a permutation of (1, 2, .., B) obtained by random shuffling, thenEx,y,z∼XY Z(·) can be approximated by 1B ∑ sit,a i t:t+k−1,s i t+k
(·) andEx,z∼XZ,y∼Y (·) by 1B ∑ sit,a j(i) t:t+k−1,s i t+k (·). Therefore, the RHS of Eq.(13) can be estimated by:
Ĵ2k (φ) := − 1
B B∑ i=1 [log(1 + e−Tk(φ i t,a i t:t+k−1,φ i t+k)) + log(1 + eTk(φ i t,a j(i) t:t+k−1,φ i t+k)], (14)
Here, we omit the “2 log 2” in Eq.(13) for simplicity.

heading: C ADDITIONAL AGGREGATE METRICS ON DMCONTROL 100K BENCHMARK
text: In this section, we provides additional comparison on DMControl 100k benchmark using aggregate metrics from rliable (Agarwal et al., 2021b).
We normalize all scores to [0, 1] by dividing by the maximum score (=1000). The metrics are shown in Figure 6, which include:
• Ranks distribution, where the ith column in the plot shows the probability that a given method is assigned rank i, averaged across all tasks. We can see that LCER and StateSAC shares (almost) same probability (about 40%) of ranking 1st, while LCER has a higher probability (35% v.s. 18%) of ranking 2rd
• Score distribution, where a point (τ, fτ ) in the plot means the fraction of runs with score > τ is fτ . The curve of LCER is above StateSAC’s when τ < 0.5, showing than LCER are more stable than StateSAC (regarding the worst performance).

Review summaries:
There has been a school thought lately in RL that good representations for control should focus on abstracting out on only the task-relevant features. In high-dimensional tasks, this is especially useful in making optimal use of a limited neural network capacity and may even help the agent be robust to distractors in the environment. The authors motivate an approach to learning task-relevant representations by having the representations capture only the controllable elements of the environment. Their approach is based on a mutual information loss between actions sequences and final states, and they validate their approach on a series of DMControl task where both high-dimensional pixel observations as well as ground truth latent state representations are available. The idea is to compare the performance of agents that attempt to learn representations of latent states from pixels and compare against an agent that has direct access to the ground truth latent state. In their empirical results they show that: - their new approach (called LCER) is able to outperform alternative approaches on a set of these tasks. On some tasks, it even outperforms an SAC agent with access to the ground truth states. - LCER is able to outperform alternatives when distracting visual features are introduced into these tasks - LCER's representations learned on a source task are able to transfer better to downstream tasks.

The paper proposes a method (LCER) for agents to develop compact task representations from high-dimensional inputs (e.g., images) during unsupervised exploration that are effective for downstream tasks. The paper hypothesizes that encouraging an agent’s representation to focus on changes related to the agent’s actions will be more useful to downstream tasks than representations that try to naively model all changes. In order to allow some non-controllable elements in the representation the paper proposes a blended objective which trades off representing action controllable aspects of state with non-controllable elements of state using an application specific empirically tuned hyperparameter \beta. The paper is set in a fully observable world where a state at time t, S_{t}, is projected to an instantaneous latent state Z_{t}. The idea of controllable elements is implemented by computing mutual information between the latent state projection in the future, Z_{t+k} and the actions taken by the agent I(  Z_{t+k}; A_{t:t+k-1} ).  The uncontrollable information is uncontrollable elements is implemented by computing mutual information between current and future latent states I( Z_{t}; Z_{t+k} ).  Mutual information between latent state and actions uses the neural InfoNCE bound trick. For the second term, the paper instead optimizes Jensen-Shannon divergence between current and future latent representations instead of mutual information, again using a neural bound as the paper claims this is more stable. The paper then shows that this objective can be combined with a standard soft actor-critic architecture to learn representations without extrinsic rewards. In empirical studies on the PlaNet, the paper’s proposed LCER method finds representations that can be used by a policy learner to converge faster than competing methods. The proposed LCER method also compares well to an SAC agent that has access to underlying state representation of the domain (e.g., joint angles, velocities). A second set of experiments on the Distracting  Control Suite which has a lot of extraneous background noise show wider separation between LCER and competing methods. A study shows that increasing the look ahead K improves the quality of representations. 

This paper proposes a novel representation learning flamework, LCER, which enables the agent to extract task-relevant representations from image-based observations. It optimizes the surrogate loss of conditional mutual information (CMI) between future-state embedding and action sequences given past-state embedding. The key contribution seems the balancing between InfoNCE objective (similar to CURL) and JSD-based CMI. The experimental comparison shows that LCER performs best among some image-based RL algorithms (PI-SAC, SLAC. CURL, PlaNet); in 100k DM Control benchmark (6 tasks), and distracted-background settings.