Actions Speak Louder Than States: Going Beyond Bayesian Inference in In-Context Reinforcement Learning

Weaknesses:

W1: They are simply representations, and they enable us to compare the limitations of Bayesian inference on finite number of tasks to the Bayesian inference on a coverage task distribution in terms of regret. These are not representations that we can attain exactly in practice, however, as it is said in Theorem 3.1, using log likelihood for the loss function in equation 1, will result in a Bayesian inference decision making.

W2: We thank the reviewer for noticing this possible disconnect. We implied that low task diversity training results in Bayesian inference on finite number of tasks $M_\theta^{\text{F-PS}}$, whereas, high task diversity shifts from this pattern to $M_\theta^{\text{E-PS}}$. We have updated Section 4 to include these points, the changes are in blue text.

W3: We understand, however, we still believe the explanation in Section 4.1, is better than any figure we could include. We feel perhaps a short video would help better in understanding, that we can include to the supplementary materials for the camera-ready version.

W4: We are not focusing on a new method in this paper, the reason for MAML inclusion is to highlight the "challenge" of the problem setting, not as a comparison. The papers you suggest as baselines, again, do not make sense for the same reason.

Questions:

Q1: There is an explicit assumption in Section 3 that the task identities are latent. We believe it would be difficult to find an ICL setting where the latent task identities are not available, perhaps utilizing a transformer with memory on a single RL task. There are numerous examples we can give. Before that, we would like to point out that the definition of task in RL setting can be very nuanced in real life. For instance, a self-driving car that drives on a wet vs dry road are two different tasks, since different actions will have different rewards, even if the final step observation is the same.

Q2: Out-of-distribution tasks cannot be considered, because in that section we are giving theoretical results when the training dataset has the exact distribution for the task distribution.

Q3: They are the same. $M^{\text{E-PS}}_\theta$ does not have access to $\boldsymbol{\mathbf{\mathcal{T}}}$. It is possible to have priors for tasks that a model has never seen before. The simplest example is image classification task, where an image has never been included in the training dataset. However, the model still able to give us a distribution over the classes for that same image. The same is also happening here, although the model has never some specific task in the training distribution, it is still possible for that model to have priors for that task.

Q4: We really thank you for asking these details when they don't seem right, instead of simply discarding our paper. $\hat{g}(\cdot | \cdot)$ is the posterior over the task definition after observing $H_{t-1}$, as given on line 212. The prior as we denoted as $\hat{P}_{\text{pre}}$, emerges throughout the training, it is not something we specifically train for. In the next parts, we show that as long as what we say in lines 207-208 is true (simply prior being non-zero for all tasks that have non-zero task distribution), our Theorem 3.2 holds. We definitely understand the intuition that the quality of match between the estimated prior and the true prior being important (and it is indeed important), however, in Theorem 3.2, we show that as we increase the context episodes, we can get arbitrarily good regret. For the same $n$ number of episodes, the quality of priors will of course matter.

Q5: They only relate to finetuning of the ResNet Model, not the full stack.

We thank the reviewer for taking the time to review our paper, and as it can be understood from the detailed questions, they have spend a good amount of time on our paper which makes our job much easier. Let us address some key points before responding to specific concerns.

Our paper's primary objective is to investigate the emergence of ICRL in relation to task diversity, rather than proposing a new method or creating the most challenging meta-RL environment. While there may be existing meta-RL environments with high task diversity only promising other alternative we are aware of is [1], we focused on creating an environment that allows for reproducible experiments, feasible offline data generation within academic constraints, and accessibility with limited computational resources.

To our knowledge, this critical topic remains largely unexplored in the literature. Our contribution starts by identifying limitations in current RL environments, introducing a purposefully simple yet effective environment for studying ICRL with task diversity, providing theoretical analysis of observed behaviors, and demonstrating the shift from Bayesian Inference to a more general learning method enabling true ICRL. As far as we are aware, this work is the only work to demonstrate the latter.

Regarding our environment's complexity, we used MAML to demonstrate that simple episode-based finetuning cannot achieve optimal performance, confirming the presence of genuine ICRL. While other methods could have been used for this demonstration, this is not the paper's central focus.

[1] Nikulin, Alexander, et al. "XLand-minigrid: Scalable meta-reinforcement learning environments in JAX.”

