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

• W1: You are right that adopting an explicit task-representation like in [1] may provide better task discrimination. In fact, Q-values may not even be equally discriminating (though extremely discriminating in practice). But the point we are making is that it is not necessary to discriminate between tasks that share the same Q-values since the optimal action is the same in each case. As Q-estimates converge, they represent just the "right amount" of belief collapse (even if not full collapse) while directly indicating the optimal action. At that point, belief can be entirely ignored for decision making since the optimal action is already indicated by Q-values, which is a simpler mapping to learn. This mechanism works even if there is a distribution shift at deploy time -- a scenario where Q-values might suggest an incorrect belief, but still suggest the correction action. By contrast, task-representation based approaches rely heavily on belief for decision making, which is not just a more complex mapping, but it also creates a single point of failure if the belief is incorrect (can happen due to distribution shift or even due to error accumulation during belief updates).
• W2: Answered during response to W1
• W3: While we agree with you, we emphasize that our custom implementation of RL^2, even without transformers, is known to be comparable to other state-of-the-art meta-RL baselines [2] (like VariBAD, a context-based approach), as mentioned in section 5 introduction.
• W4: Point noted. Suggestions are welcome.
• W5: Object-level and meta-level are terms defined in foundational work on principles of metareasoning for bounded optimality (Russell & Wefald, 1989). In our case, object-level process refers to the task-specific RL subroutine and meta-level process to the meta-RL policy. Regarding any permutation of transitions leading to the same Q-estimates, it is because in a Markovin system, the same transition model $T(s,a,s')$ and reward model $R(s,a,s')$ would be inferred regardless of the order in which the data tuples $<s,a,r,s'>$ arrive. And same $T, R$ => same MDP => same Q-values. For testing OOD generalization, the "parameters" (for procedural generation of tasks, e.g. density of obstacles or stochasticity of actions in GridWorlds) for each domain are mentioned in section 5 (and in more detail in Appendix E).
• W6: Yes, $\cdot$ (cdot) denotes concatenation here. We have mentioned that explicitly now in the updated manuscript. Thank you for pointing it out.

[2] Recurrent model-free RL can be a strong baseline for many POMDPs. Ni et al.

Dear authors,

I appreciate the authors' clarification. The following are my responses:

1. Using Q-values to guide policy training is the core claim of this paper, but the authors' response highlights the advantages of Q-value guidance on just intuition (while I acknowledge identifying the optimal action value may be a potential advantage against task representation, this needs to be further validated through theory or experiments and it seems that the task representation can also fit in your proposed theoretical proof) rather than providing further theoretical or experimental evidence. Additionally, the inability of Q-value guidance to handle continuous settings is, in my view, a more significant drawback. I also disagree with the claim that Q-values can indicate the correct action during deployment while task representations cannot. Q-values can also fail to provide accurate estimates under distribution shifts (even their relative magnitudes may become incorrect). I suggest that the authors further strengthen the comparison between Q-values and representations in their submission.

2. The authors' literature review appears to focus only on works up to 2022, while many current SOTA algorithms have already surpassed VariBAD, such as [1,2].

[1] Meta-Reinforcement Learning via Exploratory Task Clustering [2] ContraBAR: Contrastive Bayes-Adaptive Deep RL

Dear authors,

I thank the authors for my response. Here, I'd like to restate my core concern:

As said in my initial response and also seems to be agreed by the authors, the Q-value and the task representation can both meet the proof of Equation 4 in the Appendix, as the task representation would also collapse to zero for tasks j≠i, or face the condition to be identical to multiple MDPs.

From this perspective, Q-value is just a special case for instantiating the proposed “proof-of-concept”, $**with the reason to choose the Q-value missing**$ (this needs experimental or theoretical evidence). As stated by the authors, adopting the Q-value would lead to discrete scenarios. Obviously, this is a disadvantage, thereby emphasizing the comparison between the Q-value and task representation to demonstrate the advantage of the Q-value. Despite the author's response that it deserves further consideration, it is, in my view, a significant part of the submission. Hence, I strongly recommend the authors compare the Q-value and the task representation to highlight their claimed edge.

Additionally, I would like to restate that I remain positive about the potential of this paper. However, since the unique advantage of Q-value against task representation is missing, this indeed causes the paper logically incomplete, which makes me feel below the accepted bar. Nevertheless, I still thank the authors for resolving my other concerns, I decide to raise my score from 3->5.

I hope that the authors will consider my comments and make the paper more beneficial!

"I also disagree with the claim that Q-values can indicate the correct action during deployment while task representations cannot."

We appreciate the suggestion that other task representations could be beneficial, but we also note that Q-values have been extensively studied and the paper provides an argument in favor of their use (e.g., Q-values provide a simpler mapping to correct actions upon convergence), while alternative representations are worth further consideration. We should also note that Q-values can always be used in conjunction with task representations.

"Q-values can also fail to provide accurate estimates under distribution shifts (even their relative magnitudes may become incorrect)."

