LMAct: A Benchmark for In-Context Imitation Learning with Long Multimodal Demonstrations

We thank the reviewer for their thorough assessment and insightful feedback. We are pleased that they think that our paper is particularly relevant to the current trend of reasoning models and that it is great that this paper identifies areas where people can evaluate their reasoning models on longer horizon tasks.

The paper overlooks the relevant line of work on in-context learning for control.

Indeed, the papers the reviewer listed are relevant to the broader context of our work, and we have expanded the related work section to discuss all of these references. While the methods explored in that exciting line of work focus on developing novel agents with strong in-context learning capabilities, our paper's primary goal is to benchmark the current state of existing, general-purpose frontier LMs on such tasks. Understanding how to best leverage insights from specialized agent research to enhance these large, pre-trained models remains an open question, which underscores the importance of establishing clear benchmarks like LMAct.

The paper only has 6 environments, which could lead to quick saturation given the current speed of AI research.

As discussed in Section 4, we designed our benchmark such that it can easily be made more difficult as the capabilities of frontier models evolve. While our benchmark is currently in its “easiest” form (e.g., grid world without obstacles, chess against Stockfish level 0 (less than 1300 Elo)), our results demonstrate that even this version reveals significant limitations in current frontier models. Therefore, we believe that our benchmark provides a strong and high-resolution signal to measure the progress of frontier models. Once our benchmark is saturated (which is far from being the case right now), we can easily construct a more challenging v2.

Consider including partially observable environments to expand the benchmark.

We tried to cover a range of environments that are relatively easy for humans but where current models still struggle to reach expert performance. As a result, we chose fully observable tasks (technically, Atari requires at least two frames to determine the velocities) since they have simple, reactive policies and do not require exploration.

Partially observable tasks require consistent integration of information across several potentially non-adjacent time steps, which is a great test of a model’s ability to densely attend to the information in the context. Moreover, there is the theoretical problem of self-delusions in imitation of an expert that has hidden information (e.g., the hidden belief state of the agent), which is necessary under partial observability (see Ortega et al. (2021)). We want to keep this benchmark free from these complicating issues and think that the right time to move to such harder tasks is when frontier models easily solve simple, fully observable tasks at expert level (we are currently at beginner level — see the previous response). Independent of that, we agree with the reviewer that it is great to see benchmarks on interactive decision-making tasks under partial observability (perhaps better suited for in-context RL, which avoids the self-delusion problem) and refer to, e.g., the BALROG benchmark.

We have added a brief discussion to acknowledge the value of partially observable benchmarks while clarifying our focus on fully observable tasks in the revised manuscript.

Shaking the foundations: delusions in sequence models for interaction and control. Pedro Ortega et al. https://arxiv.org/abs/2110.10819.

Does the temperature affect the in-context learning capabilities?

We deliberately set the temperature to 0 to ensure the reproducibility of our results wherever possible (the API for the o1 models does not support this option). Whether the temperature also affects the performance is an interesting direction for future work. Note that since OpenAI’s models perform relatively well in most of our tasks, the effect of the temperature value does not seem to be very significant in our benchmark.

We thank the reviewer for their careful assessment and constructive feedback. We are pleased that they think our paper addresses an important problem, our paper covers a wide range of tasks, and our paper considers a wide variety of frontier models, making for a comprehensive study.

Could you include more analysis about failure modes and potential improvements?

We agree that one of the most important follow-ups is to analyze why models fail on our benchmarks and how to address these failures. We have already conducted a substantial set of such experiments in Appendix C, including:

• Checking whether models can successfully “replay” (or copy) an episode, which most models can (Appendix C.2)
• Checking to which degree bad performance can be attributed to the percentage of illegal vs. legal but suboptimal actions (Appendix C.3)
• Investigating the fraction of repeated actions in Atari – Phoenix, which reveals that models have a high tendency to repeat the previous actions, which leads to few shots being fired, explaining why the performance is worse than the random action baseline (Figure A14)
• Investigating whether failures can be attributed to various hyperparameter settings, e.g., checking whether showing legal actions alleviates the problem of having to infer the correct action format from the demonstrations (Appendix C.4)

