XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning

## Abstract

Abstract Following the success of the in-context learning paradigm in large-scale language and computer vision models, the recently emerging field of in-context reinforcement learning is experiencing a rapid growth. However, its development has been held back by the lack of challenging benchmarks, as all the experiments have been carried out in simple environments and on small-scale datasets. We present XLand-100B , a large-scale dataset for in-context reinforcement learning based on the XLand-MiniGrid environment, as a first step to alleviate this problem. It contains complete learning histories for nearly 30 , 000 30 000 30,000 30 , 000 different tasks, covering 100 100 100 100 B transitions and 2.5 2.5 2.5 2.5 B episodes. It took 50 , 000 50 000 50,000 50 , 000 GPU hours to collect the dataset, which is beyond the reach of most academic labs. Along with the dataset, we provide the utilities to reproduce or expand it even further. With this substantial effort, we aim to democratize research in the rapidly growing field of in-context reinforcement learning and provide a solid foundation for further scaling. The code is open-source and available under Apache 2.0 licence at https://github.com/dunno-lab/xland-minigrid-datasets .

## 1 Introduction

In-context learning, i.e. the ability to learn new tasks purely based on examples given in the context during inference and without any weight updates, was initially thought to be an emergent property of large language models, such as GPT-3 (Brown et al., ). However, it was quickly discovered that small transformers are also capable of in-context learning (Kirsch et al., ; Von Oswald et al., ), and even many non-transformer models have such abilities (Bhattamishra et al., ; Akyürek et al., ; Park et al., ; Grazzi et al., ; Vladymyrov et al., ; Tong & Pehlevan, ). More importantly, driven by properties of data, rather than the architecture (Chan et al., ; Gu et al., ), in-context learning is not specific to language modeling and has been found in other domains, e.g. image generation (Bai et al., ; Najdenkoska et al., ; Doveh et al., ; Tian et al., ).

However, despite the rapid adoption of the transformer architecture in reinforcement learning (RL) after the release of Decision Transformer (DT) (Chen et al., ; Hu et al., ; Agarwal et al., ; Li et al., ), models with in-context learning capabilities appeared only recently. This delay is caused by a number of reasons. Firstly, to transition from in-weights to in-context learning, a model should be trained on tens of thousands of unique tasks (Kirsch et al., ). Unfortunately, even the largest RL datasets currently contain only hundreds or tasks (Padalkar et al., ). Secondly, it was necessary to determine the right way to provide a context to a transformer and develop the data collection pipeline, which for many methods (Laskin et al., ; Lee et al., ; Shi et al., ) was different from what is commonly available in existing datasets (see Section 3).

Due to the lack of suitable datasets and the high cost of collecting data in existing environments, the recent wave of in-context RL research (Laskin et al., ; Lee et al., ; Norman & Clune, ; Kirsch et al., ; Sinii et al., ; Zisman et al., ) used environments with very simple task distributions, where it was feasible to collect datasets with hundreds of tasks. While these benchmarks are affordable, they are not suitable for the comparison of methods at scale on tasks with high diversity and difficulty, which is essential for real-world applications. Because of this, the development of in-context RL is currently hindered by these factors. We believe it is crucial to address these barriers, given the essential role of in-context learning in the path to foundation models and truly generalist agents (Team et al., , ; Kirsch et al., ; Lu et al., ; Liu et al., ).

We release XLand-100B, a large-scale dataset for in-context RL based on the XLand-MiniGrid (Nikulin et al., ) environment. It contains complete learning histories for nearly $30,000$ different tasks, covering $100$B transitions and $2.5$B episodes. It took $50,000$ GPU hours to collect the dataset, which is beyond the reach of most academic labs. In contrast to most existing datasets for RL, our dataset is compatible with the most widely used in-context learning RL methods (see Section 3). With this substantial effort, we aim to democratize the research in the rapidly growing field of in-context RL and provide a solid foundation for further scaling.

Along with the main dataset, we provide a smaller and simpler version for faster experimentation, as well as utilities to reproduce or extend the datasets even further. We carefully describe the entire data collection procedure (see Section 4), providing all necessary details about the algorithm used for collection, filtering and relabelling with expert actions (see Section 4.2). We analyse the resulting dataset to ensure that we have met all the requirements for in-context RL (see Section 4.3). In addition, we conducted preliminary experiments with the common baselines on the collected datasets, showing there is still a lot of research needed to improve the in-context adaptation abilities on complex tasks (see Section 5).

Table: Table 1: Comparison with other RL datasets.

## 2 Background

### 2.1 In-Context Reinforcement Learning

Multiple methods for in-context RL have come out, each offering a different way of training and organising the context (Laskin et al., ; Lee et al., ; Mirchandani et al., ; Liu & Abbeel, ; Raparthy et al., ; Shi et al., ). We focus on Algorithm Distillation (AD) (Laskin et al., ) and Decision-Pretrained Transformer (DPT) (Lee et al., ), which we chose as the main methods for our work due to their simplicity and generality.

Algorithm Distillation. AD (Laskin et al., ) was one of the first to show that in-context learning was possible in RL, and captures the details of many other more recent methods (Mirchandani et al., ; Liu & Abbeel, ; Shi et al., ) while remaining very simple. It trains a transformer, or any other sequence model, to autoregressively predict next actions given the history of previous interactions, i.e. observations, actions and rewards. To transition from in-weights to in-context learning, it is essential that the context should contain multiple episodes ordered by an increasing return, which is different from the way it is done in DT-like methods (Chen et al., ; Janner et al., ; Lee et al., ).

Decision-Pretrained Transformer. DPT is an alternative approach inspired by the Bayesian inference approximation (Müller et al., ). Unlike AD, it trains a transformer to predict the optimal action for a query state given a random, task specific, context. That is, the context that does not have to be ordered, but only has to contain transitions belonging to the same task. Thus, DPT requires access to optimal actions, but does not require a dataset of learning histories.

In addition, the theoretical analyses of AD and DPT methods (Lin et al., ; Wang et al., ) showed that they can implement near-optimal online RL algorithms such as Lin-UCB, Thompson sampling or even temporal difference (TD) methods solely during the forward pass.

### 2.2 XLand-MiniGrid

Figure: Figure 1: Visualization of a generic XLand-MiniGrid environment. Grid layout should be selected in advance, while the positions of the objects are randomized on each reset.
Refer to caption: x1.png

