MAMBA: an Effective World Model Approach for Meta-Reinforcement Learning

We thank you for your time and constructive review.

1. Regarding harder environments, we first clarify why the Reacher environment is complex and then describe three additional experiments (found in the new additional experiments supplementary material)

   • We emphasize that we chose Reacher because we wanted to demonstrate that Dreamer style architectures are able to handle tasks where the difficulty is attributed to the meta learning aspect (variation of the dynamics and rewards) and not necessarily the dynamics themselves because previous works showed that Dreamer is able to solve tough control problems effectively. Besides the inherent difficulty in Reacher1, when adding multiple goals the problem becomes much harder as can be seen by the difficulty of VariBAD \ HyperX to adapt to those environments.
   • We added two challenging environments, Humanoid-Dir and PandaReach. In both we see that MAMBA and the Dreamer versions are superior to VariBAD and HyperX both in return and in required environment steps to reach conversion. For full details about these new scenarios see the additional experiments supplementary file we added.
   • Finally, we also added additional experiments for a harder version of Rooms, with seven and eight rooms, for the purpose of discovering Mamba’s performance limit in this domain. We find that even here only Mamba is able to find Bayes optimal behaviors and the drop in return is due to the fact that more exploration steps are required for the additional rooms.

2. We also considered more common meta-RL benchmarks such as HalfCircle and Humanoid-Dir. Another important aspect is the sparsity of rewards – many meta-RL environments are dense, i.e. the task could be inferred from all \ many actions available to the agent, making the Bayes-optimal behavior trivial! Many of the scenarios we investigated (Rooms, Reacher) are sparse (i.e. sub-goals can only be inferred when encountered) thus require non-trivial exploration policies. Nonetheless, we conduct additional experiments on Humanoid (see above) as we appreciate the need for standard benchmarking.

3. Regarding “In the theory section, it is discussed that the local reconstruction could be harmful in high dimensional non-decomposable task distributions”: Previous works [2],[3] have shown that scenarios of high-dimensional task distributions are hard for every meta-RL algorithm. Therefore, it is not our method specifically that is limited, but every meta-RL approach. Despite that, some methods might be more tailored towards specific aspects of the meta-RL problem. In our paper we demonstrated that our algorithm handles decomposable tasks better than previous approaches, even though they are high dimensional, but we also showed that it is better in more common scenarios (e.g. Rooms-1 and Reacher-1).

4. Some examples of high dimensional task distributions are the decomposable tasks in our work; while each sub-task in itself is of lower dimension, solving a sequence of them is high-dimensional. Other examples are domains such as Meta-World or RLBench; these are hard for all meta-RL approaches, and the current SOTA for these domains are based on fine tuning, which requires orders of magnitude more samples at test time than the meta RL approach we consider.

5. We incorporated suggestions on presentation and proposed related work into the relevant sections in the paper. Please see the revised version (revisions are in blue text). We acknowledge that the specific paper mentioned by the reviewer, Pasukonis et al. [1], bears some similarities to our approach: the Dreamer (TBTT) baseline in the paper is a modification to Dreamer-v2 designed to consider long-term memory effects, and the MemoryMaze environment is somewhat reminiscent of our Rooms-N environment. However, Pasukonis et al. [1] aims to establish a baseline for testing memory capabilities of agents, rather than proposing an algorithmic solution to meta-RL in general. Therefore, the modifications they propose for Dreamer are only geared at extending its memory capabilities, disregarding other modifications we found important.

[1] Pasukonis et al., 2022, Evaluating long-term memory in 3D mazes

[2] Mandi et al., 2022, On the effectiveness of fine-tuning versus meta-reinforcement learning

[3] Rimon et al., 2022, Meta reinforcement learning with finite training tasks-a density estimation approach.

We thank the reviewer for the thoughtful insights and comments.

1. “The authors provide theoretical results comparing VariBAD (a meta-RL algorithm) with Dreamer (a model-based algorithm) on decomposable tasks which gives grounds for why such tasks are hard for model-free meta-RL algorithms.”

   To clarify our findings: our theoretical results show that while a general high-dimensional task distribution is inherently hard for any meta-RL approach (see [1] and [2]), the special case of decomposable task distributions (which are also high dimensional) have better guarantees. Further, we showed that in order to achieve those guarantees one needs to use a local reconstruction approach instead of a global one. As a final clarification, our theory doesn’t compare between VariBAD and Dreamer in particular, and doesn’t analyze the differences in guarantees between model-free and model-based approaches.

