RL$^3$: Boosting Meta Reinforcement Learning via RL inside RL$^2$

## Abstract

Abstract Meta reinforcement learning (meta-RL) methods such as RL 2 have emerged as promising approaches for learning data-efficient RL algorithms tailored to a given task distribution. However, they show poor asymptotic performance and struggle with out-of-distribution tasks because they rely on sequence models, such as recurrent neural networks or transformers, to process experiences rather than summarize them using general-purpose RL components such as value functions. In contrast, traditional RL algorithms are data-inefficient as they do not use domain knowledge, but they do converge to an optimal policy in the limit. We propose RL 3 , a principled hybrid approach that incorporates action-values, learned per task through traditional RL, in the inputs to meta-RL. We show that RL 3 earns greater cumulative reward in the long term, compared to RL 2 , while maintaining data-efficiency in the short term, and generalizes better to out-of-distribution tasks. Experiments are conducted on both custom and benchmark discrete domains from the meta-RL literature that exhibit a range of short-term, long-term, and complex dependencies.

## 1 Introduction

Reinforcement learning (RL) has been shown to produce effective policies in a variety of applications including both virtual Mnih et al. () and embodied Schulman et al. (); Haarnoja et al. () systems. However, traditional RL algorithms have three major drawbacks: they often have difficulty generalizing beyond the exact task they were trained on, can be slow to converge, and require a large amount of data. These shortcomings are especially glaring in settings where the goal is to learn policies for a collection or distribution of problems that share some similarities, and for which traditional RL must start from scratch for each problem. For example, many robotic manipulation tasks require interacting with an array of objects with similar but not identical shapes, sizes, weights, materials, and appearances, such as mugs and cups. It is likely that effective manipulation strategies for these tasks will be similar, but they may also differ in ways that make it challenging to learn a single policy that is highly successful on all instances. Recently, meta reinforcement learning (meta-RL) has been proposed as an approach to mitigate these shortcomings by deriving RL algorithms (or meta-RL policies) that adapt efficiently to a distribution of tasks that share some common structure Duan et al. (); Wang et al. ().

**Table 1: RL^3 combines the strengths of meta-RL (e.g., RL^2) and traditional RL. Like RL^2, RL^3 uses finite-context sequence models to represent data-efficient RL algorithms, optimized for tasks within a specified distribution. However, RL^3 also includes a general-purpose RL routine that distills arbitrary amounts of data into optimal value-function estimates during adaptation. This improves long-term reasoning and out-of-distribution generalization.**
|  | RL | RL^2 | RL^3 |
| --- | --- | --- | --- |
| Short-Term Efficiency | x | \textpdfrender<br>TextRenderingMode=FillStroke,<br>LineWidth=.6pt, ✓ | \textpdfrender<br>TextRenderingMode=FillStroke,<br>LineWidth=.6pt, ✓ |
| Long-Term Performance | \textpdfrender<br>TextRenderingMode=FillStroke,<br>LineWidth=.6pt, ✓ | x | \textpdfrender<br>TextRenderingMode=FillStroke,<br>LineWidth=.6pt, ✓ |
| OOD Generalization | \textpdfrender<br>TextRenderingMode=FillStroke,<br>LineWidth=.6pt, ✓ | x | \textpdfrender<br>TextRenderingMode=FillStroke,<br>LineWidth=.6pt, ✓ |
|  | (Universal) |  | (Improved) |

While meta-RL systems represent a significant improvement over traditional RL in such settings, they still require large amounts of data during meta-training time, can have poor asymptotic performance during adaptation, and although they “learn to learn,” they often generalize poorly to tasks not represented in the meta-training distribution. This is partly because they rely on black-box sequence models like recurrent neural networks or transformers to process experience data. These models cannot handle arbitrary amounts of data effectively and lack integrated general-purpose RL components that could induce a broader generalization bias.

Hence, we propose RL^3, a principled approach that embeds the strengths of traditional RL within meta-RL. Table 1 highlights our primary aims and the foremost insight informing our approach. The key idea in RL^3 is an additional ‘object-level’ RL procedure executed within the meta-RL architecture that computes task-specific optimal Q-value estimates as supplementary inputs to the meta-learner, in conjunction with the sequence of states, actions and rewards. In principle, our approach allows the meta-learner to learn how to optimally fuse raw experience data with the summarizations provided by the Q-estimates. Ultimately, RL^3 leverages the general-purpose nature of Q-value estimates, their ability to compress large amounts of experiences into useful summaries, their direct actionability, and their asymptotic optimality in order to enhance long-term performance and out-of-distribution (OOD) generalization.
In this work, RL^3 is implemented by injecting Q-value estimates into RL^2 Duan et al. ().

The primary contribution of this paper is a confirmation of the hypothesis that injecting Q-estimates obtained via traditional object-level RL alongside the typical experience histories within a meta-RL agent leads to higher long-term returns and better OOD generalization, while maintaining short-term efficiency. We further demonstrate that our approach can also work with an abstract, or coarse, representation of the object-level MDP. Our conclusions are based on the results of experiments on the Bandits and MDPs domains used in previous work Duan et al. (); Mishra et al. () as well as on a more challenging custom GridWorld domain that requires long-term reasoning. FInally, we examine the key insights that inform our approach and show theoretically that object-level Q-values are directly related to the optimal meta-value function.

## 2 Related Work

Although meta-RL is a fairly new topic of research, the general concept of meta-learning is decades old Vilalta and
Drissi (), which, coupled with a significant number of design decisions for meta-RL systems, has created a large number of different proposals for how systems ought to best exploit the resources available within their deployment contexts Beck et al. (). At a high level, most meta-RL algorithms can be categorized as either parameterized policy gradient (PPG) models Finn et al. (); Li et al. (); Sung et al. (); Al-Shedivat et al. (); Gupta et al. (); Yoon et al. (); Stadie et al. (); Vuorio et al. (); Zintgraf et al. (); Raghu et al. (); Kaushik et al. (); Ghadirzadeh et al. (); Mandi et al. () or black box models Duan et al. (); Heess et al. (); Wang et al. (); Foerster et al. (); Mishra et al. (); Humplik et al. (); Fakoor et al. (); Yan et al. (); Zintgraf et al. (); Liu et al. (); Emukpere et al. (); Beck et al. (). PPG approaches assume that the underlying learning process is best represented as a policy gradient, where the set of parameters that define the underlying algorithm ultimately form a differentiable set of meta-parameters that the meta-RL system may learn to adjust. The additional structure provided by this assumption, combined with the generality of policy gradient methods, means that typically PPG methods retain greater generalization capabilities on out-of-distribution tasks. However, due to their inherent data requirements, PPG methods are often slower to adapt and initially train.

In this paper we focus on black box models, which represent the meta-learning function as a neural network, often a recurrent neural network (RNN) Duan et al. (); Heess et al. (); Wang et al. (); Humplik et al. (); Fakoor et al. (); Yan et al. (); Zintgraf et al. (); Liu et al. () or a transformer Mishra et al. (); Wang et al. (); Melo (). There are also several hybrid approaches that combine PPG and black box methods, either during meta-training Ren et al. () or fine-tuning Lan et al. (); Xiong et al. (). Using black box models simplifies the process of augmenting meta states with Q-estimates and allows us to retain relatively better data efficiency while relying on the Q-value injections for better long-term performance and generalization.

Meta-RL systems may also leverage extra information available during training, such as task identification Humplik et al. (); Liu et al. (). Such ‘privileged information’ can of course lead to more performant systems, but is not universally available. As our hypothesis does not rely on the availability of such information, we expect our approach to be orthogonal to, and compatible with, such methods. Black box meta-RL systems that do not use privileged information still vary in several ways, including the choice between on-policy and off-policy learning and, in systems that use neural networks, the choice between transformers Vaswani et al. () and RNNs Elman (); Hochreiter and
Schmidhuber (); Cho et al. ().

The most relevant methods to our work are end-to-end methods, which use a single function approximator to subsume both learner and meta-learner, such as RL^2 Duan et al. (), L2L Wang et al. (), SNAIL Mishra et al. (), and E-RL^2 Stadie et al. (), and methods that exploit the formal description of the meta-RL problem as a POMDP or a Bayes-adaptive MDP (BAMDP) Duff (). These methods attempt to learn policies conditioned on the BAMDP belief state while also approximating this belief state by, for example, variational inference (VariBAD) Zintgraf et al. (); Dorfman et al. (), or random network distillation on belief states (HyperX) Zintgraf et al. (). Or, they simply encode enough experience history to approximate POMDP beliefs (RL^2) Duan et al. (); Wang et al. ().

Our proposed method is an end-to-end system that exploits the BAMDP structure of the meta-RL problem by spending a small amount of extra computation to provide inputs to the end-to-end learner that more closely resemble important constituents of BAMDP value functions. Thus, the primary difference between this work and previous work is the injection of Q-value estimates into the meta-RL agent state at each meta-step, in addition to the state-action-reward histories. In this work, our approach, RL^3, is implemented by simply injecting Q-value estimates into RL^2 alongside experience history, although any other meta-RL algorithm can be used.

## 3 Background and Notation

In this section, we briefly cover some notation and concepts upon which this paper is built.

### 3.1 Partially Observable MDPs

Markov decision processes (MDPs) underpin reinforcement learning, and in this paper we use the standard notation defining an MDP as a tuple $M=\langle S,A,T,R\rangle$, where $S$ is a set of states; $A$ is a set of actions; $T$ is the transition and $R$ is the reward function. A *partially observable Markov decision process* (POMDP) extends MDPs to settings with partially observable states. A POMDP is described as a tuple $\langle S,A,T,R,\Omega,O\rangle$, where $S,A,T,R$ are as in an MDP. $\Omega$ is the set of possible observations, and $O:S\times A\times\Omega\rightarrow[0,1]$ is an observation function representing the probability of receiving observation $\omega$ after performing action $a$ and transitioning to state $s^{\prime}$. POMDPs can alternatively be represented as continuous-state belief-MDPs where a belief state $b\in\Delta^{|S|}$ is a probability distribution over all states. In this representation, a policy $\pi$ is a mapping from belief states to actions, $\pi:\Delta^{|S|}\rightarrow A$.

