XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning
Authors: Alexander Nikulin
Ilya Zisman
Alexey Zemtsov
Vladislav Kurenkov
Venue: ICLR 2025 Poster (2025)
Keywords: in-context reinforcement learning
reinforcement learning
xland
minigrid
Primary area: reinforcement learning
TL;DR: We present XLand-100B, a large-scale dataset for in-context reinforcement learning based on the XLand-MiniGrid environment

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$ 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. Along with the dataset, we provide the utilities to reproduce or expand it even further. We also benchmark common in-context RL baselines and show that they struggle to generalize to novel and diverse tasks. 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.

Published as a conference paper at ICLR 2025
XLAND- 1 100B: A LARGE-SCALE MULTI-TASH DATASET FOR IN-CONTEXT REINFORCEMENT LEARNING
Alexander Nikulin* AIRI, MIPT
Ilya Zisman' AIRI, Skoltech
Alexey Zemtsov" NUST MISIS, T-Tech
Vladislav Kurenkov AIRI, Innopolis University
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 reinforce- ment learning is experiencing 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 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 different tasks, covering 100B transitions and 2.5B episodes. It took 50, 000 GPU hours 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. We also benchmark common in-context RI. baselines and show that they struggle to generalize to novel and diverse tasks. With this substantial effort, we aim to democratize research the rapidly growing field of in-context reinforcement learning and provide a solid foundation for further scaling.
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. 2020). However, it was quickly discovered that small transformers are also capable of in-context learning (Kirsch et al.] [022] [Von Oswald
et al. (023)) and even many non transfrrffrrrr models have such abilities (Bhattamishra al.) 2023
Akyürek et al., 2024] Park et al., 2024] Grazzi et al. 2024) Vladymyrov et al. [0024] Tong & evan
2024). More importantly, driven by properties of data, rather than the architecture (Chan et al. 2022
Gu et al. 2023). in-context learning is not specific to language modeling and has been found in other
domains, e.g. image generation (Bai et al. 2023 Najdenkoska al. 2023 Doveh et al. 2024 Tian
cet al. 2024).
However, despite the rapid adoption of the transformer architecture in reinforcement learning (RL) after the release of Decision Transformer (DT) (Chen et al., 2021 Hu et al. 2022 Agarwal et al.
2023
Li et al. 2023), models with in-context learning capabilities appeared only recently. This delay caused by number of reasons. Firstly, to transition from in neiwegttts to in co-tontext learning, a model should be trained on tens of thousands of unique tasks (Kirsch et al. (022). Unfortunately, even the largest RL datasets currently contain only hundreds of tasks (Padalkar et al.) 2023). 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.) 2022 Lee et al. 2023 Shi et al. 2024)
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.) 2022 Lee et al.) [2023] Norman & Clune
2023, Kirsch et al. 2023; Sinii al. 2023 Zisman al. 2023 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
"Equal contribution.
1 "Corrsponndecce to: nikulin @airi.net Work done by dunnolab.ai started T-Tech.

---

Published as a conference paper at ICLR 2025
Table 1: Comparison with other RL datasets.
Data
# tasks
# transitions
# episodes
Size
Open-Source
Enables ICRL
XLand-100B (ours)
28,876
100B
2.5B
320GB
JAT
157
323M
N/A
ITB
X
GATO
596
1.5T
63M
N/A
Open X- X-Embodiment
527
N/A
2.4M
9ТB
AlphaStar Unplugged
21B
2.8M
N/A
X
NetHack
3.5B
110K
97 GB
X
X
D4RI.
11
N/A
40 M
4.8GB
X
V-D4RL
4
2.4M
N/A
17GB
DSRL.
50
N/A
2500
N/A
RL Unplugged
90
8OM
N/A
54.ITB
X
Procgen
16
37M
N/A
50UGB
X
with high diversity and difficulty, which essential for real-world applications. Because of this, the development of in-context RL is currently hindered by these factors. We believe is crucial to address these barriers, given the essential role in-context learning the path to foundation models
and truly generalist agents (Team et al. 2021 2023 Kirsch al. 2023 2024).
Lu et al. 2024 Liu et al.)
We release XLand-100B, scale dataset for in-context RL based on the XLand MiniGrid (Nikulin et al. 2023) environment. It contains complete learning histories for nearly 30, 000 different tasks, covering T00B transitions and 2.5B episodes. It took 50, 000 GPU hours 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 smaller and simpler version for faster experimentation. as well utilities 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 complex tasks (see Section 5).
2
BACKGROUND
2.1 IN-CONTEXT REINFORCEMENT LEARNING
Multiple methods for in-context RL have come out, cach offering a different way of training and organising the context (Laskin
2022 Lee al. 2023 Mirchandani al., 2023 Liu & Abbeel
2023 Raparthy et al.) 2023 Shi et al. 2024). We focus on Algorithm Distillation (AD) (Laskin et al.
2022) and Decision- Pretraiindd Transformer (DPT) (Lee al. 2023). which we chose as the main
methods for our work due to their simplicity and generality.
Algorithm Distillation. AD (Laskin al. 2022) was one of the first show that in-context learning was possible in RI., and captures the details of many other more recent methods (Mirchandaní et al. 2023 Liu Abbeell 2023) [hhi et al., 2024) while remaining very simple. trains a transformer. any other sequence model, to autoregressively predict next actions given the history of previous interactions, i.c. 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.) 2021 I Janner et al.
2021 Lee al. 2022).
Prrcsiin-Prettaineed Transformer. DPT is an alternative approach inspired by the Bayesian inference approximation (Müller et al. 2021). Unlike AD, it trains transformer predict the optimal action
2

---

