Learning Representations for Reinforcement Learning with Hierarchical Forward Models
Authors: Trevor McInroe
Lukas Schäfer
Stefano V Albrecht
Venue: Submitted to ICLR 2023 (2023)
Keywords: Reinforcement learning
Representation learning
Continuous control

Abstract: Learning control from pixels is difficult for reinforcement learning (RL) agents because representation learning and policy learning are intertwined. Previous approaches remedy this issue with auxiliary representation learning tasks, but they either do not consider the temporal aspect of the problem or only consider single-step transitions, which may miss relevant information if important environmental changes take many steps to manifest. We propose Hierarchical $k$-Step Latent (HKSL), an auxiliary task that learns representations via a hierarchy of forward models that operate at varying magnitudes of step skipping while also learning to communicate between levels in the hierarchy. We evaluate HKSL in a suite of 30 robotic control tasks with and without distractors and a task of our creation. We find that HKSL either converges to higher episodic returns or optimal performance more quickly than several current baselines. Furthermore, we find that HKSL's representations capture task-relevant details accurately across timescales (even in the presence of distractors) and that communication channels between hierarchy levels organize information based on both sides of the communication process, both of which improve sample efficiency.

Under review as a conference paper at ICLR 2023
LEARNING REPRESENTATIONS FOR REINFORCEMENT LEARNING WITH HIERARCHICAL FORWARD MODELS
Anonymous authors Paper under double-blind review
ABSTRACT
Learning control from pixels is difficult for reinforcement learning (RL.) agents because representation learning and policy learning are intertwined. Previous ap- proaches remedy this issue with auxiliary representation learning tasks, but they either do not consider the temporal aspect of the problem or only consider single- step transitions, which may miss relevant information if important environmental changes take many steps to manifest. We propose Hierarchical k-Step Latent (HKSL), an auxiliary task that learns representations via a hierarchy of forward models that operate at varying magnitudes of step skipping while also learning communicate between levels in the hierarchy. We evaluate HKSL in a suite 30 robotic control tasks with and without distractors and task of our creation. We find that HKSL either converges to higher or optimal episodic returns more quickly than several alternative representation learning approaches. Furthermore, we find that HKSL's representations capture task-relevant details accurately across timescales (even in the presence of distractors) and that communication channels between hierarchy levels organize information based on both sides of the commu- nication process, both of which improve sample efficiency.
INTRODUCTION
Recently, reinforcement learning (RL) has had significant empirical success in the robotics do- main (Kalashnikov et al., 2018: 2021 Lu et al.) 2021 Chebotar et al., 2021). However, previous methods often require a dataset of hundreds of thousands or millions of agent-environment interac- tions to achieve their performance. This level of data collection may not be feasible for the average industry group. Therefore, RL's widespread real-world adoption requires agents to learn a satisfac- tory control policy in the smallest number of agent-environmen interactions possible.
Pixel-based state spaces increase the sample efficiency challenge because the RL algorithm is re- quired to learn useful representation and control policy simultaneously. A recent thread of research has focused on developing auxiliary learning tasks to address this dual-objective learning problem. These approaches aim to learn a compressed representation of the high-dimensional state space upon which agents learn control. Several task types have been proposed such as image re- construction (Yarats et 2020) Jaderberg al. 2017), contrastive objectives (Laskin et al., 20020a] Stooke et al.) 2021), image augmentation Yarats et al.) 2021 Laskin et al.) 2020b). and forward models (Lee al.) [2200a] Zhang et al. 2021/Gelade et al. 2019 Hatner al. 2020 2019).
Forward models are natural fit for RL because they exploit the temporal axis by generating repre- sentations of the state space that capture information relevant to the environment's transition dynam- ics. However, previous approaches learn representations by predicting single-step transitions, which may not capture relevant information efficiently if important environmental changes take many steps to manifest. For example, if we wish to train soccer-playing agent to score a goal, the pertinent portions of an episode occur at the beginning. when the agent applies a force and direction, and how close the ball came the goal. Using multi-step transitions
at the end when the agent in this situation could lead to more efficient learning, as we would focus more on the long-term consequences and less the large portion of the trajectory where the ball is rolling.

---