### 3.2 Reinforcement Learning

Reinforcement learning (RL) agents learn an optimal policy given an MDP with unknown dynamics using only transition and reward feedback. This is often done by incrementally estimating the optimal action-value function $Q^{*}(s,a)$ Watkins and Dayan (), which satisfies the Bellman optimality equation $Q^{*}(s,a)=\mathbb{E}_{s^{\prime}}[R(s,a)+\gamma\max_{a^{\prime}\in A}Q^{*}(s^
{\prime},a^{\prime})]$. In large or continuous state settings, it is popular to use deep neural networks to represent the action-value functions Mnih et al. (). We denote the vector representing the Q-estimates of all actions at state $s$ as $Q(s)$, and after $t$ feedback steps, as $Q^{t}(s)$. Q-learning is known to converge asymptotically Sutton and
Barto (), provided each state-action pair is explored sufficiently. As a rough general statement, $||Q^{t}(s)-Q^{*}(s)||_{\infty}$ is proportional to $\approx\frac{1}{\sqrt{t}}$, with strong results on the convergence error available Szepesvári (); Kearns and Singh (); Even-Dar et al. (). The theoretical objective in RL is to optimize the value of the final policy, i.e. the cumulative reward per episode, disregarding the data cost incurred and the cumulative reward missed (or regret) during learning due to suboptimal exploration.

### 3.3 Meta Reinforcement Learning

Meta reinforcement learning seeks action selection strategies that minimize regret in MDPs drawn from a distribution of MDPs that share the same state and action spaces. Therefore, the objective in meta-RL is to maximize the cumulative reward over the entire interaction (or adaptation) period with an MDP, which may span multiple episodes, in order to optimize the exploration-exploitation tradeoff. Formally,

$$ $\displaystyle\mathcal{J}(\theta)=\mathbb{E}_{M_{i}\sim\mathcal{M}}\Big{[}\sum_ {t=0}^{H}\gamma^{t}\mathbb{E}_{(s_{t},a_{t})\sim\rho_{i}^{\pi_{\theta}}}[R_{i} (s_{t},a_{t})]\Big{]}$ (1) $$

where the meta-RL policy $\pi_{\theta}$ is interpreted as a ‘fast’ or ‘inner’ RL algorithm that maps the experience sequence $(s_{0},a_{0},r_{0},...,s_{t})$ within an MDP $M_{i}$ to an action $a_{t}$ using either a recurrent neural network or a transformer network. $\rho_{i}^{\pi_{\theta}}$ is the state-action occupancy induced by the meta-RL policy in MDP $M_{i}$, and $H$ is the length of the adaptation period, or *interaction budget*. The objective $\mathcal{J}(\theta)$ is maximized using a conventional ‘slow’ or ‘outer’ deep RL algorithm, given the reformulation of the interaction period with an MDP as a single (meta-)episode in the objective function, which maximizes the cumulative reward throughout this period. We will use the term ‘experience history’, denoted by $\Upsilon$, to refer to the state-action-reward sequence within a meta-episode, which spans across multiple episodes $\{\tau_{0},\tau_{1},...\tau_{n}\}$. Fig. 1 illustrates how these components interconnect.

Another way to conceptualize this problem is to recognize that the meta-RL problem may be written as a meta-level POMDP, where the hidden variable is the particular MDP (or task) at hand, $M_{i}$, which varies across meta-episodes. This framing, known as Bayesian RL Ghavamzadeh et al. (), leverages the fact that augmenting the task-specific state $s$ with belief over tasks $b(i)$ results in a Markovian meta-state $[s,b]$ for optimal action selection, a model known as the Bayes Adaptive MDP (or BAMDP) Duff (). That is, this belief state captures all requisite information for the purpose of acting. We will revisit this concept to develop intuition on the role of object-level Q-value estimates in the meta-RL value function.

## 4 RL 3

To address the limitations of black box meta-RL methods, we propose RL^3, a principled approach that leverages (1) the inherent generality of action-value estimates, (2) their ability to compress experience histories into useful summaries, (3) their direct actionability & asymptotic optimality, (4) their ability to inform task-identification, and (5) their relation to the optimal meta-value function, in order to enhance out-of-distribution (OOD) generalization and performance over extended adaptation periods. The central, novel mechanism in RL^3 is an additional ‘object-level’ RL procedure executed within the meta-RL architecture, shown in Fig. 1, that computes task-specific optimal Q-value estimates $Q^{t}_{i}(s_{t})$ and state-action counts as supplementary inputs to the meta-RL policy in conjunction with the sequence of states, actions and rewards $(s_{0},a_{0},r_{0},...,s_{t})$. The Q-estimates are computed off-policy, and may involve model estimation and planning for greater data efficiency. The estimates and the counts are reset at the beginning of each meta-episode as a new task $M_{i}$ is sampled. In all subsequent text, Q-value estimates used as input entail the inclusion of state-action counts as well. We now present a series of key insights informing our approach.

Figure: Figure 1: Overview diagram of RL^3 . Black entities represent standard components from RL^2, and purple entities represent additions for RL^3 . $M_{i}$ is the current MDP; $s$ is a state; $r$ is a reward; $t_{i}$ and $t_{\tau}$ are the amount of time spent experiencing the current MDP and current episode, respectively; $Q_{i}^{t}$ is the Q-value estimate for MDP $i$ after $t$ actions; $\nabla\mathcal{J}$ is the policy gradient for meta-training.

First, estimating action-values is a key component in many universal RL algorithms, and asymptotically, they *fully* inform optimal behavior *irrespective of domain*. Strategies for optimal exploration-exploitation trade-off are domain-dependent and rely on historical data, yet many exploration approaches use estimated Q-values and some notion of counts *alone*, such as epsilon-greedy, Boltzmann exploration, upper confidence bounds (UCB/UCT) Auer (); Kocsis and
Szepesvári (), count-based exploration Tang et al. (), curiosity based exploration Burda et al. () and maximum-entropy RL Haarnoja et al. (). This creates a strong empirical case that using Q-value estimates and state-action counts for efficient exploration has inherent generality.

Figure: Figure 2: Sub-figure (a) shows a meta-episode in a shortest-path environment where the goal position (green circles) and the obstacles (black regions) may vary across tasks. In this meta-episode, after the meta-RL agent narrows its belief about the goal position of this task (dark-green circle) having followed a principled exploration strategy ($\tau_{0}$), it explores potential shorter paths in subsequent episodes ($\tau_{1},\tau_{2},\tau_{3},\tau_{4}$). Throughout this process, the estimated value-function $\hat{Q}^{*}$ implicitly “remembers” the goal position and previous paths traversed in a finite-size representation, and updates the shortest path calculation (highlighted in bold) using Bellman backups when paths intersect. Sub-figures (b) and (c) illustrate the many-to-one mapping of object- and meta-level data streams to Q-estimates, and thus their utility as compression and summarization mechanisms for meta-learning.

Second, Q-estimates summarize experience histories of arbitrary length *and order* in one constant-size vector. This mapping is many-to-one, and any permutation of transitions ($\langle s,a,r,s^{\prime}\rangle$ tuples) or episodes in a history of experiences yield the same Q-estimates. Although this compression is lossy, it still “remembers” important aspects of the experienced episodes, such as high-return actions and goal positions (see Fig. 2) since Q-estimates persist across episodes. This simplifies the mapping the meta agent needs to learn as Q-estimates represent a smaller and more salient set of inputs compared to all possible histories with the same implication.

Third, Q-estimates are actionable. Estimated off-policy, they explicitly represent the optimal exploitation policy for the current task given the data insofar as the RL module is data-efficient, relieving the meta-RL agent from performing such calculations inside the transformer/RNN. Over time, Q-estimates become more reliable and directly indicate the optimal policy whereas processing raw data becomes more challenging. Fortunately, by incorporating Q-estimates the meta-RL agent can eventually ignore the history in the long run (or towards the end of the interaction period) and simply exploit the Q-estimates by selecting actions greedily.

Fourth, Q-estimates are excellent task discriminators and serve as another line of evidence vis-à-vis maintaining belief over tasks. In a simple domain like Bernoulli multi-armed bandits Duan et al. (), Q-estimates and action-counts combined are sufficient for Bayes-optimal behavior even without providing raw experience data – a result surprisingly unstated in the literature to the best of our knowledge (see Appendix A.1). However, Q-estimates and action-counts may not always be sufficient for Bayes-optimal beliefs. For example, in Gaussian multi-armed bandits, the sufficient statistics include the variance in rewards for each action (see Appendix A.2). In more complex domains, it is hard to prove the sufficiency of Q-estimates regarding task discrimination. However, via empirical analysis in Appendix D, we argue that i) it is highly improbable for two tasks to have similar Q^∗ functions and ii) Q-estimates tend to become accurate task predictors in just a few steps. This implies that the meta-agent may use this finite summary for task inference rather than relying completely on arbitrarily long histories, potentially contributing to enhanced performance over long adaptation periods.

