Response to Reviewer hTvu (2/2)
> Stochastic vs non-stochastic dynamics
We choose non-stochastic environments mainly because there exists a single groundtruth future sequence for evaluating the generation consistency. In stochastic environments, there are multiple possible futures at each time step, causing the total number of correct future sequences to grow exponentially with the prediction length. This makes evaluation of generation consistency much harder. We believe it is worth future investigation to find proper ways to evaluate generation consistency in stochastic environments.
> Why would action-dependent modeling affect the ability to evaluate the memorization capabilities of a model?
First, we focus on world models which aim to approximate the state transition $p(s_{t+1}|s_t, a_t)$ and therefore have to be action-conditioned. Second, action conditioning makes it easier to obtain non-stochastic transitions, and in turn simplifies evaluation.
> Is there a set of synthetic tasks that may be a sufficient indicator for the MBRL problem itself?
We don't have a clear answer yet. There are quite a few desired abilities of MBRL agents, such as exploration, long-horizon planning, skill learning, robustness to distractions, etc. It is extremely challenging for an MBRL agent to possess all these abilities. Therefore, multiple benchmarks have been developed, each focusing on one or several aspects. In this work, we focus on long-term memory capacity with extremely long sequences of up to 2000 steps.
> What makes it harder for other models to reconstruct the correct color information?
To predict the correct color, the model needs to remember what it has seen hundreds of time steps ago (please refer to the supplementary videos for typical agent paths). RSSM and TSSM fail mainly because they have limited memory capacity. RSSM is based on RNN, which is prone to the vanishing gradients problem, making it hard to store information over long time horizons. TSSM is based on Transformer, which has quadratic complexity with respect to its context window length (denoted $m$ in the paper). In practice, we need to limit $m$ for computational efficiency, and it will be hard to propagate information for much longer than $m$ steps.
> Set query steps beyond the context steps
Thank you for your suggestion. We note that in the Multi Doors Keys environment, we have used query steps that are more than twice the context steps (Table 3). We find that S4WM significantly outperforms RSSM and TSSM in this setting. To further answer your question, we provide the generation result in the Two Rooms environment with 300 context steps and 700 query steps. Please refer to Figure 2 in the uploaded pdf. We note that during training, the model has only seen sequences of length 500 (300 context + 200 query). Our results show that S4WM can successfully generalize to sequence lengths that are beyond the training regime.
> Does good recall in the two rooms correlate to good recall in the four and proportionally in the ten as well?
We don't think there is a correlation. Good recall becomes more challenging when there are more rooms, as the model needs to remember more information. We use two/four/ten rooms as three difficulty levels to thoroughly compare the memory capacity of different models.