PaperlayerView exact revision on arXiv

Source-linked AI summary

Decision Titan: Test-Time Training for Long-Term Memory in Offline Reinforcement Learning

Jude Waide, Robert Lieck

arXiv:2610.01513v1cs.AIcs.LG

TL;DR

Long-term dependencies are difficult for sequential models because RNNs have hidden-state and gradient limits, while Transformers’ quadratic attention constrains context. This paper applies Titans’ test-time training memory to Decision Transformer for offline RL and evaluates it in X-Maze. Decision Titan learns dependencies beyond its context window and can extrapolate beyond training lengths, but recall depends on environment cues and embedding choices.

  • Problem

    RNN hidden-state expressivity and vanishing gradients, and Transformers’ quadratic attention cost, limit practical modeling of long-term dependencies.

    Source: §1 · p. 1

  • Method

    Decision Titan combines Titans blocks with Decision Transformer and is evaluated on long-term memory tasks in X-Maze.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    Decision Titan learns dependencies with ranges at least 20x its context window and extrapolates to waiting phases at least 1.7x its training range, with outcomes depending on environment encoding and embedding strategy.

    Source: §4.1.2 · p. 5, §4.3.2 · p. 7, §5 · p. 9

  • Takeaways & Limitations

    Sequential information can be remembered when environment cues reliably associate with what must be recalled, while memory struggles with event order otherwise.

    Source: §4.1.2 · p. 6

    The long-term memory module struggles to remember event order unless the environment provides reliable cues associated with the information to recall.

    Source: §4.1.2 · p. 6

Abstract

from arXiv · show

Long-term dependencies remain a major challenge for sequential decision-making in the field of AI: RNNs suffer from vanishing gradients and the limited expressivity of vector-based hidden states, whilst Transformer-based models are limited by the quadratic scaling of attention. Recent work has proposed tackling this problem with the Test-Time Training (TTT) framework, which stores episodic memories in the parameters of a neural network through gradient descent at both train and test-time. This approach has seen success in the domain of Natural Language Processing, however, to the best of our knowledge it has not yet been applied to the domain of Reinforcement Learning (RL), nor has there been a study analysing how this memory practically functions. In this paper, we study the potential of the TTT framework for offline RL by augmenting a Decision Transformer with TTT layers, dubbed the Decision Titan. We analyse performance and properties of the model in the X-Maze environment, an extension of T-Maze designed to test sequential memory, and investigate how the memory mechanism learns by visualising gate values over time. Our key findings are that Decision Titan can learn long-term dependencies with ranges 20x longer than the context window, generalises to lengths 1.7x the training data, but crucially temporal generalisation depends on the time embeddings used, and the ability to learn long-term dependencies depends on how the relevant information is encoded.

1. Introduction

Long-term dependencies require identifying which past observations matter and retaining them, but RNNs and Transformers have practical limits on memory. This work evaluates Test-Time Training for offline RL with Decision Titan and X-Maze, including how its memory behaves.

1. Introduction

  • RNNs face finite hidden-state expressivity and vanishing gradients, while Transformers’ quadratic attention cost limits context length and modeled dependency range.Linear-scaling attention still has limited context length and does not address the need for episodic long-term memory.Source: §1 · p. 1
  • The paper introduces Decision Titan by merging Titans with Decision Transformer to investigate Test-Time Training for long-term dependencies in reinforcement learning.Decision Transformer applies Transformers to offline RL as a sequence-modelling task.Source: §1 · p. 2
  • The study uses X-Maze to test ordered recall over long horizons and examines extrapolation, positional embeddings, gate values, and momentum ablations.X-Maze extends T-Maze, and the analyses investigate how the memory mechanism operates and where it may be limited.Source: §1 · p. 2
  • Decision Titan can learn dependencies with ranges 20x longer than its context window, while momentum helps propagate dependencies during training.The study also identifies catastrophic forgetting as a likely major limitation and finds that sequential recall depends on the conditions and information encoding.Source: §1 · p. 2

2. Related Work

Prior work develops TTT and Decision Transformer extensions, while this paper’s related-work discussion situates its approach among efficiency and memory designs. These directions address distinct trade-offs and leave room for combinations and further study.

2.1. Test-Time Training Developments

  • TTT efficiency has a quality-efficiency trade-off in chunk size; treating chunks as unordered can improve GPU utilisation but is ill-suited when intra-chunk order matters.The paper identifies improving TTT efficiency as important for wider adoption.Source: §2.1 · p. 2
  • Sparse product-key memory reduces access complexity to the square root of storage slots, while this paper uses MLP memory for tasks not requiring large capacity.The authors suggest sparse access as a promising direction for improving extrapolation.Source: §2.1 · p. 2

