Learning Hierarchical World Models with Adaptive Temporal Abstractions from Discrete Latent Dynamics
Authors: Christian Gumbsch
Noor Sajid
Georg Martius
Martin V. Butz
Venue: ICLR 2024 spotlight (2024)
Keywords: world models
temporal abstraction
hierarchical learning
model-based reinforcement learning
hierarchical planning
Primary area: reinforcement learning
TL;DR: We propose an algorithm to learn a hierarchy of world models from sparse latent state changes for explainable, long-horizon planning.

Abstract: Hierarchical world models can significantly improve model-based reinforcement learning (MBRL) and planning by enabling reasoning across multiple time scales. Nonetheless, the majority of state-of-the-art MBRL methods employ flat, non-hierarchical models. We propose Temporal Hierarchies from Invariant Context Kernels (THICK), an algorithm that learns a world model hierarchy via discrete latent dynamics. The lower level of THICK updates parts of its latent state sparsely in time, forming invariant contexts. The higher level exclusively predicts situations involving context changes. Our experiments demonstrate that THICK learns categorical, interpretable, temporal abstractions on the high level, while maintaining precise low-level predictions. Furthermore, we show that the emergent hierarchical predictive model seamlessly enhances the abilities of MBRL or planning methods. We believe that THICK contributes to the further development of hierarchical agents capable of more sophisticated planning and reasoning abilities.

Published as a conference paper at ICLR 2024
LEARNING HIERARCHICAL WORLD MODELS WITH ADAPTIVE TEMPORAL ABSTRACTIONS FROM DISCRETE LATENT DYNAMICS
Christian Gumbsch Noor Sajid', Georg Martius & Martin V. Butz²
Veuro-Cognitive Modeling, University of Tübingen, Tübingen, Germany
2 Autonomous Learning, Max Planck Institute for Intelligent Systems, Tübingen, Germany
3 Wellcome Centre for Human Neuroimaging, University College London, London, U.K. corresponding author: christian, www.sch.umssnnooocooocc....m.... hwwpubscchhu--tbiiieeee... tuebingen.dee
ABSTRACT
Hierarchical world models can significantly improve model-based reinforcement learning (MBRL) and planning by enabling reasoning across multiple time scales. Nonetheless, the majority of state-of-the-art MBRL methods employ flat, non- hierarchical models. We propose Temporal Hierarchies from Invariant Context Kernels (THICK), an algorithm that learns world model hierarchy via discrete la- tent dynamics. The lower level of THICK updates parts of its latent state sparsely in time, forming invariant contexts. The higher level exclusively predicts situa- tions involving context changes. Our experiments demonstrate that THICK learns categorical, interpretable, temporal abstractions on the high level, while main- taining precise low-level predictions. Furthermore, we show that the emergent hierarchical predictive model seamlessly enhances the abilities of MBRL or plan- ning methods. We believe that THICK contributes to the further development of hierarchical agents capable more sophisticated planning and reasoning abilities.
a,
Level 2
Level
(a)
(b)
Figure 1: THICK world models predict
levels. (a) Level predicts the next input + 1).
Level 2 predicts future state (t + 7) expected to change an otherwise constant latent state. (b) Ex- emplary low- (bottom) and high-level predictions (top) for: opening door (Multiwou) pushing a boulder into water (Minihace -River) or activating pad Multtiwrldd-Doz))
INTRODUCTION
The intricate hierarchical representations formed in brains through sensorimotor experience (Lee & Mumford, 2003; Rougier et al., 2005; Botvinick & Weinstein, 2014; Rohe & Noppeney, 2015: Lake al., 2017: Friston al., 2018; 2021; Radvansky Zacks, 2014; Butz, 2016; Tomov et al., 2020) serve a useful blueprint for enhancing the planning abilities of artificial agents through hierarchical world models (Schmidhuber, 1992; LeCun, 2022). Humans, for example, can plan their behavior various time scales and flexibly switch between them, such as picking up pen to write an invitation when organizing party.
Despite recent advances in equipping MBRL agents with the capacity to learn world models, i.e., au- tonomously learned forward models that encode the interaction of an agent with its environment (Ha & Schmidhuber, 2018; Hafner et al., 2019b;a; 2020; 2023), these models lack hierarchical struc- ture. Consequently, they are restricted to predictions on predefined time scales, hampering their capability for long-horizon planning. The main challenge lies in formalizing suitable methods for

---

Published as a conference paper at ICLR 2024
learning higher-level abstractions (Sutton, 1988; Sutton et al., 1999; Eppe et al., 2022; Precup, 2000; van Seijen et al., 2014). Importantly these abstractions should be tied neither to particular tasks nor to fixed nested time scales. Confext-conditioned, wvent-predietive structures offer themselves as temporally flexible, basic compositional unit (Butz, 2016; Heald et al., 2021; 2023).
We present a deep learning architecture that learns hierarchical world models, which we call Tempo- ral Hierarchies from Invariant Context Kernels (THICK'). THICK adaptively discovers higher- level time scales by guiding the lower-level world model to update parts of its latent state only sparsely in time. The high-level model is then trained to predict scenarios involving changes in these low-level latent states, depiction of THICK world models can be found in Fig. la.
We make the following key contributions:
We introduce the Contcxt-specific Recurrent State Space Model (C-RsSM), which enhances Dreamer's (Hafner et al., 2019h; 2020) Recurrent State Space Model (RSSM) by encoding context-sensitive dynamics via sparsely changing latent factors, labeled context.
We introduce THICK, which learns hierarchy of world models. The high-level runs at an adaptive time scale developing higher-level actions that anticipate lower-level context changes.
We demonstrate the effectiveness of THICK two planning scenarios: i) using THicK's hierarchical predictions to enhance MBRL in long-horizon tasks, and ii) using THick's high-level predictions to set subgoals for hierarchical mreee-predictive planning (MPC).
2
METHOD
2.1 C-RsSM WORLD MODEL
The RSSM proposed in Hafner et al. (2019b) is a recurrent neural network (RNN) that is used for model-based reinforcement learning (Hafner al., 2019a; 2020; 2023; Sekar al., 2020; Mendonca et al., 2021; Sajid et al., 2021). RSSM embeds input images i, and actions a, into latent state and predicts dynamics exclusively within this state. All aspects of the latent state evolve continuously. We require sparse latent state changes to establish hierarchical world models. Accordingly, our Context-specific RSSM (C-RsSM) integrates a sparsely changing latent state C, as context with a coarse prediction pathway (cf. Fig. 2). Our C-RsSM with trainable parameters 6 is computed by:
Latent state: St + [c,, h4, zc]
(1) Pre. Prior:
2 P (2"
[c,, h2)
(4)
Coarse Dyn.: C = (ar11- Ct-1, 21-1)
(2) Coa. Prior:
2. ~ 22
a: Ct, z1-1)
(5)
Pre. Dyn.: h, folat-1, 14-1, 21-1) (3) Posterior:
2t ~ (z) Ct2, h4, i:)
(6)
Equations in red are exclusive C-RSSM? We separate RSSM'S latent state 8, into three parts (Eq. 1): a stochastic state 2;, a continuously updated, high-dimensional, deterministic state h2, and a sparsely changing, low-dimensional context cy. At time the C-RSSM first updates the context c2 (Eq. 2), where actual c changes only occur sparsely in time. Next, C-RSSM updates he via a GRU (Chung al., 2014) cell fg (Eq. 3). The C-RsSM makes two predictions about the next stochastic state z/: i) precise prior based both h2 and C2 (Eq. 4), and ii) coarse prior 2 based on the context, stochastic state, and action, ignoring h1 (Eq. 5). Given the input image é; C-RSSM updates its posterior 2, (Eq. 6). Following DreamerV2 (Hafner et al., 2020), we sample 2, from vector of categorical distributions. Note that Eq. 2 and Eq. 5 do not depend on h,-1, thus creating coarse processing pathway independent of h. This enables predictions using only G as determin- istic memory, which is crucial because it i) encourages encoding prddiciion-reeevan information in Ct and ii) allows predictions without he which we will use later (details in Suppl. D.1).
Besides encoding latent dynamics, C-RsSM is trained to reconstruct observable variables of the outside world from its latent states 8; Two output heads os generate precise and coarse predictions:
Precise prediction: ÿe (3e 8t) (7) Coarse prediction: Of 0g Ce, 2t). (8)
We predict the input image i, the reward rt,, and reward discount 76.) i.e., Vt (úr, 7c).
'In philosophy, the term 'thickness' refers concepts that combine descriptions with evaluative context (Roberts, 2013). A THICK world model fuses representations of the world with contextual interpretation.
*Removing e in all black equations recovers the equations for the RSSM (Eqns. 1,3,4,6).
"The discount y to if an episode terminates and fixed value y otherwise.
2

---

