Return-conditioned sequence modeling

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 2205.15241 §3.1. The approach extends the decision transformer idea. That work has no card in this wiki and is only named here 2205.15241 §3.1. Card: Multi-Game Decision Transformers.
• 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 2409.05493 §3.1, §3.2. The paper does not sample a return at play time. Card: DexDiff.

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 2205.15241 §3.4. The diffusion variant is trained with a noise-prediction mean squared error instead 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.

Option B of the return choice is the authors' appendix alternative. They say the two return choices perform similarly 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 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 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 2409.05493 §3.2.

The paper names Decision Transformer as a baseline 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

Evidence limits: the headline 126% human-normalised IQM is author-reported, and the sentence does not name the model size 2205.15241 §1 2205.15241 fig. 1. The 5-seed result is for the 40M model only: IQM 0.95, standard deviation 0.05 2205.15241 app. B.1. All other results are single runs 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 2409.05493 §4.1. The real-robot setup differs from simulation, and the paper reports no sim-to-real transfer result 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.

# 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

# 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 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" 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 2409.05493 §3.2 2409.05493 app. A.4. Record the choices you make.

# 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 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 2205.15241 app. B.1; 121 logits in the algorithm 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 2205.15241 §4.5; about 500,000 transitions in the appendix 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 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 2205.15241 §3.1. The diffusion paper derives return-to-go values from goal-conditioned grasp poses of a plan, and gives no formula 2409.05493 §3.1. A value tuned for one does not transfer to the other.

Papers

• Multi-Game Decision Transformers: 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: 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.