PaperlayerView exact revision on arXiv

Source-linked AI summary

From Weak Data to Strong Policy: Q-Targets Enable Provable In-Context Reinforcement Learning

Yichen Lin, Xuyuan Xiong, Xue Wang, Xiangfu Meng, Mike Mingcheng Wei, Tao Yao

arXiv:2609.30391v1cs.LG

TL;DR

In-context RL methods commonly use supervised behavior prediction, leaving their learning signal tied to the quality of offline actions. QTPT instead trains a context-conditioned Transformer with Bellman-style Q-targets, and the paper reports stronger robustness to weak data than supervised behavior prediction.

  • Problem

    Supervised behavior prediction for in-context RL uses offline actions as targets, which can encourage imitation of weak decisions when those actions are suboptimal.

    Source: §1 · p. 1

  • Method

    QTPT replaces behavior cloning with Bellman-style Q-target pretraining to estimate context-conditioned action values from rewards and transitions.

    Source: §1 · p. 2

  • Results

    QTPT shows strong gains over supervised behavior prediction on random data and improves robustness under weak offline data.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Controlled experiments support Q-target pretraining as an objective-level alternative to supervised behavior prediction when behavior actions are weak.

    Source: §6 · p. 11

    The theoretical guarantees require offline-data support and coverage, so they do not apply to arbitrary low-coverage datasets.

    Source: §5 · p. 11

Abstract

from arXiv · show

Existing in-context reinforcement learning methods mainly pretrain Transformers with supervised behavior-prediction objectives. This enables task inference from context, but makes the learned policy strongly depend on the quality of offline actions: when trajectories are weak or suboptimal, imitation itself becomes a biased learning signal. We propose Q-Target Pretrained Transformers (QTPT), which keeps the context-conditioned Transformer architecture but replaces behavior cloning with a Bellman-style Q-target objective. QTPT therefore learns to use rewards and transitions in the context to estimate action values, rather than simply imitating the behavior policy. We theoretically analyze QTPT in stochastic linear bandits and finite-horizon MDPs, showing stronger robustness to data quality than supervised pretraining. Empirically, QTPT improves over supervised behavior prediction on controlled RL benchmarks with random or suboptimal data, and we examine extensions to D4RL Kitchen and AntMaze. Supplementary experiments evaluate backbone robustness, meta-RL comparisons, task-coherent context, and unsupported-action value overestimation. These comparisons distinguish the benefits of Q-target pretraining from the remaining limitations of offline coverage.

1 Introduction

QTPT replaces behavior-action supervision with Bellman-style Q-target pretraining to address weak offline actions. The paper analyzes this objective and evaluates its robustness against supervised behavior prediction.

1 Introduction

  • Supervised ICRL methods include Algorithm Distillation [Laskin et al., 2022] and Decision Pretrained Transformer [Lee et al., 2023], but their supervision sources differ.The matched SPT predicts offline behavior actions, so weak or suboptimal recorded actions encourage imitation of weak decisions.Source: §1 · p. 1
  • QTPT replaces behavior cloning with a Bellman-style Q-learning objective that estimates context-conditioned values from rewards and transitions rather than behavior actions alone.This reduces reliance on expert trajectories or optimal-action labels and is intended to improve robustness under weak offline data.Source: §1 · p. 2
  • The theoretical analysis derives suboptimality bounds for stochastic linear bandits and finite-horizon MDPs, separating sample bias from model bias.The two components distinguish bias from behavior-policy data from bias associated with model complexity and dataset size.Source: §1 · p. 2
  • QTPT shows strong gains over behavior prediction on random data, with supplementary tests spanning decoder families, meta-RL comparisons, and unsupported-action value overestimation.Figure 1 also contrasts SPT's reliance on expert demonstrations with QTPT's robustness to suboptimal training data.Source: §1 · p. 2, Figure 1 · p. 2

1.1 Related Work

  • QTPT asks how changing the pretraining objective affects robustness to weak offline behavior, unlike work focused on supervised sequence models learning from high-quality trajectories.QTPT uses offline trajectories to generalize to unseen tasks, whereas offline RL generally targets the tasks from which its datasets were collected.Source: §1.1 · pp. 3–4

2 Model Setting

The model treats interaction history and the current state as context for decisions across related environments. QTPT trains a shared Transformer Q-function with Bellman targets, and the analysis separates statistical and approximation effects.

2.1 Preliminaries

  • A context state combines the partial interaction history with the current state, and evolves as actions, rewards, and next states are observed.The setting includes environments with shared state and action spaces but distinct transition models, initial-state distributions, and reward functions.Source: §2.1 · p. 4