2.2. Decision Transformer Extensions

  • Other Decision Transformer extensions enable online fine-tuning, planning-conditioned actions, or Atari-scale training, but do not tackle long-term memory and are described as orthogonal.The paper identifies combining them, especially forward- and backward-looking mechanisms, as future study.Source: §2.2 · p. 2
  • RATE stores memory tokens in the context window and ELMUR uses a separate vector store, with the paper arguing its approach has better theoretical capacity scaling and expressivity.The passage does not provide the underlying comparison details.Source: §2.2 · p. 2

3. Methodology

Decision Titan applies Titans blocks to Decision Transformer while retaining its trajectory representation, tokenization, action head, and training objectives. Its methodology also specifies sequential full-trajectory training, positional embeddings, and X-Maze’s controlled ordered-memory task.

3.1. Decision Titan

  • Decision Titan uses Titans blocks with Decision Transformer’s trajectory representation, tokenization, action head, and training objectives.The method also uses chunked inference, with divergences in inference approach, training strategy, and positional embeddings.Source: §3.1 · p. 3

3.1.1. ARCHITECTURE

  • Memory as a Layer separates long- and short-term memory, making it the selected Titans architecture for evaluating TTT in reinforcement learning.The authors chose MAL over MAG for its simpler separation; they note it performed worst among Titans architectures in Behrouz et al.’s reported comparison, while outperforming the other tested techniques.Source: §3.1.1 · p. 3
  • Each chunk passes through a TTT memory layer followed by self-attention; attention spans 3k+n tokens, with n persistent memory tokens prepended to chunks.The architecture diagram describes trajectory embeddings, chunking, and MAL blocks; the TTT update occurs every b tokens.Source: Figure 1 · p. 3, §3.1.1 · p. 3

3.1.2. TRAINING

  • Training processes full episode trajectories in ordered chunks because randomly sampled length-k subsequences do not capture long-range dependency dynamics.Memory state is reset to its learned initial state for each trajectory and is neither shared across batch episodes nor persisted across training iterations; full-trajectory training scales memory and compute with trajectory length.Source: §3.1.2 · p. 3

3.1.3. EMBEDDINGS

  • The model tests global sinusoidal, chunk-relative sinusoidal, and learned axial embeddings that combine intra-chunk position with chunk position.Chunked embeddings reuse the same positional embedding in each chunk, while axial embeddings provide both global and local positional information.Source: §3.1.3 · p. 4

3.2. X-Maze Environment

  • X-Maze extends T-Maze to require following instruction sequences, with the number of possible instructions configurable to test sequential memory beyond a single bit.The agent receives instructions, waits, then repeats them; this nonphysical setup isolates memory capacity and permits control of the number of timesteps to remember.Source: §3.2 · p. 4
  • X-Maze varies its repeat cues: a single-bit signal arrives at the end of waiting, one-hot signal identifies each repeat step, and one-hot repeat also provides ids during instruction presentation.Observations encode instructions as I-dimensional one-hot vectors, and the environment uses a discrete action for each instruction plus a null waiting action with penalties for invalid actions or incorrect repeats.Source: §3.2 · p. 4

4. Experiments

Across X-Maze experiments, Decision Titan learned sequential dependencies when the environment provided usable recall cues, and generalized beyond its training lengths. Results also depended on embedding choice and momentum; the analyses suggest memory learning is stochastic and vulnerable to forgetting.

4. Experiments

  • The experiments varied waiting-phase lengths, embedding strategies, and memory-head counts; each X-Maze dataset contained 10,000 trajectories from a perfect hard-coded policy.Generalization tests evaluated waiting lengths from 5 to 200, while gate analyses compared momentum and no-momentum models using chunked embeddings and four heads.Source: §4 · p. 5, §4.3.1 · p. 7, §4.4.1 · p. 8

4.1.1. SETUP

  • The sequential-memory test used instruction sequences of lengths [6,10], waiting periods [6,10], and ten possible instructions, with evaluation at length 10.The experiment crossed all nine combinations of three environment encodings and three embedding strategies; varying sequence and waiting lengths prevented global token position from identifying the repeat phase.Source: §4.1.1 · p. 5

4.1.2. RESULTS AND DISCUSSION

  • Across Table 3's nine combinations, one-hot repeat achieved perfect mean returns, one-hot signal was close to perfect, and single signal scored [-5.76, -5.6] versus random chance (-9).Performance was the same across embedding strategies for each environment encoding.Source: §4.1.2 · p. 5, Table 3 · p. 6
  • Reliable cues that associate observations with later recall support sequential memory; without such cues, DTi struggles to remember event order.Random waiting lengths made global token positions uninformative, while the one-hot settings supplied associations between instruction identity, position, or the repeat signal.Source: §4.1.2 · pp. 5–6