Starting from the seminal work of Wang et al. (); Duan et al. (); Finn et al. () on meta-RL, much of the subsequent work (Zintgraf et al., ; Melo, ; Grigsby et al., ; Lu et al., ; Shala et al., ; Beck et al., ) has focused on environments that either have very simple task distributions, or have very small and limited distributions of hard tasks. This is because, to generalize in meta-RL, training needs to be performed on many different tasks, significantly increasing the cost and time required for experimentation. Recently, Nikulin et al. () released XLand-MiniGrid, a GPU-accelerated environment and million-task benchmarks that significantly lowered the entry barrier for meta-RL research. We will describe it shortly here.

Environment. XLand-MiniGrid is a complete rewrite of MiniGrid (Chevalier-Boisvert et al., ) in JAX (Bradbury et al., ), incorporating a notion of rules and goals from XLand (Team et al., ). Leveraging JAX, it can run on a GPU or TPU accelerators at millions steps per seconds. At its core, it is a goal-oriented grid-world environment with simple underlying dynamics, partial observability and sparse rewards. The action space is simple, consisting mainly of navigation and interaction with game objects, such as opening a door or picking and placing items. Observations are symbolic "images" encoding the agent surrounding as tile and color ID’s. Rules are functions that can change the state of the environment based on some conditions, e.g. when two specific objects are places near each other, they both disappear and one new object is placed. Goals are similar, except they only validate a predefined conditions and do not change anything. Composing different rules and goals together we can create new tasks with varying reward and dynamics functions. For more detailed description we refer to Nikulin et al. ().

Benchmarks. Along with the environment itself, Nikulin et al. () released a tool for the procedural generation of a vast number of unique tasks with varying levels of difficulty. Each task is represented by a binary tree, where the root is the goal to be achieved and rest of nodes define rules of the environment to be triggered in a recurring sequence. To standardize comparisons, four pre-sampled benchmarks with increasing diversity were provided: trivial, small, medium, high, each with one million unique tasks. For this work, we chose medium as a middle ground between yet unsolved high and less challenging small benchmarks. We also use trivial for smaller and simpler dataset version (see Section 4).

## 3 The Missing Piece For In-Context RL

In order to successfully train an in-context agent, training data shall meet certain criteria. To start with, the data should be comprised of actual learning histories (Laskin et al., ) or their approximations (Zisman et al., ), that contain enough exploration and exploitation phases of learning. Learning with just expert trajectories would not be sufficient for in-context ability to emerge, since an agent needs to know the history of policy improvement (Laskin et al., ; Kirsch et al., ). Another approach is to learn from optimal actions as proposed by Lee et al. (), but it is unclear how to access the optimal policies to get them. Besides, the data needs to contain thousands of different tasks to learn from (Kirsch et al., ). That is, for a simple task to find two squares on a $9\times 9$ grid, an agent needs to see around $2000$ different combinations of goals to start adapting for unseen locations (Laskin et al., ). Since such data was never collected and put into a single dataset, all current in-context RL practitioners were forced to generate data on their own, which inevitably added more complications in reproduction of the methods.

Besides, collecting thousands of different in-context episodes requires a considerable commitment, as training numerous RL agents is expensive in terms of time and resources. To that matter, the data used in current research is collected in very simplistic environments with straightforward goals, like reaching a specific target on a map (Laskin et al., ; Lee et al., ; Zisman et al., ) or to apply forces to actuators in order to walk a robot (Kirsch et al., ). This significantly slows down the pace of in-context RL research, as it is not only hard to test the applicability of proposed methods, but also yet unfeasible to determine the scaling laws in these environments.

To provide a complete picture for the reader, we briefly discuss the existing datasets and highlight why they are unfit for training in-context RL agents. For simplicity, we categorize them into two groups: classical datasets designed for offline-RL and datasets collected for large-scale supervised learning. Note that this categorization is fuzzy in nature and serves only for better understanding of the current structure in RL data.

Offline RL Datasets.
The datasets in this category can be considered classical, as some of them exist for more than four years now (Fu et al., ). They were initially proposed for offline RL, containing simple tasks with a flat structure, e.g. perform locomotion with different robots (Lu et al., ) or path finding in a maze. Some of them also contain data from robotic manipulators (Fu et al., ; Rafailov et al., ), or even Atari frames (Gulcehre et al., ). Another datasets collect data for more sophisticated environments, such as NetHack Learning Environment (Hambro et al., ; Kurenkov et al., ). However, the aforementioned datasets offer $<100$ different tasks with a fixed policy (except for the -replay datasets, which have limited coverage of various policies). This limitation makes it difficult for in-context RL to emerge from such data. To overcome this pitfall, we collected almost $30,000$ tasks with a deep ruleset structure, that are a challenging problem to solve.

Large-Scale Supervised Pretraining.
Recent progress in generalist agents, which can solve a multitude of environments, has been made possible thanks to large datasets. GATO dataset (Reed et al., ), however not being released to the public, consists of 1.5 trillion transitions along with 596 tasks, which makes it one of the largest dataset in RL. The open-sourced analogue, the JAT dataset (Gallouédec et al., ), is smaller in size with 157 tasks and 300 million transitions, but it provides comparable performance on most of the benchmarks. Both datasets contain expert RL demonstrations from BabyAI (Hui et al., ), Atari games (Bellemare et al., ), Meta-World (Yu et al., ) and more.

Another large dataset, Open X-Embodiment (Padalkar et al., ), is a combination of more than 60 datasets from different robotics research labs. It consists of 527 different tasks in robotics with the demonstrations from mostly from human experts. Despite the large quantity of transitions in these datasets, they do not contain learning histories with improving policies, making their application for in-context RL quite challenging. On the contrary, our XLand-100B dataset consists of $100$ billions of transitions of RL agents’ learning histories, making it possible for in-context abilities to emerge.

The only potentially applicable dataset to use for in-context RL is AlphaStar Unplugged (Mathieu et al., ). Although the authors did not initially plan to collect a suitable dataset, the data can be sorted by players’ MMR (analogous to Elo rating). This sorting can be considered a steady policy improvement, thus enabling the in-context RL ability. For more details on the datasets, refer to Table 1.

**Table 2: Descriptive statistics of XLand datasets.**
| Dataset | XLand-Trivial-20B | XLand-100B |
| --- | --- | --- |
| Episodes | 868,805,556 | 2,500,152,898 |
| Transitions | 19,496,960,000 | 112,598,843,392 |
| History length | 60,928 | 121,856 |
| Num tasks | 10,000 | 28,876 |
| Max task rules | 0 | 9 |
| Observation shape | (5, 5) | (5, 5) |
| Num actions | 6 | 6 |
| Mean final return | 0.915 | 0.894 |
| Median final return | 0.948 | 0.925 |
| Median episode transitions | 22.45 | 57.75 |
| Disk size (compressed) | 60 GB | 326 GB |

