Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient

Some results are unclear. For example, could you go into details on why the model is not doing well on the Breakout game? What is it producing? Some figures from the game can be also useful--not just this one, but the ones in which model does well too.

We assume the reviewer's question arises because Drama performs well in the game Pong but not in Breakout, despite the two games sharing some similar features. The reason is that Breakout is more visually complex due to its colorful bricks, causing the encoder to fail in effectively encoding the ball. As requested, we have added a subsection in the appendix to explain this in detail with the experiment figures. Please refer to Section A.7 in the revised version.

It is not clear why imagination context length is kept to 8. I would suggest providing experiments, both involving time-complexity and performance, for different imagination context length. Also, explaining why this is done would be useful.

MBRL involves numerous hyperparameters, and conducting a comprehensive hyperparameter search demands substantial computational resources and time. Consequently, we adopt the hyperparameters established in prior research studies. This approach is commonly employed in the literature. For example, several studies such as TWM (Robine et al., 2023), STORM (Zhang et al., 2023), Hieros (Mattes et al., 2024), and DreamerV3 (Hafner et al., 2024) utilise the same imagination horizon, originally introduced in DreamerV1. Specifically, the imagination context length in this work aligns with the hyperparameter used in STORM.

Reference:

Zhang, Weipu, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang.
“STORM: Efficient Stochastic Transformer Based World Models for Reinforcement Learning.”
In Thirty-Seventh Conference on Neural Information Processing Systems, 2023.

Robine, Jan, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling.
“Transformer-Based World Models Are Happy With 100k Interactions.”
In International Conference on Learning Representations, 2023.

Mattes, Paul, Rainer Schlosser, and Ralf Herbrich.
“Hieros: Hierarchical Imagination on Structured State Space Sequence World Models.”
In Forty-First International Conference on Machine Learning, 2024.

Hafner, Danijar, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap.
“Mastering Diverse Domains through World Models.”
arXiv, 17 April 2024.

What is the value of h?

We assume that $h$ refers to the dimension of the hidden state in Mamba. Specifically, $h$ is 16 for the XS model and 32 for the S model.

"A key difference between Mamba-based and transformer-based world models in the ‘imagination’ process is that Mamba updates inference parameters independent of sequence length."--can you explain it more?

Yes, during imagination, Mamba (both versions 1 and 2) utilises a hidden state to summarize all past information. The hidden state has a fixed dimensionality (16 or 32 in the examples above), the model updates the hidden state at time step $t$ without reprocessing the past token $x_{t-1}$. Consequently, the inference time scales linearly with the sequence length. This scalability is illustrated in Section A.8, Figure 10 (A) of the revised version, showing that Mamba-based world models are faster for "imagination" compared to Transformer-based world models.

[W1] The extent to which the Mamba architecture contributes to the model's performance remains unclear. Specifically, it is unclear how DFS impacts scores across all games. Extending ablation study 3.2.1 to cover more games, or conducting a new study that replaces Mamba with an RNN or transformer, would clarify these contributions.

As noted above, to fully address this concern, we conducted a study comparing uniform sampling and DFS in DRAMA across all games. The results demonstrate the effectiveness of DFS, which outperformed uniform sampling in 11 games, underperformed in 2 games, and tied in 13 games.

[W2] While the paper emphasizes Mamba's computational efficiency, there is a lack of exact wall-clock training and inference times. The abstract claims the model can be trained on a standard laptop, so providing specific runtime metrics would substantiate this claim.

We did not explicitly state that Drama is computationally efficient in the paper, as MBRL is naturally more computationally complex than model-free RL due to the involvement of the world model. However, we mentioned that SSMs achieve $O(n)$ memory and computational complexity, where $n$ represents the sequence length. We interpret the reviewer's question as a request for proof of this claim.

Since we are using a shared server for training, it is difficult to ensure a clean and consistent environment to test the wall-clock training time. To address this, we evaluated the training time solely for the model using the grid world environment and have included the results in the revised version. Please refer to Table 2 for the result and section 3.2.3 for the detail in the revised version.

The paragraph in lines 99–105 shifts from general statements about model-based RL to specific details about the paper's world model. This transition could lead readers to infer that every world model relies on a variational autoencoder or linear heads, which isn't necessarily the case. Additionally, other model-based RL methods exist that don't utilize world models, such as those using lookahead search.