4.2.2. RESULTS AND DISCUSSION

  • More MLP heads consistently reduced training time, while runs across all 27 configurations either stayed near chance or rapidly reached return 0.Table 4 reports the fraction of training completed before first reaching return 0, averaged across runs and embeddings.Source: Table 4 · p. 6, §4.2.2 · p. 6, Figure 3 · p. 6

4.3.2. RESULTS AND DISCUSSION

  • DTi generalized to at least 1.7x its training lengths; chunked embeddings trained on [96,100] achieved perfect mean returns through waiting lengths of 175.Chunked embeddings generalized most, followed by axial and global embeddings, but some configurations also failed at shorter lengths.Source: Figure 4 · p. 7, §4.3.2 · p. 7

4.4.2. RESULTS AND DISCUSSION

  • Without momentum, success rates were 0% at waiting lengths 50 and 100, compared with 100% and 40% with momentum.Successful momentum runs maintained a high momentum gate in one head, around 0.95.Source: §4.4.2 · p. 8
  • A run with forgetting gate α ≈0.3 retained only about 0.056% of a first-chunk update's magnitude after 21 steps; the authors conclude momentum helps propagate dependencies.Successful no-momentum runs instead had a lower α, around 0.05; performance decay was unchanged, leaving catastrophic forgetting the likely limit on longer generalization.Source: §4.4.2 · pp. 8–9

5. Discussion

Decision Titan extends dependency learning beyond the context window while reducing memory-operation scaling, but its training setup and sequential recall remain limited. The discussion identifies catastrophic forgetting as a likely extrapolation constraint and points to sparse memory updates as a direction for further study.

5. Discussion

  • DTi learns dependencies at least 20x the context window and extrapolates to 1.7x training episode lengths, showing a clear advantage over standard DT.Its total complexity is O(k2 + p), where k is context length and p is the number of MLP memory parameters.Source: §5 · p. 9
  • Memory reads and writes are constant in sequence length, yielding O(k2 + p), but DTi is not directly compared with RATE or ELMUR.The discussion notes that parameter-based storage may hold more information than memory tokens only in theory; practical storage scaling remains untested.Source: §5 · p. 9
  • Training on entire episodes constrains batch size, slows training relative to DT, and complicates efficient setup through chunk-dependent memory states.The authors also report limited ability to remember sequential information and suggest a different memory module or relative positional embeddings as possible research directions.Source: §5 · p. 9
  • At 200 steps in X-Maze, DTi falls off, while ELMUR has been shown to generalise to at least 1,000,000 T-Maze steps; the models were not directly compared.The discussion identifies catastrophic forgetting from repeated chunk updates as the most likely extrapolation limit and proposes sparse memory access for future RL study.Source: §5 · p. 9
  • The analysis mechanistically explains momentum, training-time dependency propagation, and memory-head effects, motivating scheduled gate biasing and other extensions.Other proposed directions include persistent memory tokens and evaluation in noisier environments.Source: §5 · p. 9

6. Conclusion

The paper presents Decision Titan and X-Maze as tools for studying long-term sequential memory in offline RL, concluding that TTT is a viable foundation with substantial room to improve. Its preliminaries situate TTT, Titans, and Decision Transformer as the components behind this approach.

6. Conclusion

  • Decision Titan merges Decision Transformer with Titans to address long-term dependencies in RL, and X-Maze tests ordered recall over long time horizons.The paper describes X-Maze as a controlled environment for evaluating the order of long-term dependencies.Source: §6 · p. 9
  • The paper concludes that TTT is a viable foundation for long-term memory in offline RL, while substantial room for improvement remains.The authors present their mechanistic insights and X-Maze as tools for future research.Source: §6 · p. 9

A. Detailed Preliminaries

  • The preliminaries introduce TTT as a long-term memory framework, Titans as its combination with Transformers, and Decision Transformer as a Transformer policy learner for offline RL.Source: A. Detailed Preliminaries · p. 11

A.0.1. TEST-TIME TRAINING

This section introduces test-time training as a fast-weight memory mechanism, describes how Titans combine it with Transformer attention, and summarizes Decision Transformer and experiment setup.