Q1: The goals are replicating how to "write" a specific character. We cannot write a character without knowing what to write.

Q2: Please refer to our clarifications. Briefly, we are not interested in ICRL beating MAML or any other Meta-RL method (btw AMAGO is a scalable way to implement ICRL, they don't do weight updates at test time), we simply want to provide proof that the task is not easy, and the results we get in our main results, are not just some side-effect of simple regulation/capacity increase.

Q3: We could certainly do this, but we really don't think it would provide any important information. We provide many runs in different configurations. Does the reviewer see any instability in any one of them?

Q4: See already existing Table 1.

Q5: We agree, observing more episodes would be interesting to observe. However, we are already observing very good returns after observing a single episode. Therefore, there is no reason to test with more than a single episode as context. Again, this is interesting, and we would also wish to see future works investigating this in possibly other environments.

Q6: We believe it would be Figure 6, we updated it. Let us know if you have problems with other figures.

W1: For the claim about task preconditioning, specifying a task does not eliminate execution uncertainty. Consider writing the letter 'A': even with this clear goal, significant uncertainties remain about the required force, optimal pen angle, and surface characteristics. If task specification truly eliminated all uncertainty, we would see high rewards/low losses from the first episode in Figure 3, which is not the case.

W2: Total trajectory return may not be the most appropriate metric for Meta-RL environments, since episode lengths and reward scales can naturally vary between tasks. An agent could achieve identical performance quality across tasks, yet receive different cumulative rewards simply due to varying episode lengths. In our environment, where episodes can be either short or long, using average reward per step provides a more consistent and meaningful measure of performance that normalizes these duration differences.

W3: Success rate is primarily used as a metric in robotics-related RL tasks, not as a standard measurement across RL algorithms. The assumption that Meta-RL benchmarks should be limited to environments that can report success rates is unfounded. Consider OpenAI Gym, one of the most widely used RL environment libraries - none of its environments use success rate as a benchmark metric. These environments are frequently adapted for Meta-RL research by modifying their parameters. This established practice demonstrates that success rate reporting is not a prerequisite for Meta-RL benchmarks.

W4: Let's clarify the distinction between overfitting and our observed behavior. Overfitting occurs when a model performs well on training data but poorly on new, similar inputs. This is not what we observe in our results - Figure 3's training losses show consistent performance across new instances of previously seen task types, despite never reusing the exact same instance during training. What we actually observe and theoretically demonstrate in Section 3 is a different phenomenon: strong performance on in-distribution tasks (new instances of familiar task types) but weaker performance on out-of-distribution tasks. This behavior aligns with our theoretical analysis of Bayesian inference under finite task distribution during training. Importantly, once task diversity exceeds a critical threshold, the transformer transitions from operating in a finite Bayesian inference regime to behaving more like $M^{\text{E-PS}}_\theta$.

W5: We should clarify our statement regarding "unprecedented task diversity." While you mention a dataset with high task diversity, we note that it was introduced just three months before submission (after May 28th). Following reviewer guidelines, it would be more appropriate to consider this as concurrent work. Moreover, there's a practical consideration regarding dataset accessibility. The size of such datasets often presents significant computational barriers for many researchers. We are curious how much of a computation power would be needed to repeat all our experiments on this dataset. This speaks directly to our goal of creating benchmarks that are accessible within typical academic resource constraints.

W6: We believe there has been an issue with naming. We meant Meta-RL methods that require weight updates to the model when we meant "Meta-RL", which is evidenced by our previous sentence in line 84, and we believe the reviewer also understood what we meant. However, we believe the comparison between Ad and VariBAD is still not a direct example of Meta-RL vs ICRL. Both AD and VariBAD are Meta-RL methods that do not require weight updates (finetuning) at test stage (So, they are both ICRL).

We appreciate your thorough review of our paper. This makes our job much easier. Let us address some key points before responding to specific concerns.

Our paper's primary objective is to investigate the emergence of ICRL in relation to task diversity, rather than proposing a new method or creating the most challenging meta-RL environment. While there may be existing meta-RL environments with high task diversity only promising other alternative we are aware of is [1], we focused on creating an environment that allows for reproducible experiments, feasible offline data generation within academic constraints, and accessibility with limited computational resources.

To our knowledge, this critical topic remains largely unexplored in the literature. Our contribution starts by identifying limitations in current RL environments, introducing a purposefully simple yet effective environment for studying ICRL with task diversity, providing theoretical analysis of observed behaviors, and demonstrating the shift from Bayesian Inference to a more general learning method enabling true ICRL. As far as we are aware, this work is the only work to demonstrate the latter.

Regarding our environment's complexity, we used MAML to demonstrate that simple episode-based finetuning cannot achieve optimal performance, confirming the presence of genuine ICRL. While other methods could have been used for this demonstration, this is not the paper's central focus.

[1] Nikulin, Alexander, et al. "XLand-minigrid: Scalable meta-reinforcement learning environments in JAX.”

[1] Beck, J., Vuorio, R., Liu, E. Z., Xiong, Z., Zintgraf, L., Finn, C., & Whiteson, S. (2023). A survey of meta-reinforcement learning. arXiv preprint arXiv:2301.08028.

[2] Agarwal, R., Singh, A., Zhang, L. M., Bohnet, B., Rosias, L., Chan, S., ... & Larochelle, H. (2024). Many-shot in-context learning. arXiv preprint arXiv:2404.11018.

[3] Laskin, Michael, et al. "In-context reinforcement learning with algorithm distillation.”

[4] Lee, Jonathan, et al. "Supervised pretraining can learn in-context reinforcement learning.”

[5] Nikulin, Alexander, et al. "XLand-minigrid: Scalable meta-reinforcement learning environments in JAX.”

[6] Zintgraf, L., Shiarlis, K., Igl, M., Schulze, S., Gal, Y., Hofmann, K., & Whiteson, S. (2019). Varibad: A very good method for bayes-adaptive deep rl via meta-learning. arXiv preprint arXiv:1910.08348.

[7] Mehta, B., Deleu, T., Raparthy, S. C., Pal, C. J., & Paull, L. (2020). Curriculum in gradient-based meta-reinforcement learning. arXiv preprint arXiv:2002.07956.

[8] Rakelly, K., Zhou, A., Finn, C., Levine, S., & Quillen, D. (2019, May). Efficient off-policy meta-reinforcement learning via probabilistic context variables. In International conference on machine learning (pp. 5331-5340). PMLR.

[9] I suggest the authors do a scholar search with a query "in-context reinforcement learning", limiting the results up until 2021 as a proof to my point. AD was released in 2022.

[10] Choshen, E., & Tamar, A. (2023, July). Contrabar: Contrastive bayes-adaptive deep rl. In International Conference on Machine Learning (pp. 6005-6027). PMLR.

[11] Park, S., Frans, K., Eysenbach, B., & Levine, S. (2024). OGBench: Benchmarking Offline Goal-Conditioned RL. arXiv preprint arXiv:2410.20092.

While re-reading the paper, in line 81:

In ICRL, the demonstrations are expert trajectories and the hope is that the pretrained mode [...]

This is not always the case. For example, in AD [3] (the method that pioneered the field of ICRL) authors use learning histories of single-tasked RL algorithms to extract a policy improvement operator from them. Thus, I would refrain from using such strong propositions.

Q1: The goals are replicating how to "write" a specific character. We cannot write a character without knowing what to write.

Given there is a signal from the environment, it is possible for the model to understand which character to write. In the aforementioned ICRL and Meta-RL papers the goal location is not known for the agent, but with trial and error (which includes getting the signal from the environment) the model successfully navigates to the right place or chooses the direction for the agent to go.

Q2: [...] we are not interested in ICRL beating MAML [...] we simply want to provide proof that the task is not easy, and the results we get in our main results, are not just some side-effect of simple regulation/capacity increase.

As I understand, for the main experiments authors use a transformer with 8 layers, 8 heads with 512 hidden dimensions (Table 1), the sequence consists of 2 episodes. For MAML baseline, authors take 4 layers MLP with 1024 hidden dimensions (Appendix B2) and the sequence consists of the last 5 strokes. Roughly, the transformer that authors use should have approximately 30M parameters, while MAML MLP backbone consists of 2.5M parameters at most, not mentioning the transformer's greater sequence length. To my mind, such a comparison is not conclusive or fair.

Q4 & Q6

Thank you, my concerns are addressed.

Overall, my main concern is that the proposed benchmark is suitable neither for Meta-RL nor for ICRL, since the goal for the model is given as an input and no other uncertainty is left for the model to estimate.

As an example, consider Dark Room environment (one of the most common in ICRL now), but instead of keeping the goal hidden from the model, one trains a simple MLP that takes as an input the position of an agent + a goal grid it has to navigate to. Given enough pre-training examples, this MLP could navigate to any grid (even those that the model did not see during training), but we should not call such a behavior "Meta" (rather, it is called goal-conditioned RL, which is a separate field). The example of such a pre-training method may be found in [5, 11], where authors train a goal-reaching policy to later use it for collecting a dataset of diverse behaviours.

Besides, authors the authors reject to address some of my concerns, i.e.

We could certainly do this, but we really don't think it would provide any important information.

Therefore, there is no reason to test with more than a single episode as context.

Considering this, and the fact that my main concern still holds, I am inclined to stand by my score.

I thank the authors for their thoughtful rebuttal.

W1: For the claim about task preconditioning, specifying a task does not eliminate execution uncertainty. [...]

Consider the definition of Meta-RL from [1]:

Meta-RL is most commonly studied in a problem setting where, given a distribution of tasks, the goal is to learn a policy that is capable of adapting to any new task from the task distribution with as little data as possible"

Following this definition, the setting where the model is conditioned on the goal cannot be viewed from a Meta-RL perspective. I understand that apart from uncertainty in the final goal, there may be some other sources of it (i.e. different physics of the simulator), but I do not spot it in the author's work. Besides, as I mentioned in my review, next-token prediction losses may not be a good evidence of in-context abilities [2], therefore I do not see the results in Figure 3 conclusive enough.

W2: Total trajectory return may not be the most appropriate metric for Meta-RL environments, since episode lengths and reward scales can naturally vary between tasks. [...]

My point with reference to [2] still holds. In the light of this work, I believe the authors should come up with a more conclusive metric that is not known ambiguous in predicting in-context performance.

Besides, the total return over a trajectory (averaged over tasks) is considered standard in Meta-RL literature [1, 3, 4, 6, 8, 10], so I do not see why this benchmark should deviate from it.

W3: Success rate is primarily used as a metric in robotics-related RL tasks, not as a standard measurement across RL algorithms. [...]

The success rate metrics or similar to it (counting how many timesteps the agent stands in the goal position after reaching it) are common in current ICRL research [3, 4, 5]. As such, I am not proposing to exclusively use these metrics, but I want to again highlight that MSE is not a good metric for a benchmark, especially for comparing images. I would be happy to know what authors may propose (maybe using a simple cnn that classifies images) for solving this problem.

W4: Let's clarify the distinction between overfitting and our observed behavior. [...]

I see the point authors make, I am grateful for the clarification.

W5: We should clarify our statement regarding "unprecedented task diversity." While you mention a dataset with high task diversity, we note that it was introduced just three months before submission (after May 28th).

To be precise here: the authors rightfully notice that datasets which were generated from [5] were indeed released after May 28, but the benchmark itself was published on arxiv 19 Dec 2023. In this light, the claim still sounds too strong, I would advise to ease it a little bit.

W6 [...] we believe the reviewer also understood what we meant.

Unfortunately, I did not. Please, make the text unambiguous, it is vital to be consistent with generally accepted terminology in the field.

W6 [...] Both AD and VariBAD are Meta-RL methods that do not require weight updates (finetuning) at test stage (So, they are both ICRL).

However I see the difference between Meta-RL and ICRL minimal, it is important to stick to the terminology. I have never yet seen that VariBAD [6] was called an ICRL method. I understand what the authors try to communicate here (gradient-based Meta-RL [7] vs. context representation Meta-RL [8]). In-context RL (ICRL) as a term was introduced in [3] first [9], where it was proposed to implicitly let the in-context abilities of a transformer to estimate the uncertainty from the history of interactions, rather than explicitly learn the representation of it. Thus, I would refrain from calling VariBAD one of the ICRL method.

Thank you for taking the time to review our paper. We are really encouraged by your high view of our paper.

Responses to Questions:

Q1: We agree, observing more episodes and different context lengths would be interesting to explore, especially given observations of positive results from previous studies such as AD. We would also wish to see future works investigating this in possibly other environments.

Q2: We supply the previous observations, rewards and the actions in the trajectory.