Summary
This work leverages transformers to solve decision making problems in a few-shot manner. In particular, they train transformers to predict optimal actions from a task given a query state and a few-shot dataset of state, action, reward transitions. The approach is tested in both multi-armed-bandit problems and POMDPs, showing superior performance than other RL and Meta-RL approaches. It also provides theoretical analysis of the proposed policy, providing regret bounds and guarantees of performance when different algorithms are used to generate the pre-training data.
Strengths
Method:
The work conducts a thorough study of the effectiveness of transformer for in-context learning in decision making problems. While the domains chosen are quite simple, there are some valuable learnings (especially if they transfer to more complex domains), such as:
- The value of sampling actions from the transformer's likelihood in multi-armed bandits
- The effect of using expert datasets for in-context learning
- Better generalization than competing baselines, when tested on new tasks
The work also includes examples of POMDPS with higher dimensional observations (25x25 pixels) showing the proposed pretraining paradigm achieves better performance than competing baselines.
Valuable analysis of the performance of in-context learning when having access to PPO-trained policies, bringing the setting closer to AD. I particularly value how authors evaluate the effect of different sources for the context states and actions. I am wondering whether PPO, PPO is directly comparable to AD though since in their case they provide in-context actions and trajectories for different stages of training.
By chosing simple domains, the work can thoroughly explore the behavior of in-context learning for different properties of the training data, particularly in the multi-armed bandit section. To my knowledge, this is the first work to build connections of in-context learning in decision making problems and posterior sampling, which allows to:
1. Provide an upper bound on the regret of the proposed transformer-based policy.
2. Show that if the pretraining in-context dataset comes from policies that were only trained on the data and task present in the dataset, the final policy will be the same.
Clarity:
Generally clear writing and background. The experiments are simple but informative, and the work is easily reproducible. More notation and clarity should be given in the theory section, such as stating what is H, and providing more intuition in the propositions.
Weaknesses
Novelty: while this paper conducts a reasonable study on the advantages of in-context learning in decision making problems, this is a feature that has been recently studied in more complex domains and tasks (see related works next). My main concern in this work is that there is not much novelty or learnings compared to those works, and that despite making a reasonable study of the capabilities of transformers for few-shot decision making problems, they are studied in very simple scenarios.
Related work: Several works have studied in context-learning for decision making problems, in more complex scenarios [1, 2]. These works should be cited. I am also not clear on what novelty this work provides with respect to the aforementioned papers. I would like a clarification on this as well.
[1] https://arxiv.org/pdf/2206.13499.pdf
[2] https://arxiv.org/abs/2301.07608
Baselines: While RL2 is a fair baseline, I think authors should look at other meta-RL algorithms, or tune the existing baselines to be more comparable. I understand that PPO is used as a reference, but it should at least be finetuned with the few-shot examples given - it is hard to compare methods when some are actually seeing less data. AD is designed to learn to do RL and therefore works best when having a curriculum of trajectories, which is not the case here.
Questions
It is not intuitive to me that a model trained with expert data (DPT-Exp) behaves poorly when the % of expert data is high (Figure 3.a) since that is what is seen during training, I would expect it to b heave the best. Could authors elaborate on why this is the case?
The setting in lines 312-317 makes a lot of sense. Why not using that in all the experiments, if as claimed "the original pretraining method can be seen as a simpler approximation
316 of this modified method"?
I don't understand assumption 1.1, if we are following that assumption, why not directly using P_{pre} as out policy? Some intuition would be helpful.
Proposition 6.3 is very unitnuitive to me, if I understand it correctly. Is it stating that a pretraining dataset coming from PPO and one coming from a random policy would perform the same? I don't see anything in the proposition that implies otherwise, but this doesnt match in the empirical experiments.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.