2.2 Algorithm

  • QTPT trains a shared, history-conditioned Transformer Q-function to match predicted values to bootstrapped targets, then selects greedy actions at test time with fixed parameters.Its test-time context grows with observed interactions; the analysis setup permits task-coherent episodes, not unrelated tasks as one trial's context.Source: §2.2 · pp. 4–5

2.3 Analysis Framework

  • The analysis separates finite-data and behavior-policy sample bias from model bias due to Transformer approximation and optimization limits.A nonzero approximation floor can remain for a fixed misspecified class; the approximation result establishes capacity, not optimizer recovery.Source: §2.3 · pp. 5–6

3 Main Analysis and Results

The analysis separates behavior-induced sample bias from model bias and establishes QTPT bounds under bounded-reward and coverage assumptions. For random-policy data, Bellman pretraining uses reward information where supervised imitation may retain linear cumulative regret.

3 Main Analysis and Results

  • The analysis assumes bounded rewards, optimal-policy concentrability, and lower-bounded behavior occupancy; coverage factor L is the reciprocal of the positive-occupancy lower bound, and larger L weakens the guarantee.Source: §3 · p. 6

3.1 Bound of Sample Bias

  • The sample-bias guarantee has a conservative explicit T^2 dependence and retains Bellman approximation error and class complexity, rather than establishing polynomial complexity for unrestricted histories.Source: §3.1 · p. 7

3.2 Bound of Model Bias

  • In bandits, the model-bias bound assumes λmin(Gt) ≥ αnt and leaves fixed-model approximation error ϵtf; in MDPs it includes n^-1/2 and n^-1/4 terms plus 2T√ϵtf.Source: §3.2 · p. 7

3.3 Final Bound

  • On random-policy data, supervised imitation can retain linear cumulative regret; QTPT's model-bias component decreases as O(p T log T/n), while its total guarantee still includes sample bias.Source: §3.3 · p. 8

3.4 A Dataset Criterion for When QTPT Helps

  • When actions outside G have a sufficiently positive optimal-Q gap under technical coverage conditions, Proposition 3.7 gives a sufficient condition for QTPT improvement over SPT; inside G, imitation may cause little harm.Source: §3.4 · pp. 8–9

4 Numerical Experiments

Across the evaluated contextual MDPs, QTPT generally outperforms supervised behavior prediction, particularly with weak or exploratory data. Controlled ablations, D4RL and math-reasoning extensions, and supplementary diagnostics probe the conditions and scope of these results.

4 Numerical Experiments

  • The experiments compare matched supervised and Q-target pretraining, then separately examine context, backup targets, and performance beyond controlled environments.The controlled stochastic linear bandit uses random-action and LinUCB datasets, each containing 100,000 trajectories.Source: §4 · p. 9, §4.1 · p. 9

4.1 Stochastic Linear Bandit

  • QTPT consistently outperforms SPT on both random-action and LinUCB bandit datasets.The finite-horizon curves illustrate the comparison rather than measuring the theoretical exponents; random actions can provide broad reward and transition evidence despite being poor imitation targets.Source: §4.1 · p. 9

4.2 Controlled Ablations

  • In the matched Darkroom ablations, point estimates decrease when context is removed and when bootstrapped TD targets are replaced by Monte Carlo returns.The reported confidence intervals are over evaluation episodes, not independent training seeds.Source: §4.2 · p. 10, Table 1 · p. 10

4.3 Markov Decision Processes: Darkroom and Dark Key-to-door

  • QTPT outperforms SPT in most contextual MDP settings, especially with weak or broadly exploratory offline data.Random-data pretraining can outperform expert-data pretraining in Darkroom and Miniworld, while extremely sparse expert-only rewards remain difficult for bootstrapped training in Dark Key-to-door.Source: §4.3 · p. 10, Figure 3 · p. 10

4.4 D4RL Extension

  • QTPT has the largest reported mean under the matched Transformer D4RL protocol, including 35.94% active-stage success on antmaze-umaze-v2.The active-stage checkpoint statistic is descriptive and conditional, so these runs do not establish statistically significant superiority or a comparison with published D4RL SOTA scores.Source: §4.4 · pp. 10–11, Table 2 · p. 11

4.5 More Complex Task: Math Reasoning

  • On GSM8K, QTPO improves over behavior cloning across mixture ratios by an average of 1.81%.QTPO is QTPT adapted with policy optimization.Source: §4.5 · p. 11