Published as a conference paper at ICLR 2025
for a query state given 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 dataset of learning histories.
In addition, the theoretical analyses of AD and DPT methods (Lin et al. 2023 Wang et al. 2024) showed that they 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
Starting from the seminal work of Wang et al. (2016): Duan et al. (2016); |Finn et al. (2017) on meta-RI,, much of the subsequent
work (Zintgraf et al. 2019 Melo 2022; Grigsby et al. 2023 Lu et al. 2024: Shala et al. 2024 Beck et al. 2024) has focused 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. (2023) released XLand-MiniGrid, APUU-accllerate environment and million-task benchmarks that significantly lowered the entry barrier for meta-RL research. We will describe shortly here.
Figure 1: Visualization of generic "Land-MiniGrie envi- ronment. Grid layout should be selected advance, while the positions of the objects are randomized on each reset. For the dataset we use simpler lay- out with just one room, see Ap- pendix M
Environment. (Land-MiniGria is a complete rewrite of MiniGrid (Chevalie--Boiseer al., 2023) in JAX (Bradbury et al.] 2018), incorporating notion of rules and goals from XLand (Team et al.) 2023). Leveraging JAX, it can run on GPU or TPU accelerators at millions of steps per seconds. At its grid-world environment with simple underlying dynamics, partial observability and sparse rewards. The action space is simple, sisting mainly navigation and interaction with game objects, such opening door picking and placing items. Observations 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 places near each other, they both disappear and one new object is placed. Goals are similar, except they only validate some 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. (2023).
it is coal-oriented
Benchmarks. Along with the environment itself, Nikulin et al. 2023 released a tool for the procedural generation a vast number of unique tasks with varying levels of difficulty. Each task represented by 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 high, each with one million unique tasks. For this work, we chose medium um as a middle ground between yet unsolved h gh and less challenging small benchmarks, We also trivi al for smaller and simpler dataset version (see Section
provided: trivial, small, medi um,
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 actual learning histories (Laskin et al. 2022] or their approximations (Zisman et al.) 2023), 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. 2022 Kirsch et al. 2023). Another approach is to learn from optimal actions as proposed by Lee et al. (2023), 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. 2022). That is. for a simple task to find two squares on 9 x grid, agent needs to around 2000 different combinations of goals start adapting for unseen
3

---

Published as a conference paper at ICLR 2025
Table 2: Descriptive statistics of XLand datasets.
Dataset
XLand-Trivial-008
XLand-100B
Episodes
868,805,556
2,500, 152,898
Transitions
19,496,960,008
112,598,843.399
History length
60,928
121,856
Num tasks
10,000
28,876
Max task rules
9
Observation shape
(5, 5)
(5,
Num actions
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
locations (Laskin et al. 2022). 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 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 map (Laskin et al.) 2022 Lee et al. 2023 Zisman et al. (023) or to apply forces to actuators in order walk robot (Kirsch al. 2023). 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 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. 2020). They were initially proposed for offline RL. containing simple tasks with flat structure, e.g. perform locomotion with different robots (Lu et al. 2023) or path finding in maze. Some of them also contain data from robotic manipulators (Fu et al.
2020 Rafailov et al. 2023), or even Atari frames (Gulcehre et al. 2021). Other datasets collect data for more sophisticated environments, such as NetHack Learning Environment (Hambro et al.) 2023 Kurenkoy et al.) 2024) or ProcGen (Cobbe et al. 2019 Mediratta et al.) 2024). However, the aforementioned datasets offer 100 different tasks with fixed policy (except for the -replas aa datasets, which have limited coverage of various policies). This limitation makes it difficult for in-context RL emerge from such data. To overcome this pitfall, we collected almost 30, 000 tasks with deep ruleset structure, that are challenging problem to solve.
Large-Scale Supervised Pretraining. Recent progress in generalist agents, which can solve multitude of environments, has been made possible thanks to large datasets. GATO dataset (Reed al. 2022). 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 RI. The open-sourced analogue, the JAT dataset (Gallouédec al. 2024), is smaller size with 157 tasks and 300 million transitions, but it provides comparable performance
most of the benchmarks. Both datasets contain expert RL demonstrations from BabyAI (Hui et al. 2020), Atari games (Bellemare et al. 2013), World Yuu et al. 2020) and more.
Another large dataset, Open E-Eooddiment Padalkar et al. 2023 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 human experts. Despite the large quantity transitions in these datasets, they do not contain learning histories with improving policies, making their application for in-context
4

---

