Learning Temporally-Consistent Representations for Data-Efficient Reinforcement Learning
Authors: Trevor McInroe
Lukas Schäfer
Stefano V Albrecht
Venue: ICLR 2022 Submitted (2022)
Keywords: Reinforcement learning
representation learning
continuous control

Abstract: Deep reinforcement learning (RL) agents that exist in high-dimensional state spaces, such as those composed of images, have interconnected learning burdens. Agents must learn an action-selection policy that completes their given task, which requires them to learn a representation of the state space that discerns between useful and useless information. The reward function is the only supervised feedback that RL agents receive, which causes a representation learning bottleneck that can manifest in poor sample efficiency. We present $k$-Step Latent (KSL), a new representation learning method that enforces temporal consistency of representations via a self-supervised auxiliary task wherein agents learn to recurrently predict action-conditioned representations of the state space. The state encoder learned by KSL produces low-dimensional representations that make optimization of the RL task more sample efficient. Altogether, KSL produces state-of-the-art results in both data efficiency and asymptotic performance in the popular PlaNet benchmark suite. Our analyses show that KSL produces encoders that generalize better to new tasks unseen during training, and its representations are more strongly tied to reward, are more invariant to perturbations in the state space, and move more smoothly through the temporal axis of the RL problem than other methods such as DrQ, RAD, CURL, and SAC-AE.

Under review as a conference paper at ICLR 2022
LEARNING TEMPORALLY-CONSISTEN REPRESEN- TATIONS FOR DATA-EFFICIENT REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Deep reinforcement learning (RL) agents that exist in high-dimensional state spaces, such as those composed of images, have interconnected learning burdens. Agents must learn an action-selection policy that completes their given task, which requires them to learn representattin of the state space that discerns between use- ful and useless information. The reward function is the only supervised feedback that RL agents receive, which causes representation learning bottleneck that manifest in poor sample efficiency. We present k-Step Latent (KSL), a new rep- resentation learning method that enforces temporal consistency of representations via a self-supervised auxiliary task wherein agents learn to recurrently predict action-conditioned representations of the state space. The state encoder learned by KSL produces low-dimensional representations that make optimization of the RL task more sample efficient. Altogether, KSL produces state-of-the oo-tte--tttttttt results in both data efficiency and asymptotic performance in the popular PlaNet bench- mark suite. Our analyses show that KSL, produces encoders that generalize better to new tasks unseen during training. and its representations
more strongly tied to reward, are more invariant to perturbations the state space, and move more smoothly through the temporal axis of the RL problem than other methods such as DrQ, RAD, CURL, and SAC-AE.
INTRODUCTION
Learning sample-efficient continuous-control directly from pixels is challenging task for reinforce- ment learning (RL) agents. Despite recent advances, modern approaches require hundreds thou- sands, if not millions, of agent-environment interactions to attain acceptable performance (Kalash- nikov et al. 2018 Jaderberg et al., 2019). real-world settings, this magnitude of data collection may take prohibitive amount of time. Therefore, advances in learning efficiency for continuous- control in pixel-based state spaces would significantly accelerate the success of real-world RL. ap- plications such industrial robotics and autonomous vehicles (Dulac-Arnold al. 2021
The standard RL loop ties policy optimization and representation learning together with single su- pervisory signal. This dual-objective, single-signal learning scheme results in representation learn- ing bottleneck that ultimately produces agents that understand their environment slowly throughout training (Shelhamer et 2017), thereby reducing performance in the early stages of learning. Ap- proaches that have seen the most success in alleviating the bottleneck rely on methods that explicitly encourage the distillation of useful information from the state space.
Recently, data augmentation (Chen et al. 2020a b) Hénaff et al. 2020) has been studied in the
context of RI as way to improve sample efficiency (Yarats et al. 2021b Laskin et al., 2020b). Other approaches to improve sample efficiency in RL add auxiliary tasks to the learning process. For example, image reconstruction (Yarats et al.] 2020), latent variable models (Lee et al. 2020a Gelada et al. 2019), variational objectives (Lee al.) | (020b), and contrastive learning (Laskin al. 2020a Stooke et al. 2021). While the aforementioned approaches improve performance over base RL algorithms, they do not fully leverage the long-term temporal connection between actions, future states, and rewards. For example, some approaches do not interact with the temporal axis the RL problem's underfying Markov decision process (MDP) Yarrts et al. 2020 2021b: Laskin et al.

---

