🤖 machine learning

Large Language Models for Sequential Decision-Making: Improving In-Context Learning via Supervised Fine-Tuning

This paper demonstrates that supervised fine-tuning of pretrained large language models on offline, oracle-labeled trajectories significantly enhances their in-context learning capabilities for sequential decision-making across MDPs, POMDPs, and APOMDPs, achieving substantially lower optimality gaps than baseline methods, particularly in complex, long-horizon, and partially observable environments.

Original authors: Minmin Zhang, Sina Aghaei, Soroush Saghafian

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Minmin Zhang, Sina Aghaei, Soroush Saghafian

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, well-read librarian (the Large Language Model or LLM). This librarian has read almost every book in the world and is incredibly smart. However, they have never actually played a video game or managed a complex strategy before. They know the theory, but they haven't practiced the specific moves needed to win.

This paper asks: How do we teach this librarian to make a sequence of smart decisions in a new, tricky game, using only a few examples?

Here is the breakdown of the paper's approach and findings, using simple analogies:

1. The Challenge: The "One-Shot" vs. The "Coach"

Usually, if you want the librarian to play a game, you might just show them a few examples of someone else playing (this is called In-Context Learning).

The Problem: If the game is simple, the librarian does okay. But if the game is long, confusing, or has hidden rules (like a foggy battlefield where you can't see the enemy), the librarian gets lost. They might guess the wrong move because they are just "reading" the examples in the moment without truly understanding the underlying strategy.

2. The Solution: The "Intensive Coaching Camp" (Supervised Fine-Tuning)

Instead of just showing the librarian a few examples right before the game starts, the authors put the librarian through a coaching camp.

The Method: They take a massive library of "perfect play" videos (trajectories) where an expert (an "oracle") has already won the game. They feed these videos to the librarian and say, "Learn the pattern of how the expert thinks and acts."
The Twist: They don't retrain the librarian from scratch (which would be like teaching a toddler to read). Instead, they use a technique called QLoRA (a "smart adapter"). Think of this as giving the librarian a special pair of glasses or a specific playbook that tweaks their existing knowledge just enough to master this new type of game, without rewriting their entire brain.

3. The Three Types of Games Tested

The researchers tested this coaching method on three types of decision-making scenarios, getting progressively harder:

MDP (The Clear Room): Imagine a game where you can see everything perfectly. You know exactly where you are and what will happen if you move left or right.
- Result: The coached librarian was much better than the uncoached one, especially in long games.
POMDP (The Foggy Room): Imagine the same game, but now there is a thick fog. You can't see the whole board; you only see a small patch around you. You have to guess where the enemy is based on limited clues.
- Result: The coaching helped the librarian handle the fog much better. The uncoached librarian got confused easily, but the coached one learned to trust its "gut feeling" (the pattern it learned during training).
APOMDP (The Foggy Room with a Trickster): This is the hardest level. Not only is there fog, but the rules of the game might be slightly different depending on who is running the simulation. It's like playing a game where the physics might change slightly every time you play.
- Result: The coached librarian still outperformed everyone else. It learned to be robust enough to handle the uncertainty and the "trickster" rules.

4. The "Why" (The Theory)

The authors didn't just say "it works"; they tried to explain why using math.

The Analogy: They compared the librarian's attention mechanism (how it focuses on parts of the text) to a calculator.
When the librarian is "coached" (fine-tuned), it learns to use its attention layer to secretly calculate the "best move" (a Q-function) based on the few examples you show it during the game.
They proved mathematically that the librarian's mistakes come from two sources:
1. Confusion from the few examples: If you only show it one bad example, it gets confused.
2. Bias from the training: If the "coaching camp" wasn't long enough, it might have learned a shortcut that isn't perfect.
- The paper shows that by training on enough "perfect play" videos, you can minimize these errors.

5. The Results

Better than Random: The coached librarian didn't just guess; it played strategically.
Better than "Just Reading": The coached librarian beat the librarian who was only shown examples right before the game (In-Context Learning).
The Big Win: The coaching helped the most in the hardest situations: long games, foggy environments, and games with tricky, changing rules. In some long games, the coached librarian cut the "mistake rate" in half compared to the uncoached version.

Summary

The paper demonstrates that if you take a smart, pre-trained AI and give it a specific "coaching camp" using offline data (videos of perfect play), it becomes a master of sequential decision-making. It learns to handle long-term planning, hidden information, and uncertain rules much better than if you just asked it to "figure it out" on the fly.

Note on Scope: The authors specifically tested this on synthetic games (simulated environments like energy management or grid worlds). While they mention that this could be useful for real-world fields like healthcare (where data is abundant but experiments are expensive), the paper itself only presents results from these simulated games. They did not test this on actual patients or real-world hospital data in this study.

Technical Summary: Large Language Models for Sequential Decision-Making via Supervised Fine-Tuning

Problem Statement

Sequential decision-making under uncertainty is critical in high-stakes domains like clinical treatment planning and resource management. While Dynamic Treatment Regimes (DTRs) offer a framework for learning from observational data, they often rely on the strong assumption of sequential ignorability (all confounders are observable). In reality, unobservable confounders necessitate Ambiguous Dynamic Treatment Regimes (ADTRs), which evaluate policies against a set of plausible models (model ambiguity) rather than a single one. This corresponds to Ambiguous Partially Observable Markov Decision Processes (APOMDPs).

Although Large Language Models (LLMs) have demonstrated remarkable In-Context Learning (ICL) capabilities, their application to sequential decision-making has been limited. Prior work typically trains transformers from scratch, focuses on fully observable Markov Decision Processes (MDPs), and neglects the challenges of partial observability and model ambiguity found in real-world settings. This paper addresses the gap in leveraging pretrained LLMs for sequential decision-making in MDPs, POMDPs, and APOMDPs using offline, oracle-labeled data.

Methodology

The authors propose a framework that adapts Supervised Fine-Tuning (SFT) to enable pretrained LLMs to perform few-shot sequential decision-making.

1. Framework Overview

Instead of training a transformer from scratch, the authors fine-tune an open-source pretrained LLM (specifically Llama-2-7B) using parameter-efficient QLoRA (Quantized Low-Rank Adapters). The model is trained on offline trajectories generated by an optimal (or near-optimal) action oracle.

Training: The model learns to predict the next action token given a history of observations, actions, and rewards, serialized into a textual format. The loss function is the standard next-token prediction loss over these serialized trajectories.
Inference: At test time, the model is provided with a few-shot context (demonstration trajectories from a new task) and the current history. It generates actions for the new task without further parameter updates, effectively acting as a policy conditioned on the in-context examples.

2. Theoretical Analysis

The paper provides a theoretical analysis focusing on linear MDPs to understand why SFT improves performance.

Interpretation: The authors interpret a trained linear self-attention (LSA) layer as implicitly estimating the parameters of the optimal Q-function from in-context data.
Suboptimality Bound: They derive an end-to-end suboptimality bound for the induced policy. This bound decomposes the error into two distinct components:
1. In-context estimation error: Decreases as the number of few-shot support trajectories ( $K_{test}$ ) increases.
2. Training-length bias: Decreases as the number of training trajectories ( $K$ ) increases.
Key Insight: The analysis suggests that while increasing support trajectories helps estimate the specific task, increasing the volume of training data reduces the bias introduced by the finite length of the training sequence, which is crucial for long-horizon tasks.

Key Contributions

SFT Framework for Sequential Decision-Making: The authors introduce a method to fine-tune pretrained LLMs on offline trajectories, enabling them to generalize to new tasks via few-shot demonstrations without additional parameter updates.
Utilization of Offline Oracle Data: The framework leverages offline datasets combined with an optimal action oracle to synthesize high-quality training trajectories. This is particularly relevant for domains like healthcare where experimentation is costly or unethical, and expert-derived supervision can be constructed from logged data.
Theoretical Decomposition of Error: The paper provides a theoretical bound that separates in-context estimation error from training-length bias, offering a formal explanation for the performance gains observed with SFT.
Empirical Evaluation across Complex Settings: The study extends evaluation beyond standard MDPs to include POMDPs (partial observability) and APOMDPs (model ambiguity), demonstrating robustness in these more realistic and challenging environments.

Experimental Results

The authors evaluate their approach on synthetic tasks across MDP, POMDP, and APOMDP settings, comparing the SFT-LLM against random baselines, ICL-only baselines, and the Decision-Pretrained Transformer (DPT).

Performance Gains: Fine-tuned LLMs achieve substantially smaller optimality gaps compared to both random policies and ICL-only baselines.
- In MDPs, SFT roughly halves the optimality gap in long-horizon settings (e.g., reducing the gap from ~43% to ~15% at horizon $T=15$ ).
- In POMDPs and APOMDPs, the gains are even more pronounced, with SFT-LLMs maintaining low optimality gaps even under partial observability ( $q=0.8$ ) and model ambiguity.
Comparison with DPT: The SFT approach consistently outperforms DPT across all settings. Notably, under the pessimistic criterion in APOMDPs ( $\alpha=1.0$ ), the performance gap between DPT and the SFT-LLM widens as the ambiguity set size increases, whereas the SFT-LLM remains stable.
Robustness: The model demonstrates robustness to out-of-distribution (OOD) test conditions, maintaining low optimality gaps when tested on parameters outside the training distribution (e.g., different transition probabilities or observation noise levels).
Darkroom Task: On the challenging Darkroom task (sparse rewards, hidden goals), the SFT-LLM achieves approximately 95% of the oracle's performance, significantly outperforming random policies.

Significance and Claims

The paper claims that supervised fine-tuning provides an effective route to endowing pretrained LLMs with sequential decision-making capabilities derived from offline data. This is significant because:

It leverages the broad representational capacity acquired during pretraining, potentially improving generalization across diverse tasks.
It offers a practical solution for domains like healthcare where offline data is abundant but online experimentation is limited.
It demonstrates that LLMs can handle complex decision-making scenarios involving partial observability and model ambiguity, which are often overlooked in existing ICL decision-making literature.
The theoretical decomposition of error offers a principled understanding of how training data volume and in-context examples interact to determine policy quality.

The authors acknowledge limitations, including the computational cost of fine-tuning, the reliance on an optimal oracle (which may be hard to compute for complex problems), and the focus on linear MDPs for theoretical analysis. They suggest future work should focus on applying this framework to real-world clinical datasets, reducing reliance on oracle supervision, and exploring multimodal settings.