Learning Hierarchical World Models with Adaptive Temporal Abstractions from Discrete Latent Dynamics

We thank the reviewer for their great feedback.

Performance improvements

Sparse reward tasks such as MiniHack pose the challenges of exploration and learning long-horizon behavior. With THICK Dreamer we only tackle the latter through long-horizon reward propagation. This amplifies reward signals once rewards are discovered but does not necessarily help discover rewards earlier. This is reflected in Fig. 7. A stronger performance gain is achieved when the high-level guides the low-level by setting subgoals (THICK PLaNet, Fig. 9).

Director in MiniHack

Thank you for this question. To analyze when Director fails in MiniHack, we carefully checked our implementation and analyzed the performance of Director in more detail. This lead us to find a small inconsistency in our Director implementation affecting the MiniHack tasks. We reran all MiniHack experiments. In the simplest problem (KeyRoomFixed-S5), in which Director sometimes managed to solve the task, Director now consistently solves the task. Thank you for making us revisit Director to catch this.

However, the problem of Director never learning to solve most MiniHack tasks remains and we identified two problems that impede Director’s performance in MiniHack, which we now outline and discuss in Suppl. F.3.

One problem of Director is that its learning can be slowed down during early stages of training depending on the diversity of initial data. Director trains a goal encoder, which defines the input space for the goal-conditioned policy, on its replay buffer. Thus, if in the beginning not enough diverse data is 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). We exemplify this for variants of the KeyDoor problem in MiniHack in Suppl. F.3 and Fig. 19. By increasing the diversity of collected data, e.g., randomizing spawn positions and increasing level layout, the performance of Director during initial training strongly improves.

A more severe problem of Director is unobservable information for setting goals. The RSSM encodes both observable as well as unobservable information within its deterministic latent state $h_t$. If the unobservable information, e.g., item pick-up in MiniHack, does not directly affect rewards or substantially influence image reconstruction, it might be encoded only locally in $h_t$ and wash out over time. In Dreamer, this is not a problem because the policy can amplify task-relevant information in $h_t$. Director, however, compresses $h_t$ into a low-dimensional goal-encoding. Thereby, task-relevant latent information in $h_t$ could get lost. Note that all novel tasks proposed in the Director paper [Ref a] avoid long-horizon memory, for example by coloring walls in a maze (Egocentric Ant Maze) or by providing a visible history of past button presses (Visual Pin Pad).

In smaller MiniHack tasks, e.g., KeyRoomFixed-S5, this 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, in which the worker can first learn from such simpler situations and later learn 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.

We thank the reviewer for their review, and we are happy to answer further questions.

References

[Ref a] Hafner et al. Deep hierarchical planning from pixels. NeurIPS 2022

We thank the reviewer for their super helpful comments and suggestions and for thoroughly checking our equations and catching some mistakes.

Design choice 1: Predicting states before context transitions

Our high-level model is a generative model that learns to predict situations in which latent generative factors are assumed to change. Such situations can be related to context transitions [Ref b] or event boundaries [Ref c], as well as to context switches when considering the RL options framework [Ref d]. To predict such situations, two challenges arise: (1.) learning in which situation such transitions are currently possible and (2.) how these transitions affect the latent generative factors. The low-level world model already learns to encode (2). Accordingly, we train the high-level model to only learn (1) and then prompt the low-level model for (2).

We further emphasize the reasons for this design choice now in Sec 2.2 and explain this in detail in a new “design choices” subsection in Suppl. D.1.

Design choice 2: Coarse predictions

Thank you for pointing out that we may have not highlighted this sufficiently. The coarse prediction are important for THICK for two main reasons:

(1.) Prediction-relevant contexts $c_t$: Context $c_t$ is designed to encode latent information that is necessary for prediction and reconstruction. If we would omit the coarse prior predictions (Eq. 5) and coarse output reconstructions (Eq. 8) the C-RSSM would have no incentive to encode prediction-relevant latent information in $c_t$, reducing the system to the RSSM. As a result, no hierarchy would develop and hierarchical planning would stay impossible.

(2.) Predictions without $h$: The high-level model attempts to predict a 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 high-dimensional hidden state $h$ many time steps in the future. Thus, it is advantageous that the C-RSSM can make predictions without $h_t$. After a high-level prediction we can simply feed the high-level outputs ($\hat z_{\tau(t) -1}, \hat a_{\tau(t) -1}$) into the low-level world model. This allows us to predict rewards or episode terminations/discounts at particular situations (used in THICK Dreamer and THICK PlaNet) or reconstruct images to visualize predictions (see Sec 3.1).

We emphasize the reasons for this design choice now more clearly in Sec 2.1 and 2.2 and add a paragraph to the Suppl. D.1 to explain these critical design choices in further detail.

Common RL benchmarks

THICK world models attempt to tackle partially-observable problems with long task horizons. The problems in the DMC suite have a short task horizon. The problems in Atari are almost (via frame stacking) fully observable. Only a few problems have longer task horizons. Thus, we believe that these suites would not showcase the advantage of using THICK world models particularly well.

Exploration data

Please note that in all MiniHack experiments we do not use exploration data. However, with THICK Dreamer we only tackle exploitation through reward propagation and do not use the hierarchy for exploration. This is why we sidestepped the hard challenge of exploration in Visual Pin Pad.