It can be theoretically argued that since the meta agent is a BAMDP *policy*, it is meta-trained to select greedy actions w.r.t. the BAMDP meta-value function and thus should not require constructing a task-specific plan internally. However, the optimality of the meta action-value function depends on implicitly (or explicitly in some approaches, such as Humplik et al. (); Zintgraf et al. (); Dorfman et al. (); Zintgraf et al. ()) maintaining a Bayes-optimal belief over tasks in the transformer/RNN architecture. This may be challenging if the task distribution is too broad and the function approximator is not powerful enough to integrate experience histories into Bayes-optimal beliefs, or altogether impossible if there is a distribution shift at meta-test time. This latter condition is common in practice and is a frequent target use case for meta-RL systems. Incorporating task-specific Q-estimates gives the agent a simple alternative (even if not Bayes-optimal) line of reasoning to translate experiences into actions. Incorporating Q-estimates thus reduces susceptibility to distribution shifts since the arguments presented in this section are domain independent.

Finally, Q-estimates often converge far more quickly than the theoretical rate of $\frac{1}{\sqrt{t}}$, allowing them to be useful in the short and medium term, since i) most real-world domains contain significant determinism, ii) it is not necessary to estimate Q-values for states unreachable by the optimal policy, and iii) optimal meta-RL policies may represent active exploration strategies in which Q-estimates converge faster, or evolve in a manner leading to quicker task identification. This is intuitively apparent in shortest-path problems, as illustrated in Fig. 2(a). In a deep neural network, it is difficult to know exactly how Q-estimates will combine with state-action-reward histories when approximating the meta-value function. However, as we show below, we can rather straightforwardly write an equation for the meta-value function in terms of these constituent streams of information, which may explain why this function is seemingly relatively easy to learn compared to predicting meta-values from histories alone.

### 4.1 Theoretical Justification

Here, we consider the interpretation of meta-RL as performing RL on a partially observable Markov decision process (POMDP) in which the partially observable state factor is the identity of the object-level MDP. Without loss of generality, all analysis assumes the infinite horizon setting. We will denote meta-level entities, belonging in this case to a POMDP, with an overbar. For example, we have a meta-level value function $\bar{V}$ and a meta-level belief $\bar{b}$.

First, we show a basic result, that the optimal meta-level value function is upper bounded by the object-level Q-value estimates in the limit.

*Proof:* Given a task distribution $\mathcal{M}$, then for state $s$, there exists a maximum object-level optimal value function $V^{*}_{max}(s)$, corresponding to some MDP $M_{max}\in\mathcal{M}$, such that for all MDPs $M_{i}\in\mathcal{M}$, $V^{*}_{max}(s)\geq V^{*}_{i}(s)$. Observe that the expected cumulative discounted reward experienced by the agent cannot be greater than the most optimistic value function over all tasks, since $\bar{V}^{*}(\bar{b})$ is a weighted average of individual value functions $V^{\pi_{\theta}}(s)$, which are themselves upper bounded by $V^{*}_{max}(s)$. Thus,

$$ $\max_{M_{i}\in\mathcal{M}}V^{*}_{i}(s)\geq\bar{V}^{*}(\bar{b})\hskip 14.22636 pt\forall s\in S.$ (2) $$

Next, we see that combining the asymptotic accuracy of Q-estimates and Equation (2) gives us

$$ $\lim_{t\rightarrow\infty}\max_{a\in A,M_{i}\in\mathcal{M}}Q_{i}^{t}(s,a)\geq \bar{V}^{*}(\bar{b})\hskip 14.22636pt\forall s\in S.\qed$ (3) $$

Furthermore, it follows if the meta-level observation $\bar{\omega}$ includes Q-value estimates of the current task $M_{i}$, it can be shown that as $t\rightarrow\infty$, the optimal meta-value function approaches the optimal value function for the current task, i.e., for any $\epsilon>0$, there exists $\kappa\in\mathbb{N}$ such that for $t\geq\kappa$,

$$ $\Big{|}\max_{a\in A}\Big{[}Q_{i}^{t}(s,a)\Big{]}-\bar{V}^{*}(\bar{b})\Big{|} \leq\epsilon\hskip 14.22636pt\forall s\in S.$ (4) $$

Equation (4) (see full proof in Appendix A.3) shows that for $t\geq\kappa$, acting greedily w.r.t. $Q_{i}^{*}$ leads to Bayes-optimal behavior, and knowing the Bayes-optimal belief over tasks is not required, which also implies that the entire experience history can be ignored at that point. Moreover, it follows from equation (4) that for $t<\kappa$,

$$ $\displaystyle\bar{V}^{*}(\bar{b})=\max_{a\in A}\Big{[}Q_{i}^{t}(s,a)\Big{]}+ \varepsilon_{i}(\Upsilon)$ (5) $$

where error $\varepsilon_{i}(\Upsilon)$ is the error in Q-value estimates, which is a function of the dynamics of MDP $M_{i}$ and data (both amount and quality) seen by the agent so far in the experience history $\Upsilon$. While this error will diminish as $t\rightarrow\infty$, in the short run, a function $f(\Upsilon)$ could be learned to either estimate the error or estimate $\bar{V}^{*}(\bar{b})$ entirely.

The better performance of RL^3 could be explained by either error $\varepsilon_{i}(\Upsilon)$ being simpler to estimate, or, the meta-agent behavior being more robust to errors in estimates of $\varepsilon_{i}(\Upsilon)$ when Q-estimates are supplied directly as inputs, than to errors in a more complicated approximation of $\bar{V}^{*}(\bar{b})$. Moreover, this composition benefits from the fact that the convergence rate for Q-estimates suggests a natural, predictable rate of shifting reliance from $f(\Upsilon)$ to $Q^{t}_{i}(s)$ as $t\rightarrow\infty$. However, we do not bake this structure into the network and instead let it implicitly learn how much to use the Q-estimates.

Finally, we note that near-perfect function approximation of $\bar{V}^{*}(\bar{b})$ as $t\rightarrow\infty$ reduces error in meta-value function approximation for all preceding belief states, as meta-values for consecutive belief states $\bar{b}$ and $\bar{b}^{\prime}$ are linked through the Bellman equation for BAMDPs (see details in Appendix A.3):

$$ $\bar{V}^{*}(\bar{b})=\max_{a\in A}\Big{[}\!\sum_{M_{i}\in\mathcal{M}}\!\!\bar{ b}(i)R_{i}(s,a)+\gamma\!\sum_{\bar{\omega}\in\bar{\Omega}}\bar{O}(\bar{\omega} |\bar{b},a)\bar{V}^{*}(\bar{b}^{\prime})\Big{]},$ (6) $$

This dependency helps in meta-training in RL^3, as most (outer) RL algorithms used for training the meta-critic have a temporal-difference based learning component. Without conditioning on Q-estimates, error in $\bar{V}^{*}(\bar{b})$ would instead increase as $t\rightarrow\infty$, as the meta-critic would be conditioned on a larger history, which could destabilize the meta-value learning for all preceding belief states during meta-training.

Figure: Algorithm 1 Value-Augmenting Wrapper for Discrete MDPs

### 4.2 Implementation

Implementing RL^3 involves simply replacing each MDP in the task distribution with a corresponding value-augmented MDP (VAMDP) and solving the resulting VAMDP distribution using RL^2. Each VAMDP has the same action space and reward function as the corresponding MDP. The value augmented state $\hat{s}_{t}\in S\times\mathbb{R}^{k}\times\mathbb{I}^{k}$ includes the object level state $s_{t}$, $k$ real values and $k$ integer values for the Q-estimates ($Q^{t}(s_{t},a)$) and action counts ($N^{t}(s_{t},a)$) for each of the $k$ actions. When the object-level state space $S$ is discrete, $s_{t}$ needs to be represented as an $|S|$-dimensional one-hot vector. Note that the value augmented state space is continuous. In the VAMDP transition function, the object-level state $s$ has the same transition dynamics as the original MDP, while the dynamics of Q-estimates are a function of $T$, $R$, and the specific object-level RL algorithm used for estimating Q-values. An episode of the VAMDP spans the entire interaction period with the corresponding MDP, which may include multiple episodes of the MDP, as Q-estimates continue to evolve beyond episode boundaries. In code, a VAMDP RL environment is implemented as a wrapper over a given MDP environment. The pseudocode is provided in Algorithm 1 and additional implementation details and hyperparameters for RL^2 and RL^3 are mentioned in Appendix B.

## 5 Experiments

We compare RL^3 to a modified version of RL^2. First, we replace LSTMs with transformers in both the meta-actor and meta-critic for the purpose of mapping experiences to actions and meta-values, respectively. This is done to improve RL^2 ’s ability to handle long-term dependencies instead of suffering from vanishing gradients. Moreover, RL^2 -transformer trains significantly faster than RL^2 -LSTM. Second, we include in the state space the total number of interaction steps and the total number of steps within each episode during a meta-episode (see Fig. 1). Third, we use PPO Schulman et al. () for training the meta actor-critic, instead of TRPO Schulman et al. (). These modifications and other minor-implementation details incorporate the recommendations made by Ni et al. (), who show that model-free recurrent RL is competitive with other state-of-the-art meta RL approaches such as VeriBAD Zintgraf et al. (), if implemented properly. RL^3 simply applies the modified version of RL^2 to the distribution of value-augmented MDPs explained in section 4.2. Within each VAMDP, our choice of object-level RL is a model-based algorithm to maximize data efficiency – we estimate a tabular model of the environment and run finite-horizon value-iteration using the model.

Figure: Figure 3: An RL^3 policy on a selected meta-episode visualized using a sequence of snapshots. ‘S’ is the starting tile, ‘G’ is the goal tile and the black circle shows the current position of the agent. Blue tiles marked ‘W’ are wet tiles. Wet tiles always lead to the agent slipping to one of the directions orthogonal to the intended direction of movement. Entering wet tiles yield an immediate reward of -2. Yellow tiles marked ‘!’ are warning tiles and entering them causes -10 reward. Red tiles marked ‘X’ are fatally dangerous. Entering them ends the episode and leads to a reward of -100. Black tiles are obstacles. White tiles yield a reward of -1 to incentive the agent to reach the goal quickly. On all tiles other than wet tiles, there is a chance of slipping sideways with a probability of 0.2. The object-level state-values $v^{t}(s)=\text{max}_{a}Q^{t}(s,a)$, as approximated by object-level RL, is represented using shades of green (and the accompanying text), where darker shades represent higher values.
Refer to caption: extracted/2306.15909v4/figures/meta-explore-4-left.png

