Learning Controllable Elements Oriented Representations for Reinforcement Learning 
Authors: Qi Yi
Jiaming Guo
Rui Zhang
Shaohui Peng
Xing Hu
Xishan Zhang
Ke Tang
Zidong Du
Qi Guo
Yunji Chen
Venue: ICLR 2022 Submitted (2022)
Keywords: reinforcement learning
representation learning
mutual information

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.

Under review as a conference paper at ICLR 2022
LEARNING CONTROLLABLE ELEMENTS ORIENTED REPRESENTATIONS FOR REINFORCEMENT LEARNING
Anonymous authors
Paper under double-blind review
ABSTRACT
Deep Reinforcement Learning (deep RI.) has been successfully applied to solve various decision-making problems in recent years. However, the observations 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, representation learning method that aims to provide RL. al- gorithms with compact and sufficient descriptions of the original observations. Specifically, LCER trains representations to retain the controllable elements 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- 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 mate-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.
INTRODUCTION
Deep Reinforcement Learning (deep RL) has proven its ability solve difficult sequential decision- making problems such as Dota2 (Berner et al. (2019) and StarCraft (Vinyals et al. (2019). However, it still challenge to apply deep RL 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 RI. 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 represen- tations 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., 20206 Schwarzer et al. 2020 Gelada et al. 2019). However, the dynamics of the environment are not al- ways task-relevant(e.gg 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 ob- tained representations can miss important task-relevant information (Gelada al. (019)) 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 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 '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, formally define 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

---

Under review as a conference paper at ICLR 2022
well as sufficient information. By emphasizing the importance of the controllable elements, LCER can capture more rels-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 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. 202T), showing LCER's ability to filter out distracting factors. The implementation of our work is available htttps:/ attps::/anonymous.
4open. scwence/eeeccccccccccccc /oteee//LCCe E/LCER- FC
2 RELATED WORK
In this section, we review the common approaches of representation learning in RL briefly.
Predictive representations in RL. Predictive representations are typically trained 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 model- based RL algorithms. The representations are prone to collapse if without further constraints (Gelada al. 2019). thus a reconstruction loss is usually applied ensure non-trivial representations. How- ever. this approach tries to encode everything, unable to filter out task-irrelevan information. SPR (Schwarzer et al. 2020p introduces a target encoder to avoid collapse and reconstruction, but its ef- fect lacks theoretical guarantee, and it does not emphasize the controllable elements LCER does. There are some methods that train predictive representations from Eiformation-theoreee perspec- tives, such as (van den Oord et al.) 2018 Hjelm et al. 2019 Anand et al. 2019 Lee 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 the action- related 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 representation learning by introducing various action-related objectives such as policy selectivity(Bengio et al. 2017) and contingency awareness awareness(Bellemaan et al. 2012). However, their methods are often built on heuristies and do not give a clear definition of controllable elements. Compared with these works, LCER firstly introduces novel measure of controllable ele- ments (Eq.()) and then develops an efficient method to optimize it without the aid of reconstruction loss.
Bisimulation-based Representations RL. Bisimulation based representations group those states into cluster which are indistinguishable w.r.t. future return, given any action sequence tested. The indistinguishabiiity 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 dis- tracting environments. However, these methods rely on 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 knowl- edge to constraint the representations (Jonschkowski Brock 2015 Jonschkowski et al., 2017 Thomas et al.) 201 8). 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.
2

---

Under review as a conference paper at ICLR 2022
3 PRELIMINARIES
3.1 NOTATION
We assume the underlying environment is Markov decision process (MDP), described by the tuple M = (S, A, P. R, 7). where S is the state space, A the action space, P : St x A₁ x S - [0, 1] the transition probability function which determines the distribution of next state given current state and action, and [0, discount factor. Given the current state E S, an agent choose its action a E according to a policy function a ~ <-(ss). This action will update the system state to new state s' according to the transition function P. and then reward 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 variable of state if there no ambiguity. We use $ to present the state sequence from to + k i.e. St.tk == (st. St-1, ..., S(+k). These notations be extended to other variables such actions and rewards. We use Φ 5 + Z to denote the embedding function of LCER, which maps 8 E S to latent space Z.
of (i.e. S) to refer to the random
3.2 MUTUAL INFORMATION
The mutual (MfOrmation(MII) of two random variables X, Y is defined as:
(( = (xI(Pxr|IPxxP))
(1)
where DEL is the Beilback--eibbes divergence that measures the difference of two distributions.
The conditional mutual information(CMII) I(X; Y)) is defined the expected value of I(X; Y) given the value of Z:
I(X; Z|Z) Ez~ 122).
LEARNING CONTROLLABLE ELEMENTS ORIENTED REPRESENTATIONS
Given an embedding function: ø S + Z, we argue that the number of controllable elements that & contains can be measured by the conditional mutual information(CMII):
(Ø) == Autt--oo))
Intuitively, 1y measures the information that ((S++&) and Act+s-1 shares given o(S)). Note that the non-controllable elements captured by o do not contribute to 13 because they are not related to the actions. Training representations to maximize (3) is an efficient way to make the rep- resentations capture controllable elements. However, representations obtained in this way may be too compressive to provide sufficient information for solving tasks in some cases, because not all task-relevant elements 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 surrogate objective derived from Eq. which additionally provides a mechanism to control the degree of compression. In Section describe how optimize this objective, leading to our final algorithm LCER.
we
4.1 THE OPTIMIZATION OBJECTIVE FOR LCER
By the chain rule of mutual information, we can break the T₁ defined in Eq. into the subtraction of two terms: (d) = T ((()())) A ø(S++x)) - (()())) @(S++*)). The first term tries to capture predictive information, while the second term serves as compression term that aims to filter out action-irrelevan information. To make á capture the controllable elements and maintain sufficient information the same time, we introduce hyper-parameter 3 to control the degree of compression:
( 3) == (((((S)) Ar++k A (anil)))_(((((65 (ō(S)) ((()(()))(((+)))))
(4)
By saying "compressive", we mean that these representations only encode controllable elements and filter any other information.
3

---

Under review as a conference paper at ICLR 2022
By choosing different 3 when training to maximize Eq.44), control the degree of compres- sion, making it possible to adapt to different environments by choosing different B.
However, directly maximizing (Ø; 3) is intractable, because maximizing the first term and mini- mizing the second term in Eq. at the same time have an antagonistic effect to each other, thus difficult tune in practice. In fact, we find this will often lead to unstable training process in our early experiments.
To tackle this problem, substitute [(((()) (() with Ia (0; 3 0) (x 3 1):
[x( 3) = ((t) Aut+s-1]; - (Tx( 0) (0:
3) + [(0(S++)); Amt4k-1l0(55) .3. (5)
4
In Eq.(5), we maximize two terms (J) and JBP at the same time, and their coefficients and | respectively. The first term J is closely related to the predictive information encoded by Φ. The second term Jž is related the controllable elements, which will play an important role in distracting environments as shown in Section 5.3
In practice, sum over different choices of ki leading the final objective maximize:
K
K (): =
())
(6)
4.2 HOW TO OPTIMIZE THE OBJECTIVE FOR LCER
To maximize 1" (o. 3), 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 J and J respectively that are used in this paper, and finally put them together to get our final loss for training
Optimizing J.. In order to maximize JL((), we propose to maximize the InfoNCE (van den Oord
al. 2018 lower bound of (Ø):
1 E log B
exp(f(o(s)) (($5
E (xp(((o((s) ,f ot. 0( (((++)))))
where f is any scalar function, B is the batch size, and s -1 is transition sequence from st to s' There are other alternatives to InfoNCE such NWJ (Poole et al. 2019) and MINE (Belghazi et al. 2018), however InfoNCE can often lead to better representations as shown in Tsschannen et al. 2020 In order to efficiently optimize Jr for all < k S K, we recursively define vi to encode the information of (d(s)), a d(a++-1)::
(s())) (h(W (h( ( otherwisc
kf
(8)
and set f to be the bi-linear inner product of the and (5 (65
[(bss), alatt-ii (+(+4))
Note that the dit is shared across all JJ. thus we do not need to re-compute it. The h in Eq. quite similar to "prediction" network, but is trained together with á instead of prediction errors.
Optimizing J2. In order to maximize J2(Ø), we turn to maximizing a Sensen-Shannon CMI estimator Eq. because it is proven to be more stable in our experiments, which is also sug- gested in prior works (Hjelm et al., 2019, Sanchez et al. 2020). We introduce a statistics network Tk S x A S + R and maximize:
1 8
jE(0) = B
Σ
1og(1 + t Te0h + log(1 + Talof.a20) 19(+)))
(10)
1
4

---

Under review as a conference paper at ICLR 2022
Table 1: The DMControl 100K benchmarks, which report the performance 100000 environment steps in the PlaNet benchmarks. Our method LCER achieves the best performance on of 6 tasks, and is the only one that outperforms StateSAC w.r. the average returns.
Environment
LCER(OurS)
PI-SAC
SLAC
CURL
Dreamer
StateSAC
-aripole-swingup
790 40
772 + 49
327 + 44
582 + 146
234
163
860 =
cheetah-run
505 + 38
271 57
413 - 67
299 = 48
158 = 133
206 = 28
walker-walk
622 + 127
431 + 80
528 41
403 + 24
216 + 124
664 + 111
reacher-easy
704 160
687 + 133
342 + 96
538 + 233
147 117
774 + 111
finger-spin
961 + 56
942 - 84
951
+ 32
767 = 56
33 - 42
749 + 234
ball in cup-catch
933 25
878 + 117
917 +
769 43
172 215
951 + 18
Average Returns
752
663
579
559
160
700
where ((()....jBB)) are obtained by randomly shuffling 1,.., B. The derivation of Eq. 10 is pro-
vided Appendix B In practice, we choose T'( 01:44--11 Pc+*) T(b+4k: Φctk). where th+4 is a latent variable defined in Eq. and T : Z x Z + R is network shared across different Tμ.
Final loss for training 6. Inspired by (He et al. 2020), we propose to utilize a target encoder otorg stabilize the training process, whose weights the exponential moving average(EMA of Φb Specifically, the embeddings of si+ a (A 2 1) are calculated by Starg.
pieces together, we obtain our loss for & to minimize:
K
LK (0: w,T)= =
Σ (1-
P(0)-
B MΣ [1 - -1-))
1
exp((Vi+H) )² (rarg(((+4))
E/= expl( + F Grarg (6 (0(+x)
+ (log(1 + exp (-T Peora (3 + log(1 + exp (T Drarn ( ))),
(11)
where B is the batch size, h, vit are defined in Eq 8). W is defined in Eq. 9..9 t is latent defined in the same way with whf but it encodes (())) and ( (instead of ( 1); T is
j(t)
scalar function, and Otorg is the target encoder. The parameters Φ, h W, T are training together
by minimizing I,K_ In practice, Eq.011 can be used as an auxiliary loss for any RL algorithms such as SAC (Haarnoja et al. 2018).
EXPERIMENTS
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 's strength to filter out distracting factors on the Distract- ing Control Suite, and show that it not enough to obtain good representations just by encoding predictive information. In Section 5.3 we investigate the effect of K and 3, which are the most im- portant 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 equipped with SAC (Haarnoja et al. 2018 Most of the parameters in LCER the same with PI-SAC (Lee al. 2020b), including the batch size, learning rate, action repeat, and the network architecture of o. see Appendix A for more details. Following (Laskin et al.) Lee al. 22200a b), we report the performance using the true environment steps. which is invariant to the choice of action repeat. Unless specially specified, We set 5, = 0.1 and run every experiment using 10 random seeds.
FFMA is widely used trick in RI. For example, in SAC, the weights of target Q the F.MA of Q.
5

---

Under review as a conference paper at ICLR 2022
cartpole-swinguu
cheetah--un
reacher-ears -eay
kymr
B--ACC
1000
Eskoode Reeuum
Epteods
nasa noss nala 1:5 nT 6.5 nin
Enviro
Environment Staps
Environment
ball catch
walker
1000
LEET 5.AC Pancdae
1000
SAC
SLAC
M
Eppootee Ruttur
600
Fyycodde Surturn
BLAC
Stracdac
Eppsole
200
-000
200
-200
0000 0022 nasa 0,75 0d00 6.ss 0.59 0.75 0.20
8.00
2.04
6.12
Emvtronment Steps (x10))
Environmmnt Steps (x:1PP)
Environment Steps (-11P)
The PlaNet Benchmarks. The performance of StateSAC can be seen as an upper bound because it learns directly from internal physical states rather than pixels. Our algorithm LCER consistently achieves better or comparable performance than other methods(except StateSAC) on all environments
Figure 1:
5.1 EVALUATION ON SAMPLE EFFICIENCY
In this section, We evaluate LCER on the PlaNet benchmark, which consists of six challenging con- trol tasks from the DMControl Suite TTssa 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. al. 2020b
We equip LCER with SAC (Harrojj et al. 2018) and compare LCER to 3 leading representation PI-SAC (Lee et al. 2020b). algorithms CURL (Laskin at. a...)..S..A... SLL. (Lee al. a1.20)))
CURL introduces contrastive learning into RL and actually trains representations to encode infor- mation that invariant under data augmentation. SLAC efficiently combines both representation learning and control into single objective via Bayesian inference, but it relies on reconstruction loss. PI-SAC tries to encode predictive information, while LCER focuses on encoding controllable elements. All these methods Kostrikov et al. 2020b). we use the version of PI-SAC that performs one gradient update per envi- ronment step ensure a fair comparison. We also compare LCER with 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 ternal physical states provided by DMControl Suite. The performance of At can be seen as an upper bound since 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.
also built SAC (Haarnoja et al. 2018). Following (Laskin et al.)
In Figure [ we report the learning curves of LCER, PI-SAC (Lee et al., (0206)) SLAC (Lee al.) 2020a), and Dreamer (Hafner et al.] 2020). LCER consistently achieves better or comparable per- formance than other methods. In Table We report the performance each algorithm 100k environment steps. LCER is the only algorithm that enables pixel-based SAC to outperform state- based 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 met- rics on DMControl 100k benchmark, see Section C for details.
In practice, pixel-based observations often equipped with data augmentation to larizer (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, proper data augmentation method is not always available for general forms observations. Therefore, We further test our method LCER in the
regu-
"hhe performance data of SLAAC Dreamer, CURL. provided by the authors of the corresponding papers.

---

Under review as a conference paper at ICLR 2022
cartpole-wwingup
cheelethrun It-um
finger-spin
000
1100)
Epteote Retur
Epteede Sentum
000
200
SBC
0.2
0.3
02
Environment Steps (-:1PP)
Envtro ment Stops (x:1P))
Environment Steps (-:1P))
Figure 2: The learning curves when data augmentation is not available. LCER can still continuously improve the performance to great extent without the help of data augmentation, showing LCER's ability to handle different data formats (beyond pixels) of observations.
D.cheelllh--uun
walkerrwwalk
D.finger-spin
D. cheeetah-rn
D. welkerr-welk
D. fifiger-ssin
2001
P-A.CC
Epledde Renurn
600
200
000
0.6
0.2
0.6
Enviromment Steps (v:17))
Envrro mont Stops (x:11))
Environment Steps (-:17)
The Learning curves on Distracting Control Suite, challenging benchmark that contains several kinds of visual distractions. Top row: the observations each task. Bottom learning curves on each task. The relatively poor performance of PI-SAC indicates that encoding predictive information does not ensure representations of high quality. Although these tasks are so difficult that SAC fails to solve, LCER can still improve the agent's 's performance gradually, showing LCER's robustness to distractions.
Figure 3:
the
situation where no data augmentation available. As shown in Figure 2 the performance to a great extent without the help of data augmentation. Note that some sample- efficient RL algorithms such as DrQ (Kostrikov et al. (2020b) and RAD (Laskin et al. 2020 that rely data augmentation can not be extended such situation naturally.
LCER can still improve
Our experiments on the PlaNet benchmark show that our proposed method LCER can obtain repre- sentations that help solve the tasks, and thus improve the sample efficiency.
5.2 ROBUSTNESS TO DISTRACTIONS
In this section, we test LCER on the Distracting Control Suite (Stone et al. 2021 }. 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 great extent (Stone et al.) 2021). We aim to show LCER's '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 shown in (Zhang al. 2020).

---

Under review as a conference paper at ICLR 2022
cortpole-wwingup
cheetah-run
bell Li__n,nnnnnn cup--cophcc
2000
700
Ealvod Nerumm
Eplioda Batturn
600
30:0
<00-
300
400
spert
spent
agert
LCS
110
200
UC8 kel
L00
UCC
3.00
0.30
2.
0.30
0.0
22
0.233 0.32 0.00 0.333 0:32 0.35 0.00
Environment Steps (-237)
Environment Steps (U
Enviromment Steps (x179)
(a) The effect of on performance.
D.cheetal-run
D.walkarwwalk
D.finger-spin
400
700
JCBN
Espi0se Retumm
25.0
LLBN
640
JUM
000
400
300
100-
100
Steps (-127)
Steps
(b) The effect of performance.
The effect of K and B. Figure ((a) shows that larger K generally leads to better perfor- mance. The effect of 3 is more complicated as shown in Figure 4(b) Larger 3 does lead to better performance within certain range, however when 3 goes beyond this range(i.e.8 hinder the performance. In practice, we can gradually increase the value of ,3 to search for optimal performance.
Figure 4:
1). it may
In Figure 3] we run experiments on three control tasks from the Distracting Control Suite: D.cheetah- run, D. D..walee-vve walk, D.finger-spir" In all tasks, we set the difficulty to 'easy' and use a dynam- ically varying background (refer to (Stone et al. [021) for more details). We choose 3 = 0.4 for D.cheetah-run and D.finger-spin, and 3 = for D.walker-wallf) 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 inrceinrelevant 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 's performance, LCER suffer less than other methods, showing LCER's 's robustness these distractions.
5.3 THE EFFECT OF K AND 3
In this section, we try to specify the effect of K and B, which are the most important hyper- parameters in LCER.
To investigate the effect of K, we set 9 = 0 and report the performance for K = 0, 1, 5 in Figure 4(a) Note that K 0 is 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 our special design of Eq.(1p. To investigated the effect of (3, report the performance for 3 = 0.1, 0.2, 0.4, 1 in Figure 4(b) As shown in Figure 4(b) larger 3 does lead to better performance
*We use "D." to distinguish these tasks from ones in the DMControl Suite.
'We sweep over (0.1,0.2,0.4.1) to get the optimal B. The learning curves for these f3' are provided in
Figure 4(b)
8

---

Under review as a conference paper at ICLR 2022
Caltpole-belance
walker-run
hopper-hop
Apper
synr LOUetad
Edioode Recur
Epicoda Benum
125
200
422
300
Epieote
220
300
-2
0.30
2.04
0.00
0.00
0.23
0.12
0.14
0.
0.2
0.0
0.2
Environment Steps 1-10
Enwircmment Steps (1F)
Environment Steps (x:97)
Figure 5: The learning curves of transferring obtained representations to target tasks. From left to right, the source tasks are cartpole-swinguy hopper-stand, walker-walk respectively. "LCER" and "SAC transferred" refer to using the representations obtained by LCER and vanilla SAC respec- tively, and "SAC" refers training from scratch using vanilla SAC on target tasks. The represen- tation obtained by LCER can encode more useful information than vanilla SAC, leading to better performance on target tasks.
within a certain range, however when 3 goes beyond this range (e.g.8 = 1), it may hinder the performance in some tasks. The result in Figure ( that the optimal 3 should be chosen according to the given environment. By introducing hyper-parameter B3 LCER is able to strike balance the proportion of controllable elements that á captures, allowing for searching for better performance.
5.4 THE QUALITY OF REPRESENTATIONS
In this section, try to measure the quality of by generalizing 6 to unseen tasks. Specifically, we train & on a source task, froze é, and then apply it to target task within the same domain. If LCER can capture the controllable elements efficiently, then the representations obtained 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). Figure we report learning curves on target t SAC, (CER((sf.) and LCER(self tsf.). SAC is trained directly on target tasks. (tf)( and (cE.(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(seff.ssh) 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 well LCER(self tsf.) target tasks. However, when solving target tasks needs to exploit more environment dynamics which are not explored in source tasks (such as hopper-stand", (CER((sf..) shows inferior performance. We conclude that the transfer ability of L.ER 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-wall walk and cartpole-swingup LCER((sf.) can learn as effective representations as (seR(self.tst) does even source tasks.
6 CONCLUSION
In this paper, we introduce the LCER, an efficient framework to train representations that can re- flect the action-related dynamics of the environment. By focusing 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.
'including cartpole-swingup + cartpole-balance, walker-walk -+ walker-run, and hopper-stand + hopper- hop.
This because hopper-hop the agent needs move forward, however in hopper-stand it does

---

Under review as a conference paper at ICLR 2022
REFERENCES
Learning to poke by poking: Experiential learning of intuitive physics. In Advances in Neural Infor- mation Processing Systems 29: Annual Conference on Neural Information Processing Systems, NeurIPS, 2016.
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Marc G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. CoRR, 2021a.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare G. Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. ArXiv, 2021b.
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Mare-Alexandre Côté, and Devon Hjelm. Unsupervised state representation learning in atari. Annual Conference on Neural Infor- mation Processing Systems 2019, NeurlPS, 2019.
Adrià Puigdomènech Badia, Pablo Sprechmann, Alex Vitvitskyi, Zhaohan Daniel Guo, Bilal Piot, Steven Kapturowski. Olivier Tieleman, Martín Arjovsky, Alexander Pritzel, Andrew Bolt, and Charles Blundell. Never give up: Learning directed exploration strategies. In 8th International Conference on Learning Representations, ICLR, 2020.
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, R. Devon Hjelm, and Aaron C. Courville. Mutual information neural estimation. Proceedings of the 35th International Conference Machine Learning, ICML, 2018.
Mare G. Bellemare, Joel Veness, and Michael Bowling. Investigating contingency awareness using atari 2600 games. In Proceedings of the Twenty-Sixth AAAI Conference on Arlificial Intelligence, 2012.
Emmanuel Bengio, Valentin Thomas, Joelle Pineau, Doina Precup, and Yoshua Bengio. Indepen- dently controllable features. CoRR, 2017.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Christopher Hesse, Rafal Józefowicz, Scott Gray, Catherine Olsson, Jakub Pachocki, Michael Petrov, Henrique Pondé de Oliveira Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, Ilya Sutskever, Jie Tang, Filip Wolski, and Susan Zhang. Dota with large scale deep reinforcement learning. CORR, 2019.
Pengyu Cheng, Weituo Hao, Shuyang Dai, Jiachang Liu, Zhe Gan, and Lawrence Carin. CLUB: A contrastive log-ratio upper bound mutual information. Proceedings of the 37th International Conference on Machine Learning, ICML, 2020.
Norman Ferns and doina Precup. Bisimulation metrics are optimal value functions. Proceedings of the Thirtieth Conference Uncertainty in Artificial Intelligence, UAI, 2014.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. Proceedings of the 36th International Conference on Machine Learning, ICML, 2019.
Zhaohan Daniel Guo, Bernardo Ávila Pires, Bilal Piot, Jean nan-Bastien Grill, Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar. Bootstrap latent-predictive representations for mul- titask reinforcement learning. Proceedings of the 37th International Conference on Machine Learning, ICML, 2020.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine. Soft actor-critic algo- rithms and applications. CoRR, 2018.
Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. Proceedings of the 36th International Conference Machine Learning, ICMI,,0155
10

---

Under review as a conference paper at ICLR 2022
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learn- ing behaviors by latent imagination. In 8th International Conference on Learning Representa- tions, ICLR, 2020.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick. Momentum contrast for unsupervised visual representation learning. 2020 IEEE/CVF Conference on Computer Vision and Puttern Recognition, CVPR, 2020.
R. Devon Hjelm, Alex Fedorov, Samuel Marchhidon, Karan Grewal, Philip Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. 7th International Conference on Learning Representations ICLR, 2019.
Rico Jonschkowski and Oliver Brock. Learning state representations with robotic priors. Auton. Robots, 2015.
Rico Jonschkowski, Roland Hafner, Jonathan Scholz, and Martin A. Riedmiller. Pves: Position- velocity encoders for unsupervised learning of structured state representations. CoRR, 2017.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. CoRR, 2020a.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. CoRR, 2020b.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: contrastive unsupervised representa- tions for reinforcement learning. Proceedings of the 37th International Conference on Machine Learning. ICML 2020.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Re- inforcement learning with augmented data. Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS, 2020.
Alex X. Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS, 2020a.
Kuang- Huei Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama. Predictive information accelerates learning in RL. Advances in Neural Informa- fion Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurlPS, 2020b.
Xiao Ma, Siwei Chen, David Hsu, and Wee Sun Lee. Contrastive variational reinforcement learning for complex observations. 4th Conference on Robot Learning, CoRL. 2020, 2020.
Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and R. Devon Hjelm. Deep reinforcement and infomax learning. Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020.
Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural sam- plers using variational divergence minimization. In Advances in Neural Information Processing Systems 29; Annual Conference on Neural Information Processing Systems 2016, NeurlPS, 2016.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS, 2019.
Ben Poole, Sherjil Ozair, Aäron den Oord, Alex Alemi, and George Tucker. On variational bounds mutual information. Proceedings of the 36th International Conference on Machine Learning, ICML. 2019.
11

---

Under review as a conference paper at ICLR 2022
Eduardo Hugo Sanchez, Mathieu Serrurier, and Mathias Ortner. Learning disentangled representa- tions via mutual information estimation. 2020.
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bachman. Data-efficient reinforcement learning with momentum predictive representations. CoRR, 2020.
Rui Shu, Tung Nguyen, Yinlam Chow, Tuan Pham, Khoat Than, Mohammad Ghavamzadeh, Stefano Ermon, and Hung Bui. Predictive coding for locally-linear control. Proceedings of the 37th International Conference Machine Learning, ICML, 2020.
Alessandro Sordoni, Nouha Dziri, Hannes Schulz, Geoffrey J. Gordon, Philip Bachman, and Remi Tachet des Combes. Decomposed mutual information estimation for contrastive repre- sentation learning. In Proceedings of the 38th International Conference Machine Learning, ICML, 2021.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite challenging benchmark for reinforcement learning from pixels. CoRR, 2021.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmalcki, Josh Merel, Andrew Lefrannq, Timothy P. Lillierap, and Martin A. Riedmiller. Deepmind control suite. CoRR, 2018.
Valentin Thomas, Emmanuel Bengio, William Fedus, Jules Pondard, Philippe Beaudoin, Hugo Larochelle, Joelle Pineau, doina Precup, and Yoshua Bengio. Disentangling the independently controllable factors of variation by interacting with the world. CoRR, 2018.
Michael Tschannen, Josip Djolonga, Paul K. Rubenstein, Sylvain Gelly, and Mario Lucic. On mutual information maximization for representation learning. 8th International Conference on Learning Representations, ICLR, 2020.
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. CoRR, 2018.
Oriol Vinyals, Igor Babuschkin, Wojciech M. Czarnecki, Michaël Mathieu, Andrew Dudzik, Juny- oung Chung, David H. Choi, Richard Powell, Timo Ewalds, Petko Georgiev, Junhyuk Oh, Dan Horgan, Manuel Kroiss, Ivo Danihelka, Aja Huang, Laurent Sifre, Trevor Cai, John P. Agapiou, Max Jaderberg, Alexander Sasha Vezhnevets, Rémi Leblond, Tobias Pohlen, Valentin Dalibard, David Budden, Yury Sulsky, James Molloy, Tom I.. Paine, Çaglar Gülçehre, Ziyu Wang, To- bias Pfaff, Yuhuai Wu, Roman Ring, Dani Yogatama, Dario Wünsch, Katrina McKinney, Oliver Smith, Tom Schaul, Timothy P. Lillicrap. Koray Kavukcuoglu, Demis Hassabis, Chris Apps, and David Silver. Grandmaster level in starcraft II using multi-agent reinforcement learning. Nat., 2019.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improv- ing sample efficiency in model-free reinforcement learning from images. 2019.
Amy Zhang, Harsh Satija, and Joelle Pineau. Decoupling dynamics and reward for transfer learning. In 6th International Conference on Learning Representations, ICLR, 2018.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invari- ant representations for reinforcement learning without reconstruction. CoRR, 2020.
12

---

Under review as a conference paper at ICLR 2022
A IMPLEMENTATION
In this section, we explain the implementation in details.
Action Repeat Following Lee 2020b; Laskin al.j Lee et 2020a Hafner al. 2019),
we use different action repeat to difterent environments. The action repeat is parameter that deter- mine how frequently the agent is supposed to make a decision. For fair comparison, we adopt the action repeat from (Lee et al., 2020b a). Specifically, for cheetah-run, cartpole-swingup reacher- easy, and ball -upp-ccccccccch we use action repeat of 4. For walker-walk and finger-spin, we use and respectively. Note that the action repeat parameters can directly affect the performance.
Initial collecting step Following (Lee et al. 2020b a). we pre-train our model using a small amount of random data. Specifically, first interact with the environment using random pol- icy to collect initial experience, and then using the experience to pre-train our encoder Φ. Note that the amount of initial random data has been taken into account in all experiments. For cheetah-run, walker-walk, reacher-easy, and ba cup-cctth, we collect random data for 10000 true environ- ment steps and then train for 10000 iterations. For finger-spin, 10000 true environment steps and 50000 iterations respectively. For cartpole-swingun 4000 true environment steps and 1000 itera- tions.
Network Architectures We use the encoder(i.e. σ) that shares the same architecture with (Laskin et al., Lee et al., 2020b), which consists convolution layers(32 x filters with stride 2 for the first layer, and 32 3x3 with stride for the rest 3 layers) and fully-connected layer that flattens the output of convolution layers and then maps it into a latent space with 50 dimensions. For the h in (8, we use MLPs with 128-d hidden layers, whose input and output are both of 50 dimensions. For the actor and critic, we both use MLPs with two 512-d hidden layers. We use tanh to squash the output of the actor make sure it fits the action bound of the environment.
Other Hyperparameters We list the other hyperparamters in Table are the same with (Lee et al. 2020b).
Most of these paramters
Hyperparameter
Value
Batch size
256
Learning rate
3e-4
Replay buffer size
100000
Dimension of latent space
50
Stacked frames
Optimizer (81, 82) Adam
Adam (0.5, 0.999) for a, (0.9, 0.999) for others
target update interval
1
actor log std
[-10, 2]
Q EMA TÜ
0.005
Encoder EMA T
0.05
Initial temperature
0.1
Discount
0.99
Non-linearity
ReLU
Entropy Target
-dim(A)
Table The hyperparameters of LCER.
13

---

Under review as a conference paper at ICLR 2022
B ESTIMATION OF CONDITIONAL MUTUAL INFORMATION
In this section, we describe how we optimize JJSD defined in Section 4.2 In the followings, the (Se+*)), A1 (S₁) denoted by X, Y, Z respectively.
We first borrow some theoretical results from Nowozin et al. 2016
JSD ( (x; = E Px Prie)
=
Priz) 2 sup Ex..
(12)
Here, fUSD is the generator function of Jensen-Shannon divergence: fusD(u) = - 1 ulog u, and fjsp its corresponding convex conjugate function: (T) sup, (ut
fusD(u)) = 1og(2 - exp(t)). We choose V(x, v, = log(2) - log(1 + exp( - ( v,,) then:
JSD (X:YIZ) 2 sup Xxy
(2 YexxxX2en-.e//sssiie......
sup Ex y XYZ Xx log (x..))))) X log + Tx..,,))) log
Tk
sup - E. E Xy, log(1 + -((x)))))))))) - E log(1 + (2(2,22)) log 2.
Ta
(13)
In 1 we suppose p(ylz) = p(y) Aut+k-1 independent from ((S))) for simplicity. Such an assumption is appropriate in our case, because (St, 1. St+k) is sampled from large replay buffer(= 10a) whose action policy is quite non-stationary due to policy updates.
To estimate the RHS of Eq.(1)) we first randomly sample B transitions {sf, a 1: + from the replay buffer. Assume shuffling, then Ex,y, 2 (-) can be approximated by 1
is a permutation of (1, B) obtained by random
B 2 (-) and E xx,
1 Egpopte 26) by B litee tt (-). Therefore, the RHS of Eq. (1) can be estimated by:
1 Σ log(1 +e - Tulheaaa wilw + log(1 + Faloh (14)
P( =
Here, omit the "2 log 2" in 13 for simplicity.
14

---

Under review as a conference paper at ICLR 2022
LCER)pur))
100k steps
1.00
100
100k steps
Fraction of rurn with score
Fraction (in %)
LCER/ours)
PLSAC
0.75
Distribution
60
0.50
0.25
0.00
0
0.0
0.2
0.4
0.8
1.0
1
Normalized Score (x)
Ranking
(a) Score Distributions
(b) Rank comparisons
Figure 6: Additional aggregate metrics on DMControl 100K benchmark.
C ADDITIONAL AGGREGATE METRICS ON DMCONTROL 100K BENCHMARK
In this section, provides additional comparison DMControl 100k benchmark using aggregate metrics from rliable (Agarwal et al. 20216).
We normalize all scores to [0, 1] by dividing by the maximum score (=1000). The metrics shown in Figure 6 which include:
Ranks distribution, where the 1th column in the plot shows the probability that given method is assigned rank i, averaged across all tasks. We can see that LCER and Sta shares (almost) same probability (about 40%) of ranking st while LCER has a higher probability (35% v.s. 18%) of ranking 2nd
Score distribution, where point fr) in the plot means the fraction of runs with score > fr. The curve of LCER is above StateSAC's when 7 0.5, showing than LCER are more stable than Sta (regarding the worst performance).
15


Reviews:
[{"review_id": "xtZ9z7kyl0V", "reviewer": "Reviewer_uWv6", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n- This paper is well structured and easy to follow.\n- The empirical results in 100k DM Control benchmark (6 tasks), and even distracted-background settings, show LCER performs best among some image-based RL algorithms, which successfully reveals that learned representation of LCER captures the controllable features in the image observations.\n\n## Weaknesses\n- I think the derivation of LCER seems a bit misleading. The objective that LCER actually optimizes may not be introduced from the CMI directly; it starts with CMI between future-state embedding and action sequences given past-state embedding (Eq. 3), introduces the bottleneck parameter $\\beta$ (similar to Lee et al. 2020) (Eq. 4), and deformulates the objective again (Eq. 5). In Eq. 5, $J_k^2$ is the same term as Eq. 3, which might be confusing. It can be simply argued that LCER introduces the CMI term into CURL objective ($J_k^1$) with some ratio $\\beta>0$, and also introduces the surrogate form of CMI based on JSD (Eq. 10). In addition, balancing $\\beta$ seems very important hyper-parameters for LCER (I guess even in standard DM Control, where $\\beta=0.1$ is used. If LCER learns controllable elements from CMI, $\\beta$ should be 1, but isn't).\n- I am also a bit confused about the connection and difference between LCER and PI-SAC, since both of them optimize the CMI objective as an auxiliary loss for RL. While the authors just say, \"PI-SAC tries to encode predictive information, while LCER focuses on encoding controllable elements\", further clarification about this is required to emphasize the contribution of this paper.\n- For PI-SAC baselines, the authors can include the \"gradient\\_step=2, 4\" results for a fair comparison, since it improves the performance according to its codebase (https://github.com/google-research/pisac). In contrast, LCER can also utilize such extra gradient steps. Moreover, I find that while LCER uses $K=5$ sequence length for InfoNCE + CMI objective, PI-SAC only uses $K=3$ (from Appendix A of Lee et al. 2020), which might be \"unfair\" comparisons, considering the experiments excepting multiple gradient steps of PI-SAC.\n- I think the authors may also include the results of Dreamer (Hafner et al. 2020) for 100k benchmark since it generally achieves better results than PlaNet; in addition, CURL and PI-SAC compared against it in their original papers (+ code is available online).\n- In distracting-environments settings, the authors may also include Contrastive Variational  Reinforcement Learning (CVRL) as a baseline, proposed by Ma et al. (2020), which achieves better results compared to Dreamer, image-based SAC, and image-based D4PG (+ code is available online). In addition, Ma et al. (2020) also take a contrastive approach to handle distracting background, so it should be added as related work.\n\n\n### Reference\nLee et al. Predictive Information Accelerates Learning in RL. Advances in Neural Information Processing Systems (2020).\n\nHafner et al. Dream to Control: Learning Behaviors by Latent Imagination. International Conference on Learning Representations (2020).\n\nMa et al. Contrastive Variational Reinforcement Learning for Complex Observations. Conference on Robot Learning (2020).\n", "summary_of_the_paper": "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.", "summary_of_the_review": "As discussed in Main Review in detail, I think the derivation of LCER seems a bit misleading, and further clarification compared to PI-SAC is required (since both PI-SAC and LCER optimize the CMI objective). In addition, some important baselines (PI-SAC w/ gs=2 or gs=4, Dreamer, CVRL) seem missing in the current manuscript. Considering these aspects, I larn towards rejection.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "naO6eSGOW1b", "reviewer": "Reviewer_pqPt", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n\n1.  The paper is generally well written and well motivated. While section 4 required multiple passes (likely due to multiple approximations that are used), the high-level motivation and ideas are easy to grasp.\n\n2. I like the idea - it is quite simple and elegant - learn representations that retain information about the controllable elements of the environment. \n\n3. I liked the experiments and ablations that the paper includes. While I have several issues with the claims/interpretation-of-results, I think the experiment design is well thought.\n\n## Areas for improvement\n\n1. Several claims are not well-substantiated. \n\n    * In some cases, it could just be a matter of rewording. For example \"LCER is the first representation learning algorithm that enables the pixel-based SAC to outperform state-based SAC on the DMControl100K benchmark\". From Table 1, LCER lags behind StateSAC on 5 (out of 6 environments). While its performance is impressive, the claim itself is not supported. Comparing the average performance (across 6 envs) is misleading because of 1 outlier environment. \n\n    * In other cases, the claims are not tested. For example, there is no experiment to show that the representations indeed \"retain the controllable elements of the environment\" (which is the key motivation behind the work).\n\n    * In several cases, the results are quite close to each other (eg Figure 1 walker walk, reacher-easy, ball-in-cup, cartpole...) In absence of statistical significance test, it is difficult to conclude that the proposed approach is indeed better. Same applies to several results in the tables as well.\n\n2. When evaluation robustness to distractors, better baselines (like Zhang et al 2020) should be used.\n\n## Further questions\n\nNote that I did not consider these points as weakness as I did not fully understand these aspects.\n\n1. In equation 4, what is the meaning of \"$I([\\phi(S_t), A_{t:t+k-1}]$\" ie what is the \"[...]\" operator doing?\n2. In equation 3, is only A conditioned on $\\phi(S_t)$ or both $\\phi(S_{t+k})$ and $A$ are conditioned? \n3. In page 3, Section 4, the authors correctly point out \"For example, in navigation tasks, the position of the goal\nis task-relevant but not controllable by the agent\" as a motivation for controlling the degree of compression. However, their experiments did not have any goal conditioned environments. In that case, would it make sense to directly use equation 3 (atleast in the context of experiments considered in the paper).\n4. In page 4, the authors mentions \"To tackle this problem, we substitute I...\" (right before equation 5). Could they describe why is this approximation valid?\n5. What is the runtime cost of LCER compared to the baseline? I expect it to be much higher than the baseline (given it is used k-step lookaheads to compute the losses. Reporting the runtime comparison will be useful for a holistic comparison. Similarly, I expect LCER to be using much more params and would like to see a comparison of those numbers.\n6. What is the effect of using the target encoder? How severe is effect if target encoder isnt used?\n7. In page 6, the paper mentions \"All parameters of StateSAC..\" do you mean all hyper-parameters of StateSAC... ?\n8. In figure 2, cartpole results seems to be stopped too early. The two baselines may be able to overtake the red curve.\n9. In page 7, the paper mentions \"In LCER, we also randomly shift the observations by [−4, 4] pixels\". Is this aimed to provide translation invariance? If yes, isnt a conv net already providing that? Is this augmentation applied to the baselines as well?", "summary_of_the_paper": "The paper proposes the LCER algorithm that 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. The main experiments use DMControl suite and includes several ablation experiments.\n", "summary_of_the_review": "The paper proposes a simple and elegant idea and shows that the approach \"works\" in practice. However, it is difficult to evaluate the effectiveness/significance of the approach. I look forward to the author's response and interacting with them to understand their approach better.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "KOuMH4SWCA", "reviewer": "Reviewer_WDLZ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths:**\n- paper was clearly written and well-motivated\n- approach can be applied to any task\n- clear mathematical exposition of the loss function used\n- demonstration of better sample complexity against other baselines including CURL, PlaNet, SLAC, and PI-SAC.\n- demonstration of robustness to distractors compared to a PI-SLAC baseline\n- the authors have explored the robustness of the algorithms to the two new hyperparameters introduced, $\\beta$ and $K$\n\n\n**Comments:**\nI do like the idea of using controllability to inform compression of state representations -- after all, what is important to the agent is what it can control about the environment. Yet there may be situations where this approach can be too compressive. Perhaps the most important situation is where there are non-controllable task-informative features. The authors do point out one example, the goal location. But there are in principle many others. For instance, the directions for solving a task (e.g. a map showing the solution to a maze) may be painted on a wall in the environment -- the agent could still solve the task by completely ignoring the directions, but it would do much better if it could exploit that information. Another limitation with this approach is if the environment dynamics change. For instance, suppose an agent wanted to navigate a drone to a goal. Controllability will differ depending on wind conditions. So if the agent learned a representation under e.g. no wind, that representation may be less effective if the agent had to navigate in the presence of strong and volatile wind. These limitations are by no means reasons to disqualify the paper, but it would be worth pointing this out in their discussion of their algorithm.\n\n\n**Weaknesses:**\n- the idea of using mutual info to guide representation learning has been considered previously in various guises. Perhaps the closest version to the current approach is by Lee et al, 2020, which uses mutual information between the initial and final states of a trajectory to guide state compression. The authors' main modification is really to introduce action sequences into the picture, motivated by the idea of using controllability to guide compression, and in this reviewer's view, this contribution seems borderline incremental. Indeed, the improvement over PI-SAC (Table 1) seems to be within statistical error for several of the tasks (cartpole swing-up, reacher easy, finger spin, and ball-in-cup), which makes me wonder to what extent is the improvement attributable to implementational differences instead.\n- mutual information between two random variables depends on their probability distributions. One of the variables in their objective is the action sequence, but this depends, in particular, on the policy learnt (training data are sampled from the replay buffer). This begs the question to what extent are the representations learnt optimised for the policy (and hence the particular task). It may be helpful to freeze the encoder and evaluate the representations on different tasks in the same environment (for instance, by changing the reward function to encourage different behaviour). In this setting, what is controllable about the environment would remain the same, but the optimal policy for the task will have to differ. This would also shed light on how transferable these representations would be to different tasks in the same environment. Moreover, it would be helpful to compare against a ground truth baseline where a LCER agent learns representations directly on the target task and is then made to \"transfer\" to the same task -- the idea is to see how well the agent that transfers representations compares to one that has access to the \"ground truth\" representation.\n- Section 5.4 presenting the evaluation of representation transfer was especially weak. The difference between the source and target task was not made clear, making it difficult to interpret what transfer is taking place between tasks. Moreover, why did the authors only compare against SAC baselines in this section? Clearly SAC was not designed to extract out task irrelevant features, so this seems like a weak baseline to compare against.\n- as mentioned above, to what extent are these representations robust to perturbations in the state dynamics. \n- a limitation of the approach is the introduction of extra hyperparameters $\\beta$ and $K$, which need to be tuned for each task, particularly the weighting hyperparameter $\\beta$.\n- minor typographical errors including spurious capitalizations and places where the authors wrote \"dose\" where they meant \"does\" (e.g. in the middle of the 1st paragraph of the \"Predictive representations in RL\" subsection and in the middle of the captions for Figs 3 and 4)\n", "summary_of_the_paper": "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:\n- 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.\n- LCER is able to outperform alternatives when distracting visual features are introduced into these tasks\n- LCER's representations learned on a source task are able to transfer better to downstream tasks.", "summary_of_the_review": "While I liked the motivation of their approach, my main hesitation for recommending this paper lies both on the strength of their contribution and issues with their empirical results. On the strength of their contribution, to me it seems borderline incremental, so I would be curious to hear the opinion of the other reviewers. Moreover, the results seem marginally better than PI-SAC, which makes me wonder to what extent the better results can the better results be conclusively attributed to the better objective and not to implementational factors, like the choice of $h()$ function architecture or a better approximation of the mutual information objective.\n\nOn the empirical side, their explanation of transfer in Section 5.4 could have been clearer, and they could have compared against stronger baselines. I would also have appreciated insights into how much the representations are fine-tuned to the policy being used. Finally, it would strengthen the paper if the authors could provide insight into how well do these representations transfer to different objectives as well as to perturbations of dynamics in the same environment. The main motivation behind the authors' work is to have an agent learn state representations that capture controllable elements of the environment only. Evaluating transfer on different goals in the same environment would be a test of this, since an agent that truly captured the controllable features should not be adversely affected by a change in goal or reward function.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "CGkod7dFIz4", "reviewer": "Reviewer_vvtm", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The high-level idea seems sound and practical and the implementation strategy builds on popular ideas in the literature and seems reasonable. \n\nDoes not cite DIYAN (Eysenbach 2018) https://arxiv.org/abs/1802.06070 which also uses mutual info between latent state projection Z and actions A, though not for representation learning.  \n\nIn equation 6, the information in the sequence is given as the sum of information at each horizon value k which is probably a sensible approximation, but I am guessing not technically correct as the states at each horizon value are not mutually independent. \n\nPage 5 middle paragraph starting “where B is the batch size” contains “is a latent defined” which I am guessing should be “is a latent variable defined”. \n\nWhy is JSD more stable than MI? Is it possible to give an intuition? \n\nI didn’t find the section on augmentation directly relevant to the paper’s point and I think it could be omitted from the main text if space was tight. It might allow the paper to expand section 4.2 a bit more to explain the recursion better.  \n\nAlternative: LACER: Learning Agent Controllable Elements Representations ", "summary_of_the_paper": "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. ", "summary_of_the_review": "I think the idea of explicitly trading off the action focused vs. state focused loss to generate representations from unsupervised experience is new but closely related to precedents. It speeds convergence on easier PlaNet tasks and seems empirically to make a difference in achievable expected reward for the harder tasks in the Distracting Control Suite.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces an objective for representation learning that captures "controllable elements" in the environment (i.e., things that are affected by the agent's actions). In their reviews and discussion, the reviewers agreed this idea was intuitive, well-motivated, and the paper well written. However, multiple reviewers raised concerns about the evaluation and the extent to which LCER is truly an improvement over PI-SAC. Although many of the reviewer's concerns were addressed in the rebuttal period, at the end of the discussion they were still unconvinced or confused about how much LCER really helps over PI-SAC. Based on this, my assessment is that this paper is a promising piece of work, and that with some more controlled comparisons (see suggestion below) it would be a useful contribution to the literature. However, given that the claims are not fully supported as it currently stands, I recommend rejection.

Specific suggestion to improve the paper: based on reading the paper and the discussion, it seems to me (as per the authors' own statement in response to Reviewer uWv6) that the most valid/controlled comparison between LCER and PI-SAC is in Figure 4, where LCER w/ $\beta=0.1$ "can be seen as a variant of PI-SAC with the same embedding choices as LCER" (author's words). However, when taking into account the error bars of the training curves, other values of $\beta$ are only clearly better than $\beta=0.1$ in 1/3 environments (D.walker-walk). This does not make for a particularly convincing result that LCER is better than PI-SAC. To improve the paper, I'd encourage the authors to run further well-controlled comparisons such as this in a larger number of environments. If they can show via such controlled comparisons that LCER is generally better than PI-SAC (i.e. LCER w/ $\beta=0.1$) then that would be a much more compelling demonstration of LCER's superiority.
Decision: Reject