We have revised the text for clarity and included references to additional approaches.

"There are various approaches to obtaining a world model, including Monte Carlo tree search (Schrittwieser et al., 2020), offline imitation learning (DeMoss et al., 2023) and latent dynamics models (Hafner et al., 2019). In this work, we focus on learning a world model $f(O_t, a_t; \omega)$ from actual experiences to capture the dynamics of the POMDP in a latent space."

References:

Schrittwieser, Julian, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, et al. “Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model.” Nature 588, no. 7839 (24 December 2020): 604–609. https://doi.org/10.1038/s41586-020-03051-4.

Hafner, Danijar, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. “Learning Latent Dynamics for Planning from Pixels.” In Proceedings of the 36th International Conference on Machine Learning, 97:2555–2565. PMLR, 2019.

DeMoss, Branton, Paul Duckworth, Nick Hawes, and Ingmar Posner. “DITTO: Offline Imitation Learning with World Models.” arXiv, 6 February 2023. http://arxiv.org/abs/2302.03086.

Since DFS is agnostic to most baselines, it is important to see a comparison of either Drama with Uniform Sampling or baselines with DFS sampling to understand if the architecture is helping or the sampling.

As noted above, to fully address this concern, we conducted a study comparing uniform sampling and DFS in Drama across all games. The results demonstrate the effectiveness of DFS, which outperformed uniform sampling in 11 games, underperformed in 2 games, and tied in 13 games.

At line 222, it is not clear what targets mean.