4.6 Supplementary Tests of Adaptation and Value Calibration

  • Support-controlled diagnostics find that Double-target and softmax backups mitigate, but do not eliminate, harmful unsupported-action overestimation; CQL adds no statistically reliable improvement.Comparisons with RL2 and PEARL test complementary data-acquisition protocols, while context experiments distinguish task-coherent histories from unrelated extra examples.Source: §4.6 · p. 11

5 Limitations

The theoretical guarantees rely on offline-data support and coverage, and the D4RL extension is scoped to matched Transformer baselines rather than specialized offline-RL systems.

5 Limitations

  • The guarantees do not cover arbitrary low-coverage datasets; the intended weak-data setting permits suboptimal or noisy behavior only with sufficient state-action coverage.The experiments are fixed-protocol D4RL extensions, not claims of state-of-the-art performance against specialized offline-RL pipelines.Source: §5 · p. 11
  • Alternative Bellman backups, larger-scale environments, and broader sequence-generation tasks remain open directions.The ablations identify context conditioning as necessary for in-context adaptation and Bellman/TD targets as an additional source of robustness.Source: §5 · p. 11

6 Conclusion and Future Work

The paper presents QTPT as an objective-level alternative to supervised behavior prediction, with controlled gains when behavior actions are weak. It also identifies limits in offline coverage, unsupported-action value estimates, and the scope of its D4RL evidence, while outlining further evaluation and backup-design work.

6 Conclusion and Future Work

  • QTPT uses reward and Bellman-target supervision to learn a context-conditioned value function from offline trajectories, with analysis separating finite-data error from Transformer approximation.The analysis includes explicit coverage, confidence, and horizon dependence.Source: §6 · p. 11
  • Controlled experiments show gains over supervised behavior prediction when behavior actions are weak.The support-controlled diagnostic finds that Double-target and softmax backups mitigate, but do not eliminate, harmful unsupported-action overestimation; a standard CQL penalty yields no statistically reliable incremental improvement.Source: §6 · pp. 11–12
  • The D4RL extension is descriptive evidence under a matched sequence-model protocol, not a standardized state-of-the-art claim.The paper calls for broader evaluation on larger D4RL suites, newer offline-RL benchmarks, and industrial-scale datasets.Source: §6 · p. 12
  • Future work includes alternative Bellman backups, distribution-shift tests, and broader evaluations across environments and reasoning tasks.The supplementary states that anonymous source code is available and the main-theory proofs are in Appendix A.Source: §6 · p. 12

A.3.1 Proof of Proposition 3.7

The proof compares QTPT and SPT when trajectories contain only a few informative context states, under stated coverage and gap conditions. The appendix also describes how a Transformer can realize a QTPT Bellman update and how the method uses context during pretraining and rollout.

A.3.1 Proof of Proposition 3.7

  • Under the stated conditions, Proposition 3.7 compares QTPT and SPT when trajectories contain only a few informative context states.The proof separates states outside the non-informative context set from states within it; its bound depends on the stated coverage and gap conditions.Source: §A.3.1 · pp. 19–20
  • The proof’s comparison contrasts an exponential-in-n term for QTPT with a 1/n term for SPT.The passages state these rates in fragments and provide no complete bound values.Source: §A.3.1 · p. 20

A.4 Transformer Realization of the QTPT Update

  • A Transformer block of appropriate scale can construct the TD target and implement a relaxed QTPT update for bounded Q-value predictions and tokenized transition histories.This is an architectural realizability result, not a guarantee of zero approximation error or convergence to the ideal Bellman minimizer.Source: §A.4 · p. 21

Pretraining

Rollout

  • At rollout, the Transformer extracts Q-values from the accumulated context and current state, selects a greedy action, then appends the observed reward and next state.The rollout context grows after each action and transition.Source: Rollout · pp. 21–22

B.3–B.5 Environment Details

  • The appendix describes Darkroom and Dark Key-to-Door data collection using random and expert policies, and Miniworld data collection using expert and random trajectories.Darkroom and Dark Key-to-Door each collect 80k trajectories per policy; Miniworld collects 24k trajectories per policy.Source: §B.3 · p. 24, §B.4 · p. 25, §B.5 · p. 25

B.6 Implementation Detail

The appendix gives the linear-bandit implementation and ablation setup, then tests QTPT in stationary and non-stationary environments. QTPT performs better in the non-stationary setting, where the reported gap over SPT widens.

B.6 Implementation Detail

  • The linear-bandit model uses a GPT-2-style decoder with 8 layers, 4 attention heads, and embedding dimension 256, trained for 50 epochs.The original setup uses batch size 32, Adam, a 2,000-step warmup, and peak learning rate 5 × 10−6.Source: §B.6 · pp. 25–26

