---
kind: concept
slug: return-conditioned-sequence-modeling
title: Return-conditioned sequence modeling
summary: >-
  Conditioning a policy on a target return: a transformer predicts return, observation, action and reward tokens and samples a high target return at play time; a diffusion policy can take a return-to-go embedding instead. Two papers, both with author-reported results. The main design choices are how often the return appears, how it is sampled and how it is defined.
tags: [offline-rl, sequence-modeling, atari, diffusion-policy]
papers: ["2205.15241", "2409.05493"]
related: [imitation-learning, multi-task-learning, transfer-learning, diffusion-models, offline-reinforcement-learning]
status: draft
updated_at: 2026-10-25
---

## Overview

The problem: learn a policy from logged trajectories only, with no live play
during training, and ask it for good behaviour at play time. The shared idea is
to give the policy a target return, the value it should aim for, as an input,
and to ask for a high one at play time.

Two papers do this in different ways:

- **Token-sequence model.** A decoder-only transformer models a trajectory as
  one sequence of observation patches, a target return, the action and the
  reward for each step, in that order. The sequence is trained with
  next-token prediction, so the model can be asked for a return first and an
  action second [arxiv:2205.15241 §3.1]. The approach extends the decision
  transformer idea. That work has no card in this wiki and is only named here
  [arxiv:2205.15241 §3.1]. Card: [Multi-Game Decision Transformers](/papers/2205.15241).
- **Diffusion policy with return-to-go input.** A transformer diffusion policy
  predicts short action chunks. Its conditioning includes a return-to-go
  embedding, derived from a plan written by a vision-language model
  [arxiv:2409.05493 §3.1, §3.2]. The paper does not sample a return at play
  time. Card: [DexDiff](/papers/2409.05493).

The invariant holds only for the token-sequence model: training is ordinary
cross-entropy on the next return, action and reward tokens, with no change to
the usual infrastructure. Only the sampling at play time differs between
recipes [arxiv:2205.15241 §3.4]. The diffusion variant is trained with a
noise-prediction mean squared error instead [arxiv:2409.05493 §3.2], and its
return-to-go is an input, not a token.

## Variants

### Token-sequence model (decision transformer style)

The table below covers this variant only.