It has been addressed. We expanded the writing to make it clearer. Updated lines in the revised version: 221-225 (Some latex formular is not supported here so we can't copy the revised sentence to here.)

Section 2.3 is not described well.

We have updated the text to clarify that we sample $b_{img}$ trajectories, each of length $l_{img}$.

"The behaviour policy is trained within the `imagination', an autoregressive process driven by the dynamics model. Specifically, a batch of $\displaystyle b_{img}$ trajectories each of length $l_{img}$ is sampled from the replay buffer. "

What is $h_t$ in behavior policy learning? The deterministic variable is defined as $d_t$ in Eq 5.

Thank you for pointing this out. It was a typo, which has now been corrected.

While training Dreamer, the method uses the whole sequence of $(b_{img}, l_{img})$ to compute a good hidden state and uses all sampled to generate trajectories in the future. I am curious to know why only the last hidden state $l_{img}$ is used for learning the policy and not the whole sequence like Dreamer? (describe around line 243).

Dreamer samples one batch of trajectories, typically in the shape (64, 16). These samples are then used to generate rollouts of length 15, resulting in a total of (1024, 16) samples to train the behavior model. However, since some starting points are consecutive, this can lead to overlapping and correlated imagined trajectories. To address this, we resample $b_{img} = 1024$ trajectories directly from the buffer to increase the diversity of the training samples. To ensure the same batch size of rollout, we only use the last hidden state $l_{img}$ to generate rollout with the horizon $h=15$, while while the sequence preceding $l_{img}$ is used solely to bootstrap the hidden state of Mamba2.

The terminal flag is defined as $e_t$ at line 93 ...

Thank you for pointing this out. We have updated the notation in the figure and legend to align with the main text.

Was any experiment conducted to see if the behavior model underestimates rewards especially with limited data?

We did not conduct any experiments specifically to examine this phenomenon. However, it is well explained with a simple yet convincing example in Chapter 8.3 of [Sutton & Barto, 2018]. In MBRL, the behavior model may overestimate rewards in states where the world model is underfitting—a phenomenon known as the model exploitation problem. This issue remains a significant challenge in model-based RL.

Theoretically, DFS can help mitigate both problems. It increases the likelihood of sampling fresh trajectories to train the world model and decreases the likelihood of sampling trajectories where the world model is underfitting. However, DFS does not fully resolve these issues.

Reference: Sutton, Richard S., and Andrew G. Barto. Reinforcement Learning: An Introduction. MIT Press, 2018.

For the ablation presented in Sec 3.2.2, were both variants Mamba-1 and Mamba-2 based WM trained with DFS?

Yes, We stated it in the main text line 398. We changed Figure 2's legend to further clearify it.

How does the proposed method compare with R2I [1]. Since they propose using SSMs, should it be included as another baseline?

Thank you for highlighting this reference; we will include it in our related work. As mentioned earlier, we currently lack sufficient computational resources to evaluate DRAMA on the POPgym and Maze baselines. However, in the Atari100k benchmark, DRAMA outperforms R2I. We believe Atari100k provides a reasonable benchmark to demonstrate the model's capabilities. That said, we agree that including POPgym and Maze baselines would enhance the evaluation, and we will consider these benchmarks in future work.

Unsupported claims about capturing long-term dependencies

Our approach is motivated by evidence in the literature showing that State-Space Models (SSMs) have the ability to capture long-term dependencies, making them particularly effective for long-range modeling tasks, such as those in the Long Range Arena [Tay et al., 2021; Gupta et al., 2022; Smith et al., 2023]. Sequence lengths in this domain range from 1,024 to over 16,000. Mamba1 and Mamba2 inherit this capability as they are SSMs. Related work [Deng et al., 2023] has demonstrated that SSMs excel as world models, effectively capturing dynamics in specially tailored environments designed to measure long-term memory capabilities.

To address the reviewer’s concern, we conducted an ablation study (in Sec 3.2.3) focusing on the critical components of the world model: the dynamics model. In a simple yet representative grid-world scenario, our results confirm that Mamba2 effectively captures long-term dependencies as expected. The long-term in the ablation experiment is refered as 1664 training sequence length. The result can be seen in Table 2 above.

Reference:

Smith, Jimmy T. H., Andrew Warrington, and Scott W. Linderman. "Simplified State Space Layers for Sequence Modeling." In The Eleventh International Conference on Learning Representations, 2023.

Tay, Yi, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. "Long Range Arena: A Benchmark for Efficient Transformers." In International Conference on Learning Representations, 2021.

Gupta, Ankit, Albert Gu, and Jonathan Berant. "Diagonal State Spaces Are as Effective as Structured State Spaces." In Advances in Neural Information Processing Systems, 35:22982–22994, 2022.

Deng, Fei, Junyeong Park, and Sungjin Ahn. "Facing Off World Model Backbones: RNNs, Transformers, and S4." In Advances in Neural Information Processing Systems, 36:72904–72930, 2023.

Scaling DRAMA for Direct Comparison

Scaled model comparisons require substantial computational resources and time, which were not available to us. One key motivation for presenting Mamba-2 as a world model is its parameter efficiency. While scalability is undoubtedly important, smaller yet efficient models offer significant advantages. For instance, model-based reinforcement learning (MBRL) often faces the challenge of model exploitation, where the behavior model exploits imperfections in the world model to achieve higher rewards by repeatedly reaching states where the world model is underfitting. A potential solution to this issue is training multiple models to estimate uncertainty in predictions. However, this approach requires smaller and more efficient models, making Mamba-2 well-suited for such future directions.

To specifically address concerns and provide a ‘like-for-like’ comparison, we trained a 12M version of DreamerV3 on the Atari100k benchmark and reported the results in the appendix. The results demonstrate that Mamba-2 achieved a significant advantage over this variant of DreamerV3 in the domain of small models on the Atari100k benchmark. The results are presented in Table 1 above, with the detailed table and training curves provided in Appendix A.1 of the revised version.

DFS Method Clarity

To address this concern, we conducted a study comparing uniform sampling and DFS in Drama across all games. The results demonstrate the effectiveness of DFS, which outperformed uniform sampling in 11 games, underperformed in 2 games, and tied in 13 games.

Hyperparameter Sensitivity

We agree that this is an important concern. In our evaluation, we used the default hyperparameters for Mamba-2, while the other components of Drama were configured similarly to DreamerV3, with one exception: the actor was set to half the size of the critic. However, model-based RL inherently requires significant computational resources for hyperparameter tuning due to the complexity of its components, including the autoencoder, dynamics model, and behavior policy. Given these constraints, we prioritised using the available computational resources for the other requested ablation experiments instead. We plan to evaluate the sensitivity of the model to other hyperparameter values in future work.

Code and Reproducibility

We will release the code repo once the anonymity is no longer applied.