C.1 Ablation Study on the Stochastic Linear Bandit Problem

  • The ablations vary Q-target and attention choices as well as embedding dimension, layer count, and attention-head count.The original configuration combines softmax, ReLU, and double DQN, with 8 layers, 4 heads, and dimension 256; the supplied tables identify configurations but provide no result values.Source: §C.1 · p. 26, Table 3 · p. 26, Table 4 · p. 26

C.2 Experiments on Non-stationary Environment

  • QTPT outperforms SPT in the non-stationary environment, where the performance gap widens as SPT’s effectiveness diminishes.The comparison also reports QTPT’s superior performance there relative to the stationary setting.Source: Figure 4 · p. 27, §C.2 · p. 27

C.3 More Complex Task: Math Reasoning

The paper adapts QTPT to mathematical reasoning, treating problems as environments and token selection as actions. On GSM8K and the described reasoning tasks, QTPT improves over supervised baselines, while direct Q-function token generation produced degenerate outputs in preliminary experiments.

C.3 More Complex Task: Math Reasoning

  • For mathematical reasoning, problems are environments, generated tokens are states, token choices are actions, and reward reflects correctness or reasoning accuracy.The approach showed consistent improvements over supervised learning baselines on these tasks.Source: §C.3 · pp. 27–28
  • Directly using an offline-trained Q-function to generate tokens frequently produced degenerate outputs, including repetitive tokens.The authors hypothesize that Bellman updates distort linguistic coherence when applied to token probabilities, and instead use QTPT to refine a policy in an actor-critic framework.Source: §C.3 · p. 28

GSM8K Evaluation

  • QTPT outperformed behavior cloning across the tested expert/suboptimal data mixtures on GSM8K by an average of 1.81% accuracy.The evaluation used an unseen test set and a static offline dataset.Source: Figure 5 · p. 28, §C.3 · pp. 28–29

C.4 D4RL Extension and Controlled Ablations

The paper reports a D4RL extension and matched-baseline comparisons, with explicit cautions about how those results should be interpreted. It also describes target-construction choices and limits on what the implementation and diagnostics establish.

C.4 D4RL Extension and Controlled Ablations

  • The D4RL extension reports active-stage means of 13.23 ± 4.03 for kitchen-partial-v0, 14.62 ± 3.36 for kitchen-mixed-v0, and 35.94 ± 13.52 for antmaze-umaze-v2.These are descriptive 3–4-seed matched-Transformer results, not standard final-policy scores or evidence of statistically significant superiority; active-stage means exclude failures.Source: §C.4 · p. 29
  • Table 1 compares TD, Monte Carlo, and no-context variants, but its intervals cover evaluation episodes rather than independent training seeds.The reported point estimates decrease when context is removed and when TD targets are replaced by Monte Carlo returns; they do not establish seed-level significance.Source: §C.4 · p. 29

C.5 Comparison with Other Baselines

  • QTPT has lower cumulative regret than the matched baselines, with a more pronounced gap on random than expert data.The figure is a matched-implementation comparison, not a universal ranking of original offline-RL algorithms; the shared-critic IQL-style policy is not a full independent IQL reproduction.Source: §C.5 · p. 29, Figure 6 · p. 30

D.1 Target Construction, Stochasticity, and Cost

  • The implementation uses a separate target network for evaluation, but the Double-target/softmax combination provides no conservative lower-bound guarantee for unsupported actions.For fixed targets, stochastic rewards and transitions alter target variance, not the population regression minimizer; this identity does not claim variance stays constant across changing target networks.Source: §D.1 · p. 30
  • The target-network value difference bound is a one-step comparison, not a guarantee about a complete nonlinear training trajectory.The attention-normalization comparison offers a plausible representation-level explanation for the ablation result but does not prove a unique cause.Source: §D.1 · p. 30

D.2 Unsupported-Action Overestimation

The support-controlled diagnostic finds policy-relevant max-selection errors on unsupported actions, while stabilization reduces harmful selection relative to naive hard-max QTPT. Meta-RL comparisons vary with the data-acquisition protocol and budget.

D.2 Unsupported-Action Overestimation

  • The support-controlled bandit randomizes seven supported and three unsupported actions per task, with unsupported actions never executed offline.The ten action vectors remain visible, and the setup tests behavior-support OOD rather than a fixed action-index rule.Source: §D.2 · pp. 30–31
  • QTPT’s max-gap, harmful selection, and OOD regret have 95% intervals strictly above zero, despite an OOD offset near zero.The diagnostic identifies policy-relevant max-selection error, not uniform inflation of every unsupported value.Source: §D.2 · p. 31
  • Naive hard-max QTPT increases harmful selection by 0.0524 and OOD regret by 0.1028 versus QTPT, with positive paired 95% intervals.The combined stabilization reduces harm, but the experiments do not isolate each stabilizer’s contribution; CQL-induced changes have intervals containing zero in either backup branch.Source: §D.2 · p. 31

