Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unified token representations for sequential decision models

Zhuojing Tian, Yushu Chen

TL;DR

The paper addresses inefficiencies in transformer-based offline RL caused by triple-token (RTG, state, action) representations that inflate sequence length and compute. It introduces Unified Token Representation (UTR), which encodes each timestep with a single token, and develops two backbones, Unified Decision Transformer (UDT) and Unified Decision Conv (UDC), the latter using a Gated CNN-based decision module. The authors provide a theoretical argument and bounds for tighter generalization under merged tokens via Rademacher complexity, and substantiate the approach with experiments on MuJoCo and AntMaze showing competitive or superior performance with substantially reduced computation. This work demonstrates that UTR generalizes across architectures and offers a scalable foundation for future large decision models in offline RL. The practical impact lies in enabling efficient, real-time or resource-constrained decision-making without sacrificing performance.

Abstract

Transformers have demonstrated strong potential in offline reinforcement learning (RL) by modeling trajectories as sequences of return-to-go, states, and actions. However, existing approaches such as the Decision Transformer(DT) and its variants suffer from redundant tokenization and quadratic attention complexity, limiting their scalability in real-time or resource-constrained settings. To address this, we propose a Unified Token Representation (UTR) that merges return-to-go, state, and action into a single token, substantially reducing sequence length and model complexity. Theoretical analysis shows that UTR leads to a tighter Rademacher complexity bound, suggesting improved generalization. We further develop two variants: UDT and UDC, built upon transformer and gated CNN backbones, respectively. Both achieve comparable or superior performance to state-of-the-art methods with markedly lower computation. These findings demonstrate that UTR generalizes well across architectures and may provide an efficient foundation for scalable control in future large decision models.

Unified token representations for sequential decision models

TL;DR

The paper addresses inefficiencies in transformer-based offline RL caused by triple-token (RTG, state, action) representations that inflate sequence length and compute. It introduces Unified Token Representation (UTR), which encodes each timestep with a single token, and develops two backbones, Unified Decision Transformer (UDT) and Unified Decision Conv (UDC), the latter using a Gated CNN-based decision module. The authors provide a theoretical argument and bounds for tighter generalization under merged tokens via Rademacher complexity, and substantiate the approach with experiments on MuJoCo and AntMaze showing competitive or superior performance with substantially reduced computation. This work demonstrates that UTR generalizes across architectures and offers a scalable foundation for future large decision models in offline RL. The practical impact lies in enabling efficient, real-time or resource-constrained decision-making without sacrificing performance.

Abstract

Transformers have demonstrated strong potential in offline reinforcement learning (RL) by modeling trajectories as sequences of return-to-go, states, and actions. However, existing approaches such as the Decision Transformer(DT) and its variants suffer from redundant tokenization and quadratic attention complexity, limiting their scalability in real-time or resource-constrained settings. To address this, we propose a Unified Token Representation (UTR) that merges return-to-go, state, and action into a single token, substantially reducing sequence length and model complexity. Theoretical analysis shows that UTR leads to a tighter Rademacher complexity bound, suggesting improved generalization. We further develop two variants: UDT and UDC, built upon transformer and gated CNN backbones, respectively. Both achieve comparable or superior performance to state-of-the-art methods with markedly lower computation. These findings demonstrate that UTR generalizes well across architectures and may provide an efficient foundation for scalable control in future large decision models.
Paper Structure (17 sections, 15 equations, 1 figure, 2 tables)

This paper contains 17 sections, 15 equations, 1 figure, 2 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Offline Reinforcement Learning with Transformer Variants
  4. Gated CNNs for Conditional Sequence Modeling
  5. Methodology
  6. Preliminaries
  7. Unified Token Representation
  8. Gated CNN Decision Module
  9. Theoretical Analysis
  10. Rademacher complexity.
  11. Connection to generalization.
  12. Discussion of assumptions and their plausibility.
  13. Experiments
  14. Experimental Setup
  15. Results and Analysis
  16. ...and 2 more sections

Figures (1)

  • Figure 1: $\textbf{Left:}$ Decision Transformer. $\textbf{Middle:}$ Decision Convformer. $\textbf{Right:}$ Gated CNN Decision Module.