Multi-Game Decision Transformers

In plain words

Most game-playing AI systems learn one game at a time. This paper asks whether one neural network can learn many old Atari video games at once, using only recordings of other players and without playing live while it learns.

The trick is to treat each game as a long sequence of tokens, the same way a language model treats text. Each step holds the picture cut into small squares, a target score to aim for, the action taken and the reward received. A transformer, the kind of model behind chat assistants, learns to predict the next token. When it plays, it first asks for a high but plausible target score, then picks an action that fits that target.

The authors report that bigger models do better, and that a model trained on many games can pick up a new game quickly from a small amount of data. They release the pre-trained models and code. All results are their own.

TL;DR

• Offline, single-set-of-weights agent for 41 training Atari games (46 in total, 5 held out), trained on DQN replay trajectories §3.3.
• Decision-transformer sequence model over image patches, return, action and reward tokens §3.1 §3.2.
• Expert-level behaviour comes from sampling a high target return before each action, a classifier-guided idea borrowed from language generation §3.4.
• Author-reported: 126% human-normalised IQM across 41 games, scaling with size, best fine-tuning among the methods tested §1 §4.5.
• Released code is inference-only and differs from the main text in several places. Read the Paper vs code section before reimplementing.

Problem

Offline RL normally needs one policy per task, and multi-environment training on Atari has been studied less than single-game or single-environment RL §1. Training data mixes expert and non-expert play §3.3. Plain behaviour cloning imitates that mix and rarely produces the best episodes §4.7. Temporal- difference (TD) methods such as C51 and CQL become unstable as the model grows in the multi-game setting §4.4. The authors want a single model that is trained on all data but can still act at expert level, and that scales §1.

Method

Sequence formulation §3.1. The agent sees observation o^t, emits action a^t and gets reward r^t. The model is a decoder-only transformer trained with next-token prediction over the sequence \langle ..., o^t_1, ..., o^t_M, \hat{R}^t, a^t, r^t, ... \rangle, where \hat{R}^t is the target return for the remainder of the episode §3.1. The order puts the return before the action, so the model can sample a return first and then an action. Observation patches are not predicted, because the authors judge them non-discrete and expensive to generate §3.1.

Tokenization §3.2 app. B.1.

1. Each 84x84 grayscale frame app. B.1 is cut into a 6x6 grid of 14x14 pixel patches, giving M = 36 patches §3.2 app. B.1.
2. Each patch is linearly projected and gets a trainable position encoding §3.2.
3. Returns are uniformly quantised into integer bins from -20 to 100 §3.2.
4. Actions are already discrete §3.2.
5. Rewards are mapped to a ternary class {-1, 0, +1} §3.2. See Paper vs code for the four-class version in the release.
6. One timestep is therefore 36 patch tokens, one return, one action and one reward token, 39 in all. Four frames make a 156-token sequence §4.1.

Attention mask app. B.1. Ordinary causal masking would block the patches of one frame from attending to each other. The authors allow full attention within the observation tokens of a timestep and keep causal masking across timesteps.

Training data §3.3. The data come from an existing set of Atari trajectories collected while DQN agents trained §3.3. Per training game there are 2 runs of 50 checkpoints, each with 1M environment steps §3.3. Across 41 games that totals 4.1 billion steps and about 160 billion tokens §3.3. There is no filtering or balancing §3.3. Five games are held out for transfer tests §3.3 app. B.2.

Training objective §3.1 §3.2. Cross-entropy on the next return, action and reward tokens. Training is ordinary next-token prediction and needs no change to the usual infrastructure. The expert-level inference step below only changes sampling at play time §3.4.

Expert action inference §3.4 alg. 1 app. B.3. The authors model "expert-level" as a binary variable whose probability grows with the return, P(\mathrm{expert}^t | R^t, ...) \propto \exp(\kappa R^t), with inverse temperature \kappa. Applying Bayes' rule to the learned return distribution gives a biased sampling rule. The biased logits add \kappa times a normalised return, (R - R_{\mathrm{low}})/(R_{\mathrm{high}} - R_{\mathrm{low}}) §3.4. The main text sets \kappa = 10 §3.4 app. B.3. The return is resampled at every step app. B.3, which avoids storing past target returns app. B.3.

# Expert return and action inference  [alg. 1] [§3.4]
# Input:  environment env, trained model; kappa = 10
# Output: one action per environment step

def play_episode(model, env):
    # Start from the observation tokens and the history  [alg. 1]
    context = start_context(env)
    while not env.done:
        # Return logits over the 121 bins from -20 to 100  [alg. 1] [app. B.3]
        return_logits = model.return_logits(context)
        # Add kappa * normalised return to each logit  [§3.4]
        normalised = (return_bins + 20) / 120
        biased = return_logits + kappa * normalised
        # Sample a target return; top 85th percentile  [app. B.3]
        target = sample(biased, temperature=1, top_percentile=85)
        # Sample an action conditioned on that target  [alg. 1]
        action_logits = model.action_logits(context, target)
        action = sample(action_logits, temperature=1, top_percentile=85)
        reward, observation = env.step(action)
        # Slide the window; the history is regenerated, not stored  [app. B.3]
        context = slide(context, observation, action, reward)

