q-transformer: Lucidrains' Reimplementation of Autoregressive Q-Learning
Implementation of Q-Transformer, Scalable Offline Reinforcement Learning via Autoregressive Q-Functions, out of Google Deepmind
At a glance
- What is it?
- A PyTorch implementation of Google Deepmind's Q-Transformer, aimed at engineers who want to study or extend offline RL with autoregressive Q-functions. It ships a model, a learner, an environment loop and a memmapped replay dataset, but you bring the environment and the compute.
- Who is it for?
- Adopt q-transformer if you already have a robot or simulator that can implement the BaseEnvironment interface and you want to study autoregressive Q-learning without rebuilding the architecture. Do not adopt it if you need a turnkey training pipeline, a benchmark suite, or a maintained set of pretrained weights: none of those are in the repository.
- Can I use it commercially?
- Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 35 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The gap q-transformer fills: discrete actions, one at a time
Standard deep Q-learning assumes a single scalar or low-dimensional action. A robot arm with several joints has a combinatorial action space, and regressing a Q-value over that joint space is where naive approaches tend to stall. Q-Transformer, the Google Deepmind paper this repository implements, sidesteps that by treating each action dimension as a token and predicting them autoregressively, with the reward attached only after the final action token. The README states the repository will keep the single-action Q-learning logic around "for final comparison with the proposed autoregressive Q-learning on multiple actions", which tells you the author intends this as a study artefact as much as a library. The audience is therefore narrow: researchers and engineers who already understand offline RL, want to read or modify the architecture, and are willing to supply their own environment. It is not a framework that will train a policy for you out of the box.
Autoregressive action decoding inside QRoboticTransformer
The model constructor in the README takes a MaxViT-style vision backbone through the vit dictionary (dim_conv_stem, dim, dim_head, depth as a four-tuple, window_size, mbconv rates, dropout), plus num_actions, action_bins, depth, heads, dim_head, cond_drop_prob and a dueling flag. The action space is discretised into bins, so a continuous controller has to be quantised before it reaches the model. Actions are generated in sequence rather than as one vector, and the todo list records that the main proposal was built as "autoregressive discrete actions until last action, reward given only on last". Two decoding paths exist: concatenating previous actions at the frames plus learned tokens stage, and an encoder-decoder variant with cross attention to the frame and learned tokens. The todo list also records kv caching for action decoding, which matters because autoregressive decoding over action dimensions is otherwise wasteful. Text conditioning is handled through a null conditioner, so the model can run with no instructions at all, and cond_drop_prob supports classifier-free guidance during training.
Agent, environment and the replay dataset contract
The Agent class is the piece that turns a model into data. It takes the model, an environment, num_episodes and max_num_steps_per_episode, and calling it runs the loop. The environment contract is explicit in the README: env.init() returns instructions and an initial state as Tuple[str, Tensor[*state_shape]], and calling env(actions) returns rewards, next state and a done flag as Tuple[Tensor[()], Tensor[*state_shape], Tensor[()]]. A MockEnvironment is provided with state_shape (3, 6, 224, 224) and text_embed_shape (768,) so the loop can be exercised before a real simulator is wired up. The dataset side is a folder of memmapped files produced by the environment loop, consumed by ReplayMemoryDataset. The README notes both a one time step option and an n-time steps option. This is a deliberate design: memory-mapped files keep large replay buffers off the heap, but it also means you are managing a directory of binary artefacts rather than a database, and nothing in the README describes compaction, versioning, or what happens when the folder grows past the filesystem's limits.
Getting it running: install, construct, collect, learn
Installation is a single pip command: pip install q-transformer. The README's usage block then imports QRoboticTransformer, QLearner, Agent and ReplayMemoryDataset from q_transformer. You construct the model with the vit and action arguments above, instantiate an environment (MockEnvironment for a dry run), build an Agent with num_episodes and max_num_steps_per_episode, and call it to populate the replay memory. QLearner then takes the model, dataset = ReplayMemoryDataset(), num_train_steps, learning_rate, batch_size and grad_accum_every, and calling it runs training. Note the example values: batch_size 4 with grad_accum_every 16 gives an effective batch of 64, and 10000 train steps against a 1000-episode, 100-step-per-episode collection. Those are illustrative, not tuned. Inference is a separate call: model.get_optimal_actions(video, instructions) with a video tensor of shape (2, 3, 6, 224, 224) and a list of natural-language instructions. There is no CLI, no config file and no checkpointing helper documented in the README, so persistence and resumption are your responsibility.
Where the implementation is thin, and where it will bite
The README is honest about what is unfinished. Gumbel-based action sampling with annealed noise is listed as unchecked, so exploration currently relies on randomising a subset of actions, which the todo marks as done. Beam search for optimal actions is unchecked. Delusional bias, the known failure mode of Q-learning with function approximation, is listed as an open question pending consultation with RL experts, not a solved problem in this code. The conservative regularisation and n-step Q-learning are checked off, but their hyperparameters are not documented in the README beyond what appears in the constructor. The most practical limitation is the environment contract: if your simulator cannot return rewards, next state and done in exactly that tuple form, you are writing an adapter before you write any research. A second constraint is scale. The example uses 1000 episodes of up to 100 steps, and the model is a MaxViT backbone at 224x224 resolution; that is a GPU workload, and nothing in the repository suggests a CPU path is practical.
How this differs from Stable-Baselines3 and other RL libraries
Stable-Baselines3 targets online RL with a stable, documented set of algorithms (PPO, SAC, TD3, DQN) and a Gymnasium environment interface that most simulators already satisfy. It gives you training scripts, callbacks, evaluation helpers and logging. q-transformer does none of that. It implements one paper's method, assumes offline data collected by its own Agent loop, and exposes the model internals rather than a training harness. The trade is control for convenience: with Stable-Baselines3 you get a working pipeline in an afternoon and limited ability to change the Q-function's structure; with q-transformer you get an autoregressive, bin-discretised Q-function you can modify, and you build the surrounding pipeline yourself. If your problem is a standard continuous-control benchmark and you have a simulator, Stable-Baselines3 is the shorter path. If your problem is multi-dimensional discrete action selection from a fixed dataset and you want to experiment with the Q-Transformer formulation, this repository is closer to the thing you are trying to study.
Maintenance, licence and what to check before you depend on it
The project is MIT licensed, which permits commercial use, modification and redistribution provided the copyright notice and permission notice are retained. That is the whole of the licence implication, and it is not legal advice. The repository is not archived, the last push is dated 2026-08-11, and the most recent release listed is 0.4.5 from 2025-05-20, so the code is being touched. Maintenance risk is not abandonment but shape: the todo list contains unchecked items that could change public APIs, including the decoder head variant and potential randomised action ordering with an ordering conditioning token. Pinning a version is sensible. There is no homepage, so the README and the paper are the documentation. Before depending on it, verify three things against your own setup: that your environment can satisfy the BaseEnvironment tuple contract, that your action space quantises sensibly into the action_bins you pick, and that the ReplayMemoryDataset folder format survives whatever storage you point it at.
Editorial conclusion
Adopt q-transformer if you already have a robot or simulator that can implement the BaseEnvironment interface and you want to study autoregressive Q-learning without rebuilding the architecture. Do not adopt it if you need a turnkey training pipeline, a benchmark suite, or a maintained set of pretrained weights: none of those are in the repository. Before committing, verify that your environment can return the exact tuple shapes the README specifies, that your action space can be discretised into the action_bins you choose, and that you can afford the episode count and train steps the example implies. The MIT licence removes most legal friction, but the paper's algorithm and the cited prior work remain the authoritative reference for what the code is supposed to be doing.
Community notes