In our test domains, each meta-episode involves procedurally generating an MDP according to a parameterized distribution, which the meta-actor interacts with for a fixed interaction budget $H$. This interaction might consist of multiple object-level episodes of variable length, each of which are no longer than a maximum task horizon. For a given experiment, each approach is trained on the same series of MDPs. For testing, each approach is evaluated on an identical set of 1000 MDPs distinct from the training MDPs. For testing OOD generalization, MDPs are generated from distributions with different parameters than in training. We select three task domains for our experiments, which cover a range of short-term, long-term, and complex dependencies.

Bernoulli Bandits: We use the same setup described in Duan et al. () with $k=5$ arms. To test OOD generalization, we generate bandit tasks by sampling success probabilities from $\mathcal{N}(0.5,0.5)$.

Random MDPs: We use the same setup described in Duan et al. (). The MDPs have 10 states and 5 actions with stochastic transitions and rewards, and the task horizon is 10. The mean rewards and transition probabilities are drawn from a normal and a flat Dirichlet distribution, respectively. To test OOD generalization, the rewards are generated deterministically and initialized from $\mathcal{U}(0,2)$.

GridWorld Navigation: A set of navigation tasks in a 2D grid environment. We experiment with 11x11 (121 states) and 13x13 (169 states) grids. The agent always starts in the center of the grid and needs to navigate through obstacles to a single goal location. The grid also contains slippery tiles, fatally dangerous tiles and warning tiles surrounding the latter. See Fig. 3(a) for an example of a 13x13 grid. Instead of one-hot vectors, we use the 2D $(x,y)$ grid location to represent agent state. To test OOD generalization, we vary parameters including the stochasticity of actions, density of obstacles and the number of dangerous tiles. For this domain, we consider an additional variation of RL^3, called RL^3 -coarse where a given grid is partitioned into clusters of tiles (or abstract states), each of size 2. Abstract states are comprised only of adjacent, traversable cells, and are used *solely* for the purpose of estimating the object-level Q-values. Our goal is to test whether coarse-level Q-value estimates are still useful to the meta-RL policy. The domains and the abstraction strategy used for the RL^3 -coarse approach are described in greater detail in Appendices E and B.3, respectively.

## 6 Results

In this section, we present several experiments that demonstrate the effectiveness of RL^3 . Beyond matching or exceeding the performance of RL^2 in all test domains, RL^3 also shows better OOD generalization, which we attribute to the increased generality of the Q-value representation. More striking, the advantages of RL^3 increase significantly with longer interactions periods and less stochastic tasks. We hypothesize that this is due to the increased accuracy of the Q-value estimates in these cases. Last, we find that RL^3 performs well even with coarse-grained object-level RL over abstract states, showing minimal drop in performance in most cases and even occasional increases, while the computational savings are substantial.

We emphasize that the core of our approach, which is augmenting MDP states with action-value estimates, is not inherently tied to RL^2 and is orthogonal to most other meta-RL research. VAMDPs can be plugged into any base meta-RL algorithm with a reasonable expectation of improving it.

Bandits: Table 2 shows the results for the Bandits domains. For $H=100$ and $H=500$, both approaches perform comparably. However, the OOD generalization for RL^3 is slightly better. We also experiment with a Markovian version of RL^3, where we use a feed-forward neural network that is conditioned only on the Q-estimates and action-counts, since those are sufficient for Bayes-optimal behavior in this domain. As expected, the results are similar to those with regular RL^3 .

MDPs: Table 3 shows the results for the MDPs domains. Once again, for $H=100$ and $H=500$, both approaches perform comparably, and once again, OOD generalization is slightly better for RL^3 . We suspect that for such short budgets on this highly stochastic domain, Q-estimates do not converge enough to be very useful. To test this hypothesis, we test both approaches by applying the models trained for $H=500$ to $H=1000$ (row 4 in Table 3) using a moving window context for the transformer models (see Appendix B for details). Here, RL^3 generalizes significantly better, demonstrating the utility of Q-estimates when they are allowed to converge for more iterations. In fact, the score achieved by RL^3 when trained for $H=500$ and tested on $H=1000$ is similar to that of the original RL^2 implementation Duan et al. () when trained specifically for $H=1000$.

GridWorlds: Table 4 shows the results for the GridWorld Navigation domain. On 11x11 grids with $H=250$, RL^3 significantly outperforms RL^2. On 13x13 grids with $H=350$, the outperformance margin is even greater, showing that while RL^2 struggles with a greater number of states, a longer adaptation period and more long-term dependencies, RL^3 can take advantage of the Q-estimates to overcome the challenge. We also test the OOD generalization of both approaches in different ways by increasing the obstacle density (Dense), making actions on non-water tiles deterministic (Deterministic), increasing the number of wet ‘W’ tiles (Watery), increasing the number of danger ‘X’ tiles (Dangerous) and having the goal only in the corners (Corner). On all variations, RL^3 continues to significantly outperform RL^2. In a particularly interesting outcome, both approaches show improved performance on the Deterministic variation. However, RL^3 gains 80% more points than RL^2 does, which is likely because Q-estimates converge faster on this less stochastic MDP and therefore provide greater help to RL^3 . Conversely, in the Watery variation, which is more stochastic, both RL^2 and RL^3 lose roughly equal number of points. Finally, in each case, RL^3 -coarse significantly outperforms RL^2. In fact, it performs on par with RL^3, even outperforming it on Corner variation, except on the canonical 13x13 case and its Deterministic variation, where it scores about 90% of the scores for RL^3 .

**Table 2: Test scores (mean $\pm$ standard error) for Bandits domain and the ^†OOD variation.**
| Budget $H$ | RL^2 | RL^3 | RL^3 (Markov) |
| --- | --- | --- | --- |
| 100 | $76.9\pm 0.6$ | $77.5\pm 0.5$ | $75.2\pm 0.5$ |
| 500 | $392.1\pm 2.5$ | $393.2\pm 2.7$ | $391.75\pm 2.6$ |
| 500^† | $430.2\pm 2.8$ | $\mathbf{434.9}\pm 2.8$ | $433.7\pm 2.8$ |

**Table 3: Test scores (mean $\pm$ standard error) for MDPs domain and the ^†OOD variation. ^+Time limit doubled without retraining policy.**
| Budget $H$ | RL^2 | RL^3 |
| --- | --- | --- |
| 100 | $159.5\pm 0.8$ | $158.9\pm 0.8$ |
| 500 | $927.8\pm 3.7$ | $926.9\pm 3.7$ |
| 500^† | $772.8\pm 1.7$ | $\mathbf{775.9}\pm 1.7$ |
| 1000^+ | $1871.8\pm 7.4$ | $\mathbf{1916.8}\pm 7.4$ |

**Table 4: Test scores (mean $\pm$ standard error) for GridWorlds domain and the ^†OOD variations (13x13, $H$=350).**
| Variation | RL^2 | RL^3 | RL^3 -coarse |
| --- | --- | --- | --- |
| 11x11, $H$=250 | $524\pm 22$ | $\mathbf{630}\pm 22$ | $\mathbf{613}\pm 21$ |
| 13x13, $H$=350 | $584\pm 28$ | $\mathbf{902}\pm 27$ | $831\pm 28$ |
| Dense^† | $384\pm 26$ | $\mathbf{691}\pm 28$ | $\mathbf{674}\pm 28$ |
| Deterministic^† | $960\pm 38$ | $\mathbf{1574}\pm 35$ | $1464\pm 36$ |
| Watery^† | $513\pm 26$ | $\mathbf{826}\pm 27$ | $\mathbf{822}\pm 28$ |
| Dangerous^† | $283\pm 28$ | $\mathbf{647}\pm 30$ | $\mathbf{657}\pm 30$ |
| Corner^† | $319\pm 23$ | $508\pm 23$ | $\mathbf{645}\pm 23$ |

Fig. 3 shows a sequence of snapshots of a meta-episode where the trained RL^3 agent is interacting with an instance of a 13x13 grid. The first snapshot shows the agent just before reaching the goal for the first time. Prior to the first snapshot, the agent had explored many locations in the grid. The second snapshot shows the next episode just after the agent finds the goal, resulting in value estimates being updated using object-level RL for all visited states. Snapshot 3 shows the agent consequently using the Q-estimates to navigate to the goal presumably by choosing high-value actions. The agent also explores several new nearby states for which it does not have Q-estimates. Snapshot 4 shows the final Q-value estimates.
A short video of the GridWorld environment, showing both RL^2 and RL^3 agents solving a set of problem instances side-by-side, is included in the supplementary material.

Computation Overhead Considerations: As mentioned earlier, for implementing object-level RL, we use model estimation followed by finite-horizon value-iteration to obtain Q-estimates. The computation overhead is negligible for Bandits (5 actions, task horizon = 1) and very little for the MDPs domain (10 states, 5 actions, task horizon 10). For 13x13 GridWorlds (up to 169 states, 5 actions, task horizon = 350), RL^3 takes approximately twice the computation time of RL^2 per meta-episode. However, RL^3 -coarse requires only 10% overhead while still outperforming RL^2 and retaining more than 90% of the performance of RL^3 . This demonstrates the utility of state abstractions in RL^3 for scaling. Our implementation is available in the supplementary material.

## 7 Conclusion