Under review as a conference paper at ICLR 2023
In this paper, we introduce Hierarchical k-Step Latent (HKSL an auxiliary task for RL agents that explicitly captures information in the environment at varying levels of temporal coarseness. HKSL accomplishes this by leveraging hierarchical latent forward model where each level in the hierarchy predicts transitions with a varying number of steps skipped. Levels that skip more steps should capture coarser understanding of the environment by focusing changes that take more steps to manifest and vice versa for levels that skip fewer steps. HKSL also learns to share information between levels via communication module that passes information from higher to lower levels. As a result, HKSL learns set of representations that give the downstream RL algorithm information on both short- and long-term changes in the environment.
We evaluate HKSL and various baselines in a suite of 30 DMControl tasks (Tssa et al. 2018 8 Stone jet al. 202T) that contains environments without and with distractors of varying types and intensities. Also, we evaluate our algorithms in "Falling Pixels", task of our creation that requires agents to track objects that move at varying speeds. The goal in our study is to learn a well perfforeronnnng control policy in the smallest number of agent-envirroment interactions as possible. We test our algorithms with and without distractors because real-world RL-controlled robots need to work well in controlled settings (e.g., laboratory) and uncontrolled settings (e.g., a public street). Also, distractors may change at speeds independently from task-relevant information, thereby increasing the challenge of relating agent actions to changes in pixels. Therefore, real-world RL deployments should explicitly learn representations that tie agent actions to long- and short-term changes in the environment.
In DMControl experiments, HKSL reaches an interquartile mean of evaluation returns that is 29% higher than DrQ (Yarats et al. 2021), 74% higher than CURL (Laskin et al., 2020a), 24%
higher than PI-SAC (Lee et al. 2020b). and 359% higher than DBC (Zhang et al. our experiments in Falling Pixels show that HKSL converges to an interquartile mean of evaluation returns that is 24% higher than DrQ. 35% higher than CURL, 31% higher than PI-SAC, and 44% higher than DBC. We analyze HKSL's hierarchical model and find that its representations more accurately capture task-relevant details earlier on in training than our baselines. Additionally, we find that HKSL's communication manager considers both sides of the communication process, thereby giving forward models information that better contextualizes their learning process. Finally, we provide data from all training runs for all benchmarked methods.
2021 Also,
2 BACKGROUND
We study an RL formulation wherein an agent learns a control policy within a partially observable Markov decision process (POMDP) (Bellman 1957 Kaelbling al. 1998). defined by the tuple (S, O, A. Ps oo R, 7). S is the ground-truth state space, O is a pixel-based observation space, A is the action space, P S x A x S + [0, 1] is the state transition probability functioon, p S x A x O -> [0. 1] is the observation probability function, R S x A - IR is the reward function that maps states and actions to a scalar signal, and y E [0, 1) is a discount factor. The agent does not directly observe the state 8, E S at step I, but instead receives specify as a stack of the last three images. At each step 4, the agent samples an action a, E A with probability given by its control policy which is conditioned Given the action, the agent receives a reward r; = R(s,, az), the POMDP transitions into a next state 8441 S with probability Pa P'((t,, 0t 8441), and the next observation (stack of pixels) 0++1 O is sampled with probability (* at, 0t+1). Within this POMDP, the agent must learn a control policy that maximizes the sum of discounted returns over the time horizon T' of the POMDP's episode: arg max, Ex E r'r.].
observation or O which we
the observation at time t, T T oo)).
3 RELATED WORK
Representation learning in RL. Some research has pinpointed the development of representation learning methods that can aid policy learning for RL agents. In model-free RL, using represen-
tation learning objectives as auxiliary tasks has been explored in ways such as contrastive objec-
tives [Laskin et al. 2020a Stooke et al., 2021). image augmentation (Yarats et al. 202 Laskin
al. 2020b), image reconstruction (Yarats al. 2020), information theoretic objectives (Lee al.
https::/anccyynous. 4open. www.n.scceneeeooeoooeccooooeeeee wcw.encer//eie www.n///////slleeeeeeeer 01000 README. md
2

---

Under review as a conference paper at ICLR 2023
2020b), and inverse models (Burda et al. 2019 Pathak et al. 2017). HKSL fits within the auxiliary task literature but does not use contrastive objectives, image reconstruction, information theoretic objectives, nor inverse models.
Forward models and hierarchical models. Forward models for model-free RL approaches learn representations that capture the environment's transition dynamics via next-step prediction objec- tive. Some methods learn stochastic models that are aided with image reconstruction (Lee et al.) 2020a) reward-prediction objectives (Gelada et al., 2019). Other methods combine forward mod- els with rewa rd prediction and bisimulation metrics (Zhang et al. 202T) or momentum regression targets (Schwarzer al. 2021). Outside of the purpose of representation learning, forward models are used extensively in model -ased RL approaches to learn control policies via planning proce-
dures (Hafner et al. 2020 2019
Ha & Schmidhuber 2018 Zhang et al. 2019).
Stacking several forward models on top of one another forms the levels of hierarchical model. This type of model has been studied in the context of multiscale temporal inference (Schmidhuber, 1991), variational inference CCung al.] (2017). and pixel-prediction objectives (Kim et al., [019 [axena
ct al. 2021). Additionally, hierarchical models have been used for speech synthesis (Kenter et al,
2019 learning graph embeddings (Chen et al. 2018), and decomposing MDPs (Steccanella et al.,
2021 Sequence prediction literature has explored the use of hicrarchical models via manually-
defined connections between levels (Saxena et al. 2201] Koutnik et al. 2014) and using levels with
uniform time-step skipping (Kumar et al.) 2020 Castrejon et al. 2019).
Unlike the aforementioned forward model approaches, HKSL combines set of forward models that step in the latent space with independent step sizes without additional prediction objectives. Also, HKSL contains a connection between forward models that learns what to share by using the context from the entire rollout from higher levels and the current timestep of lower levels, which leads to faster learning.
4 HIERARCHICAL STST LATENT
HKSL's hierarchical model is composed of forward models that take steps in the latent space varying levels of temporal coarseness. We define temporal coarseness as the degree to which a level's forward model skips environment steps. For example, if a forward model predicts the latent representation of a state five steps into the future, it is considered more coarse than a forward model that predicts only one step forward. Coarser levels should learn to attend to information in the environment that takes many steps to manifest in response to an agent's action. In contrast, finer levels should learn to attend to environmental properties that immediately respond to agent actions. This is because coarser levels need to make fewer predictions to reach steps further into the future than finer levels.
At each learning step, a batch of B trajectories of length l are sampled from the replay memory T = (o,, de...., ac+k-1 )) The initial observation of each trajectory or is uniformly randomly sampled on per-episode basis t ~ U1, : B In the following, we will denote the first and last timestep of the batch with /. = and / = k, respectively.
HKSL's components. See Figure I for visual depiction of the HKSL. architecture. HKSL's hier- archical model is composed of h levels. Each level / has a forward model f", a nonlinear projection module w', an online image encoder e',, and a momentum image encoder é,, that is updated as an exponential moving average of the online encoder (e.g., (He et al. 2020)) Between consecutive levels there a communication manager dl1-1 to pass information from one level | to the level below it [ 1. The number of steps skipped by a given level n' is independent of the coarseness of other levels in the hierarchy.
Forward models. HKSL's forward models are modified version of the common GRU recurrent cell (Cho et al. 2014) that allows for multiple data inputs at each step. See Appendix C.3 for detailed mathematical description. At step = 1, the forward models take the representation produced by the level's encoder of = (o)) along with a concatenation of n' action vectors āi = [-./a_ predict the latent representation of future state z z/=n' f' For any
Ending the range of numbers T guarantces that trajectories do not overlap episodes.
3

---

Under review as a conference paper at ICLR 2023
carat
[au]as|as]
[oufas/ad
f2
fz
#
22
Caras
Cluss
Figure I: Depiction of HKSL architecture with an "unrolled" two-level hierarchical model where the first level moves one step n' 1 and the second level moves three steps n² First, the online encoders (blue) encode the initial observation 0₁ of the sampled trajectory. Next, the forward models (red) predict the latent representations of the following observations, with level predicting single steps ahead conditioned on the level's previous representation and applied action. The for- ward model of second level predicts three steps ahead and receives the previous representation and concatenation of the three applied actions. The communication manager (green) forwards infor- mation from the representations of the coarser second level to each forward model step of the first level as additional inputs. All models are trained end-to-end with normalized l loss of the differ- ence between the projected representations of each level and timestep and the target representations of observations at the predicted timesteps. Target representations obtained using momentum encoders (purple) and projections are done by the projection model (orange) of the given level.
following timestep > 1, the forward models take the predicted latent representation zf as input instead of the encoder representation.
Communication managers. Communication managers d'1-1 pass information from coarser to finer levels in the hierarchy 1- 1) while also allowing gradients to flow from finer to coarser levels (1 - 1 + l). A communication manager d takes all latent representations produced by level I and one-hot-encoded step as inputs and extracts information that is relevant for the forward model in level 1 1 at step f. For all levels other than the highest level in the hierarchy, the forward models also receive the output of c.
Loss function. HKSL computes loss value at each timestep within each level in the hierarchy as the normalized & distance between a nonlinear projection of the forward model's prediction and the "true" latent representation produced by the level's momentum encoder. Using this "noisy" approximation of the target ensures smooth changes the target between learning steps and hypothesized reduce the possibility of collapsed representations (Grill et al. 2020: Tarvainen & Valpola 2017). We denote the projection model of level with w and the HKSL loss of level across the minibatch of trajectories can be written
CHKSL Σ Engour || (f d (+)) e/n (o++n)))3
where is the number steps that a given level can take in T.
HKSL and SAC. We make few adjustments to the base SAC algorithm to help HKSL fit naturally. critics. Each critic and target critic in given level's encoder and momentum
For one, we replace the usual critic with an ensemble the ensemble receive the latent representations produced encoder, respectively. We allow critics' gradients to update their encoders' weights, and each critic is updated using n-step returns where corresponds to the of the level within which the critic's

---

