RL + Transformer = A General-Purpose Problem Solver (2025)

Abstract: What if artificial intelligence could not only solve problems for which it was trained but also teach itself to tackle novel tasks? In this paper, we finetune Llama 3.1 using reinforcement learning on the grid-world game Frozen Lake and investigate its ability to solve maps it has never encountered—a phenomenon recently termed In-Context Reinforcement Learning (ICRL). Without additional training, the transformer demonstrates the capacity to adapt to both in-distribution and out-of-distribution environment parameterizations. Moreover, it remains effective when trained on data that blends optimal and suboptimal behavior, combines strategies from its context (behavior-stitching), and dynamically adapts to non-stationary environments. These proof-of-concept findings suggest that in-context learning via reinforcement-tuned transformers may form the basis of a promising general-purpose problem-solver.

## introduction imagine a mars mission in which a robot's appendage suddenly loses functionality. an adaptive agent capable of rapidly learning new behaviors could continue the mission successfully. while reinforcement learning (rl) has excelled in stationary environments @xcite , realworld applications frequently demand quick adaptation to unexpected changes-something traditional rl struggles to achieve efficiently @xcite . recent advances leveraging transformer architectures demonstrate remarkable generalization capabilities through in-context learning (icl), enabling rapid adaptation to novel tasks without retraining @xcite . inspired by this, we explore whether transformers finetuned via reinforcement learning can adapt to changing conditions without additional training, a in this work, we finetune llama 3.1 @xcite on the grid-world game frozen lake (farama foundation, 2022) using a reinforcement learning objective @xcite . we show that our model not only solves previously unseen maps from the same distribution (as seen in figure 1 ) but also generalizes to out-of-distribution environments, robustly learns from varying data quality, dynamically adapts to non-stationary conditions, and effectively combines learned behaviors. these capabilities underscore the potential of reinforcement-tuned transformers as versatile, general-purpose problem solvers capable of human-like adaptability. ## background and related work recent efforts combining transformers with reinforcement learning (rl) have significantly advanced solving complex sequential decisionmaking tasks. the decision transformer (dt) reformulates rl as a sequence modeling problem by predicting future actions conditioned on past trajectories and desired returns, effectively leveraging transformer architectures for improved performance. however, dt relies heavily on high-quality, diverse training trajectories, limiting its applicabil-ity to new or out-of-distribution scenarios @xcite . to overcome the limitations of dt, algorithm distillation (ad) trains transformers to emulate rl algorithms themselves rather than directly modeling trajectory sequences. this allows ad to capture underlying patterns in algorithmic decisionmaking, enabling sample-efficient generalization. yet, this approach requires that the problem first be solved with traditional rl, making it dependent on the hyper-parameters and quality of the algorithm it is imitating @xcite . the decision-pretrained transformer (dpt) further addresses limitations of dt and ad by training transformers to directly imitate an action oracle, thus learning near-optimal policies without explicitly requiring future reward predictions. despite its improved trajectory stitching capabilities, dpt's reliance on oracle-provided optimal actions during training restricts its practical application to scenarios where such optimal solutions are readily available @xcite . meta-rl approaches utilizing transformers have demonstrated promising capabilities for rapid task adaptation via in-context learning (icl). these methods efficiently generalize learned behaviors to novel tasks based on minimal contextual experience @xcite . however, their robustness to dynamically changing environments and resilience against imperfections in training data quality have not been extensively explored. building upon the foundational works that merge transformers with reinforcement learning, we focus on harnessing the potential of icrl. while previous studies have established that transformers can be finetuned using reinforcement learning to solve novel problems, our work advances this integration by uncovering and demonstrating several novel advantages of icrl that have not been previously explored. specifically, we show: these findings suggest that icrl offers significant advantages in developing versatile ai systems capable of human-like adaptability. ## methodology to explore the capabilities of icrl, we employ the open-source large language model (llm) called llama 3.1 8b instruct @xcite . we finetune this model using the deep q-network (dqn) reinforcement learning algorithm @xcite , which enables the model to learn optimal actions through trial and error. our training data are collected from the parametric game frozen lake (farama foundation, 2022), a dynamic environment where the game parameters can be changed between episodes. rather than focusing on solving a single, specific version of frozen lake, our objective is to enhance the model's performance across multiple episodes with varying game configurations. by doing so, we aim to improve the model's ability to generalize and find better solutions over time, thus highlighting the benefits of the icrl approach. this section aims to provide a clear understanding of our experiments and results. we begin by explaining the general problem formulation for a partially observable markov decision process (pomdp). then, we review how reinforcement learning is applied to solve a pomdp. next, we demonstrate how reinforcement learning can be applied to a pre-trained transformer model. finally, we document our environment setup and data collection procedures. ## experiments having finetuned llama 3.1 with the generated data, in this section, we evaluate its performance across a variety of tasks to demonstrate its capabilities as a general-purpose problem solver. we examine its ability to solve both in-distribution and outof-distribution examples, its capacity for in-context icrl-trained transformer improves its score on unseen and out-of-distribution environments. generated maps are larger than anything ever seen during training. improvement can be observed (though not as significant as in the in-distribution case) as the agent demonstrates that it has learned useful behaviors even for environments outside the distribution of its training data. ## conclusion in this study, we have demonstrated that fine-tuning llama 3.1 with reinforcement learning enables it to function as a general-purpose problem solver. within the frozen lake environment, it is capable of adapting to situations it has never encountered before. while in-context reinforcement learning (icrl) may not always find the perfect solution, it can enhance performance in unforeseen scenarios. this progress indicates that agents capable of human-like adaptability and continuous improvement are within reach.