In this paper, we introduced RL^3, a principled hybrid approach that combines the strengths of traditional RL and meta-RL and provides a more robust and adaptable reinforcement learning algorithm for complex and diverse environments. We advanced intuitive and theoretical arguments regarding its suitability for meta-RL and presented empirical evidence to validate those ideas. Specifically, we demonstrated that RL^3 outperforms RL^2 across a wide range of tasks, enhancing long-term performance and generalization on out-of-distribution tasks, while maintaining efficiency in the short term. In future work, we plan to explore extending RL^3 to handle continuous state spaces, possibly using state abstractions for discretization.

## Acknowledgments

This work was supported in part by the National Science Foundation grant numbers 1954782, 2205153, and 2321786.

## A Proofs

### A.1 Bayes Optimality of Q-value Estimates in Bernoulli Multi-armed Bandits

Given an instance of a Bernoulli multi-armed bandit MDP, $M_{i}\sim\mathcal{M}$, and trajectory data $\Upsilon_{1:T}$ up to time $T$, we would like to show that the probability $P(i|\Upsilon_{1:T})$ can be determined entirely from $Q$-estimates $Q_{i}^{T}$ and action-counts $N_{i}^{T}$, as long as the initial belief is uniform or known.

In the following proof, we represent an instance $i$ of $K$-armed Bandits as a $K$-dimensional vector of success probabilities $[p_{i1},...,p_{iK}]$, such that pulling arm $k$ is associated with reward distribution $P(r=1|i,k)=p_{ik}$ and $P(r=0|i,k)=(1-p_{ik})$.

Let the number of times arm $k$ is pulled up to time $T$ be $N_{ik}^{T}$, and the number of successes associated with pulling arm $k$ up to time $T$ be $q_{ik}^{T}$. Given that this is an MDP with just a single state and task horizon of 1, the $Q$-estimate associated with arm $k$ is just the average reward for that action, which is the ratio of successes to counts associated with that action i.e., $Q_{ik}^{T}=\frac{q_{ik}^{T}}{N_{ik}^{T}}$. To reduce the clutter in the notation, we will drop the superscript $T$ for the rest of the subsection.

Now,

$$ $\displaystyle P(i|\Upsilon_{1:T})$ $\displaystyle=\alpha P(i)\cdot P(\Upsilon_{1:T}|i)$ (7) $$

where $\alpha$ is the normalization constant, $P(i)$ is the prior probability of task $i$ (which is assumed to be known beforehand), and $\Upsilon_{1:T}$ is the sequence of actions and the corresponding rewards up to time $T$. Assuming, without loss of generality, that the sequence of actions used to disambiguate tasks is a given, $P(\Upsilon_{1:T}|i)$ becomes simply the product of probabilities of reward outcomes up to time $T$, noting that the events are independent. Therefore,

$$ $\displaystyle P(\Upsilon_{1:T}|i)$ $\displaystyle=\prod_{k=1:K}\prod_{t=1:T}([r_{tk}=1]p_{ik}+[r_{tk}=0](1-p_{ik}))$ (8) $\displaystyle=\prod_{k=1:K}p_{ik}^{q_{ik}}\cdot(1-p_{ik})^{N_{ik}-q_{ik}}$ (9) $\displaystyle=\prod_{k=1:K}p_{ik}^{Q_{ik}N_{ik}}\cdot(1-p_{ik})^{N_{ik}-Q_{ik} N_{ik}}$ (10) $$

Putting everything together,

$$ $\displaystyle P(i|\Upsilon_{1:T})$ $\displaystyle=\alpha P(i)\cdot\prod_{k=1:K}p_{ik}^{Q_{ik}N_{ik}}\cdot(1-p_{ik} )^{N_{ik}-Q_{ik}N_{ik}}$ (11) $$

This equation proves that $N_{i}^{T}$ and $Q_{i}^{T}$ are sufficient statistics to determine $P(i|\Upsilon_{1:T})$ in this domain, assuming that the prior over task distribution is known. ∎

### A.2 Non-Bayes Optimality of Q-value Estimates in Gaussian Multi-armed Bandits

Given an instance of a Gaussian multi-armed bandit MDP, $M_{i}\sim\mathcal{M}$, and trajectory data $\Upsilon_{1:T}$ up to time $t$, here we derive the closed-form expression of the probability $P(i|\Upsilon_{1:T})$ and show that it contains terms other than $Q$-estimates $Q_{i}^{t}$ and action-counts $N_{i}^{t}$.

In the following proof, we represent an instance $i$ of $K$-armed Bandits as a $2K$-dimensional vector of means and standard deviations $[\mu_{i1},...,\mu_{iK},\sigma_{i1},...,\sigma_{iK}]$, such that pulling arm $k$ is associated with reward distribution $P(r|i,k)=\frac{1}{\sqrt{2\pi\sigma_{ik}}}\exp(\frac{r-\mu_{ik}}{\sigma_{ik}})^
{2}$.

Let the number of times arm $k$ is pulled up to time $T$ be $N_{ik}^{T}$. Given that this is an MDP with just a single state and the task horizon is 1, the $Q$-estimate associated with arm $k$ is just the average reward for that action $\text{Avg}[r_{k}]$ up to time $T$. To reduce the clutter in the notation, we will drop the superscript $T$ for the rest of the subsection.

As in the previous subsection, we now compute the likelihood $P(\Upsilon_{1:T}|i)$.

$$ $\displaystyle P(\Upsilon_{1:T}|i)$ $\displaystyle=\prod_{k=1:K}\prod_{t=1:T}\frac{1}{\sqrt{2\pi\sigma_{ik}}}\exp( \frac{r_{tk}-\mu_{ik}}{\sigma_{ik}})^{2}$ (12) $$

Therefore, the log likelihood is

$$ $\displaystyle\log P(\Upsilon_{1:T}|i)=\sum_{k=1:K}\sum_{t=1:T}\frac{(r_{tk}- \mu_{ik})^{2}}{\sigma_{ik}^{2}}-\log{(2\pi\sigma_{ik})}/2$ (13) $\displaystyle=\sum_{k=1:K}N_{ik}\frac{\text{Avg}[(r_{tk}-\mu_{ik})^{2}]}{ \sigma_{ik}^{2}}$ $\displaystyle\hskip 90.00015pt-N_{ik}\log{(2\pi\sigma_{ik})}/2$ (14) $\displaystyle=\sum_{k=1:K}N_{ik}\frac{\text{Avg}[r_{k}^{2}]-2\mu_{ik}\text{Avg }[r_{k}]+\mu_{ik}^{2}}{\sigma_{ik}^{2}}$ $\displaystyle\hskip 90.00015pt-N_{ik}\log{(2\pi\sigma_{ik})}/2$ (15) $\displaystyle=\sum_{k=1:K}N_{ik}\frac{(\text{Var}[r_{k}]+\text{Avg}[r_{k}]^{2} )-2\mu_{ik}\text{Avg}[r_{k}]+\mu_{ik}^{2}}{\sigma_{ik}^{2}}$ $\displaystyle\hskip 90.00015pt-N_{ik}\log{(2\pi\sigma_{ik})}/2$ (16) $\displaystyle=\sum_{k=1:K}N_{ik}\frac{\text{Var}[r_{k}]+(Q_{ik})^{2}-2\mu_{ik} Q_{ik}+\mu_{ik}^{2}}{\sigma_{ik}^{2}}$ $\displaystyle\hskip 90.00015pt-N_{ik}\log{(2\pi\sigma_{ik})}/2$ (17) $$

Therefore, computing this expression requires computing the variance in rewards, $\text{Var}[r_{k}]$, associated with each arm up to time $T$, apart from the $Q$-estimates and action-counts. This proves that $Q$-estimates and action-counts alone are insufficient to completely determine $P(i|\Upsilon_{1:T})$ in Gaussian multi-armed bandits domain. ∎

### A.3 Object-level Q-estimates and Meta-level Values

*Proof of Equation (4):* In standard meta-RL, the only observed variable in the POMDP state $\bar{s}_{t}=[s_{t},i]$ at time $t$ is the state $s_{t}$ of the current MDP i.e., $\bar{\omega}_{t}=s_{t}$, while the task identity $i$ is hidden. However, in RL^3, $\bar{\omega}_{t}$ includes the vector of $Q$-estimates $Q^{t}_{i}(s_{t})$ for the hidden task, which means that the meta-level observation function $\bar{O}(\bar{\omega}|\bar{b},a)$ factors in the probability that a particular $Q$-esimate will be observed following an action $a$ given an initial belief $\bar{b}$ state. (Note that we will use $\bar{b}(\bar{s})$ and $\bar{b}(i)$ interchangeably since $i$ is the only hidden variable in $\bar{s}$). In practice, such $Q$-value estimates provide excellent evidence (see Appendix D) for task identification. This allows for robust belief recovery even if the initial belief is not Bayes-optimal (or altogether not maintained), especially as the Q-estimates converge and stabilize in the limit, leading to two cases:

Case 1: The observed $Q$-values are unique to MDP $M_{i}$. In this case, the belief distribution will collapse rapidly to zero for tasks $j\not=i$, and thus $\max_{a\in A}Q_{i}(s,a)=\bar{V}^{*}(\bar{b})$.

Case 2: The observed $Q$-values are not unique. In this case, belief will not collapse to a single MDP. However, belief will still reduce to zero for tasks not compatible with the observed $Q$-values. The meta-level value function $\bar{V}^{*}(\bar{b})$, which will be an expectation over object-level values, will simplify to $\max_{a\in A}Q_{i}(s,a)$ since $Q$-values for all remaining tasks are identical, where $i$ may represent any of the (identical $Q$-valued) tasks with non-zero belief.

This proves equation (4). Note that in the limit, the task can be identified perfectly from the stream of experiences as all state-action pairs are explored, and the meta-level value function becomes equivalent to the optimal object-level value function of the identified (or current) task. However, the above proof demonstrates that RL^3 can infer this equivalency implicitly in the limit without relying on the stream of experiences or identifying the task fully, and furthermore, directly model the meta-value function in terms of the supplied object-level value function. ∎