Published as a conference paper at ICLR 2025
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 for in-context RL is AlphaStar Unplugged (Mathieu ct al. 2023). 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 steady policy improvement, thus enabling the in-context RL ability. For more details on the datasets, refer to Table 1
4 XLAND-100B DATASET
We present XLand-100B, large dataset for in-context RI., and its smaller and simpler version -and-Trivial-2000 for faster experimentation. Together they contain about 3.5B episodes, 130B transitions and 40,000 unique tasks (see for detailed statistics). Datasets are hosted on public S3 bucket and will be freely available for everyone under CC BY-SA 4.0 licence. Next, we describe the data format, collection and evaluation.
4.1 DATA FORMAT
Storage format. We chose store the datasets in HDF: file format based on its popularity and convenience. It allows 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 ~600G for our main dataset. Using a little trick described later, were able to reduce the size more just 326 GB (see Table 2 However, naive of compression can dramatically increase batch sampling time and slow overall training time down. We tuned HDF5 cache chunk size specifically maximize sampling throughput for large sequence lengths. After tuning, we achieved fourfold speedup over the naive compression, and were only two times slower compared of 15, we see this as good trade-off, which increases the overall affordability of the dataset. See Appendix for throughput benchmarks.
compression. Given that we reduced the dataset size by factor
Data and metadata format. observations, actions, rewards and dones encountered during agent training in separate HDF5 groups with unique IDs per history (see Appendix B for more details). To be compatible with DPT-like methods (Lee et al., 2023), we also store expert actions for each transition (see Section 4.2). Unlike popular formats such as RLDSSFann and Minar 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 collect complete learning histories, i.c. for each history we store all
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 di vmod function. In addition, for each history
store the "Land-MiniGin 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 set up the environment for evaluation.
4.2 DATA COLLECTION
At a high level, data collection was organized 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 collect the full dataset. Collecting dataset of this size for any other environment suitable for in-context RL, e.g. Meta-Wast World (Yu et al. 2020), would take much longer, which is unlikely be feasible for most practitioners (Nikulin et al.
ds
https///gitthbb..om/aaaaaaa
o nari
5

---

Published as a conference paper at ICLR 2025
1.0
From Scratch Pretrained
Intial
0.4
Return
0.3
Return
0.5
4k
0.2
2k
0.0
0.1
0
12.58
25B
0.5B
1B
0
4 5 7 Number of rules
2
8
Transitions
Transitions
Figure 2: Evaluation return for multi-task goal-conditioned receurent PPO pretraining on 65k tasks. Pretrained agent was further used a starting point for single-task finetuning dur- ing dataset collection.
Figure 3: Single-task evalua- tion curves on 36 hard tasks for policies trained from scratch or fine-tuned from multi-task pre- trained checkpoints. See Ap- pendix for curves on tasks of all difficulty.
Figure 4: Distribution of the tasks by difficulty sampled initially and in the resulting dataset. To ensure the quality, we filtered tasks where the fi- nal return was below 0.3 or the data was corrupted due to some errors during training.
2023). 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 O
Base algorithm. For our datasets we chose PPO (Schulman et al. 2017), due its high scalability and compatibility with massively parallel environments. We ported the implementation from recurrent PPO provided by (Nikulin et al.) (023), 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. 2014) for memory. it showed satisfactory performance on preliminary experiments. Since the base algorithm was implemented in JAX (Bradbury et al. 2018), were able 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-effiiient algorithm, it was still necessary train it on billions of transitions. Fortunately, trained it on thousands of environments in parallel, learning history for each particular environment quite short, i.e. around 120k transitions.
Pretraining. For our main XLand-100B dataset we uniformly sampled tasks from medi um- 1 benchmark from -Laan-MMnnGria 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 multi-task cask-conditioned manner and also use simpler grid layout with just one room (see Appendix M). 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 additional input to the agent. After that, train the agent on 65k tasks simultancously for 25B transitions. As Figure showw, such an agent learns generalize zero-shot new tasks quite well, although do not aim push it to the limit, as zero-shot generalization does not produce a smooth learning history during fine-tuning. We skip this stage for benchmark.
dataset due the simplicity of the tasks in trivial-1m
Finetuning. This is 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 IB transitions on 30k uniformly sampled tasks from medium-1n 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 controlled manner. For example, we can train AD (Laskin et al., 2022) on all 30k tasks using one history per task or ~ 900 tasks using all histories per task to equalize the number of training tokens. For the -LLan--Tiivil--000
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 are able to show strong performance even on the hardest.

---