2. Thank you for pointing out these references. As mentioned in the general comment to all reviewers above, the field of meta-RL is rich and diverse, and we attempted to focus our related work section on what we deemed most relevant to our paper. For the sake of completeness, we have incorporated these additional papers into an extended related work section supplied in the appendix (Section F) of the revised version of the paper.

   We also compared MAMBA with $RL^2$, and a recent related method, DREAM. Both experiments can be found in the additional experiments supplementary PDF, and demonstrate that MAMBA is preferable over these two algorithms as well.

3. We have followed previous work ($RL^2$, PEARL, VariBAD, HyperX) and focused our investigation on common meta-RL benchmarks. We acknowledge the importance of standardized benchmarking in meta-RL; as it seems like Alchemy is a benchmark that is gaining traction, we will be happy to test MAMBA on it and report our results in a future version of the paper. That said, our current investigation already makes extensive comparisons between MAMBA and other meta-RL algorithms on well-established benchmarks.

Finally, we thank the reviewer for the helpful wording and styling suggestions, we revised some of the text as suggested.

[1] Mandi et al., 2022, On the effectiveness of fine-tuning versus meta-reinforcement learning

[2] Rimon et al., 2022, Meta reinforcement learning with finite training tasks-a density estimation approach.

Thank you for your time spent reviewing our work, and for your valuable insights.

1. Regarding the task decomposability assumption:

   We would like to clarify that MAMBA does not assume task decomposability, and works well for many meta-RL tasks that are not decomposable. The focus on decomposable domains in the paper is for analysis purposes: for this class of problems, we can understand why our approach is more suitable. Obviously, real world problems are not necessarily decomposable, and we provide sufficient evidence that MAMBA is beneficial even without this structure in the task. See our response to this effect in the general comment to all reviewers above.

2. “Isn't the fact that they should be performed in a sequence make the tasks dependent, and thus not independently decomposable?”

   Thank you for this excellent question, you are correct – in this case, every subtask may affect the initial distribution of future sub-tasks. While this is true, our theoretical results still stand as they are not affected by the initial distributions of the sub-tasks. For completeness and clarity we modified our theory and discussion sections to refer to tasks which are decomposable into independent sub-tasks up to the initial distribution.

   We further clarify that although the initial distribution is affected by previous tasks, the rewards and dynamics attributed to every sub-task are independent and are determined only according to the goals.

3. “The proposed algorithm is a minor modification over Dreamer”

   Correct. Our main contribution is showing that with relatively minor, yet tailored, modifications of Dreamer we achieve state-of-the-art results in meta-RL (with significant improvement over Dreamer-Vanilla). Moreover, we view the fact that we only needed minor modifications as a central contribution of this paper rather than a drawback, as one of our objectives is to show that Dreamer is successful in meta-RL settings as-is.

4. Regarding choice of the number of sub-episodes:

   For environments borrowed from VariBAD \ HyperX, we use the same number of sub-episodes as in these works. In new environments, we choose the number of episodes to be sufficient to explore all goals in the goal distribution.

5. Regarding stochastic environments:

   Indeed, all of the environments we have tested had deterministic dynamics and reward functions. However, there is no reason for stochasticity to make the task identification harder; to test this, we added an investigation of a variant of the HalfCircle environment, where a stochastic wind force is added at each step. We added the results to the supplementary experiments PDF and we also present them here for convenience:

As can be seen in the table, the return did not drop significantly from the original results in Table 1 of the paper, and the Dreamer variants are still superior to the baselines. All of the Dreamer variants perform close to Bayes-optimally. Some samples from the learned policy are presented in our anonymous website

6. “In Figure 1 (left most) why isn't the agent exploring the right portion of the semi-circle during the first episode?”

   As can be seen in our anonymous website in the far left example of the Point Robot section, after missing the goal the agent finds it in later episodes. It is likely that it is suboptimal to “waste time” exploring the far right of the half circle first as it will lower the return for the rest of the location for the off-chance the goal is actually there. Also, from the HyperX paper (figure 7c) we can see the same behavior appears in HyperX and VariBAD - the agent does not explore the far end of the half circle in the first exploration episode, suggesting that it is indeed suboptimal to do so.