*Proof of Equation (6):* We first write the Bellman equation for the optimal meta-level POMDP value function in its belief-MDP representation:

$$ $\bar{V}^{*}(\bar{b})=\max_{a\in A}\Big{[}\sum_{\bar{s}\in\bar{S}}\bar{b}(\bar{ s})\bar{R}(\bar{s},a)+\gamma\sum_{\bar{\omega}\in\bar{\Omega}}\bar{O}(\bar{ \omega}|\bar{b},a)\bar{V}^{*}(\bar{b}^{\prime})\Big{]}.$ (18) $$

However, given that in the POMDP state $\bar{s}=[s,i]$, the only hidden variable is the task $i$, we can re-write this as

$$ $\bar{V}^{*}(\bar{b})=\max_{a\in A}\Big{[}\sum_{M_{i}\in\mathcal{M}}\bar{b}(i)R _{i}(s,a)+\gamma\sum_{\bar{\omega}\in\bar{\Omega}}\bar{O}(\bar{\omega}|\bar{b} ,a)\bar{V}^{*}(\bar{b}^{\prime})\Big{]},$ (19) $$

where $\bar{b}(i)$ denotes the meta-level belief that the agent is operating in MDP $M_{i}$, and $R_{i}(s,a)$ is the reward experienced by the agent if it executes action $a$ in state $s$ in MDP $M_{i}$. Here, $\bar{b}^{\prime}$ may be calculated via the belief update as in §3.1. ∎

## B Architecture

### B.1 RL 2

Our modified implementation of RL^2 uses transformer decoders Vaswani et al. [] instead of RNNs to map trajectories to action probabilities and meta-values, in the actor and the critic, respectively, and uses PPO instead of TRPO for outer RL. The decoder architecture is similar to Vaswani et al. [], with 2 layers of masked multi-headed attention. However, we use learned position embeddings instead of sinusoidal, followed by layer normalization. Our overall setup is similar to Esslinger et al. [].

For each meta-episode of interactions with an MDP $M_{i}$, the actor and the critic transformers look at the entire history of experiences up to time $t$ and output the corresponding action probabilities $\pi_{1}...\pi_{t}$ and meta-values $\bar{V}_{1}...\bar{V}_{t}$, respectively. An experience input to the transformer at time $t$ consists of the previous action $a_{t-1}$, the latest reward $r_{t-1}$, the current state $s_{t}$, episode time step $t_{\tau}$, and the meta-episode time step $t$, all of which are normalized to be in the range $[0,1]$. In order to reduce inference complexity, say at time step $t$, we append $t$ new attention scores (corresponding to experience input $t$ w.r.t. the previous $t-1$ experience inputs) to a previously cached $(t-1)\times(t-1)$ attention matrix, instead of recomputing the entire $t\times t$ attention matrix. This caching mechanism is implemented for each attention head and reduces the inference complexity at time $t$ from $\mathcal{O}(t^{2})$ to $\mathcal{O}(t)$. Note that this caching mechanism is possible only when the input to the transformer is always the entire history, built up incrementally, instead of a moving window, since removing early experiences from the moving window would invalidate the cached attention scores as they attend to the entire history.

[A note on using a moving context to apply a model trained $H=500$ to $H=1000$: A naive implementation would use a moving window context of 500 time steps. However, with this choice, a model highly biased towards exploiting rather than exploring near the end of the window would only exploit from time steps 501 through 1000 even though there is room more exploration. As a workaround to this problem, we use a context length of 250 steps until time step 750, and increase it gradually thereafter to 500 steps towards the end. In our experiments, this heuristic rule led to a significant improvement over the naive method.]

Note that our results are not significantly better than those in the original RL^2 paper, which uses a recurrent neural network, for the Bandits and the MDPs domain. However, our transformer version performs better on the GridWorld domain. Moreover, the transformer version drastically reduces real-world training time and is more stable and easier to tune.

### B.2 RL 3

The transformer input in RL^3 includes of a vector of $Q$-estimates and a vector of action counts at each step $t$ for the corresponding state. As mentioned in Section 4.2, this is implemented in our code by simply converting MDPs in the problem set to VAMDPs using a wrapper, and running our implementation of RL^2 thereafter. The Markov version of RL^3 uses a dense neural network, with two hidden layers of 64 nodes each, with ReLU activation function.

For object-level RL, we use model estimation followed by value iteration (with discount factor $\gamma=1$) to obtain $Q$-estimates. The transition probabilities and the mean rewards are estimated using maximum likelihood estimation (MLE). Rewards for unseen actions in a given state are assumed to be 0. States are added to the model incrementally when they are visited, so that value iteration does not compute values for unvisited states. Moreover, value iteration is carried out only for iterations equal to the task horizon (which is 1, 10, 250, 350 for Bandits, MDPs, 11x11 GridWorld, 13x13 GridWorld domains, respectively), unless the maximum Bellman error drops below 0.01.

### B.3 RL 3 -coarse

During model estimation in RL^3 -coarse, concrete states in the underlying MDP are incrementally clustered into abstract states as they are visited. When a new concrete state is encountered, its abstract state ID is set to that of a previously visited state within a ‘clustering radius’, unless that previous state is already part of a full cluster (determined by a maximum ‘cluster size’ parameter). If multiple visited states satisfy the criteria, the ID of the closet one is chosen. If none of the visited states that satisfy the criteria, then the new state is assigned a new abstract state ID, increasing the number of abstract states in the model. It is worth noting that this method of deriving abstractions does not take advantage of any structure in the underlying domain. However, this simplicity makes it general purpose, efficient, and impartial, while still leading to excellent performance. For our GridWorld domain, we chose a cluster size of 2 and a clustering radius such that only non-diagonal adjacent states are clustered (Manhattan radius of 1).

The mechanism for learning the transition function and the reward function in the abstract MDP is the same as before. For estimating $Q$-values for a given concrete state, value iteration is carried out on the abstract MDP and the $Q$-estimates of the corresponding abstract state are returned.

**Table 5: RL^2 /RL^3 Hyperparameters**
| Hyperparameter | Value |
| --- | --- |
| Learning Rate (Actor and Critic) | 0.0003 (Bandits, MDPs) |
|  | 0.0002 (GridWorlds) |
| Adam $\beta 1,\beta 2,\epsilon$ | $0.9,0.999,10^{-7}$ |
| Weight Decay (Critic Only) | $10^{-2}$ |
| Batch size | 32000 |
| Rollout Length | Interaction Budget ($H$) |
| Number of Parallel Envs | Batch Size $\div$ $H$ |
| Minibatch Size | 3200 |
| Entropy Regularization Coeff | 0.001 (Bandits $H=100$) |
|  | 0.04 (GridWorlds) |
|  | 0.01 otherwise |
| PPO Iterations | See training curves |
| Epochs Per Iteration | 8 |
| Max KL Per Iteration | 0.01 |
| PPO Clip $\epsilon$ | 0.2 |
| GAE $\lambda$ | 0.3 |
| Discount Factor $\gamma$ | 0.99 |
| Decoder Layers | 2 |
| Attention Heads | 4 |
| Activation Function | gelu |
| Decoder Size ($d_{model}$) | 64 |

Figure: Figure 4: Average meta-episode return vs PPO iterations for Bandits $H=100$ and $H=500$
Refer to caption: extracted/2306.15909v4/figures/training_curves/Bandits100.png

Figure: Figure 5: Average meta-episode return vs PPO iterations for MDPs $H=100$ (left) and $H=500$ (right)
Refer to caption: extracted/2306.15909v4/figures/training_curves/MDPs100.png

Figure: Figure 6: Average meta-episode return vs PPO iterations for GridWorld 11x11 (left) and 13x13 $(\textit{right})$.
Refer to caption: extracted/2306.15909v4/figures/training_curves/gw11x11.png

Figure: Figure 7: The task-identification power of $Q$-estimates. Left: Fraction of $\delta$-duplicates, with $\delta=0.1$, as a function of time steps in a set of 5,000 random MDPs. Right: Accuracy of a simple multi-class classifier in predicting task ID given $Q$-table estimates, as function of time step. Both figures are generated using the same policy.
Refer to caption: extracted/2306.15909v4/figures/qanalysis/fraction_of_duplicates_without_ribbon.png

## C Training Curves

Figs. 4, 5, and 6 show the training curves for Bandits, MDPs, and GridWorld environments, respectively, across 4 random seeds. The results in the main text may differ slightly since the actor models were evaluated greedily, whereas the training curves reflect the actors’ stochastic policies.

We ran each experiment on an NVIDIA GeForce RTX 2080 Ti GPU, which took approximately 12-24 hours for Bandits and MDPs ($H=100$); and took approximately 3-4 days for other domains.

## D Additional Analysis

In this section, we show that $Q$-estimates, though imperfect, produce reasonable signals for task identification. Here, we test this claim thoroughly with 3 analyses.

### D.1 Requirements for a Unique Q ∗ superscript 𝑄 Q^{*} italic_Q start_POSTSUPERSCRIPT ∗ end_POSTSUPERSCRIPT -Function

Throughout, we assume fixed state space and action space. Below, we show that if the transition function is fixed, then two $Q^{*}$-tables will be identical if and only if both reward functions are also equal. First, we show that identical $Q^{*}$ functions imply identical reward functions. Given the Bellman equations,