Published as a conference paper at ICLR 2025
tasks, increasing the diversity and coverage of the resulting dataset. For the same results on tasks of all levels of difficulty, see Appendix(]
Postprocessing. After fine-tuning, it necessary to additionally label the transitions with the expert actions to support DPT-like
1.00
0.75
methods (Lee et al. 2023). To do this, walk through the entire learning history with the final policy, starting from the initial hidden state for the RNN. We evaluate the validity of such 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 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 Ik 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
0.50
0.25
r
Traee rules
0.0
0.5
1.0
Normaliized #episodes
Figure 5: Learning histories for the XLand- 100B dataset separated by number rules. For visual clarity, we show only sample of the possible number of rules and normal- ize the number of episodes, as they may vary considerably.
4.3 DATA EVALUATION
In this section we validate that the resulting XLand-1009 dataset actually fulfils the two most important requirements for in-context RI., namely 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 Soandd-riiial-200 in the Appendix E
Improvement history. In the Figure from the learning histories separated by the number of rules. 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 al. 2023 Shi et al. 2024). For learning history averaged over full dataset see Appendix
we show the averaged return
1.0
Actions
0.7
0.6
0.5
0k
25k
50k 75k 100k 125k
Transitions
Figure Agreement between actions predicted by the expert and the actual actions in the learning history. We use final PPO policy as an expert for ac- tions labeling. As can be seen agreement is growing towards the end of learning history.
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.] (2023) 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 labeling is adequate in our case and the expert at the end predicts actions close to what the policy did in reality 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.
5 EXPERIMENTS
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. 2022) and DPT (Lee et al. 2023) for experiments, the exact implementations details are in Appendix F and Appendix G Both methods were trained on (land-Triviad -20B and XLand- 100B with (512, 1024, 2048, 4096)
and (1024, 2048, 4096) context lengths respectively. We do not include 512 context length for the
7

---

Published as a conference paper at ICLR 2025
rules
# rules
riles
rules
1.0
1.00
Maze
20t3
bata
0.75
0.8
0.75
0.50
0.6
0.25
0.2
0.4
0.25
0.00
0.2
0.00
-0.2
100
300 500
200
400
150
200
300
500
200
400
500
Episodes
Eplaodes
Egisades
Esisedas
Figure 7: Comparison of the learning histories in -100B dataset vs. AD performance 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
-1 -0B0 dataset 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 both datasets. Figure g demonstrates the performance of the method for different context lengths. On -Trivial-208 dataset it is able show stable policy improvement from about 0.28 0.4 during the evaluation. For -100B the performance is similar, but the pace of improvement is faster. We hypothesise that happens due wider data-coverage, since the agent sees more complex tasks and is able to learn faster from them.
XLand-Trivial-222
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 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 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.
0.6
0.5
3124 leg
2048 sequien
Return
0.4
0.3
0.2
0.1
0.0
0
100
200
300
400
500
Episodes
XLand-100B
0.6
3524
sog
0.5
4198
DPT. Decision-Preerainndd Transformer (Lee al. 2023 is another method that exhibit in-context RL capabilities. However, in our experiments we were unable train it so that these abilities emerge. 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 investigation, we refer the reader to Appendix
Retur
0.4
0.3
Figure 18 in Appendix K
0.2
0.1
0.0
100
300 400 Episodes
Figure 8: AD performance on datasets for different se- quence lengths. Both datasets lead to the emergence of in- context ability. We report the average return on 1024 unseen tasks across 3 seeds.
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-MiniGris library (Nikulin et al. 2023). 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, there many learning histories that start from high reward. We hope to improve the RL baseline for data collection to avoid the need for multi-task pre-training in the future.
8

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
Legal Compliance and Licensing. The dataset will be released under the CC BY-SA 4.0 license, and the accompanying code will be available under the Apache 2.0 license. have ensured that our data and code releases comply with all relevant legal and ethical guidelines.
Data Privacy and Security. The dataset does not include any human subjects or personally identifi- able information. All data generated by artificial agents interacting within simulated environment. Consequently, there are no privacy or security concerns associated with the dataset.
Potential Misuse and Harmful Applications. The dataset and methods are intended advance research in in-context reinforcement learning. We do not anticipate any immediate risks of misuse or harmful applications arising from work. Nonetheless, encourage users apply the dataset responsibly and adhere to ethical standards in their research.
Environmental Impact. Collecting the dataset required substantial computational resources, totaling approximately 50,000 GPU hours. We acknowledge the environmental footprint associated with large-scale computations. By openly releasing this dataset, we aim to minimize redundant data collection efforts by other researchers, potentially reducing the overall environmental impact in this field.
REPRODUCIBILITY STATEMENT
To ensure reproducibility share the important details throughout the paper and appendix. We explain how we pretrain, finetune models for collecting data and the post-processing step in Sec- tion 4.2 well discussing the dataset details (size, compression ratios, etc.) in Appendix [ and Appendix C For experiments section, we discuss the details of baselines implementation in Appendix [i and Appendix G We provide hyperparameters for collecting the dataset and for baselines in Appendix O
We also release the codebase with tools for creating and expanding the dataset in the following repository: Xland- moanl-miiigidd-aatten
REFERENCES
Pranav Agarwal, Aamer Abdul Rahman, Pierre-Luc St-Charles, Simon JD Prince, and Samira Ebrahimi Kahou. Transformers in reinforcement learning: a survey. arXiv preprint M.Xiv:2007.05979 2023.
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Arhitec- tures and algorithms. arXiv preprint arXiv:2401.12977 2024.
Yutong Bai, Xinyang Geng, Karttikeya Mangalam, Amir Bar, Alan Yuille, Trevor Darrell, Jitendra Malik, and Alexei Efros. Sequential modeling enables scalable learning for large vision models. arXiv preprint arXiv:2312.00788 2023.
Jacob Beck, Risto Vuorio, Zheng Xiong, and Shimon Whiteson. Recurrent hypernetoorks are surprisingly strong in meta-rl. Advances in Neural Information Processing Systems, 36, 2024.
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, jun 2013.
Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade, Understanding in-context learning in transformers and Ilms by learning to learn discrete functions. arXiv preprint arXiw:2310.03016, 2023.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
com/ goom/goollao /e/oolee/aa

---

Published as a conference paper at ICLR 2025
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang. Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 35: 15:18878-1888 18891, 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-15099 15097, 2021.
Maxime Bhoiselir-Boiivertt, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306.1383) 2023.
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint arXiv: /40: /299 2014.
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. arXiv preprint arXiv: 1912.0118, 2019.
Tri Dao. F1ashAttention-2:: Faster attention with better parallelism and work partitioning. In International Conference Learning Representations (ICLR), 2024.
Sivan Doveh, Shaked Perek, M Jehanzeb Mirza, Amit Alfassy, Assaf Arbelle, Shimon Ullman, and Leonid Karlinsky. Towards multimodal in-context learning for vision & language models. arXiv preprint arXiv:2403. 2024.
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 1.02779, 2016.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pp. 1126-1135. PMLR, 2017.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
Quentin Gallouédec, Edward Beeching, Clément Romac, and Emmanuel Dellandréa. Jack of All Trades, Master of Some, a Multi-Purpose Transformer Agent. arXiv preprint arXiv:2402.098.4 2024. URL httLhhtts:://// http:s://arriiii http:///////w....ooo/aaaba//aa wttp:///g.ba/....../../// 2402. 09844)
Riccardo Grazzi, Julien Siems, Simon Schrodi, Thomas Brox, and Frank Hutter. Is mamba capable of in-context learning? arXiv preprint arXiv:2402.03170, 2024.
Jake Grigsby, Linxi Fan, and Yuke Zhu. Amago: Scalable in-context reinforcement learning for adaptive agents. arXiv preprint 2000v2231
Yuxian Gu, Li Dong. Furu Wei, and Minlie Huang. Pre-training to learn in context. arXiv preprint arXiv:2305.09117, 2023.
Caglar Gulcehre, Ziyu Wang. Alexander Novikov, Tom Le Paine, Sergio Gomez Colmenarejo, Konrad Zolna, Rishabh Agarwal, Josh Merel, Daniel Mankowitz, Cosmin Paduraru, Gabriel Dulac-Arood, Jerry Li, Mohammad Norouzi, Matt Hoffman, Ofir Nachum, George Tucker, Nicolas Heess, and Nando de Freitas. RI unplugged: A suite of benchmarks for offline reinforcement learning, 2021.
Eric Hambro, Roberta Raileanu, Danielle Rothermel, Vegard Mella, Tim Rocktäschel, Heinrich Küttler, and Naila Murray. Dungeons and data: A large-scale nethack dataset, 2023.
Shengchao Hu, Li Shen, Ya Zhang, Yixin Chen, and Dacheng Tao. On transforming reinforcement learning by transformer: The development trajectory. arXiv preprint arXiv:22 141114, 2022.
David Yu- Tung Hui, Maxime Chevalier- Boisvvrr, Dzmitry Bahdanau, and Yoshua Bengio. Babyai 1.1, 2020.
10

---

Published as a conference paper at ICLR 2025
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. Advances in neural information processing systems, 34:1273-122 1286, 2021.
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in-context learning by meta-learning transformers. arXiv preprint arXi::2212.04455, 2022.
Louis Kirsch, James Harrison, C. Freeman, Jascha Sohl-Dickstein, and Jürgen Schmidhuber. Towards general-purpose in-context learning agents. NewrIPS 2023 Workshop on Generalization in Planning, 2023. URL https: hnttps::/openrewviee new.ewweeeeeeeeeeeee..ee......e. forumm 2tt.:eommmmmddeeee00dddd
Vladislav Kurenkov, Alexander Nikulin, Denis Tarasov, and Sergey Kolesnikov. Katakomba: Tools and benchmarks for data-driven nethack. Advances in Neural Information Processing Systems, 36, 2024.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:22 /4215, 2022.
Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. arXiv preprint arXiv:2306.14889 14892, 2023.
Kuang-Huce Hueii Lee, Ofir Nachum, Mengjiao Sherry Yang. Lisa Lee, Daniel Freeman, Sergio Guadar- rama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. Advances in Neural Information Processing Systems, 35:27921-27995 2022.
Wenzhe Li, Hao Luo, Zichuan Lin, Chongjie Zhang, Zongqing Lu, and Deheng Ye. A survey on transformers in reinforcement learning. arXiv preprint arXiv:2301.08 303.00044,, 2023.
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforce- ment learning via supervised pretraining. arXiv preprint arXiv:2310.08566 2023.
Hao Liu and Pieter Abbeel. Emergent agentic transformer from chain of hindsight experience. In International Conference on Machine Learning, pp. 213 223744. PMLR, 2023.
Xiaoqian Liu, Xingzhou Lou, Jianbin Jiao, and Junge Zhang. Position: Foundation agents as the paradigm shift for decision making. arXiv preprint arXiv:2405. /7009, 2024.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Cong Lu, Philip J. Ball, Tim G. Rudner, Jack Holder, Michael Osborne, and Yee Whye Teh. Challenges and opportunities in offline reinforcement learning from visual observations, 2023.
Michaël Mathieu, Sherjil Ozair, Srivatsan Srinivasan, Caglar Gulcehre, Shangtong Zhang, Ray Jiang, Tom Le Paine, Richard Powell, Konrad ZoIna, Julian Schrittwieser, David Choi, Petko Georgiev, Daniel Toyama, Aja Huang, Roman Ring, Igor Babuschkin, Timo Ewalds, Mahyar Bordbar, Sarah Henderson, Sergio Gómez Colmenarejo, Aäron van den Oord, Wojeiech Marian Czarnecki, Nando de Freitas, and Oriol Vinyals. Alphastar unplugged: Large-scale offline reinforcement learning, 2023.
Ishita Mediratta, Qingfei You, Minqi Jiang, and Roberta Raileanu. The generalization gap in offline reinforcement learning, 2024.
Luckeciano Melo. Transformers
reinforeeooo learners. In International Conference Machine Learning, pp. 15340-15359. PMLR, 2022.
Suvir Mirchandani, Fei Xia, Pete Florence, Brian Ichter, Danny Driess, Montserrat Gonzalez Arenas, Kanishka Rao, Dorsa Sadigh, and Andy Zeng. Large language models as general pattern machines. arXiv preprint 20Xiv:2307.0472, 2023.
11

---

Published as a conference paper at ICLR 2025
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. arXiv preprint arXiv:2112.10510, 2021.
Ivona Najdenkoska, Animesh Sinha, Abhimanyu Dubey, Dhruv Mahajan, Vignesh Ramanathan, and Filip Radenovic. Context diffusion: In-context aware image generation. arXiv preprint arXiv:2312.03584 2023.
Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, Artem Agarkov, Viacheslav Sinii, and Sergey Kolesnikov. Xland-minigrid: Scalable meta-reinforcemeen learning environments in jax. arXiv preprint (rxX::23122//2444 2023.
Ben Norman and Jeff Clune. First-explore, then exploit: Meta-leamiin intelligent exploration. arXiv preprint arXiv:2307.02227, 2023.
Abhishek Padalkar, Acorn Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Anthony Brohan, et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint arXiv:2310.08864 2023.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kang- wook Lee, and Dimitris Papailiopooloss Can mamba learn how learn? comparative study in-context learning tasks. arXiv preprint 2t0:2400,04248,, 2024.
Ofrr Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation, 2022.
Rafael Rafailov, Kyle Beltran Hatch, Anikait Singh, Aviral Kumar, Laura Smith, llya Kostrikov, Philippe Hansen-Estruch. Victor Kolev, Philip Ball, Jiajun Wu, et al. D5rl: Diverse datasets for data-driven deep reinforcement learning. 2023.
Sharath Chandra Raparthy, Eric Hambro, Robert Kirk, Mikael Henaff, and Roberta Raileanu. Gen- cralization to new sequential decision making tasks with in-context learning. arXiv preprint arXiv:2312.03801, 2023.
Jeff Rasley, Samyam Rajbhandari. Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimiza- tions enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery Data Mining, pp- 3505-3506, 2020.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent, 2022.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint urXiv:1707.065 2017.
Gresa Shala, André Biedenkapp. and Josif Grabocka. Hierarchical transformers are efficient meta- reinforcement learners. arXiv preprint wrXiv:2402.06402. 2024.
Lucy Xiaoyang Shi, Yunfan Jiang, Jake Grigsby, Linxi Fan, and Yuke Zhu. Cross-episodic curriculum for transformer agents. Advances in Neural Information Processing Systems, 36, 2024.
Viacheslav Sinii, Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, and Sergey Kolesnikov. In-context reinforcement learning for variable action spaces. arXiv preprint 23X1v:2312.13322 2023.
Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmiee Michael Chang. Natalie Clay. Adrian Collister, al. Human-timescul adaptation in an open-ended task space. arXiv preprint arXiv:2301.07600, 2023,
Open Ended Learning Team, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, Max Jaderberg, Michael Mathieu, et al. Open-ended learning leads to generally capable agents. arXiv preprint arXiv:2107. / /8008 2021.
12

---

Published as a conference paper at ICLR 2025
Keyu Tian, Yi Jiang, Zehuan Yuan, Bingyue Peng. and Liwei Wang. Visual autoregressive modeling: Scalable image generation via next-scale prediction. arXiv preprint TrXiv:2404002905, 2024.
William L Tong and Cengiz Pehlevan. Mips learn in-context. arXiv preprint arXiv:2405.15613 2024.
Max Vladymyrov, Johannes von Oswald, Mark Sandler, and Rong Ge. Linear transformers are versatile in-context learners. arXiv preprint arXiv:2402. 2024.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174. PMLR, 2023.
Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning reinforcement learn. arXiv preprint arXiv:1611.05765 2016.
Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, and Shangtong Zhang. Transformers learn temporal difference methods for in-context reinforcement learning. arXiv preprint arXiv:2405. 13861, 2024.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Merl-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100 1100. PMLR, 2020.
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl. Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta-learning. arXiv preprint wwXiv:1110008345. 2019.
llya Zisman, Vladislav Kurenkov, Alexander Nikulin, Viacheslav Sinii, and Sergey Kolesnikov.
Emergence of in-context reinforcement learning from noise distillation. arXiv preprint 122755 2023.
13

---

Published as a conference paper at ICLR 2025
A DOWNLOADING THE DATASETS
Both Mhand-100m and XLand-Trivine -rii---vviialee datasets hosted on public S3 bucket and are 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.
XLand- -rind-Trivial-2000 approx 60GB size
curl -I, -o
Mland-100B: approx 325GB size
curl -L -o xland-100b. hdf5
B WHAT IS INSIDE DATASET?
Both Trivial and 100B dataset are HDF5 files holding the same structure. The dataset is grouped the following way:
data (key_ _i] / (entity_naat
where key_id id an ordinal number task in dataset, learning learning_history _dd a learning history number from 0 32 and entity name is of the names mentioned in Table B
NB! Do not confuse key_ ji with the task ID, which should accessed via
data [" _iey_idd_ attrs "rulesse-id"]
Table 3: Data description
Name
Type
Shape
Description
states
np.uint8
(5, 5)
8t, colors and tiles from agent's POV
actions
uintin
scalar
dr, from to NUM_ACT TITM_ACTTCC TOS
rewards
fllalo
scalar
which agent recieved at timestep t de terminated truncated episode flag same as but from a generating policy
dones
np.bool
scalar
expert_actions
np.uint8
scalar
C COMPRESSION CHUNK SIZE TUNING
Table 4: Throughput with Porch dataloader with different HDF5 compression chunk size settings. We used 2048 sequence length, 64 batch size and 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
14

---

Published as a conference paper at ICLR 2025
D ADDITIONAL FIGURES OF DATA COLLECTION
1.0
1.0
Return
Return
0.5
0.6
0.5
0.2
0.0
1 2
3
4 5 6 7 8 9
3
7
Transitions
Number of Rules
Number of rules
Figure 9: Single-task evalua- tion curves on 256 tasks for policies trained from scratch or fine-tuned from multi-task pre- trained checkpoints.
Figure 10: Final return by num- ber rules on 256 tasks for policies trained from scratch or fine-tuned from multi take pre- trained checkpoints.
Figure 11: Final return by number of rules in the final -Land-100B dataset after post- processing.
ADDITIONAL FIGURES OF DATA EVALUATION
-aadd-Triial--00B
XLand-1008
XLand-Trivial-208
1.0
0.8
Return
Return
0.8
0.6
0.5
Artions
0.4
0.4
0.0
0k
2k
4k
6k
8k
Ok
3k
6k
9K
12k
15k
0k
10k
20k 30k 40k 50k
60k
Eessodes
Episodes
Transitions
Figure 12: Return averaged all learning histories in the final dataset.
Return averaged over all learning histories in the final XL dataset.
Figure 13:
Figure 14: Agreement between actions predicted by the expert and the actual actions the learning history. We use final PPO policy as an expert for ac- tions labeling.
15

---

Published as a conference paper at ICLR 2025
F AD IMPLEMENTATION
Table 5: Time for training and evaluation.
sequence length
1024
2048
4096
train eval
15 hrs
10 hrs
11.5 hrs
20 min
30 min
50 min
We implement AD following the original paper of Laskin et al. (2022). To optimize the speed of training and inference we use FlashAttention-2 (Daol (202) with KV-caching and ALiBi positional embeddings (Press et al. 2022). 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., (0200) 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 -0BBB dataset and evaluation on 1024 tasks on H100 GPUs shown in the Table The computations were done on an internal cluster.
The hyperparameters was copied from (Laskin et al. 2022) except for the size of the network, it was scaled up 25 M. The exact hyperparameters can found in Table X
We also show training logs for both datasets. Trivial: wandb medium: wandb
G DPT IMPLEMENTATION
Our implementation is based on the original one (Lee et al. 2023). Compared to AD implementation (Laskin et al. 2022), during training phase, the model context is generated with decorrelated dataset transitions increase data diversity and model robustness: given query observation and 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, is lot easier analyze this policy than 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. 2023
The training consisted of 3 epochs due computational and time limitations 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 8 A100 GPUs. The computations were done on an internal cluster.
The hyperparameters was copied from (Lee al. 2023) except for the size of the network, was made up to 25 M, and sequence length, it was increased due to complexity of the tasks. The exact hyperparameters can found in Table
We also show DPT training logs. Trivial: wandb
medium: wandb
H ON DPT LIMITATIONS IN POMDP
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. (022). The agent is required to find an invisible key and then open 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- 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.
16

---

Published as a conference paper at ICLR 2025
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 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 [3] in the Markovian case the model converges the optimal return, finding both key and a door. In the latter case, the model reaches plateau reward of 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 key was already found from the random context. Without this knowledge, impossible for the agent to know whether should search for the key or for the door.
We also show logs for Key-To-Door Experiments: Markovian: wandb, non-Markovian:
wandb
Table 6: Q-Learning Hyperparameters
Hyperparameter
Value
Num. Train Goals
2424
Num. Histories
5000
Num. Updates
50,000
Learning Rate
3e-4
2.0
20
2.0
2.0
1.5-
200
307
100
200
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 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
DETAILS OF IN-CONTEXT EVALUATION
The evaluation process resembles the standard in-weight learning, with the key difference being that the learning itself happens during evaluation. When interacting with the environment, the agent populates the Transformer's context with the latest observations, actions and rewards. the start of the evaluation, the context is empty. Note that the agent's 's context is cross-episodic, which makes it possible for the agent access transitions previous episodes. We report the cumulative reward that the agent achieved at the end of each episode. We evaluation for each model for 500 episodes, reporting mean return across 1024 unseen tasks with standard deviation across 3 seeds. We claim that the in-context emerged when the mean return rises up until some level. This means the agent improves its policy from episode to episode, learning how to solve task.
17

---

Published as a conference paper at ICLR 2025
J ADDITIONAL FIGURES OF AD PERFORMANCE
1.0
data model
0.8
Retum
0.6
0.4
0.2
0.0
1
2
3
4
5
8
9
Number of Rules
Figure 16: AD performance on different complexities of rulesets. AD is evaluated on 1024 training tasks from -10 10B.. Sequence length is 1024.
nile
ride
0.
1.6
3.4
32
200
200
220
300
400
300
223
200
353
406
330
106
250
200
403
200
205
323
404
355
126
250
300
403
200
2.0
200
1.0
4.0
2.0
02
3.
500
Spacdss
Spaodes
Spandss
Spacdks
Figure 17: AD's performance on different task complexities, a full variant of Figure 7
18

---

Published as a conference paper at ICLR 2025
K ADDITIONAL FIGURES OF DPT PERFORMANCE
XLaad-Trivial-20B
XLand-100B
0.6
0.6
1024
0.5
0.5
0.4
4096
Retur
Retum
0.3
0.3
0.2
0.2
0.0
0.0
100
200
300
400
500
100
200
300
400
500
Episodes
Episodes
Figure 18: DPT performance on both datasets. Evaluation parameters are the same as in Figure
1.0
data model
0.8
Return
0.4
0.2
3
4
8
9
Number of Rules
Figure 19: DPT performance different complexities of rulesets. DPT is evaluated on 1024 training tasks from -100B. Sequence length is 1024.
goals MMMMAMA
goals
# goals
goals
8
mode
mede
model de'a
deta
5a0h
0.6
Return
0.6
Return
0.6
Return
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.2
500
500
Episodes
Epiiodes
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 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 in Figure
19

---

Published as a conference paper at ICLR 2025
L RULESETS VIZUALIZATION
Ruleset 22
Task
Ruleset 68
Tree
Ayer the S0
Rightti B
B
Exaeple
Layour
Ecsngle Stare Lapouc
Osstractor Balae
Olstractor Rales
Nlstractor Sales
Sperber on
Aper pewrsoo BI-B
Ave Irue
Figure 21: Rulesets with increasing difficullyy (i.e. with 3, 5, and 9 rules respectively) sampled from medium- Im benchmark. They can be accessed by ID's via the XLand-MiniGrid package.
M
GRID LAYOUTS USED FOR DATASET COLLECTION
Figure 22: Layout used for -rind-Trivial-208 Corresponds to LLaan-MiniGriddRI -1-9x9 environ- ment from the Xinnn-rindr suite. Game ob- jects are task dependent.
Figure 23: Layout used for 100mn Cor- responds to XI Lland-MnnGGridRRI-133. environ- ment from the Minirrri suite. Game ob- jects are task dependent.
20

---

Published as a conference paper at ICLR 2025
N ADDITIONAL EXPERIMENTS ON NUMBER OF TASKS
We run additional experiments with different number of tasks to show how it can affect the perfor- mance. Experiments were done with the sequence length of 2048. To make the number of gradient updates closer to the original experiments, we progressively increased the number of epochs for each smaller subset. Note that the return increases from 100 goals to 3k goals, indicating the importance of scaling the number of goals in the dataset. Running the AD for a single epoch on full goals dataset clearly not enough, however, large scale training too computationally expensive and is left for future research.
0.5
0.4
Return
0.3
0.2
tasks; epoch
3k tasks; 5 epochs
0.1
300 tasks; epochs
100 tasks; 20 epochs
0.0
100
200
300
400
500
Episodes
Figure 24: Training AD for different number of goals.
21

---

Published as a conference paper at ICLR 2025
O HYPERPARAMETERS
Table 7: DPT Hyperparameters
(a) DPT Hyperparamelers for
(b) DPT Hyperparameters Key-to- Door
datasets
Hy perpaaaraeee
Value
Ryperperameter
Valur
Embedking Dim.
64
Embedding Dim.
S4
Number of Layes
Number of Layers
Number Hexes
Redloorwrdd Dim.
256
Peedforward Dim.
Layernom Placoment
Pre Norm
Layemorm Placement
Pre Norm
Embeeking Dropout Rate
0.1
Residual Dropout
Batch
[512. 256, 128
Sequence Length
100, 250, 350, 500)
Sequence Length
[1224
Ratch sizp
126
Optimizer
Adam 00.9. 0.99)
Optimizer
Betas
Betas
(0.9, 0.99)
Learning Rote
le-1
Learning Rote
le-3
Smoothing
Learnin Rane Schodale
Cusing Devay 0.05 25 M
Learuing Rase Schedale
Costne Decay 0,05
Wentup Ratio
Warmup Ratio
Purameters
Parumeters
200 K.
Table 8: AD Hyperparameters
Hyperparameter
Value
Embedding Dim.
64
Number of Layers
8
Number of Heads Feedforward Dim. Layernorm Placement
512
Pre-norm
Embedding Dropout
0.1
Batch size
[256, 128, 64] [1024, 2048, 4096]
Sequence Length
Optimizer
Adam (0.9, 0.99)
Betas
Learning Rate
le-3
Learning Rate Schedule
CosineI
Warmup Steps
500
Parameters
25
22

---

Published as a conference paper at ICLR 2025
Table 9: PPO hyperparameters used in multi-task pre-training from Section 4.2
Hyperparameter
Value
env id benchmark_id
HLand-MiniGriddRRI 13x11
medium- Im
use_bf16
True
cetrain_multitask
True
context emb dim context hidden dim context context_dropout dropout obs bos__emb_din
16
64
0.0
16
action_ent lam_eembinn rn rnn hon_hidden_ddee dim
16
1024
mn mn_ mnn_num__ layers
256
num_envs num_steps update epochs num minibatches
65536
256
64
total timesteps
25,000,000,000
optimizer
Adam
decay_lr Ir clip_eps
True
0.0005
0,2
gamma
0.995
gae_lambda
0.999
ent_coef
0.001
cf_
0.5
max grad norm
0.5
eval_episodes
256
eval_sced train_sced
42
42
23

---

Published as a conference paper at ICLR 2025
Table 10: PPO hyperparameters used in single-task fine-tuning from Section 4.2
Hyperparameter
Value
env id benchmark_id
MLand-MiniGriddRR
medium- Im
use_bf16
True
cetrain_multitask
False
context emb dim
16
context hidden dim
64
context context_dropout dropout bos__emb_din
0.0
16
action_ent seiei_em__diiii_
16
rn rnn hon_hidden_ddee dim
1024
mn mn_ mnn_num__ _ayers
256
num_envs num_steps update epochs num minibatches
8192
256
total timesteps
1,000,000,000
optimizer
Adam
decay_lr Ir clip_eps
True
0.0005
0,2
gamma
0.995
gae_lambda
0.999
ent_coef
0.001
cf_
0.5
max grad norm
0.5
eval_episodes
256
eval_sced train_sced
42
42
24


Reviews:
[{"review_id": "wgJYLREXmW", "reviewer": "Reviewer_4NsP", "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.", "questions": "Line 350\n\n> 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.\n\nI 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?\n\n# References\n\n[1] D4RL: Datasets for Deep Data-Driven Reinforcement Learning, https://arxiv.org/abs/2004.07219", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.\n2. The paper is nicely structured and well-written, giving insights and reasoning to the community.\n3. 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.\n4. 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!", "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.\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tYQLhB3hPC", "reviewer": "Reviewer_gcgV", "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. \n\nThe 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.\n\nThe 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).\n\nFinally, 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.", "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)?\n2. 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 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. \n\nSome strengths worth highlighting:\n\n1. I particularly appreciate the level of detail regarding data collection, the dataset format, and other implementation/engineering methodology.\n2. 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).\n3. Including the lighter XLand-Trivial-20B dataset should make the benchmark significantly more accessible to researchers with limited resources.\n\nI would be happy to increase the score, pending clarification on some of my concerns below.", "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.\n\nRegarding the AD & DPT baselines:\n\n1. 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).\n2. 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?\n3. 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.\n4. 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.)?\n\nMinor:\n\n5. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fClfnstlDk", "reviewer": "Reviewer_HFJ5", "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.", "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?\n\n- How do you ensure that filtered tasks with low returns are adequate proxies for high-quality demonstrations in ICRL?\n\n- How might labeling inaccuracies impact the performance of models relying on these demonstrations?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- The paper is well written.\n- There do not seem to be any major issues with the method or the data set.\n- Authors appear to have taken reproducibility seriously.\n- The technical presentation of the data set, including relevant specifications, is comprehensive and clear.\n- The work addresses an important gap that is hindering progress in an area of reinforcement learning of growing interest.", "weaknesses": "The paper has no critical weaknesses. However, for the sake of constructive academic discussion, I include some drawbacks below:\n\n- 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.\n- 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.\n- 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.\n\nThese 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rm5eMAvdM1", "reviewer": "Reviewer_SXgP", "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.", "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. \n- 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? \n- 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?\n- In Figure 8, AD performance decreases with a larger context length. Why is that? Can you clarify how many episodes the context comprises?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 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.", "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. \n\nTo 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.\n\nWhile 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.\n\nIn 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)