The authors also describe a simpler alternative. It draws N = 128 return samples from the model, keeps the largest one, and conditions the action on it app. B.3. This needs no \kappa. The reported settings for it are return temperature 0.75, no return percentile cut-off, and action logits in the top 50th percentile at temperature 1 app. B.3. The authors say the two choices perform similarly app. B.3. The released code implements this alternative, not the \kappa version (see Paper vs code).

Baselines §4.2 app. C. Behaviour cloning is the same transformer with the return token removed app. C. C51 is online multi-game DQN with an Impala CNN §4.2. CQL adds a behaviour-cloning term to C51's categorical loss §4.2. Representation baselines are CPC, BERT and ACL, each with its own pre-training objective and then a behaviour-cloning head §4.2 app. C.

Image augmentation app. B.5. Random crop (4 zero pixels of padding on each side, then crop back to 84x84) and random rotation worked best. The authors also tried channel permutation, pixel permutation and flips. The same random parameters are used for each window, and augmentation is applied in both pre-training and fine-tuning app. B.5.

Evaluation app. B.4. Atari 2600 with the preprocessing of Agarwal et al. 2020, 18 shared discrete actions, no sticky actions, and each game scored over 16 rollouts app. B.4. Observations are JPEG-compressed in the dataset and at evaluation app. B.4.

Hyperparameters

Results

Author-reported. Setup: 41 training Atari games §3.3. Human-normalised score (HNS) and inter-quartile mean (IQM) across games §4.1 app. B.4. Baselines are behaviour cloning, online C51, offline CQL, and representation methods §4.2. Single training seed except for DT-40M app. B.1.

Raw per-game scores for the best multi-game models (41 games) tab. 2 app. I. DT (200M), BC (200M), online DQN (10M) and CQL (60M), as reported:

Paper vs code

Official code: google-research/google-research/multi_game_dt, linked from the authors' project page. The page is the project site named in the paper §1 [abstract, footnote]. The repository is Apache-2.0 LICENSE. The folder holds a README (marked "work in progress"), an AUTHORS file and one Colab notebook Multi_game_decision_transformers_public_colab.ipynb. The notebook loads a pre-trained checkpoint from a Google Cloud Storage path Multi_game_decision_transformers_public_colab.ipynb, which was not downloaded for this card. code_commit is the head of the whole google-research repository at clone time. The folder itself was not versioned separately.

What is confirmed by the code:

• Architecture. The released model uses d_model 1280, 10 layers and d_model // 64 heads, which gives 20 heads. That matches DT-200M app. B.1 Multi_game_decision_transformers_public_colab.ipynb Multi_game_decision_transformers_public_colab.ipynb.
• Patches. A stride-14 convolution gives a 6x6 grid of 36 patches on 84x84 inputs §3.2 Multi_game_decision_transformers_public_colab.ipynb.
• Causal mask. Observation tokens within a timestep attend to each other and the rest is causal app. B.1 Multi_game_decision_transformers_public_colab.ipynb.
• Alternative return selection. The released call matches the appendix alternative: 128 samples, return temperature 0.75, action top-50 percentile app. B.3 Multi_game_decision_transformers_public_colab.ipynb.

Differences:

• Main-text expert inference is not what is released. The paper's main method uses \kappa = 10 and a top-85 action cut-off §3.4 app. B.3. The released call sets the optional bias weight opt_weight = 0, which disables it Multi_game_decision_transformers_public_colab.ipynb, and uses the appendix alternative instead.
• The kappa term exists in code but is off. When opt_weight > 0, the code adds opt_weight times a linear ramp from 0 to 1 over the return bins Multi_game_decision_transformers_public_colab.ipynb. With 121 bins this equals the normalised return in §3.4, so opt_weight plays the role of \kappa.
• Action cut-off. The paper uses the top 85th percentile in the main method app. B.3. The released call uses 50 Multi_game_decision_transformers_public_colab.ipynb. The code keeps logits strictly above the given percentile Multi_game_decision_transformers_public_colab.ipynb.
• One return token per sequence. The paper puts a return token in every timestep §3.1. The released model sets single_return_token = True Multi_game_decision_transformers_public_colab.ipynb. The return mask keeps only the first return token Multi_game_decision_transformers_public_colab.ipynb and the return loss uses only that token Multi_game_decision_transformers_public_colab.ipynb. At inference only position 0 of the return sequence is written Multi_game_decision_transformers_public_colab.ipynb. The per-step resampling in Algorithm 1 is not what this code does.
• Reward classes. The paper says three reward classes {-1, 0, +1} §3.2. The release embeds four classes: no reward, positive, terminal and negative Multi_game_decision_transformers_public_colab.ipynb Multi_game_decision_transformers_public_colab.ipynb. The encoder only emits 0, 1 and 3; class 2 (terminal) is not produced by it.
• Return definition. The Colab defines return-to-go as positive rewards minus negative rewards, with a terminal reward turning the value into -1 Multi_game_decision_transformers_public_colab.ipynb. §3.2 says returns are sums of future rewards §3.2 §3.1. The training data pipeline is not released, so the definition used in training is not verified.
• Return bins. Returns are clipped to [-20, 100] and shifted to 0..120, so 121 tokens Multi_game_decision_transformers_public_colab.ipynb. The paper says 120 buckets in appendix B.1 and 121 logits in Algorithm 1 app. B.1 alg. 1.
• Dropout. The released model sets dropout 0.1 Multi_game_decision_transformers_public_colab.ipynb. The paper does not state a dropout rate.
• Pixel scaling. Pixels are divided by 255 before the patch convolution Multi_game_decision_transformers_public_colab.ipynb. The paper does not state this.
• Training code and training hyperparameters. The repository has no training code, no Jaxline configuration, no optimiser code and no data pipeline. The LAMB settings, learning rates and batch sizes in §4.1 cannot be checked against code.
• Inference cost. The released optimal_action calls the model twice per step, once for return logits and once for action logits Multi_game_decision_transformers_public_colab.ipynb Multi_game_decision_transformers_public_colab.ipynb.