$$ $\displaystyle Q^{*}_{1}(s,a)$ $\displaystyle=R_{1}(s,a)+\gamma\sum_{s^{\prime}}T(s,a,s^{\prime})\text{max}_{a ^{\prime}}Q^{*}_{1}(s^{\prime},a^{\prime})$ (20) $\displaystyle Q^{*}_{2}(s,a)$ $\displaystyle=R_{2}(s,a)+\gamma\sum_{s^{\prime}}T(s,a,s^{\prime})\text{max}_{a ^{\prime}}Q^{*}_{2}(s^{\prime},a^{\prime})$ (21) $$

Substituting $Q^{*}_{2}=Q^{*}_{1}$ in Equation (21), we get

$$ $\displaystyle Q^{*}_{1}(s,a)$ $\displaystyle=R_{2}(s,a)+\gamma\sum_{s^{\prime}}T(s,a,s^{\prime})\text{max}_{a ^{\prime}}Q^{*}_{1}(s^{\prime},a^{\prime})$ (22) $$

Subtracting Equation (20) from Equation (22), we get $R_{1}(s,a)=R_{2}(s,a)$. Thus, $(Q^{*}_{1}=Q^{*}_{2})\land(T_{1}=T_{2})\implies(R_{1}=R_{2})$.

Now, if two MDPs have the same reward and transition function, they are the same MDP and will have the same optimal value function. So, $(R_{1}=R_{2})\land(T_{1}=T_{2})\implies(Q^{*}_{1}=Q^{*}_{2})$.

Since encountering similar $Q^{*}$-tables is thus dependent on both transitions and rewards ‘balancing’ each other, the question is then for practitioners: How likely are we to get many MDPs that all appear to have very similar $Q^{*}$-tables?

### D.2 Empirical Test using Max Norm

Given an MDP with 3 states and 2 actions, we want to find the probability that $||Q^{*}_{1}-Q^{*}_{2}||_{\infty}<\delta$, where $Q^{*}_{1}$ and $Q^{*}_{2}$ are 6-entry (3 states $\times$ 2 actions) $Q^{*}$-tables. The transition and reward functions are drawn from distributions parameterized by $\alpha$ and $\beta$, respectively. Transition probabilities are drawn from a Dirichlet distribution, Dir($\alpha$), and rewards are sampled from a normal distribution, $\mathcal{N}(1,\beta)$. In total, we ran 3 combinations of $\alpha$ and $\beta$, each with 50,000 MDPs, a task horizon of 10, and $\delta=0.1$. To get the final probability, we test all $((50,000-1)^{2})/2$ non-duplicate pairs and count the number of max norms less than $\delta$.

Results: For $\alpha=1.0$, $\beta=1.0$, we found the probability of a given pair of MDPs having duplicate $Q^{*}$-table to be $\epsilon=2.6\times 10^{-9}$. For $\alpha=0.1$, $\beta=1.0$, which is a more deterministic setting, we found $\epsilon=4.6\times 10^{-9}$. Further, with $\alpha=0.1$, $\beta=0.5$, where rewards are more closely distributed, we found $\epsilon=1.1\times 10^{-7}$. Overall, we can see that even for a set of very small MDPs, the probability of numerically mistaking one $Q^{*}$-table for another is vanishingly small.

### D.3 Predicting Task Families

The near uniqueness of $Q^{*}$-functions is encouraging, but max norm is not a very sophisticated metric. Here, we test whether a very simple multi-class classifier (1 hidden layer of 64 nodes), can accurately identify individual tasks based on their $Q$-*estimates*. Moreover, we track how the classification accuracy improves as a function of the number of steps taken within the MDP as the estimates improve. In this experiment, the same random policy is executed in each MDP for 50 time steps. As before, our MDPs have 3 states and 2 actions.

We instantiate 10,000 MDPs whose transition and reward functions are drawn from the same distribution as before: transitions from a Dirichlet distribution with $\alpha=0.1$ and rewards sampled from a normal distribution $N(1,0.5)$. Thus, this is a classification problem with 10,000 classes. *A priori*, this exercise seems relatively difficult given the number of tasks and the parameters chosen for the distributions. Fig. 7 shows a compelling result given the simplicity of the model and the relative difficulty of the classification problem. Clearly, $Q$-estimates, even those built from only 20 experiences, provide a high signal-to-noise ratio w.r.t.  task identification. And this is for a *random* policy. In principle, the meta-RL agent could follow a much more deliberate policy that actively disambiguates trajectories such that the $Q$-estimates evolve in a way that leads to faster or more reliable discrimination.

## E Domain Descriptions

### E.1 Bernoulli Multi-Armed Bandits

We use the same setup described in Duan et al. []. At the beginning of each meta-episode, the success probability corresponding to each arm is sampled from a uniform distribution $\mathcal{U}(0,1)$. To test OOD generalization, we sample success probabilities from $\mathcal{N}(0.5,0.5)$

### E.2 Random MDPs

We use the same setup described in Duan et al. []. The MDPs have 10 states and 5 actions. For each meta-episode, the mean rewards $R(s,a)$ and transition probabilities $T(s,a,s^{\prime})$ are initialized from a normal distribution ($\mathcal{N}(1,1)$) and a flat Dirichlet distribution ($\alpha=1$), respectively. Moreover, when an action $a$ is performed in state $s$, a reward is sampled from $\mathcal{N}(R(s,a),1)$. To test OOD generalization, the reward function is made deterministic and initialized from $\mathcal{U}(0,2)$

Each episode begins at state $s=1$ and ends after task_horizon = 10 time steps.

### E.3 GridWorlds

A set of navigation tasks in a 2D grid environment. We experiment with 11x11 (121 states) and 13x13 (169 states) grids. The agent always starts in the center of the grid and needs to navigate through obstacles to a single goal location. The goal location is always at a minimum of min_goal_manhat Manhattan distance from the starting tile.
The grid also contains slippery wet tiles, fatally dangerous tiles and warning tiles surrounding the latter. There are num_obstacle_sets set of obstacles, and each obstacle set spans obstacle_set_len tiles, in either horizontal or vertical configuration. There are num_water_sets set of wet regions and each wet region always spans water_set_length, in either a horizontal or vertical configuration. Entering wet tiles yields an immediate reward of -2. There are num_dangers danger tiles and entering them ends the episode and leads to a reward of -100. Warning tiles always occur as a set of 4 tiles non-diagonally surrounding the corresponding danger tiles. Entering warning tiles causes -10 reward. Normal tiles yield a reward of -1 to incentivize the agent to reach the goal quickly. On all tiles, there is a chance of slipping sideways with a probability of 0.2, except for wet tiles, where the probability of slipping sideways is 1.

The parameters for our canonical 11x11 and 13x13 GridWorlds are:
num_obstacle_sets = 11, obstacle_set_len = 3, num_water_sets = 5, water_set_length = 2, num_dangers = 2, and min_goal_manhat = 8. The parameters for the OOD variations are largely the same and the differences are as follows. For Deterministic variation, the slip probability on non-wet tiles is 0. For Dense variation, obstacle_set_len is increased to 4. For Watery variation, num_water_sets is increased to 8. For Dangerous variation, num_dangers is increased to 4. For Corner variation, min_goal_manhat is set to 12, so that the goal is placed on one of the corners of the grid.

There is no fixed task horizon for this domain. An episode ends when the agent reaches the goal or encounters a danger tile. In principle, an episode can last through the entire meta-episode if a terminal state is not reached.

When a new grid is initialized at the beginning of each meta-episode, we ensure that the optimal, non-discounted return within a fixed horizon of 100 steps is between 50 and 100. This is to ensure that the grid both has a solution and the solution is not trivial.

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Could cumulative reward plots over timesteps for each method be added to visualize meta-training progress?
2. What rollout length was used for RL^{2} in the experiments?
soundness: 3
strengths: - By adding Q-estimates, RL^{3} demonstrates improved efficiency over RL^{2} by requiring fewer PPO iterations and delivering better OOD performance on the selected benchmarks.
summary: This paper proposes RL^{3}, a novel method designed to enhance out-of-distribution (OOD) generalization and reduce meta-training time in meta-reinforcement learning (meta-RL) by integrating Q-estimates as additional input information for the meta-RL agent. The approach is implemented within the RL^{2} framework using transformers and is evaluated against RL^{2} across several benchmark tasks, including bandits, random MDPs, and gridworld.
weaknesses: - The experimental benchmarks focus on tasks with limited state spaces, which may restrict the applicability of the findings.
- Additional experiments on benchmarks with continuous state spaces, such as parameterized MuJoCo environments, are necessary to fully evaluate RL^{3}’s practical benefits over RL^{2} in terms of sample efficiency. The results based on tabular Q-function estimation in small state spaces do not directly indicate performance gains in more complex, real-world settings.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: 1. Several places in the paper seem to suggest that one of the benefits of RL^3 will be asymptotic guarantees of convergence.  Does RL^3 in fact, gain asymptotic performance guarantees?  

2. In Section 4.2 it states, "In practice, we provide action advantages along with the max Q-value (value function) and optionally, standard errors, instead of just Q-estimates".  How critical is this to RL^3's success?  How much does this affect the results?

3. Sec 5 talks about using PPO as the metalearner, but earlier the paper talks about using a blackbox/NN for the metalearner.  I'm a bit confused here---is PPO used just for augmented RL^2?

