lucidrains/q-transformer

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

Q‑Transformer – offline RL with autoregressive Q‑functions

What it is – A PyTorch implementation of the Q‑Transformer paper (Google DeepMind, 2023). It treats the Q‑function as an autoregressive model that predicts a sequence of discrete actions, enabling scalable offline reinforcement learning for robotic tasks.

Key components

  • QRoboticTransformer – a vision‑transformer backbone (based on MaxViT) that ingests video frames and optional text instructions, and outputs Q‑values for a discretised action space.
  • QLearner – training loop that performs Q‑learning on a replay‑memory dataset generated by the agent.
  • Agent – runs the model in a user‑provided environment (subclass of BaseEnvironment) to collect trajectories.
  • ReplayMemoryDataset – a simple dataset wrapper for the stored (state, action, reward, next‑state, done) tuples.

How to get it

pip install q-transformer

Typical workflow

  1. Define an environment – implement BaseEnvironment or use the supplied MockEnvironment which expects video tensors and instruction embeddings.
  2. Create the model – instantiate QRoboticTransformer with vision hyper‑parameters (e.g., a small MaxViT) and specify the number of discrete actions and bins.
  3. Collect data – Agent(model, environment, ...)() runs episodes, stores transitions in a replay buffer, and writes them to disk.
  4. Train – QLearner(model, dataset=ReplayMemoryDataset(), ...)() runs Q‑learning (supports dueling heads, n‑step returns, gradient accumulation, etc.).
  5. Infer actions – after training, call model.get_optimal_actions(video, instructions) to obtain the best discrete action sequence for new video‑instruction pairs.

What makes it special

  • Implements the autoregressive Q‑function idea, allowing the model to handle multi‑dimensional discrete action spaces efficiently.
  • Includes optional dueling architecture, n‑step returns, and a conservative regularisation placeholder.
  • Designed for robotics: the model processes multi‑camera video ((3, 6, 224, 224)) and language instructions.
  • The repo tracks a detailed TODO list that mirrors ongoing research directions (e.g., beam‑search decoding, cross‑attention to past actions, gumbel‑based exploration).

Who might use it

  • Researchers experimenting with offline RL for manipulation or embodied agents.
  • Practitioners who need a ready‑to‑run baseline for discretised robotic control from video and language.
  • Anyone looking to reproduce or extend the Q‑Transformer paper (the implementation has been reproduced by Kotb et al., 2024).

Citations – The README provides BibTeX entries for the original Q‑Transformer paper and related works (FlashAttention, continual learning regularisation, etc.).

All details above are taken directly from the repository’s README.

Related

  • Dispatch
    Decision Transformer Integration in Hugging Face Transformers

    Hugging Face has integrated the Decision Transformer, an offline reinforcement learning method that treats RL as a conditional-sequence modeling problem, into the transformers library and Hugging Face Hub.

  • Dispatch
    Deep Q-Learning with Space Invaders

    Hugging Face explains Deep Q-Learning (DQN), a reinforcement learning method that uses neural networks to approximate Q-values for high-dimensional state spaces like Atari games.

  • Dispatch
    Training Decision Transformers for Offline Reinforcement Learning

    Hugging Face provides a guide and implementation for training an offline Decision Transformer from scratch to solve the HalfCheetah environment using the transformers Trainer and a custom data collator.

  • Dispatch
    Hugging Face Introduction to Q-Learning

    Hugging Face provides a technical guide to Q-Learning, an off-policy value-based reinforcement learning method that uses a Q-table to determine optimal state-action pairs.

  • Project
    lucidrains/transfusion-pytorch

    A PyTorch implementation of the Transfusion architecture that unifies next-token prediction for text and flow matching for continuous modalities like images in a single model.