from Under review as a conference paper at ICLR 2023
given encoder resides. By matching encoders and critics in this way, ensure encoder weights updated by gradients produced by targets of the same temporal coarseness.
Second, the actor receives concatenation of the representations produced by all online encoders. HKSL's actors will make better-informel action selections because they can consider information in the environment that moves at varying levels of temporal coarseness. Finally, modify the actor's loss function to use a sum of Q-values from all critics:
Luctor = - -aass, 20n'
Σ (o, a)] - a log mlo
(2)
EXPERIMENTS
We evaluate HKSL with series of questions and compare it against several relevant baselines. First, is HKSL more sample efficient in terms of ggent-environment interactions than other representation learning methods 5.2 Second, what is the efficacy of each of HKSL's components ($ 5.3 Third, how well do HKSL's encoders capture task-relevant information relative to our baselines' encoders? 5.4 Finally, what does does dhl-1 consider when providing information 1 (S5 5,4)
5.1 EXPERIMENTAL. SETUP
Baselines. We use DrQ (Yarats et al. 2021). CURL (Laskin et al. 2020a). PI-SAC (Lee et al. 20206), DBC (Zhang et al. 2021 and DreamerV2 Hafner et al. (2021) our baselines. DrQ regularizes Q-value learning by averaging temporal difference targets across several augmentations of the same images, CURL uses contrastive loss similar to CPC (van den Oord et al. 2018) to
learn image embeddings. PI-SAC uses Conditional Entropy Bottleneck (Fischer) 20200 auxiliary loss with both forward and backward model to learn a representation of observations that capture the environment's transition dynamics. DBC uses a bisimulation metric and a probabilisite forward model to learn representations invariant to task irrelevant features. DreamerV2 is model- based method that performs planning a discrete latent space. All model-free methods use SAC (Haarnoja al. 2018a b) the base RL algorithm, while DreamerV2 leverages an on-policy actor-critic method with a A-target critic (Schulman et al.) 2016). All methods use the same encoder, critic, and actor architectures ensure fair comparison. Additionally, each method uses the same image augmentation. See Appendix for hyperparameter settings.
Environments. We use six continuous-control environments provided by MuJoCo (ooorov al. 2012) via the DMControl suite (Tassa et al. 2018 2020), popular set of environments for testing robotic control algorithms. Each of the six environments uses episodes of length 1 environment steps and set number of action repeats that controls the number of times the environment is stepped forward with a given action. We use five variations of each DMControl environment for total of 30 tasks. Four of the variations use distractors provided by the Distracting Control Suite API (Stone) jet al. 202T), and the fifth variation uses no distractors. We use the "color" and "camera" distractors on both the "easy" and "medium" difficulty settings. The color distractor changes the color of the agent's 's pixels on each environment step, and the camera distractor moves the camera in 3D space each environment step. The difficulty setting controls the range of color values and the magnitude of camera movement each task
environment of our design, which call "Falling Pixels". In Falling Pixels, the agent controls a platform at the bottom of the screen and is rewarded +1 for each pixel catches. Pixels fall from the top of the screen and are randomly assigned speed when spawned, which controls how far they travel downwards with each environment step. See Appendix B for further information on the environments.
Additionally, we
5.2 SAMPLE EFFICIENCY
Training and evaluation procedure. In our training scheme, agents perform an RL and representa- tion learning gradient update once per action selection. Every 10k environment steps in DMControl
RRfer to Stone al. 2021 for details.
5

---