The trends of these investigations are complex (often model and task-dependent), and definitive answers would require sophisticated further analysis, which we consider outside the scope of our benchmark paper. We have added a compact discussion of this analysis to Section 4 in the updated manuscript.

Do the o1 models perform better because they use a larger context length (8192 tokens) than the other models (2048 tokens)?

Section 2.1 states that the o1 models use a larger output sample length (i.e., the maximum number of tokens to generate, including the “thinking tokens” for o1), not a larger context length. The (input) context length is the same across all models. We ablated the maximum sample length in Figure A7 and showed that only the o1 models benefit from a larger sample length. This is because the o1 models cannot finish their reasoning traces if they do not have sufficient “thinking tokens” at their disposal. Since our goal is to evaluate the models’ best possible performance, we set the output sample length to 8192 for the o1 models (which is expensive) and 2048 for all other models (since this is cheaper and does not affect their performance).

Do you also provide the rewards from the demonstrations?

No, we only perform in-context imitation learning, so we do not provide reward information to the models. We only use the rewards to compute the model scores for our benchmark results. We consider in-context reinforcement learning an interesting direction for future work.

Are the 100 steps per episode sufficient for most tasks? Can this number be attributed to the limited context length of current models?

Yes, for chess, crossword puzzles, gridworld navigation, and tic-tac-toe, 100 steps are more than sufficient to complete an episode (e.g., for chess, the average number of steps per game is 38 — see Section 2.2). However, for Atari and DM Control, we would ideally want to evaluate more steps (e.g., we only evaluate roughly 6 seconds of Atari play — see Section 2.2), but the context size limitations of current frontier models (32 Atari demonstrations episodes with 100 steps already hit the 1M token limit) do not allow this evaluation. We chose a maximum of 100 steps to be consistent across tasks while still providing a meaningful numerical signal (e.g., the models already fail in the early stages of Atari).

Why are there a lot of illegal actions in chess and crosswords? Could adding the rules of the game to the context be helpful?

The action space for these two tasks is more complex than for the other tasks (maybe except for Cheetah Run, which most models seem to be able to match, though). Adding game-specific explanations would probably help increase the performance, but our study is to investigate task-agnostic in-context learning behavior since such game-specific behavior may not always be available for real-world tasks (e.g., videos of human experts performing tasks in dynamic environments).

Why does the performance degrade with increasing demonstrations for Cheetah Run?

We do not know with certainty. Figure A18 shows that it cannot be attributed to an increased number of illegal actions, i.e., the models’ actions are increasingly suboptimal rather than illegal. A definitive answer to what causes in-context interference would require a thorough investigation, which is beyond the scope of this benchmark paper, but we have added a brief discussion of these observations and open questions to the paper.

We hope our responses and revisions effectively address the reviewer’s concerns and clarify the paper's contribution.

We thank the reviewer for their positive feedback and interesting questions. We are pleased that they think that our claims are clear and supported with convincing evidence, that our experiments are well designed and presented, and that our results are an important sanity check of the current capability of frontier models for interactive decision making.

Why do more demonstrations have little and sometimes negative impact?

At least in theory, our fully observable tasks do not require an optimal policy to attend to information more than one step in the past. Accordingly, an LM’s performance on our tasks must lie somewhere between two extremes:

• If the LM has stored the optimal policy in its weights, then any historical information would only inform the model about the task at hand and which of its in-weights policies to select (for which even a single observation might suffice).
• If the LM has no suitable pretrained policies in its weights and, therefore, would have to rely purely on distilling the expert behavior in the context, more demonstrations should increase performance.