| Design axis | Option A | Option B |
|---|---|---|
| Return tokens | One per timestep, as in the paper's formulation [arxiv:2205.15241 §3.1] | One per sequence, as in the released code; written to position 0 only (see the card's Paper vs code section) |
| Return choice at play time | Bias the return logits toward high returns with a weight kappa = 10, keep logits above the 85th percentile [arxiv:2205.15241 §3.4] [arxiv:2205.15241 app. B.3] | Draw 128 return samples, keep the largest, sample with temperature 0.75; action logits kept above the 50th percentile [arxiv:2205.15241 app. B.3] |
| Reward tokens | Three classes {-1, 0, +1} [arxiv:2205.15241 §3.2] | Four classes in the release (see the card's Paper vs code section) |

Option B of the return choice is the authors' appendix alternative. They say
the two return choices perform similarly [arxiv:2205.15241 app. B.3]. The
released code implements Option B, not the kappa version.

Behaviour cloning is the same model with the return token removed. It is the
baseline, not a return-conditioned variant [arxiv:2205.15241 app. C].

### Return-to-go as a diffusion conditioning input

The return-to-go values are not sampled. The plan is turned into a series of
values by computing goal-conditioned grasp poses, and the paper gives no
formula, reward scale or reference frame for them [arxiv:2409.05493 §3.1].
An embedding of the values is added to the transformer's embedding layer, so
the policy models actions given observation and return-to-go, not the joint
distribution with future states [arxiv:2409.05493 §3.2].

The paper names Decision Transformer as a baseline [arxiv:2409.05493 §4.1].
Its return-to-go is not the same kind of quantity as the target return above;
see Contradictions.

## When to use which

| Situation | Use | Why |
|---|---|---|
| Logged data mixes expert and weak play and you want the best episodes | Token-sequence model with a high target return | Behaviour cloning imitates the mix and rarely produces the best episodes [arxiv:2205.15241 §4.7]. Per game, the return-conditioned model beat behaviour cloning in 31 of 41 games [arxiv:2205.15241 §4.7] [arxiv:2205.15241 fig. 8] |
| You already have expert-only data | Either; return conditioning on the full data is the paper's comparison | In the paper, the full-data model beat behaviour cloning on the top 10% expert subset [arxiv:2205.15241 §4.8]. Values are only in a figure |
| Many tasks, little data on a new one | Pre-train on many tasks, fine-tune on about 1% of the new data | Fine-tuning ran 100k steps at learning rate 1e-4 [arxiv:2205.15241 §4.1]. The pre-trained model was best on each of the 5 held-out games [arxiv:2205.15241 §4.5] |
| Temporal-difference training is unstable for you | Sequence modelling, no TD target | The paper reports that C51 and CQL become unstable as the model grows in the multi-game setting [arxiv:2205.15241 §4.4] |
| Limited accelerator budget | A smaller model; the 200M model used 64 TPUv4 for 8 days [arxiv:2205.15241 app. B.1] [arxiv:2205.15241 tab. 1] | Smaller models were trained in 1 to 2 days [arxiv:2205.15241 tab. 1]. No single-GPU setting is given |
| A high-level planner already gives a goal-derived value per step, and you need a low-level action chunk | Diffusion policy with return-to-go input | Under the same VLM planner, the average success rate was 0.70 against 0.23 for Diffusion Policy-T across four simulated scenes [arxiv:2409.05493 tab. 1]. The comparison is between whole systems; the reviewed sections report no ablation of the return-to-go input [arxiv:2409.05493 §4.2] |
| Little demonstration data or compute | Neither variant as reported | The diffusion paper says its model needs much demonstration and is computationally expensive [arxiv:2409.05493 §5]. The 200M sequence model used 64 TPUv4 for 8 days [arxiv:2205.15241 app. B.1] |
| Need long memory | Not covered by the evidence | Only 4 frames per sequence were used [arxiv:2205.15241 §4.1]. The diffusion policy uses a 2-step observation history [arxiv:2409.05493 app. A.6] |

Evidence limits: the headline 126% human-normalised IQM is author-reported, and
the sentence does not name the model size [arxiv:2205.15241 §1] [arxiv:2205.15241 fig. 1].
The 5-seed result is for the 40M model only: IQM 0.95, standard deviation 0.05
[arxiv:2205.15241 app. B.1]. All other results are single runs
[arxiv:2205.15241 §4.1]. The diffusion results are author-reported on four
simulated scenes with 35 trials each, and on a real UR5e with 10 trials per
scenario [arxiv:2409.05493 §4.1]. The real-robot setup differs from simulation,
and the paper reports no sim-to-real transfer result [arxiv:2409.05493 §4.5].

## Implementation sketch

Same format as the other concept pages: pseudocode in our own words, one
citation per step. The token-sequence recipe is shown first; the alternative
differs only in the return-choice block. The diffusion variant follows.

```pseudocode
# Build one timestep of the training sequence  [§3.2]
def timestep_tokens(frame, target_return, action, reward):
    # Cut the 84x84 frame into a 6x6 grid of 14x14 patches  [§3.2]
    patches = patchify(frame, size=14)            # 36 tokens
    # Quantise the return into 121 bins from -20 to 100  [§3.2] [app. B.1]
    ret = clip(round(target_return), -20, 100) + 20
    # Map the reward to a class  [§3.2]
    rew = reward_class(reward)
    # The return comes before the action, so it can be sampled first  [§3.1]
    return patches + [ret, action, rew]           # 39 tokens per timestep
```

```pseudocode
# Play time: choose a return, then an action  [§3.4] [alg. 1]
def act(model, context, kappa=10):
    # Logits over the 121 return bins  [alg. 1]
    logits = model.return_logits(context)
    # Main recipe: add kappa times the normalised return  [§3.4]
    logits = logits + kappa * normalised_bins    # values in [0, 1]
    # Keep logits above the 85th percentile, then sample  [app. B.3]
    target = sample(keep_top(logits, percentile=85))
    # Action conditioned on the chosen return  [alg. 1]
    action_logits = model.action_logits(context, target)
    return sample(keep_top(action_logits, percentile=85))
```

The alternative replaces the return block with: draw 128 return samples at
temperature 0.75 and keep the largest; then sample the action from logits
kept above the 50th percentile [arxiv:2205.15241 app. B.3]. Read the
percentile rule as "keep logits above the p-th percentile", which is how the
released code applies it, not as "keep p percent of logits"
[arxiv:2205.15241 app. B.3] (see the card's Pitfalls).

The diffusion variant trains on a noise-prediction loss. The paper writes the
noisy action as the clean action plus noise, with no scaling coefficient, and
does not give the noise schedule or the number of denoising steps
[arxiv:2409.05493 §3.2] [arxiv:2409.05493 app. A.4]. Record the choices you
make.

```pseudocode
# Diffusion variant: one training step  [§3.2] [eq. 4]
def diffusion_training_step(obs, actions, return_to_go):
    # Encode each observation step on its own  [app. A.4]
    obs_emb = visual_encoder(obs)
    # Embed the return-to-go values as an extra input  [§3.2]
    rtg_emb = rtg_embedding(return_to_go)
    # Draw a denoising step k; the sampling law is not stated  [§3.2]
    k = random_step(batch_size=actions.shape[0])
    eps = sample_normal(shape=actions.shape)
    # Noisy actions as written in eq. 4, with no scaling coefficient  [eq. 4]
    eps_hat = noise_net(obs_emb, rtg_emb, actions + eps, k)
    # Mean squared error between predicted and true noise  [eq. 4]
    return mse(eps_hat, eps)
```

## Contradictions

Within the sequence-model paper, several numbers disagree. Across the two
papers, the word "return" does not mean the same thing; neither is resolved here.

- Action percentile: the main recipe uses 85 [arxiv:2205.15241 app. B.3]; the
  released call uses 50 (see the card's Paper vs code section).
- Return bins: 120 buckets in the appendix [arxiv:2205.15241 app. B.1]; 121
  logits in the algorithm [arxiv:2205.15241 alg. 1]. The released code uses
  121 for the closed interval -20 to 100.
- Fine-tuning data: "1M steps" from a 50M-step dataset [arxiv:2205.15241 §4.5];
  about 500,000 transitions in the appendix [arxiv:2205.15241 app. B.2]. The
  appendix figure is consistent with 1% of 50M.
- Return choice: the main text samples a return at every step
  [arxiv:2205.15241 alg. 1]; the released code writes one return per sequence
  (see the card's Paper vs code section).
- Meaning of the return: the sequence-model paper defines the target return
  as a sum of future rewards [arxiv:2205.15241 §3.1]. The diffusion paper
  derives return-to-go values from goal-conditioned grasp poses of a plan,
  and gives no formula [arxiv:2409.05493 §3.1]. A value tuned for one does not
  transfer to the other.

## Papers

- [Multi-Game Decision Transformers](/papers/2205.15241): one transformer trained offline on Atari trajectories from 41 games, with return-conditioned action sampling; author-reported scaling and fine-tuning results.
- [DexDiff: Towards Extrinsic Dexterity Manipulation of Ungraspable Objects in Unrestricted Environments](/papers/2409.05493): a VLM plan supplies return-to-go values to a transformer diffusion policy for grasping large flat objects; author-reported simulated and real success rates.