🤖 machine learning

TabQL: In-Context Q-Learning with Tabular Foundation Models

This paper proposes TabQL, a reinforcement learning framework that replaces the parametric Q-network in Deep Q-Learning with a tabular foundation model to enable rapid, in-context adaptation and improved sample efficiency through zero- or few-shot Q-value inference.

Original authors: Qisai Liu, Zhanhong Jiang, Timilehin Ayanlade, Ashutosh Kumar Nirala, Yang Li, Aditya Balu, Soumik Sarkar

Published 2026-05-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Qisai Liu, Zhanhong Jiang, Timilehin Ayanlade, Ashutosh Kumar Nirala, Yang Li, Aditya Balu, Soumik Sarkar

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 are teaching a robot to navigate a maze. The old way of doing this (called Deep Q-Learning or DQN) is like forcing the robot to memorize every single turn, mistake, and reward through a process of trial and error, constantly rewriting its own "brain" (neural network) with math-heavy calculations every time it takes a step. It works, but it's slow, requires millions of steps, and the robot often forgets what it learned if the maze changes slightly.

The paper introduces a new method called TabQL (Tabular Q-Learning). Think of TabQL as giving the robot a super-smart, pre-trained tutor that doesn't need to be retrained from scratch. Instead of rewriting its brain, the robot simply shows the tutor a short "cheat sheet" of its most recent experiences, and the tutor instantly figures out the best move.

Here is a breakdown of how TabQL works using simple analogies:

1. The Two-Step Dance: Warm-up and Switch

TabQL doesn't jump straight into the new method. It uses a two-phase approach:

Phase 1: The Warm-up (The "Training Wheels" Phase):
First, the robot uses the old, standard method (DQN) for a short while. Imagine a student doing basic math drills. The goal here isn't to become a master immediately, but to generate a decent set of notes. The robot explores the maze a bit, makes some mistakes, and collects a small "replay buffer" of data (state, action, reward). This ensures the robot has a rough idea of where it is going before switching to the new system.
Phase 2: The Switch (The "Context" Phase):
Once the robot has enough notes, it switches to TabQL. Instead of doing complex math updates, the robot takes its most recent "notes" (a small window of recent experiences) and feeds them to a Tabular Foundation Model (TFM).
- The Analogy: Imagine you are playing a video game. Instead of calculating the physics of every jump, you look at your last 10 moves in a notepad. A super-intelligent assistant (the TFM) reads those 10 moves and says, "Based on what you just did, the best move right now is this."
- The TFM is "pre-trained" on millions of generic data problems (like a tutor who has seen every type of math problem before). It doesn't need to be retrained for the maze; it just needs to see the specific context of your current situation to give a smart answer.

2. How It Learns: "In-Context" vs. "Gradient Descent"

Old Way (DQN): Like a student trying to learn by solving a problem, getting it wrong, and then painfully adjusting their entire understanding of math to fix the error. This happens millions of times.
TabQL Way: Like a student who has already mastered the concept of math. When faced with a new problem, they just look at a few similar examples from their notebook (the context) and instantly apply the pattern. They don't need to relearn math; they just need to see the specific examples relevant to now.

The paper calls this "In-Context Learning." The robot learns by looking at the context (recent history) rather than by changing its internal weights (retraining).

3. The "Cheat Sheet" Quality Control

The paper notes a critical detail: The robot still uses the old DQN method to generate the "labels" (the answers) for the cheat sheet during the process.

The Risk: If you switch to the new method too early (before the robot has done enough warm-up), the "cheat sheet" will be full of bad notes. The super-smart tutor will read bad notes and give bad advice.
The Fix: The paper proves there is a "threshold." You must do enough warm-up so the notes are decent. Once you cross that line, the new method takes over and learns much faster than the old method.

4. Why It's Better (The Results)

The authors tested this on several grid-world games (like Taxi, CliffWalking, and FrozenLake).

Speed: TabQL reached the "perfect score" much faster than the standard DQN. It needed far fewer steps to learn the maze.
Stability: Because it relies on looking at patterns in recent data rather than noisy math updates, it was less likely to get confused or "forget" things.
Generalization: When the starting conditions of the maze changed, TabQL adapted better than the old methods, likely because the "tutor" could recognize patterns across different scenarios more easily.

