MAMBA: an Effective World Model Approach for Meta-Reinforcement Learning
Authors: Zohar Rimon
Tom Jurgenson
Orr Krupnik
Gilad Adler
Aviv Tamar
Venue: ICLR 2024 poster (2024)
Keywords: Meta Reinforcement Learning
World Models
Model Based Reinforcement Learning
Primary area: reinforcement learning

Abstract: Meta-reinforcement learning (meta-RL) is a promising framework for tackling challenging domains requiring efficient exploration. Existing meta-RL algorithms are characterized by low sample efficiency, and mostly focus on low-dimensional task distributions. In parallel, model-based RL methods have been successful in solving partially observable MDPs, of which meta-RL is a special case.
In this work, we leverage this success and propose a new model-based approach to meta-RL, based on elements from existing state-of-the-art model-based and meta-RL methods. We demonstrate the effectiveness of our approach on common meta-RL benchmark domains, attaining greater return with better sample efficiency (up to $15\times$) while requiring very little hyperparameter tuning. In addition, we validate our approach on a slate of more challenging, higher-dimensional domains, taking a step towards real-world generalizing agents.

MAMBA: AN EFFECTIVE WORLD MODEL APPROACH FOR METAAREINNFRREMENT LEARNING
Zohar Rimon Tom Jurgenson' 1+
Orr Krupnik'
Gilad Adler?
Aviv Tamar'
'Technion Israel Institute of Technology 'Ford Research Center Israel
ABSTRACT
Meta-reinforcement learning (meta-RL) is a promising framework for tackling challenging domains requiring efficient exploration. Existing meta-RL algorithms are characterized by low sample efficiency, and mostly focus on low-dimensional task distributions. In parallel, model-based RL methods have been successful in solving partially observable MDPs, of which meta-RL is special case. In this work, we leverage this success and propose a new model-based approach meta- RL, based on elements from existing state-of-the-art model-based and meta-RL methods. We demonstrate the effectiveness of our approach on common meta- RL benchmark domains, attaining greater return with better sample efficiency (up to 15x) while requiring very little hyperparameter tuning. In addition, we validate our approach on a slate of more challenging, higher-dimensional domains, taking step towards real- world generalizing agents.
INTRODUCTION
There is a growing interest in applying reinforcement learning (RL. agents to real-world domains. One requirement of such applications is for the agent be able generalize and solve scenarios beyond its training data. Meta-reinforcement learning (meta-RL) is line of work aiming to allow RL agents to generalize novel task instances, attempting to learn general control policy over family of tasks. The term meta-RL covers wide range of approaches (Duan et al. 2016 Finn et al., 2017 Zintgraf et al. 2019) which have produced increasingly impressive results in recent years.
One line of work aiming to solve meta-RL problems is context- basex meta-RL (Rakelly et al. 2019 Zintgraf et al. 2019 2021). These algorithms encode trajectories of training tasks
context or belief variable, and infer it from data at test time in order to identify the task at hand and solve it. This approach is Bayesian in nature, attempting to balance exploration (to identify the task parameters) and exploitation (to succeed in the current task), While state-of-thed
results have been reported on simulated robotic domains (Zintgraf' et al. 202T; Hiraoka et al. 202T), previous meta-RL methods tend to be sample inefficient, as they rely on model-free RL algorithms. In addition, they have mostly been shown to succeed only in dow-dimensional task distributions.
Model-based RL is known to be more samplc-efficient than its model-free counterpart (Janner al. 2019 Ye et al. (02T). By their nature, model--ased RL algorithms may also be more flexible, they learn model of the environment and use it to plan or learn policies, instead of learning control policies by directly interacting with the environment. A recent line of work in model-based RL proposed Dreamer, an approach which has shown versatility a variety of domains (Hafner et al. 2019a 2020: 2023). Dreamer has delivered the sought-afte after sample efficiency of model-based RL on variety of domains, while requiring very little hyperparameter tuning. addition, due to its recurrent latent encoding of trajectories, it has been successful in solving partially observable Markov decision process (POMDP) scenarios (Hafner al., 2020: 2023). POMDPs can be viewed as a general case of the -aatext-based nI. Formulation (where the context identifying the task is unobserved by the agent). In light of its favorable propertics, we believe that Dreamer a good fit for meta-RL. Therefore, in this work we propose to use Dreamer to solve meta-RL scenarios. To the best of our knowledge, this is the first use of the Dreamer architecture in meta-RL.
We compare Dreamer to oftteoof-hhe-rr neta-RL algorithms (Zintgraf al. 2019 2021), and pro- Algorithm), an algorithm leveraging the best compo- and tomjurgen-
pose MAMBA (for MetA-RL
"Contributed equally, correspondence y @ @omggmail.cmm Code available at: htwp://ps./gggttbbb.moo////aaaae

---

Behaviors learned by MAMBA in various domains. In all of the figures, the first, second and third episodes shown in red, green and blue, respectively. From left to right: Point Robot Navigation, Escape Room, Rooms-N and Reacher-N (plotted point rep- resents the end-effector; see Sec. 4 for details). As the environments have sparse rewards, MAMBA evidently learns near- near-Bayes Bayes-optial behavior: it explores the environment in the first episode and exploits the information in the subsequent ones. For additional visualizations https://s si tes.
Figure 1:
nents of each. MAMBA outperforms meta-RL and model-based RL baselines on several well- studied meta-RL benchmarks, both terms of task success and in sample efficiency. In addition, we propose a new class of challenging meta-RL domains with high- -igeessinnii task distributions, which can be decomposed into lower-dimensional sub-tasks. By bounding the number of tasks re- quired to obtain near- Baaa-Bays--ptimal behavior (based on the analysis of Rimon et al. 2022), we show that the ability to decompose the task into sub-tasks crucial for solving such domains efficiently. We then empirically show that MAMBA can decompose the problem better than existing meta-RL approaches, which leads to significantly better performance.
To summarize, our main contributions are:
1. We present shared formulation of context- meta-RL algorithms and the Dreamer line of work, highlighting their similarities and differences.
2. We develop MAMBA, a sample-efficient meta-RL approach which outperforms a variety of baselines on meta-RI. RI. benchmarks, while requiring very little hyperparameter tuning.
3. We highlight an interesting class of meta-RI. environments with many degrees of freedom that is difficult for existing meta-RI. methods to solve. We analyze these domains theoreti- cally, and demonstrate the ability of MAMBA to solve them efficiently.
BACKGROUND
In this section we present background on partially observable Markov decision processes (POMDP) and meta-RL (Sec.) 2.1. Next, we describe context- based algorithms and model- based RL (Sec. 2.2 which are popular approaches for solving meta-RL and POMDPs respectively.
2.1 PROBLEM FORMULATION
Partially observable Markov decision processes (POMDP, Cassandra et al. 1994) are sequential decision making processes defined by tuple (S.A, Po. P. R, T.2. O), where S and A are the state and action spaces respectively. The process initialized at a state sampled from the initial state distribution s ~ Po and evolves according to the Markovian state transition probability ((s'(s,a) and reward function R(rls, a) for T time steps. The states are unobserved by the agent, which has access to observations from the observation space Qt sampled each time step from O(oss). Let = (Oy, 01, 71, o1.... de, Te, oz) the history up to time step t. The objective to learn a policy (at+1/T)) which maximizes the expected return J E.P..O Σr r't where E,,P,0 denotes the expectation with respect the policy, observation function and POMDP transitions.
In reiffreiiforremmnt learning (meta-RL) we consider distribution p(M) over a space of MDPs M, where every MDP is tuple (S, A, po, PM, RM. T, K). Here S and A are the state and ac- tion spaces respectively, shared among all MDPs, while (ar (s'|s, a) and RM (r|s, a) are the MDP- specific transition and reward functions. The agent plays K consecutive episodes in the same MDP, each consisting of time steps. The full horizon = K T comprises the meta-episode, with each of its K sub-trajectories denoted sub-episodes. The objective is to learn policy N (00+11/T)
2

---

ri Meta-RL can be seen as a special case of a POMDP, where the unobserved components are Pas and RMs, and the transitions are episodic - every T time steps the state is reset.
which maximizes the expected return J EM-p
E
E/=
Bayes-Adaptive MDP: Solving POMDPs is known to be intractable (Papadimitriou & Tsitsiklis, 1987]. While an optimal policy is generally history-dependent, a convenient representation of the optimal policy uses a sufficient summary of the history (also called sufficient statistic or information state; Subramanian & Mahajan 2019) One such representation is the belief, b, = p(R. Rt P(ra), the posterior probability over R. P given the history of transitions up to time t. The POMDP can be cast as an MDP over the belief space, known as Bayes- adyes--d MDPs (BAMDP, Dufi 2002). In BAMDP, the state 8, is augmented with the belief by and is denoted as the augmented state 87 (st, be). Formally, the BAMDP is tuple (S+, AA,P P+, R+,T), T), where S is the space of augmented states, A is the original action space, and PT. P+, and R+ operate over augmented states. An optimal policy iT To E arg max, J is termed Bayes-optimal, and is history-dependent in the general case.
2.2 CONTEXT-BASE BASED META-RL RL AND NODEL-BASED ALGORITHMS FOR POMDP
Context-based meta-RL: One popular approach to solving meta-RL. is encode the observed history into a latent representation which approximates the belief at time I, b, and learn a belief- conditioned policy in order maximize the return J
A recently successful context-based algorithm, VariBAD (Zintgraf et al. 2019), learns b, using Variational Autoencoder (VAE, Kingma & Welling 2013): recurrent neural network (RNN) encoder encodes T3 into the latent space of the VAE which is interpreted as b,. Unlike other VAE formulations which only reconstruct the input (namely re). VariBAD reconstructs the full trajectory Tr from b,. Meanwhile, a policy is trained on the augmented state (s,, 2,) using model-free RL algorithm (specifically PPO. Schullmn al. 2017). For stability, the gradients do not flow from the RL objective through 2, into the encoder, i.e., the belief is shaped only by the reconstruction signal.
Recurrent latent model-based algorithms for POMDP: Model-based RL. (MBRL) algorithms model the dynamics and rewards R given data, and then use the learned models P and ft to imag- ine future outcomes and optimize a policy. This approach has been shown empirically to be very sample-efficient (Schrittwieser et al. 2020: Ye et al.) 2021| Hafner et al. (023). In particular, latent MBRL algorithms are appealing approach to handle partial observabilitg lity by first encoding a quence of previous observations т: into a latent space 2, E Z, and learning dynamics P(z [ze, ar) and rewards A(relz,, at) directly in the latent space. This approach offers two main advantages: first, by encoding sequences of observations we allow to recall information from previous time steps which might help with the partial observability aspect. Second, imagined rollouts used to optimize the policy are defined over much smaller space (|Z) << [2]), allowing more efficient planning.
2023). In Dreamer, the world model w is a type of RNN dubbed recurrent state space model (RSSM, Hafner et al. 2019b), which predicts observable future outcomes from sequences of past observations. Concretely, random sequence of transitions of length I., TEs+L is encoded into a latent space 2.++ using the RSSM. Then, for every & E [ [i...., L.,, Dreamer predicts from 21 the next reward r++1 and observation 0r-11 and minimizes Lpred = - log w(rr+ - log w(ot+1lze).
One such family of algorithms is Dreamer (Hafner al. 2019a 2020
In parallel, Dreamer trains policy conditioned z, which maximizes the returns of imagined future latent-space trajectories of length N/aa. Finally, as in VariBAD, gradients from the policy do not propagate through to the world model
3 METHOD
In this section we present MAMBA. We start by motivating two main reasons for selection of Dreamer (Hafner et al.] (023) as the backbone of our method. First, in Sec.B.II] we demonstrate that model-based algorithms and Dreamer in particular in fact structurally similar to state-of-the thett
3

---

meta-RL algorithms. Then, in Sec. 3.2 explain why in some meta-RL RL tasks the world model learning in Dreamer is advantageous. Finally, we present MAMBA in Sec. 3.3
3.1 COMPARING VARIBAD AND DREAMER
VariBAD meta-RL algorithm, while Dreamer is a model-based RL method intended to solve POMDPs. Despite that, they share many similar attributes. We next describe the commonalities and differences between them.
At the core of both algorithms lie two fundamental components: a recurrent predictor for observable outcomes and condit-conditione policy. The recurrent predictor is realized by the belief encoder in VariBAD and by the world model in Dreamer. Both encode sequence of observations into latent space representation 2 using a RNN structure. Then, 2; is used to predict observable outcomes such as future observations and rewards. The latent-couditioned policy is trained via online RL objectives by predicting actions given Zt.
Despite their similarities, there are crucial differences between VariBAD and Dreamer. First and foremost, during training, VariBAD uses its recurrent predictor to encode the full history ти starting at the very first time step of the meta-episode, while Dreamer encodes sub-trajectories that do not necessarily start at the first time step. From a computational standpoint, this makes Dreamer less computationally demanding. However, in meta-RL scenarios it could create unnecessary ambi- guity as hints regarding the identity of the MDP (namely Par and RM), collected during time steps predating the sub-trajee aub-rrajectory are ignored. We refer to this as the encoding context.
A second key difference is the recurrent predictor reconstruction window. In Dreamer. the world model performs local reconstruction (predicts single step into the future), while VariBAD performs global reconstruction (predicts entire trajectory, past and future, from every 2,).
The differences in encoding context and local vs. global reconstruction inspire our algorithmic improvements in MAMBA (see Sec. 3.3 Other technical differences are described in Appendix B
3.2 DECOMPOSABLE TASK DISTRIBUTIONS
In this section we focus on the difference in reconstruction between VariBAD and Dreamer. To this end we consider specific family of task distributions and explain why they are difficult for meta-RL RL algorithms to solve, by extending the theoretical analysis of Rimon et al. (002). We then show that utilizing the decomposability of these tasks makes them amenable to much more efficient solution.
Finally, we relate this decomposability to the local reconstruction horizon, and posit that Dreamer should solve these tasks more effectively. We corroborate this in experiments (Sec.
3.2. PAC BOUNDS
Meta-RL has been shown scale badly with respect the number of degrees of freedom (DoF) in the task distribution (Mandi et al.) 2022 Rimon et al. (2022) devised PAC (Probably Approximately Correct) upper bounds on the suyyess-uoopiimalttyy with respect to the number of training tasks. They showed exponential dependency the number of DoF of the task distribution.
Many real-world domains with a high number of DoF can be decomposed into sub-tasks, each with a low number of DoF. For example, robot may be required to solve several independent problems in a sequence. In this section we analyze this family of task distributions, which we term decomposable task distributions (DTD), and prove that they have significantly better PAC bounds than the general case. In addition, we demonstrate that on DTD the local reconstruction approach taken by Dreamer advantageous compared the global reconstruction in VariBAD.
Bounds on Bayes-suboptimality for DTD: For task that can be decomposed to K independent tasks (up to their initial state distribution), each with D DoF, a naive approach (such as Rimon et al. 2022) would lead to a bound that depends exponentially on O(D x K). However, as we demonstrate in the next theorem, a method that takes task independence into account obtains an O(D) and linear dependency independently (proofs and full details in Appendices A A.2). We follow the formulation in Rimon et al. (2022) and assume access to N samples from the prior task distribution f. We estimate
exponential dependency
K, equivalent solving each task
4

---

using density estimation approach, obtaining f, and find the Bayes-optimal policy with respect to f. i.c Ti E arg max J = arg max E M- [B,, M >
Theorem (informal) Under mild ussumptions, there exists a density estimator f which maps N samples from decomposable task distribution distribution over the task space, such that with probability at least L/N:
(('8) log N N
R(az)) S.
Where R is the sub-optimality with respect to the true prior distribution and is constant which depends on the smoothness of the prior task distribution.
The key to proving Theorem is to estimate each of the sub-task distributions individually, and not the entire task distribution at once, as was suggested by Rimon et al. [222] (which leads to bounds with exponential dependency on D K instead of D). This theoretical result is in line with our hypothesis - concentrating on the current sub-task (or performing local reconstruction) at each time step crucial for DTD.
For example, in the Rooms-N environment (see Sec 4. although the total number of DoF is N (number of rooms), the difficulty of the task does increase, as the agent should solve each of the rooms separately. In comparison, our theorem states that a standard N N-dimensional navigation goal, which has the same number of DoF, should be much harder since all DoF must be estimated once.
3.2.2 DECOMPOSABILITY AND LOCAL RECONSTRUCTION
As described in Sec. 3.1 to learn latent representation, VariBAD utilizes global reconstruction approach, whereas Dreamer utilizes a local one. Based on Theorem we hypothesize that the latter approach is better for DTD (as the former approach reconstructs multiple sub-tasks once).
To test our hypothesis, we compare VariBAD with the original reconstruction loss to a variant of VariBAD reconstructing only small window of 10 transitions around the current time step. We evaluate the two variants on the Rooms-N environment with 4 and rooms, where the agent should explore consecutive rooms, passing to the next room after solving the current one (see Sec. for details). We observe that using the local reconstruction is favorable in terms of total return in both the Rooms-4 and Rooms-5 scenarios (see Sec.
To further analyze this phenomenon, we visualized the reward prediction using the learned decoder
in VariBAD in each state of the state space and for every encoded belief (Fig. 2). We observe that VariBAD (without local reconstruction) wrongfully predicts rewards in adjacent states to the agent (Fig. 2 top (b), (f), and thus the agent stays in the same location without a goal (as indicated by the constant state index). In comparison, introducing local reconstruction makes the local reward prediction much more accurate, as indicated in Fig. 2 bottom (d), (e). However, comparing top (f) and bottom (f), local reconstruction hurts the prediction of rewards in distant states. This helps the local reconstruction agent to better focus on the goal and follow nparmoptimal behavior. For further quantitative analysis Appendix E This analysis shows the efficacy of local reconstruction in comparison to a global one on DTD, which motivates our design choices in the next section.
3.3 MAMBA: META-RI. MODEL BASED ALGORRITHM
In this section we describe MAMBA, our approach incorporating the model-based based Dreamer into the meta-RL settings. We base MAMBA on DreamerV3 (Hafner et al. (223) since it was shown be high-peeffrmance, sample-efficient and robust algorithm applicable out-of-the-box to many partially observable scenarios with various input modalities. Moreover, described in Sec. 3.1 the RSSM (Hafner et al. 2019b) of the Dreamer architecture plays the same role as the belief encoder in context-based meta-RL algorithms ZZntgraff al. 2019 2021 allowing easy application of tools and conclusions from these algorithms MAMBA. Finally, the world model Dreamer learns by reconstructing single step into the future, a preferable property compared to algorithms which use a long reconstruction horizon (Sec. 3.2). Next, we describe the modifications on top of Dreamer to create MAMBA:
5

---

Modifying trajectories to support model-based meta-RL algorithms: We augment the observations o, with the reward and the current time step of = [ot, r't, f]; both readily avail- able in any environment and do not require any special assumptions. In meta-RL it is common to incorpo- rate the rewards into the observation (Duan et al.] 2016) Rakelly et al.] 2019 Zintgraf et al. 2019) as it helps in identifying the sampled MDP. As for adding t, it allows to model the reset of the environment after every T steps (see Sec. 2
Seaccsr
stop
(b)
Sampling full meta-episodes: To avoid missing signals about the iden- tity of the current MDP (see encod- ing context in Sec. (.I). we mod- ify Dreamer keep track of all information from the entire meta- episode when encoding 24 using the RSSM. The original Dreamer only encodes 2t for limited random se- quences of actions; it samples a tra- jectory: t (meta-episode) from the re- play buffer, samples a random start- ing index t, and then encodes fixed- length sequence of L, 64 transi- tions obtaining L- We emphasize that using this approach, information gathered at steps 1, ..... t- is discarded. Instead, in MAMBA, we randomly sample the meta-episode т, and feed it in its entirety into the RSSM, obtaining =1:H- Although this creates computational burden, as experiments show, this allows MAMBA to obtain superior performance to Dreamer (see Sec, 4.
Time stop
(P)
Figure 2: Reward prediction at current location (green), state identifier (bluc), and goal state indication (red) through time on the Rooms-4 environment. Top: vanilla VariBAD (global reconstruction), bottom: VariBAD with local recon- struction. At the highlighted time steps (indicated by rows) we plot the reconstructed reward for every grid cell, where positive reward prediction is represented green, and the agent trajectory is drawn in blue.
To alleviate some of the computational load, after obtaining 21:H when optimizing the RSSM and policy we sample L indices ii.....iz iL ~ frri)))((H)) and optimize (maximize reconstruction the RSSM, and maximize the return for the policy) only according to Ziy...., 2,, Formally, let L. and L- be the loss functions for the world model w and the policy æ from sin- gle latent vector z. Then, in Dreamer the world-model loss is Eit/ Lau((e) and the policy loss DitL L (21), while in MAMBA the world-model loss El= (2i,) and the policy loss is EL (. Thus. despite the difference in RSSM input sequence lengths (L in Dreamer and H in MAMBA), during optimization Dreamer and MAMBA are updated according to the same number of latent vectors 24-
Scheduling the world model horizon: We found that during the initial update steps, unrolling the world model for the entire meta-episode length H is inaccurate, due the inaccuracy of the yet-untrained RSSM. In light of this inaccuracy, we reduce the computational load by learning only from prefixes of meta-episodes TH' (where H' is the prefix length). As the world model improves, we can increase H' and roll the model out for more steps. We therefore introduce a schedule: let Ny be the total number of environment steps, and define NwM << Ny. The world model horizon H' increases linearly from some initial value of W Minstial at the start of training to the full meta- episode length when Nw as environment steps are taken.
Although the three modifications from Dreamer to MAMBA may appear technical, they are mo- tivated by the concrete differences between general POMDP and the meta-RI. problem. In the following section, we show that they also lead to dramatic performance improvement in meta-RI., over both the original Dreamer and state-of-the-an meta-RL methods.

---

In this section we provide an investigation of the effectiveness of MAMBA. We first show that MAMBA is a high-performing meta-RL. algorithm: compared to baselines it obtains high returns with less environment steps. Next, we focus on several meta-RL scenarios including ones with sub- tasks, varying dynamies and visual observations, and show that MAMBA outperforms all baselines in each. We also investigate properties of MAMBA compared to the baselines, such as robustness to hyperparameters and the ability utilize information from previous sub-episodes. For visualiza- and for 4 EXPERIMENTS
tions of our results see https: s sites.etseee google com com/view http:///eew//mmaaaaeaeaee additional environments, experiments and details see Appendices and
Environments: We use two common 2D environments in meta- RLL Point Robot Navigation
(PRN), and Escape Room (Zintgraf et al. 2019 Dorfman et al.] 2021; Rakelly et al., 2019). In both, goal is randomly chosen on half-circle and the agent must explore the circumference until reaches it. Point Robot Navigation this point is a source of positive reward, and Escape Room, it is a door that leads to high rewards located outside the circle. These represent basic sce- narios for varying rewards and varying dynamics in meta-RL.
Next, to test our DTD hypothesis (Sec. 3.2) design two scenarios composed of sub-tasks, one discrete (Rooms-N), and one continuous (Reacher-N) where N represents the number of sub- tasks. In both tasks, set of N goals are selected independently, and the challenge is to reach them in the correct order and memorize them between sub-episodes. Crucially, the first N - 11 goals only provide positive reward once, thus encouraging the agent to find all N goals and stay on the last one as long as possible. In Rooms-N the minigrid agent starts in the leftmost room and the i-th goal in the i-th room from the left (thus the last goal on the rightmost room). In Reacher-N, the agent (adapted from the DeepMind Control Suite, Tuunyasuuunakoo et al. 2020) is a two-link robot and goals are placed in the reachable regions around Full details can be found Appendix C See Fig. D for visualizations of the scenarios and behaviors learned by MAMBA.
Baselines and evaluation protocol: To demonstrate the effectiveness of MAMBA, we compare it with the following baselines: two state-of-the-art meta-RL algorithms, VariBad and HyperX; an unmodified DreamerV3 which we dub Preamer-Vanilla. and an otherwise unmodified Dream- erV3 with hyperparameters identical to MAMBA, which call Dreamer-Tune. Because tried to make MAMBA as light as possible, the hyperparameter difference between Dreamer-Tume and Dreamer- Preamer-Vanilla is in network sizes and policy imagination horizon (see Appendix D for details).
In all our experiments, we report the mean and standard deviation of four random seeds, as well as the number of environment steps each algorithm took to converge. In every experiment test the best model seen during evaluation on a held-out test set of 1000 tasks.
MAMBA is an effective meta-RL algorithm: Our main results are described in Table 1 First, we observe that both VariBAD and HyperX, state-of-the-art algorithms specifically designed for meta-RL, perform worse than the model-based methods (both versions of Dreamer and MAMBA), both in final return and in convergence rate. This reinforces our observation regarding the similarity between both families of algorithms, and we attribute the performance gap to better algorithmic design choices in the Dreamer architecture (such as RSSM instead of simple RNN, and model-based based policy optimization).
Next, we find that further fine-tuning Dreamer to meta-RL scenarios again boosts performance; fo- cusing on Dreamer- Vanill--Vanilla, Dreamer- Tune, and MAMBA, we see that inconclusive whether Dreamer- Vanilla and Dreamer-Tune is the better model. However, it is clearly evident that both are inferior to MAMBA. We note that the differences between Dreamer-Tune and MAMBA are algo- rithmic only (see Sec. 3.3 clearly demonstrating favorable algorithmic design choices in MAMBA.
Finally, we note that for the sub-tasks scenarios in partieular, MAMBA achieves much better results than all basclines, in line with our hypothesis in Section 3.2
We evaluate robustness to hyperparameters on the Reacher- sceraio. Here, we take three VariBAD and HyperX hyperparameter configurations from the original pa- pers and we compare them to both Dreamer and MAMBA. The hyperparameter configuration of Dreamer- Vanillr-Vanilla taken from Hafner al. (202)), which is almost identical Dreamer-Tune
Robustness of MAMBA:
7

---

Total Raturn
Total Roturn
Return
Eval Retur
Hypert
MANSA
Frames
Time[h]
Figure 3: Returns in Reacher-1: Left: sample efficiency meeta-pisoddeeee returns against the number of frames played. Since VariBAD and HyperX took 15x frames to converge we show their limit instead of full plot. Right: time efficiency meta-episode returns against training time (hours).
Eploode Rerurn
Prent (100 hamns
Pramt (10P Srarts
Figure Sub-episode Returns of MAMBA, VariBAD, Hyper and Dreamer-Vanie Vanilla in Reacher-1 (the meta-episode contains three sub episodes). Return of episode 0, 1, and 2, are left, middle, and right respectively. The gap between MAMBA and Dreamer demonstrates that MAMBA retains information from previous sub-episodes better than Dreamer.
and MAMBA (see Appendix D) Results for their respective configurations are in Table T As for VariBAD and HyperX, Table I contains the best configuration which was also used for the other Reacher tasks (473.7 27.5 for VariBAD, and 505.9 + 36.0 for HyperX). The other two configura- tions of VariBAD achieved 104.6 + 34.3 and 20.4 + 2.4. In HyperX, one configuration resulted in unstable policies obtaining arbitrarily negative returns, and the last configuration obtained 14.3++0.4. We conclude that these algorithms are highly sensitive to their hyperparameters.
This experiment hints that for new meta-RL scenarios, MAMBA and Dreamer provide good per- formance out-of-the-box without extensive hyperparameter tuning, while VariBAD and HyperX are more affected by hyperparameter tuning (see Figure 3 left).
Is MAMBA better at exploitation than Dreamer? In Sec. 3.3 we hypothesize that when al- lowing the RSSM to encode the full meta-episode, information regarding the MDP identity can be effectively retained between sub-episodes. In Fig. 4 we plot the return per sub-episode in the meta- episode. Returns for the first episode (left) are similar (hinting that both Dreamer and MAMBA learn to explore well), but MAMBA better utilizes the information gathered in previous episodes in the subsequent episodes reaching higher returns (middle and right).
In this scenario we leverage the ability of the Dreamer architecture handle visual inputs. We consider two scenarios, a single goal Reacher scenario with two sub-episodes with image observations, and Panda Reacher (Choshen & Tamar 2023). where a 7-DoF robot searches for goals with its end-effector. It is worth noting that VariBAD and HyperX do not support visual inputs. We ran each variant in Reacher and Panda Reacher for 15M and 20M environment steps respectively. Returns are 62.7+ 63.4 4 and 26.3-+266.6 in Dreamer- Vanilla, 62. 62.74-63.4 and 1 in Dreamer- Treamer-Tune and 215.6 64.8 and 137.6 0.7 using MAMBA, correspondingly. Both scenarios show that MAMBA maintains superior returns in visual settings as well.
Visual domain:
Runtime limitation: One limitation of our approach is runtime to recall information from the start of the meta-episode, we unroll the RSSM over the entirety of the meta-episode. In some scenarios, this requires us sequentially encode orders of magnitude more transitions.
8

---

VariBAD
HyperX
Dreamer- Vanilla
Dreamer-Ture 239.7 3.1
MAMBA (Ours)
PRN:
218.0 + 26.0
204.0 + 19.9
241.1 9.9
242.2 + 7.9
~10M
~IM
~2M
~2M
~2M
Rooms-3:
158.7 2.6
111.8 + 20.2
137.5 6.2
142.6 4.1
156.2 + 1.7
Rooms-4:
88.1 39.2
14.1 + 3.0
1.7.
119.8 3.1
136.7 1.4
Rooms-5:
0.9 = 2.5
-15.8
96.5 2.5
93.9 + 5.3
113.1 + 5.7
Rooms-6;
-11.0 0.9
0.7 ~1OM
69.5 6.3
71.0 6.3 ~6M
94.5 = 1.2
~100M
~6M
~6M
Reacher-1:
27.55
503.9 36.0
503.4 + 69.0
655.5 12.3
Reacher-2: 2:
33.2
30.0
247.4 + 80.5
217.6 + 64.3
285.8 + 89.6
Reacher-3:
0.2
183.6 100.0
76.9 = 80.5
325.0 + 47.0
Reacher-4:
0.0 0.0
-0.5 1.1
0.0
0.1 0.2
77.7 61.1
~100M
~100M
~00
~10
~TO
EscapeRoom:
70.7 1 5.3
66.9 1 6.5
68.2
73.2 1 7.8
73.9 3.1
200
~20
~4M
~4M
~MM
Table 1: Total return comparison of VariBad, HyperX, DreamerV3, and MAMBA on different meta environments. We report the final reward (mean + std), and the number of time steps until conver- gence (below dashed line).
We investigate the runtimes of each of the baselines on the Reacher scenario with a single goal. Fig. 3 (right) shows the return of each algorithm compared to the runtime of the algorithur First, we note that both Preamer-Tune and Dreamer- -anaile-Vannlla are faster than MAMBA (Dreamer-Tuus Tune is slightly faster than Dreamer-Vaniee better than both VariBAD and HyperX. The time gap between Dreamer and MAMBA could be further reduced, which is an interesting direction for future research. Another direction for future rescarch is to "chunk" the meta-episode transitions in order to reduce the length of the RNN rollout. This requires further investigation, as it removes contextual information which may be required.
due to reduced NN sizes). However, MAMBA is on-par or
RELATED WORK
Meta- Reittreeeifrrr Learning (meta-RL) has been extensively studied in recent years DDuan al., 2016: Wang al.) 2017; Finn al., 2017). Of specific interest to our work are context- based meta-RL, RL approaches (Rakelly et al.
et al., 2019 2021), which attempt to learn Bayes-optimal policy, and are model freee thus suffering from high sample complexity.
Model-based RL has been shown to be more sample efficient than its model-free counterparts Nagabandi al., 2018 Janner et al., 2019 Schrittwieeer al., 2020 Ye et al. 2021). While there
are many different approaches to model learning. Dreamer is prominent line of work (Hafner et al.,
2019a by 2020 2023 Mendonea et al.) 2021). which learns recurrent latent space world model and uses rollouts in latent space train policy. This approach has been shown to work well in POMDP environments, owing to the inherent RNN history aggregation ability (Ni et al. 2022).
Model Based Meta-RL: Recent approaches (Nagabandi et al. 2018 Perez et al. 2020 Lin et al., 2020 Hiraoka et al. 202I all continuously adapt to a varied task distribution, making no distinction of when new task is attempted. In accordance, they all use short histories of states and actions, limiting usability when long horizons are required for task identification. Rimon et al." (2022): Lee & Chung 1202T) learn model of the task distribution, sampling from it to generate tasks and using these to generate full trajectories. Conversely, in our work, imagined rollouts are history-dependent allowing the model incorporate information from online data. An extended related work section
is available in Appendix F
CONCLUSION
We present MAMBA, figh-performing and versatile meta-RI. algorithm based on the Dreamer architecture. We show that MAMBA significantly outperforms existing meta-RL approaches in cu- mulative returns, sample efficiency and runtime. We find this result to be consistent in common meta-RL benchmarks, in tasks with visual inputs and in a class of decomposable task distributions, for which we also provide PAC bounds. In addition, MAMBA requires little to no hyperparame- ter tuning. Finally, we believe that MAMBA can be used as a stepping stone for future meta-RL. RL research, generalizing to more challenging domains and real-world applications.
All experiments were conducted using an Nvidia T4 GPU, with 32 CPU cores and 120GB RAM.

---

REPRODUCIBILITY
Our implementation of MAMBA can be found in htttps://gttttub.ccm//ooobb aaaaa
and based on an open-source implementation of DreamerV3 (Hatner et al.) 2023)
found at ements from the open-source implementations of Zintgraf et al.) (2019) (021) at https://githe: com/ 1mz http://minnnrraa/aaaaa
https://gitttb..ooom 2tcp://Mi1//dreeeerrrrttoooo
borrowing el-
All hyperparameters and details modi-
fied from the original implementation are described in the Experiments section (Sec. [ as well as
in Appendix D
ACKNOWLEDGEMENTS
This work received funding from the European Union (ERC, Bayes-RL, Project Number 101041250). Views and opinions expressed are however those of the authors only and do not essarily reflect those of the European Union or the European Research Council Executive Agency. Neither the European Union nor the granting authority be held responsible for them.
REFERENCES
Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shi- mon Whiteson. survey of meta-reinffrrcement learning. arXiv preprint WxXi:2301108002 2023.
Anthony Cassandra, Leslie Pack Kaelbling, and Michael L Littman. Acting optimally partially observable stochastic domains. In Aaai, volume 94, pp. 1023-1028, 1994.
Era Choshen and Aviv Tamar. Stfi:23006.024188 2023.
Contrabar:
Contrastive bayes-adaptive deep rl. arXiv preprint
Ignasi Clavera, Jonas Rothfuss, John Schulman, Yasuhiro Fujita, Tamim Asfour, and Pieter Abbeel. Model-based reinforcement learning meta-policy optimization. In Conference on Robot Learning, pp. 617-629. PMLR, 2018.
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning- identifiability challenges and effective data collection strategies. Advances in Neural Information Processing Systems, 34:4607-4618, 2021.
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning, 2016.
Michael O'Gordon Duff. Optimal Learning: Computational procedures for Bayes-adaptive Markov decision processes. University of Massachusetts Amherst, 2002.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic neta-learning for fast adaptation of deep networks. In International conference on machine learning, pp. 1126-1135. PMLR, 2017.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learn- ing behaviors by latent imagination. In International Conference on Learning Representations, 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pp. 2555-2565 PMLR, 2019b.
Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010002199. 2020.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.041 2023.
10

---

Takuya Hiraoka, Takahisa Imagawa, Voot Tangkaratt, Takayuki Osa, Takashi Onishi, and Yoshi- masa Tsuruoka. Meta-model- based meta-policy optimization. In Asian Conference on Machine Learning. pp. 129-144. PMLR, 2021.
Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv: arXiv: 2e00100550.444, 2019.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. Advances in neural information processing systems, 32, 2019.
Heinrich Jiang. Uniform convergence rates for kernel density estimation. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1694-1703. PMLR, 06-11 Aug 2017.
Diederik P Kingma and Max Welling. arXiv:1312.6114 2013.
Auto-encoding variational
bayes.
arXiv preprint
Kimin Lee, Younggyo Sco, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynam- ics model for generalization in mhodel-based reinforcement learning. In International Conference on Machine Learning, pp. 5757-5766. PMLR, 2020.
Suyoung Lee and Sae- Young Chung. Improving generalization in meta-rl with imaginary tasks from latent dynamics mixture. Advances in Neural Information Processing Systems, 34:27222-27235, 2021.
Zichuan Lin, Garrett Thomas, Guangwen Yang, and Tengyu Ma. Model-based adversarial meta- reinforcement learning. Advances in Neural Information Processing Systems, 33:10161-10173, 2020.
Evan Z Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn. Decoupling exploration and ex- ploitation for meta-reinforcemen learning without sacrifices. In International conference on ma- chine learning, pp. 6925-6935. PMLR, 2021.
Zhao Mandi, Pieter Abbeel, and Stephen James. On the effectiveness of fine-tuning versus meta- reinforcement learning, 2022.
Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak. Discover- ing and achieving goals via world models. Advances in Neural Information Processing Systems, 29949999-2439,, 2021.
Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments through meta- reinforcement learning. arXiv preprint 1rxhv:1803.11344 2018.
Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free rl can be strong baseline for many pomdps. In International Conference Machine Learning, pp. 16691- 16723. PMLR, 2022.
Christos H Papadimitriou and John N Tsitsiklis. The complexity of markov decision processes. Mathematics of operations research, 12(3):441-450 1987.
Jurgis Pasukonis, Timothy Lillicrap, and Danijar Hafner. Evaluating long-term memory in 3d mazes. arXiv preprint arXiv:2210.13386 2022.
Christian Perez, Felipe Petroski Such, and Theofanis Karaletsos. Generalized hidden parameter mdps: Transferable model- in handful of trials. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pp. 5403-5411, 2020.
Brieuc Pinon, Chan-Charles Delvenne, and Raphaël Jungers. A model-based approach to meta- reinforcement learning: Transformers and tree search. arXiv preprint arXiv:2208. /535, 2022.
11

---

Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off off-policy mett-reiinforemment learning via probabilistic context variables. In International conference on machine learning, pp. 5331-5340. PMLR, 2019.
Zohar Rimon, Aviv Tamar, and Gilad Adler. Meta reinforcement learning with finite training tasks-a density estimation approach. Advances in Neural Information Processing Systems, 35:1 35:13640- 13653, 2022.
Murray Rosenblatt. Remarks on Some Nonparametric Estimates of Density Function. The Annals of Mathematical Statistics, 27(3):832 - 837, 1956. doi: 10.1214/aoms/1 17777///917777
Julian Schrittweserr, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, (88 609, 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Jayakumar Subramanian and Aditya Mahajan. Approximate information state for partially observed systems. In 2019 IEEE 58th Conference on Decision and Control (CDC), pp. 1629-1636. IEFE, 2019.
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm control: Software and tasks for continuous control. Software Impacts, 6:100022, 2020.
Jane Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn, 2017.
Qi Wang and Herke Van Hoof. Model-based meta reinforcement learning using graph structured surrogate models and amortized policy search. In International Conference Machine Learning, pp.
PMLR, 2022.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data, Advances in Neural Information Processing Systems, 34:25476-25488, 2021.
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl. Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta- learning. In International Conference on Learning Representations. 2019.
Luisa M Zintgraf, Leo Feng, Cong Lu, Maximilian Igl. Kristian Hartikainen, Katja Hofmann, and Shimon Wh Whiteson. Exploration approximate hyper-state space for meta reinforcement learning. In International Conference on Machine Learning, pp. 12991-13001 PMLR, 2021.
12

---

A APPENDIX
In Section A. We define the term Decomposable Task Distribution and in Section A.2 we show (un- der some mild assumptions) that much more favorable bounds can be obtained for these distributions compared to the general case.
A.1 DECOMPOSABLE TASK DISTRIBUTIONS
follow the notation from Rimon et al. (2022) and assume that a task distribution defined by parametric space θ, prior distribution over this space ((@) and mapping function from the parametric space to the MDP space : θ + MM. To sample an MDP from the distribution, a parameter is sampled from f and is passed through g. Essentially, over MDPs can be defined as a parametric variation. This formulation is very broad and can be used formalize all the environments considered in this paper; refer the reader to/Rimon et al. (2022) for examples and motivation. Using this formulation we can write the expected return from Section 2.1 as:
assumes that the distribution
H
J EB-s BR,M=
Σ
Ti.k
We now formally define a Decomposable Task Distribution, which as its name suggests, task distribution where each task can be decomposed to subtasks.
Definition Decomposable Task Distribution. Given Oz. where t; N, C Ra,d d, E N. fs E P( 9: 0 + M and M the MDPs space. To sample an MDP from this sampled from the priors and subtasks are being generated 0; fü. M; = ge(O;), Vi E [1, Ny). The resulting full MDP is defined by playing in M, between time steps and f;-1 vi E [1, Nr), without resampling from the initial state distribution when passing from one MDP to another.
distribution the parameters
A.2 PAC BOUNDS USING SUB-TAN TASK DECOMPOSITION
General Regret Bounds We first define the regret the difference in the return of policy the return of the Bayes-optimal policy as follows:
R(T)= =0 J Eows
EmBo. M=9())
Σ
F. P...Mpgh ()(6)
Σ
Where тOO is the Bayes-optimal policy, TB arg max J Note that the regret is iff Bayes-optimal. Since the regret is taken with respect to the unkown prior distribution, it is measure of the generalization of т.
In their work, Rimon et al. (0022) suggest a density estimation approach to achieve regret bounds. Given a set of training samples, they first estimate the parametric distribution of tasks f (0) with
Gaussain Kernel Density Estimation (KDE, Rosenblatt 1956) fo (0). By using KDE error bounds Jiang 2017) they achieved probably almost correct (PAC) bounds the regret:
Theorem (Theorem 6 in Rimon et al. 2022) For prior f(0 bounded parametric space θ that satisfies |l < do and is a-Höller continuous, and a Gaussian KDE, we have that with probability least - 1/N:
log N
R(T < 2 Ca
Where N are the number of sampled tasks, Rmax is the maximum possible reward, |θ] and d are the volume and dimension of the space O respectively, Ca is polynomial of order d and ni e
max, J = arg max, Eg-io Ta [Er.M-n(0) 9(0) Σ# r [] ie the Bayes-optimal policy under the assumption that the estimator fa is correct.
13

---

This theorem shows that, given general parametric space, the number of task samples needed to guarantee a constant regret scales exponentially with the dimension of the parametric space. This observations shows inherit limitation of meta-RL.
Bounds for Decomposable Distributions The obvious approach to achieve regret PAC bounds for decomposable task distributions, is to take a similar approach as done Theorem
2 We can
handle the decomposable distributions normal distributions where O full Oo x θ x X θ
and the prior distribution over the parametric space is ffult (0 (01..... NNT) )) n Nr ((:).. In that case we can approximate ffull directly using a KDE f over the samples {0, i= and using Theorem get regret bounds of:
loog m2 20
R(TT)ru)) S 22 @falll
Carsu
Where d'uli = Σnn d,, with d, being the dimensionality of θ. This means that if we approach decomposable task distributions regular distributions, to guarantee constant regret we 'll need exponentially growing number of samples with respect to d i.e depend on both the number of sub-tasks and the number of DoF in each sub-task distribution.
We suggest a different approach: estimate each distribution f. by a KDE ft separately, i. = (f.) For a policy π and a decomposable task distribution defined by {ti, fi, 6, I we define {o.] } Ny Fi=1" where y," (0) 9s (0), except for the initial state distribution, which is set be the induced initial distribution of the i' i'th MDP, as defined in Definition I v0 E O2 and Vi E Nr].
Next we follow the derivation of Theorem in Rimon et (022)):
NT.
Ny
J7
T -M (0)
Σ
E Er.M=o (0)
Σ
f((0d00
Nr
4 Σ
(1(0)- ao
57
J
Σ
(x
Tt
lsus
And:
Nr
JJ
j
(a,M=
Σ n
11(0)-
Tanb
t=
Nr Rmany (t, tz-1) / f1(0)-F(00) de
And following the same lines as the proof Lemma 3 of Rimon et al. (2022) we get:
Ny
Rmass (t: - 1)-1) / [i(() T((0) do
R(T 7
5 2.
1=1
Using the PAC bounds for the KDE error of Jiang (011)) as was done in Rimon et al. (0022) we receive the final regret PAC bounds:
Theorem 3 For a decomposable prior distribution (4i, fú, O, i where every parametric space is bounded, [| < co and f o-HHöler continuous and for set of Gaussian KDEs as defined above, we have that with probability at least 1 - 1/NN:
Ny
log N N
To
R(T²" ( ) 2Rmaa H 1 maazl Odmax
14

---

Where N the number of sampled tasks, Rmar is the maximum possible reward (over all the subtasks), 1mmazl and dmaz are the maximal volume and dimensions of the spaces [0: Nr. respectively, Cdmar is polynomial of order dmaz and T [E
isws E arg max, J Ssub
i.e the Bayes-optimal policy under the assumption that the es- re []]
timator fuu is correct.
As we can see, using the fact that the distribution can be decomposed into sub-distributions and ap- proximating each one individually, results in much more favourable bounds. With the new approach, to guarantee constant regret need exponentially growing number of samples with respect to dmaz instead of duull This essentially means that if we approach this kind of tasks by solving each sub-task individually, the difficulty of the task will be determined by the sub-task with the highest number of DoF, and not by the number of sub-tasks.
For example, looking at the multi- goal Rooms-N N environment from experiments, although the total number of degrees of freedom is N, L. grows as we increase the number of rooms, the difficulty of the task does not change, the agent should solve each of the rooms separately. In comparison, our theorem states that a standard navigation goal in N dimensions, which has the same number of DoF will be much harder since we must estimate all the DoF at once.
Therorem B also sheds light on our findings from Section 3.2 which show that VariBAD works better once the reconstruction window is smaller, which helps the belief (estimated by the latent vector) focus on the current sub-task.
B DIFFERENCES BETWEEN VARIBAD AND DREAMER
For simplicity, we denote VariBAD as VB, and Dreamer as D.
History of optimization at timestep t: In VB, when executing the optimization at time t, the entire meta-episode gets encoded into b, Conversely, in D, short (L 64) sequences of transitions are sampled from the replay buffer and encoded into latent vectors 2 using the RSSM world model. This is problematic in the meta-RL context, since it does not allow components (policy, model, value) to depend on the full history of the meta-episode and may lead to ambiguity that should have already been resolved.
World model prediction horizon: D, the world model predicts a single step, but forces the latent states of subsequent transitions to be consistent. In VB, the decoder directly predicts outcomes (next states and rewards) for every possible observation by feeding the belief, current observation and action as inputs to the decoder
Horizon differences: In VB, episodes were short (about 100 time steps), while in D, POMDPs with hundreds of time steps were considered. This technical consideration is probably one of the factors why D samples uub-trajectories of length L, and why VB were able to run on a GPU without out-of memoy exceptions. However, it is a limiting factor for many meta-RL tasks that are characterized with a long horizon due the combination of several sub-episodes in every meta-episode.
Policy optimization: In VB online model-free algorithm PPO (Schulman et al. 2017) was used while D use a combination of the REINFORCE (Sutton et al.) 1999) algorithm and backpropogation through trajectories imagined by the world model.
Policy input: In VB, the policy is conditioned on both the latent 2, and the current state St. However, in D the policy is conditioned only on 2t.
Hidden sizes in RNN: D has much larger hidden layer sizes than VB.
C FULL ENVIRONMENTS TECHNICAL DETAILS
Here we provide the full details regarding the environments we used in our experiments (Section
In VariBAD, observations were used, so observation in this context is equal state.
15

---

Point Robot Navigation A commonly used 2D continuous control benchmark for meta- RL (Zint-
graf et al. 2019 Rimon et al. 2022 Dorfman et al. 2021 Choshen & Tamar 2023). where goal is sampled semi-cirele with radius of 1.0. The reward is sparse and defined as 1- Pagent P || when the agent is inside the goal radius (0.2) and 0 otherwise (where Pagent. Pgoat the agent and goal positions, respectively). Here, a meta-episode consists of two sub-episodes of K = 100 steps each. The Bayes-optimal behavior is to explore the semicircle the first episode and exploit the information acquired by going directly to the goal on the second.
Point Robot Navigation Wind This is an extension of the Point Robot Navigation environment, where a stochastic wind force is added at each step. More explicitly, at each step, a 2D Gaussian random vector with controllable std (e.g. 0.3) is sampled and added to the selected action. Other than that, the environment details and Bayes-optimal behavior are the same as the in the Point Robot Navigation environment.
Escape Room In this environment, taken from Dorfman al. 2021), 2D point robot is ini- tialized in a circular room of radius 1.0. An arched door of length x is randomly sampled on the upper half of the room, providing the only exit from the room. The reward is 0 inside the room and 1 outside of it. Here, a meta-episode consists of two sub-episodes of = 60 steps each. A Bayes-optimal agent should explore the upper half circle, pushing against the wall until successfully exiting.
8
Humanoid-Dir First proposed in Rakelly et al. (2019), in this environment the Humanoid MuJoCo (a high dimensional humanoid agent) task is initialized with a random target walking direction e [0, 2m]] on the 2D plane. The agent receives reward proportional to the projection of its speed on the unknown goal direction. Since in meta-RL we observe the environment rewards, the Bayes-optimal behavior should be able detect the correct direction after few steps (theoretically one step) and consequently walk as fast as possible towards this direction.
Reacher This environment, available the DeepMind Control Suit (DMC, Tunyasuvunakool et 20200), is planar two-link robot tasked to reach various goals the plane with the tip of its second link. A goal is sampled on the 2D plane within ball around the agent, which when reached (to certain small distance), provides a reward proportional
the distance from the goal. For this environment we use 3 sub-episodes of length K = 300 in every meta-cpisode. In contrast the point robot navigation task, since the goal can be sampled anywhere within the circle, the Bayes-optimal policy should explore in spiral form to cover the entire goal space.
We also consider a multi-goal variant of this task, Reacher-N, with observation and action spaces identical to the original Reacher task. Here. N goals are sampled i.i.d. (according to the distribution above) which the robot must reach in a predetermined order. The reward for each goal [1...., N] may be obtained only after all the previous goals were reached. For the first N agent receives positive reward only once when the agent first penetrates the goal radius. In the N-th goal, the positive reward is given every step, until the sub-episode terminates. The Bayes-optimal behavior in this case is to explore until the last goal is found, and then stop moving. Once a goal was identified, in subsequent sub-episodes the policy should head to that identified goal directly. This serves as a high-dimensional environment decomposable into independent sub-tasks (up to the initial distribution), matching our analysis in Sec, [3.2] We note that for 3 and 4 goals increased the number of steps per sub-episode to 400 instead of 300.
goals. the
Panda Reacher Proposed by Choshen & Tamar (203). Panda Reacher is dynamically complex system with visual inputs and sparse rewards which requires a non-trivial Bayes-optimal behavior. In this environment, 7-degree-of-freee ofreeeom (DoF) robotic arm is required reach goals with the robot end-effector (i.e. the tip of the robot). The goals are hidden and must be discovered through
Inside the goal radius the rewards are positive.
16

---

Hyper-paranee parammeeeettrr
Dreamer- Vanilla
Dreamer-Tune
Description
dyn_discrete
16
Numher of discrete latent vectors in world model.
dyn hidden
512
64
MI. size the world model.
dyn stoch
32
16
Size of each latent vector in the world model.
imag horizon
15
10
Length of the imagination horizon (NiMe ef. Sec 2.2
units
512
128
Size hidden MLP layers.
Table 2: Hyper parameters differences between Dreamer-Vanilla and Dreamer-Tune.
Rooms-4
Rooms-5
VariBAD Vanilla
14.1 3.0
0.9+-2.5
VariBAD Local
88. 1-.1223
15.1+-13.1
MAMBA
136. 136.741.4
113.1#5.7
Table VariBAD reconstruction ablation study. Using local reconstruction helped the Rooms-N environment, which contains sub-tasks.
exploration. The rewards are sparse, meaning only when the end-effector is within a short distance from the goal, a reward of is obtained (otherwise the reward is O). The observation is an image of size 84x 84 x4, and the action space is the direction towards which the end-effector is moved.
Rooms-N: As an additional multi- goal task, we consider a discrete grid-world environment con- sisting of sequence of N rooms of size 3 x separated by 3 corridor. At each location the agent observes its (I, y) location and can choose to go right/eeftuup/down (will stay in place if the resulting state is out of the map). A goal is sampled the corner of each room and the agent starts in the left-most room. The goals in rooms , .]] are obtainable if all the previous goals were visited. The agent receives reward of 1 when goal is visited, and -0.1 otherwise. For the first N - rooms, once goal is visited in sub-episode, it no longer emits a positive reward. On the first sub-episode, a Bayes-optimal policy should explore each of the rooms, continuing to the next room upon finding the goal. In the second sub-episode, it should go directly to the goal in each room. This environment serves as an additional sub-task decomposable environment (see Sec. 3.2)
D HYPERPARAMETERS OF DREAMER-TUN -UNNE
Our parameters of Dreamer- Vanilla are identical to DreamerV3 (Hafner et al. 2023) In Table 2 we list the differences between Dreamer- Preamer-Vanilla and Dreamer-Tie Tune.
E EMPIRICAL RESULTS FOR VARIBAD WITH LOCAL RECONSTRUCTION
In Sec. 3.2 we introduced variant of VariBAD that uses a local reconstruction (reconstructing window of adjacent transitions instead of the full trajectory). In this section we evaluate it against the original VariBAD on two sub-task decomposable environments: Rooms-4 and Rooms-5. Results can be found in Table We also provide the results of MAMBA from Sec. as reference.
As we can see in Table local reconstruction helped VariBAD these scenarios, which is in line with our theoretical findings (cf. Sec. 3.2
F EXTENDED RELATED WORK
Meta-Reine Reinforcement Learning (meta-RL) has been extensively studied in recent years (Duan ctt al., 2016 Wang et al.) 2017 Finn et al. 2017: Beck et al. 2023). Of specific interest to our work context- based meta-RI RL approaches (Rakelly et al. 2019 Humplik al.) 2019 Zintgraf et al. 2019; 2021 Wang & Van Hoof, (022). which attempt to lear a Bayes-optimal policy. Most of these model-free thus suffering from high sample complexity.
Others, such as Humplik et al. (201)), utilize access to privileged information during the meta- training phase and thus do not apply to the general setting of meta-RI..
17

---

Model-based RL has been shown to be more sample efficient than its model-free counterparts Nagaaannii et al. 2018 Clavera et al. 2018 Janner et al. 2019 Schrittwieser et al. 2020; Ye et al, 2021). While there are many different approaches to model learning and usage, prominent line of work is the Dreamer framework (Hafner et al., 2019a b 2020) 2023) Mendonca et al. 2021), which learns recurrent latent space world model and uses rollouts in latent space to train policy. This approach has been shown to work well in POMDP environments, owing to the inherent history aggregation ability of the RNN (Ni al., 2022).
Model Based Meta-RLa RL: Recent approaches Nagabandi et al. 2018 Perez et al. 2020 Lin et al.,
2020 Lee et al.] 2020; Hiraoka al. [202]] Wang Van Hoof 2022 Pinon et al. 2022 propose
model- based meta-RL methods. [Ngabandi et al. (2018) use dynamics model and adapt it online
to varying dynamics in robotic scenarios. Perez et al. (2020) propose to encode the parameters controlling the reward and transition functions into latent variables of learned model. To tackle the problem task distribution shift, [in et al. (2020) train a shared dynamies model on a distribution of tasks, but train policy at test-time given new reward function, preventing fast adaptation to new tasks. Lee et al. (2020) trains a context vector to be consistent with forward and backward dynamics, conditioning
this context to enable dynamics model to generalize quickly. The above approaches all continuously adapt to a varied task distribution, making no distinction of when a new task instance is attempted. In accordance, they all use short histories of states and actions, limiting usability when long horizons are required for task identification.
An approach that does focus on long horizons and is designed specifically to test model memory is baseline proposed in Pasukonis et al. (202)), where a variation of Dreamer-v2 (Hafner et al. 2020) is used benchmark new environment. While this baseline bears some similarities to our approach, Pasukonis et al. (0022) aims to establish a baseline for testing memory capabilities of agents, rather than proposing an algorithmic solution to meta-RL in general. In accordance, the modifications they propose for Dreamer are only geared at extending its memory capabilities, disregarding other modifications we found important in this work.
Wang & Van Hoof (2022) propose GNN-based model learned for each task of the training distribu- tion, along with an amortized policy selection method to quickly find the correct model at meta-test time. In contrast, our method learns a single model and single policy for the entire task distribu- tion. Pinon et al. 2022) propose a Transformer-bart the model test time. While general, this approach may be limited to discrete- action domains; in addition, tree search is a time consuming process. Our method, on the other hand, can operate on discrete and continuous domains (as it is based the Dreamer architecture), and learns a policy so that it can act quickly at test time without additional planning.
world model, using tree search to plan with
Rimon et al. 2022) and Lee & Chung) 2021 learn model of the task distribution, sampling from to generate tasks and using these to generate full trajectories. Conversely, in our work, imagined rollouts are history-dependent, allowing the model incorporate information from online data.
G ABLATION STUDY OF MAMBA
In this appendix we ablate the three different modifications proposed in Section 3.3 to identify the contribution of each.
First, we note that it is infeasible to run experiments without adding the reward into the observation as it prevents any possibility of the agent distinguishing between tasks; to validate this, we removed the reward observation and tested on the Rooms-3 environment. Results are as expected: com- pared to MAMBA, which obtained a return 156.2+ 1.7, MAMBA w/o reward in the observation obtained only 11.5= 1.4.
Next, MAMBA with the original batch sampling (i.e. batches do not start from the start of the meta- episode) is in fact Dreamer- Tune; therefore, we did not run separate ablation test for this design choice.
Finally, tested the effects of removing the horizon scheduling in the Point Robot Navigation scenario, and indeed found that the results obtained are on-par but runtime is slower (on the same hardware), see Table
18

---

MAMBA
MAMBA w/o Horizon Scheduling
Average Total Return:
156.2 + 1.7
156.8 + 1.2
Total Time [H]:
34.1
44.0
Table 4: Ablating the horizon schedule in the Point Robot Navigation environment.
RL22
VariBAD
HyperX
Dreamer- Vanilla
Dreamer- Tune
MAMBA (Ours)
PRN:
239.443.11
218.0 = 26.0
204.0 + 19.9
241.1 + 9.9
239.7 = 3.1
242.2 + 7.9 ~2M
~00
~00M
~15M
~2M
~22
Rooms-3: 3:
108.0=31.7
158.7 2.6
111.8 + 20.2
137,5 + 6.2
142.6 4.1
156.2 1.7
Rooms-4:
85.1
88.1 39.2
14.1 + 3.0
115.0 + 6.7
119.8 + 3.1
136.7 + 1.4
Rooms-5:
72.4+:36.8
0.9 + 2.5
-15.8 + 0.4
96.5 + 2.5
93.9 + 5.3
113.1 + 5.7
Rooms-6:
64.9+115.
-11.0 0.9
-15.7 0.7
69.5 + 6.3
71.0 6.3
94.5 + 1.2
Rooms-7:
42.9122.9
50.2
23.2::2.9
Rooms-8: 8:
29.0=11..0
29.1
55.613.0
~000
100M
100M
6M
~6M
~6M
Reacher-1:
595.8 21,4
473.7 27.5
505.9 + 36.0
552.0 27.8
503.4 69.0
655.5 12.3
Reacher-2:
404.0+335.9
46.6 33.2
30.0 48.5
247.4 80.5
217.6 64.3
285.8 + 89.6
Reacher-3:
390.6 + 49.1
0.2 0.2
0.5 + 0.2
183.6 + 100.0
76.9 + 80.5
325.0 + 47.0
Reacher-4:
0.0 1
0.0 = 0.0
-0.5 + 1.1
0.4 4 0.0
0.1 0.2
77.7 + 61.1
~100
~500
~OM
~
~TM
EscapeRoom:
79.914.4
70.7 5.3
66.9 1 6.5
68.2 2.4
73.2 1.2888
73.9 E 3.1
~20
~20
-20M
54M
24M
<4M
Table 5: Total return comparison of RI.2, VariBad, HyperX, Dreamer- Vanilla, Dreamer-Tume and MAMBA on different meta environments. We report the final reward (mean + std), and the number of time steps until convergence (below dashed line).
H ADDITIONAL EXPERIMENTS
H.1 ADDITIONAL META-RL. BASELINES
H.1.1
RL2
Duan et al. 2016) is common meta-RL baseline; our additional experiments (Table show that MAMBA superior RI.2 environments except EscapeRoom, where the difference is slightly in favor of RI. (although RI.2 used 5x more data than MAMBA).
RL
H.1.2 DREAM
DREAM (Liu et al. 2021 is meta-RI. algorithm with a more sophisticated out-of-the-box explo- ration mechanism. Since the experiments in DREAM deal with goal conditioned meta-RL tasks in grid-world environments, we tested it in our Rooms environment.
Our experiments in Rooms-3 show that DREAM fails to obtain any meaningful policy, and meta-episode rewards remain negative (MAMBA, Dreamer, VariBAD and HyperX all managed to obtain returns greater than 100). Results from experiments with more rooms behave similarly. This may result from bad Rooms; however, this is another advantage of using a Dreamer-based architecture as in MAMBA (same hyperparameters fit different scenarios out-of-the-box).
adaptation between the scenarios in the DREAM paper and
H.2
ADDITIONAL ENVIRONMENTS
H.2. ROOMS-7 AND ROOMS-8
We added more rooms to challenge MAMBA. We observe that MAMBA still manages to find a Bayes-optimal behavior and the drop in return is attributed to the exploration required to find the additional goals. VariBAD or HyperX were not tested as they already failed to solve scenarios with less rooms. See Table
19

---

VariBAD
HyperX
Dreamer- Vanilla
Dreamer- Tune
MAMBA
Return:
1369.3 = 75.3
2068.3 + 156.7
2096.3 = 79.8
2405.9 + 119.0
Steps:
100M
30M
30M
30M
Table 6: Comparison of VariBad, HyperX, Dreamer-Vanilla Dreamer-Tune, and MAMBA in the Humanoid-De Dir environment.
VariBAD
HyperX
Dreamer- Vanilla
Dreamer-Tune
MAMBA
Return:
194.5 44.7
177 118.5
226. 133.5
224.5+4.4
224.1=5.2
Steps:
20M
20M
2M
2M
2M
Table 7: Comparison of VariBad, HyperX, Dreamer-Vanille Dreamer-Tune, and MAMBA the Point Robot Navigation- Wind environment.
H.2.2 HUMANOID-DII
First proposed in Rakelly et al. (2019). in this environment the Humanoid MuJoCo task is initialized with random target walking direction 0 E 00, 2 the 2D plane. Table summarizes summaizss the results of this experiment. As in the results described in Sec. 4 using less environment steps (30M vs. 100M) the algorithms based on the Dreamer architecture perform better than VariBAD and HyperX. Among these, MAMBA is substantially better, reaching the highest return.
H.2.3 POINT ROBOT NAVIGATION WIND
This is a modification of the Point Robot Navigation (PRN) environment, where a stochastic wind force is added at each step. Results in Table that the conclusions from Point Robot Navigation also extend to this stochastic version.
20


Reviews:
[{"review_id": "ffaTK8roUe", "reviewer": "Reviewer_cu12", "summary": "The authors introduce MAMBA (MetA-RL Model-Based Algorithm), a meta-reinforcement learning algorithm developed using the Dreamer framework. The authors extend the Dreamer framework which was designed for POMDPs to the general case of meta-reinforcement learning.  Additionally, the authors introduce two novel environments specifically for meta-reinforcement learning with high-dimensional task distributions, which can be decomposed into lower-dimensional sub-tasks. They demonstrate the performance of their algorithm on these new environments as well as standard meta-RL environments.", "questions": "1. What determines the number of sub-episodes for each environment? It appears that each environment has a different number of sub-episodes. How does altering the number of sub-episodes influence the algorithm's performance?\n\n2. The algorithms are tested in deterministic environments (although I could be mistaken). How do you anticipate the performance would be affected if the environments were stochastic? Could the introduction of randomness potentially complicate the task identification process?\n\n**Minor Clarifications**\n\n1. In Figure 1 (left most) Why isn't the agent exploring the right portion of the semi-circle during the first episode? \n\n2. To enhance understanding, it would be beneficial to have comprehensive descriptions of the new environments (Multi Goal Rooms and Multi Goal Reacher), including details about their state-space, action space, and reward structure.\n\n3. I would encourage the authors to add markers to all plots to enhance readability.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The algorithm is sample efficient when compared to other meta-RL algorithms. \n\n2. The authors conduct a good number of simulations to explain their algorithm, and evaluate its performance. \n\n3. The paper is generally well written and easy to follow.", "weaknesses": "**The assumption of task decomposability and task independence is strong, vague, and confusing**\n\nThe paper assumes scenarios of task decomposability, where each task is decomposed into independent tasks, I think this is a pretty strong assumption, and not many environments will satisfy this criteria. The example quoted for task decomposability by the authors is a little confusing, the authors  provide an example of a robot being required to solve several independent problems in a sequence.  Isn't the fact that they should be performed in a sequence make the tasks dependent, and thus not independently decomposable? Further, in the Appendix (section C multi-goal Reacher-N) the authors claim that the environment is decomposable into *nearly independent sub-tasks*, which is a little vague and confusing.\n\n**Algorithmic contribution is minimal**\n\nThe proposed algorithm is a minor modification over Dreamer, with the only significant change being sampling full-meta episodes instead of a smaller fixed length", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Bgtois6d2c", "reviewer": "Reviewer_fKkT", "summary": "This work draws a connection between model-based RL and a common meta-RL approach, VariBAD. Specifically, both encode the trajectory history into a latent variable, which is then used to predict observations, e.g., the state and rewards. This work leverages this connection to propose a new meta-RL approach based on Dreamer called MAMBA. Compared to Dreamer, MAMBA differs in three ways:\n(1) Like many meta-RL algorithms, MAMBA augments the state with the current timestep and the reward.\n(2) Whereas Dreamer only computes the latent variable over the past 64 steps of the history, MAMBA computes it over the entire past history. This helps in the cases where important information was discovered more than 64 steps ago.\n(3) MAMBA introduces a curriculum where the horizon of the episode is increased over time to help combat the effect where long-range model predictions are inaccurate at the beginning of training.\n\nThis work evaluates MAMBA and finds that it performs favorably compared to VariBAD and Dreamer.", "questions": "Please see previous section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "*Clarity*\n- This work is well-written and easy to understand. The problem this paper is attempting to solve, and its proposed algorithm are clearly presented, which makes it easier to reproduce.\n\n*Originality and Significance*\n- The proposed changes appear to be fairly minor: augmenting the state with these additional observations is something that happens in a fair number of other papers; and increasing the history length for computing the context variable is akin to increasing a hyperparameter in Dreamer, which is reduced for computational reasons. However, they seem fairly sensible and result in fairly good performance gains.\n- Additionally, this work makes interesting connections between model-based RL and VariBAD, and shows that Dreamer can outperform existing meta-RL algorithms.\n\nOverall, this work unproblematically provides several contributions that are interesting to the meta-RL community, so I would be in favor of acceptance.", "weaknesses": "I think this work already provides valuable contributions, but can primarily be strengthened significantly by shedding more light on its results.\n- MAMBA proposes 3 changes over Dreamer. It would be very helpful to perform an ablation study on these 3 changes and understand which ones are most important and how they impact performance.\n- This work finds that Dreamer performs better than VariBAD and HyperX out-of-the-box on meta-RL tasks. I find the paper's claim that this performance gap probably results from architectural differences (VariBAD uses a very simple architecture, whereas Dreamer's is heavily tuned), but it would be interesting to ablate this and understand this better. At face value, non-meta-RL algorithms outperforming meta-RL algorithms on meta-RL tasks is somewhat surprising, and warrants further investigation. Further, it would be worth understanding in what ways these algorithms differ out-of-the-box. Does Dreamer result in better exploration? Or just exploitation? Computing the results in Figure 4 for VariBAD could be very helpful. Additionally, on what sorts of tasks do we expect Dreamer to be better than VariBAD? 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? What happens when this structure is violated?\n- Further analysis of when it's important to predict the whole past + future (as in VariBAD) vs. local reconstruction (as in Dreamer) would be helpful. It's quite believable that local reconstruction could result in better empirical performance due to optimization, though the reason why VariBAD predicts the whole trajectory is as a proxy for predicting a distribution over dynamics, which is intractable. Predicting long-past events could in principle be really important for learning a good representation if some observation explains the dynamics seen early on, though this structure is not present in the current tasks, which may contribute to the conclusions drawn in this work. Some experiments or discussion on this would be helpful.\n- Similarly, discussion about when it is necessary / possible to use the entire past to compute the context variable would be helpful. For really long horizon tasks, rolling out a recurrent policy is simply computationally intractable, so chunking is necessary, e.g., along the lines of R2D2 (https://openreview.net/pdf?id=r1lyTjAqYX)\n- DREAM (https://arxiv.org/abs/2008.02790) reports better exploration than VariBAD on tasks requiring more sophisticated exploration. It \nand subsequent papers (https://arxiv.org/pdf/2211.08802.pdf) also provide more complex tasks. Given the results of Dreamer outperforming VariBAD out-of-the-box, it would be interesting to consider discussion or comparison with DREAM or on these more complex tasks.\n- Despite meta-RL being core to this work, the related works section glosses over a rich literature on meta-RL, including model-based meta-RL e.g., see https://arxiv.org/abs/2301.08028, https://arxiv.org/abs/1905.06424, https://arxiv.org/abs/1803.11347, https://arxiv.org/abs/1809.05214\n\nThese open questions are what justify my score of a 6 over an 8.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yI2WezKS2l", "reviewer": "Reviewer_X5No", "summary": "In this paper, the authors tackle the problem of meta-RL: learning across a range of tasks at the same time. They combine the sample efficiency of model-based RL with an existing meta-RL framework to produce a method (Mamba) which is more sample efficient and powerful than existing methods on a range of meta-RL tasks including a set of novel tasks which are benchmarked against other methods.\n\nThe 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. They then provide experimental results over a number of meta-RL tasks providing evidence for the improvements that Mamba has over existing methods.\n\nThe appendix goes into more detail on the theoretical and technical details.", "questions": "The questions are all based on the weaknesses in the previous section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Based on the evidence provided, it is clear that within the domains examined, Mamba is a stronger meta RL-algorithm than those tested against both in terms of returns and sample efficiency. The results are clearly presented and the explanations are generally sufficient.", "weaknesses": "While overall the paper is good, and provides good evidence within the context, it seems that there are a number of approaches within the literature which have not been covered. These include:\n\n1) Pinon et al, A model-based approach to meta-reinforcement learning:transformers and tree search, https://arxiv.org/pdf/2208.11535.pdf\n2) Wang and Hoof, Model-based meta reinforcement learning using graph structured surrogate models and amortized policy search, https://proceedings.mlr.press/v162/wang22z/wang22z.pdf\n3) Clavera et al, Model-Based Reinforcement Learning via Meta-Policy Optimization, https://arxiv.org/pdf/1809.05214v1.pdf\n4) Lee et al, Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning, https://arxiv.org/abs/2005.06800\n\nWhile there is a short section with a few papers on model-based RL, it seems that there are some important pieces of research which have been missed out here. Given how much has been missed, it is not clear that the results are as persuasive as they at first seem.\n\nIn addition, there exists an extensive meta-RL benchmarking suite produced by Wang et al, Alchemy: A benchmark and analysis toolkit for meta-reinforcement learning agents, https://arxiv.org/abs/2102.02926. I believe that it is vital for the community to come together to standardise benchmarking, and so such a toolkit seems ideal to truly show the applicability and strengths of Mamba.\n\nOn a stylistic note, there are many typos throughout, particularly in the Appendix. These, along with points of clarity are explained here:\n\n1) Section 3.2: family task-> family of task\n2) Section 3.2.1 PAC not defined\n3) S 3.2.1: with high number->with a high number\n4) S 3.3: singnals->signals\n5) The plot in figure 2 is not clear given that the environment has not been clearly explained by this point.\n6) Sometimes Dreamer-tune is written and sometimes Dreamer-tuned\n7) In section 4, figure 1 is mentioned, which appears very close to the beginning of the paper, without context and is unclear at this point. It should be later in the paper when it is introduced.\n8) Figure 4: amd->and\n9) Table 1: Bold font on the 7 in 73.9±3.1\n10) As discussed above, I believe that a lot of model-based meta-RL work has been left-out.\n11) A1 folmulation->formulation\n12) Definition 1 \\forall 1<=i<=N_T should be written \\forall i \\in [1,N_T]\n13) A2: well known Regret->Regret\n14) Bayes optimal does not have consistent hyphenation or capitalisation\n15) Above theorem 2: boudns->bounds\n16) Theorem 2: polynomial with d->polynomial of order d\n17) assumption the estimator->assumption that the estimator\n18) Top of page 14: bounds in the private case-> ?\n19) First equation on page 14: C_d should be C_{d^full}\n20) guerntee-> guarantee\n21) Proof of Lema 3->Proof of Lemma 3\n22) Theorem 3: distriubtion->distribution\n23) Theorem 3: Something is wrong after alpha-Holder continuous\n24) polynomial with d_max should again be \"of order\"\n25) assumption the estimator->assumption that the estimator\n26) using the fact the->using the fact that\n27) much favourable->much more favourable\n28) this kind of tasks->these kinds of tasks\n29) same amount of DoF->same number of DoF\n30) we have estimate->we have estimated\n31) Appendix B: short L=64->short (L=64)\n32) Check consistency of hyphenation in world-model throughout\n33) short 100 timesteps -> short (L=100)\n34) Top of page 16: p_agent and p_goal not defined\n35) Appendix E: It's not clear if there should be more in this section. \n\nAs can be seen, there are a lot of typos in the appendix which have taken away time from more important aspects of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "22BNg2QNME", "reviewer": "Reviewer_tjN7", "summary": "This paper proposes a model-based method for Meta-RL. The method is based on the Deamer-v3 algorithm, which is already good at solving POMDPs. They propose three changes to Dreamer-v3, that make it better suited to meta-RL style POMDPs specifically. The changes are training the model on full episodes, scheduling the training horizon length, and adding rewards and timesteps to the observations. They test the proposed model on meta-RL tasks and show that it works well. Additionally, a theoretical result is presented, which shows that decomposable task distributions can enable much lower sub-optimality than non-decomposable task distributions. Using that theoretical result they suggest a simple empirical modification applicable to VariBAD, which achieves better results in the decomposable task distributions. The proposed method naturally enjoys the same benefit.", "questions": "- The proposed method seems quite close to [1]. It would be good to add this in the related work section and discuss how they differ.\n- What are some examples of task distributions where the task dimension is high?\n\n[1] Pasukonis et al., 2022, Evaluating Long-Term Memory in 3D Mazes", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "## Contribution\n- Dreamer-style model-based RL algorithms seem like a natural fit for meta-RL problems. This paper shows that with few simple modifications, Dreamer-v3 performs well in meta-RL POMDPs.\n- The theoretical result about decomposable task distributions is interesting and nicely supported by the experiments.\n- The empirical results support the design choices and the theoretical results\n\n## Presentation\n- The paper is well written for the most part.", "weaknesses": "## Benchmark selection\n- The environments used in the empirical section nicely illustrate the theoretical points and the advantage from the design choices, but they still feel a little limited considering the POMDPs dreamer-v3 is capable of solving. I would not want you to drop any of the domains already included, but it would make the paper stronger to run the algorithms on something more complex than reacher, which is from the easier end of mujoco tasks.\n- It would be good to run the algorithm on at least some common meta-RL environments to help people who are familiar with that literature to ground the performance. I have no doubt that this works well there but it would be reaffirming to see those results. Consider running this for example in Walker and Humanoid tasks.\n- In the theory section, it is discussed that the local reconstruction could be harmful in high dimensional non-decomposable task distributions. It would have been good to include a demonstration of such failure mode in the experimental section.\n\n## Presentation\n- I found it confusing that the results from the room experiments were discussed in 3.2.2. before the room experiments were described.\n- Figure 2 is very small and busy. Neither 3.2.2 nor the caption contains enough detail to understand what is going on in the figure.\n- The sentence in parenthesis in the first paragraph of 3.2.2 is hard to interpret. Maybe a typo?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)