## 4 XLand-100B Dataset

We present XLand-100B, a large dataset for in-context RL, and its smaller and simpler version XLand-Trivial-20B for faster experimentation. Together they contain about 3.5B episodes, 130B transitions and 40,000 unique tasks (see Table 2 for detailed statistics). Datasets are hosted on public S3 bucket and freely available for everyone under CC BY-SA 4.0 licence. The code is open-source and available under Apache 2.0 licence at [https://github.com/dunno-lab/xland-minigrid-datasets](https://github.com/dunno-lab/xland-minigrid-datasets). Next, we describe the data format, collection and evaluation.

### 4.1 Data Format

Storage format. We chose to store the datasets in HDF5(^1^11[https://github.com/HDFGroup/hdf5](https://github.com/HDFGroup/hdf5)) file format based on its popularity and convenience. It allows to work with large amounts of structured data without loading it to memory, store arbitrary metadata, and customise compression and chunk size to maximize the sampling throughput. We used gzip compression with default compression strength of 6, which reduced dataset size from almost 5TB+ to just $\sim$600GB for our main dataset. Using a little trick described later, we were able to reduce the size even more to just 326 GB (see Table 2). However, a naive use of compression can dramatically increase batch sampling time and slow overall training time down. We tuned HDF5 cache chunk size specifically to maximize sampling throughput for large sequence lengths. After tuning, we achieved a fourfold speedup over the naive compression, and were only two times slower compared to no compression. Given that we reduced the dataset size by a factor of 15, we see this as a good trade-off, which increases the overall affordability of the dataset. See Appendix F for throughput benchmarks.

Data and metadata format. We collect complete learning histories, i.e. for each history we store all observations, actions, rewards and dones encountered during agent training in separate HDF5 groups with unique IDs per history (see Appendix D for more details). To be compatible with DPT-like methods (Lee et al., ), we also store expert actions for each transition (see Section 4.2). Unlike popular formats such as RLDS(^2^22[https://github.com/google-research/rlds](https://github.com/google-research/rlds)) and Minari(^3^33[https://github.com/Farama-Foundation/Minari](https://github.com/Farama-Foundation/Minari)), we store transitions as one sequential array per history per modality. The rationale here is that under compression, it is much cheaper to sample slices of long episodes in sequence rather than sampling across different groups.

We also store observations efficiently to further reduce dataset size. Instead of storing two channels for tile and color, we map their indexes into Cartesian product of colors and tiles, halving the storage size. They can be decoded easily during sampling without any overhead with divmod function. In addition, for each history we store the XLand-MiniGrid environment ID, benchmark ID and ruleset ID, which can be used later to filter the dataset, e.g. based on the complexity of the tasks, split into train and test or set up the environment for evaluation.

### 4.2 Data Collection

Figure: Figure 2: Evaluation return for multi-task goal-conditioned reccurent PPO pretraining on 65k tasks. Pretrained agent was further used a starting point for single-task finetuning during dataset collection.
Refer to caption: x2.png

At a high level, data collection was organised into three stages, namely multi-task task-conditioned pre-training; single-task fine-tuning to collect learning histories; and finally post-processing and filtering. Although we used highly optimised GPU-accelerated implementations of the base RL algorithm and environment, it still took $50,000$ GPU hours to collect the full dataset. Collecting a dataset of this size for any other environment suitable for in-context RL, e.g. Meta-World (Yu et al., ), would take much longer, which is unlikely to be feasible for most practitioners (Nikulin et al., ). Next, we describe the collection process, including the selection of the base RL algorithm and all subsequent steps. We provide exact hyperparameters for each stage in Appendix E.

Base algorithm.
For our datasets we chose PPO (Schulman et al., ), due its high scalability and compatibility with massively parallel environments. We ported the implementation from recurrent PPO provided by (Nikulin et al., ), customizing it to meet our needs. We added callbacks for saving transitions during training and extended agent architecture to take ruleset encoding as an optional condition for pre-training. We used GRU (Cho et al., ) for memory, as it showed satisfactory performance on preliminary experiments. Since the base algorithm was implemented in JAX (Bradbury et al., ), we were able to just-in-time compile the entire training loop, achieving 1M steps per second during training on one GPU with mixed precision enabled. As PPO is not the most sample-efficient algorithm, it was still necessary to train it on billions of transitions. Fortunately, as we trained it on thousands of environments in parallel, learning history for each particular environment is quite short, i.e. around 120k transitions.

Pretraining. For our main XLand-100B dataset we uniformly sampled tasks from medium-1m benchmark from XLand-MiniGrid. It contains tasks of various difficulty, ranging from zero to nine rules.
Unfortunately, on many hard tasks our base algorithm could not manage to converge in the time budget allocated for a single training run. On the hardest tasks, it was not even possible to get a non-zero reward at all, due to the exploration challenge that such tasks poses. In order to speed up convergence and exploration on harder tasks, we pre-train an agent in a multi-task task-conditioned manner. We expose the ruleset specification, which is usually hidden from the agent, and encode the goal and rules via embeddings, concatenating resulting encodings and passing it as an additional input to the agent. After that, we train the agent on 65k tasks simultaneously for 25B transitions. As Figure 2 shows, such an agent learns to generalize zero-shot on new tasks quite well, although we do not aim to push it to the limit, as zero-shot generalization does not produce a smooth learning history during fine-tuning. We skip this stage for XLand-Trivial-20B dataset due to the simplicity of the tasks in trivial-1m benchmark.

Finetuning. This is a key stage in the data collection process, during which we finetune a pretrained agent while recording the transitions encountered into the dataset. We finetune the agent using 8192 parallel environments for 1B transitions on 30k uniformly sampled tasks from medium-1m benchmark. We mask out the task-conditioning encoding to prevent zero-shot generalization. We record transitions only from first 32 parallel environments. This way, we can keep the size of the dataset manageable, still leaving the possibility to study scaling laws and generalization in a controlled manner. For example, we can train AD (Laskin et al., ) on all 30k tasks using one history per task or on $\sim$ 900 tasks using all histories per task to equalize the number of training tokens. For the XLand-Trivial-20B dataset, instead of fine-tuning, we train the agent from scratch on 10k uniformly sampled tasks from the trivial-1m benchmark, keeping all other hyperparameters the same. In the Figure 3 we show the effect of finetuning on hard tasks (with more than seven rules) compared to training from scratch. It can be seen that we are able to show strong performance even on the hardest tasks, increasing the diversity and coverage of the resulting dataset. For the same results on tasks of all levels of difficulty, see Appendix G.

Figure: Figure 5: Learning histories for the XLand-100B dataset separated by number of rules. For visual clarity, we show only a sample of the possible number of rules and normalize the number of episodes, as they may vary considerably.
Refer to caption: x5.png

Postprocessing. After fine-tuning, it was necessary to additionally label the transitions with the expert actions to support DPT-like methods (Lee et al., ). To do this, we walk through the entire learning history with the final policy, starting from the initial hidden state for the RNN. We discuss the validity of such a labelling scheme later in the Section 4.3. Finally, all individual learning histories from different tasks were combined into one large dataset. To ensure quality, we filtered out any task with a final return below $0.3$ as an unrepresentative learning history. There were some failures, such as GPU crashes, which are inevitable during large-scale training. So any runs with corrupted data were also filtered out. In total, we filtered out about 1k tasks, leaving almost 29k tasks in the final dataset. We provide detailed statistics for each dataset in Table 2 and the final distribution of tasks by number of rules in Figure 4.

### 4.3 Data Evaluation

In this section we validate that the resulting XLand-100B dataset actually fulfils the two most important requirements for in-context RL, namely it contains learning histories with distinct policy improvement pattern and has expert
actions for each transition (see Section 3 for a discussion). We provide analogous results for XLand-Trivial-20B in the Appendix H.

Improvement history. In the Figure 5 we show the averaged return from the learning histories separated by the number of rules. In order to better show the speed of learning on the same scale, we have normalized the x-axis for each level of difficulty, as the number of episodes can vary greatly (as it takes more time to solve complex tasks). One can see that the dataset provides a whole range of learning speeds, from very fast on easy problems to much slower on the hardest, which may be important for methods based on AD (Zisman et al., ; Shi et al., ). For a learning history averaged over full dataset see Appendix H.

Figure: Figure 6: Agreement between actions predicted by the expert and the actual actions in the learning history. We use final PPO policy as an expert for actions labeling.
Refer to caption: x6.png

Expert actions relabeling. In contrast to AD, which predicts next actions from the trajectory itself, DPT-like methods require access to optimal actions on each transition for prediction. However, for the most nontrivial or real-world problems, obtaining true optimal actions in large numbers is unlikely to be possible.
Recently, Lin et al. () introduced approximate DPT, scheme where expert actions are estimated from the entire history by some algorithm. We implemented this scheme for lack of evident alternatives. However, we had to make sure that such a labeling is adequate in our case and the expert at the end predicts actions close to what the policy did in reality near the end of training. This is not obvious, as during the labeling it can diverge into out-of-distribution hidden states for RNN. In the Figure 6 we show that on XLand-100B the agreement between the predicted actions by the expert and the actual actions increases closer to the end of the learning history, meaning that the expert does not diverge during relabeling.

Figure: Figure 7: Comparison of the learning histories in -100B dataset vs. AD performance on the same training tasks. AD is able to solve simple tasks, however its performance degrades as the rulesets get deeper. The context length of the model is $1024$. The evaluation parameters, except training tasks, are the same as in Figure 8.
Refer to caption: x7.png

## 5 Experiments

Figure: Figure 8: AD performance on our datasets for different sequence lengths. Both datasets lead to the emergence of in-context ability. We report the average return on $1024$ unseen tasks across $3$ seeds.
Refer to caption: x8.png

In this section, we investigate whether our datasets can enable an in-context RL ability. Additionally, we demonstrate how well current in-context algorithms perform across different task complexities and outline their current limitations. We take AD (Laskin et al., ) and DPT (Lee et al., ) for our experiments, the exact implementations details are in Appendix I and Appendix J. Both methods were trained on XLand-Trivial-20B and XLand-100B with $\{512,1024,2048,4096\}$ and $\{1024,2048,4096\}$ context lengths respectively. We do not include a 512 context length for the -100B dataset as we consider it too short, given the nearly 3x increase in median episode length between the two datasets. For evaluation, we run three models on $1024$ unseen tasks for $500$ episodes.

AD. Algorithm Distillation shows an emergence of in-context ability during training on both datasets. Figure 8 demonstrates the performance of the method for different context lengths. On -Trivial-20B dataset it is able to show a stable policy improvement from about $0.28$ to $0.4$ during the evaluation. For -100B the performance is similar, but the pace of improvement is faster. We hypothesise that it happens due to wider data-coverage, since the agent sees more complex tasks and is able to learn faster from them.

To further examine the performance on the -100B dataset, we evaluate AD on the training tasks from the dataset and separate the performance based on the complexity of the tasks. The complexity is defined by the number of rules an agent needs to trigger before the successful completion. As shown in Figure 7, AD is able to demonstrate in-context abilities on simple tasks, but it struggles with more complex ones. We speculate that our results are not final, as AD is not yet fully capable of learning to solve the training tasks. We believe there is a need for further research to discover new and more sample-efficient architectures capable of solving the more complex rulesets of our dataset.

DPT. Decision-Pretrained Transformer (Lee et al., ) is another method that exhibit in-context RL capabilities. However, in our experiments we were unable to train it so that these abilities emerge. Figure 18 demonstrates the lack of performance even on the simplest trivial tasks. We believe this is closely connected to the inability of DPT to reason in POMDP environments. For a detailed explanation, we refer the reader to Section J.1.

## 6 Limitations and Future Work

There are several limitations to our work, some of which we hope to address in future releases. Despite the size and diversity of the datasets provided in terms of tasks, we do not provide diversity in terms of the domains, as all tasks share the same observation and action spaces. In addition, the tasks also share an overall latent structure, i.e. it is always a form of binary tree. This can be addressed with more diverse benchmark generators in the XLand-MiniGrid library (Nikulin et al., ). All learning histories were collected on grids with only one room, which may limit the transfer to the harder layouts with multiple rooms containing doors. Finally, the effect of fine-tuning from pre-trained checkpoints is underexplored and could potentially hurt performance, as there are many learning histories that start from a high reward. We hope to improve the RL baseline for data collection to avoid the need for multi-task pre-training in the future.

## Acknowledgments and Disclosure of Funding

At the time of writing, all authors were employed by AIRI or/and Tinkoff. All computational resources were provided by AIRI and Tinkoff.

## Checklist

- 1.
For all authors…
(a)
Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] See Section 4.
(b)
Did you describe the limitations of your work?
[Yes] See Section 6.
(c)
Did you discuss any potential negative societal impacts of your work?
[N/A] We provide reasons for this answer in Appendix B
(d)
Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] We have read the ethics guidelines and checked our paper satisfies them
- 2.
If you are including theoretical results…
(a)
Did you state the full set of assumptions of all theoretical results?
[N/A]
(b)
Did you include complete proofs of all theoretical results?
[N/A]
- 3.
If you ran experiments (e.g. for benchmarks)…
(a)
Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See Section 4.
(b)
Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Refer to Appendix I and LABEL:apndx:hyper.
(c)
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] They are showed and explained in the captions.
(d)
Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Refer to Appendix I and Appendix J.
- 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets…
(a)
If your work uses existing assets, did you cite the creators?
[N/A]
(b)
Did you mention the license of the assets?
[N/A]
(c)
Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d)
Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A]
(e)
Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A]
- 5.
If you used crowdsourcing or conducted research with human subjects…
(a)
Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
(b)
Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A]
(c)
Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]