Summary

TabQL is a new way to teach robots to make decisions. Instead of forcing the robot to painfully relearn its brain every time it takes a step, TabQL gives the robot a pre-trained "tutor." The robot shows the tutor a few recent examples of what happened, and the tutor instantly predicts the best next move.

It works best if you let the robot do a little bit of "warm-up" practice first to ensure the examples are good. Once that's done, the robot learns the maze significantly faster and more efficiently than before.

Technical Summary: TabQL – In-Context Q-Learning with Tabular Foundation Models

Problem Statement

Deep Q-Learning (DQN) has successfully scaled reinforcement learning (RL) to high-dimensional state spaces by approximating the action-value function using deep neural networks. However, DQN relies on massive amounts of online interaction, is sensitive to distribution shifts, and often generalizes poorly beyond its training regime. These limitations stem from its reliance on slow, interaction-heavy stochastic gradient descent (SGD) to perform repeated Bellman updates.

Conversely, recent advances in foundation models have demonstrated "in-context learning" (ICL), where large sequence models can adapt to new tasks by conditioning on a small number of examples without explicit parameter updates. While sequence-based RL methods (e.g., Decision Transformers) exist, they primarily focus on offline or imitation settings and do not explicitly model Bellman optimality or provide convergence guarantees for online Q-learning.

The central question addressed by this work is: Can action-value learning be carried out through in-context inference rather than repeated gradient-based Bellman updates? The authors argue that Q-learning can be reframed as a conditional prediction problem over experience, where estimating a Q-value for a new state-action pair is a prediction conditioned on a history of past tuples.

Methodology: Tabular Q-Learning (TabQL)

TabQL is a reinforcement learning framework that replaces the parametric Q-network in DQN with a Tabular Foundation Model (TFM) capable of performing in-context Bellman inference. The framework operates through a hybrid architecture combining a warm-up phase and an online inference phase.

1. Core Components