Q-values are always derived per task, without any assumptions about the task. Just like in standard RL, upon convergence, they are always correct for that task and $\text{argmax}_a Q(s,a)$ is always the best action, making distribution shift irrelevant. That said, we agree that it is possible that if the meta-RL agent does insufficient exploration initially due to distribution shift, the Q-value estimates could end up being relatively less useful. However, insufficient exploration is not a problem unique to RL$^3$.

"Q-value guidance on just intuition rather than providing further theoretical or experimental evidence."

Results shown in Fig 3a and 3c provide experimental evidence that RL$^3$ performs better under distribution shift for the respective domains. The videos provided in the supplementary (most prominently "rl3-4.mp4", of which we provide snapshots in figure 4) also suggest that the agent is largely following the value function once the goal is found.

"inability of Q-value guidance to handle continuous settings is, in my view, a more significant drawback"

RL$^3$ is indeed (currently) a discrete-action space meta-RL algorithm. However, several advanced approaches for extending it to continuous actions are underway and we believe they would be better presented in a separate paper. At a high level, the most basic approach involves using a simplified variation of DDPG for task-level RL and providing the Q-values of the best action and some other actions in its vicinity.

Regarding citations, thank you for the pointers. We will include the following papers in our discussion of related work: 1) Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning; 2) Meta-Reinforcement Learning via Exploratory Task Clustering; 3) ContraBAR: Contrastive Bayes-Adaptive Deep RL.

• W1: It is true that both RL$^2$ and RL$^3$ optimize the same meta-RL objective (as do most meta-RL algorithms), and that this objective considers optimizing behavior across multiple object-level episodes. We believe the confusion here lies in the fact that while many meta-RL algorithms, including RL$^2$, are described as 'performing object-level RL' within an 'inner loop', they are doing so implicitly. It is not possible to directly extract an actual object-level Q-estimate from RL$^2$. RL$^3$ on the other hand uses explicit Q-estimation (in this case, via value iteration) to produce additional inputs for the meta-agent. The VAMDP introduced later in the paper is done so in order to formalize this process.

• W2: We agree that tests on more complex domains would be more convincing, although our specific hypothesis regarding the utility of object-level Q-estimates does not include SOTA performance on continuous domains as a condition. However, we do plan to extend RL$^3$ to larger/continuous state spaces, anticipating similar performance gains as the arguments made in favor of RL$^3$ would still hold.

• W3: RL$^3$ differs from RL$^2$ only in that it takes an additional input: object-level Q-estimates. There are many ways one might compute these Q-estimates. We provide an example, which is to use value iteration (or value iteration on abstracted states), though this is not required.

• W4: We agree that "simply feeding the entire sequence of states and actions from a whole task continuously into an LSTM or transformer" is likely not a good strategy for strong OOD or long-term performance. This is why our proposed method augments this strategy with object-level Q-estimates, which we argue in the paper provide (in practice) several nice properties not attainable from trajectories alone.

• Q1: A VAMDP is a Value-Augmented MDP (outlined in section 4.2 and used within an algorithm in Appendix B). This construction is similar to a regular MDP except that the state space of the VAMDP includes possible Q-estimates. For example, an agent may visit the same state in the MDP multiple times, but as it performs Q-learning, its Q-estimate upon each visit will be different. This will result in a different VAMDP state.

• Q2: Both RL$^2$ and RL$^3$, and indeed most meta-RL methods, solve the same objective function (Equation 1). The fundamental difference between RL$^2$ and RL$^3$ is that RL$^3$ takes an additional input (object-level Q-estimates). Importantly, this additional input can be derived from the same input given to RL$^2$, therefore this does not constitute "extra'' or "privileged'' information.

• Q3: There is no singular reason that RL$^2$ (or any other learning algorithm) has difficulty with OOD generalization. The most common understanding is that learning algorithms are optimized to achieve high performance in-distribution by exploiting patterns specific to that data. Thus, when presented data beyond the training distribution, the new data may not follow the patterns the algorithm has learned to exploit, and the algorithm performs worse.
  RL$^2$ seems to have worse long-term performance because the problem of mapping longer and longer (state, action, reward) trajectories to actions becomes more difficult over time, while RL algorithms performing Q-learning become more accurate over time. RL$^2$ does not, in practice, inherit the nice asymptotic properties of RL because it does not maintain object-level Q-estimates (or some equivalent) explicitly and thus it cannot benefit from this more easily exploitable representation.

• Q4: In the case that you are asking about whether Q-values are different than Q-estimates: We are using these terms interchangeably.
  In the case that you are asking for which value function the different Q-values are referring to: In RL, Q-estimates (say in DQN or in Q-learning) refer to the object-level value function of the MDP at hand, i.e. expected return-to-go within an episode of the MDP. In meta-RL (RL$^2$), the meta-value function (in this paper we use $\bar{V}$) is estimated by the meta-critic (PPO critic) during meta-training and refers to the expected return-to-go in the meta-episode. RL$^3$ computes Q-values of the MDP in a meta-episode using RL and then directs the result into the meta-RL actor and the meta-critic.