## Appendix A License

We distribute our dataset under CC BY-SA 4.0 license.

## Appendix B General Ethic Conduct and Potential Negative Societal Impact

To the best of our knowledge, our work does not present any direct potential negative societal impact, besides the carbon emissions produced by the GPU accelerators during datasets collection.

## Appendix C Downloading the Datasets

Both XLand-100B and XLand-Trivial-20B datasets hosted on public S3 bucket and freely available for everyone under CC BY-SA 4.0 Licence. We advise starting with Trivial dataset for debugging due to smaller size and faster downloading time.

Datasets can be downloaded with the curl (or any other similar) utility: curl -L -o xland-trivial-20b.hdf5 https://sc.link/A4rEW for smaller version and curl -L -o xland-100b.hdf5 https://sc.link/MoCvZ for main version.

## Appendix D What is Inside Dataset?

Both -Trivial and -100B dataset are HDF5 files holding the same structure. The dataset is grouped the following way:

where key_id is an ordinal number of a task in dataset, learning_history_id is a learning history number from 0 0 to $32$ and entity_name is one of the names mentioned in Table 3.

NB! Do not confuse key_id with the task ID, which should be accessed via

**Table 3: Data description**
| Name | Type | Shape | Description |
| --- | --- | --- | --- |
| states | np.uint8 | (5, 5) | $s_{t}$, colors and tiles from agent’s POV |
| actions | np.uint8 | scalar | $a_{t}$, from 0 0 to NUM_ACTIONS |
| rewards | np.float16 | scalar | $r_{t}$, which agent recieved at timestep $t$ |
| dones | np.bool | scalar | $d_{t}$, terminated or truncated episode flag |
| expert_actions | np.uint8 | scalar | same as $a_{t}$ but from a generating policy |