Tabular Foundation Model (TFM): A frozen, pretrained sequence model (e.g., TabPFN, TabDPT) that maps a context of tabularized transitions and a query state-action pair $(s, a)$ to a scalar Q-value prediction. The TFM is not fine-tuned; adaptation occurs entirely through the context window.
Experience Representation: Transitions are augmented into tuples $x_i = (s_i, a_i, r_i, s'_i, \hat{Q}^{DQN}(s_i, a_i))$ . The final coordinate, $\hat{Q}^{DQN}$ , serves as a "soft label" derived from a DQN, acting as the target for the TFM's regression.
Memory Decomposition:
- Warm-up DQN: Encodes global value information and provides a coarse prior. It supplies approximate Q-labels for transitions not present in the immediate context.
- Rolling Context ( $C_t$ ): Encodes local, recent value information. The TFM uses this context to perform fine-grained, in-context inference.

2. Algorithmic Framework

The algorithm proceeds in two distinct phases:

Phase 1: Warm-up (Lines 2–4): A standard DQN trains for $T_0$ steps. This phase populates the replay buffer with informative transitions and ensures the DQN has converged sufficiently to provide high-quality "soft labels" ( $\hat{Q}^{DQN}$ ). These labels are appended to the replay buffer alongside standard transition data.
Phase 2: In-Context Bellman Inference (Lines 10–17):
- Context Construction: A context window $C_t$ of size $K$ is sampled from the most recent transitions in the replay buffer.
- Q-Value Inference: The TFM predicts Q-values for the current state $s_t$ and all possible actions $a$ by conditioning on $C_t$ : $\hat{Q}^{TFM}_t(s_t, a) = f_\phi(C_t, s_t, a)$ .
- Action Selection: Actions are selected via $\epsilon$ -greedy policy based on the TFM's predictions.
- Label Generation & Update: Crucially, while the TFM selects actions, the DQN (still running in parallel) is used to predict the Q-values for the executed transitions. These DQN-predicted values serve as the labels for the new transitions, which are then added to the buffer. This maintains the quality of the context for future inferences.
- DQN Update: The DQN parameters are updated via standard TD-loss on mini-batches from the buffer, ensuring the "teacher" labels remain accurate.

3. Theoretical Mechanism

TabQL treats Q-learning as a conditional inference problem. The TFM approximates the fixed point of an empirical Bellman operator induced by the context $C_t$ .

Amortization: Instead of iterative gradient updates, Bellman updates are "amortized" into a single forward pass of the TFM.
Error Structure: The total error is decomposed into:
1. Bellman Contraction: Standard geometric decay of error.
2. Statistical Error: Arising from finite context size and label noise.
3. ICL Error: The approximation error of the TFM in predicting the fixed point.
  Unlike DQN, where per-step errors compound via SGD noise accumulation, TabQL's per-step errors are bounded by fixed constants (assuming the TFM capacity and context quality are sufficient).

Key Contributions

Conceptual Framework: The paper reframes Q-learning as an in-context inference problem using tabular foundation models, bridging the gap between classical tabular methods, deep RL, and foundation models.
Practical Algorithm: Introduction of TabQL, which combines a DQN warm-up phase with online contextual Q-estimation. A key innovation is the use of the DQN to generate supervision signals (labels) for the TFM during online learning, ensuring context quality without requiring the TFM to be retrained.
Theoretical Analysis: The authors provide rigorous convergence and sample complexity guarantees. They demonstrate that TabQL achieves near-optimal control when the ICL approximation error and label error are small. The analysis shows that TabQL avoids the variance accumulation inherent in SGD-based DQN, leading to improved sample efficiency under specific conditions (e.g., when exploration concentrates on relevant state-action subsets).
Empirical Validation: Extensive experiments on Gymnasium environments (Taxi-v3, CliffWalking-v1, FrozenLake-v1, CartPole-v1) demonstrate that TabQL reaches near-optimal performance with significantly fewer environment interactions compared to DQN, Double DQN, and Dueling DQN.

Experimental Results

Sample Efficiency: In discrete environments (Taxi, CliffWalking, FrozenLake), TabQL variants (using TabPFN and TabDPT backbones) reach near-optimal returns substantially faster than deep RL baselines, often requiring an order of magnitude fewer interactions.
Switching Point Sensitivity: The experiments reveal a threshold behavior for the warm-up length $T_0$ . If the switch to in-context learning occurs too early (before the DQN provides accurate labels), TabQL converges to a suboptimal policy and cannot recover. Once $T_0$ exceeds a specific threshold, performance stabilizes at the optimal level.
Context Size: Increasing the context size $K$ accelerates training and stabilizes learning curves up to a saturation point, beyond which diminishing returns are observed.
Generalization: TabQL demonstrates superior generalization across unseen initial conditions compared to DQN, which often fails to transfer. Aggregating experience across multiple seeds allows the TFM to learn transferable value inference patterns rather than memorizing specific trajectory statistics.
Continuous Spaces: The framework was also tested on CartPole-v1 (continuous observation), where TabQL with TabPFN outperformed DQN baselines, suggesting applicability beyond purely discrete state spaces.

Significance and Claims

The paper claims that TabQL represents a new learning paradigm where adaptation occurs primarily through context construction rather than parameter modification. By decoupling learning from online optimization, TabQL offers:

Improved Sample Efficiency: Reducing the need for massive online interaction, which is critical for real-world domains like robotics and healthcare where data collection is costly.
Theoretical Stability: Avoiding the compounding variance of stochastic gradient updates, providing a more robust theoretical structure for convergence.
Interpolation: The framework successfully interpolates between vanilla tabular Q-learning, deep Q-learning, and in-context reinforcement learning.

The authors acknowledge limitations, noting that the switching point remains a tunable hyperparameter and that the current approach is best suited for low-dimensional state-action spaces. Extending the method to high-dimensional inputs (e.g., images) would require more expressive models, which is identified as a future direction. The work positions TabQL not as a replacement for all RL, but as a specialized, efficient approach for interaction-limited settings where reducing environment interaction is prioritized over minimizing inference cost.