Under review as a conference paper at ICLR 2022
2020b), or are only concerned with next-step dynamics (Laskin et al. 2020a Stooke et al. 2021 Gelada et al. 2019). As we will show with our work, igmoring the long-ter temporal axis of the MDP withholds information that can help agents to understand consequences of their environment interactions more quickly. Approaches that do attempt to capture longer-term information either require significant amount of random exploration and pre-training (Lee al., 2020a b), which counteracts the goal of minimizing agent-envirooment interactions, or condition their predictions of the future on the full-dimensional state space (Lee et al. 2020b), thereby not encouraging latent representations themselves to be predictive of latent representations multiple steps into the future.
In this paper, we revisit the problem of sample-efficiet continuous-comtrol in image- baage- state spaces. We introduce k-Step Latent (KSL a representation learning module that aids RL agents with continuous control on pixels. In contrast to aforementioned prior methods, KSL's auxiliary task directly exploits the RI. problem's underlying MDP, requires no pre-training, and encourages latent representations to be predictive of latent representations that are multiple steps into the future. KSI. agents learn to create temporally-cohere such that latent representations, when conditioned upon a series of actions, are predictive of latent representations that are multiple steps into the future. KSL's auxiliary task propels agents to new state-of-the thrt results on both the data efficiency (100k steps) and asymptotic performance (500k steps) checkpoints in the popular PlaNet benchmark suite (Hafner et al. 2019).
representations by projecting states into the latent space
This paper makes the following contributions. First, KSL. augments the learning process of the popular Soft Actor-Critic (Haarnoja et al.) 2018a) algorithm produce in the PlaNet benchmark suite. KSL significanth, outperforms current methods such as DrQ (Yarats) al. 2021b) and RAD (Laskin et al. (020b). Across all six tasks in PlaNet, KSL produces an average improvement over DrQ of 46.9% and 14.5% at the 100k and 500k steps mark, respectively, and an average improvement over RAD of 58.1% and 21.6% at the 100k and 500k steps mark, respectively. Second, we analyze the learned encoders and latent projections that are produced by several methods. Our analyses show that KSI. creates representations that are more strongly tied to reward, are less sensitive to perturbation in the state space, and move more smoothly over time than other baseline methods. Also, our analyses show that KSL creates encoders that generalize to unseen tasks better than other baseline methods.
state-of-the-arr results
BACKGROUND
We consider the usual formulation of RL agents that act within a Markov decision process (MDP), defined by the tuple (S, A, P. 7). where S is the state space, A is the action space, P(* (st, at) the transition probability distribution, R: : x A + is reward function that maps states and actions scalar feedback signal, and y E [0, 1) is a discount factor. The agent's 'ss goal is to learn an action-selection policy x (a(s) that maximizes the sum of discounted rewards over the time horizon T of the given task: arg max PP Σ R (st.
Specifically, we examine learning agents in state spaces comprised of images o E O. Image-based states are nigh-dimensional and may give an incomplete description of the environment. For exam- shows occluded objects. As such, it is common to stack series of n consecutive images to help alleviate the problems associated with incomplete information: 8t = {or.... 0 (Mnih et al. 2015).
ple, incomplete information may arise when the camera's 's point-of-viee
2.1
SOFT ACTOR-CRITIC
Soft Actor-Critic (SAC) (Haarnoja et al. 2018a b) a current state-of-the-ar off-policy, model-free RL algorithm for continuous control. SAC learns state-action value-function critic Qø stochastic actor T&, and a temperature a that weighs between maximizing reward and entropy in an augmented RL objective:
Es E,, R(s,, a.) (+((T(-(ss))..
SAC's critic is learned by minimizing the squared Bellman error over trajectorics
(St, 0e, Te, 81+1) sampled from replay memory D:
Lerati (0)))()))))n(((((sseeee - (r: + ?1))
(I)
Code available at: httppe///tthhb..ccroooooeeeeeeer repo/kk
2

---

Under review as a conference paper at ICLR 2022
8]
%
12
11
ll
8j-1
V/n
Figure 1: Depiction of KSL's architecture. Online and momentum paths are highlighted in orange and green, respectively. Recurrent operation shown with red arrow. EMA updates shown with dashed arrows. Gradient flow shown with dotted, purple arrows. Stop gradient (sg) shown with //.
where the temporal-differencee target, y. is computed by sampling an action from the current policy:
En' QQ d') - a log ग ( [s(41)]-
(2)
The target critic Qa updated as an exponential moving average (EMA): 0 + ÇŌ + (1 ()0 where controls the "speed" which 0 tracks 9 (c.g.. He et al. (2020)) SAC's actor parameterizes multivariate Gaussian (H, or) where is a vector means and a is the diagonal of the covariance matrix. The actor is updated via minimizing:
Lactor (
Ea Q0(ss, a) - log Ta(a[s=)]
and a is simply learned against a target entropy value.
k-STEP LATENT
3.1 MOTIVATION
We design KSL to address the representation learning bottleneck of base RL algorithms by providing an auxiliary task that encourages representational consistency of states that are nearby in time. States that are nearby in time are likely share high levels of mutual information (Mobahi et al.] 2209)) which implies that the states' representations should as well. KSL. contains a recurrent mechanism that encourages representations of states, when conditioned upon actions, be predictive of states that are multiple steps into the future. Intuitively, we hypothesize that agents who understand the long-term consequences of their actions would perform better than agents who do not.
3.2 TRAINING DETAILS
A visual depiction of KSL can be seen in Figure and pseudocode can be found in Appendix B At each learning step, we uniformly sample M trajectories of length k, from the replay memory, where 8 can be any state in the episode. For each state in each trajectory, we apply a random translation augmentation (Laskin et al. (2220b). Using augmentations within KSL encourages representations be invariant to perturbations over the temporal axis. KSL contains an online and momentum pathway, which are highlighted in Fig- of trajectory") sy is fed through an online encoder 2, = (8,)) and 8j+1 1 is fed through a momentum encoder Zm &m (8)4))1), Next, 2, is concatenated with dj and fed through learned transition model = = =[[==la;)), This transition model enables KSL to learn temporal projections from latent representations of state sy and applied action ay representation of the following state 8j+1 the form of 2/2- Then, 2/ and 2m are fed through their respective projection modules qo = V/ (=(2)), Dm = Vi (2..). Such modules after an encoder have been shown to improve the learned representations of the encoder (Chen et al. 2020a). Finally, is fed through a prediction head yo P.g.).
= {(81, o1:..., (1-1,18)))
ure I orange and green, respectively. At step
Note that representations through both pathways computed for every step and trajectory. We omit batch notation for steps and trajectories in the paragraph for clarity.
3

---

Under review as a conference paper at ICLR 2022
KSL's loss for step j is computed the distance between the lz normalized versions of the final two vectors from both paths: [1/b//1lyoll) - (ym/ ll l). This operation is performed in a loop over all & steps per each trajectory in with one notable difference at each step after j 1. Instead of encoding every along the online path, we recursively feed z from the previous step j into the learned transition model T. This recurrent operation is shown with the red arrow Therefore, Po is only used to compute direct representation of the first state within each trajectory. The minimization of KSL's loss requires &o to create latent projections such that multi-step consequences can be inferred from any arbitrary starting state using the transition model T. We formulate the loss function as an average over the computed distance for every state and trajectory in given batch:
in Figure [
(i.)) No ((.5)
(i.))
Vm
LKSL
MR
II
ym
(F.)
with 3/5-0) denoting the final vector of pathway p for trajectory at step j.
We implement stop gradient to disallow backpropagation from Equation through the momentum pathway. Instead, we update the momentum encoder and momentum projection module as an EMA of their online counterparts. We use the state encoding of the online encoder & as input into the critic Qs and actor Ts- While we allow the gradient of Qa to update weights within $ we implement stop gradient that disallows signal from the actor's computation graph (Equation and and Equation 3) to propagate into D., Gradient flow depicted in Figure as purple, dotted lines. 2m an EMÄ of tb and Qg is EMA Qs. it is natural for T... to transform states for input into Q8- See Appendix(Ffite an expanded description of how KSL fits within SAC. We also provide ablations across architecture choices and values of &o in Appendix F and Appendix G
We note that KSL is not restricted to SAC and can be applied alongside any RL algorithm for representation learning in tasks with continuous discrete action spaces.
3.3 OPTIMIZATION STRATEGY
The weights of 4, are affected by two task-specific losses (Equation I and Equation Usually, learning algorithms that attempt to optimize parameters across multiple losses use a weighted linear combination of each loss. However, doing causes the sharing of inductive biases across each task, which may be an issue if the two tasks compete (Yu et al., 20200 /ener & Koltun| 2018). The representation learning bottleneck of the RL signal may cause disruptions in b, during optimization of KSL's loss. Additionally, this problem may be exacerbated with optimizers that accumulate gradient statistics (Jean et al., 2018), such Adam (Kingma & Ba, 2015). Much of the current research in multi-objective learning has focused on methods that actively learn to balance between multiple losses in the context of single optimizer (Kendall et al. 2018 Groenendijk et al., 2021). Instead, we explicitly separate the optimization of the RL and KSL tasks by sequestering the weights and losses of each task into separate optimizers. This division is not meant to solve the conflicting gradient problem. The division avoids conflation of the Adam optimizer's momentum estimates between the two tasks. As a consequence of our multi-optimizer choice, & is granted a higher diversity in gradient- update directions. In our experiments, we performance boost due to this separation. See Appendix [ for more details.
EXPERIMENTAL EVALUATION
We investigate several facets of KSL and other baseline methods. Most importantly, does KSL result in agents that converge to asymptotic performance in fewer agent-environment interactions than other current methods? We examine KSL's learning progress in environments from the DeepMind Control Suite (DMControl) (Tassa et al. 2018 2020 (Section 4. against several relevant baselines (Section 4.2) and show these results in Section 4.3 Also, we examine the characteristics of encoders and latent representations that KSL and other methods produce. We inspect the relationship between representations and reward, the robustness of encoders to perturbations in the state space, the ability of encoders to generalize to other tasks in the DMControl suite, and the temporal coherence of representations.
4

---

Under review as a conference paper at ICLR 2022
4.1 DMCONTROL
We examine the learning speed of agents by measuring performance within the popular PlaNet benchmark suite (Hafner al., 2019). This suite contains environments from DMControl (Tasaa et al. (018)) and has been widely used to measure the sample efficiency of RL agents on image- based states. DMControl presents a wide variety of contnuouss-ontrol tasks, including rigid- body control, control of loose appendages with generated momentum, and bipedal traversal. There are six tasks that range from one-dimensional to six-dimensional action spaces and agents in DMControl are exposed to both dense and sparse reward signals. See IxpendixJfor more details.
Two checkpoints of interest have emerged in recent literature Laskin et al. (2020a)): the 100k steps mark and the 500k steps mark. The former is considered to be the data efficiency mark and the latter the asymptotic performance mark. The term "steps" refers to the number of environment transitions and not the number of learning updates. Each task has an episode length of Ik steps and has a specific number of action repeats. The agent's 's chosen action is repeated several times, each of which results in an environment transition. For example, Cheetah, Run has an action repeat of four, which results 250 actions selections per episode.
4.2 BASELINES
We compare KSL against six baseline methods in the data-efficiency experiment: RAD (Lakin et al., 2020b), DrQ Yarrtt et al., 2021b), CURL (Laskin et al., 2020a), SAC-AE (Yarats et al. 2020 SAC State, and SAC Pixel. We specifically choose the first four methods because they were (i state-of-the-art in the PlaNet benchmark suite at one point in time and (ii) all use the exact same encoder and SAC architecture as well as experimental design in PlaNet. Likewise, KSL uses the same encoder and SAC architecture as these four methods to ensure fair comparison. SAC State and SAC Pixel are commonly chosen as experimental control agents. SAC State learns on proprioceptive versions of the state space (e.g., vectors that contain "ground truth" information such as joint angles and velocity). SAC Pixel learns on the pixel version of the state space but with no image augmentation and no representation learning routine. We produce the results for baseline methods using code provided by the original authors. See Appendix D for a comparison between KSL and other methods that do not meet criteria (i) and (ii).
4.3 DATA EFFICIENCY EXPERIMENT
During training, all agents perform random actions for the first Ik steps to collect experience tuples for the replay buffer. Each agent type performs one gradient update per action selection by using randomly sampled experience tuples. We perform an agent evaluation every 20k steps. Agent evaluation uses current snapshot of the learning agent and deterministically samples actions. The agent's 's final score for each evaluation is the mean of total rewards across 10 episodes. For all methods, we use batch size of 128. A full table of hyperparameters can be found in Appendix AA] We also provide a wall clock time comparison between KSL, RAD, DrQ, CURL, and SAC-AE in Appendix C
Figure 2 presents the evaluation curves for KSL, RAD, DrQ. CURL, and SAC-AE over 10 seeds. Table highlights the 100k and 500k steps mark as averages + one standard deviation, with best mean evaluation returns between mmage-based methods bolded. We use Welch's t-test to probe for statistical significance between the highest average result and all other results. We highlight that KSL achieves the best average result in all six tasks at both checkpoints of interest and especially outperforms all other methods at the data efficiency checkpoint.
4.4 ANALYZING LATENT REPRESENTATIONS AND LEARNED ENCODERS
We analyze the latent representations and the learned encoders of all image-based methods from Section 4.2 We collect 10 episodes of state-reward pairs in the Walker, Walk task with a fully- trained agent, which we refer to as the evaluation states and rewards. We choose to focus on Walker, Walk as it is one of the more difficult tasks in PlaNet, as judged by the dimension of its action space, the number of action selections per episode, and the difficulty of learning bipedal movement. For the following analyses, we train encoders for each method in Walker, Walk for 100k steps across five seeds, saving the encoders' weights every 5k steps. Specifically, we use the encoders that produce
5

---

Under review as a conference paper at ICLR 2022
$C:44
50.11
80
Ran
Walber,
1209
036
Figure 2: Episodic evaluation returns for agents trained in DMControl. Mean as bold line and + one standard deviation as shaded area.
Table I: Episodic evaluation returns (mean + one standard deviation) for PlaNet benchmark. High-
est mean results per task shown in bold. significant at p 0.05, ** significant p 0.01.
500k Steps
KSL
RAD
DrQ
CURL
SAC-AE
SAC State
SAC Pixel
Finger, Spin
976
970 + 28
954 + 47
748 + 108
699 169
927
192
166
Cartpole, Swingup
+:
207
189
1
419
Reacher, Easy
963 28
817 + 261
938 + 63
760 + 106
643 136
975 +
143
30
Cheetah, Run
802 + 30
406 660
459 + 69
166 + 24
484 +: 79
772 60
197
15
Walker, Walk
953 8
858 1: 90
914 1:
689 1: 90
865 1: 59
964 1
42
12
Ball in Cup. Catch
+
106
312
100k Steps
Linger, Spin
#5
781
1: 101
735 1 146
699 105
677 147
672 76
224
101
Carppole, Swingup
841
33**
620
114
651 + 172
274 55
278 39
812 45
200
72
Reacher, Fasy
751
137**
429
175
513
139
473 + 184
334 55
919
123
136
15
Cheetah, Run
54**
1:
278
228
130
Walker, Walk
730
133**
378
164
442 + 208
337 62
364 54
604
317
127
24
Ball in Cup, Catch
945
+ 12*#
744
+ 172
683 +: 131
462 4: 324
465 128
957
26
97 27
representations for SAC's critic Qe and actor Tg. which is 4, for KSL. this section, we refer to these encoders the evaluation encoders.
The reward function the only supervised signal in RL, so follows that learned representations of the state space should relate to reward. We first observe this relationship qualitatively. Figure shows the first two components of PCA projections of latent representations produced by KSL (top row) and DrQ (bottom row). Here, we focus on the early stages of training to observe emergent behavior in learned representations. The left, middle, and right columns of plots use encoders after 5k, 10k, and 15k environment steps, respectively. Point colors denote the reward. We highlight that KSL's plots show a clear separation between low-reward and high-reward states more quickly than DrQ's plots. By 10k steps. KSL's encoder has grouped the highest-reward states together and by 15k steps, high-reward and low-reward states seem to be linearly separable within the PCA projections. In contrast, DrQ's projections show little sign of reward-based organization by 15k steps. All other baseline methods show the same pattern DrQ. See Appendix [ for PCA plots of other baseline methods. Despite the concrete connection between states and rewards in the context of RL, KSL is the only method that we observe discovering this relationship in the early stages of training.
Organization Around Reward:
We also quantify the relationship between learned representations and reward by measuring the mean squared error (MSE) of a linear regression that fitted to estimate reward from latent projections. Figure 4 displays the mean MSE for all evaluation encoders all evaluation states. We note

---

Under review as a conference paper at ICLR 2022
FCA, KS..
XSL
21100
460
32004
5200
4609
-300
08
335
034
Figure 3: PCA projections of latent representations learned by KSL (top row) and DrQ (bottom row) at 5k (left column), 10k (middle column), and 15k (right column) environment steps. Colors indicate the reward received in the original state.
USE of Lineor Prod coor Thre Traniin
VNE Lineer hrdlow Thrresheet Trapann
M
0.-
44-
30%
Figure 4: MSE of linear regression fitted to estimate reward from latent representations. The right plot re-scales the y-axis to show the difference between Mon-SAC-Pixel methods. Lower is better.
that KSI. consistently produces representations that are better linear predictors of reward while also displaying lower variance from one environment step checkpoint to the next than other baseline methods. Also, we note that SAC Pixel produces results that are significantly worse than other baseline methods. This suggests that the RL signal, alone, is not enough to produce representations that are coherently tied to reward.
We measure the invariance of encoders to perturbations the state space. Specifically. we measure the & distance between latent representations of several augmented ver- sions of states. We produce three versions of each evaluation state per augmentation using random translation, random erasing, and Gaussian noise (O00-115) Then, we project representations with all evaluation encoders for all augmented states. Figure 5 displays the average & distance for all evaluation encoders. We note that y-axes are scaled highlight the curves of baseline methods, which results in SAC Pixel being off-screen. See Appendix M for un-scaled y-axes. These results show the importance of combining image augmentation and auxiliary tasks. Methods that use nei- ther (SAC Pixel) show high sensitivity to all three augmentations. Methods that only use auxiliary tasks (SAC-AE) show a high sensitivity to geometric translations. Methods that use only augmenta- tion (RAD and DrQ) sometimes show instability to the visual perturbance of erasing and Gaussian noise. Methods that use both augmentation and auxiliary tasks (CURL. and KSI.) show stability across all three augmentations. We highlight that KSL's encoders generally the most consistent. For example, the change in sensitivity from one evaluation checkpoint to the next is generally much lower for KSL than other baseline methods. See Appendix K for more details.
Perturbation Invariance:
We test the ability of learned encoders to generalize their projections from one task to another. To this end, we first train each agent in each task for 100k steps five times and save the encoders' weights. Then, we freeze each encoder and train a base SAC agent for 100k steps in every task other than the original training task. Table presents the mean + one
Generalization of Encoders:
7

---

Under review as a conference paper at ICLR 2022
Tendiaten
Traske twatoncs Shandon Oralng
Turdelies Ewarannc Gazadan Neaed
Figure 5: l2 distance between latent representations of augmented states. From left-to-right, random translate, random erase, Gaussian noise.
Table 2: Mean + one standard deviation of encoder generalization experiment. Best results bolded. significant at 0.05, ** significant at p = 0.01.
100k Steps
KSL
RAD
DrQ
CURL
SAC-AE
Finger, Spin
705 179
689 T 128
636 = 210
722 + 108
693 = 127
Cartpole, Swingup
325 + 64
258 + 62
255 39
309 + 62
240 = 29
Reacher, Easy
355 189**
227 115
206 150
163 + 107
198 139
Cheetah, Run
251 +
232 232-:43
21 60
235 +M
228
Walker, Walk
373 + 85*
289
283 95
320 +
327 = 66
Ball in Cup. Catch
566 + 219*
329 + 216
320 205
440 + 182
297 = 200
standard deviation of evaluation returns at 100k steps of training. We highlight that KSL's encoders outperform other methods in five of the six tasks. See Appendix Hor full evaluation curves.
Temporal Coherence: Control of robotic components, such as those similar to the objects in DM- Control, involves important elements that move smoothly over time, such as joint angles. Ideally, our learned representations should capture this important information from the state space. Therefore, our representations should also move slowly and smoothly across time. We quantify this character- istic by measuring the lz distance between the latent representations of consecutive state pairs in evaluation set. Figure 6 displays the mean across all evaluation encoders. We highlight that KSL learns representations that move more slowly and consistently than the representations learned by other methods.
RELATED WORK
REPRESENTATION LEARNING AND REINFORCEMENT LEARNING
KSL builds off of the success of Bootstrap Your Own Latent (BYOL) Grill et al. 2020) by adapting its architecture to work over the temporal axis of the RL problem's underlying MDP. BYOL is self-supervised representation learning method for images that does not require negative examples to avoid collapsed representations.
Adding auxiliary tasks base RL algorithms is a common method encourage the learning of useful representations. Yarats et al. (2020) introduce SAC-AE, which adds an auxiliary image- reconstruction task with deterministic auto-encoder (Ghosh et al., 2020) and show that gradients from the actor harm the learning process of encoders, Laskin et al. (2020a) introduce CURL, which augments the RL loop with a projection head akin contrastive predictive coding (van den Oord et al. 2018) with a momentum encoding procedure (e.g., Chen et al. (202cc)). Alternative tasks have been explored such as aligning the latent- and state-space with bisimulation metrics (Zhang et al. 2021). maximizing mutual lemporal-informatioo via learning hidden activations (Anad ctt al. 2019 using contrastive learning for task generalization (Agarwal et al., 2021). and using
latent dynamics models (Lee al. 2020a Hafner et al. 2020 2021 Gelada al. 2019 Guo al.
2020).
Other approaches have relied using tricks from computer vision, such image augmentation.
RAD (Laskin et al. (220b) explores how simply augmenting incoming images can help SAC im- prove without any other alterations. DrQ (Yarats et al. 2021b), published concurrently with RAD,
8

---

Under review as a conference paper at ICLR 2022
Dezance Ltber
Eewbwrred Stees (427)
Btua (001)
Figure 6: l2 distance between consecutive latent states that are produced by the evaluation encoders.
extends RAD to include regularization via averaging Q-function values over several image augmen- tations during training.
KSL is similar to SPR (Schwarzer et al. 2021 in that both use architecture similar to BYOL. However, there are several key differences. Unlike SPR, KSL explicitly separates the weight-update process of the RL task and the representation learning task. This separation results in independent tracking of accumulated gradient information within optimizers between tasks. Due to the overlap in parameters between tasks, this separation leads to differences in weight updates. We show in Appendix|] [ that the update direction of the KSL and RL task often conflict throughout training. A high diversity of gradient directions across multiple tasks may be beneficial in the context of non- i.id data streams (Aljundi et al., 2019), which is a common scenario in RL. Also, KSL and SPR have significant architectural differences. For one, SPR was designed for discrete-action settings and its architecture relies on one-hot encoded actions. In contrast, KSL's architecture is general enough to be applied in both discrete- and continuous-action domains. Additionally, KSL's transition module operates vectors, and therefore it contains convolutional layers. Unlike SPR, this allows KSL to fully share the entire spatial learning process between both the RL and representation learning task. Also, SPR shares more layers than just encoder between the representation learning and RL task (see note in Figure in [Schwarzer et al.] (202T)) We refer to this choice as "knowledge sharing". We find that both knowledge sharing (Appendix and convolutional transition model (Appendix Gi are detrimental for KSL.
5.2 ANALYSIS OF LEARNED REPRESENTATIONS
Most of the focus in RL research is on developing algorithms that provide state-of-the-art perfor- mance. Comparatively, less attention has been spent on analyzing the learned representations them- selves. Laskin et al. (2020b) qualitatively analyze spatial attention maps (Zagoruyko & Komodakis, 2017) of the learned encoder various image augmentations. Zhang et al. (2021) provide plots of L-SNE-projectedd representations to show organization around learned state values. Hafner et al. 2019), and Yarats et al.) (2020) observe the relationship between learned representations and the proprioceptive versions of states. We refer the interested reader to the survey work of Bengio et al.) 2013), which provides a thorough overview of hypotheses around the question, "what makes for good representation, in general?".
6 CONCLUSION
In this work, we addressed the problem of improving the sample efficiency of reinforcement learning (RL) agents for continuous control in image-based state spaces. To this end, we introduced k-Step Latent (KSL), an auxiliary learning task for RL agents. KSL leverages advancements in represen- tation learning and data augmentation to produce state-of-the-art results in the PlaNet benchmark suite. We also analyzed the encoders and latent representations produced by KSL. and other baseline methods. These analyses showed that KSL produces encoders that are more robust to perturbations in the state space and generalize to other tasks better than other baseline methods. Also, the analyses showed that KSL's representations are more strongly tied to reward and move more steadily across the temporal axis of the RL problem's 's underlying MDP. Despite the sample-efficiengg gains of KSL, it requires increased wall clock time compared to previous methods. It may be possible to reduce this burden via parallel processing by reducing the size of KSL's neural networks.

---

Under review as a conference paper at ICLR 2022
7 REPRODUCIBILITY STATEMENT
Please see the footnote on page 2 for a link to an anonymized code repository. Also, see Ap- pendix A for implementation details such as hyperparameters and architecture choices, Finally, see Appendix for KSL's pseudocode.
REFERENCES
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In International Conference Learning Represenlutions, 2021.
Rahaf Aljundi, Min Lin, Baptiste Goujaud, and Yoshua Bengio. Gradient based sample selection for online continual learning. In 33rd Conference on Neural Information Processing Systems (NeurIPS), 2019.
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Mare-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems (NeurIPS), pp. 7766-8779, 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arXiv preprint arXiv: 17X7.06607.064550 06450, 2016.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE Transactions Pattern Analysis and Machine Intelligence, 1788))9988833 2013.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. Proceedings of the 37th International Conference on Machine Learning (ICML), 2020a.
Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey Hinton. Big self- supervised models are strong cuperuipd learners. In 34th Conference on Neural Information Processing Systems (NeurIPS), 2020b.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.04297, 2020c.
Gabriel Dulac-Amold, Nir Levine, Daniel J. Mankowitz, Jerry Li. Cosmin Paduraru, Sven Gowal, and Todd Hester. Challenges of real-world reinforcement learning:definitions, benchmarks & analysis. Machine Learning, 10:2419- 2468, 2021.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G. Bellemare. DeepMDP: Learning continuous latent space models for representation learning. In Proceedings of the 36th International Conference on Machine Learning (ICML), 2019.
Partha Ghosh, Mehdi S. M. Sajjadi, Antonio Vergari, Michael Black, and Bernhard Scholkopf. From variational deterministic autoencoders. In International Conference Learning Representa- tions (ICLR), 2020.
Jean- Bean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Ghesh- laghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent: A new approach to self-supervised learning. In 34th Conference on Neural Information Processing Systems (NeurIPS). 2020.
Rick Groenendijk. Sezer Karaoglu, Theo Gevers, and Thomas Mensink. Multi-loss weighting with coefficient of variations. In IEFE Winter Conference on Applications of Computer Vision, 2021.
Daniel Guo. Bernardo Avila Pires, Bilal Piot, Jean- nan-bbstien Grill, Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar. Bootstrap latent-predictes representations for multitask reinforce- ment learning. In International Conference Machine Learning (ICML), 2020.
10

---

Under review as a conference paper at ICLR 2022
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference Machine Learning (ICML). volume 80, pp. 1861-1870, 2018a.
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. arXiv preprini arXiv: /8/ /1/50000......... 2018b.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Proceedings of the 36th Inter- national Conference on Machine Learning (ICML), pp. 2555-2565, 2019.
Danijar Hafner, Timothy Lillierap, Jimmy Ba, and Mohammad Norouzi. Dream control: Learn- ing behaviors by latent imagination. In International Conference on Learning Representations (ICLR), 2020.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In Internation Conference on Learning Representations (ICLR), 2021.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 9726-9735, 2020.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In Proceedings of The Thirty-Second AAAI Conference on Artificial Intelligence (AAAI-18), 2018.
Olivier J. Hénaff, Aravind Srinivas, Jeffrey De Fauw, Ali Razavi, Carl Doersch, S. M. Ali Eslami, and Aaron van den Oord. Data-efficient image recognition with contrastive predictive coding, In Proceedings of the 37th International Conference on Machine Learning (ICML), 2020.
Riashat Islam, Peter Henderson, Maziar Gomrokchi, and Doina Precup. Reproducibility of bench- marked deep reinforcement learning tasks for continuous control. In Proceedings of the ICML 2017 workshop on Reproducibility in Machine Learning (RML), 2017.
Max Jaderberg, Wojeiech M. Czarnecki, lain Dunning, Luke Marris, Guy Lever, A. Castañeda, Charlie Beattie, Neil C. Rabinowitz, Ari S. Morcos, Avraham Ruderman, Nicolas Sonnerat, Tim Green, Louise Deason, Joel Z. Leibo, D. Silver, D. Hassabis, K. Kavukcuoglu, and T. Graepel. Human-level performance in 3d multiplayer games with population -based reinforcement learning- Science, 364:859 865, 2019.
Sébastien Jean, Orhan Firat, and Melvin Johnson. Adaptive scheduling for multi-task learning. In Continual Learning Workshop at the 32nd Conference on Neural Information Processing Systems (NeurIPS), 2018
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, and Sergey Levine, QT-Opt: Scalable deep reinforcement learning for rision-based robotic manipulation. In 2nd Conference on Robot Learning (CoRL), 2018.
Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), 2018
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd Interna- tional Conference on Learning Representations (ICLR), 2015.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representa- tions for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learming (ICML), volume 119, pp. (650, 2020a.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto. Pieter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. In 34th Conference on Neural Information Processing Systems (NeurIPS), volume 33, pp. 19884-19895. 2020b.

---

Under review as a conference paper at ICLR 2022
Alex X. Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pp. 741-752, 2020a.
Kuang-Huee Huei Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama. Predictive information accelerates learning in rl. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pp. 11 1180--11911, 2020b.
Shikun Liu, Edward Johns, and Andrew J. Davison. End-to-end multi-task learning with attention. In /EEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 1871-1880, 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei Rusu, Joel Veness, Marc Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518:529-533 533, 2015.
Hossein Mobahi, Ronan Collobert, and Jason Weston. Deep learning from temporal coherence in video. In Proceedings of the 26th International Conference on Machine Learning (ICML), pp- 737-744, 2009.
Max Schwarzer, Ankesh Anand, Rishab Goel, Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self-predictive representations. In International Conference on Learning Representations (ICLR), 2021.
Ozan Sener and Vladlen Koltun. Multi-task learning as multi-objective optimization. In 32nd Conference on Neural Information Processing Systems (NeurIPS), 2018.
Evan Shelhamer, Parsa Mahmoudich, Max Argus, and Trevor Darrell. Loss is its own reward: Self- supervision for reinforcement learning. In International Conference on Learning Representations (ICLR), 2017.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In Proceedings of the 38th International Conference on Machine Learning (ICML), volume 139, pp. 9870-9879, 2021.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy Lillicrap, and Martin Ried- miller. DeepMind control suite. arXiv preprint axXiv:1001.0069. 2018.
Yuval Tassa, Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, and Nicolas Heess. control:: Software and tasks for continuous control. arXiv preprint arXiv:2006. 12:0006..2990 /2833, 2020.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint arXiv: / 007700374.... 2018
Haowen Xu, Hao Zhang, Zhiting Hu, Xiaodan Liang, Ruslan Salakhutdinov, and Eric Xing, Au- toloss: Learning discrete schedules for alternate optimization. In International Conference on Leanring Representations (ICLR), 2019.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Im- proving sample efficiency in model-free reinforeement learning from images. arXiv preprint arXiv: 1910,0174/ 741, 2020.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with proto- typical representations.. In Proceedings of the 38th International Conference on Machine Learning (ICML), volume 139, pp. 1 11920-11931. 2021a.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representa- fions (ICLR), 2021b.
12

---

Under review as a conference paper at ICLR 2022
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. In 34th Conference on Neural Information Processing Systems (NeurIPS), 2020.
Sergey Zagoruyko and Nikos Komodakis. Paying more attention to attention: Improving the per- formance of convolutional neural networks via attention transfer. In International Conference on Learning Representations (ICLR), 2017.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learn- ing invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations (ICLR), 2021.
13

---

Under review as a conference paper at ICLR 2022
Table Hyperparameters used to produce paper's main results.
Hyperparameter
Value
Image padding
4 pixels
Initial steps
1000
Stacked frames
Evaluation episodes Optimizer (81, + (0, Φ. v) Learning rate (0, Φ, V, a)
10
Adam
(0.9, 0.999)
2e 4 Cheetah, Run
1e otherwise
Batch size function EMA C
128
0.01
T EMA Ç
0.05
Target critic update freq
Actor learning freq
Critic learning freq
Auxiliary learning freq dim(z)
50
ς Initial a Target a Replay memory capacity
0.99
0.1
- -|A|
100,000
Actor log stddev bounds
[-10,2]
A IMPLEMENTATION DETAILS
Each of KSL's encoders o., & are the same used in SAC-AE, CURL, DrQ, and RAD. Each have four convolutional layers with 32 3x3 kernels. All layers have stride of one, except for the first, which has a stride of two. The convolutional layers are followed by a single fully-connected layer with 50 hidden units and a layer norm (Ba et al., 2016) operation.
KSL's transition module T begins with two fully-connected layers with 1024 and 512 hidden units, respectively. These layers are followed by layer norm operation and final fully-connected layer that outputs a vector of length 50. Both the momentum and online versions of KSL's projection modules V o. V m have two fully-connected layers with 512 and 50 hidden units, separated by layer norm operation. Finally, KSL's prediction head P a single 50x50 learnable matrix.
All actor and critic networks three-layer dense neural networks with 1024 hidden units in the first two layers. All networks use ReLU nonlinearities except for the final output of p and 4m which uses tanh.
p is used to project states for the critic Qo and the actor тe. pm is used to project states for the target critic Qg. We implement stop gradients that disallow propagation from from n.. computation graph into Φ-.
Hyperparameters and algorithm settings be found in Table 3 0 refers to all of the parameters in both critics, & refers to the actor's ' parameters, and Vi refers to all parameters within KSL. The learning frequency refers to the number of action-selections between each learning update. During training, we use uniform random sampling of the replay memory. Initial steps are used to fill the replay memory and are collected by selecting actions randomly.
B KSL PSEUDOCODE
Algorithm presents pressents Py Torch-like pseudocode for KSL. For clarity, the same mathematical symbols here as the main body of the paper. Also, we include PyTorch-like tensor-indexing for clarity. For example, line 3 shows sf; 1, :, i, :], which means all batches, 1 out of k all color
14

---

Under review as a conference paper at ICLR 2022
channels across the stack of images, all pixels in the width, all pixels in the height. There is one minor difference between the notation in Algorithm| I and the notation in the main body of the paper. To allow for the for loop, 2 is kept 2, in the pseudocode.
Algorithm 1: PyTorrc--liks pseudooode for KSL's loss.
Input Trajectories ((51, 01, -.. Ah-1, 8k) J sampled from replay buffer with
augmented states
loss +-
20 f- Po(s]: 1, :I)
for j in 1...k do
2m t- Pm j 7+1, 1,1,, ;,
[[
Zo
Go < (z(zo)
Im V m (m (2m)) ddetach)
No - P()))
10
loss + loss +
voll
Vo
Mon
Vm
mean()
11 end
12 wwLoopmmzze.zzro.geee gradO
13 loss.backwaad
14 xsl.oppmizer.step(
C WALL CLOCK TIME
To compare wall clock time, average the total number of seconds to complete 15k episodes of training of Cartpole, Swingup across three seeds. Table 4 shows this average as a value that is normalized to SAC Pixel. All metrics were collected AMD Threadripper 1900x @ 3.8GHz, and 128GB of RAM. Codebases for all methods are written in Porch..
machine with 1080Ti GPU, 8-core
D COMPARISON OF KSL TO PI-SAC AND PROTO-RL
Here, we include performance comparison between KSI, PI-SAC (Lee al. 2020b) and Proto- RL (PRL) Yarats et al., 2021a). PI-SAC uses both forwards and backwards encoder to learn a compressed representation of image-based states that maximizes the predictive information be- tween representations and future states while minimizing useless information in the state space. We choose not include this method in Section 4.2 due to significant differences in architecture, hy- perparameters, and experiment settings in the PI-SAC paper and papers of methods presented in For one, PI-SAC's encoder and SAC network architectures are different than those in explain significant differences in performance [Islam et al. 2017] (enderson et al.) 2018). Second, PI-SAC uses a significant amount reported results, while all methods presented in Section 44. do not. For example, in Cheetah, Run, Lee et al. (2020b) use 10k pretraining steps, which results in PI- SAC agents receiving the same number of gradient updates that the agents described in Section 4.2 would only receive by the 40k step mark. Thirdly, the hyperparameters used for tasks (e.g., the action repeat) in the PI-SAC paper are not the same the ones used in previous PlaNet studies. For example, Lee et al. (2020b) use an action repeat of 4 for Cartpole, Swingup, while all methods in Section 4.2 use This results in PI-SAC agents receiving 2x the number of gradient updates for baseline methods.
Section 4.2
Section 4.2 Even small differences in network architecture
of pre-training to achieve
any given environment-steps checkpoint
PRL introduces an auxiliary task that encourages intrinsically -ootivvate agent exploration. PRL agents are trained in two stages. First, the state encoder is pre-trained via task-agnostic method that attempts to maximize the entropy of the distribution of agent-visited states. Second, the state encoder is frozen and standard SAC agent is trained on top of the learned representations. We choose to not include this method in Section 4.2 as there are differences in architecture and evaluation procedure between the PRI. paper and the papers of methods shown in Section 4.2 For one, PRL's
15

---

Under review as a conference paper at ICLR 2022
Table 4: Training time of agent methods, normalized to Pixel SAC.
Method
Normalized Time
SAC Pixel
1.0
KSL
3.9
RAD
1.2
DrQ
1.8
CURL
1.6
SAC-AE
1.5
Table 5: Episodic evaluation returns (mean - one standard deviation) for PlaNet benchmark. High-
est mean results per task shown in bold. significant at p 0.05, ** significant p 0.01.
500k Steps
KSL
PI-SAC
PRL
Finger, Spin
976 + 14
936 = 84
868 + 105
Cartpole, Swingup
871 + 10**
844 19
843 = 21
Reacher, Easy
963 + 28**
888 = 58
599 + 190
Cheetah, Run
802 + 30**
374 + 68
392 + 58
Walker, Walk
953 8
939 = 16
755 + 109
Ball in Cup, Catch
973 + **
956 + 10
949 + 25
100k Steps
Finger, Spin
899 + 61 **
772 = 84
658 + 88
Cartpole, Swingup
841 + 33
722 - 94
366 + 235
Reacher, Easy
751 + 137**
337 + 143
126 + 99
Cheetah, Run
566 + 54**
227 + 55
243 + 39
Walker, Walk
730 + 133**
494 65
332 - 141
Ball in Cup, Catch
945 + 12**
836 + 106
639 + 183
encoder produces larger latent vectors than the encoders' of methods in Section 4.2 (50 vs. 128). Second, Yarats et al.) (202Ia) use an action repeat of two for all tasks and evaluate agents over 1M steps of training. The evaluations in PRL's paper are done in two-part 500k/500k split. First, 500k steps of auxiliary-task pre-training. Second, 500k steps of RL-task training. We keep this one-to- one ratio in our evaluations. For the 100k steps checkpoint, we train PRL agents in a 50k/50k split. Similarly, we train PRL agents in a 250k/250k split for the 500k steps checkpoint.
We adjust the encoder and SAC networks of PI-SAC and PRL to be congruent with methods in Section 4.2 all hyperparameters to be congruent with Table 3 in Appendix [ and all environ- ment settings to be congruent with Table 7 in Appendix []] Otherwise, in their original state. We make these adjustments within code provided by the original authors. Table 5 shows the evaluation returns across 10 seeds using the same methodology described in Section 4.3 Figure 7 shows the evaluation curves. We highlight that KSL agents outperform both PI-SAC and PRI. agents in all six tasks for both the 100k and 500k steps mark.
keep PI-SAC and PRL
E KNOWLEDGE SHARING AND CHOICE OF K
Figure shows potential options for the knowledge sharing (KS) mechanism within KSL. On the far left, we show the Q-function of the critic Qo- Although the KS mechanism uses layers from both Q-functions, we only show one in the "Q-function layers" diagram for simplicity. When 1, only the first layer of the Q-functions are used for KS, and when = 2, the first two layers are used. When = 0, the KS section of KSL's transition module is made of single linear layer that not part of the Critic's O-functions.
The choice & creates a trade-off in terms of algorithm performance and compute expense. A larger & gives KSL access deeper levels of temporal information. However, larger & results larger amount of compute, both terms of data sampling and the recurrent loop within KSL. With every incremental k; KSI. needs sample and perform the translation augmentation batch
16

---

Under review as a conference paper at ICLR 2022
Cheeral, Run
We bur. Maae
Rallin Cup Conce
180
Ewitommen SupERIaTT
Draicesed Stas (41P))
SHIpE-
Figure 7: Episodic evaluation returns for KSL, PI-SAC, and PRL agents trained in PlaNet bench- mark suite. Mean as bold line and + one standard deviation as shaded area. Purple and green dotted lines indicate when 50k/50k PRL and 250k/250k PRL, respectively, begin to receive task-specific rewards.
Q function layers
Linear
1=1
1=
KS
Linear
T
Linear
TP
Figure 8: Depiction of options for knowledge sharing. Critic's Q-functions shown on far left. KSL's transition module T shown on far right with two components: knowlege sharing (KS) and transition predictor (TP). Data-flow is shown with dashed lines.
Table Ablation experiment: 100k steps (left) and 500k steps (right).
0
1=1
/=2 2
0
-=
= 2
493
+ 33
493 + 59
512 + 40
h:
725
15
728
43
711 31
575
+
58
538 + 13
479 + 73
k 3
798
+
30
782
+
36
703 + 138
k=5
535 +
14
536 + 37
517 1 36
b = 5
763 + 40
775
32
712 = 35
size more states from the replay memory. Then, each additional state needs to perform several additional passes through neural network modules as well as compute loss, therefore leading to an exponential growth in compute requirements.
We perform full grid-sweep across all combinations of E {1,3,5 5} and E (0, 1, in the task of Cheetah, Run. Figure g below, shows the mean (bold line) and one standard deviation (shaded area) across five random seeds. (uble these results for the 100k and 500k checkpoints.
17

---

Under review as a conference paper at ICLR 2022
220
18
Figure 9: Grid search performance for KSL for k E (1,3,5) (left right) and / E (0.1.2). Mean performance shown in bold and + one standard deviation with shaded area.
Q+
Axc
Figure 10: Depiction of how KSL fits within the greater SAC algorithm.
These results show that knowledge sharing provides no benefit and may actually harm the RL learn- ing process. In general, when l 1, >> 0 leads to lower levels of performance with higher variance between runs. Also, the results show that there little reason to go beyond k 3. As such, the main results in the paper use = 0 and k = 3.
F How KSL FITS WITHIN SAC
Figure 10 shows how KSL fits within the greater SAC algorithm. t takes in current states and creates latent projections %o for both the critic Qs (bluc) and the actor To (red). 4m takes in future states $t+1 and creates latent projections Zm for the target critic Qa (green).
G CONVOLUTIONAL TRANSITION MODEL
We test the performance of convolutional transition model T within KSL.. To make convolutional T fit within KSL, we adjust its architecture. First, we remove the final dense layer from the encoders
Φ. T is then made of convolutional layers with 32 and 16 filters, respectively, and both have 3x3 kernel with a stride of two. Following the convolutional layers is single dense layer with 1024 +dim(AA) hidden units that outputs vector of length 50, followed by a layer norm operation. Action vectors are concatenated to the flattened feature maps from the transition model's convolutional layers before being fed through its dense layer. Both the critics and the actor have two convolutional layers added to their networks, all with the same architecture as described here for T. Additionally, they contain one dense layer with 1024 hidden units that outputs a vector of length 50. The remaining architectures of d m. &b P. Q9. Qg. and Rg remain unchanged. The additional layers in the critic target are updated as an EMA of their counterparts in the critic.
Figure II shows the evaluation results for both KSL with (five seeds) and without (ten seeds) con- volutional T. We highlight that the version KSL with a convolutional performs significantly worse than the version of KSL with fully-dense T. Also, bility of representational collapse in the version of KSL with a convolutional T As a result, we use the version of KSL with fully-dense T for all other results in this paper.
notice significantly higher proba-
18

---

Under review as a conference paper at ICLR 2022
C24
470
8.00
Figure 11: Mean + one standard deviation of evaluation returns for convolutional (conv) and fully- dense (no conv) transition model.
Fngger Span
Caripsto Swinuup
Neacher Easy
313
100
530
340
11
Cheebali Ran
We knc,
Bal Cup, Catch
Gralsnmeet Sepe
1.34
1:0
0.
112
Figure 12: Evaluation curves for the encoder encoder-generaliznn generalization experiment.
H GENERALIZATION LEARNING CURVES
Figure shows the learning curves for the generalization experiment.
I OPTIMIZERS VS. TASKS
The Adam optimizer (Kingma & Ba 2015 performs optimization updates by using exponentially- decayed first moment in and second moment i estimates of historical gradients. Adam updates parameters & at timestep as
Update direction
O++1 f- 0r
(5)
Vor +
where is small constant.
19

---

Under review as a conference paper at ICLR 2022
Chestah, Rua
Walker, Walk
550
300
Oy
23
5.22
3.92
2.90
0.55
00
032
0.44
0.08
Dru romment Shaps to
Enstrormnnt Steps (x:17)
Figure 13: Performance of agents on Cheetah, Run (left) and Walker, Walk (right) for scenario (i) (red) and (ii) (blue). Plots depict averages (bold line) and one standard deviation (shaded area) over five runs.
Angle Accumulated Gradient
1.50
0.50
0.45
00.25
-0.50
Ceetn Run.
Rescwr, Zay
Flreer Sch
-1.00
0.02
0.23
0.24
0.25
3.26
0.07
0.00
0.02
9.2
Envtronment Stops (x107)
Figure 14: Øi for for all tasks in the PlaNet benchmark suite throughout training.
In KSL's formulation, the RL task and the representation learning task both send feedback from their loss functions into 4 Two unique situations can arise with the optimization process KSL: (i single optimizer handling updates from both tasks (ii) separate optimizers handling updates for either task. The differences in optimizer-hed held gradient statistics results in different optimization steps for T in either scenario. Using the optimizer for both tasks may result in updates that conflate the accumulated gradient statistics of O between tasks.
To confirm this, perform two evaluations. First, measure the learning performance of KSL agents for scenario (i) and (ii) in both Cheetah, Run and Walker, Walk. Figure [3] below, shows the average and standard deviation across five seeds. We highlight that KSI. agents that use separate optimizers for either task result, on average, in better learning performance. Also, we note that these results are statistically significant at p 0.05.
Second, we query m and è from either optimizer in case (ii) throughout training for all tasks in the PlaNet benchmark suite. We extract the accumulated statisties for 4 and measure the angle Øig = (gi. 9t) ///9:/1lll9g)) where and 9; are the "update direction". as denoted in Equation 5 for the RL- and KSL-task optimizers, respectively. When dig < O, task task-speeific updates are considered to be conflicting. When gradients are conflicting, updates move weights in "opposing" directions. Figure[T4depiccs diy for all tasks PlaNet over three runs of training. We highlight that, on average, Φig < 0 for all tasks. This implics that maintaining combined m and è for both the RL and KSL task conflates update information from both tasks.
The RL task is unlike other optimization problems in that its sequential nature and non-i.i.d data create a moving target. Given that KSL's auxiliary task improves the performance on the RL task, hypothesize that the pull from KSL's conflicting update helps to regularize learning of the RL task. Also, the independent tracking of gradient statistics helps to increase the diversity of gradients, which has been shown improve learning in streams of non-i.i.d data (Aljundi al. 2019).
20

---

Under review as a conference paper at ICLR 2022
Table 7: Dimensions of action spaces, action repeat values, and reward function type for all six
environments in the PlaNet benchmark suite.
Environment, Task
dim((A)
Action Repeat
Reward Type
Finger, spin
2
Dense
Cartpole, swingup
8
Dense
Reacher, easy Cheetah, run
4
Sparse
4
Dense
Walker, walk
2
Dense
Ball in Cup, catch
4
Sparse
Figure 15: Depiction of the six tasks tasks in the PlaNet benchmark suite.
Alternatives to our multi-optimizer choice may include altering the gradients of the tasks before they combined and passed the optimizer (c.g.. Yu et al. (2020)), changing the architecture of 4 to include sask-fpec attention (Liu et al. 2019), or applying additional learning procedures that control the learning schedule between tasks (e.g.. Xu al. (209))
J ENVIRONMENT DETAILS
Table 7 describes the dimensions of the action space, the number of action repeats, and the type of reward (dense vs. sparse) for each task in the PlaNet benchmark suite. Figure [5 depicts the six tasks within the PlaNet benchmark suite.
K EXTENDED ENCODER SENSITIVITY
We measure the change in sensitivity perturbations the absolute difference between coder checkpoint to the next. A lower value is desireable, as it indicates smooth changes in latent representations throughout the training process. Smooth changes in latent representations results in less variability in input for the downstream RL training task. Best results are bolded. & this metric.
L EXTENDED PCA PLOTS
Here we provide PCA plots for all methods not shown in the main body of the paper. Figure [6 Figure 17
Figure 18
Figure 19/sow plots for SAC Pixel, RAD, CURL, and SAC-AE, respectively.
M EXTENDED AUGMENTATION PLOTSS
Figure shows plots for the random erasing augmentation and Figure 21 shows plots for the Gaus- sian noise augmentation.
21

---

Under review as a conference paper at ICLR 2022
Table 8: Change in sensitivity from one encoder checkpoint to the next.
Method
Translation
Erasing
Flipping
KSL
.008
.0044
.001
RAD
.016
.001
.006
DrQ
.015
.004
.018
CURL
.004
.001
.005
SAC-AE
.012
.002
.004
SAC Pixel
.032
.088
.452
Peal
175
15
4.4
Figure 16: First two principal components of SAC Pixel.
PCA.
PCA.
4.0
1.00
545
1.00
-441
Figure 17: First two principal components of RAD.
PCA,
PCA
13t
0.811
0.21
120
309
020
0
+000
-244
0.6
0.5
Figure 18: First two principal components of CURL..
SACHE.
PCA. SAGAE, 15k
SAC-AE 15k
010
M13
10
615
120
020
Figure 19: First two principal components of SAC-AE.
22

---

Under review as a conference paper at ICLR 2022
Invanance Standom Crasing)
RRandom Ereaing)
400
6.33
6.00
Figure 20: Random erase augmentation.
Tranalation (Cousiaa
Invariance Rolao)
400
6.00
Figure 21: Gaussian noise augmentation.
23


Reviews:
[{"review_id": "kBhysgOm7i1", "reviewer": "Reviewer_pTps", "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 paper addresses a very important problem in RL: data efficiency. The author's perspective on leveraging long-term temporal connection is not exactly novel (see Self-Predictive Representation [Schwarzer et al. 2021], Successor Features [Kulkarni, et al. 2016, Barreto et al. 2017]), but the specific method introduced seems to be novel, to the best of my knowledge.  \n\nOn the method:\nThe motivation is clearly stated and makes sense to me. I would have like to see discussion on how this differ/relates to successor features, when learned jointly or separately. \n\nA claimed \"that learned representations of the state space should relate to reward\" is not really justified: there are data-efficient method that disentangle the reward from the representation. It also seems to contradict another desired property: generalisation from one task to another. \n\nOn the experiment:\nThe experiments are extensive and the method is compared against sensible baselines. The results with respect to data-efficiency are promising.\n\n\nminor comments: \n- the font is very small on most figures axes\n- figures 5.2, 5.3 and 6 would make more sense with a different y axis scale.", "summary_of_the_paper": "The paper tackles the problem of sample inefficiency in continuous control, by noting that standard RL methods deal with both policy optimisation and representation learning jointly with a single supervisory signal, namely the reward. Consequently, authors propose to leverage long-term temporal connections between actions in the representation learning, and introduce k-Step Latent (KSL), a representation learning module for learning temporally consistent representations of the state space. Authors show that KSL improves previous state-of-the-art methods in PlaNet benchmark suite and provide some analysis of the representations learned by KSL.", "summary_of_the_review": "The main idea behind the paper is not novel, but the implementation is, and the results are promising. Some claims are not founded and somehow contradictory, and some related works are missing. But overall an interesting contribution!", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "hPFzN9KfUb1", "reviewer": "Reviewer_Qru8", "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": "Pros:\n\n - The paper is well-written and easy to comprehend;\n\n - The empirical evaluation on dm-control suite indeed shows that KSL yields state-of-the-art results on the six presented tasks;\n\n - The choice of using independent optimisers for training the representation learning module given the signals from the predictive latent supervision (Eq. 4) and critic training (Eq. 1) respectively raises a good point for representation learning in RL with auxiliary tasks, from multi-task learning perspective. Such inductive bias is additionally substantiated with empirical comparisons.\n\nConcerns:\n\n - The main evaluation on the overall performance is rather limited, it might worth showing more (especially the middle/hard tasks as defined in Yarats, et al. 2021);\n\n - The KSL model combines many existing techniques in representation learning for RL, such as image augmentation (Laskin, et al. 2020) and multi-step latent predictive supervision auxiliary task (Schwarzer, et al. 2021), leading to limited novelty of the proposed KSL model.\n\n - The multi-step latent predictive supervision objective for representation learning in KSL is highly similar to the SPR model (Schwarzer, et al. 2021), it seems like the authors acknowledge the similarity hence spent a paragraph discussing the difference, but discussions are mainly based on the architectural/training differences. From this perspective, KSL appears as an adaptation of the SPR model to the continuous control tasks (despite the admirable engineering efforts). The difference argument could be made convincing given that the authors could provide further clarifications of the differences between KSL and SPR, from an algorithmic perspective, and provide further empirical comparisons between the two agents (e.g., the authors state that \"KSL’s architecture is general enough to be applied in both discrete- and continuous-action domains\", hence it would be interesting to see an adaptation of KSL to discrete control and test it on atari benchmarks, and compare with the SPR results).\n\n - The arguments of the improved representation learning of KSL in terms of \"Permutation Invariance\" and \"Temporal Coherence\" seems poorly justified. The reported figures do not show significant or consistent improvement of the KSL over the baseline methods, only by largely zooming in the y-axis could one observe some differences. I doubt if such minuscule differences are consequential in the overall learning.\n\n - I think using the momentum encoder to provide input to the target Q-network in SAC critic-training is an interesting choice, but lacks further (theoretical) justification, i.e., why would this be better than simply using the online encoder as the inputs to the target Q-network, is it possibly because of the consistent temporal lags?\n\nMinor Points:\n\n - The main empirical evaluations of the learned representations is based on the walker-walk task, which is a simple task with dense reward structure(Yarats, et al. 2021). It would be more interesting to see how the learned representations are indicative of the reward structure in the sparse reward tasks, such that Cartpole-Swingup.\n\n - KSL is motivated by the inductive bias that \"States that are nearby in time are likely to share high levels of mutual information\", another similar work that utilises multi-step action-dependent latent predictions by Whitney, et al. (2019) is based on the inductive bias that the similarities between the embeddings for the states and/or action sequences should be based on their successor outcomes (e.g., successor representations). It would be nice to see some discussions on the relationship between the two seemingly independent inductive biases.\n", "summary_of_the_paper": "Summary:\nThe authors introduced k-Step Latent (KSL), a representation learning method for visual-based continuous control tasks. KSL utilises multi-step latent action-dependent predictive supervision for training the representation. The empirical evaluations are based on the dm-control suite benchmarks, KSL demonstrates improved sample efficiency (100k evaluation) and asymptotic performance (500k evaluation) across the six tasks presented, comparing to the baseline algorithms (mainly based on image augmentation). The authors further empirically examined the properties of the learned representations, and showed that the trained encoder quickly learns to be representative of the reward structure. The authors also argue that KSL supports more robust representation learning and stronger generalisability.\n", "summary_of_the_review": "Scores:\n\nI suggest marginal rejection (5/10). KSL indeed show state-of-the-art performance on the presented tasks and I like the way the authors assessed the quality of representation learning. However, KSL seems like a combination of existing methods, but lacks comprehensive empirical evaluation in that sense. Moreover, the high similarity with SPR is concerning without further clarification and empirical justification. Some arguments of the improvement learned representation is over-stated.\n", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "cTLf0waB7SN", "reviewer": "Reviewer_38iT", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "**Clarity.**\nFor the most part, the writing is clear. One minor concern is that the description of KSL (sec 3.2) is difficult to understand, and could be made much more clear with better choices of notation, including the pseudocode in the main text, and an explanation of the algorithm that follows the pseudocode. \n\nThat said, the main issue with the clarity and writing of the paper is that the contribution is not made clear both in the paper nor in the experimental analysis. Is the paper about using k-step latents plus an auxiliary loss tying these to observations? Or is the paper about the specific form this auxiliary loss takes? Most of the writing and the naming (KSL) imply the former but the evaluation speaks to the latter.\n\n**Novelty and significance.**\nClarity aside, the paper does not propose a sufficiently novel contribution for acceptance. There is a very large body of work in model-based reinforcement learning that uses k-step model-based predictions and corresponding losses to improve data efficiency and performance, and reduce model rollout errors. Some of this is even cited in this paper, but others that are not cited include Recurrent environment simulators by Chiappa et al 2017, TreeQN by Farquhar et al 2018, MuZero by Schrittwieser et al 2019, and Muesli by Hessel et al 2021. These methods all use k-step latents with a loss tying them to (encodings of) observations.\n\nThe specific form of the auxiliary loss here, adapted from BYOL (which is not even mentioned until the related work for some reason), is very similar to that of SPR, which is cited by the paper. There are some minor differences but these mainly seem like implementation details and since there are no empirical comparisons I have to assume this is the case.\n\nFurther, the experiments themselves are extremely limited. Evaluating only on 6 tasks from the DM control suite is not enough to show that this is a compelling and useful contribution, especially given the high overlap with prior work. The additional analysis of the learned representations are nice, but are not enough without showing the strength of the proposed approach, or else doing a much more thorough analysis. Finally, I’d like to see ablations of the components and choices made for the proposed method. Why have both \\Psi_o and P? Is using the EMA for the momentum pathway the best choice? Is a normalized L2 loss the best choice?\n", "summary_of_the_paper": "This paper proposes a representation-learning method (k-step latent, or KSL) that uses a self-supervised auxiliary loss between recurrently-predicted action-conditioned representations of the state space and non-recurrently predicted target representations, in the style of BYOL. The method is trained using two separate optimizers on different parts of the models, to avoid interference in the statistics maintained by the optimizers. Results at 100k and 500k steps on 6 tasks from the DM control suite (from pixels) compared to methods using alternative self-supervised auxiliary losses show that the proposed method improves data efficiency. Analysis of the learned latent representations shows that those from KSL produce more robust encoders and are more consistent with the underlying MDP.", "summary_of_the_review": "Overall, this paper lacks sufficient novelty for acceptance. It recombines existing techniques in a slightly different way than previously and shows improvements on a very small and narrow set of environments without comparing to the most relevant related work.", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "cAD_Wo6r5s", "reviewer": "Reviewer_TKuY", "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": "Strenghts:\n* Simple approach\n* Clearly written\n* Representation learning for better generalization or sample efficiency is an important topic in RL\n* Positive experimental results\n\nWeaknesses: \n* My main worry with this application of BYOL to RL is that the introduction of the transition model T changes the 'support' of z_m away from the support of z_o. In other words, while psi_o expects an output of T, psi_m gets the direct output of phi_m for which it was not trained and which might be entirely different from the output of T. Based on the experiments, this still seems to do something useful, but I would argue that psi_m should be seen more as random mapping than as the slow moving state-encoding. In any case, I think this out-of-distribution problem for psi_m should be addressed in the paper. \n* For Figure 3: Why not use t-SNI instead of only the first two dimensions of PCA? In particular, while it's not wrong to say that \"DrQ's projections show little sign of reward-based orgnaization by 15k steps\", that is slightly misleading as it doesn't say anything about the latent representation as we're only looking at 2 principal axes. \n\nAdditional Questions:\n* How is translation augmentation applied?\n* Nit: How would the results change when removing the sg before the policy?", "summary_of_the_paper": "The paper applied \"Bootstrap your own latent\" (BYOL) to the case of RL by introducing an additional learned transition model and show that this can improve sample efficiency.", "summary_of_the_review": "An interesting direct application of BYOL to RL. However, the necessity to include action-conditioned transition models in RL raises additional complications compared to BYOL which have not yet been addressed (or discussed) and I believe these should be included in the paper before publication. ", "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": "8GbrIxj5Ovh", "reviewer": "Reviewer_J6YX", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "Data-efficiency of RL algorithms is an important research area, and this paper explores how auxiliary self-supervised learning can improve data-efficiency in continuous control domains. The paper is very rich with experimental details, the implementation choices are carefully ablated and the paper is overall well written and explained. \n\n### Similarities with SPR\n------\nMy main issue with this paper is the fact that it positions the K-step latent objective as a “new representation learning method”, whereas in fact the K-step latent is exactly the representation learning method used in SPR [1]. Throughout the abstract, intro and the methods section, the paper positions KSL as a new representation method: but in practice it’s an adaptation of SPR to continuous control that requires some implementation changes. The entirety of section 3.2 and the Figure 1 is exactly SPR (see Section 2.2 and Figure 2 in the SPR paper), but the paper makes no references to it except in the related work section. The related work section again fails to acknowledge that KSL and SPR share the same representation learning objective, and not only the architecture. \n\nThe authors could very well have positioned this paper as “We applied SPR to continuous control, here’s some implementation changes we needed to make along the way to make it work”, and it would have been a much more honest and accurate description of the work.  The empirical contribution here and the detailed analyses itself would have been valuable on its own. \n\n### Questions on Experiments\n------\nThe description of the “Generalization of Encoders” experiments is very sparse. Can you specify what exactly the training tasks and evaluation tasks are in the generalization experiment? \n\nFor the invariance experiment, it would have been nicer to see invariance to real-world distractors and not artificial noise. I would recommend the Distracting Control Suite [2] for more convincing experiments around these.\n\nFor a lot of results, the performance seems to be under-reported than the original results in papers, especially for RAD and DrQ. Here’s a link to the raw performance scores for baseline methods used in [3] https://console.cloud.google.com/storage/browser/rl-benchmark-data/dm_control, and these were reportedly obtained from the corresponding authors. Can you clarify the discrepancy in the performance data? This seems to be a major issue. \n\nAdditionally, in a lot of performance curves, the standard deviation regions overlap, making it harder to establish stochastic dominance of one method over another. It would be nicer to see a better stochastic analysis using stratified CIs on multiple normalized metrics (see Figure 11 in [3]). You can do this easily via the colab: https://bit.ly/statistical_precipice_colab\n\n[1]  Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A., & Bachman, P. (2020). Data-efficient reinforcement learning with self-predictive representations. ICLR 2021. https://arxiv.org/abs/2007.05929\n\n[2] Stone, A., Ramirez, O., Konolige, K., & Jonschkowski, R. (2021). The Distracting Control Suite—A Challenging Benchmark for Reinforcement Learning from Pixels. https://arxiv.org/abs/2101.02722\n\n[3] Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A., & Bellemare, M. G. (2021). Deep reinforcement learning at the edge of the statistical precipice. NeurIPS 2021 https://arxiv.org/abs/2108.13264", "summary_of_the_paper": "This paper shows that an auxiliary self-supervised task that enforces temporal consistency of latents improves sample efficiency in continuous control environemnts, and identifies important implementation details that make it work. ", "summary_of_the_review": "Data-efficiency of RL algorithms is an important research area, and this paper explores how auxiliary self-supervised learning can improve data-efficiency in continuous control domains. The paper is very rich with experimental details, the implementation choices are carefully ablated and the paper is overall well written and explained. \n", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a reinforcement learning architecture that uses an auxiliary k-step step loss in the context of continuous control from image-based states.

While the topic is relevant and potentially impactful, several reviewers have major concerns about the manuscript. Among these, I highlight:
- Reviewers J6YX, 38iT and Qru8 have concerns about the novelty and contribution of the approach compared to existing literature.
- Reviewers J6YX, TKuY, 38iT and Qru8 have concerns about the experimental evaluation and the quality of comparisons to baselines.

Overall, it seems that the paper would benefit from further polishing.
Decision: Reject