RL + Transformer = A General-Purpose Problem Solver

TL;DR

This work demonstrates that a pre-trained transformer can acquire meta-learning capabilities by fine-tuning with reinforcement learning to enable In-Context Reinforcement Learning (ICRL). By training LLaMA 3.1 8B Instruct with a DQN objective on a parametric Frozen Lake environment and evaluating on unseen maps and non-stationary changes, the model iteratively improves its performance without weight updates, stitches together previously learned skills, and adapts to non-stationary conditions. The key contributions include in-context behavior stitching, robustness to data quality, and adaptation to non-stationary environments, coupled with a practical RL-fine-tuning setup using IA3 adapters. Together, these results suggest a path toward general-purpose problem solvers that leverage context to self-improve and tackle novel challenges, with implications for scalable and adaptable AI systems.

Abstract

What if artificial intelligence could not only solve problems for which it was trained but also learn to teach itself to solve new problems (i.e., meta-learn)? In this study, we demonstrate that a pre-trained transformer fine-tuned with reinforcement learning over multiple episodes develops the ability to solve problems that it has never encountered before - an emergent ability called In-Context Reinforcement Learning (ICRL). This powerful meta-learner not only excels in solving unseen in-distribution environments with remarkable sample efficiency, but also shows strong performance in out-of-distribution environments. In addition, we show that it exhibits robustness to the quality of its training data, seamlessly stitches together behaviors from its context, and adapts to non-stationary environments. These behaviors demonstrate that an RL-trained transformer can iteratively improve upon its own solutions, making it an excellent general-purpose problem solver.

Figures (7)

• Figure 1: ICRL-trained Llama 3.1 learns to solve an unseen Frozen Lake environment. The trajectories in early (\ref{['fig:subfig1']}), mid (\ref{['fig:subfig2']}), and late (\ref{['fig:subfig3']}) interactions show solution refinement. Mistakes in early inference (e.g., falling into holes) disappear with experience in late inference.
• Figure 2: Fine-tuning LLaMA 3.1 8B Instruct with IA3 Adapters and a reinforcement learning objective. The model is fed sequences of states, actions, and (if nonzero) rewards, with every episode prefixed by the <|begin_of_text|> (BOT) token and terminated by the <|end_of_text|> (EOT) token. Tokens like <|start_header_id|> (SHI), <|end_header_id|> (EHI), and <|eot_id|> (EID) separate the state, action, and reward, mirroring how instruct models delineate user and assistant roles. The model predicts the Q-value of the current state for every action, updating the Q-values during training using the Bellman backup equation.
• Figure 3: Mean cumulative reward over 50 trials as an ICRL-trained transformer improves its score on unseen environments. Maps (i.e. environment parametrization) have never been trained on but are chosen from the same distribution as training examples. Significant improvement (approximately 900% when $\alpha = 0.1$) can be observed as the agent demonstrates that it has learned to solve unseen maps. Also, $\alpha = 0.1$ significantly outperforms $\alpha = 0.01$.
• Figure 4: Mean cumulative reward over 50 trials as an 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.
• Figure 5: Illustration of how ICRL combines different experiences to generate improved solutions. Subfigures (\ref{['fig:chaining-subfig1']}) and (\ref{['fig:chaining-subfig2']}) show two example trajectories provided as context to the ICRL transformer. The inference trials in subfigures (\ref{['fig:chaining-subfig3']}), (\ref{['fig:chaining-subfig4']}), and (\ref{['fig:chaining-subfig5']}) display the paths predicted by the transformer, which leverage information from both examples to develop an optimal solution.