## Appendix E Hyperparameters

Table: Table 4: DPT Hyperparameters

**Table 5: AD Hyperparameters**
| Hyperparameter | Value |
| --- | --- |
| Embedding Dim. | 64 |
| Number of Layers | 8 |
| Number of Heads | 8 |
| Feedforward Dim. | 512 |
| Layernorm Placement | Pre-norm |
| Embedding Dropout | 0.1 |
| Batch size | [256, 128, 64] |
| Sequence Length | [1024, 2048, 4096] |
| Optimizer | Adam |
| Betas | (0.9, 0.99) |
| Learning Rate | 1e-3 |
| Learning Rate Schedule | CosineLR |
| Warmup Steps | 500 |
| # Parameters | 25 M |

**Table 6: PPO hyperparameters used in multi-task pre-training from Section 4.2.**
| Hyperparameter | Value |
| --- | --- |
| env_id | XLand-MiniGrid-R1-13x13 |
| benchmark_id | medium-1m |
| use_bf16 | True |
| pretrain_multitask | True |
| context_emb_dim | 16 |
| context_hidden_dim | 64 |
| context_dropout | 0.0 |
| obs_emb_dim | 16 |
| action_emb_dim | 16 |
| rnn_hidden_dim | 1024 |
| rnn_num_layers | 1 |
| head_hidden_dim | 256 |
| num_envs | 65536 |
| num_steps | 256 |
| update_epochs | 1 |
| num_minibatches | 64 |
| total_timesteps | 25,000,000,000 |
| optimizer | Adam |
| decay_lr | True |
| lr | 0.0005 |
| clip_eps | 0.2 |
| gamma | 0.995 |
| gae_lambda | 0.999 |
| ent_coef | 0.001 |
| vf_coef | 0.5 |
| max_grad_norm | 0.5 |
| eval_episodes | 256 |
| eval_seed | 42 |
| train_seed | 42 |

**Table 7: PPO hyperparameters used in single-task fine-tuning from Section 4.2.**
| Hyperparameter | Value |
| --- | --- |
| env_id | XLand-MiniGrid-R1-13x13 |
| benchmark_id | medium-1m |
| use_bf16 | True |
| pretrain_multitask | False |
| context_emb_dim | 16 |
| context_hidden_dim | 64 |
| context_dropout | 0.0 |
| obs_emb_dim | 16 |
| action_emb_dim | 16 |
| rnn_hidden_dim | 1024 |
| rnn_num_layers | 1 |
| head_hidden_dim | 256 |
| num_envs | 8192 |
| num_steps | 256 |
| update_epochs | 1 |
| num_minibatches | 8 |
| total_timesteps | 1,000,000,000 |
| optimizer | Adam |
| decay_lr | True |
| lr | 0.0005 |
| clip_eps | 0.2 |
| gamma | 0.995 |
| gae_lambda | 0.999 |
| ent_coef | 0.001 |
| vf_coef | 0.5 |
| max_grad_norm | 0.5 |
| eval_episodes | 256 |
| eval_seed | 42 |
| train_seed | 42 |

## Appendix F Compression Chunk Size Tuning

**Table 8: Throughput with PyTorch dataloader with different HDF5 compression chunk size settings. We used 2048 sequence length, 64 batch size and 8 workers. Chunking was applied along the learning history dimension.**
| Compression | Chunk size | Throughput |
| --- | --- | --- |
| None | None | 1,549,619 |
| gzip | None | 173,895 |
| gzip | 256 | 423,513 |
| gzip | 512 | 549,397 |
| gzip | 1024 | 666,152 |
| gzip | 2048 | 768,706 |
| gzip | 4096 | 749,851 |
| gzip | 8192 | 737,646 |