In practice, the LMs’ behavior lies somewhere between these extremes and, as our results show, greatly varies per model and task. Possible “mechanical” explanations for this behavior are that LMs are pretrained to attend sparsely over their context (in which case more observations would only help marginally), have a recency bias (which would make attention beyond the current/previous episode even less likely), and are usually not explicitly trained for in-context learning from examples (although it’s a somewhat emergent capability of pre-training).

We mostly observe negative impact, i.e., “in-context interference” for Cheetah Run and gridworld navigation with ASCII observations. In both cases, our analysis of the percentage of illegal actions (Figures A18 and A19) suggests the problem is not due to an increased number of illegal actions but, instead, increasingly suboptimal actions. A definitive answer to what causes in-context interference would require a thorough investigation, which is beyond the scope of this paper (which introduces the benchmark), but we have added a brief discussion of these observations and open questions to the paper. Please also see the summarized findings of our investigations of failure modes (mostly appendix results) in our response to reviewer 46g3, under “Could you include more analysis about failure modes and potential improvement strategies?”.

How should we change the training of these frontier models to be more effective with in-context prompting when solving interactive decision-making tasks?

While we do not have a definite answer, it seems plausible that pretraining or finetuning with data from interactive decision-making tasks, and, in particular, in-context imitation of an expert policy, would be quite effective. Another promising direction could be to focus on models capable of general in-context reinforcement learning, which is a bit different than our in-context imitation setting, but in principle, all our tasks could easily be extended by providing additional reward observations. We consider these very interesting and important directions for future research and have updated our manuscript with the above discussion.

We thank the reviewer for their thorough review, constructive feedback, and pointers to additional relevant related work. We are pleased that they think that the soundness and validity of the experimental design are beyond doubt and that the originality of the work arises from creative combinations of existing environments to study problems in a long context.

Why did you not evaluate the two other state-of-the-art models QWEN and Llama 3?

We had the dilemma of choosing a set of models to evaluate and ultimately settled on 8 models but did not mean to imply that any models outside of this set are not state-of-the-art. Unfortunately, we were unable to evaluate Llama 3 due to licensing issues, which is why we decided to focus on closed-weights models (which we explicitly stated in Section 2.1). Nevertheless, we fully agree that both QWEN and Llama 3 would be interesting additions to our benchmark and have discussed them in our revised manuscript. Since the landscape of state-of-the-art models is rapidly evolving, we open-sourced our benchmark to enable the community to evaluate QWEN, Llama 3, and any other future models using our approach.

Why did you not evaluate transformers created for long-context tasks (such as RMT and RATE) on your benchmark?

As the reviewer points out, RMT and RATE are specialized transformer models that were specifically developed for long-context tasks. However, as stated in Section 1, our goal is to evaluate state-of-the-art LMs in dynamic environments. Nevertheless, we agree that innovations around augmenting transformers with recurrency (and forms of more explicit retrieval) are very interesting approaches to potentially overcome LMs’ performance limitations in our benchmark, and we have expanded the related work section to discuss these approaches. Similar to QWEN and Llama 3, we are excited to see how the community will leverage the open-source nature of our benchmark to improve their models.

What is the practical usefulness of this benchmark?

We thank the reviewer for raising the important point regarding practical usefulness and transferability. While the LMAct environments are simpler than complex real-world scenarios like robotics, they require fundamental agentic capabilities (long-context multimodal understanding, in-context imitation learning) that are often tested in robotics research and likely necessary for real-world success. LMAct serves as a controlled testbed to evaluate these core skills, diagnose current model limitations, and guide research toward more generally capable agents. Although direct transfer is beyond this paper's scope, these findings can inform future work addressing that challenge. We have added a brief discussion clarifying this focus in the revised manuscript.

Why did you not compare your benchmark to BabiLong?

The BabiLong benchmark is relevant in the wider context of our work, and we have added a brief discussion in the related work section. Thanks for the pointer!