Unified token representations for sequential decision models

TL;DR

This paper introduces a unified token representation for sequential decision models in reinforcement learning, reducing complexity and improving scalability while maintaining or enhancing performance.

Contribution

It proposes a novel unified token representation that merges multiple elements into a single token, reducing sequence length and complexity in RL models.

Findings

UTR reduces sequence length and model complexity.

UTR-based models achieve comparable or better performance.

Theoretical analysis indicates improved generalization.

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.