D.3 Meta-RL Comparisons

  • QTPT has lower mean pseudo-regret at 1M transitions, while the 5M-transition means are almost identical.QTPT uses a fixed random-behavior corpus while RL2 collects fresh on-policy interactions, so the comparison is budget-dependent rather than a universal ranking.Source: §D.3 · p. 32
  • Under the fixed weak-corpus protocol, the QTPT–PEARL paired difference is −1.503 ± 0.173, with an exact two-sided sign-test p = 0.0156.This is not a reproduction of PEARL’s original benchmark; the two meta-RL tables use different task and action spaces and should not be pooled.Source: §D.3 · p. 32

D.4 Backbone and Mixture Robustness

Across nine matched comparisons, QTPT has lower mean gap than SPT, with a larger margin on random data and a narrower margin on expert data. These supplementary tests compare three decoder choices, not all model scales.

D.4 Backbone and Mixture Robustness

  • The supplementary study compares QTPT and SPT within small randomly initialized GPT-2-, Llama-, and Qwen2-style decoders, not fine-tuned pretrained language models.Its seeds and training budget differ from the original bandit setting, so its absolute GPT-2 values should not be equated with Figure 2b.Source: §D.4 · pp. 32–33
  • The backbone study uses 3-layer, width-128, 4-head models and evaluates a fixed 3,000-update endpoint rather than the submitted 8-layer setup.The supplementary configuration uses 10,000 training trajectories and Polyak target-network updates with coefficient 0.01.Source: §D.4 · p. 33
  • QTPT has lower mean gap in all nine comparisons, with a large random-data margin and a narrower expert-data margin.Similar random-data SPT means are consistent with imitation of nearly uniform behavior, but do not prove exact convergence to a random policy.Source: §D.4 · p. 33

D.5 Task-Coherent Context and Extra Episodes

Synthetic diagnostics find that task-coherent history helps identify the task, whereas unrelated extra context does not. The paper’s checklist says the abstract and introduction reflect the stated contributions and scope.

D.5 Task-Coherent Context and Extra Episodes

  • The diagnostics use a hidden-goal reward task with two observed transitions per query and add 24 transitions in the support-augmented condition.The true goal is visible only in the oracle condition, and the diagnostic encoder uses bidirectional attention over observed context rather than the submitted causal decoder.Source: §D.5 · p. 33
  • Same-task support improves the diagnostic, whereas unrelated support does not; removing history or replacing it with equally long unrelated history loses task information.These results support the task-coherence explanation, but do not establish that arbitrary extra context must hurt or that the method beats every baseline on long-horizon policy tasks.Source: §D.5 · p. 34

NeurIPS Paper Checklist

  • The paper’s checklist answers yes that the abstract and introduction accurately reflect its contributions and scope.The justification names the theoretical and empirical claims, including the D4RL extension and controlled ablations.Source: NeurIPS Paper Checklist · p. 35
  • The checklist says the conclusion discusses remaining limitations and future work, including broader benchmarks, alternative Bellman backups, OOD overestimation, and distribution shift..Source: NeurIPS Paper Checklist · p. 35

2. Limitations

The checklist identifies boundaries in statistical reporting, reproducibility materials, compute details, and broader-impact discussion. It also points to broader benchmarks, alternative Bellman backups, OOD overestimation, and distribution shift as remaining limitations or future work.

2. Limitations

  • An anonymized public code release was not provided at submission time, although the experiments use standard public benchmarks and described data-generation procedures.Source: §2 · p. 37
  • Some larger benchmark summaries, including the D4RL extension, are point estimates because of computational cost.The checklist marks the paper’s experiment-significance reporting as no, while noting standard-deviation bands for simulation results where applicable.Source: §2 · p. 37
  • The paper does not fully specify hardware type, memory, runtime, or total compute for each experiment.The checklist marks experimental resource reporting as no.Source: §2 · p. 38
  • The paper does not provide a detailed separate discussion of both positive and negative societal impacts, describing the work as primarily foundational.The checklist also identifies broader benchmarks, alternative Bellman backup designs, OOD overestimation, and distribution shift as remaining limitations and future work.Source: §2 · pp. 35, 38
  • The manuscript cites existing benchmarks and models but does not explicitly list all license terms.Source: §2 · p. 39

Related papers

Loading 2609.30391v1…