No discrepancy was checked for the baseline implementations (C51, CQL, CPC, BERT, ACL, Impala); they are not in the repository.

Pitfalls

• Two inference recipes. Main text: kappa = 10, top-85 cut-off §3.4 app. B.3. Release: no kappa, max of 128 return samples at temperature 0.75, top-50 cut-off Multi_game_decision_transformers_public_colab.ipynb. The paper does not say which recipe produced the 126% headline or the per-game raw scores. Pick one and report it.
• Percentile semantics. Code keeps logits above the p-th percentile Multi_game_decision_transformers_public_colab.ipynb. Top 85th percentile therefore keeps about the top 15% of logits. A reader reading "top 85th percentile" as "keep the top 85%" would get a different policy.
• Return is written to position 0 only. The release writes one sampled return into position 0 Multi_game_decision_transformers_public_colab.ipynb. Algorithm 1 resamples the return at every step alg. 1. How the released wrapper carries the return between steps is not verified; it uses a window of return values Multi_game_decision_transformers_public_colab.ipynb Multi_game_decision_transformers_public_colab.ipynb.
• Return bins: 120 or 121. The paper says 120 buckets app. B.1 and 121 logits alg. 1. The code uses 121 Multi_game_decision_transformers_public_colab.ipynb. Use 121 for the closed interval -20..100.
• Fine-tuning data size. §4.5 says 1% is "1M steps" from a 50M-step dataset §4.5. App. B.2 says roughly 500,000 transitions app. B.2. 1% of 50M is 500,000, so B.2 is consistent with the dataset size and §4.5 is probably a typo.
• Single seed. Only DT-40M was rerun, over 5 seeds: IQM 0.95 with sd 0.05 app. B.1. All other results are single runs §4.1. The authors call the variance small compared with the gaps they report app. B.1; no significance tests are given.
• Expert data is a filtered subset. The expert set keeps the top 10% of trajectories by episodic return §4.8 app. F. It is used only for the expert-data comparison; the main method trains on all data §4.8.
• Reward tokens and the terminal convention. The reward embedding has four classes, and the return target can be -1 after a terminal reward (see Paper vs code). Copy the convention you train with; the paper does not define it.
• Dropout. A rate of 0.1 appears only in the released model Multi_game_decision_transformers_public_colab.ipynb. The paper says nothing about it.
• Context is short. Four frames per sequence in all experiments §4.1. No longer memory is evaluated.
• Compression. Observations are JPEG-compressed in the dataset and at evaluation app. B.4. Reproduced scores depend on this step.
• Compute. 64 TPUv4 per model, 8 days for DT-200M app. B.1 tab. 1. Reimplementation needs a large accelerator budget; the paper gives no single-GPU setting.
• Figure and equation numbers. The numbers here follow LaTeX source order, not the compiled PDF: figure 1 is the IQM figure, figure 2 the teaser, and appendix tables are 1 (model sizes) and 2 (raw scores). Check against the compiled PDF before quoting a number.
• Atari dataset. The training trajectories come from Agarwal et al. 2020 §3.3 app. B.4. The repository does not include them; they must be obtained separately.
• Expert action inference is independent of training. Only the sampling changes; a model trained with the usual cross-entropy can be used with either recipe §3.4. Keep the training objective and the inference recipe separate when you debug.

Lineage

• Builds on: decision transformer, treating offline RL as sequence modelling §3.1 (Chen et al. 2021, cited in the paper as chen2021decision; arXiv id not given in the source text). Upside-Down RL §2. GPT-2 architecture §4.1.
• Uses ideas from: discriminator-guided generation in language models (cited as krause2020gedi) for the expert classifier §3.4; nucleus sampling (cited as holtzman2019curious) for the action cut-off app. B.3.
• Compared against: C51 (Bellemare et al. 2017), CQL (Kumar et al. 2020), Impala (Espeholt et al. 2018), CPC, BERT, ACL, behaviour cloning §4.2.
• Concurrent work: Gato (cited as gato2022deepmind) trains a transformer generalist on near-optimal data and needs expert prompts at inference §2.
• Superseded by: not stated in content on this wiki.