## Appendix G Additional Figures of Data Collection

Figure: Figure 9: Single-task evaluation curves on 256 tasks for policies trained from scratch or fine-tuned from multi-task pre-trained checkpoints.
Refer to caption: x9.png

## Appendix H Additional Figures of Data Evaluation

Figure: Figure 12: Return averaged over all learning histories in the final XLand-Trivial-20B dataset.
Refer to caption: x12.png

## Appendix I AD Implementation

We implement AD following the original paper of Laskin et al. (). To optimize the speed of training and inference we use FlashAttention-2 (Dao, ) with KV-caching and ALiBi positional embeddings (Press et al., ). We also concatenate observations, actions and rewards along the embedding dimensions, as it reduces the context size by the factor of three. We also use DeepSpeed (Rasley et al., ) to enable distributed training. The total number of parameters of our model is 25 M.

The approximate time of training for single epoch on a -100B dataset and evaluation on $1024$ tasks on 8 H100 GPUs is shown in the Table 9. The computations were done on an internal cluster.

The hyperparameters was copied from (Laskin et al., ) except for the size of the network, it was scaled up to 25 M. The exact hyperparameters can be found in Table 5.

We also show training logs for both datasets. Trivial: [wandb](https://wandb.ai/dunnolab/xminigrid-datasets?nw=96pyqrtxwuu); medium: [wandb](https://wandb.ai/dunnolab/xminigrid-datasets?nw=5u8wa1myqlw)

**Table 9: Time for training and evaluation.**
| sequence length | 1024 | 2048 | 4096 |
| --- | --- | --- | --- |
| train | 15 hrs | 10 hrs | 11.5 hrs |
| eval | 20 min | 30 min | 50 min |

## Appendix J DPT Implementation

Our implementation is based on the original one (Lee et al., ). Compared to AD implementation (Laskin et al., ), during training phase, the model context is generated with decorrelated dataset transitions to increase data diversity and model robustness: given a query observation and a respective expert action for it, an in-context dataset is provided by random interactions within the same ruleset. During evaluation phase, multi-episodic contextual buffer consists only of previous episodes and updates after the current one ends. The intuition behind this approach is the observed policy during any given episode is fixed, so it is a lot easier to analyze this policy than a dynamically changing one while it is executing.

Both AD and DPT shares the same observation encoder and transformer block, except there is no positional encoding in DPT, as stated in (Lee et al., ).

The training consisted of $3$ epochs due to computational and time limitations as $1$ epoch approximately lasted $12$ hours, while the evaluation on $1024$ rulesets on $500$ episodes could take from $5$ to $21$ hours, depending on the model’s sequence length. All experiments ran on $8$ A100 GPUs. The computations were done on an internal cluster.

The hyperparameters was copied from (Lee et al., ) except for the size of the network, it was made up to 25 M, and sequence length, it was increased due to complexity of the tasks. The exact hyperparameters can be found in Table 4.

We also show DPT training logs. Trivial: [wandb](https://wandb.ai/dunnolab/xminigrid-datasets?nw=smrpeqzlu6a); medium: [wandb](https://wandb.ai/dunnolab/xminigrid-datasets?nw=pfg0umx9c5d)

### J.1 Dark Key-To-Door Realization

We additionally demonstrate the inability of DPT to learn in-context in Partially Observable MDP (POMDP) on the example of a toy Dark Key-To-Door environment (Laskin et al., ). The agent is required to find an invisible key and then open an invisible door. The reward of $1$ is given when the agent first reaches the key and then the door. Note that the door cannot be reached until the key is found. This way Key-To-Door can be considered a POMDP. However, the environment can be reformulated as an MDP by providing additional boolean indicator of reaching a key in addition to the agent’s position. This way, algorithms that work only with MDPs are able to solve this environment.

Based on this fact, we learn two different Q-tables for both environments: with and without the key indicator. The learning histories of Q-Learning algorithm are stored together with optimal actions computed via the oracle.

For clarity, we call DPT training and evaluation Markovian when "reached key" indicator is provided for every state. We trained DPT on Key-To-Door for $150,000$ updates in Markovian and non-Markovian setups to show the difference in performance. As it can be seen in Figure 15, in the Markovian case the model converges to the optimal return, finding both a key and a door. In the latter case, the model reaches a plateau reward of $1$ which means it finds a key. As we empirically observe, without the indicator DPT reaches only a suboptimal return. We believe it happens due to DPT inability to reason whether a key was already found from the random context. Without this knowledge, it is impossible for the agent to know whether it should search for the key or for the door.

We also show logs for Key-To-Door Experiments: Markovian: [wandb](https://wandb.ai/dunnolab/xminigrid-datasets?nw=jv26obunfp); non-Markovian: [wandb](https://wandb.ai/dunnolab/xminigrid-datasets?nw=gmm7b681xea)

**Table 10: Q-Learning Hyperparameters**
| Hyperparameter | Value |
| --- | --- |
| Num. Train Goals | 2424 |
| Num. Histories | 5000 |
| Num. Updates | 50,000 |
| Learning Rate | 3e-4 |

Figure: Figure 15: The performance of DPT on Key-To-Door environment. From left to right: first two plots indicate returns on $20$ train and $50$ test goals, respectively, for Key-To-Door as MDP by providing additional indicator of reaching the key to the model. Second two plots indicate returns with the model that has no access to the fact of reaching a key. The results are averaged across four seeds
Refer to caption: x15.png

## Appendix K Additional Figures of AD Performance

Figure: Figure 16: AD performance on different complexities of rulesets. AD is evaluated on $1024$ training tasks from -100B. Sequence length is $1024$.
Refer to caption: x16.png

Figure: Figure 17: AD’s performance on different task complexities, a full variant of Figure 7.
Refer to caption: x17.png

## Appendix L Additional Figures of DPT Performance

Figure: Figure 18: DPT performance on both datasets. Evaluation parameters are the same as in Figure 8.
Refer to caption: x18.png

Figure: Figure 19: DPT performance on different complexities of rulesets. DPT is evaluated on $1024$ training tasks from -100B. Sequence length is $1024$.
Refer to caption: x19.png

Figure: Figure 20: Comparison of the learning histories in -100B dataset vs. DPT performance on the same training tasks. DPT is not able to solve simple tasks and there is no observation in-context learning emerges, model’s performance also degrades as the rulesets get harder. The context length of the model is $1024$. The evaluation parameters, except training tasks, are the same as in Figure 8.
Refer to caption: x20.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Line 350

> In contrast to AD, which predicts next actions from the trajectory itself, DPT-like methods require access to optimal actions on each transition for prediction. However, for the most nontrivial or real-world problems, obtaining true optimal actions in large numbers is unlikely to be possible.

I am uncertain about this line. As shown by [1], one can use a fully finetuned RL policy for a task, that achieves good performance, to collect “expert” data. Why is that not the case for this paper’s tasks?

# References

[1] D4RL: Datasets for Deep Data-Driven Reinforcement Learning, https://arxiv.org/abs/2004.07219
soundness: 3
strengths: 1. The paper is a dataset paper, and I think it clearly motivates the need for such a dataset in the open-source community. In that way, it is a strong work, since the release of a dataset/hard enough benchmark can truly boost the research productivity in this important field. The release of the ImageNet benchmark accelerated the pace at which computer vision grew as a field, and a well-designed benchmark can do that for in-context RL as well, hence I support the acceptance of this paper.
2. The paper is nicely structured and well-written, giving insights and reasoning to the community.
3. As the authors claim, researchers in this field often need to generate their own data. This results in lack of reproducibility and wastage of computational resources. This paper can help mitigate some of those challenges.
4. The authors have spent significant effort in trying to compress the dataset while maintaining throughput for loading the dataset from the compressed version. This effort is highly appreciated!
summary: The paper introduces a dataset, XLand-100B, for training and testing in-context RL algorithms. In-context RL is the problem of predicting how to act in a task, given trajectories from the task, without any updates to the model weights. There is prior research in this field, but future progress is hindered by lack of open-source datasets/benchmarks where algorithms can be tested. The paper introduces a large scale dataset containing 100B transitions, 2.5B episodes from 30,000 (in effect, around 29,000 after filtering). The paper also tests two current algorithms, AD and DPT, to show that more progress/research is needed in the field.
weaknesses: 1. Despite the enormous size of the dataset, it only contains one type of environment, based on the Mini-Grid set of tasks. While this is a starting point for research in this direction, adding tasks from more realistic domains, like robotics/other tasks that strictly require RL, would be appreciated. This is my main concern for not giving a higher score.
2. Adding more visualizations of the tasks, trajectories, etc in the appendix might be helpful to understand the precise nature of the tasks, what does # rule 3 → # rule 9 actually mean for the hardness of the task, etc.
3. The idea of collecting such a dataset is not novel, despite the importance of doing such work that can empower future novel research in this direction.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. Why did you choose to use recurrent PPO with GRU for data generation over using a transformer (as in the two ICRL experiments)?
2. Evaluation is performed for 500 episodes and the score averaged across 1024 tasks. Am I correctly understanding this as performing 1024 separate single-task evaluations (each 500 episodes long) and then averaging scores across tasks? I.e., the transformer context never contains data from multiple tasks.
soundness: 3
strengths: Overall, the paper makes a substantial contribution to in-context RL research (and RL research in general). The authors argue clearly what distinguishes their work from the many existing datasets and benchmarks, and as such, the design of the benchmark is very well-motivated. The paper is also well-written and the presentation is good. 

Some strengths worth highlighting:

1. I particularly appreciate the level of detail regarding data collection, the dataset format, and other implementation/engineering methodology.
2. The authors made significant effort to make the benchmark easy to use (e.g. tuning the data compression to balance dataset size and sampling speed, and providing estimated optimal actions which are needed by methods such as DPT).
3. Including the lighter XLand-Trivial-20B dataset should make the benchmark significantly more accessible to researchers with limited resources.

I would be happy to increase the score, pending clarification on some of my concerns below.
summary: This paper introduces new datasets “XLand-100B” and “XLand-Trivial-20B” intended for research on in-context reinforcement learning (ICRL). XLand-100B consists of 100B transitions spanning 30k different tasks from the XLand-MiniGrid environment. 

The authors argue that research in ICRL has been hampered by the lack of large & diverse datasets that are compatible with ICRL, as well as by the significant compute requirements in this domain, and discuss how their contribution aims to make this research more accessible, in particular for researchers with limited resources. They emphasize that compared to prior work (NetHack, Procgen, GATO,...), their dataset is larger, more diverse, and better suited for ICRL research since it contains full learning histories instead of unordered or solely optimal/expert trajectories.

The data was collected by first pretraining a PPO multi-task policy on 65k tasks for 25B transitions and then finetuning it on a subset of 30k tasks. The dataset was only collected during the latter phase (for XLand-100B). The authors show that the pretrained policy performs well in the zero-shot setting, and that using it as a starting point for finetuning yields a very large performance boost compared to training a single-task policy from scratch (on hard tasks).

Finally, the authors benchmark Algorithm Distillation (AD) and Decision-Pretrained Transformer (DPT) on their dataset. They show promising results with AD, which displays an emergent ICL ability. However, they note that they were unable to show the same for DPT. The authors hypothesize that DPT performs poorly in POMDP settings as its lack of positional encoding prevents it from leveraging useful historical information from the current episode.
weaknesses: I don’t believe there to be any significant weaknesses relating to the dataset (the main contribution) itself, apart from the limitations mentioned in the paper.

Regarding the AD & DPT baselines:

1. My understanding is that DPT, in contrast to AD, is unable to reason in POMDP settings because its transformer doesn’t include a positional encoding. Have you considered the modification of applying a positional encoding to just the transitions from the current (ongoing) episode? This wouldn’t significantly change the DPT formulation but would remove the limitation of not being able to use historical information from the current episode. The DPT experiment as it stands does not seem like a fair comparison to AD as its poor performance likely doesn’t come from the method itself but the observation space (lack of historical context).
2. It is stated that evaluation is performed for 512 episodes with a context length of k=1024/2048/4096. My understanding is that the context is initially empty and that at the end of evaluation, the context contains the most recent k transitions (potentially from multiple episodes). It’s unclear how often the context actually includes data from multiple episodes. What is the median episode length?
3. Have you performed experiments (with AD) where the context only contains data from the current episode? From Appendix H, it is clear that knowledge of the current episode history is essential for solving some of the partially-observable tasks (e.g. knowing whether you have already picked up a key). Including data from previous episodes could make it more difficult since the policy would need to distinguish whether a key pick-up happened in the current or previous episode.
4. Do you have any insight on why longer context lengths perform strictly worse with AD? Could it be related to the issue mentioned in my previous question (3.)?

Minor:

5. Figure 8 is a bit unclear as there is no label/legend for the colors. The 1024 context length could be confused with the 1024 unseen tasks used for evaluation.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: - Could you clarify the rationale for choosing PPO as the sole algorithm for generating learning histories? How might the dataset be affected if other RL methods with different exploration and learning characteristics were incorporated?

- How do you ensure that filtered tasks with low returns are adequate proxies for high-quality demonstrations in ICRL?

- How might labeling inaccuracies impact the performance of models relying on these demonstrations?
soundness: 3
strengths: - The paper is well written.
- There do not seem to be any major issues with the method or the data set.
- Authors appear to have taken reproducibility seriously.
- The technical presentation of the data set, including relevant specifications, is comprehensive and clear.
- The work addresses an important gap that is hindering progress in an area of reinforcement learning of growing interest.
summary: This paper introduces XLand-100B, a large-scale data set for in-context reinforcement learning containing 100B transitions across 30,000 different tasks. The dataset includes complete learning histories of agents trained with PPO in the XLand-MiniGrid environment. The authors evaluate two common in-context RL algorithms on the dataset, showing the limitations of these methods with complex tasks. The work aims to democratize in-context RL research by providing a standardized, large-scale benchmark.
weaknesses: The paper has no critical weaknesses. However, for the sake of constructive academic discussion, I include some drawbacks below:

- The dataset exclusively uses PPO to generate learning histories, without theoretical or empirical justification for why PPO is especially suited for ICRL compared to other RL algorithms. Given the diversity of RL methods—each with distinct exploration strategies, convergence behaviors, and learning dynamics—relying solely on PPO risks biasing the dataset toward a specific style of learning history. This may inadvertently limit the dataset’s utility for evaluating ICRL methods that could benefit from more varied demonstration patterns.
- The dataset’s approach to ensuring high-quality demonstrations is primarily based on filtering out tasks with low final returns, which does not fully address whether all included histories are genuinely informative or relevant for ICRL.
- The use of approximate expert actions for labeling could compromise action fidelity, particularly for complex tasks, which may impact models that rely on high-quality expert demonstrations.

These weaknesses may affect the extent to which this benchmark is a good proxy for in-context RL. While benchmarks are very important for research progress, misaligned benchmarks can be actively counterproductive. On balance, I do not expect this to be a major problem, but it would be desirable if the authors strengthened the paper on this respect.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - Can you clarify whether the dataset is collected on a 5x5 grid or on the grid visualized in Figure 1? If so, it is important to update Figure 1 to show the actual environment. 
- In Figure 5, it seems that there is not much difference in learning performance from 1 to 7 rules, but a larger jump to 9 rules. Can you clarify why this is the case? 
- Why does AD not exhibit in-context improvement in Figure 7 (only slightly for 1 rule)? Can you provide empirical evidence if this is a property of the data or the algorithm?
- In Figure 8, AD performance decreases with a larger context length. Why is that? Can you clarify how many episodes the context comprises?
soundness: 2
strengths: The paper highlights an important limitation of current ICRL research, which is often focused on simple benchmarks such as gridworlds and small datasets. To this end, the authors generate a large-scale dataset and make it publicly available. The dataset contains episodes from 28K tasks, which is larger than other ICRL datasets. The authors make good design decisions in terms of their storage format (hdf5, compression levels, chunk size) and data collection strategy. As such, the dataset may be useful to the broader ICRL community.
summary: In recent years, there have been some efforts on in-context reinforcement learning (ICRL). According to the authors, most ICRL research has been done on simple environments, and the lack of challenging benchmarks limits progress of the field. To address this, the paper proposes XLand-100B, a large scale dataset for in-context learning in gridworlds, comprising 30K tasks and 2.5B episodes.
weaknesses: While the paper highlights, that current ICRL research is focusing on simple environments, the gridworld environment they consider also seems simplistic. The generated dataset seems to be specific for a 5x5 single-room grid (not the multi-room in Figure 1? See questions). Generally, the environment seems similar to Dark-Room/Key-Door, and can be seen as a more advanced grid-world (with more than 2 objects) to test ICRL in toy environments. Therefore, it is unclear how well future findings on this dataset would transfer to other settings like robotics. 

To the best of our understanding, there is a maximum of 9 rules in this environment, which correspond to reaching different objects in order. It would help the paper to provide a visual illustration of those rules for the reader to get an understanding of how diverse those tasks are (especially as it is a 5x5 grid). Due to those 9 rules, many different tasks can be produced. However, while there are 28K tasks in the generated dataset, it is unclear how much overlap or diversity there is between tasks. Therefore, it is possible that ICRL agents do not benefit from training on additional tasks. We suggest conducting an ablation study, in which agents are trained on 100/1000/10000/all tasks to evaluate the effect on ICRL abilities. This is currently missing.

While the authors release a large dataset with lots of tasks, from the perspective of a user it is unclear where to start. This is due to the lack of a clear benchmark split. The authors should clarify on what tasks to train on and what to evaluate on. This could be by number of pre-training tasks (as mentioned above) or number of rules etc. This point is important and needs significant revision and empirical evaluation of the considered methods.

In Figures 12 and 13 the learning curves (over 8K and 15K episodes) across all tasks are provided. The agents reach optimal performance quickly, only after a coupled hundred steps. Therefore, there is little learning progress in the remaining 14K updates. This may be a reason why AD and DPT do not exhibit meaningful in-context improvement during evaluation (Figure 7, 17, 20). Generally, there seems to be little in-context improvement by any ICRL method, which is different from previous ICRL works. This raises the question, whether the problem lies in the algorithms or the datasets, and should be discussed in the paper.

Meta-review: The authors present a new dataset for training and testing in-context RL algorithms that goes beyond existing tasks in simple environments and on small-scale datasets. They also provide results from running common baselines, showing that existing methods struggle to generalize to novel and diverse tasks. 

Reviewers found the paper well-written and and containing a significant amount of detail about how the dataset is collected. The benchmark seems easy to use and made efforts to ensure reproducibility. However, the reviewers also raised concerns that the contribution is limited in the sense that this benchmark is built on only one type of environment, Mini-Grid and all baselines implemented using the same underlying RL algorithm, PPO. Finally, there are concerns that no ICRL method seems to perform well on the benchmark, raising the question if the issue is with ICRL methods or the benchmark itself.

Because reviewers agree that the benchmark is thoughtfully designed, well-motivated, and timely for research in ICRL, I vote to accept this paper.

Decision: Accept (Poster)