A.0.1. TEST-TIME TRAINING

  • TTT replaces a conventional hidden state with neural-network fast weights, updated by an inner-loop self-supervised reconstruction objective during training and inference.The module learns to compress a sequence into its hidden state.Source: §A.0.1 · p. 11
  • Sun et al. (2025)'s formulation adds projections and an input-dependent loss gate for selective storage and query-based retrieval, with O(1) storage and retrieval for stateless M.The learned gate controls gradient-update magnitude; retrieval depends on queries matching keys associated with relevant prior information.Source: §A.0.1 · pp. 11–12

A.0.2. TITANS

  • Titans combines TTT long-term memory with Transformer short-term attention, placing memory as context, gate, or layer in its MAC, MAG, and MAL variants.The memory module is appended before attention in MAC, gates attention in MAG, and precedes attention in MAL.Source: §A.0.2 · p. 12

A.0.3. DECISION TRANSFORMER

  • Decision Transformer learns offline policies by autoregressively predicting actions from return, observation, and action sequences within a fixed k-timestep context.The method casts policy learning as sequence modelling; its context-window size limits the dependencies it can capture.Source: §A.0.3 · pp. 12–13

B. Hyperparameters

C. Additional Experiments

The T-Maze experiment tests whether Decision Titan can solve a task whose dependency exceeds both models' context window.

C.1. T-Maze

  • The comparison is designed to test whether TTT memory lets DTi solve tasks that standard DT cannot at the same context-window size.Source: §C.1 · p. 13
  • With a context length of 5, DTi solves T-Maze while DT does not, despite the task's 6-timestep dependency range.Source: Figure 7 · p. 13

C.1.1. SETUP

  • Both models use a context window of 5 and face a 5-timestep waiting phase, making long-term memory necessary to solve the task.The experiment uses the single-signal X-Maze setting with instruction sequence length 1 and instruction count 2, equivalent to T-Maze.Source: §C.1.1 · p. 13

C.1.2. RESULTS AND DISCUSSION

In the T-Maze experiment, DTi reaches mean return 0 while DT remains at approximately -0.5.

C.1.2. RESULTS AND DISCUSSION

  • Approximately -0.5 for DT versus 0 for DTi in mean return across 25 runs; the authors say this result requires learning long-term dependencies.The 25 runs comprise five runs per seed across five seeds.Source: §C.1.2 · p. 13

C.2. Control Tasks

The MuJoCo control experiments check whether adding DTi's memory mechanism has unintended effects where long-term dependencies are not expected.

C.2. Control Tasks

  • The authors do not expect long-term memory to benefit MuJoCo tasks without long-term dependencies, and use them as a sanity check for unintended effects.Source: §C.2 · p. 13

C.2.1. SETUP

  • The comparisons include standard DT and DTi without TTT layers, separating the memory mechanism from DTi's other architectural choices.Source: §C.2.1 · p. 14
  • Tests use three MuJoCo environments with expert and medium datasets; DTi uses chunked embeddings, while DT uses learned global timestep embeddings.Each model uses the same base parameters; evaluation averages five episodes on five seeds per environment.Source: §C.2.1 · p. 14

C.2.2. RESULTS AND DISCUSSION

MuJoCo control results show no significant differences among the three models, while the additional gate visualizations include both successful and failed X-Maze runs.

C.2.2. RESULTS AND DISCUSSION

  • MuJoCo control results show no significant differences among the three models; full-state visibility explains why disabling memory does not reduce return versus standard DT.The authors note that memory-disabled DTi may suffer in non-Markovian environments, where prior context matters.Source: §C.2.2 · p. 14

C.3. Additional Gate Visualisation Figures

  • Figure 8 shows a failed 100-step-waiting run with momentum, where no head learns to maintain momentum or minimize the forget gate.The appendix identifies Figure 8 as a failure case with momentum.Source: Figure 8 · p. 14, §C.3 · p. 14
  • Figure 9 shows a successful run with momentum at a 50-step waiting phase, with observations reported as matching the 100-step case.The appendix identifies Figure 9 as a success case with momentum.Source: Figure 9 · p. 14, §C.3 · p. 14
  • Figure 10 shows a successful 50-step-waiting run without momentum, where one head learns a small forgetting-gate value.The appendix identifies Figure 10 as a success case without momentum.Source: Figure 10 · p. 14, §C.3 · p. 14

Training

The training material lists task-specific return targets and an effective batch size, alongside a table of mean returns on control tasks.

Training

  • Return targets are 6000 for HalfCheetah, 3600 for Hopper, and 5000 for Walker2D.The effective batch size is 100.Source: Training · p. 15
  • The effective batch size is 100.Source: Training · p. 15
  • Table 7 is labeled as reporting mean return on control tasks.Source: Table 7 · p. 15

Related papers

Loading 2610.01513v1…