4. Is Fig 3b comparing only meta-training time (as the last sentence of the MDP results subsection seems to suggest)? Or does it include all required computation?  How are the q-values and counts obtained? How much does this cost? How would this scale for larger/real-world problems?
soundness: 2
strengths: 1. The paper is *very* well-presented.  It was easy to understand and enjoyable to read.
2. The work is well-situated in the literature.
3. The paper provides both empirical and theoretical analyses that both motivate and justify the approach and design decisions.
4. Reported results are good and convince the reader that the proposed approach does provide (some of) the improvements hypothesized.  In particular, rewards are comparable or better than RL^2 while (meta)training time is decreased and OOD generalization seems also to be improved to some extent.
summary: The paper presents a modification of meta-RL that incorporates additional information in the form of Q-estimates and state counts (and possibly other things as well) in order to improve training times, data requirements and possibly asymptomatic performance.  The approach is motivated with both empirical and theoretical arguments and validated on three artificial test domains.
weaknesses: 1. Experimental results are obtained on only toy problems.  It seems likely that the approach may not scale to more difficult/larger/real-world problems.  How does this work on a modestly more difficult domain like Atari, for example?
2. Some claims are not really supported.  For example, while the following statement from Sec. 5 seems like it could be true, there is no supporting evidence given: "VAMDPs can be plugged into any base meta-RL algorithm with a reasonable expectation of improving it."  As another example, it is also suggested that RL^3 might enjoy convergence guarantees, but this is also not supported in any way.
3. Performance improvement results seem like they don't include the overhead for computing q-values and state counts and are therefore potentially somewhat misleading.  This is addressed to some degree in the Sec 6, "Computation Overhead Considerations" subsection, but the discussion is too brief to be convincing; in particular, it is not clear that the claim will hold if scaled to larger/real-world problems.
4. The paper doesn't compare RL^3 with any other meta-RL approaches, other than RL^2, the one it modifies.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: See Weakness.
soundness: 3
strengths: - Novel approach: The paper introduces an innovative method that combines the strengths of traditional RL and meta-RL, potentially addressing some limitations of existing meta-RL approaches.
- Theoretical foundation: The authors provide theoretical insights into why incorporating Q-value estimates can be beneficial, linking them to the optimal meta-value function.
- Improved performance: RL3 demonstrates better long-term returns and out-of-distribution generalization than RL2, while maintaining short-term efficiency.
- Flexibility: The approach can be applied to various meta-RL algorithms, not just RL2.
summary: This paper presents RL3, a novel approach to meta reinforcement learning that combines traditional RL techniques with meta-RL methods. The authors propose incorporating Q-value estimates from traditional RL as additional inputs to a meta-RL architecture, aiming to improve long-term performance and out-of-distribution generalization while maintaining short-term efficiency.
weaknesses: - Limited baselines: The paper would benefit from comparing RL3 to more state-of-the-art meta-RL approaches beyond just RL2. In particular, comparing RL3 to hypernetwork-based approaches would provide valuable insights into its relative performance.
- Scope of experiments: The experiments are limited to discrete domains. As suggested, it would be valuable to see how RL3 performs in high-dimensional state spaces and with continuous action spaces. This limitation restricts the applicability of the current results to more complex, real-world scenarios.
- Compute budget dependency: The performance of RL3 relative to RL2 seems to depend on the compute budget (H), which appears to be environment-dependent. The paper would benefit from providing guidelines or heuristics for selecting an appropriate range of H values for different environments.
- Generalizability: While the paper demonstrates improved out-of-distribution generalization, it's unclear how well this generalizes across a wider range of task distributions or more complex domains.
- Computational overhead: The paper doesn't thoroughly discuss the potential computational overhead of computing Q-value estimates alongside the meta-RL process, which could be a significant factor in practical applications with high dimensional state spaces.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1) What is a VAMDP? Is there a clear and formal definition?

2) Could the authors outline the fundamental differences in the implementation of RL2 and RL3? Are there any precise mathematical formulations that highlight these differences? As presented, Equation 1 only expresses the meta-learning objective but does not distinguish RL2 from RL3.

3) Why is RL2 unable to address OOD generalization problems? Additionally, why does RL2, compared to RL, lack long-term performance? If RL methods have OOD generalization capabilities, why does RL2 not inherit them?

4) What is the difference between the Q-value in RL, RL2, and the Q-estimates used in RL3?

5) In Figure 3b, why are the performance results of RL2 and RL3 equal when the budget is set to 128? Which specific setting in the experiments is responsible for this outcome?
soundness: 2
strengths: 1) RL3 may reduces meta-training time, which has potential practical implications for scaling meta-RL systems.

2) The authors provide theoretical foundation, including proofs that support the efficacy of incorporating Q-value estimates into meta-RL, explaining the potential advantages in handling OOD tasks and long-term planning.
summary: The paper proposes a novel hybrid approach, RL3, that integrates value-augmented MDP into meta reinforcement learning (meta-RL) frameworks RL2. The authors aim to address limitations in current meta-RL models, such as poor asymptotic performance and weak out-of-distribution (OOD) generalization. RL3 introduces object-level Q-value estimates computed using traditional RL, which are integrated into the meta-RL model to improve cumulative reward, speed up meta-training, and enhance OOD generalization.
weaknesses: 1) The author’s understanding of RL2 appears to remain questionable. Specifically, the claim made in the discussion of RL3 suggests that 
"For object-level RL, we use model estimation followed by value iteration (with discount factor
γ = 1) to obtain Q-estimates."
However, in reality, RL2 also employs Q-value estimation in a similar manner to predict the future expectation of reward using the discount factor γ. Specifically, both RL2 and RL3 optimize the same equation mentioned as Equation (1) in the paper. They both employ the idea of optimizing across multiple episodes simultaneously, and there does not appear to be a fundamental difference between the two approaches.

This similarity raises concerns about the novelty of the author's argument regarding RL3's distinct capabilities.

2) Authors mentioned "To test OOD generalization, we vary parameters including the stochasticity of actions, density of obstacles, and the number of dangerous tiles." Although this level of variation is commonly considered a new task in meta-learning settings, the actual changes introduced are relatively minor. Furthermore, GridWorld is a relatively simple environment. 

Whether you have plans to test on more complex domains in future work, such as vision-based RL tasks or robotic control tasks (for example DMC or Atari)? 

3) The authors did not provide a clear definition of the RL3 algorithm in a formalized manner, nor did they present a direct comparison with the formulation of the RL2 algorithm. This lack of precise algorithmic definitions makes it difficult for readers to fully understand how RL3 operates and differs from RL2.

Could authors  include a side-by-side comparison with RL2 to highlight the key differences.

4) I not sure that simply feeding the entire sequence of states and actions from a whole task continuously into an LSTM or transformer, as done in RL3, will necessarily enhance OOD generalization and long-term performance. A more promising approach would involve training on trajectories of both short and long lengths. Whether authors have considered alternative methods combining varying trajectory lengths.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: See Weakness above.
soundness: 2
strengths: 1. The paper constructively combines the previous meta-RL algorithms called RL$^2$ and general RL signal to absorb the advantage and avoid the disadvantage of both sides.
2. The paper adopts extra action value as the auxiliary input and provides some theoretical evidence to show rationality.
3. Compared to RL$^2$, which uses the history trajectory rather than the RL$^3$ using the action value, the latter can learn more efficiently and easily. I appreciate the authors' thorough explanation.
summary: Discerning that RL$^2$ shows poor asymptotic and ood performance while general RL demonstrates short-term efficiency, this paper constructively combines the general RL with the previous meta-RL algorithm RL$^2$ and calls it RL$^3$. To be specific, the paper seeks to add an extra signal as input, namely the action-value function. Though seems to be straightforward, adding this input would give a useful compression compared to trajectory history. To validate the effectiveness, the paper provides the theoretical proof for recovering the optimal meta-level value function. Finally, some experiments on discrete domains are performed.
weaknesses: 1. The biggest concern for me is that as shown in Appendix A.4, $\textbf{based on the extra input action-value}$, the iteration process would lead to two cases either the belief distribution collapses rapidly to zero for tasks $j\ne i$ or the belief distribution will reduce to zero for tasks not compatible with the observed Q-values. $\textbf{Does this mean that whatever is capable of distinguishing tasks as input should meet this proof and the action value is only a feasible case?}$ If so, what is the advantage of this over adopting the task representation as input like the context-based meta-RL algorithms [1] ? Note that these algorithms can also hold a good compression w.r.t the transition / trajectory. Also, these algorithms can resolve the continuous domain rather than be limited to the discrete domain.
2. In my view, why RL$^3$ holds better OOD performance can not be interpreted by the theoretical justification.
3. RL$^2$ is a very early meta-RL algorithm, incrementally building upon RL$^2$ may still have a large gap with SOTA algorithms. As some context-based meta-RL algorithms hold the SOTA, I believe the authors can move RL$^3$ onto the context-based meta-RL algorithms to show further superiority or at least add some comparison.
4. This is related to problem 3, some lately-published papers regarding advanced meta RL should be discussed.
5. Some unclear definitions and typos: what are object-level and meta-level? What is Figure b) trying to show? Why does any permutation of transitions yield the same Q-estimates? For testing OOD generalization, MDPs are generated from distributions with different parameters than in training, what does the "parameters" mean here? How do the authors generate the OOD data? VeriBAD -> VariBAD
6. Unclear pseudo-code: Does the $ONEHOT(s)\cdot Q(s) \cdot N(s)$ denote the results after concatenation?

[1] Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning. Lanqing Li, et.al.

$\textbf{I would be happy to raise the score if the above problems are properly addressed.}$

Meta-review: The authors propose a meta-RL framework based on prior work RL^2 that incorporates Q-estimates for better sample efficiency, asymptotic performance, and OOD generalization. They instantiate this algorithm as a modification of RL^2 with transformers and compare against RL^2.

Reviewers found the paper very well written and clear, properly situated within the related literature. There were also both empirical and theoretical analyses to motivate and justify the approach and design. The empirical results are significant compared to the incorporated baselines. However, reviewers also found that those empirical results were only evaluated in toy environments and some competitive baselines were missing. Further, some claims were not supported, including generalization claims given the limitations of the environments used. 

While the idea, motivation, and algorithm are interesting, I agree with reviewers that more compelling justification is needed by incorporating results from more complex environments and compared with more recent meta-RL baselines. For these reasons, I vote to reject this paper in the current form, and encourage the authors to revise and re-submit.

Decision: Reject