Under review as a conference paper at ICLR 2023
Sample Efficiency
IQM
plimallty Gap
DreamerV2
0.5
DBC
HKSL
Evalustio Retums
PI-SAC
DrQ
0.4
CURL
CURL
0.3
HKSL
PI-SAC
DBC
0.2
DreamerV2
0.1
0.15
0.30
0.45
0.60
0.75
0.0
Normalized Score
6 Environment Steps (xx10)
10
Figure 2: IQM (left) and optimality gap (middle) of evaluation returns at 100k environment steps, and IQM throughout training (right) across all DMControl tasks. Shaded areas are 95% confi- dence intervals.
and 2.5k environment steps in Falling Pixels, we perform an evaluation checkpoint, wherein the agent's 's policy is sampled deterministically as the mean of the produced action distribution, and we compute the average performance across 10 episodes. All methods are trained with a batch size of 128. We train agents for 100k and 200k environment steps for five seeds in DMControl and Falling Pixels, respectively.
Results. We use the "rliable" package (Agarwal al. 2021 to plot statistically robust summary metrics in our evaluation suite. To produce aggregate metrics, we normalize all DMControl returns to the maximum per-episode returns, which is 1k for all tasks. Specifically, Figure 2 shows the interquartile mean (IQM) (left) and the optimality gap (middle) along with their 95% confidence intervals (CIs) that are generated via stratified bootstrap sampling at the 100k steps mark in DM- Control. Optimality gap measures the amount by which a given algorithm fails to achieve a perfect scord") Additionally, Figure 2 shows IQM and 95% CIs as a function of environment steps (right) in DMControl. Both of these results show that HKSL significantly outperforms our baselines across our 30 environment DMControl testing suite. See Appendix E for individual environment results. We note that simply using forward model does not guarantee improved performance, suggested by the comparison between HKSL, PI-SAC, and DBC.
Due to the randomness in Falling Pixels, the maximum per-episode return is difficult to calculate. Therefore, we do not aggregate Falling Pixels with DMControl returns, but instead show the IQM and 95% CIs for Falling Pixels function of environment steps in (leftt. We highlight that HKSL significantly outperforms all of our basclines, converging to performance of collecting over 20% more pixels per episode than the next-best- berf--rronnnnng algorithm. Collecting large number of pixels in Falling Pixels requires agents to keep track of environment objects that move at varying speeds. HKSL explicitly achieves this with its hierarchy of forward models. Also, we note that DreamerV2 struggles to outperform a random policy. We hypothesize that this is due to Falling Pix- els' observation space characteristics: the important information is single-pixel-sizer |Hafner al.) (202)) show that image- maage-rrconssruution gradients are important to Dreamer V2's success (Figure Hafner et (2021)). and the small details in Falling Pixels cause an uninformative reconstruc- tion gradient"
5.3 COMPONENT ABLATIONS
We probe each component of HKSL to determine its contribution to the overall RL policy learning process. Specifically, we test SAC without the hierarchical model but with HKSL's ensemble of critics (No Repr), HKSL where each level in the hierarchy moves with a single step (All n = 1), HKSL without e (No c), HKSL where each level in the hierarchy shares encoders (Shared Encoder), and single-level HKSL (h = 1). The No Repr ablation tests whether HKSL's performance boost
*For all plots, we performed at least 5,000 samples. 'We note that perfect score (optimality gap = 0) is technically impossible in the DMControl suite. As such, only the relative positioning of CIs should he considered.
Hafner al. (2021) also give this reason for why DreamerV2 does poorly in the VVideo Pinhall" environ- ment.

---

Under review as a conference paper at ICLR 2023
Evaluation Falling
Evaluation Falling
DrQ HKSL
14
Evaluation Returns
12
Evaluation Returns
DBC
12
10
PI-SAC
100
0.5
1.0
1.5
2.0
0.5
1.0
1.5
2.0
Environment Steps
Environment Steps (x111)
Figure 3: IQM and 95% Cls of evaluation returns for all algorithms in Falling Pixels (left) and ablations over HKSL's (right).
Abiiations in Carpols, Swingup
Ablations Ball in cup, Catch
Aclations in Walser Walk
daps
13N Normalter Score
(CB1 hormaltoet Gromm
108 Notoleee
Shardd caccile
80
0.0
Envronment (4387)
Environmmnt Steps
Endronment Stece
Figure 4: IQM 95% CIs of evaluation returns for HKSL. ablations in Cartpole, Swingup (left), Ball in Cup, Catch (middle), and Walker, Walk (right).
is due to the ensemble of critics or the hierarchical model itself. The All = ablation tests our hypothesis that only learning representations at the environment's presented temporal coarseness can miss out on important information. The No ablation tests the value of sharing information between levels. The Shared Encoder ablation tests if one encoder can learn information at varying temporal coarseness. Finally, the h ablation tests the value of the hierarchy itself by using standard forward model (e.g., (Schwarzer et al. 2021 Mclnroe et al. 2021)).
See Figure for the performance comparison between these ablations and full HKSL in the distractors setting of Cartpole, Swingup, Ball in Cup, Catch, and Walker, Walk. All results are reported IQMs and 95% CIs over five seeds. We highlight that variations without all components perform worse than full HKSL. This suggests that HKSL requires each of the individual components to achieve its full potential.
Also, we ablate across the number of levels in HKSL's hierarchy in Falling Pixels. Figure (right) depicts IQMs and 95% CIs over five seeds for values of in the 21, 3, 4} with temporal coarse- ness of levels set to [1, 3, 5, 7] for levels one through four, in order. We highlight that increasing achieves a monotonic improvement in evaluation returns up to when h 4. We hypothesize that setting h = 3 captures all relevant information in Falling Pixels, and increasing to h = leads to similar returns as when h 3 and does not destabilize learning.
5.4 REPRESENTATION ANALYSIS
How well do representations align with task-relevant information? To the ability of encoders to retrieve task-relevant information from pixel input, we save the weights of the encoders for each method throughout training in our evaluation suite. We then use the representations produced by varying
these encoders to train linear projection (LP) to predict task-relevant information timescales. This process akin to linear probing (Alain & Bengio 2017). a method used to analyze representations (e.g.. (Anand et al. 2019)). We note that the encoders' weights are frozen, and the gradient from the prediction task only updates the LP's weights.
In the Cartpole, Swingup task, the objective is to predict the cart's and pole's coordinates. In the Ball in Cup, Catch task, the objective is to predict the ball's s coordinates. We collect 10 and five episodes of image-coordinate pairs in each environment for I.P training and testing, respectively. We repeat

---

Under review as a conference paper at ICLR 2023
Cartpole, Swingup
No Distractors
Color
Camera
3.0
FKSL
2.2
DVQ
3.0
3.8
CURL
DBC
PISAC
0.4
DreameN2
9:6
0.2
1++1
2
2+ 10
2++1
f++55
+ 10
5.0
+++1
++
f+ 10
Ball in Cup, Catch
No Distractors
Color
Camera
9.030
3.225
3.025
9.629
9.500
5.015
0.015
0.012
3.005
0.035
8,000
[+1
(+2
[+33
3,000
[+1
(+2
1
0.000
2++1
2++2
1-433
Figure 5: MSE on task-relevant information in unseen episodes for Cartpole, Swingup (top) and Ball in Cup, Catch (bottom) at the 100k environment steps mark. Non-distraction, color distractor, and camera distractor settings shown from left-to-right. Lower is better.
Cartpole, Swingup
No Distractors
Color
Camera
3.4
HISL
1.2
3.0
Drg
CHRI
0.8
DOC
P-SAC
0.8
0.6
N2
2.6
2.4
0.4
0.2
0.0
11-1
24:55
Z- 10
1+-
f-1-5
2-+ 10
K4 10
Ball Cup, Catch
Distractors
Colar
Camera
2.020
0.033
3.005
0.030
2.013
8.615
2.010
0.015
0.020
3.500
2.003
0.012
9.005
0.005
3.000
t+2
t+33
3,000
f++1
0.000
6++2
++33
Figure 6: MSE task-relevant information in unseen episodes for Cartpole, Swingup (top) and Ball in Cup, Catch (bottom) at the 50k environment steps mark. Non-distraction, color distractor, and camera distractor settings shown from left-to-right. Lower is better.
this data-collection exercise for both environments' non-distraction, easy color distractors, and easy camera distractors versions. After fitting the LP on the training sets, we measure the mean squared error (MSE) on the unseen testing set. Figure]]shows the average MSE and + one standard deviation over the testing episodes using encoders trained for 100k environment steps our benchmark suite. In Cartpole, Swingup (top row), we use the LP to predict coordinates from one (f + 1), five (t + 5) and 10 (t + 10) steps into the future. Ball Cup, Catch (bottom row), we use the LP to predict coordinates from one (t + 1), two (t 2) and three (t + 3) steps into the future. We highlight that HKSL's encoders produce representations that more accurately capture task-relevant information with the lowest variance in nearly every case. Also, this accuracy carries over into the distraction settings, giving a reason for HKSL's 's relatively strong performance the presence of distractors, despite not addressing distractors explicitly.
8

---

Under review as a conference paper at ICLR 2023
& Distance Between Vectors from
PCA Projections Output
10
1
0.0
1.3
2.1
3.2
4.2
5.1
0.0
0.0
1.5
2.8
5
0.0
1.4
-5
-10.0
-7.5
-5.0
-25
0.0
2.5
7.5
Figure 7: Average distance between vectors produced by (left). The numbers along the side and bottom correspond to the value of t. PCA projections of representations produced by for multiple timesteps across 18 trajectories (right) with colors corresponding to trajectories.
We repeat this process using encoders from earlier in the agent-training process. MSE and + one standard deviation over the testing episodes using encoders trained for 50k environ- ment steps in our benchmark suite. We note that the same pattern from the 100k environment steps encoders persists. These results suggest that HKSL agents benefit from more informative represen- tations in earlier stages of training than our basclines, which leads to better sample efficiency.
the
What does e consider? We hypothesize that the communication manager d provides wide diversity of information for f by taking into account the current transition of the below level - as well the representations from the above level 1. To check this hypothesis, we perform two tests. First, we measure the la distance between the vectors produced by when the step / is changed and other inputs are held fixed. completely ignores 6, the distance between c(-, 1) and c(,, 4), for trajectory-
example, would be zero. Second, examine the separability of c's outputs basis. If two sampled trajectories are very different, then the representations produced by the above level should change c's output such that either trajectory should be clearly separable.
We first train an HKSL agent where h = 2, n² 1, and n² = in Cartpole, Swingup for 100k environment steps and collect 50 episodes of experiences with random policy. Then, randomly sample a trajectory from this collection and step through the latent space with both forward models. We repeat this 100 times and measure the pairwise & distance between c's outputs for every value of I within sampled trajectories. ( reports the average distance between each pair. We note that the distance between c's output grows the steps between the pairs grows. This suggests that e considers the transition of the level below it when deciding what information to share. Additionally, we highlight that the distance increases consistently where pairs that are the same number of steps apart are about the same distance apart. For example, pairs (2,5) and (3, 6) are both three steps apart and share roughly the same average lz distance. This suggests that e produces representations that are grouped smoothly in the latent space. Figure 7 (right) visualizes the PCA projections of e's outputs from 18 randomly sampled trajectories, where each trajectory
different color. This figure confirms our second intuition, as the representations are clearly separable trajectory basis with representations smoothly varying across steps within the same trajectory.
CONCLUSION
This paper presented Hierarchical k-Step Latent (HKSL), an auxiliary task for accelerating control learning from pixels via a hierarchical latent forward model. Our experiments showed that HKSL's representations can substantially improve the performance of downstream RL agents in pixel-based control tasks, both in terms of converged returns and sample efficiency. We also showed that HKSL's representations more accurately capture task-relevant information than baselines and do early in training. Finally, we showed that the communication manager organizes information in response to the above and below levels.

---

Under review as a conference paper at ICLR 2023
7 REPRODUCIBILITY STATEMENT
We open-source all code required to train HKSL. The anonymous URL on page two links to this code. Also, we open-source the code for our Falling Pixels environment by placing it within the supplementary .zip file. Finally, we release all data from the training runs for each of the algorithms used this study, which can be found the supplementary file.
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare G Bellemare. Deep reinforcement learning the edge of the statistical precipice. 2021.
Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes. International Conference on Learning Representations (Workshop Track), 2017.
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. 33rd Conference on Neural Infor- mation Processing Systems (NewrIPS), 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arXiv preprint arXiv:1607.064450, 2016.
Richard Bellman. markovian decision process. Indiana University Mathematics Journal, 6:679- 684, 1957.
Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexci A. Efros. Large-scale study of curiosity- learning. In International Conference Learning Repre- sentations (ICLR), 2019.
Lluis Castrejon, Nicolas Ballas, and Aaron Courville. Improved conditional vrnns for video predic-
tion. In International Conference on Computer Vision (ICCV), 2019.
Yevgen Chebotar, Karol Hausman, Yao Lu, Ted Xiao, Dmitry Kalashnikov, Jacob Varley, Alex Irpan, Benjamin Eysenbach, Ryan Julian, and Chelsea Finn andyou Sergey Levine. Actionable models: Unsupervised offline reinforcement learning of robotic skills. In Proceedings of the 38th International Conference on Machine Learning (ICML), 2021.
Haochen Chen, Bryan Perozzi, Yifan Hu, and Steven Skiena. Harp: Hierarchical representation learning for networks. In The Thirty-Second AAAI Conference on Artificial Intelligence (AAAI- 18), 2018.
Kyunghyun Cho. Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder decoder approaches. In Proceedings of SSST-8, Eighth Workshop on Syntux, Semantics and Structure in Statistical Translation, pp. 103-111, 2014. 2014,
Junyoung Chung, Sungjin Ahn, and Yoshua Bengio. Hierarchical multiscale recurrent neural net- works. In International Conference on Learning Representations (ICLR), 2017.
Ian Fischer. The conditional entropy bottleneck. Entropy, 2020.
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 Machine Learning (ICML), 2019.
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 sule-supe learning. In 34th Conference Neural Information Processing Systems (NewrIPS), 2020.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 1003v 10/22, 2018.
10

---

Under review as a conference paper at ICLR 2023
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: aetii0/92...0. 2018b.
Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning (ICML), pp. 2555 2555-2565, 2019.
Danijar Hafner, Timothy Lillicrap. 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-9733 2020.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In Proceedings The Thirty-Second AAAI Conference on Artificial Intelligence (AAAI-18), 2018
Riashat Islam, Peter Henderson, Maziar Gomrokehi, 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, Volodymyr Mnih, Wojeiech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations, 2017.
Leslie Pack Kaelbling, Michael I.. Littman, and Anthony R. Cassandra. Planning and acting in partially observable stochastic domains. Artificial Intelligence, 101(1):99-134, 1998.
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 reinforeement learning for vision-based robotic manipulation. In 2nd Conference on Robot Learning (CoRL), 2018.
Dmitry Kalashnikov, Jake Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski, Chelsea Finn, Sergey Levine, and Karol Hausman. Scaling up multi-task robotic reinforcement learning. In Proceedings of the 5th Conference on Robot Learning (CoRL.), 2021.
Tom Kenter, Vincent Wan, Chun-An Chan, Rob Clark, and Jakub Vit. CHiVE: Varying prosody in
speech synthesis with linguistically driven dynamic hierarchical conditional variational network. In Proceedings of the 36th International Conference on Machine Learning (ICML), 2019.
Taesup Kim, Sungjin Ahn, and Yoshua Bengio. Variational temporal abstraction. In 33rd Conference on Neural Information Processing Systems (NeurIPS), 2019.
Jan Koutnik, Klaus Greff, Faustino Gomez, and Juergen Schmidhuber. clockwork mnn. In Pro- ceedings of the 31st International Conference Machine Learning (ICML), 2014.
Manoj Kumar, Mohammad Babaeizadeh, Dumitru Erhan, Chelsea Finn, Sergey Levine, Laurent Dinh, and Durk Kingma. Videoflow: A conditional flow- model for stochastic video gener- ation. In International Conference Learning Representations (ICLR), 2020.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representa- tions for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning (ICMI), volume 19, pp. 5639-5650) 2020a.

---

Under review as a conference paper at ICLR 2023
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.
Alex X. Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pp. 741-752, 2020a.
Kuang-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 890-1199,,, 2020b.
Yao Lu, Karol Hausman, Yevgen Chehotar, Mengyuan Yan, Eric Jang, Alexander Herzog, Ted Xiao, Alex Irpan, Mohi Khansari, Dmitry Kalashnikov, and Sergey Levine. Aw-opt: Learning robotic skills with imitation and reinforcement at scale. In roceedings of the 5th Conference on Robot Learning (CoRL), 2021.
Trevor McInroe, Lukas Schäfer, and Stefano V. Albrecht. Learning temporally-consistent represen- tations for data-efficient reinforcement learning, arXiv preprint: arXiv:2110.04935,, 2021.
Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International Conference on Machine Learning (ICML), 2017.
Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork variational autoencoders. In 35th Conference on Neural Information Processing Systems (NeurIPS), 2021.
Jürgen Schmidhuber. Neural sequence chunkers. Technical report, 1991.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Picter Abbeel. High- dimensional continuous control using generalized advantage estimation. In International Con- ference on Learning Representations (ICLR), 2016.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self- eeff-rredittive representations. In International Conference on Learning Representations (ICLR), 2021.
Saurabh Singh and Shankar Krishnan. Filter response normalization layer: Eliminating batch de- pendence in the training of deep neural networks. In Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
Lorenzo Steccanella, Simone Totaro, and Anders Jonsson. Hierarchical representation learning for markov decision processes. arXiv preprint: arXiv:2106.01655 2021.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting con- trol suite - challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.02722 2021.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In Proceedings of the 38th International Conference Machine Learning (ICML), volume 139, pp. 9870-9879, 2021.
Antti Tarvainen and Harri Valpola. Mean teachers are better role models: Weight-averaged con- sistency targets improve semi-supervised deep learning results. In 31st Conference on Neural Information Processing Systems (NeurIPS), 2017.
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 1r01v
2018.
Yuval Tassa, Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, and Nicolas Heess. dm control: Software and tasks for 002006/29900 /2983, 2020.
continuous control. arXiv preprint arXiv:2006.
12

---

Under review as a conference paper at ICLR 2023
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033, 2012.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint arXiv:1807.03748. 2018.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Im- proving sample efficiency in model-free reinforcement learning from images. arXiv preprint 1FXX::1910001741 2020.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference Learning Representa- tions (ICLR), 2021.
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.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew J. Johnson, and Sergey Levine. Solar: Deep structured representations for model-based reinforcement learning. In International Conference on Machine Learning (ICML), 2019.
A EXTENDED BACKGROUND
Soft Actor-Critic Soft Actor-Critic (SAC) (Haarnoja et al. 201 8a is popular off-policy, model- free RL algorithm for continuous control. SAC uses state-action value-function critic Q and target critic Q, stochastic actor and a learnable temperature that weighs between reward and entropy: Eo |E, R(or, az) + (=)o(=-
SAC's critic is updated with the squared Bellman error over historical trajectories
T (0r, re, 0(+1) sampled from replay memory D:
fooitic E,,D[((((r 0z) (re + +3))
(3)
where y is computed by sampling the current policy:
= Bn [Q(0++ - log
(4)
The target critic Q does not receive gradients, but updated an exponential moving average (EMA) of Q (e.g., He et al. (2020). SAC's actor parameterizes multivariate Gaussian ((H,)) where z is a vector of means and a is the diagonal of the covariance matrix. The actor is updated via minimizing
Lactor = -Eo-s,r-Dl00 a) a log (a/oo)))
(5)
and a is learned against a static value.
B ENVIRONMENTS
Table the action space, the action repeat hyperparameter, and the reward function type of each environment used in this study. The action repeat hyperparameters that are displayed in the table the standards defined by (Hatner al. 2019) and the same used in most studies in DMControl. The versions of each environment with distractors follow the presented information as well.
The Falling Pixels enviromment is rendered as a 35 x 15 grayscale image. The agent is confined to the bottom row and pixels are spawned the top row. The agent is placed randomly along the bottom row and the top row is filled with pixels at the beginning of each episode. With each environment step, the pixels travel downwards until they reach the bottom row. the agent is occupying a pixel's column when it reaches the bottom row, that pixel is "collected" and the agent is rewarded +1. Regardless of whether a pixel is collected, disappears from the board once it reaches the bottom row. When column does not have pixel within it, there is 2.5% chance for new pixel to be spawned in that row each environment step. When spawned, the pixel is assigned a speed from the set (1,3, 5) uniformly random. Each episode is 250 environment steps.
13

---

Under review as a conference paper at ICLR 2023
Table 1:
Dimensions of action spaces, action repeat values, and reward function type for all six environments in the DMControl benchmark suite and Falling Pixels.
Environment, Task
dim((A)
Action Repeat
Reward Type
Finger, spin Cartpole, swingup
2
2
Dense
8
Dense
Reacher, easy Cheetah, run Walker, walk Ball in Cup. catch Falling Pixels
2
4
Sparse
4
Dense
2
Dense
4
Sparse
1
Dense
Table 2: SAC Hyperparameters used to produce paper's main results.
Hyperparameter
Value pixels
Image padding
Initial steps
1000
Stacked frames
3
Evaluation episodes
10
Optimizer
Adam
(Bi, (2)) Optimizer Learning rate
(0.9, 0.999)
1e-3
Batch size
128
Q function EMA
0.01
Encoder EMA
0.05
Target critic update freq
2
dim(z)
50
? Initial a
0.99
0.1
Target
-A
Replay memory capacity Actor log stddev bounds
100,000
[-10,2]
C ARCHITECTURE AND HYPERPARAMETERS
C.I SAC SETTINGS
All encoders follow the same architecture as defined by (Yarats et al., 2020 made of four convolutional layers separated by ReLU nonlinearities, linear layer with 50 hidden units, and final layer norm opertion (Ba et al. 2016). Each convolutional layer has 32 3 x kernels and the layers have stride of 1, 1. and respectively. This contrast the encoder used in the PI-SAC study (Lee et al. 2020b), which uses Filter Response Normalization (Singh & Krishnan 2020) layers between each convolution.
These encoders are
The architectures used by the SAC networks follow the same architecture deinfed by Yarats) jet al., 2020). Both the actor and critic networks have two layers with 1024 hidden units, separated by ReLU nonlinearities. This is in contrast to the networks used in the PI-SAC study, which uses different number of hidden units in the actor and critic networks.
Several studies have shown that even small differences in neural network architecture statistically signficant differences in performance (Islam et al. 2017; Henderson et al., 2018). As such, avoid using the original PI-SAC encoder and SAC architectures to ensure fair study between all methods.
Table shows the SAC hyperparameters used by all methods in this study. For method-specie hyperparameters (e.g., auxiliary learning rate, architexture of auxiliary networks, etc.), we defaulted to the settings provided by the original authors.
14

---

Under review as a conference paper at ICLR 2023
C.2 HKSL HYPERPARAMETERS
Table shows the hyperparameters that control HKSL. h represents the number of levels, con- tains a list of the skips of each level from lowest highest level, k shows the length of the trajectory sampled at each training step, learning rate corresponds to the learning rate of all HKSL's compo- nents, and actor update freq corresponds the number steps between each actor update. These hyperparameters were found with a brief search over the non-distractor setting of each environment.
HKSL's communication manager is simple two-layer non-linear model. The first layer has 128 hidden units and the second has 50. The two layers are separated by ReLU nonlinearity.
Table 3: Hyperparameters used for HKSL for each environment.
Environment, Task
h
n
Learning rate
Actor Update Freq
Finger, spin
2
[1.3]
Ie-4
2
Cartpole,
[1.3]
le-3
Reacher, easy Cheetah, run
2
[1,3]
le-4
2
[4.5]
10
lc-4
2
Walker, walk
2
[1.3]
6
le-3
I
Ball in Cup, catch
2
[1,3]
le-3
Falling Pixel
3 [1,3,5]
6
le-3
1
C.3 HKSL'S FORWARD MODELS
The usual GRU formulation at step
" yra σ (a)(((a((())))
(6)
Nyrn = tanh (o (irma @ 2 22-11 aa)))
(8)
Sgmu = - -1 wngma) 2 21-1 + u'gn @ h
(9)
where each each distinct f is affine transform, is the sigmoid nonlinearity, and is the Hadamard product. In order to allow the forward models to take the optional input from c, we add an identical additional affine transforms:
- = ((Ca((a)))))
(10)
r: ((Ck)
(11)
h = tanh(f (Ir' C((z-1))
(12)
d (1 ut) 2-1 + u' ht
(13)
where C₁ denotes the output from at step t. Finally, the output of the forward model is the average of the two pathways:
g' + Sgru 2
(14)
D ATTENTION MAPS
We examine the encoders within HKSL's hierarchy to ascertain their objects of focus. Each encoder receives gradients relating to different magnitude of temporal coarseness. Therefore, each encoder should learn to "focus" on different aspects of input images. The top row in each plot shows the unstacked frames that go into the past from right to left (e.g., the framestack depicted with images as [or-2, 0,-1, og]. The bottom row of each plot shows the attention maps from each encoder. The attention maps are generated by taking the output of the final convolutional layer and averaging across the feature map dimension. All encoders are from HSKL agents after 100k environment steps of training.
15

---

Under review as a conference paper at ICLR 2023
Figure 8: Input frame stack (top row) and corresponding attention maps (bottom row) for scenario
from Cartpole, Swingup. Encoder from first and second level shown on the left and right, respec- tively.
Figure 9: Input frame stack (top row) and corresponding attention maps (bottom row) for scenario from Cartpole, Swingup. Encoder from first and second level shown on the left and right, respec- tively.
Figure depicts scenario from Cartpole, Swingup. We that the encoder from the first level (left) attends to the pole, object that is not controlled by the agent. In contrast, the encoder from the second level (right) attends to the cart, which directly controlled by the agent. Figure 9 also depicts scenario from the Cartpole, Swingup environment. Here, the cart is offscreen for one frame the stack. Here, we see the same pattern as in Figure The encoder from the first and second level pay more attention the pole and the cart, respectively.
Figure depictts a scenario from the Ball in Cup, Catch environment. We highlight that the encoder from the first level (left) appears attend entirely to the information from the most recent frame in the input stack. In contrast, the encoder from the second level (right) gathers the full trajectory of information from each frame the stack. This phenomenon is especially apparent in Figure 11
16

---

Under review as a conference paper at ICLR 2023
U
Figure 10: Input frame stack (top row) and corresponding attention maps (bottom row) for sce- nario from Ball in Cup, Catch. Encoder from first and second level shown on the left and right, respectively.
Figure II: Input frame stack (top row) and corresponding attention maps (bottom row) for nario from Ball in Cup, Catch, Encoder from first and second level shown on the left and right, respectively.
where the encoder from the second level (right) captures the trajectory of the ball as it falls into the cup.
E
INDIVIDUAL ENVIRONMENT RESULTS
This section shows the mean (bold lines) + one standard deviation (shaded area) for every individual environment and distractor combination. Figure 12 displays the non-distractor environments, Fig-
ure 13 shows the color distractors on the easy setting, Figure shows the color distractors on the medium setting, Figure [3] shows the camera distractors the easy settings, and Figure 16 shows the camera distractors on the medium setting.
17

---

Under review as a conference paper at ICLR 2023
No Distractiors
Caripolle, Swingup
Bell Eatth
icer, Walk
IBZL 36
1030
500
830
600
SBC
came-12
430
200
0
Cheelah,
Finger, Spin
Beacher, Easy
700
4.30
600
200
200
1.0
200
0.4
2.6
1.0
0.2
5.A
2,5
0.2
0.4
0.0
Eie Fomman: S00ES Ae
Envirenment Steps Excats
Etee ronmaa Stons kett7)
Figure 12: Evaluation returns for agents trained in DMControl without distractors. Bold line depicts the mean and shaded area represents +- one standard deviation across five seeds.
Color Distractiors (Easy)
Swingup
Eatch
IKSL
300
820
600
CURL
630
420
200
Chestah, Run
Finger: Spin
Reacher, Easy
820
620
400
330
430
230
230
120
100
Crofonmanttteset..oo. icaPy
Inderomment Stepe (ouafs
Dnu Stape Exttl)
Figure 13: Evaluation returns for agents trained in DMControl with color distractors on the easy setting. Bold line depicts the mean and shaded area represents +- one standard deviation across five seeds.
Color Distractiors (Medium)
Ball Catch
Wb
IKCL
301
100
300
3remmer1
432
200
100
Cheetsh,
Finger, Spin
Beacher, Pasy
530
530
400
300
DaNistnn
230
Ora rowmamn Stace 277
Indrorment Stepe Ceurt
Cre rommary Stoce Couu7
Figure 14: Evaluation returns for agents trained in DMControl with color distractors the medium setting. Bold line depicts the mean and shaded area represents +- one standard deviation across five seeds.
18

---

Under review as a conference paper at ICLR 2023
Camera Distractiors (Easy)
Ball Catch
Walk
19CL
400
420
200
220
100
0-
Chestah,
Finger, Spin
Beacher, Easy
920
000
500
230
420
600
200
0.2
1.0
EreionmantSiesali 5o207)
Instrommen Staps tocahs
Drv romman: Stgg (4127)
Figure 15: Evaluation returns for agents trained in DMControl with camera distractors the easy setting. Bold line depicts the mean and shaded area represents +- one standard deviation across five seeds.
Camera Distractiors (Medium)
Cartpole,
Ball Eatch
ker, Walk
4:0
1901
300
3rQ
620
430
200
100
50
Cheetsh,
Finger, Spin
Beacher, Easy
720
600
2.30
320
220
120
0.2
0.1
0.6
3.5
EnwionmantS.oo 12207
Inskronment Stago 200371
Dn rormann Stage intaa)
Evaluation returns for agents trained in DMControl with camera distractors the medium setting. Bold line depicts the mean and shaded area represents +- one standard devia- tion across five seeds.
Figure 16:
19


Reviews:
[{"review_id": "go71sPs1nJ", "reviewer": "Reviewer_EsdY", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "", "summary_of_the_paper": "The main contribution of the paper is a hierarchical architecture for predicting observations at different levels of temporal coarseness, and its evaluation for RL from images.\n\nThe individual dynamic models for the different temporal coarseness levels are independent from each other and each consist of an encoder, a latent-dynamics model and a nonlinear projection network. Furthermore, except for the lowest level dynamic model, each model also contains a communication manager that summarizes all latent predictions within a trajectory segment and provides it to the latent dynamic model of the next lower level as additional input.\n\nThe models are evaluated using a modified SAC agent, which uses an example of Q function: one critic per temporal coarseness level. Furthermore, the policy gets all embeddings (1 per coarseness level) of the current observation as input. The dynamic model itself, however, is only used for representation learning, but not directly used during RL. The modified SAC agent is evaluated at the DMControl suite and on a new toy task \"falling pixels\". The result show that the proposed method can outperform suitable baselines in terms of sample efficiency.\n", "summary_of_the_review": "The paper is well written and of high quality, with sufficient novelty.\nHowever, I outlined a few issues that I would like the authors to address:\n1. Ablations of critical hyperparameters (number of hierarchies, chosen coarseness) are missing.\n2. There are some issues regarding the clarity", "strength_and_weaknesses": "- The paper tackles an important problem on how to capture and predict the effects of the agent's actions at different temporal coarseness levels. (+)\n- The method/architecture is sound and achieves good performance (+)\n- The paper is well written (+)\n- The results seem to be well reproducible, code published (+) \n- A few parts are a bit unclear (-)\n- The paper does not require theoretical analysis or derivations, so the novelty is limited to the network architecture (-)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "_qnwxv55JK", "reviewer": "Reviewer_6JA8", "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": "", "summary_of_the_paper": "This paper presents an approach towards data-efficient reinforcement learning that leverages a hierarchy of latent forward models to train representations at different levels of temporal abstraction. The models are designed to take in encoded observations at a given timestep t and a sequence of actions (where sequence length >= 1) and can predict future encodings of observations at different timescales depending on the coarseness of the corresponding model. Additionally there is a communication network that passes information from forward models at slower temporal scales to more fine-grained temporal scales. The models are trained using squared error between predictions and \"true\" latents, and are effectively used to shape the learning of observation representations; the representations from different levels are fed into the actor to predict the final action output by the policy. The approach (combined with SAC) is tested on several continuous control tasks from the DMControl suite along with variations such as distractors, and changes to camera pose or object colors; an additional task of \"falling pixels\" that potentially requires reasoning at different temporal abstractions is also used as a test environment. Comparisons are made against baselines which use different objectives to shape representations on top of SAC. Overall it performs better than the compared baselines at low sample complexity (100k transitions) thereby being more data-efficient. A few ablations and analysis of the learned representations are provided.", "summary_of_the_review": "Overall, while the approach presents a novel application of hierarchies of forward models to shape representations for data-efficient RL, the results are a bit weak. Some ablations, asymptotic performance results and analysis of the sensitivities of the approach are missing. As such I would recommend to reject this paper in its current form.", "strength_and_weaknesses": "Strengths:\n1. The paper provides an interesting approach towards training a hierarchy of forward models (albeit at explicitly specified levels of temporal abstraction) and shows that the resulting method performs well on a standard set of benchmarks. \n2. The addition of the communication module to communicate information from higher to lower levels is novel and seems to help improve performance on the tested domains.\n\nWeaknesses:\n1. The key novelty of the approach is on the addition of a hierarchy of forward models but there is no clear ablation in the paper on the sensitivity of the approach to the choice of hyper parameters of the hierarchy. Specifically there is no ablation across the number of levels used (h), and the number of steps skipped by a given level (n^l). The only relevant ablation provided was with [h=1] or [all n=1]. The appendix mentions that these parameters were chosen after a short search -- would be great if all the results across this search are presented. This would provide more empirical evidence to the stability and generality of the provided approach which is currently hard to evaluate.\n2. Table 3 of the appendix presents hyperparameters used with HKSL. There are 5 different sets of parameters for the 7 tasks presented in the paper -- parameters varied include the trajectory length (k), number of levels (h), number of skipped steps (n^l) and learning rate. Most prior work (e.g. DrQ, CURL) uses atmost one or two sets of hypers for a variety of tasks. Is there a particular reason to tune the hypers for HKSL? Relating to the previous question, is it due to the sensitivity of HKSL to hypers? Would be good to clarify this.\n3. On the topic of hypers, how does \"k\" affect the number of sequential predictions a model needs to make? For \"Finger Spin\" and \"Reacher Easy\" k = 3 but n_l = [1, 3]. Does this mean that the higher level forward model (n_1 = 3) does only a single forward prediction? Please clarify this in the text.\n4. All the results presented are only for 100k env steps. Most prior work reports results on 100k and 500k env steps as well as asymptotic performance; is there a reason why this is not done in the proposed approach? It would be useful to add these results as it would be good to know what the asymptotic performance of the method is, not just whether the data efficiency in the low sample complexity regime is good. \n5. All the presented baselines use either some form of representation loss or single-step prediction losses. Is there a reason to not compare to a baseline that uses multi-step prediction losses, or some form of reconstruction objective (e.g. Dreamer)? While reconstruction approaches can potentially fail on the distracting control suite they might do quite well on some of the analysis results presented (such as MSE prediction error of object positions). Would be a useful comparison to add.\n6. Re: Fig 7, it was mentioned that the data used for this comparison was collected using random policies. Why is this data interesting to consider for this particular analysis? Wouldn't it make sense to look at trajectories that are temporally consistent in some form for this comparison? It is also mentioned in the analysis that \"This suggests that c considers the transition of the level below it when deciding what information to share\", how do the results from Fig 7 (left) show this? From my perspective it only shows that the error scales proportionally to \"t\", not that the communicator (c) uses knowledge from the lower-level (which it doesn't have access to) to decide which information to communicate. It would be great if this point is explained better.\n7. Out of the tasks tested in the paper only the \"Falling Pixels\" task introduced in the paper explicitly requires some form of reasoning across temporal scales. It would be helpful if the approach was tested on atleast one other task (e.g. from Atari?) where temporal abstraction of some form was necessary.", "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": "VpPvEXrv-p5", "reviewer": "Reviewer_X4PZ", "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": "", "summary_of_the_paper": "This paper demonstrates how learning an ensemble of hierarchical models that predict latent representations at varying step-sizes into the future can be used for representation learning in model-free reinforcement learning. The proposed hierarchical technique is not used for planning, but rather the representations learned using their technique enable more efficient model-free reinforcement learning from pixels than baseline techniques. The paper provides compelling ablation results, baseline comparisons across 30 environments plus one newly introduced pixel-based control task, and releases all code needed to train their model and run the new environment.", "summary_of_the_review": "Overall I think this paper is clear, novel, potentially impactful, and represents sound and good science. I am recommending acceptance, but would reduce my score if the authors do not respond to my questions about model-based vs. model-free RL techniques.", "strength_and_weaknesses": "Strengths\n* Clarity of writing and presentation. The paper is very easy to read. The motivation is clearly laid out. The method is clearly explained for either a casual reader (main text) or more detailed reader (mathematically laid out in the appendix). The results are sign-posted and follow sensible questions to ask given the setup of the approach.\n* Statistics and reproducibility. The authors use interquartile means to show that their approach statistically outperforms recent baselines, as well as the various ablations of their model. The comparisons seem sound, and the diversity of experiments underscores the general utility of the hierarchical model-based representation learning technique. The authors also provide a link to code to ensure that results are reproducible. I commend the authors for their scientific openness and careful scholarship.\n* Novelty and impact. The idea of using hierarchical models for reinforcement learning is not novel for model-based reinforcement learning, where the hierarchical models are used for planning. However, I am not aware of work that uses hierarchical models in model-free reinforcement learning as a way of learning better representations that can ignore distractors. I thought this was a nice twist, and this paper would invite more investigation from others looking at hierarchical models for model-based RL to think about using the same approaches for generic representation learning in model-free RL instead.\n\nWeaknesses\n* Overall I really enjoyed reading this paper. However, I had one big question related to the framing of the method as being useful for sample efficiency. The presented approach is model-free, and results are presented after a *tiny fraction* of the number of environment steps often shown for model-free algorithms. The paper highlights the usefulness of their approach (HKSL) for exactly this kind of efficiency. BUT normally we would use model-based techniques to improve sample efficiency, which are absent from this paper. It would be particularly compelling if the presented method can outperform model-based techniques when environment samples are limited, or if the presented method still outperforms model-free baselines even in the limit of significant experience. \n  * To address this concern, could the authors provide results for one of two experiments:\n    * Compare to a model-based reinforcement learning method that uses the learned hierarchical models for planning, or really any other model-based RL technique, for the low-sample regime? I would like to understand whether the improvements here in model-free actually rival what one would expect for model-based.\n    * Compare results for massive amounts of experience as well (ie show the asymptotic behavior of all the presented baselines, not just their behavior after 50k or 100k steps). If HKSL, the presented method, still outperforms baselines in the limit of significant experience, that would ameliorate my concern about “why not just use model-based RL”.\n\n\n", "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": "Not applicable"}, "extra_text": {}}, {"review_id": "I2VyFIKnhf", "reviewer": "Reviewer_WeT7", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "Learn a hierarchy of latent models where each level takes in the latent state and a length n concatenation of actions and predicts the state after n steps, instead of simply predicting single step transitions. Information is passed from the layers predicting longer temporal distance to the lower levels. The representations are trained with l2 loss between the representations and the encoded observations. The representation is then used by SAC with h critics, corresponding to different levels of the hierarchy, and a policy representation that is based on the encodings learned by the levels.", "summary_of_the_review": "I think this paper can be accepted because it appears to provide a novel, simple change to existing model-based state representation RL methods, and it is implemented well enough with sufficient experiments to constitute a contribution.", "strength_and_weaknesses": "The reasoning that single step forward models fail to capture relevant information is a bit weak in the context of an MDP, where the markov property should ensure that all information is captured in a single state. \n\nIt isn't clear exactly how the encoder is prevented from collapse, since if the encoder is encoded to zero state then the forward models would have zero loss.\n\nWhile the results are promising and the writing is clear, the intuition behind what the different levels of the hierarchy might be encodinig. In particular, it would be nice ot have some visualization demonstrating how the higher levels capture different information from the lower ones, and which parts of the encoding the policy attends to. This might be more informative than indicating if task-relevant components are attended to, since it isn't clear why hskl would be especially effective at capturing task-relevant components. Otherwise, it is possible that a multi-head representation might get similar results, without the heads having temporally different meanings.\n\nThe experiments have many domains with only two levels, except for the falling pixels. It might be nice to look at a very long horizon tsk with a large number of levels.\n", "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": {}}]

Author rebuttal:

Decision: Reject