Source-linked AI summary
Q-Transformer: Scalable Offline Reinforcement Learning via Autoregressive Q-Functions
Yevgen Chebotar, Quan Vuong, Alex Irpan, Karol Hausman, Fei Xia, Yao Lu, Aviral Kumar, Tianhe Yu, Alexander Herzog, Karl Pertsch, Keerthana Gopalakrishnan, Julian Ibarz, Ofir Nachum, Sumedh Sontakke, Grecia Salazar, Huong T Tran, Jodilyn Peralta, Clayton Tan, Deeksha Manjunath, Jaspiar Singht, Brianna Zitkovich, Tomas Jackson, Kanishka Rao, Chelsea Finn, Sergey Levine
TL;DR
High-capacity Transformer policies are difficult to train with offline reinforcement learning on mixed-quality robotic data. Q-Transformer addresses this by tokenizing action dimensions for autoregressive Q-learning and adding conservative and return-based training choices. It outperforms imitation-learning and offline-RL baselines on a large real-world multi-task manipulation suite, while remaining focused on sparse binary-reward offline settings.
Problem
Training high-capacity Transformer policies with reinforcement learning on large, diverse real-world robotic datasets remains difficult.
Method
Q-Transformer uses per-dimension action tokenization to represent Q-functions as autoregressive discrete sequences, with conservative regularization for offline learning.
Results
Q-Transformer achieves the highest success rate and exceeds the average performance of the best prior method by about 70% on real-world manipulation tasks.
Takeaways & Limitations
Q-Transformer can improve upon human demonstrations by learning from autonomously collected sub-optimal data in large-scale multi-task robotic reinforcement learning.
Takeaways & Limitations
The method focuses on sparse binary-reward tasks, which are common in episodic manipulation but not universal.
Abstract
from arXiv · showhide
In this work, we present a scalable reinforcement learning method for training multi-task policies from large offline datasets that can leverage both human demonstrations and autonomously collected data. Our method uses a Transformer to provide a scalable representation for Q-functions trained via offline temporal difference backups. We therefore refer to the method as Q-Transformer. By discretizing each action dimension and representing the Q-value of each action dimension as separate tokens, we can apply effective high-capacity sequence modeling techniques for Q-learning. We present several design decisions that enable good performance with offline RL training, and show that Q-Transformer outperforms prior offline RL algorithms and imitation learning techniques on a large diverse real-world robotic manipulation task suite. The project's website and videos can be found at https://qtransformer.github.io
1 Introduction
Q-Transformer addresses the difficulty of training high-capacity Transformer policies with large, diverse real-world datasets for offline robotic reinforcement learning. It uses per-dimension Q-value tokenization and is evaluated on large-scale multi-task policies.
- High-capacity robotic models can generalize across tasks, environments, objects, and robot morphologies, but supervised training limits performance to demonstrator data quality.
- Training Transformers with reinforcement learning at large scale remains difficult, motivating a method that combines diverse real-world datasets with high-capacity Transformer policies.
- Q-Transformer converts Q-function estimation into discrete sequence modeling by tokenizing each action dimension separately, avoiding exponential action-cardinality growth.
- The method is evaluated on text-conditioned multi-task policies using 38,000 successful demonstrations, 20,000 failed autonomous episodes, more than 700 tasks, and 13 robots.
- Q-Transformer outperforms prior large-scale robotic reinforcement-learning architectures and Transformer-based Decision Transformer models in the reported evaluation.
2 Related Work
The related work positions Q-Transformer as an offline reinforcement-learning system that adapts conservative Q-learning and autoregressive action discretization to high-capacity robotic Transformers. It differs from sequence-modeling approaches that cannot improve or recombine demonstrated behavior.
- Q-Transformer develops a version of conservative Q-learning designed to train large Transformer-based Q-functions on mixed-quality offline data.
- Prior robotics Transformer research largely uses supervised learning, while other work applies Transformers to reinforcement learning or conditional imitation learning.
- Decision Transformer uses reward-conditioned sequence modeling but does not provide a mechanism to improve demonstrated behavior or recombine dataset components into more optimal behaviors.
- Q-Transformer discretizes each action-space dimension as a separate autoregressive time step without hierarchical MDP decomposition.
3 Background
Offline robotic reinforcement learning learns from fixed transition or episode datasets without online interaction, but must address distributional shift and sparse rewards. Q-Transformer’s architecture combines visual and language inputs for multi-task control.
- Value-based reinforcement learning represents discounted return with a Q-function and selects actions by maximizing that function.
- Offline reinforcement learning uses data generated by an unknown behavior policy without additional online interaction during training.
- Q-Transformer processes camera observations and language instructions through visual and language encoders before producing action predictions with its Transformer architecture.
- Distributional shift arises when the Q-maximizing action lies outside the dataset, motivating conservative penalties that lower values for out-of-distribution actions.
- The robotic tasks use sparse binary rewards, with success or failure assigned at the final episode time step.
4 Q-Transformer
Q-Transformer enables offline Q-learning with Transformers by autoregressively discretizing action dimensions, conservatively regularizing unseen actions, and augmenting Bellman targets with Monte Carlo and n-step returns. These choices are designed to support learning from mixed-quality robotic data.
- Q-Transformer ingredients: Q-Transformer combines per-dimension discretization and autoregression, conservative Q-function regularization, and Monte Carlo or n-step returns.
- Autoregressive discrete Q-learning: Each action dimension is treated as a sequence step, with Q-values conditioned on a state-history window and previously generated action dimensions.
- Autoregressive discrete Q-learning: The Bellman update applies reward only to the final action dimension, while discounting occurs between environment time steps rather than within a complete action.
- Conservative Q-learning: Unobserved actions are regularized toward the minimal attainable cumulative reward, which is 0 for the task setting described.
- Conservative Q-learning: The training objective combines temporal-difference error with conservative regularization weighted by multiplier α.
- Learning efficiency: Monte Carlo returns provide a lower bound on the optimal Q-function, allowing targets to use max(MCt:T, Q(st, at)) without changing Bellman convergence.
- Learning efficiency: n-step returns over action dimensions significantly improve learning speed because rewards arrive only after the complete action is executed.
5 Experiments
Q-Transformer is evaluated on diverse real-world manipulation data, including demonstrations and failed autonomous episodes, across 72 tasks. It outperforms imitation learning and offline RL baselines, while ablations identify conservatism, Monte Carlo returns, and n-step returns as important design choices.
- Real-world evaluation: Experiments evaluate Q-Transformer on 72 unique manipulation tasks using demonstrations and autonomously collected data.The training data includes about 38,000 successful demonstrations and about 20,000 failed autonomous episodes.
- Real-world evaluation: Q-Transformer exceeds the average performance of the best prior method by about 70% and outperforms RT-1, Decision Transformer, and IQL.The result demonstrates improvement over human demonstrations using autonomously collected sub-optimal data.
- Real-world evaluation: Q-Transformer can be combined with a language task planner for affordance estimation and action execution on long-horizon tasks.The reported combination outperforms prior methods for planning and executing long-horizon tasks.
- Ablations: Softmax conservatism reduces performance to around the fraction of demonstration episodes, while removing conservatism or Monte Carlo returns causes collapse.The reported softmax result is approximately 8%, and the authors attribute the ablation outcomes to overestimated values or weaker bootstrapping.
- Ablations: n-step returns achieve similar performance with 4 times fewer steps and improve performance on longer-horizon tasks.On real robots, n-step returns also provide significantly faster training with minimal performance loss compared with 1-step returns.
- Massive dataset scaling: A larger experiment uses about 300,000 trials from 13 robots, and Q-Transformer improves over RT-1 despite its 82% success rate.The dataset contains about 115,000 successful trials and 185,000 failed autonomous episodes.
6 Limitations and Discussion
The framework has three stated limitations: it targets sparse binary-reward tasks, per-dimension discretization becomes more costly as action dimensionality grows, and the current work focuses on offline RL.
- Reward setting: Q-Transformer focuses on sparse binary-reward success-or-failure tasks, which are not universal across robotic settings.The authors expect extension to more general reward settings in future work.
- Action dimensionality: Per-dimension discretization increases sequence length and inference time as action dimensionality grows, making higher-dimensional control more cumbersome.The authors suggest adaptive discretization as a possible alternative for higher-dimensional or mixed-granularity action spaces.
- Training regime: The work concentrates on offline RL, leaving online fine-tuning as a future direction.The authors associate online fine-tuning with potentially more effective autonomous improvement of complex robotic policies.
A Proof of MDP optimization consistency
The proof argues that optimizing Q-values one action dimension at a time is equivalent to optimizing the Q-function for the original full-action MDP.
- Proof setup: The analysis reformulates the full action a1:dA as a sequence of action dimensions and compares its optimization with per-dimension Bellman updates.The full-action formulation evaluates the reward after executing the complete action and bootstraps from the next state.
- Consistency result: The per-dimension optimization recovers the optimization of the original full-action MDP.The cited conclusion states that the resulting update optimizes the original full-action MDP as in Eq. 3.
B Proof of convergence
The convergence proof extends the standard Q-learning argument to per-action-dimension Q-functions by showing that the corresponding Bellman operator is a contraction.
- Proof strategy: The proof extends standard Q-learning convergence to per-action-dimension Q-functions by accounting for maximization over successive action dimensions.The argument is described as almost identical to the standard Q-learning convergence proof, with the additional per-dimension maximization.
- Action sequences: The analysis defines action sequences whose dimension may be less than or equal to the full action-space dimensionality.dA denotes the action-space dimensionality, while a may denote a partial action sequence.
- Contraction proof: The Bellman operator is shown to be a contraction by treating separately sequences shorter than and equal to the action-space dimension.The proof explicitly divides into Case 1 for shorter sequences and Case 2 for full-dimensional sequences.
- Contraction bound: The contraction bound is expressed using the supremum over action sequences, with the sup-norm defined as ||f||∞= supx[f(x)].The stated bound assumes 0 ≤γ ≤1.
C Analysis of the conservatism term
The analysis shows that the conservatism objective induces a weighted Bellman backup that downweights out-of-distribution actions and limits their Q-value overestimation.
- Analysis setup: The theoretical analysis studies the training objective in tabular Q-function settings when the objective can be minimized exactly.The minimizer is derived by differentiating the objective with respect to Q.
- Weighted backup: The resulting weighted Bellman backup assigns weights near 1 to in-distribution actions and near 0 to out-of-distribution actions.The weight m(s, a) lies between 0 and 1 and is small when the behavior-policy probability πβ(a|s) is small.
- Conservatism effect: The weighted backup is designed to prevent over-estimation of Q-values for unseen actions.This differs from the standard Bellman backup because large Q-value targets are multiplied by m(s, a).
D.1 Transformer sequence model architecture
Q-Transformer adapts a Transformer architecture to multi-task robotic offline RL by tokenizing observations, language instructions, and action dimensions. Its training computes per-dimension temporal-difference losses with conservative targets and additional return-based updates.
- Architecture: The architecture processes image observations through a convolutional encoder and tokenization, while embedding each natural-language task instruction for multi-task control.The design is adapted from RT-1 and uses a Transformer backbone.
- Action representation: Q-Transformer uses 8 action dimensions, discretizing each into N = 256 value bins.Each action dimension is represented separately for autoregressive Q-value modeling.
- Implementation: The implementation uses an exponential-moving-average target network with averaging constant 0.01 for target Q-value estimation.The reward is sparse, with value 1.0 only at the final successful step, and γ = 0.98.
- Training objective: Training computes each action dimension’s loss from next-dimension maximum Q-values, Bellman targets, and mean-squared temporal-difference errors.The algorithm also maximizes the Bellman target with the Monte Carlo return.
- Conservatism: A conservative regularizer assigns a target of 0 to discretized action bins that differ from the dataset action.The regularizer sums over unseen action bins, whose number is determined by N.
F Running training for multiple random seeds
The authors evaluate training variability by running Q-Transformer and RT-1 with five random seeds in simulation. Q-Transformer retains its improved performance across the resulting seed distribution.
- Evaluation protocol: Five random seeds are used to compare the mean and variance of Q-Transformer and RT-1 performance in simulation.Figure 7 summarizes the distribution across these training runs.
- Result: Q-Transformer retains its improved performance across the distribution of random seeds.The authors present this as evidence that the observed learning improvement is not limited to one seed.
G Q-Transformer value function with a language planner experiments
The experiments use Q-Transformer values to estimate task affordances and support long-horizon language-planned manipulation. Q-Transformer produces sharper near-robot value estimates and outperforms the RT-1 with QT-Opt comparison in planning and execution.
- Value-function comparison: Q-Transformer produces sharper Q-values for objects close to the robot than for objects farther away.The qualitative comparison is against QT-Opt in a sim-to-real setting.
- Affordance estimation: Multi-task relabeling consistently improves Q-Transformer affordance estimates in precision, recall, and F1 score.Relabeling assigns alternate task names and zero reward to sampled episodes to clarify task boundaries.
- Long-horizon evaluation: Q-Transformer outperforms RT-1 with QT-Opt in both planning and execution on SayCan-style long-horizon tasks.A language model proposes task candidates, Q-values select the highest-affordance candidate, and an execution policy performs it.
- Evaluation tasks: The evaluation covers drawer manipulation, opening and closing drawers, and moving objects near specified targets.These tasks are listed among the real-robot experiments.