q-transformer

Q-transformer

Implementation of Q-Transformer, Scalable Offline Reinforcement Learning via Autoregressive Q-Functions, out of Google Deepmind

I will be keeping around the logic for Q-learning on single action just for final comparison with the proposed autoregressive Q-learning on multiple actions. Also to serve as education for myself and the public.

The autoregressive Q-learning formulation has been reproduced by Kotb et al.

Q Learning is Not Yet Scalable

Install

$ pip install q-transformer

Usage

import torch

from q_transformer import (
    QRoboticTransformer,
    QLearner,
    Agent,
    ReplayMemoryDataset
)

# the attention model

model = QRoboticTransformer(
    vit = dict(
        num_classes = 1000,
        dim_conv_stem = 64,
        dim = 64,
        dim_head = 64,
        depth = (2, 2, 5, 2),
        window_size = 7,
        mbconv_expansion_rate = 4,
        mbconv_shrinkage_rate = 0.25,
        dropout = 0.1
    ),
    num_actions = 8,
    action_bins = 256,
    depth = 1,
    heads = 8,
    dim_head = 64,
    cond_drop_prob = 0.2,
    dueling = True
)

# you need to supply your own environment, by overriding BaseEnvironment

from q_transformer.mocks import MockEnvironment

env = MockEnvironment(
    state_shape = (3, 6, 224, 224),
    text_embed_shape = (768,)
)

# env.init()     should return instructions and initial state: Tuple[str, Tensor[*state_shape]]
# env(actions)   should return rewards, next state, and done flag: Tuple[Tensor[()], Tensor[*state_shape], Tensor[()]]

# agent is a class that allows the q-model to interact with the environment to generate a replay memory dataset for learning

agent = Agent(
    model,
    environment = env,
    num_episodes = 1000,
    max_num_steps_per_episode = 100,
)

agent()

# Q learning on the replay memory dataset on the model

q_learner = QLearner(
    model,
    dataset = ReplayMemoryDataset(),
    num_train_steps = 10000,
    learning_rate = 3e-4,
    batch_size = 4,
    grad_accum_every = 16,
)

q_learner()

# after much learning
# your robot should be better at selecting optimal actions

video = torch.randn(2, 3, 6, 224, 224)

instructions = [
    'bring me that apple sitting on the table',
    'please pass the butter'
]

actions = model.get_optimal_actions(video, instructions)

Appreciation

StabilityAI, A16Z Open Source AI Grant Program, and 🤗 Huggingface for the generous sponsorships, as well as my other sponsors, for affording me the independence to open source current artificial intelligence research

Todo

first work way towards single action support
offer batchnorm-less variant of maxvit, as done in SOTA weather model metnet3
add optional deep dueling architecture
add n-step Q learning
build the conservative regularization
build out main proposal in paper (autoregressive discrete actions until last action, reward given only on last)
improvise decoder head variant, instead of concatenating previous actions at the frames + learned tokens stage. in other words, use classic encoder - decoder
- allow for cross attention to fine frame / learned tokens
redo maxvit with axial rotary embeddings + sigmoid gating for attending to nothing. enable flash attention for maxvit with this change
build out a simple dataset creator class, taking in the environment and model and returning a folder that can be accepted by a ReplayDataset
- finish basic environment loop
- store memories to memmapped files in designated folder
- ReplayDataset that takes in folder
  - 1 time step option
  - n-time steps
handle multiple instructions correctly
show a simple end-to-end example, in the same style as all other repos
handle no instructions, leverage null conditioner in CFG library
cache kv for action decoding
for exploration, allow for finely randomizing a subset of actions, and not all actions at once
- also allow for gumbel based sampling of actions, with annealing of gumbel noise
consult some RL experts and figure out if there are any new headways into resolving delusional bias
figure out if one can train with randomized orders of actions - order could be sent as a conditioning that is concatted or summed before attention layers
- offer an improvised variant where the first action token suggests the action ordering. all actions aren't made equal, and some may need to attend to past actions more than others
simple beam search function for optimal actions
improvise cross attention to past actions and states of timestep, transformer-xl fashion (w/ structured memory dropout)
see if the main idea in this paper is applicable to language models here

