Retrieval-augmented Decision Transformer

Context. In-context reinforcement learning trains a transformer offline on many tasks so that, on a new task, it improves from episode to episode by reading its own earlier episodes, with no weight updates. The established methods, Algorithm DistillationAD (Laskin et al., 2023): a transformer trained to predict the next action of a learning RL agent given its earlier episodes, so that it imitates the learning process itself (Laskin et al., 2023) and the Decision-Pretrained TransformerDPT (Lee et al., 2023): a transformer trained to predict the optimal action for a query state given a context of interactions from the same task (Lee et al., 2023), keep whole episodes in context, so the input grows with episode length and attention cost with its square, and they have been shown mostly on bandits and small grid worlds. Schmied et al. (2025) keep past experience outside the model, in a vector index searched by nearest neighbors as in retrieval-augmented language models, and give a Decision TransformerDT (Chen et al., 2021): a transformer that treats offline RL as sequence modeling, predicting the next action from recent states, actions, rewards and the reward still to come only the stored pieces of trajectory that resemble its current situation (Retrieval-Augmented Decision Transformer, RA-DT).

An agent does not need its whole history to act well, only the few past moments that resemble the present and what followed them, and a frozen encoder with nearest-neighbor search can find those, so a context of 50 to 100 steps can replace thousands. On grid worlds with sparse rewards this beats the whole-episode methods while training up to 7 times faster. On robot control and video games it raises scores on held-out tasks, but there neither it nor any baseline improves across episodes, so in-context learning beyond grid worlds is not shown.

Setting. A trajectory $\tau=(s_0,a_0,r_0,\hat R_0,\dots)$ lists states, actions, rewards and returns-to-go$\hat R_t=\sum_{t'\ge t}r_{t'}$, the reward still to be collected from step $t$ to the end of the episode $\hat R_t$. A Decision Transformer $\pi_\theta$ predicts $a_t$ from the last $C$the context length in steps; 50 for RA-DT, 100 on the two larger grids steps of states, returns-to-go, actions and rewards, trained by cross-entropy or squared error on offline data, and at test time it is conditioned on a high target return. The training data come from PPOproximal policy optimization, a standard on-policy RL method; here one PPO agent is trained per task and all its transitions are kept agents trained separately on each training task, with every transition kept, so they mix exploratory and good behavior. Evaluation follows earlier in-context RL work: on a held-out task the trained agent plays a fixed number of episodes, called trials, each finished episode becomes available as context, and success means reward rising across trials. AD and DPT see two whole episodes (200 to 2,000 steps), and so does the plain DT on the grid worlds (up to 6,400 tokens; 50 steps on the robotics and video-game benchmarks); RA-DT sees $C$ steps plus $2C$ retrieved steps.

Method. Every $C$ steps along every stored trajectory, the index receives a key, an embedding $g(\tau_{t-C:t})$ of the last $C$ steps, and a value, the sub-trajectory $\tau_{t-C:t+C}$, the same $C$ steps followed by the next $C$: at test time only the past is available to form a query, and the continuation shows how a similar situation played out. The embedding $g$the trajectory encoder that maps a sub-trajectory to a vector for search; a DT pretrained on the same data, or a frozen language model applied after a fixed input map is either a DT pretrained on the same data or a domain-agnostic encoder: FrozenHopfielda fixed map (Paischer et al., 2022) that sends any input vector into the token-embedding space of a frozen language model by attention over its vocabulary maps each token into the token-embedding space of a frozen language model,$$\mathrm{FH}(x)=E^\top\mathrm{softmax}(\beta EPx),$$with $E$the language model's $v\times d_{\rm LM}$ token-embedding matrix, $v$ the vocabulary size the token-embedding matrix, $P$a fixed random matrix with Gaussian entries mapping inputs of size $d_{\rm in}$ to $d_{\rm LM}$ a fixed random matrix and $\beta$ a scale, and frozen BERTan encoder-only language model of 110M parameters (Devlin et al., 2019) runs on the result; the query averages the encoded state tokens. Retrieval takes the top $l$the number of nearest keys first retrieved, before rescoring keys by cosine similarity$q^\top u/(\|q\|\,\|u\|)$ between the query vector $q$ and a key vector $u$ $s_{\rm rel}$ to the query, rescores them by$$s_{\rm ret}=\underbrace{s_{\rm rel}}_{\text{relevance}}+\alpha\,\underbrace{s_{\rm u}}_{\text{utility}},$$with both scores normalized to $[0,1]$ and $\alpha$the weight of utility against relevance in the retrieval score; 1 by default a weight, and keeps the top $k$the number of retrieved sub-trajectories given to the model: one on the grid worlds and video games, four on the robotics tasks. The utility $s_{\rm u}$ is 1 for a sub-trajectory from the query's own task (0 otherwise) during training, when the index holds data from many tasks, and the return of the episode it came from at test time, when the index holds only the agent's own episodes on the new task. Retrieved sub-trajectories enter through cross-attention layers placed after each self-attention layer. During training the index is the training set itself, minus the query's own trajectory, and retrieval is regularizedwhy: retrieving near-identical sub-trajectories teaches the model to copy its context, which leaves it unable to produce high-reward actions at test time: near-duplicate entries (cosine similarity above 0.98) are removed, retrievals above that similarity are skipped, and 20% of query tokens are dropped, so that the retrieved context is relevant without being a near copy of the input.

Main findings (all empirical). On three partially observed grid worlds with sparse rewards, RA-DT ends its trials with the highest reward at every grid size and is the only method that improves across trials on the procedurally generated mazes of MazeRunner15×15 procedurally generated mazes with up to three goals to visit in order, distance-sensor and position observations and 400-step episodes (Steps 1 and 2). The domain-agnostic encoder comes close to the domain-specific one on the smallest grids, falls behind on the 20$\times$20 grids and does better on the hardest Key-DoorDark Key-Door: a grid world in which the agent must find an invisible key and then an invisible door, so a task is a pair of cells grid, so retrieval can work without an encoder trained for the domain. The gain comes from retrieving sub-trajectories that are both similar to the current one and from the same task (Step 3), and training is up to 7 times faster on the largest grid because the sequence no longer grows with the episode (Step 4). On Meta-Worlda benchmark of simulated robot-arm manipulation tasks such as opening a door or inserting a peg, DMControla suite of simulated locomotion and balance tasks such as walker-walk and cartpole-balance and Procgenprocedurally generated video games with image observations; here 12 games for training and 4 held out it has the highest average scores on held-out tasks, but no method improves across trials there (Step 5).

Caveats. Results are over three seeds, and most grid-world numbers are read from plots. The baselines are limited to two episodes of context and, on the Dark grids, 40 trials, and with more trials AD and DPT continue to improve (Step 1). RA-DT adds an index, an encoder, a reweighting that uses task labels during training, three retrieval regularizers and per-environment settings, such as blending the query with a random key on the robotics tasks and replacing retrieval with same-task samples in the first 10 steps of episodes on Dark-Room and Dark Key-Door. Its gains on held-out robotics tasks come mostly from single tasks, and in-context improvement is shown only in partially observed grid worlds and mazes, where memory of past episodes is required to solve the task at all.