Published as a conference paper at ICLR 2024
F1
F1
F2
1v
63
AS
21
22
23
min KL
min KL
min KL
AC-
Figure 2: C-RsSM world model. Left: The C-RsSM encodes dynamics within latent states with stochastic part 2, and two deterministic parts h, and The network predicts the next stochastic state 2, via two pathways: It makes coarse predictions &f based mainly on C2 and precise predictions zh based h4- Right: Internally, the sparsely changing context is updated via GateL.ORD cell Is with a sparsely operated update gate. A GRU cell fo is used to continuously change h.
Sparse context updates The latent context code e₁ is designed to change sparsely in time, ideally distinct, environment-specifee transition points. Accordingly, the coarse dynamics Is (Eq. 2) are modeled by GateLORD cell (Gumbsch et al., 2021), which learns sparsely changing latent states C via update gate, whose activity Lu-regularized via loss term Esparse Note that context ee alone is too coarse to predict the current stochastic state 2,.
Loss function Given sequence of length of input images i1. actions a1 rewards rI,T, with discounts 7IT, the parameters & of C-RSSM are jointly optimized to minimize the loss L(():
((d) Eg (pred (6) + 3KL (k) + geparse Caparse (t)].
including the prediction loss gored the KL loss CK,, and sparsity loss Espare with respective hyper- parameters gured _K,,, and geparse. The prediction loss Crred drives the system to accurately predict perceptions y via its output heads Og, including context-conditionel coarse predictions (Eq. 8). The KI. loss CKL minimizes the KI. divergences between prior predictions p/ and p and the approx- imate posterior 46- The sparsity loss Lsparse encourages consistency of context C.. The exact loss functions provided in Suppl. D.3. We set gpred and JKL to Dreamer defaults (Hafner al., 2020) and modify the sparsity loss scale geparse depending on the scenario (cf. Suppl. B).
2.2 HIERARCHICAL WORLD MODEL
To learn hierarchical world model, we leverage C-RssM's discrete context C2 updates by means of our Temporal Hierarchies from Invariant Context Kernels (THICK) algorithm. A C-RSsM world model We segments sequences into periods of stable context activity (c, Cr+1 interspersed with sparse context updates (cf. Fig. 3a). THICK uses these discrete context dynamics as adaptive timescale for training high-level network Wg. The core assumption is that states prompting context updates coincide with crucial changes in latent generative factors. These key states are predicted by the high-level network Ws. while states between context updates are ignored.
Cr-1) 1),,,
To train the high-level world model Wg. we require input-target pairs for a given sequence T images 21.7, actions 01.2, and episode termination flags dr.r. The sequence is passed through the low-level model We to obtain a sequence of contexts C1:7- Targets are defined as all time steps т with context changes, i.e., where C, * Cy the episode ends. We define the function r(-) as
7()) min(( ({r
(c, + c, Vd, =1)))
(10)
Thus, 7(-) maps every point to the next point in time r(t) with context change, effectively imple- menting variable temporal abstraction that generates target predictions =(t) for every t.
3

---

Published as a conference paper at ICLR 2024
stable contert 9us Gro
new C₁
Arli)
W
(a) Time series discretization
(b) Level 2 world model
Figure 3: High-level segmentation: (a) The low-level C-RSSM discretizes sequences into segments with constant contexts. We use this segmentation to determine inputs and targets for the high level. (b) The high-level world model predicts the states and actions that lead to a context change at time (tt) from latent states 2, and c.. High-level actions (A, or A.) distinguish high-level outcomes.
High-level targets We predict all variables at (() that may cause a context change or are needed for planning across a context change: (-()) (, (t),,, F?:7()) (cf. Fig. 3b). In particular, we predict the stochastic states (r and actions a+(t)-1 immediately before context change at time rt), because both can cause an update of CT(c) (see Eq. 2). Intuitively, this means that observations, e.g. seeing something fall, as well as actions, e.g. catching something, could contribute to a change of c, We furthermore predict the elapsed time (t and the accumulated discounted reward "2.7())>
which may account for variable duration and rewards when evaluating high-level outcomes:
(r(t)-1
Elapsed time:
AT(1)
r(t) (tt) - Accumulated rewards:
P (11)
(:7(1)
8=0
High-level inputs To predict high-level targets, we use the low-level stochastic state 2, and context G as inputs. However, we need to disambiguate different potential outcomes, which generally depend on the world and the policy pursued by the agent. Accordingly, akin to actions on the low level, we create self-organizing high-level "actions" A,, similar to skills or options (Sutton al., 1999). A, encode a categorical distribution over probable next context changes. To learn A,, the high-level world model implements a posterior action encoder Qo and a prior action encoder Po (cf. Fig. 3b). Overall, the high-level world model We with learnable parameters 0 is computed by:
Post.: A, ~ Qo(Ai ct 2t, Cr()): 2r(1))
(12)
Prior: A, (A. Ct2 c1,
(15)
Action: (,(()) ~ F (ār(0)-1 A,, Ct2 21)
(13) Time A+(t) F (Af(t) A Ct2 C2, 21)
(16)
State: 2- F (1)- 1 A Cc2 21) (14) Reward: A7 (t((t) ~ F (f Ar, C1, 21) (17)
The posterior Qo receives not only C₁ and 2, its input but also privileged information about the actually encountered next context, i.e. Cr(1) and Zr(e) (Eq. 12), which leads to the emergence of individualized, result-conditioned action encodings in A,, The prior Po learns distribution over A, approximating the posterior without the privileged information (Eq. 15). During training, THICK samples the high-level action A, from Qg. During evaluation, we sample from the prior Pa instead. We model A, and A, as one-hot encoded categorical variables.
Loss function
The high-level world model Wa with parameters is trained to minimize the loss 2(0) E pred spred (8) + a (0)).
(18)
with hyperparameters pred and a A scaling the prediction gored and action Ω loss terms, respec- tively. The prediction loss drives the system better predict the high-level targets. The action loss drives the system to minimize the KI. divergence between the posterior high-level action distribution Qo and the prior distribution Po. The exact loss functions can be found in Suppl. D.4.
Summary Our THICK world model augments traditional flat world models by a high level. which lears predictions of variable length, anticipating context transitions. This augmentation allows for seamless transitions between coarse, low-level and abstract, high-level predictions. Given
4

---

Published as a conference paper at ICLR 2024
High level:
W
Low level::
W
Figure 4: Temporal abstract predictions of THICK world models. From low-level context Cr and stochastic state 2t, the high level predicts a future stochastic state (r(t)=1 as well as the action ar(t)-1 With these predictions, the context Cr()) is updated on the low level together with the coarse prior (())- This process can be repeated (dashed line) create temporal abstract roll-out.
context Ct. stochastic state and sampled high-level action A., the high-level model We predicts a scenario (ir(()-1, 2x(1)-1) immediately prior to the next anticipated context change. By feeding this prediction into the coarse processing pathway of C-RssM, context C+(1) (Eq. 2) and a coarse prior estimate of the corresponding stochastic state (c()) (Eq. 5). Longer temporal abstract roll-outs can be created by feeding C+(() and (7()) again into We (see Fig. 4). In this way, actual context change predictions are naturally generated by C-RsSM.
predict the subsequent, new
2.3 DOWNSTREAM APPLICATIONS OF THICK WORLD MODELS
World models have been applied in many downstream tasks, including MBRL (Ha & Schmidhuber, 2018; Hafner et al., 201 90: 2020; 2023), exploration (Sekar et al., 2020; Sancaktar et al., 2022), or model-predictive control (MPC) (Hafner et al., 2019b; Vlastelica et al., 2021). With minimal changes, the hierarchical roll-outs from THICK be seamlessly integrated where flat roll-outs were previously utilized. We exemplify this integration in two key areas: MBRL and MPC.
2.3.1 THICK DREAMER: MBRL WITH HIERARCHICAL ROLLOUTS
Dreamer (Hafner et al., 2019a) learns behavior by training an actor and a critic from "imagined" roll-outs of its RSSM world model. More specifically, Dreamer imagines sequence of states 8t. H from a start state 8; given an ector-generated action sequence d::--H- Dreamer computes the general A-return (s.) (Sutton & Barto, 2018) for every 8, and its critic vig is trained to regress (*))..
In sparse reward tasks, one challenge is reward propagation for training the critic (Andrychowicz al., 2017). Here, Dreamer faces difficult trade-off: Long roll-outs (large H) speed up reward propagation but degrade the quality of the predicted roll-outs. We propose THICK Dreamer, which combines value estimates from low- and high-level predictions to boost reward propagation. THICK Dreamer maintains an additional critic ex to evaluate temporal abstract predictions. Like Dreamer, we first imagine low-level roll-out of H states St:+H- Additionally, for every time in the roll-out, we predict temporal abstract outcome Cr(t) and (r()) and estimate long horizon value Vloog as
(sl)n = =f EET( +
+ + (1() (Cr(())) Er()));
(19)
with all variables predicted via the THICK world model and immediate rewards via Eq. 8 of C-RssM given THICK's world model predictions (cf. also supplementary Alg. 1). We estimate the value of a state 8; as a mixture of short- and long-horizon estimates with
V(s:) 2 v (1 (lo
(20)
where the hyperparameter controls the trade-off between the two estimates. We set 1 0.9 in all experiments and train both critics ii and 7ix to regress the value estimate. In sum, to speed up credit assignment when learning value function, THICK Dreamer combines low-level roll-outs with temporal abstract predictions to additionally estimate the value of likely long-horizon outcomes.
2.3.2 THICK PLANET: HIERARCHICAL MPC
The original RSSM was proposed in PlaNet (Hafner et al., 2019b) as world model for MPC. PlaNet searches for the optimal action sequence di,+ to maximize the predicted returns f.++ Thereby, PlaNet employs zero-order trajectory optimization via the cross entropy method (CEM) (Rubinstein, 1999). Once a is identified, the initial action at is executed and the procedure is repeated.
5

---

Published as a conference paper at ICLR 2024
KeyRoom CerRooo-Fixede Fixed-85
Door
4
HL:
1+3 t+5
1+7
5
10
Figure 5: Context changes. We show the input images i, 16-dim. contexts c++1 and reconstructed
high-level predictions i -(())1- For KeyRoom the context changes when finding the key, picking it up, opening door (here from diagonally adjacent grid) or exiting the room. In Door the context changes when the robot grabs the handle. The high level predicts the states before the next changes.
CEM optimizes randomly sampled trajectories. Sampling good action sequence is exponentially harder for increasing task horizons. We hypothesize that such tasks could be solved with much fewer high-level actions. For this, we propose THICK PlaNet. THICKK PlaNet plans on the high level to solve the task and uses the low level to follow this plan. We define reward function R(()) to estimate the return of a high-level action sequence A1:k with length K recursively as
pm (()) + free C R(A++1KK +1))+1)
for k K, for k = K
R R(A (') = ((() + 7/
(21)
p T(t)
with all variables predicted via temporal abstract roll-out (see Sec. 2.2) starting with A: 1 and t. We search for the optimal sequence Ái. maximizing R(-) with Monte Carlo Tree Search. Based on the first action Á we sample a subgoal 2q ~ Fo A Ct zz). This subgoal is valid as long it has not been reached yet and nothing has drastically changed in the environment. Thus, we only replan on the high level when the context has changed. We apply CEM on the low level to reach z0oal while also maximizing task return with
++
arg max
Σ
for + K. sim(2z 2, ) with frr of se),
(22)
for a planning horizon H. The function sim(.) is a similarity measure between z and z,. The hyperparameter x controls the trade-off between external and internal reward. Previously, similarity between Gaussian distributed z, of the RSSM was estimated using cosine similarity (Mendonca et al., 2021). However, for the categorically distributed 2,, the cosine similarity can be low even when they stem from the same distribution. Instead we use the cosine similarity of the logits, i.e.
4 goal
sim (1,, =1 goal)
(23)
||4e
where is the dot product and L, and 1500l are the logits of the distributions that produced z; and 2; goal respectively. Compared to other similarity measures, e.g. KL divergence, our measure has the desirable property that sim 2 =00al)= e [0, 1]. which simplifies setting the hyperparameter k, which we set to = 0.025 to mainly guide the behavior in the absence of external reward.
3 RESULTS
We empirically evaluate THICK to answer the following questions:
Can THICK learn temporal abstractions? We show that the learned high-level world model in- deed discerns meaningful, interpretable temporal abstractions across various scenarios (Sec. 3.1). Can THICK's hierarchical predictions improve MBRL? We show that THICK Dreamer achieves higher returns than Dreamer in long-horizon tasks with sparse rewards (Sec. 3.2).
Can THiCk's 's world model used to plan hierarchically? We show that MPC with THICK world models is better than flat world models solving long-horizon tasks (Sec. 3.3).

---

Published as a conference paper at ICLR 2024
WandDrDesth- Advanced
PickUp
VieusIPinPadPPV
2+2
(+
(+2
[+3
Figure 6: High-level actions A., We show input images i, and predictions a for two high- level actions A and A?. Red frames depict sampled actions Á. Exemplar actions A. are: Exiting the room or attacking monster (left), grasping or pushing ball (center), activating pads (right).
We evaluate our THICK world models in various scenarios. MiniHack (Samvelyan et al., 2021) is sandbox framework for designing RL environments based on Nethack (Küttler al., 2020). We test our system on benchmark problems as well as newly created tasks. All problems, detailed in Suppl. 1, have hierarchical structures which subgoals need be achieved (e.g. fetch wand) to fulfill a task (e.g. kill a monster) to exit dungeon and receive sparse reward. The observation is pixel-based, ego-centric view 2 grid-cells around the agent. MiniHack uses discrete actions.
VisualPinPad (Hafner et al., 2022) is suite of visual, long-horizon RL problems. Here an agent (black square) needs step on a fixed sequence of pads receive sparse reward. We use three levels of difficulties based the number of pads and target sequence length (three, four, five).
MultiWorld (Pong et al., 2018) is a suite robotic manipulation tasks for visual RL. these tasks Sawyer robot has to either move an object to a goal position (puck in Pusher or ball in P PickUp) or open a door (Door). We use fixed goals and take the normalized distance between the to-be- controlled entity and the goal position as dense rewards (in Pusher-Dense, PickUp, Door) and thresholded distances as sparse rewards (in Pusher Pusher-Sparsee -Sparse.. Details are provided in Suppl. E.2.
3.1 INTERPRETABLE CONTEXTS AND HIERARCHICAL PREDICTIONS
First, we analyze the predictions of THICK world models across diverse tasks. Example sequences displayed in Fig. 5. in Suppl. F.1 and on website. In MiniHack. context updates typically coincide with item collection, map changes, area exploration, or dungeon exits. In Multiworld, context changes occur due object interactions vating pads can prompt context changes. The high-level model predicts the states preceding context changes, often abstracting details, leading blurry reconstructions, For instance, KeyRoom, the system forecasts the agent's 's level exit without knowledge of the exact room layout (Fig. 5, t + 6). Nevertheless, the lower level consistently predicts the next frames accurately, as shown Fig. 1b.
workspace boundaries. In VisualPinPad, acti-
Abstract action representations A, emerge on the high level, as illustrated in Fig. 6. These actions categorically encode different agent- world interactions, e.g., grasping or pushing ball in P i ckk.p. The prior Qo learns to sample actions based on the likelihood of their outcomes (red frames in Fig. 6). If there are more actions A, than necessary, different actions encode the same outcome.
3.2 MODEL-BASED REINFORCEMENT LEARNING
We investigate whether hierarchical roll-outs can improve MBRL in the MiniHack suite by com- paring THICK Dreamer to DreamerV2 (Hafner et al., 2020) and to Director (Hafner et al., 2022), a hierarchical RL method based on Dreamer. Fig. 7a-7d show that THICK Dreamer matches or outperforms flat Dreamer in all tasks in terms of sample efficiency or overall success rate. The ad- vantage of THICK Dreamer is more pronounced in tasks that require completing multiple subgoals (e.g. completing five subgoals in EscapeRoom Director outperforms the other methods KeyRoom but fails to learn other MiniHack tasks. We investigate the failure cases of Director in Suppl. F.3 and show MiniHack results in Suppl. F.2.
finding key to open door in KeyRoom).

---

Published as a conference paper at ICLR 2024
(a)
KayRoon-Pisss Pixen-S5
(b)
KeyCorridor-8
(c)
WandOfDeath Daahhf-eeethhddd
(d)
EscapeRoom
Dirictoo
0.5
Dreamer
0.5
0.5
0.5
Timcx
Dresmeee
0
2.5
6
10 stepps
10
steps
10 steps
10 steps
(e)
(1)
mean improvements in returrss
(g)
ineam % improovenents
0.4
20
0.2
corricor
length
5
9
corridor length
corridor length
Figure 7: MiniHack results. Top graphics (a-d) plot the mean success rate during evaluation for various MiniHack tasks using 7 seeds. For KeyCorri dor (e) we systematically vary corridor length and plot mean differences in evaluation returns (f) and percentage of task success (g) between THICK Dreamer and Dreamer over different lengths. Shaded areas depict + one standard error.
We hypothesize that task horizon length the main factor boosting THICK Dreamer's performance. To investigate this, we systematically vary the task horizon in the KeyCorridos problem (see Fig. 7e) by modifying the corridor length. Fig. 7f-7g plot the mean difference in obtained rewards and success rate over 500k steps of training between THICK Dreamer and Dreamer for different corridor lengths. The performance gain of THICK Dreamer tends increase with corridor length until at some length both approaches fail to discover rewards during training, detailed in Suppl. F.2.
We further analyze the effect of task horizon VisualPinPad. VisualPinPad poses two challenges: exploration and long-horizon behavior. To analyze the latter in isolation, we sidestep the challenge of discovering the sparse rewards by initially filling the replay buffer of all models with 1M steps of exploration using Plan2Explore (Sekar et al., 2020) (details in Suppl. F.4). Fig. 8 shows the performance of THICK Dreamer, DreamerV2. and Director. THICK Dreamer matches Dreamer in P PinPadTh ree and is slightly more sample efficient in the more challenging tasks.* 4 Thus, fusing hierarchical predictions to train a single policy in THICK Dreamer seems better suited for long- horizon learning than the hierarchical policies of Director or not employing hierarchies.
3.3 ZERO-SHOT MODEL-PREDICTIVE CONTROL
Lastly, we analyze whether our hierarchical predictions are suitable for planning by comparing THICK PlaNet and PlaNet (Hafner et al., 2019b) in Multiworld. We consider the challenging setup offline dataset of IM samples collected by Plan2Explore (Sekar
of MPC for models trained et al., 2020). Figure 9 shows the zero-shot performance over training. For Pusher Pusher-Dense, i.e. a short-horizon task" with dense rewards, there is notable difference between both methods. When rewards are sparse (Pusher-Sparsel or the task horizon long (Door and PickUp). THICK PlaNet achieves higher returns than PlaNet. Additionally, the subgoals set by the high level can be decoded, shown in Suppl. F6,, which improves the explainability of the system's 's behavior.
(a)
VisuslPinPad
(b)
VisualPinPadPaa
(c)
VisualPinPsdFive
600
400
300
Directkee
400
200
Dreamer
200
200
THICK
100
Diraamer
0
0
3
2
5
10
steps
705 steps
10a steps
Figure 8: VisualPinPad results. We plot the mean evaluation returns for 7 seeds (+: standard error).
Previiousy, Hafner et al. (2022) reported that Director outperforms Dreamer in VisualPinPad. We hypoth- esize that this improvement stems from more sophisticated exploration, which is not necessary in our setting. Since the puck starts between the gripper and goal, the task be solved by directly moving to the goal.
8

---

Published as a conference paper at ICLR 2024
(a)
Pusher-Dees -ease
(b)
SeaherrSparrse
(c)
Door
(d)
PickUp
30
30-
40
20
20
PlaNet
20
THICK
10
10
0
PlaNet
0-
0
4
8
12 16
16
0
4
8
12
16
0.5
1
1.5
2
10
updates
10a updates
10
updates
109
updates
Figure 9: MPC Multiworld results. Each graphic plots the mean returns for zero-shot planning in Multiworld over world model updates using 10 seeds. Shaded areas depict the standard deviation.
4 RELATED WORK
Learning hierarchical RNNs from sparse activity was proposed in Schmidhuber (1992), where a high level would become active based on low-level errors. Subsequently, there has been a lot of research on fostering sparsity in RNNs (Graves et al., 2014; Neil et al., 2016; Goyal et al., 2021; :; Gumbsch et al., 2021; Jain et al., 2022), which we compare in Suppl. C.
Sparsity in RNNs:
Temporal abstract predictions: One main challenge for learning temporal abstractions is segment- ing a sequence into meaningful units. Discrete latent dynamics were previously used to model proac- tive gaze behavior (Gumbsch et al., 2022). Alternative segmentation methods are identifying easy- to-predict bottleneck states (Neitz et al., 2018: 8; Jayaraman et al., 2019; Zakharov et al., 2021), using fixed time scales (Saxena et al., 2021), prediction aror-based segmentation (Gumbsch al., 2019), or regularizing boundary detectors (Kim et al., 2019; Zakharov et al., 2022) (details in Suppl. C).
Hierarchical RL (HRL): HRL is orthogonal research direction to hierarchical world models.In In HRI. a high-level policy either selects a low-level policy or provides goals or rewards for a low level (Pateria al., 2021). In contrast, agent. Typically in HRL, the high level operates on fixed time scales (Hafner et al., 2022; Nachum al., 2018; Vezhnevets
THICK Dreamer uses high-level predictions to train flat RI.
al., 2017: Gürtler et al., 2021) or aas-deppndently based on subgoal completion (Bacon et al., 2017; Levy et al., 2019). In THICK world models, the high level is learned time- and task iuseenndependentl purely from predictions and latent state regularization.
CONCLUSION
We have introduced C-RsSM and THiCK-fully self-supervised methods to construct hierarchical world models. By imposing sparsity objective, C-RsSM develops context codes that update only at critical situations, where prediction-releve aspects of the environment change. On a higher level, THICK learns to anticipate context-altering states. Categorical high-level action codes enable the an- ticipation of different outcomes, accounting for multiple lower-level context transitions. As a result, THICK world models can predict both abstract context transitions and exact low-level dynamics. Additionally, we have shown that the hierarchical predictions can improve long-horizon learning.
Limitations THICK relies on setting the hyperparameter |eparse which determines the high-level segmentation. Ideally, this hyperparameter should be tuned for every task. However, we found that the same value works well across similar tasks. Furthermore, except for improving long-horizon learning our downstream applications have similar restrictions as the method they build upon. For example, if Dreamer never discovers a solution to a task, THICK cannot decompose it.
Future directions We see great potential of THICK world models as a tool to build more so- phisticated agents that explore and plan their behavior across multiple time scales. A promising direction is combining MCTS with RL (Schrittwieser et al., 2020), e.g. for biologically plausible planning (Mattar Lengyel, 2022) by searching for high-level goals that goal-condition low-level policies (Akakzia et al., 2021). Another potential lies in integrating more active epistemic-driven exploration (Sekar et al., 2020; Sancaktar et al., 2022), which could lead to a more robust consoli- dation of context codes and transitions between them. Future extensions could also explore richer predictions purely from the context Ct. This would allow the high-level to directly predict context transitions without predicting observable state information used for intermediate queries to the low- level. Lastly, while we employed THICK to establish a two-level hierarchy of world models, THICK could be applied on multiple levels recursively build N-level world model hierarchy.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We thank Karl Friston, Marco Bagatella, and Tomáš Daniš for the valuable feedback and help- ful discussions. We acknowledge the support of the German Federal Ministry of Education and Research through the Tübingen AI Center (FKZ: 01IS18039B). This research was funded by the German Research Foundation (DFG) within Priority-Program SPP 2134 project "Development of the agentive self" (BU 1335/1 1-1, EL 253/8-1). Georg Martius and Martin Butz are members of the Machine Learning Cluster of Excellence, EXC number 2064/1 Project number 390727645. We thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for sup- porting Christian Gumbsch. Noor Sajid is grateful for support from the Medical Research Council (MR/S502522/1) and the 2021-2022 Microsoft PhD Fellowship.
REPRODUCIBILLTY STATEMENT
We provide our source code on httph: https://githus: Cotpo/nttieeooileegg///titeeeeet Ad- ditionally, we specify all our hyperparameters choices, details on the conducted hyperparameter search, and advice on how tune the hyperparameters for novel task Suppl. B.
REFERENCES
Ahmed Akakzia, Cédric Colas, Pierre Yves Oudeyer, Mohamed Chetouani, and Olivier Sigaud. Grounding language to aattnonnuusly-aqquieed skills via goal generation. In International Con- ference on Learning Reprexentation, 2021.
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. In Advances in Neural Information Processing Systems, volume 30, 2017.
Pierre-Lue Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Proceedings of the AAAI conference artificial intelligence, volume 31, 2017.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: / Ntkiv:330..30 308.3432, 2013.
Matthew Botvinick and Ari Weinstein. Model-based hierarchical reinforcement learning and hu- action control. Philosophical Transactions the Royal Society B: Biological Sciences, 369 (1655):2013045)) 304480, 2014.
Martin V Butz. Toward unified sub-symbolic computational theory of cognition. Frontiers in psychology, 7:925, 2016.
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.09483 2022.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv: 112.355 2014.
Fei Deng, Junyeong Park, and Sungjin Ahn. Facing off world model backbones: RNNs, transform- ers, and $4. In Advances in Neural Information Processing Systems, 2023.
Manfred Eppe, Christian Gumbsch, Matthias Kerzel, Phuong DH Nguyen, Martin V Butz, and Stefan Wermter. Intelligent problem-solving as integrated hierarchical reinforcement learning. Nature Machine Intelligence, 4(1):11-20, 2022.
Karl Friston, Rosalyn J Moran, Yukie Nagai, Tadahiro Taniguchi, Hiroaki Gomi, and Josh Tenen- baum. World model learning and inference. Neural Networks, 2021.
Karl J Friston, Richard Rosch, Thomas Parr, Cathy Price, and Howard Bowman. Deep temporal models and active inference. Neuroscience & Biobehavioral Reviews, 90:486-501. 2018.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine. Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations, 2021.
10

---

Published as a conference paper at ICLR 2024
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.5401, 2014.
arXiv preprint
Christian Gumbsch, Martin Butz, and Georg Martius. Autonomous identification and goal- directed invocation of event-predictive behavioral primitives. IEEE Transactions Cognitive and Developmental Systems, 13(2):298-311 -,1,, June 2019.
Christian Gumbsch, Martin V Butz, and Georg Martius. Sparsely changing latent states for predic- tion and planning in partially observable domains. In Advances in Neural Information Processing Systems, volume 34, pp. 17518-17531 17531, 2021.
Christian Gumbsch, Maurits Adam, Birgit Elsner, Georg Martius, and Martin V Butz. Developing hierarchical anticipations via neural network-based event segmentation. In 2022 IEEE Interna- fional Conference on Development and Learning (ICDL), pp. 1-8. IEEE, 2022.
Nico Gürtler, Dieter Büchler, and Georg Martius. Hierarchical reinforcement learning with timed subgoals. In Advances in Neural Information Processing Systems, volume 34, pp. 21732-21743, 2021.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10122 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learn- ing behaviors by latent imagination. In International Conference Learning Representations, 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, 2019b.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Reprexentations, 2020.
Danijar Hafner, Kuang-Huei Lee, Ian Fischer, and Pieter Abbeel. Deep hierarchical planning from pixels. In Advances in Neural Information Processing Systems, volume 35, 2022.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint 2rxxii:3301041000 2023.
James B. Heald, Máté Lengyel, and Daniel M. Wolpert. Contextual inference underlies the learning of sensorimotor repertoires. Nature, 600:489 493, 2021.
James Heald, Máté Lengyel, and Daniel M. Wolpert. Contextual inference in learning and mem- ory. Trends in Cognitive Sciences, 27 (7(1):43-66 64, 2023.
Arnav Kumar Jain, Shivakanth Sujit, Shruti Joshi, Vincent Michalski, Danijar Hafner, and Samira Ebrahimi Kahou. Learning robust dynamics through variational sparse gating. In Advances in Neural Information Processing Systems, volume 35, pp. 1612-1626, 2022.
Dinesh Jayaraman, Frederik Ebert, Alexei Efros, and Sergey Levine Time-agnostic prediction: Predicting predictable video frames. In International Conference on Learning Representations, 2019.
Taesup Kim, Sungjin Ahn, and Yoshua Bengio. Variational temporal abstraction. In Advances in Neural Information Processing Systems, volume 32, 2019.
Jan Koutnik, Klaus Greff, Faustino Gomez, and Juergen Schmidhuber. A Clockwork RNN. In Eric P. Xing and Tony Jebara (eds.), International Conference on Machine Learning, pp. 1863- 1871, 22-24 24 Jun 2014.
Heinrich Küttler, Nantas Nardelli, Alexander Miller, Roberta Railcanu, Marco Selvatici, Edward Grefenstette, and Tim Rocktäschel. The nethack learning environment. In Advances in Neural Information Processing Systems, volume 33, pp. 7671-7684, 2020.
Brenden M Lake, Tomer D Ullman, Joshua Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40:e253, 2017.

---

Published as a conference paper at ICLR 2024
Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review, 62, 2022.
Tai Sing Lee and David Mumford. Hierarchical bayesian inference in the visual cortex. JOSA A, 20 (7):1434-1448, 2003.
Andrew Levy, George Konidaris, Robert Platt, and Kate Saenko. Learning multi-level hierarchies with hindsight. International Conference on Learning Representations, 2019.
Marcelo G. Mattar and Máté Lengyel. Planning in the brain. Neuron, 110(6)91149344 2022. ISSN URL https: www. sciencedirectt.
0896-6273. doi: (000:///.nerron2222122.0000 com/sciencee http:/artcceepppppeeee /oodiii/eeeattiieee///////662222
Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak. Discover- ing and achieving goals via world models. Advances in Neural Information Processing Systems, 34:24379-24338 2021.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world mod- els. In The Eleventh International Conference Learning Representations, 2023.
Ofir Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforcement learning. In Advances in Neural Information Processing Systems, volume 31, 2018.
Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Vi- sual reinforcement learning with imagined goals. In Advances in Neural Information Processing Systems, volume 31, 2018.
Daniel Neil, Michael Pfeiffer, and Shih-Chii Liu. Phased Istm: Accelerating recurrent network train- ing for long event-based sequences. In Advances In Neural Information Processing Systems, pp. 3882-3890, 2016.
Alexander Neitz, Giambattista Parascandolo, Stefan Bauer, and Bernhard Schölkopf. Adaptive skip intervals: Temporal abstraction for recurrent dynamical models. In Advances in Neural Informa- fion Processing Systems, volume 31, 2018.
Shubham Pateria, Budhitama Subagdja, Ah-hwee Tan, and Chai Quek. Hicrarchical reinforcement learning: A comprehensive survey. ACM Computing Surveys (CSUR), 54(5):1-35, 2021.
Vihang Prakash Patil, Markus Hofmarcher, Constt-Cnnstnntin Dinu, Matthias Dorfer, Patrick M Blies, Johannes Brandstetter, Jose Arjona-Medina, and Sepp Hochreiter. Align-RUDDER: Learn- ing from few demonstrations by reward redistribution. In International Conference on Learning Representations, 2021.
Vitchyr Pong, Murtaza Dalal, Steven Lin, and Ashvin V Nair. Multiworld: Multitask environments for RL, 2018. URL https: : github. com/ ht////vtttttrrrre//e/ httecchrrmmlliiiii www.wwwwwtwwd..ommoomommmmmmwmii
Vitchyr H Pong, Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine. Skew- fit: state-covering self-supervised reinforcement learning. In Proceedings of the 37th Interna- tional Conference on Machine Learning, pp. 7783-7792, 2020.
Doina Precup. Temporal abstraction in reinforcement learning. University of Massachusetts Amherst, 2000.
Gabriel A Radvansky and Jeffrey M Zacks. Event cognition. Oxford University Press, 2014.
Debbie Roberts. Thick concepts. Philosophy Compass, 8(8):677-688, 2013.
Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world mod- els are happy with 100k interactions. In International Conference on Learning Representationss 2023.
Tim Rohe and Uta Noppeney. Cortical hierarchies perform bayesian causal inference in multisensory perception. PLoS biology, 33(2)e110207) 2015.
12

---

Published as a conference paper at ICLR 2024
Nicolas P Rougier, David C Noelle, Todd S Braver, Jonathan D Cohen, and Randall CO O'RRelly. Pre- frontal cortex and flexible cogmitive control: Rules without symbols. Proceedings of the National Academy of Sciences, 102(20):7333--33) 7343333338733,, 2005.
Reuven Rubinstein. The cross-entropy method for combinatorial and continuous optimization. Methodology and computing in applied probability, 1:127-190, 1999.
Noor Sajid, Panagiotis Tigas, Alexey Zakharov, Zafeirios Fountas, and Karl Friston. Exploration and preference satisfaction trade-off in reward-free learning. In ICMI. 202/ Workshop on Unsu- pervised Reinforcement Learning, 2021.
Mikayel Samvelyan, Robert Kirk, Vitaly Kurin, Jack Parker-Holder, Minqi Jiang, Eric Hambro, Fabio Petroni, Heinrich Kuttler, Edward Grefenstette, and Tim Rocktäschel. Minihack the planet: A sandbox for open-ended reinforcement learning research. In Neural Information Processing Systems Datasets and Benchmarks Track, 2021.
Cansu Sancaktar, Sehastian Blaes, and Georg Martius, Curious exploration via structured world models yields zero-shot object manipulation. In Advances in Neural Information Processing Systems, 2022.
Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork variational autoencoders, In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 29246 29257, 2021.
Jürgen Schmidhuber. Learning complex, extended sequences using the principle of history com- pression. Neural Computation, 4(2):234-22 242, 1992.
Julian Schritwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839)6004 609, 2020.
Ramanan Sekar, Oleh Rybkin, Kostas Danilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self supervieee world models. In International Conference on Machine Learning, pp. 8583-8592. PMLR, 2020.
Richard Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3:9-44, 1988.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard S Sutton, Doina Precup, and Satinder Singh. Between MDPs and semi-MDPs: A frame- work for temporal abstraction in reinforcement learning. Artificial intelligence, 112((12):181-
211, 1999.
Momchil S Tomov, Samyukta Yagati, Agni Kumar, Wanqian Yang, and Samuel Gershman. Dis- covery of hierarchical representations for efficient planning. PLoS computational biology, 16(4): e1007594, 2020.
Harm van Seijen, Shimon Whiteson, and Leon Kester. Efficient abstraction selection in reinforce- ment learning. Computational Intelligence, 30(4))657-699, 2014.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention all you need. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30, 2017.
Alexander Sasha Vezhnevets, Simon Osindero, Tom Schaul, Nicolas Heess, Max Jaderberg, David Silver, and Koray Kavukcuoglu. Feudal networks for hicrarchical reinforcement learning. In International Conference Machine Learning, pp. 3540-3549, 2017.
Marin Vlastelica, Sebastian Blaes, Cristina Pinneri, and Georg Martius. Risk-averse zero-order trajectory optimization. In 5th Annual Conference Robot Learning, 2021.
13

---

Published as a conference paper at ICLR 2024
Alexey Zakharov, Qinghai Guo, and Zafeirios Fountas. Variational predictive routing with nested subjective timescales. arXiv preprint 10.1:2210.11136,, 2021.
Alexey Zakharov, Qinghai Guo, and Zafeirios Fountas, Long-horizon video prediction using dynamic latent hierarchy, arXiv preprini arXiv:2212.14376, 2022.
14

---

Published as a conference paper at ICLR 2024
Supplementary Material for:
Learning Hierarchical World Models with Adaptive Temporal Abstractions from Discrete Latent Dynamics
Contents
A Pseudocode
16
Hyperparameters
18
Extended Related Work
20
D THICK World Models: Implementation Details
21
D.1 THICK Design Choices
D.2
GateLORD.
21
D.3
C-RSSM Loss
22
D.4
High-level World Model Training
23
D.5
THICK Dreamer: Details
23
D.6
THICK PlaNet: Details
24
Environment Details
25
E.1 MiniHack
25
E.2 Multiworld
26
F Extended Results and Experiment Details
27
F11
Analysis of Contexts and Predictions
27
F2
MBRL in MiniHack: Experiment details and extended results ................................
29
F3
MBRL in MiniHack: Director
F.4
MBRL in VisualPinPad: Experiment details
32
F5
MBRL in VisualPinPad: Effect of "ype............................
32
E6
MPC: Experiment details and extended results
33
E7
MBRL & MPC: Scaling models vs. hierarchy
34
F8
Ablations and Hyperparameters
34
Computation and Code
37
15

---

Published as a conference paper at ICLR 2024
A PSEUDOCODE
Algorithm outlines how THICK world models make temporal abstract predictions using both levels of the hierarchy (also visualized in Fig. 4). Blue parts are only needed for MBRL or MPC (see Sec. 2.3). For temporal abstract rollouts, which are used in THICK PlaNet, the process can be repeated K times by using the output states, i.e (x()) and d ((t)* inputs again.
Algorithm THICK Temporal Abstract Prediction
I: input: context C₁ stochastic state
2: A, ~ (A, Ct2 zx)
>> sample high-level action high-level state prediction > high-level action prediction
3: 2r(1)-1 Fo (2 Acc, z1)
4: dx(1)-1 ~ Fo( A,c,, z1)
5: AF(t) Fe (Ar(t)) |A,c,, zt)
high-level time prediction > high-level reward prediction
6: (f2-1)) ~ Fe Fe(ft.r (:((1) |Acc, zi)
7: CT(1) + 9o (0((r(1)-1) Ct, (1))
low-level context
8: (0)) ~ 1 _e 7((()_1 (x(1))) -16))-1-: CT(1), (r
> low-level coarse prior
9: Foy Pr()) ~ of (r() Pi(o Cr(E), (T))
>> coarse reward & discount prediction
10: output: Cr(1), (0)) Át, f7, (1) (9(0))
Algorithm 2 describes how to create input-target data for training the high-level world model. In continual learning environments with no early termination of an episode, we omit the red part.
Algorithm 2 THICK Training Data Generation
1: input: discount factor y, sequences of contexts C1.r. stochastic states 21.7, actions a1:7,
2: rewards 71:7, and episode termination flags d1.7.
3: initialize: train data 2 <- {}. unassigned inputs I + {
4: for 7 1- do
5: if c, + C--1 or d, then
context change episode is over at time
6: for (ct, E do
compute passed time Ar + t and accumulated rewards Tur + Ed'- 7 add input-target tuple ((c,, zt), (Zr-1, ar-1 (trr))) to D 9: remove (C,, z,) from I
7:
8:
add potential input (c,, ) I
11: output: train data
Algorithm describes the general training and generation behavior THICK world models. Red parts are only used for THICK PlaNet. Blue parts are only used for THICK Dreamer. In our zero-shot planning experiments using THICK PlaNet, we do add new data to the replay buffer and only plan and execute actions during evaluation.
16

---

Published as a conference paper at ICLR 2024
Algorithm THICK World Models
initialize neural networks and replay buffer
2: (plan -I
3: for 1 gond do
4: update low-level world model state s, ~ wo(s: 81-1, a:-1)
5:
// Behavior
6:
cf # C-1 2 pplan + 7 then
7:
plan subgoal 25cal using MCTS and temporal abstract rollouts (Alg. 1)
8:
(plan {-1
9:
plan action a, using CEM given 8, and =t (Eq. 22)
10:
sample new action a, from actor given
execute action a in environment and observe rt. is and d,
12:
add (4,, Ge, de) to replay buffer
// Tra world models
14:
draw sequence batch B +- (drT, ari,, Ter,, desr) from replay buffer
15:
embed batch in latent state Sy-r ~ Wal (S B)
update low-level world model we using B (Eq. 9)
17:
generate high-level training batch 'D from (Sr.T, ae.r, Te.T, de:r, (Alg. 2)
update high-level world model Wo using D (Fq. 18)
19:
// Train actor and critic
20:
imagine trajectory (8t".H H, Qq":H Te":Hs )H.H) using wi from random start 8:" make temporal abstract predictions for each 8 using We and ws (Alg. 1)
21:
22:
compute value V (Eq. 20)
23:
update critics Ux and Vy (Eq. 33)
24:
update actor æ
17

---

Published as a conference paper at ICLR 2024
B HYPERPARAMETERS
Table I: Hyperparameter choices. If there is only centered value it counts for all suites. Oth- erwise different values are chosen for MiniHack (MH), VisualPinPad (VPP), or Mulitworld (MW).
Name
Value VPP
MH
MW
Low-Level World Model (C-RsSM or RSSM)
Batches (size x sequence length)
x
Dimensions of C₁
16
Dimensions h2
256
Dimensions of 2,
32 32
MLP features per layer Sparsity loss scale feparse Prediction loss scale 33red
256
19/10
1
KL loss scale 8KL.
1
KL balancing ghal
0.8
Output heads og for Prioritize ends in replay
ig Tt. Ft yes
ie. 7t
it. Ft
n10
no
Learning rate
0.0001
High-Level World Model
Q9 & Pa number of layers x features Fg number of layers features
x 200
1024
Number of actions A,
3 yes
5
5 no
Use terminations d, for segmentation
no
Loss for training F (t) At, Ce z1)
CCE
CCE
NLL
Action prediction loss scale a=+(()-1
1
0.1
State prediction loss scale -1()))
Time prediction loss scale Ar(L)
I
0.1
Reward prediction loss scale a' For())
1
KL balancing a bal
0.8
Learning rate
0.0001
THICK Dreamer
Imagination horizon
15
Value estimate balance x/
0.9
0.9
X-target of V.A
0.95
0.95
Long-horizon critic Ux layers features
4x 400
4x 400
Long-horizon critic Uy learning rate
0.0002
0,0002
THICK PlaNet
CEM planning horizon H
12
Long-horizon scale k
0.025
MCTS simulations
100
MCTS discount
0.997
Common
Optimizer
Adam
MI. activation functions
ELU
Discount 7
0.99
'Unlike the other MiniHack problems, nf Hack-Corrd dor tasks are fully deterministie, which is why lower factor of gaparse 1 here.
18

---

Published as a conference paper at ICLR 2024
World model learning hyperparrmeeters For optimizing the world models, i.e. our THICK world models and the baseline models in Dreamer and Director, we use the default DreamerV2 hyperpa- rameters (Hafner et al., 2020) except for minor variations. Specifically, we reduced the model size by setting the feature size of the RSSM and the dimensionality of h, to 256. As we show in Suppl. F.7 model size does not strongly affect performance in our setting. Additionally, for THICK Dreamer and Dreamer we did not employ layer normalization for the GRU within the RSSM, be- cause in pre-tests this showed increased robustness for both approaches.
MBRL hyperparameters For training the actor and critic in THICK Dreamer and Dreamer, we use the default hyperparameters DreamerV2. For Director we mostly used default hyperpa- rameters (Hafner et al., 2022), however, made some minor adjustments to the training frequency to ensure a fair comparison. Director performs one training update every 16 policy steps instead of every 5 steps in PreamerV2. This was done to reduce wall-clock time but decreases sample effi- ciency (Hafner et al., 2022). We increase the update frequency (16 - 5) in order to fairly compare sample efficiency between approaches.
MPC hyperparameters For MPC with CEM we use the hyperparameters of PlaNet (Hafner et al., 2019b). For high-level planning with MCTS, we use MuZero's (Schrittwieser al., 2020) imple- mentation, with mostly the same hyperparameters. However, intuitively we would not expect multi- ple predictions to reach a goal. Thus, we decrease the number of simulations to 100.
Differences between environments The main difference between MiniHack and the other envi- ronments is that in MiniHack episodes can terminate based on the success of the task or the death of the agent. VisualPinPPd and Multiworld do not feature early episode termination. As is customary with DreamerV2, for environments that do not feature early episode termination, we do not pre- dict discounts Tt, nor do we prioritize the termination of episodes in the replay buffer. Importantly, we do not treat episode terminations as context changes. For action prediction, we use Categorical Cross Entropy Loss (CCE) for predicting discrete actions (Minihack and VisualPinPad),, and scale down the high-level prediction loss for predicting actions and elapsed time when training purely on task-free offline data (Multiworld)). Finally, the sparsity loss scale geparse was tuned for each suite.
Hyperparameter search For determining the sparsity loss scale geparse, the value timate balance v. and the long-horizon planning scale K, we ran a grid search using three random seeds and using tasks of each suite (MiniHack: KeyRcom-Fi Fied---5xx KeyCorridor dor-4, KeyCorridor- 8, Vi- vi suaLP intdiinPaaFFFF ve; Multiworld: Door, PickUp). We determined the best hyperparameter value for each suite depending on task performance and a qual- itative inspection of the high-level predictions (see Suppl. F.8). For simplicity and demonstrate robustness, we used the same values for each suite, whenever it was reasonable.
WandOfDeath- aer.-eaaah-vvvnneess. Advances; MiniHack-Corridae sual Pin Pad: Vi sua wtw.tlPFPPaaooooooo
How to tune When tuning THICK world models for new task, we recommend mainly searching over the sparsity loss scale gepars E (1,5, 10, 25, 50}. Typically, one random seed is sufficient to determine which geparse leads to few, but not too few, context changes.
19

---

Published as a conference paper at ICLR 2024
C EXTENDED RELATED WORK
Sparsity in RNNs: The development of hierarchical RNNs based on sparse activity was already proposed in the 90s by Jürgen Schmidhuber (Schmidhuber, 1992). The Neural History Compres- sor (Schmidhuber, 1992) uses hierarchical stack of RNNs that autoregressively predict the next inputs. The higher levels in the hierarchy remain inactive until the lower level fails to predict the next imput. Recently, there has been increasing interest in regularizing RNNs towards sparse latent updates. Alternative approaches to Lo regularization of GateLORD (Gumbsch et al., 2021) are to use sparse attention masks (Graves et al., 2014; Goyal et al., 2021), competition among submod- ules (Goyal et al., 2021); regularizing update gates towards a variational prior (Jain et al., 2022), or time-dependent updates (Koutnik et al., 2014; Neil et al., 2016).
Temporal abstractions from regularized latent dynamics: Previously, sparse changes in the la- tent states of a low-level model have been used to model temporal abstractions (Gumbsch et al., 2022: Saxena et al., 2021). In contrast to our work, the temporal abstractions in Gumbsch et al. (2022) were learned in much simpler settings with highly structured observations, instead of the high-dimensional, pixel-based observations examined in this work. Additionally, these temporal abstractions were only used to model goal-anticipatory gaze behavior of infants and have not been applied for MPC or MBRL. Separately, Saxena et al. (2021) introduced a hierarchical video pre- diction model (i.e. without action) that used different clock speeds at each level to learn long-term dependencies using pixel-based input. Although this was apt for learning slow-moving content at higher levels of the temporal hierarchy, unlike C-RSSM and THICK, it requires the temporal abstrac- tion factor to be defined explicitly.
Temporal abstractions from predictability: Adaptive Skip Intervals (ASI) (Neitz et al., 2018) is a method for learning temporal abstract autoregressive predictions. In ASI, a network is trained to predict those inputs within predefined horizon, that best allow predicting extended sequences into the future. As result, the model learns to skip number of inputs towards predictable transitions. Similarly, comporal-agnosti predictions (TAP) (Jayaraman et al., 2019) identify frames of a video within time horizon that are highly predictable. TAP is then trained to only predict the predictable "bottleneck" frames. Zakharov et al. (2021) provide a learning-free free mechanism to detect context change by evaluating how predictable future states are. Briefly, their approach detects changes in the latent representation of each layer in the model hierarchy and introduces temporal abstraction by blocking bottom-up information propagation between different contexts. This is different from THICK, where context changes are determined using a learning- based sparsity regularization. An opposing approach is use unexpected prediction errors of forward model for self-supervised time series segmentation (Gumbsch et al., 2019). Here, the idea is that in certain states, the dynamics of agent-environment interactions change, e.g. changing the terrain during locomotion, leading to a temporary increase in the prediction error.
Temporal abstractions from learning boundary detectors: In addition to using indirect mea- sures to segment sequence, straightffrrwrd approach is to train a boundary detector that signals the boundary of subsequences (Kim et al., 2019; Zakharov et al., 2022). Kim et al. (2019) train a boundary detector that is regularized by specifying the maximum number of subsequences allowed and their maximum length. This requires prior knowledge about the training data and imposes hard constraints the time scales of the learned temporal abstractions. Our sparsity loss instead implements a soft constraint. Conversely, Zakharov et al. (2022) introduced a boundary detection mechanism using non-parametric posterior over the latent states. Here, the model learns to tran- sition between states only if a change in the represented features had been observed - otherwise temporally persistent states were clustered together.
Faster credit assignment in RL: THICK Dreamer predicts long horizon outcomes via its high-level model when training the critic in imagination boost reward propagation. This allows for faster credit assignment for tasks with sparse or delayed rewards. Previously, this was tackled using reward redistribution (Patil al., 2021).
20

---

Published as a conference paper at ICLR 2024
D THICK WORLD MODELS: IMPLEMENTATION DETAILS
D.1 THICK DESIGN CHOICES
Due to space constraints, we explain some design choices in more detail here.
High-level targets Our goal is to learn high-level world model that predict situations in which latent generative factors are assumed to change, e.g. a door openings or object manipulations. Be- sides that, we want use high-level outputs to predict future rewards and reconstruct images some time (t). Thus, we at least need the context Cr(t) and the stochastic state =+(() to make these reconstructions through the coarse processing pathway (see Eq. 8). There are two potential ways to predict (r()) and (r))), either by predicting the state before or after the context transition.
We predict the states before the context transition. Our main reasoning is that the prediction of context-altering situations presents two challenges: i) learning in which situation such transitions are currently possible and ti) how these transitions affect the latent generative factors. The C-RssM already learns to encode ii). Thus, in order to reduce redundancy and simplify the challenge, we train the high-level model learn only 1) and then prompt the low-level model for ii). One example from MiniHack would be predicting the agent standing in front of closed door and performing door-opening action. We believe that this is a simpler prediction compared to predicting the ego- centric view of the agent after opening door and looking into (potentially unknown) new room.
Coarse predictions for contextual learning We want the context c to encode latent information that is necessary for prediction and reconstruction. If we omit the coarse prior predictions (Eq. 5) and coarse output reconstructions (Eq. 8) the C-RSSM would not have any incentive to encode prediction-relevant latent information in C, Instead, it could purely utilize h, and avoid a sparsity penalty in Eq. 9 via C by never updating Ct. Completely omitting h in C-RSSM impedes learning, as we show in our ablations Suppl. F.8. Thus, we instead add the coarse processing pathway. Through coarse predictions, C-RSSM must encode latent factors in C in order to reduce the KL-loss (Eq. 31) and prediction loss (Eq. 30).
Coarse predictions to omit h2 The high-level model attempts to predict future state of the system. The full latent state would contain the deterministic component h. However, for the high- level model it would be very challenging to predict the unregularized nigh-dimensional deterministic hidden state h many time steps in the future. The coarse pathway of the C-RSSM allows to update the context dynamics Ct, predict stochastic states zf, and reconstruct external variables without the deterministic hidden state h, Thus, it is advantageous that the C-RSSM can make predictions without h. After a high-level prediction, we can feed the high-level outputs (É+(1)-1, (+(0)-1 into the lewelevel world model. This brings many advantages: for example, this allows to predict rewards or discounts/episode termination at particular states after a high-level prediction, which we use in THICK Dreamer in and THICK PlaNet (see Sec. 2.3). Furthermore, to visualize predictions as shown in Sec. 3.1. Additionally, we can continue with low-level rollouts after a high-level prediction, which feature that we have not yet utilized.
reconstruct images
D.2 GATELORD
We want the context code c, change only sparsely over time. Thus, we implement the discrete context dynamics 9o as a GateLORD cell (Gumbsch et al., 2021). GateLORD is an RNN designed to maintain sparsely changing latent states et- To realize this inductive bias, GateLoRD uses two subnetworks 5% and og that control cu-updates via an internal update gate A,, GateLORD can be summarized as follows:
Candidate proposal:
é = 5 (2 (1--1,)--,,2--1)
(24)
Update gate: A, = 3 (a (4,,
(25)
Context Update: = o + (1 - (1) C
(26)
with o denoting the Hadamard product. We use the action a and the last stochastic state x, as the cell inputs. Based on this cell input and the last context c--1, GateL.ORD proposes new context
21

---

Published as a conference paper at ICLR 2024
C2 via its proposal subnetwork o (Eq. 24). Whether the context is updated depends update gate A, 10 (Eq. 26). This update gate A, is the output of the gating subnetwork & (Eq. 25) which uses rectified tanh activation function (ReTanh), with ReTanh(()) := max(0, tanh(x)). This ensures that the gate activations are E [0, 10 Note that to compute A,, the subnetwork & internally samples from a Gaussian distribution before applying the ReTTanh function. This was shown improve robustness (Gumbsch al., 2021). Thus, context updates are a stochastic process.
Originally (Gumbsch et al., 2021), GateLORD used a subnetwork to compute the cell output using multiplicative gating. We omit this here and instead feed the context directly to the GRU cell fa as shown in Fig. (right).
The centralized gate A, of GateLoRD makes it easy determine the changes in context, i.e. C2 * Ct-1- Since all the context updates depend on As, we know that the context changed if Ar > 0. This is an advantage over other RNNs that use multiple gates for sparsely changing latent states. We this measure to determine context changes when building the world model hierarchy.
D.3 C-RSSM LOss
The loss of the C-RSSM (Eq. 9) is composed of three parts: the prediction loss _pred, the KL loss CKL, and the sparsity loss Esparse Except for the sparsity loss, we adapt these loss terms from the RSSM. However, always need to account for the coarse prediction pathways of the C-RssM.
We define the prediction loss Corod as
T Σ t=1
Cered
Σ
- log oolv st) log og ce, 20)
(27)
T
pE(t,reema)
Equations in red are exclusive to the C-RSSM. Thus, the network trained to minimize the negative log likelihood for predicting the images i,, rewards т, and future discounts ?r. Here we account for both the precise predictions over the output heads 0g (Eq. 7), as well as for the coarse predictions over the output heads og (Eq. 8). Following the codebase of DreamerV2 (Hafner al., 2020), in continual learning environments when there is early episode termination, do not predict the discount 7t. and instead use a fixed discount 7 = 0.99.
The C-RsSM predicts two prior distributions for the next stochastic state é, fine predictions using the full state (Eq. 4) and coarse predictions based only on the context, last action, and stochastic state (Eq. 5). We need to account for both types of prediction in the KL loss LKL with
T M KL ao az( he. (.) 1r)
CKL (K)
h.) +KL go (z) h1, ( a. Cts 111
(28)
Thus, we want minimize the divergence between the fine prior p/ and the approximate posterior Jo well the divergence between the coarse prior p&, and qø- As in DreamerV2 (Hafner al., 2020), we use KL--balancing, which scales the prior Pe of each KI. divergence by a factor gbal = 0.8, and the posterior 9s by ghall This enables faster learning the prior to avoid that the posterior regularized towards an untrained prior.
We take the sparsity loss Lsparse from GateLORRD (Gumbsch et al., 2021), which is a Tap- regularization of the context changes Ac,. This is implemented
Caparso (6)
Σ
T
A4
1
T ΣI ((M)
(29)
TJ
T ==
where J is the dimensionality of the context c and (.) denotes the Heaviside step function. That is, an Le-relulaiizationn of the context changes is implemented as the binarization of the update gates A, (Eq.26), We estimate the gradient of the Heaviside step function using the straight straight-through estimator (Bengio al., 2013). The advantage of GateLORDDss Ie-reuularizztion to other regularization toward
22

---

Published as a conference paper at ICLR 2024
sparsity. such as using low variational prior (Jain et al., 2022), is that in fully-observable and highly predictable situations, GateLORD will shut its gates and not change the context almost regardless of the order of magnitude of geparse to avoid punishment.
D.4 HIGH-LEVEL WORLD MODEL TRAINING
The high-level world model with parameters 0 is trained to minimize both the prediction loss g of predicting the next context change state, and the action loss 24, which is the divergence of prior and posterior high-level action distributions.
For every high-level larget
(100))) (e)))) we use an appropriate prediction loss and employ a weighted sum in gpred. For the to-be-predicted action (r(()-1- passed time (r(t), and rewards r (t)) use the negative log-likelihood (NLL). For stochastic state prediction -+
1, we know the underlying distribution that generated the target variable, i.e. the posterior (-). Thus,
use the KL divergence loss for high-level state predictions. Overall, we get
gpred (0)
1 Σl T
Σ
-a log FolY Ar, Ca, ze)
YE(arw (.)(()))
+ a a'KLL 9o (Z+(0-11
(2+(0-1 A,, Ct. 2x)
(30)
a² The hyperparameters a E fa",a"" aa a"")))"))))) can be used to scale the individual loss terms. By default, we set a = for all loss terms. When training the network task-free exploration, during zero-shot MPC as described in Sec. 3.3, we found that predicting actions (() at context
i.e.
changes and elapsed time Ar(t) was challenging. To mitigate this, during task-free exploration we set a 0.1 and Arr(t) = 0.1. For predicting continuous actions we sample from a Gaussian distribution of predicted actions and compute the NLL the loss for action prediction. For discrete actions we predict a Categorical distribution from which we sample the actions, and compute the Categorical Cross Entropy (CCE) loss.
The action loss Q drives the system to minimize the divergence between the posterior high-level action distribution Qo(Ac Ct, Zt, Cr(()) Zr(0)). and the prior distribution Po Ct, 21) with
1 T
L (0
- T 1=1
Σ
KL Qo(Ac
Cr, Re, Cr((): R+(1)) ||
Po c,, C,, k0)
(31)
Like the KL loss [KL on the low level (see Suppl. D.3), we use KL balancing (Hafner et al., 2020) to scale the prior part by a 0.8 and the posterior part by - bal
D.5 THICK DREAMER: DETAILS
THICK Dreamer estimates the overall value V(s,) of a state 8, as a mixture of short- and long- horizon estimates (Eq. 20). The short-horizon value estimate is computed the general X-target as in DreamerV2 with
[11 A)ve (B+4)) for t < H, for L = H
(*(s))= Fr + 9e f (81+1)
(32)
where f. and & are sampled from the output heads og given s; (Eq. 7) and is a hyperparameter. THICK Dreamer trains both critics ve and ex to regress the overall value estimate using squared loss:
H
((0) = Ep Σ 1 (vo(s)) sg((v()))))
(33)
for the two critics uo (ug, ) with parameters (E. x}. and sg(+) the stop gradient operator.
The functions V and ylong compute value targets using the critics V/ and ux respectively. Like DreamerV2. we stabilize critic training by using a copy of the critics during value estimation (in Eq. 32 and Eq. 19). The copy is updated every 100 updates.
23

---

Published as a conference paper at ICLR 2024
D.6 THICK PLANET: DETAILS
For planning at the high level, we use an MCTS implementation based on MuZero (Schrittwieser al., 2020). We only replan at the high level if the context changes, i.e. Cr f Ct- I. Since all sub- goals 2: are situations that lead context changes, no additional criterion for subgoal completion is needed. Upon reaching subgoal, e.g. touching an object, the context can sometimes change for multiple subsequent time steps. This causes the high-level to replan multiple times in a row. To avoid high computational load from replanning and to enable smoother trajectories, we inhibit replanning for = time steps after setting new subgoal. While this could potentially degrade performance in dynamic environments, found this to work well in Multiworld.
24

---

Published as a conference paper at ICLR 2024
(a) NaterCrossiig-RRig (b) XeyRoon-Fixed-S
(c) (ttnddDDDttaaaaaaeeeee
(d) Ri ver
Figure 10: MiniHack environments. Staircases with an upward facing arrows mark the starting point of the agents. Staircases with downward facing arrows are the exits that need to be reached. In (a), (b), and (d) start points and exits are randomized.
E ENVIRONMENT DETAILS
E. MINIHACK
Here we provide detailed explanation of all the MiniHack problems we considered. In all set- tings, we restricted the action space to the minimum number of actions needed to solve the task. In all tasks, the agents receive a sparse reward of I when exiting the room and small punish- ment of -0.01 when performing an action that has no effect, egg. moving against a wall. easier tasks SaaterCrossigg-ie Ring, GeyRoom-Fixed-22 200 time steps to solve the task. In all other tasks, the time limit is set to 400 time steps. For aesthetic reasons, we different characters in different levels.
KeyCorr (e))orridor) the agent is allowed
is a newly designed, simple level in which an agent needs to fetch randomly placed ring of levitation and float When a ring is picked up in our tasks, it is automatically worn". The level is inspired by LavaCross Levit tate-Ring deadly lava blocks the exit. However, we found that Dreamer struggles to learn this task, because of the early terminations when entering the lava.
MaterCrossiin-Rinn
river get the goal (Fig. 10a).
from the MiniHack benchmark suite, where river of
KeyRoom- F Fixed-S55 is a benchmark task, in which an agent spawns in a room at random position and has fetch a randomly placed key open a door and enter a smaller room with randomly located exit (Fig. 10b). The door position is fixed. In all our tasks, using the key opens the door from any grid cell adjacent to the door, even diagonally.
KeyCorridos Nor-N is a novel task, in which an agent starts in front of a locked door in the top left corner of a 2-grid-wide corridor. the bottom right corner of the corridor is the key for the door. We vary the length N of the corridor to systematically manipulate the task horizon.
is based on the Wandof benchmark tasks, in which an exit is guarded by minotaur, which instantly kills the agent upon contact. The agent needs to pick up wand to attack and kill the monster. Thus, the agent needs to carefully select the direction of the attack, because if the attack bounces off wall, it kills the agent instead. WandOfDea th comes in multiple levels of difficulty.
(Fig. 10c) is a self-created level layout, designed to be more challenging than WandO - -Medi um but not difficult as WandOfDeath fDetth Ha rd. In SandofDeathh-eett um the agent can only walk horizontally and the location of the wand is fixed. In WandofDeath-Ha- rd the map very large, making this hard
UUually. to wear a ring in MiniHack sequence of actions needs to be performed: PUTON + RING + R LIGHT for putting the ring the right finger. We simplify this by automatically applying the action sequence when the ring is picked up.
25

---

Published as a conference paper at ICLR 2024
(a) Pusher
(b) Door
(c) PickUp
Figure 11: Multiworld environments. Goal positions for the objects are shown in yellow.
exploration problem. Our version is of intermediate difficulty, where the number of accessible grids (28) is roughly the same in WandOfDeath: - (27), while the randomly placed wand needs to be found first. The minotur is asleep as in MandOfDeath-Ment
River benchmark task, in which an agent needs get to an exit on the other side of river (Fig. 10d). In order to cross the river the agent needs to push boulders into the water to form a small land bridge. To solve the task the agent needs to move at least two randomly placed boulders into the river.
EscapeRoom is difficult new problem designed by us, which combines the challenges of many other problems (Fig. 14a). Using EscapeRoom we test the ability to learn to execute complex event sequence of five subgoals. Nonetheless, the task can be learned without extensive exploration or large action spaces. The agent starts in small room and the goal is to unlock door and escape. However, to get the key, the agent must (1.) pick up a ring of levitation and (2.) float over a small patch of water into corridor. In the corridor, the agent a key. In order to get back to the door in the first room, the agent needs to (4.) push a boulder into water. Finally, the agent can (5.) unlock the door and exit the room. During levitation, the agent is too light to push the boulder. In EscapeRoon, the agent can only carry one new item and picking up second item results in dropping the first one.
(3.) exchange the ring of levitation for
E.2 MULTIWORLD
In Multiworld (Nair et al., 2018; Pong et al., 2020). All tasks Multiworld use different action spaces and camera viewpoints for their pixel-based observation, shown in Fig. 11. In Pusher the 2-dimensional ac- tions control the x- and y y-movement robot has hook instead of a gripper at its endeffector and the 3-dimensional action controls r-, y-, and = =-movement. In P PickUp the 3-dimensional action controls the y- and =-movement and the gripper opening. We binarized the gripper opening prevent accidental object drops. In all tasks, the goal positions are fixed. In Pusher and PickUp they are visible in the video frames. In Door the goal is to open the door fully. For Pusher-Denc and PickUp we compute the reward for every time step as
tasks that have previously been used to study visual reinforcement learning
the endeffiector, whereas the gripper fixed. In Door the
5
r1
(34)
where & is the Euclidean distance between object and goal time t. For Pusher-Sparse the agent received reward of r, when the euclidean distance between the puck and the goal 5, < 0.025, otherwise = O. For Door the reward r, is the current angle of the door joint.
26

---

Published as a conference paper at ICLR 2024
KeyRoom MiRRoo--Fiee---- Fixed-85
Door
LL:
+ 10
+ 15
Figure 12: Low-level predictions. The low-level predictions accurately predict the next frames (cf. Fig. 5 for the high-level predictions and contexts C₁ of the same sequences).
F EXTENDED RESULTS AND EXPERIMENT DETAILS
F.1 ANALYSIS OF CONTEXTS AND PREDICTIONS
In this section, we provide further examples of high- and low-level predictions and context codes C--
C-RSsM predictions and contexts Figure 12 visualizes the low-level predictions for two example sequences. The low-level world model predicts the immediate next state and the reconstructions are more accurate than the abstract high-level predictions (cf. Fig. 5). Figure displays four example sequences with the corresponding contexts and high-level predictions.
High-level actions We analyze the high-level actions A, in more detail for the EscapeRoom problem. EscapeRoom a challenging MiniHack level, designed to contain diverse agent- environment interactions, shown in Fig. 14a and described in detail in Suppl. E. To illustrate the emerging high-level action representations THICK Dreamer, we show inputs é, and image recon-
WaterCroseing- -Ring
Musher-Densee
Cc+1
(0-11
L+3
L+5 5
L+7
I I+2 PecualPinPaFFie
( (+
River
G:41
HL:
700-1
2+2
t 11
8+6
Figure 13: Context changes and high-level predictions. We show the input images é, 16-dim. contexts C++1 and reconstructions &c ( of high-level predictions. For WaterCrossing-Rinn the context changes when stepping on the ring, picking up, or arriving on the other side of the shore. In Pusher the context changes when the robot moves the puck. In River the context changes when pushing a boulder into water. VisualP Minallinne. adF ive the context changes when stepping pad.
27

---

Published as a conference paper at ICLR 2024
(a)
FEE C
EscapeRoor
A
(x(())1
PICK-UP
DOWN
RIGHT
LEFT
LEFT
LEPT
79)
AQ
(x(t)-1
LEFT
UP
UP
DOWN
RIGHT
LEFT
A
UP
DOWN
DOWN
DOWN
DOWN
DOWN
6
10
22
26
Figure 14: High-level action predictions. (a) In the Es capeRoom problem an agent needs to pick up ring of levitation, hover over patch of water to get to key, exchange the ring for key, push a boulder into the water, and use the key to unlock door. (b) Visualization of high-level actions for one exemplary sequence. The upper row shows the input image i, Image reconstructions & (0)-1 and low-level action predictions a-(1)-1 are shown for the three high-level actions A,. Red outlines depict which action A, was sampled.
structions (- and predicted low-level actions a-(1)-1 for all high-level actions A, in Fig. 14b for one exemplary sequence.
At specific time steps, the three possible high-level actions A, encode particular agent-enviromment interactions: A encodes picking up the ring of levitation (1 = or exiting the level (22, 26)). A encodes crossing the water after obtaining the ability to levitate (t E {6, 10, 16)). either upward E {6. 10)) or downward
16). A encodes pushing the boulder into water (t E {6, 10, 16)) or stepping in front of the door (t = 22). For all other time steps, high-level actions produce either identity predictions (e.g. A!) or predictions that seem to encode average scene settings (cf. A3 or A 1o). These predictions account for unexpected context shifts, which can always occur with a small probability due to the stochasticity from sampling 2; and the stochastic update gates of GateLORD (see Suppl. D.2). The low-level actions predicted for these situations seem to be mostly random. The
(a)
%6 with e o 1 MiniBack
(b)
%t with # C Multivorld
(c)
%t with # Q-1 in VisualPinPad
100
100
WaterCross CraeeCCrre--nnnggg
Tueher-Spazce
Three
60
KeyBoon- Fized--5
80
Door
80
Four
WandDfDeath-
40
RacapaRoon
60
PickUp
60
Five
River
20
20
20
0
0
0
0
2
2
0
2
6
100 steps
10s updates
10
steps
Figure 15: Context changes over training. Each graphic plots the mean percentage of time steps per training batch for which the context e changes in MiniHack (a), VisualPinPad (b), and Multi- world (c). Shaded areas depict the standard deviation.
28

---

Published as a conference paper at ICLR 2024
(b)
fiequenc frequency
4
VaterCroes Ring
frequency
WandOrDeal DaattrrDaat-aaaaa
2
6
8
10
12
14
5
10
15
20
At
A
EscapeRocm
frequency
River
5
10
20
5
10
15
20
25
A
At
Figure 16: Context duration during an episode. Each graphic plots histogram of the mean number of time steps At between two consecutive context changes during an episode over 10 episodes, max 50 steps) for different MiniHack tasks seeds). Shaded areas depict the standard deviation.
prior Qo (red frames and text in Fig. 14b) typically samples reasonable high-level actions. However, occasionally the prior samples action Å, that leads to an identity average prediction (e.g.
t = 6) due to the randomness of the process.
Quantifying context changes To quantify the changes in context, we plot the mean percentage of time steps when context changes occur (i.e. C₁ # Ct-1) the course of training in Fig. 15. Importantly, context changes are somewhat consistent within the same task, but, as expected, can vary between tasks in the we analyze the time between context changes for some MiniHack tasks. We plot the histogram of time gaps between context changes Fig. 16 which illustrates that different tasks also show different distributions of context durations.
suite despite using the same hyperparameter Jeparse Additionally,
Task-relevance of contexts Lastly, we analyze whether context changes occur in task-relevant situations for some MiniHack problems. For this, we generate rollouts using the fully trained policy and identify points t", which we consider to be crucial for solving the task. For points t" before picking WaterCross-Ring, and XeyRoom-Fixed-e item. For EscapeRoom, we use points in time t" when the agent stands in front of movable boulder blocking the exit path. We compute the mean percentage of context changes occurring around t (#1 step) over 10 sequences and take the average over all randomly seeded models. The results are shown in Table 2. The C-RsSM tends to update context with high probability in the identified situations. This suggests that task-relevant aspects, such as item pickups or boulder pushes, are encoded in the contexts.
we take the time
Table 2: Task-relevance of context changes. We list the mean percentage of context changes Ct. * C+-11 for fully trained policies (7 seeds) at crucial task relevant points in time in 10 sequences. See text for criterion of t*. The standard deviation is denoted by +.
WaterCross (7.1
KeyRoom-Fixed-SS Fanthheeath-Ady
EscaveRcom
91.4
91.4(11.5)
88.57 (E 15.7)
F.2 MBRL IN MINIHACK: EXPERIMENT DETAILS AND EXTENDED RESULTS
In Fig. 17 we plot the success rate of THICK Dreamer, DreamerV2, and Director for additional MiniHack tasks not shown in the main paper.
To investigate the effect of task horizon, we compare the performance gain of THICK Dreamer over Dreamer for different corridor lengths in KeyCorridor. To compute improvements for every seed, we subtract the success rate and returns Dreamer of THICK Dreamer (visualized in Fig. 7 7g). For corridor lengths of onward, the improvements of THICK Dreamer Dreamer tend
29

---

Published as a conference paper at ICLR 2024
(a)
XeyCorridor-4
(b)
KeyCorrider-6
(c)
KeyCorridor-6
Director
Dreamer
0.5
0.5
0.5
THICK
Dreamer
2.5
2.5
10a steps
10a steps
10 steps
(d)
KeyCorridor-1
(c)
KeyCorridor-9
(
KeyCorridor 100
1
0.5
0
0
2.5
2.5
10 steps
10 steps
10 steps
(g)
KeyCorridor-11
(h)
WaterCross, -Ring
i)
0.5
0.5
0.5
0
2.5
10 steps
10 steps
10 steps
Figure 17: MiniHack success. We plot the mean evaluation success rate (7 seeds, + standard error).
(a)
KeyCorridor-5
(b)
KeyCorrider-6
KeyCorridor-7
Director
Dreamer
0
THICK
-1
Dreamer
2.5
5
2.5
5
2.5 10 steps
10a steps
10 steps
(d)
KeyCorridor-8
(e)
KeyCorrider-9
(I)
KeyCorridor-10
-1
-1
2.5
5
2.5
5
2.5 10 steps
10a steps
10 steps
(g)
KeyCorridor-11
(h)
VaterCrosa-Rlng
()
River
0
-
-2
2.5
5
10a steps
10a steps
10 steps
())
XeyRoon-F Fixeee-----
(k)
WandOrDeath-Aee
EacapeRoom
0
-2
-2
6
2
10a steps
105
steps
10a steps
Figure 18: MiniHack returns. We plot the mean evaluation returns (7 seeds, + standard error).
30

---

Published as a conference paper at ICLR 2024
(a)
(b)
(c)
Director: success rate
Director: returns
KeyRoom: random
KayRoon: random kny
0.5
XayRoom: fixed
0
ZayCorridor-4: fixed
0
0.5
1,5
0.5
1.5
10a steps
10a steps
Director in KeyDoor environments. We test Director in variants of the KeyRoom-FF xed-S5 problem, with fixed spawn positions of agent, key, and goal (i.e. KeyRoom: fixed, shown in a), random spawn points of the key (i.e. KeyRoom: random key), fully randomized spawn points (KeyRoom: random) or the similar KeyCorridor-- 4 problem. We plot the mean evaluation success rate (b) and the mean evaluation returns (c) (7 seeds for KeyRoom: random, seeds otherwise, + standard error).
Figure 19:
to increase with corridor length. However, for a corridor length of 11 most runs fail to discover the reward (see Fig. 17g), which dampens the improvement in performance. The inability solve these tasks seems to come from inadequate exploration. Our results in VisualPinPad indicate that if exploration is addressed, then the performance gain of THICK Dreamer also holds for longer task horizons.
F.3 MBRL IN MINIHACK: DIRECTOR
Director (Hafner et al., 2022) shows strong performance in the KeyRoom-Fixeds task. How- ever, Director does not learn the other MiniHack tasks that we considered Director (Fig. 7). Which crucial aspects are different across tasks and what causes Director fail? We identify two key prob- lems when applying Director in MiniHack, namely i) Director's reliance on diverse initial data and éi) problems with specifying unobservable goal information.
Diversity of initial data Director trains goal encoder on its replay buffer from which samples when training goal-conditioned policy. We hypothesize that if early in training not enough diverse data collected, this is reflected in the goal space. As a result, the manager (high-level) does not set meaningful goals for the worker (low-level) and learning is severely slowed down.
We analyze this aspect by training Director on variants of the KeyRoom-FF Fienn-iiee-SSeeee s5 problem. By default, the initial positions of the agent, key, and goal within the second room random- ized in KeyRoom- -iieo-----x......- We create additional variants of the task in which all entities are spawned in fixed positions (positions shown in Fig. 19a) only the initial position of the key is randomized. Additionally, we train Director in the KeyCorridor dor-4 task, which is very similar to KeyRoom Fixed-S5 S5 with fixed spawn points but of much smaller size (8 grid corridor vs. two rooms with 16 and 4 grids). Thus, in KeyCorridor the observations show little diversity.
Figure 19b plots evaluation success rates of Director in the KeyRocm-variants. Director needs more steps to solve the tasks when entities spawn at fixed positions. Director does not learn consistently learns solve Note that, KeyCorridor 4 is much smaller and has a shorter task hori- zon. A similar trend can be observed in the collected returns (Fig. 19c).
solve KeyCorridor-4 whereas with the same training KeyRoom-Fixed
Thus, we conclude that diversity in the initial observations drastically boosts Director's per- formance. The ego-centric views of MiniHack often contain the same similar observa- tions, especially when traversing long corridors or empty rooms, e.g. in KeyCorridor-8 or WandOfDeath- Adwatthddddd.d....e This similarity observations might impede Director's learning in the MiniHack tasks we considered here.
Unobservable aspects of goals fying unobservable information in the goals. The RSSM encodes both observable and unobservable information within its deterministic latent state h, If the unobservable information, e.g. item pick-
We hypothesize that a severe problem of Director could be speci-
31

---

Published as a conference paper at ICLR 2024
(a)
(b)
VisualPinPadIuree
VieuxIPiaPeddTures
VisualPinPadFour
(no exploration steps)
(2.5 10" exploration steps)
(2.5 10 exploration steps)
600
600
400
400
Dreamee
400
200
200
THICK
200
100
Dreamer
0
3
2
10 steps
10' stope
10 steps
(d)
VioussPsnPaaTTrre
(e)
Vimus1PinPPadoou
VizualPinPadFive
(5 10 exploration steps)
(5 10 exploration steps)
(5 10 exploration steps)
600
400
100
Drexmer
200
THICK
100
Dreamer
3
2
steps
10th steps
10 steps
Figure 20: VisualPinPad results and exploration. We plot the mean evaluation returns (7 seeds for 5 10y exploration, seeds otherwise; = standard error).
ups in MiniHack, does not directly affect rewards or substantially influence image reconstruction, it might encoded only locally in h, and wash out over time. In Dreamer this not a problem be- cause the policy can amplify task-relevant information in h2 Director, however, compresses h, into a low-dimensional goal-encoding. Thereby, task-relevant latent information could get lost. Note that all novel tasks proposed in Hafner et al. (2022), in which Director shows very strong performance, avoid long- horizon memory, e.g. by coloring walls in a maze (Egocentric Ant Maze) or by providing visible history of past button presses (VisualPinPad).
In smaller MiniHack tasks, e.g. KeyRoom-Fls Fiein---xe-s5,. memory can sometimes be circumvented by specifying goals via observable information. For example, if both the key and door are visible, a goal would be the same observation without the key visible (picked up) and an open door. This creates a curriculum to pick up a key and open the door automatically from the high-level goal of an open door. In larger task spaces, e.g. KeyCorridor- 8, Director never encounters such simpler situations to begin with.
which the worker can first learn from such simpler situations and later learn
F.4 MBRL IN VISUALPINPAD: EXPERIMENT DETAILS
For the VisualPinPad suite we generated offline training data to bypass the challenge of discovering very sparse rewards. For data collection, used Plan2Explore (Sckar et al., 2020) with the default settings of the DreamerV2 (Hafner et al., 2020) codebase. We trained two randomly initialized models of Plan2Explore for 5 (0, 250k, 500k, 1M) environment steps in each task of the Visual Pin Pad suite. For each setting, we determined the model that achieved the highest overall returns during training. We initialized the replay buffer of all new models with the S samples.
Originally, VisualPinPad has more levels of difficulty. However, in wessualPinPaaa Plan2Explore did not receive any reward during 1M steps of exploration. In addition that, the results in Hafner et al. (2022) suggest that Dreamer is also not able to discover the very sparse rewards of VisualP inP adSix on its own. Thus, we omitted VisualP PiiiPannPaiiii. and more complicated levels.
F.5 MBRL IN VISUALPINPAD: EFFECT OF EXPLORATION
We analyze the effect of exploration data by varying the number data points with which we ini- tialize the replay buffers. For this, consider exploration data collected by S E {0, 250k, 500k) environment steps of Plan2Explore and compare THICK Dreamer to Dreamer. In PinPadTThae Dreamer and THICK Dreamer always achieve the same performance regardless of exploration data
32

---

Published as a conference paper at ICLR 2024
Pusher-Denas
Door
PickUp
Figure 21: Subgoals proposed in the first time step. We reconstruct images based on the subgoals 21 znnal that THICK PlaNet the first time step. The subgoals typically pushing the puck (Pusher), moving to the door handle (Door), or grasping the ball (PickUp). For PickUp the system sometimes fails find a reasonable subgoal (center).
available (cf. Fig. 20a, Fig. 20b Fig. 20d). Without exploration data, neither THICK Dreamer nor Dreamer manage to obtain rewards in PinPadE our and PinPadF's saFivve ive within 600k steps. Sim- ilarly, both methods do not discover rewards when initialized with 250k steps exploration in PinPadFive ve. Thus, for the more complicated problems, both THICK Dreamer and Dreamer need sufficient exploration. Whenever there is enough exploration data to learn the more complicated tasks, THICK Dreamer manages achieve high rewards faster than Dreamer (see Fig. 20c, Fig. 20e, Fig. 201).
For the larger problems, i.e. PinPadFour our and PinPadFin ivad,ivee we quantify the effect of explo- ration data on sample efficiency by determining the number of environment steps needed to reach certain level of reward. We take 95% of the highest mean reward in all of our experiments threshold. This corresponds to a mean reward of 359 for ann dur and 274 for P PinP adF ive. Table 3 shows the number of environment steps needed for THICK Dreamer and Dreamer reach this threshold for particular environments over the amount of exploration data. In sum, medium amount of exploration data (500k) enables reaching the threshold the fastest. THICK Dreamer reaches the reward threshold faster than Dreamer in all experiments. This advantage increases with level diffi- culty.
Table 3: Sample efficiency in Visual Pin Pad. We list the number of environment steps needed for THICK Dreamer, Dreamer, and their difference to pass a reward threshold (95% of max. reward) during evaluation for particular environments and amount of exploration data.
PinPa Paduuuuur
exploration data
250k
500k
1M
THICK Dreamer
200k
120k
140k
Dreamer
280k
180k
200k
difference
80k
60k
60k
P AnnadFiiee
exploration data
250k
500k
1M
THICK Dreamer
260k
340k
Dreamer
/
360k
590k
difference
100k
250k
F.6 MPC: EXPERIMENT DETAILS AND EXTENDED RESULTS
To study zero-shot planning, generated offline datasets of 1M environment steps for every task. For data collection, we used Plan2Explore in the same way as described in Suppl. F.4. After deter- mining dataset for each task, trained the models purely this data.
In addition to increasing performance for long-range tasks, THICK PlaNet provides the additional advantage that the subgoals proposed by the high-level network can be directly reconstructed into
33

---

Published as a conference paper at ICLR 2024
(a)
VandCEDeat the Advanced
(b)
Dour
30
rehurn
0.5
Dreamess8
retumm
20
PlaNet-8
DreamesM
Dresmer
PlaNet
Dreame-X1.
PlaNet NL
0.5
Trick Dreemer
THICK PlaNot
2
8
3
10 steps
10 updates
Figure 22: Effect of scaling RSSM size. We modify the number of hidden units per layer and dimensionality of the deterministic hidden state h2 by factor (S 0.5,Ms = I, = 2, XL 4) for Dreamer (a) or PlaNet (b) and compare to the unmodified THICK counterparts. Each graphic plots the mean returns over training steps (5 random seeds for ablations). Shaded areas depict the standard error in (a) or standard deviation in (b).
images through the low-level output heads og The resulting goal images are easily interpretable by humans. Figure 21 shows exemplary goals selected by the high-level planner in the first time step of an episode. Thus, the behavior of THICK PlaNet is much more explainable than simply performing MPC in a flat world model.
F.7 MBRL & MPC: SCALING MODELS VS. IIBRARCHY
An alternative approach to adding hierarchies to world models is to improve performance by scaling up the model size (Hafner et al., 2023; Deng et al., 2023). Could simply increasing the capacity of the world model improve performance of Dreamer PlaNet similarly to our approach of incorporating hierarchical abstraction?
We investigate this in the fandofDeating We increase the RSSM model capacity by scaling the number of hidden units per layer (256 before) and the dimensionality of the determinstic hidden state h, (256 before) by different factors (factors: S = 0.5, = I. 2, XL = 4). Unlike the model scaling in Hafner et al. (2023), we did not increase the dimensionality simple. Figure 22 plots the returns of the different model sizes over environment steps. In both tasks, increasing the model size did not improve Dreamer (Fig. 22a) or PlaNet (Fig. 22b). Thus, for the investigated setups, scaling up model size does not bring the same advantages as THICK hierarchy.
task for MBRL and in Door for zero-shot MPC.
z, as both investigated environments are visually rather
F.8 ABLATIONS AND HYPERPARAMETERS
Ablations We ablate various components of the C-RsSM and THICK Dreamer within MBRL setup. We evaluate the resulting systems using the two exemplary tasks of and VisualPinPadFoux. Figure 23 plots the re- turns of the ablated systems over environment steps. Using the C-RSSM in DreamerV2 results in roughly the same performance SissulliPPaaooo
or slightly better performance than using the RSSM (i.e. Dreamer). However, removing the determin- istic latent state h, and the precise processing pathway from the C-RsSM (i.e. C-RssM Dreamer without h) impedes the system from learning the tasks. Omitting Ug. and only using one critic ux for both the short- and long- horizon returns (Eq. 20), slightly degrades the performance of THICK Dreamer.
Hyperparameter Jeparse Next, we compare the effect of sparsity loss scale gsparse on THICK Dreamer in Www.allinPadaaaaaaa. our and on THICK PlaNet in Multiworld-Door.
FFor this ablation we picked higher sparsity regularization gspare for both tasks (@eparse = 50 for Wan HandOfibeath aheeth-AAee Sh-Advances.. gepare 10 for visuallli w such that the number of time steps with open gate roughly matches that of the C-RsSM.
34

---

Published as a conference paper at ICLR 2024
(a)
(b)
VisusiPinPedFour
0.5
300
Dreamer
retum
retum
Dreamer
200
0
C-RSSM Dreumer
C-RSSM Deeumer (no hz)
-0.5
100
Dreanerr (no vy)
0
2
4
6
2 10a steps
3
10 steps
Figure 23: C-RsSM and THICK Dreamer ablations. Each graphic plots the mean returns over training steps. We compare Dreamer and THICK Dreamer (7 random seeds) against various ablations (5 seeds each): Dreamer using the C-RsSM (C-RsSM Dreamer), Dreamer using only the coarse processing pathway of the C-RSSM (C-RSSM Dreamer no h.), and THICK Dreamer using only one critic (THICK Dreamer rt). Shaded areas depict the standard error. error.
Figure 24a plots the mean returns of THICK Dreamer or different values for feparss in VisualPinPadFoo aaannaadFous our.. Figure 24b shows the percentage of time steps with context changes over training. For THICK Dreamer, regularizing context changes too little is not as detrimental as overly regularizing context changes. If the contexts are weakly regularized, i.e. small gaparse, then the con- text changes in most time steps. As result, the high-level learns an identity mapping, and during a temporal abstract prediction, the network simply predicts the next state at time + (see Alg. I). Stronger regularization boosts sample efficiency of learning long-horizon behavior. This is even true if, at some point after the behavior has been sufficiently learned, the context is no longer adapted (e.g. grparse 10). However, overly strong regularization, which prohibits context changes early during training. impedes the learning of the task (e.g. 3uparse 100). In this case, the high-level pre- dictions essentially average state predictions, which simply contributes noisy values for learning the critic. THICK Dreamer is very robust to the choice of gepärse in VisualP winlllinnPaaa i aduo.r.
Figure 24c plots zero-shot planning performance of THICK PlaNet for different values for geparse with the percentage of context changes shown in Fig. 24d. For THICK PlaNet both too strong and weak regularization degrade performance. However, strongly regularizing the network toward sparse context changes is slightly less detrimental for THICK PlaNet than a weak sparsity regular- ization (cf. geparse 100 and feparse 5). For weak sparsity regularization, the context changes in every time step, which prevents the high level from finding useful subgoal sequence during plan- ning. As a result, the low-level might be guided into the wrong direction by the proposed subgoals.
(a)
(c)
returns in Door
(d)
returns in VisualPinPadFour
%t with c # . in VisualPinPadFour
%t with c, * c in Door
100
100
g 0.01
30
300
80
80
g
20
g
10
60 10
200
25
40
100
10
40
100
100
20
0
20
0
0
0
1
2
3
4
1
2
4
0
3
0
2
3
10th steps
10a steps
10 updates
10 updates
Figure 24: Effect of sparsity. We plot the mean return for THICK Dreamer in (a) and the mean zero-shot planning returns of THICK PlaNet in Door (c) for different values of the hyperparameter gepars (5 random seeds). In addition, we plot the per- centage of time steps with context changes over training time for both tasks (b, d). We test different ranges of geparse for the different tasks. Shaded areas depict the standard deviation.
VisualPinPade
35

---

Published as a conference paper at ICLR 2024
(a)
Advaaaceeddd
(b)
VisusiPinPadFou
300
0.5
retum
retum
200
0
V 0.9
0..
100
V
0.8
R 0.5
1
0
2
4
6
2 10a steps
3
10 steps
Figure 25: Effect of hyperparameter sh Each graphic plots the mean returns for THICK Dreamer over training steps for different values the hyperparameter vó (5 random seeds). Shaded areas depict the standard error.
Hyperparameter v THICK Dreamer introduces a new hyperparameter 3 which balances the influence of the short-horizon value estimates V and long-horizon value estimates V/ongg on the overall value V (Eq. 20). Figure 25 shows how & affects task performance. Only con- sidering short-horizon value estimates, i.e. vb = 1, results less sample efficient learning than taking small amounts of long-horizon value estimates into consideration, i.c. + 0.9 for aeaoDDeeth-Advaaaaa
and 0.8 I e < 0.9 for VisualPinnnn PnniPPinPPFuus However, relying too strongly long-horizon estimates, i.c. &b 0.5, impedes policy learning. This effect is less pronounced for very long-horizon tasks such as vi Visual PindidFour.. We set ab 0.9 in all experiments.
Hyperparameter THICK PlaNet introduces the hyperparameter which scales the influence of the subgoal proximity on the reward estimate of the low-level planner (Eq. 22). We analyze the effect shown in Fig. 26. Incentivizing
of on THICK PlaNet's performance in Mu 1 t. subgoal proximity too strongly, i.e. к 1, can result in the agent getting stuck at a subgoal. This reduces overall task performance. Ignoring the subgoal, i.c. 0, also decreases performance for long-horizon tasks such as Door. In Door, THICK PlaNet works well across a wide range k.
Replanning strategy THICK PlaNet proposes new goals the high-level upon context transi- tions. We do this mainly to save computational cost from running MCTS at the high-level at every time step. Figure 27 compares the effect of high-level planning every time step to replanning upon context transitions in Door. The returns seem mostly the same. Thus, replanming at every time step is effective as setting new subgoals only upon context transitions and can be applied computational efficiency is not concern.
Door
Doar
30
30
return
20
==0
0.025
10
0.05
10
0.25
*
0-
0
2
2
10 updates
10 updates
Figure 26: Effect of hyperparameter plot mean returns of THICK PlaNet for zero- shot MPC for different values of x (5 seeds, standard deviation).
We
Figure 27: Effect of replanning. We plot mean returns of THICK PlaNet for zero-shot MPC when replanning upon context transition ev- ery step (5 seeds, + standard deviation).
36

---

Published as a conference paper at ICLR 2024
G COMPUTATION AND CODE
All experiments were run on an internal computing cluster. Each experiment used one GPU. Experi- ments using DreamerV2 took roughly 15-20 hours, while THICK Dreamer experiments took around 35-45 hours of wall clock time, depending on the overall number of environment steps. Experiments with Director took around 40-60 hours of wall clock time. Zero-shot MPC experiments with PlaNet took about 30-35 hours for 10t training steps, while THICK PlaNet took roughly 50-60 hours for the same number of training steps. The higher wall clock time for training THICK world models stems mainly from the greater number of trainable parameters and more detailed logging. In addition to that, THICK PlaNet takes longer evaluate, due the additional computational cost of running MCTS the high level during planning. The code for running the experiments is available online.
37


Reviews:
[{"review_id": "U8VW26MBub", "reviewer": "Reviewer_sX3c", "summary": "This paper proposes a method called THICK for learning hierarchical world models with adaptive temporal abstractions from discrete latent dynamics. The key idea is to use a context-specific recurrent state space model (C-RSSM) to create a sparsely changing context latent variable. This context encodes higher-level transitions that are used to train a high-level predictor. The high-level model predicts states that lead to context changes, enabling temporal abstract predictions. The hierarchical predictions can be integrated into model-based RL and planning methods.", "questions": "1) How does the performance compare to hierarchical RL methods like h-DQN or FeUdal Networks?\n\n2) Could the hierarchy be extended to have multiple abstraction levels instead of just two?\n\n3)  How well does THICK scale to even longer time horizons or higher-dimensional state spaces?\n\n4)  Could active exploration be used to shape useful context abstractions instead of relying on a sparsity loss?\n\n5) Is there some mechanisms could be add to make the high-level plans more reactive to execution errors or environment changes?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The C-RSSM provides a simple yet effective way to learn sparsely changing context variables from pixel observations without supervision.\n\n- The high-level model is trained in a self-supervised manner to anticipate context changes based on the C-RSSM's discrete dynamics.\n\n- The method is generally applicable across various environments with visual observations.\n\n- Experiments show that the learned hierarchies capture meaningful abstractions related to subgoals.\n\n- Integrating THICK's hierarchical predictions improves sample efficiency of model-based RL in long-horizon tasks.\n\n- The high-level plans can be visualized, enhancing model interpretability.", "weaknesses": "1) The sparsity hyperparameter for the C-RSSM needs to be tuned for each environment.\n\n2) The high-level model operates at a fixed abstract timescale, less flexible than methods with hierarchical policies. \n\n3) The high-level plans are not actively updated during execution, being replanned only at context changes.\n\n4) The expressiveness of temporal abstractions may be limited compared to methods with backing task priors or curriculum learning.\n\nHowever, I think this paper still take a step further to hierarchical world model.", "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": "", "extra_scores": {}}, {"review_id": "piUpp6xT77", "reviewer": "Reviewer_h1NE", "summary": "This paper proposes a novel way of learning a hierarchical world model. On top of Dreamer's world model implementation, RSSM, this paper introduces Context-specific RSSM (C-RSSM) with a slowly (sparsely) changing discrete context state $\\mathbf{c}_t$, which represents some static scene info that is preserved for a long time horizon. The continuously changing states $\\mathbf{h}_t$ and $\\mathbf{z}_t$ are then conditioned on this context state $\\mathbf{c}_t$. With the trained C-RSSM, a high-level transition is defined as a contiguous transition segment with the same $\\mathbf{c}_t$. The high-level world model (THICK) can be trained to predict the context and stochastic state of the next segment. The paper utilizes C-RSSM and THICK for model-based RL by combining with Dreamer and model-based planning by using MPC. The experiments on MiniHack, PinPad, and robotic manipulation environments show that THICK+Dreamer and THICK+MPC outperform the flat world model baselines when the task horizon becomes longer.", "questions": "Please address the weaknesses mentioned above.\n\n\n### Minor questions and suggestions\n\n* It might be better if Figure 2 could illustrate that $\\mathbf{c}$ is changing slowly and the high-level transition happens when $\\mathbf{c}$ has changed.\n\n* In Equation 11, $\\delta$ starting from 0 makes more sense?\n\n* In Equation 21, $t < K$ should be $\\tau(t) < K$ or $A_{\\tau(t):K}$ should be $A_{t+1:K}$?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "* This paper tackles an important problem of learning a hierarchical world model. \n\n* The idea of learning a sparsely changing context state is intuitive and using context switches to define high-level transitions is a sensible choice. Moreover, this enables THICK to naturally incorporate variable-length skills.\n\n* The learned hierarchical world models (THICK and C-RSSM) can be integrated with both model-based RL and model-based planning.\n\n* The paper is clearly written and the figures help us understand the complex concepts of the proposed method.\n\n* The experimental results support that the hierarchy in the world model improves long-horizon prediction.", "weaknesses": "* The design choices of the proposed approach are not explained and justified. Especially, the high-level world model learns to predict the future state and action just before the context switch, rather than the context and state right after the context switch. Predicting a low-level action sometime in the future sounds very difficult to me. Although this design choice also makes sense and works in practice, it would be great to explain the rationale behind this specific design choice. \n\n* Section 2.1 introduces the coarse prior of the stochastic state but it is unclear what is the rationale behind this design. It would be great to further explain why C-RSSM needs both coarse and precise priors. Similarly, more explanation about the need for coarse predictions in Equation 8 would help understand the proposed method.\n\n* Generally, the experimental results of THICK+Dreamer and THICK+PlaNet are similar to those of DreamerV2 and PlaNet. Despite its novelty, these weak experimental results may make its impact less significant. Stronger results in more diverse environments would be greatly appreciated. If possible, it would be great to see its performance on the common RL benchmarks, such as Atari and DMC tasks.\n\n* The paper claims that Thick Dreamer is more sample efficient than DreamerV2 in PinPadFour and PinPadFive but the improvement is relatively marginal to strongly claim this.\n\n* Moreover, the experiments on the PinPad environments use Plan2Explore to fill in initial exploratory data. As RL is inherently a combination of exploration and exploitation, it would be important to see how it works for exploration. Thus, it is recommended to include experiments on RL from scratch.\n\n* Although the investigation of world model hierarchies is important, many deep learning approaches seem in favor of scaling model sizes instead of injecting hierarchies. In this sense, it might be interesting to see comparisons between scaling models [a] vs. C-RSSM.\n\n\n[a] Deng et al. Facing Off World Model Backbones: RNNs, Transformers, and S4. NeurIPS 2023", "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": "", "extra_scores": {}}, {"review_id": "J3hURTNOk2", "reviewer": "Reviewer_yLBk", "summary": "The paper proposed Temporal Hierarchies from Invariant Context Kernels, a method for extracting a hierarchical world model using the proposed Context-specific RSSM (which extends the existing RSSM model (Hafner et al., 2018)) using a coarsely-updating context variable. CRSSM predicts in two levels, selectively updating parts of its latent state (context) sparsely in time. They combine the hierarchical world model (CRSSM) with Dreamer (Hafner et al., 2020) to propose THICK Dreamer, a model-based RL method that combines value estimates from low- and high-level predictions of state and context variables to compute a long-horizon value function. The authors evaluate THICK on MBRL tasks and qualitatively show context switches in their world model.", "questions": "None", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1) The paper is easy to follow and well-structured with the main contributions listed clearly in the introduction. The proposed world model, Context-specific RSSM, makes a simple yet effective addition to RSSM, and the authors contextualize it well within the existing RSSM model. Contrasting with the existing literature in this area, I find the model visualizations to be very clear and self-explanatory.\n2) The authors show interesting qualitative results using CRSSM, extensive evaluation of their world model when fitted in Dreamer and PlaNet for MBRL, and provide thorough details about their task setups in the appendix adding to the reproducibility of this work.", "weaknesses": "1) Looking at figure 7, there does not seem to be a significant difference between the performance of Dreamer and the proposed method. Can the authors justify the marginal performance improvement? \n2) The reported results using Director show close to 0 success on all the tasks. Can the authors explain why was that the case?", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)