Citations

@inproceedings{qtransformer,
    title   = {Q-Transformer: Scalable Offline Reinforcement Learning via Autoregressive Q-Functions},
    authors = {Yevgen Chebotar and Quan Vuong and Alex Irpan and Karol Hausman and Fei Xia and Yao Lu and Aviral Kumar and Tianhe Yu and Alexander Herzog and Karl Pertsch and Keerthana Gopalakrishnan and Julian Ibarz and Ofir Nachum and Sumedh Sontakke and Grecia Salazar and Huong T Tran and Jodilyn Peralta and Clayton Tan and Deeksha Manjunath and Jaspiar Singht and Brianna Zitkovich and Tomas Jackson and Kanishka Rao and Chelsea Finn and Sergey Levine},
    booktitle = {7th Annual Conference on Robot Learning},
    year   = {2023}
}

@inproceedings{dao2022flashattention,
    title   = {Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},
    author  = {Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher},
    booktitle = {Advances in Neural Information Processing Systems},
    year    = {2022}
}

@inproceedings{Kumar2023MaintainingPI,
    title   = {Maintaining Plasticity in Continual Learning via Regenerative Regularization},
    author  = {Saurabh Kumar and Henrik Marklund and Benjamin Van Roy},
    year    = {2023},
    url     = {https://api.semanticscholar.org/CorpusID:261076021}
}

@misc{Rubin2024,
    author  = {Ohad Rubin},
    url     = {https://medium.com/@ohadrubin/exploring-weight-decay-in-layer-normalization-challenges-and-a-reparameterization-solution-ad4d12c24950}
}

@article{Farebrother2024StopRT,
    title   = {Stop Regressing: Training Value Functions via Classification for Scalable Deep RL},
    author  = {Jesse Farebrother and Jordi Orbay and Quan Ho Vuong and Adrien Ali Taiga and Yevgen Chebotar and Ted Xiao and Alex Irpan and Sergey Levine and Pablo Samuel Castro and Aleksandra Faust and Aviral Kumar and Rishabh Agarwal},
    journal = {ArXiv},
    year   = {2024},
    volume = {abs/2403.03950},
    url    = {https://api.semanticscholar.org/CorpusID:268253088}
}

@inproceedings{Hasselt2015DeepRL,
    title   = {Deep Reinforcement Learning with Double Q-Learning},
    author  = {H. V. Hasselt and Arthur Guez and David Silver},
    booktitle = {AAAI Conference on Artificial Intelligence},
    year    = {2015},
    url     = {https://api.semanticscholar.org/CorpusID:6208256}
}

Metadata

Release files for q-transformer 0.5.0

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for q-transformer 0.5.0
File	Size	Uploaded
q_transformer-0.5.0.tar.gz	23.3 kB	Aug 11, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for q-transformer 0.5.0
File	Interpreter	ABI	Platform	Reset
q_transformer-0.5.0-py2.py3-none-any.whl 24.7 kB Aug 11, 2026	Python 3, Python 2	none	any	Details

Total release size: 48.0 kB

Release files / q_transformer-0.5.0.tar.gz

Download URL	q_transformer-0.5.0.tar.gz
Size	23.3 kB
Tags	Source
SHA-256 checksum How to use checksums	`a806205765874fcf94f9f25d5e18dad641ba2ff8d410543352b00cfa10fd2f1a`
BLAKE2b-256 checksum How to use checksums	`74a4b5edc199c03dbac416dc1daec7e0887b0123fbe3b4cbcdd8977d860746d3`
Upload date	Aug 11, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`

Release files / q_transformer-0.5.0-py2.py3-none-any.whl

Download URL	q_transformer-0.5.0-py2.py3-none-any.whl
Size	24.7 kB
Tags	Python 2 Python 3
SHA-256 checksum How to use checksums	`fe3c7710804d5e47376bdcb06bb95294d44b7a5b5e68cd5a9169f63bae6ec718`
BLAKE2b-256 checksum How to use checksums	`36eebcd3f37392143bf0f9f0d48108a71c9188d77fb2ca1a393670ba63aa37de`
Upload date	Aug 11, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`