We thank the reviewer for taking time to review our paper, and for their thoughtful suggestions and insights. We have added multiple experiments during the rebuttal period, the results of which are described in the general comment above and can be viewed in the rebuttal PDF (uploaded in the supplementary materials). Central results are described below:

1. Ablation Study of MAMBA:

   • We first note that it is infeasible to run experiments without adding the reward into the observation because it prevents the agent any possibility to distinguish between tasks – so we cannot remove this feature in the meta-RL context; to validate this, we removed the reward observation and tested on the Rooms-3 environment. The results are as expected:

   	MAMBA	MAMBA w/o Reward Observations
   Average Total Return	156.2$\pm$ 1.7	11.5$\pm$ 1.4

   • We note that Mamba with the original batch sampling (i.e. batches do not start from the start of the meta-episode) is in fact the second baseline (Dreamer – our params).
   • Finally, we’ve tested the effects of removing the horizon scheduling in the HalfCircle scenario, and indeed found that the results obtained are on-par but runtime is slower (on the same hardware):

   	MAMBA	MAMBA w/o Horizon Scheduling
   Average Total Return	156.2$\pm$ 1.7	156.8$\pm$ 1.2
   Total Time [H]	34.1	44.0

   • Due to space constraints, we added these experiments to the appendix (Section G).

2. We compared to DREAM as suggested by the reviewer. We tested DREAM in the Rooms environment as this was the most similar to the Minigrid environments in the original paper. We found that DREAM was unable to produce a useful policy. See more details in the additional experiments PDF in the supplementary material.

Next, we respond to the questions raised by the reviewer:

3. Regarding the cause of the gap in performance between VariBAD and Dreamer:

   • Dreamer was originally designed for POMDPs which meta-RL is a flavor of, it is not very surprising that Dreamer is effective in this context. In addition, Dreamer can be seen as an effective out-of-the-box algorithm with many moving parts, and it is difficult to understand the individual contribution of each part. Moreover, it has been used as-is in many other settings with great success. Finally, a full comparison between Dreamer and VariBAD is detailed in Section 3.3 in our paper.

   • Regarding exploration/exploitation: we have modified Figure 4, adding the final performance of VariBAD and HyperX. We note that MAMBA has both better exploration (as seen by the high rewards in the first episode of every meta-episode) and better exploitation (as seen by the higher total returns of the different sub-episodes).

   • “The tasks in this work take a very particular structure (they decompose nicely into iterated tasks). Is this a structure that clearly benefits one of the algorithms?”

     As shown by our theory Section 3.2.1, a method with local reconstruction is expected to work better than one with global reconstruction on decomposable tasks. We also showed that while local reconstruction is beneficial for VariBAD as well, dreamer still outperforms it (Section 3.2.2).

     “What happens when this structure is violated?”

     Our results show that Dreamer works better on a variety of task structures including non-decomposable meta-RL tasks. See our comment to this effect in the general response to all reviewers above.

4. Regarding the local vs. global reconstruction:

   • In Section 3.2 of the paper we do analyze, both theoretically and empirically, when and why local reconstruction should be helpful (i.e decomposable tasks).
   • There is a difference between encoding and reconstruction. Our work also encodes the whole history (from the beginning). Every meta-rl algorithm should do that in order to pass information from previous episodes to future ones.

Thanks to the authors for the thorough response. Overall, I find the additional experiments and discussions to be useful and interesting. I remain in support of accepting the paper, though I am not willing to raise my score all the way to an 8 (the next possible score).

• It’s worth noting that the most novel proposed change in this paper doesn’t seem to be extremely important for improved performance based on the authors new ablation studies. This is fine, but worth noting.
• I am somewhat skeptical of the reported DREAM results. The rooms task is of similar horizon and arguably simpler complexity than the tasks considered in the DREAM paper. Hence, I find it surprising that the authors report that DREAM fails entirely. I find the authors’ explanation of hyperparameter tuning etc to be plausible, though the comparison seems somewhat sloppy if baselines aren’t tuned for the specific domain they’re being applied to. It also seems possible to be due to implementation error, since the DREAM paper itself and follow ups also suggest relative robustness to hyper parameters. I don’t particularly hold this against the authors and still support acceptance.