Paper: Learning Hierarchical World Models with Adaptive Temporal Abstractions from Discrete Latent Dynamics (ICLR-2024)

Reviewer comment (Evaluation):
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?

Author response:
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).

Reviewer comment (Theory):
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.

Author response:
### 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.

Reviewer comment (Theory):
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.

Author response:
### 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.

Reviewer comment (Experiments):
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.

Author response:
### 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.

Reviewer comment (Experiments):
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.

Author response:
### 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.

Reviewer comment (Experiments):
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.

Author response:
### 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.

Reviewer comment (Experiments):
The sparsity hyperparameter for the C-RSSM needs to be tuned for each environment.

Author response:
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).

Reviewer comment (Theory):
The high-level plans are not actively updated during execution, being replanned only at context changes.

Author response:
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.

Reviewer comment (Theory):
The expressiveness of temporal abstractions may be limited compared to methods with backing task priors or curriculum learning.

Author response:
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.

Reviewer comment (Experiments):
How does the performance compare to hierarchical RL methods like h-DQN or FeUdal Networks?

Author response:
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.

Reviewer comment (Experiments):
Could the hierarchy be extended to have multiple abstraction levels instead of just two?

Author response:
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.

Reviewer comment (Experiments):
How well does THICK scale to even longer time horizons or higher-dimensional state spaces?

Author response:
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.

Reviewer comment (Experiments):
Could active exploration be used to shape useful context abstractions instead of relying on a sparsity loss?

Author response:
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.

Reviewer comment (Evaluation):
Is there some mechanisms could be add to make the high-level plans more reactive to execution errors or environment changes?

Author response:
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