However, we now plot the return curves for PinPadThree without exploration data in Suppl. F.5 and Fig. 20a. For PinPadFour and PinPadFive no rewards are discovered within 600k steps for THICK Dreamer and Dreamer without exploration data.

To paint a full picture,we analyzed the effect of the number of exploration samples on performance by comparing different amounts of exploration samples (250k, 500k, 1000k) in Suppl F.5, Fig. 20 and Tab. 3. The effects of slightly increased sample efficiency for THICK Dreamer in PinPadFour becomes more pronounced when fewer exploration samples are available. Regardless of the amount of exploration data THICK Dreamer always matches or outperforms its flat counterpart.

Comparison model size vs. hierarchy

Thank you for the excellent suggestion. Unfortunately, we could not compare against [Ref a] since their code is not publicly available. But in a new additional experiment (Suppl. F.7 and Fig. 22), we compare Dreamer and PlaNet with different RSSM model sizes against THICK Dreamer and THICK PlaNet in MiniHack-WandOfDeathAdv and Multiworld-Door. We change model size by modifying the number of hidden units per layer and the size of the deterministic hidden state $h_t$ by different factors ($\in \{0.5, 1, 2, 4\}$ for S-, M-, L-, XL-sized models). Increasing model size does not improve Dreamer or PlaNet in these tasks. Our unmodified THICK world models always achieve the best performance.

Thank you for the valuable feedback and questions. First, we would like to clarify potential misunderstandings before answering questions.

Sparsity loss scale

While it is true that ideally $\beta^\mathrm{sparse}$ can be set individually for each task, in practice the same value works well across similar tasks (e.g., same setting per suite, see Suppl B).

Time scale of the high-level

It is not true that the high level operates on a fixed time scale. This is one important feature of THICK - we apologize for not clarifying this sufficiently well. As we explain in Sec 2.2, “THICK uses [...] discrete context dynamics as an adaptive timescale for training a high-level world model”. The function $\tau(t)$ implements a variable temporal abstraction that maps a point in time to the next point in time with context changes. Thus, the high-level learns predictions of variable length, anticipating context transitions. As we analyze in Suppl. F.1 (e.g., Figure 16) the duration between context changes can widely vary within a task or across tasks. We highlight this now more in the text.

Reactivity of high-level plans

In the paper, we had THICK PlaNet replan on the higher level only when the lower level context changes. This decreases computational effort. In a new ablation (Suppl F.8 & Fig. 27) we show that replanning on the higher level in every step does not drastically change performance in Multiworld-Door. Thus, while THICK PlaNet may run model-based planning on the higher level in every time-step, we recommend running it sparsely or only when context changes to save computational energy.

Task priors or curriculum learning

We unfortunately do not understand this point. Could you please elaborate? We do not see curriculum learning and our learning of temporal abstractions to be exclusive. They could be easily combined.

Questions

How does the performance compare to hierarchical RL methods like h-DQN or FeUdal Networks?

h-DQN requires pre-determined abstractions [Ref a], white our method discovers them from scratch.

FUN performs worse or less sample efficiently than DreamerV2 when both methods are applied to the same problems (compare Fig. 4 of [Ref b] with Figure F.1 in [Ref c], or compare Fig. 2 of [Ref b] with Fig. B.2 in [Ref c]). Seeing that THICK Dreamer performs as good or better than DreamerV2, we believe it would outperform FUN.

Could the hierarchy be extended to have multiple abstraction levels instead of just two?

As we write in our conclusion, THICK is not restricted to learning a two-level hierarchy. By replacing level 2 with a C-RSSM, in principle, the same mechanism could be applied to learn a level 3 world model etc. In this way an N-level hierarchy of world models with nested time scales could develop.

Another option would be to foster the development of multiple coarse pathways on the low level by means of varying sparsity regularization strengths. For each pathway a separate high-level model could be trained. In this approach, the developing hierarchy would not necessarily be strictly nested.

How well does THICK scale to even longer time horizons or higher-dimensional state spaces?

We believe THICK should scale to very long time horizons, since it does not depend on the exact time scale of the environment but on the developing time scale of context updates. Similarly, we expect THICK to be also applicable to problems with high-dimensional state spaces. Potentially the dimensionality of $c_t$ would need to be increased to capture all possible context changes in the environment.

Could active exploration be used to shape useful context abstractions instead of relying on a sparsity loss?

This is an excellent point. Active exploration might support the shaping of context abstractions. We keep this as future work, also as it would be a complementary direction.

Is there some mechanisms could be add to make the high-level plans more reactive to execution errors or environment changes?

As outlined above and shown in our new ablation in Suppl. F.8 and Fig. 27, high-level planning could be active at every time step

References:

[Ref a] Kulkarni et al. Hierarchical Deep Reinforcement Learning: Integrating Temporal Abstraction and Intrinsic Motivation. NeurIPS 2016

[Ref b] Vezhnevets et al. FeUdal Networks for Hierarchical Reinforcement Learning. ICML 2017

[Ref c] Hafner et al. Mastering Atari with Discrete World Models. ICLR 2020