• Q5: There is no specific experimental setting that causes RL$^2$ and RL$^3$ to have similar performance when the horizon is low (128). Our hypothesis is that this problem setting is relatively easy and thus RL$^2$ is able to achieve near-optimal performance. RL$^3$ likely also achieves near-optimal performance in this setting and does so in a similar number of iterations. One possible reason for the similarity in performance is that with a relatively low budget, the object-level Q-estimates within RL$^3$ have less time to converge to a useful signal and thus RL$^3$ derives less (or no) benefit from having access to them in this low-budget case.

• W1: While we agree with you, we emphasize that our custom implementation of RL$^2$, even without transformers, is known to be comparable to other state-of-the-art meta-RL baselines [2], as mentioned in section 5 introduction.
• W2: We agree with your assessment. This paper is a proof of concept (as mentioned in the introduction), and our specific hypothesis does not include SOTA performance on continuous domains as a condition. However, we do plan to extend RL$^3$ to larger/continuous state spaces, anticipating similar performance gains as the arguments made in favor of RL$^3$ would still hold.
• W3: Interaction budget (or adaptation period) $H$ is part of the problem specification. It is not a hyperparameter. RL$^3$'s advantage over RL$^2$ increases with $H$.
• W5: Regarding computation overhead for scaling to larger state spaces, several strategies can be employed, like state-abstractions to reduce state space size (which we show leads to substantial computational savings in gridworlds), or using a more-efficient planner than value iteration (e.g., RTDP or MCTS). In fact, it is even possible to do value iteration on a GPU, which would compute Bellman updates for all states at once [3]. For continuous state spaces, we speculate that fitting Q-values to even a small function approximator would be sufficient to benefit RL$^3$.

[2] Recurrent model-free RL can be a strong baseline for many POMDPs. Ni et al.
[3] Markov Decision Process Parallel Value Iteration Algorithm On GPU. Chen et al. 2013.

• W1: We agree with your assessment. This paper is a proof of concept (as mentioned in the introduction), and our specific hypothesis does not include SOTA performance on continuous domains as a condition. However, we do plan to extend RL^3 to larger/continuous state spaces, anticipating similar performance gains as the arguments made in favor of RL^3 would still hold.

• W2: Our expectation (but not a definitive claim) that plugging in VAMDPs to an off-the-shelf base meta-RL algorithm could enhance its performance is based on the observation that the arguments made in favor of RL$^3$ are agnostic to the base meta-RL algorithm. Regarding convergence guarantees, we do not claim any guarantees. However, it is reasonable to expect the task-specific RL component to converge to near-optimal Q-values in the limit (subject to sufficient exploration), and RL$^3$ to be able to learn (subject to deep learning limitations) the trivial mapping from optimal Q-values to greedy actions.

• W3: See response to Q4.

• W4: While we agree with you, we emphasize that our custom implementation of RL$^2$, even without transformers, is known to be comparable to other state-of-the-art meta-RL baselines [2], as mentioned in section 5 introduction.

• Q1: See our response to W2.

• Q2: Thank you for asking this question. While re-doing some experiments to answer your question, we discovered that standard-errors are no longer necessary with the final optimized hyperparameters. In our earlier experiments, before the hyperparameters were fully optimized, we did observe significant benefits from including standard errors. Regarding including Q-values vs. including action advantages, while they carry the same information, action advantages are smaller in scale and therefore friendlier for neural networks. In the updated manuscript, we have removed the reference to standard-errors.

• Q3: PPO is used for meta-training the meta-RL policy. The meta-RL policy is a blackbox recurrent model (in our implementation of RL$^2$, just a transformer mapping experience-sequences to actions). To augment RL$^2$ inputs with Q-estimates, value-iteration (over an estimated model of the environment) is used.

• Q4: The Fig. 3b y-axis compares the number of PPO iterations in meta-training, not wall-time. However, they are highly correlated as the computation time in meta-training is bottlenecked by gradient updates. The exact wall-time overhead of computing q-values (which is done using value iteration over an estimated model of the task-specific MDP) is discussed in the subsection "Computation Overhead Considerations". For scaling to larger state spaces, several strategies can be employed, including state-abstractions to reduce state space size (which we show leads to substantial computational savings in gridworlds), or using a more efficient planner than value iteration (e.g., RTDP or MCTS). In fact, it is even possible to do value iteration on a GPU, which would compute Bellman updates for many states at once [3].

Once again, thank you for the review. It has helped us improve our manuscript.

[2] Recurrent model-free RL can be a strong baseline for many POMDPs. Ni et al.
[3] Markov Decision Process Parallel Value Iteration Algorithm On GPU. Chen et al. 2013.

We agree with your assessment of the weaknesses. This paper is a proof of concept (as mentioned in the introduction), and our specific hypothesis does not include SOTA performance on continuous domains as a condition. However, we do plan to extend RL^3 to larger/continuous state spaces, anticipating similar performance gains as the arguments made in favor of RL^3 would still hold. Regarding your questions:

1. We have added the meta-training plots to the appendix.
2. The context/rollout length for RL$^2$ (and RL$^3$) transformer is $H$ i.e., the full history of experiences in the meta-episode.