# Multi-State-Action Tokenisation in Decision Transformers for Multi-Discrete Action Spaces

arXiv: https://arxiv.org/abs/2407.01310 · HTML: https://mlsanity.com/paper/2407.01310/

## Key facts

- No citing papers in the map yet (data through 7 Oct 2026).
- Placed in Decision transformers and sequence modeling, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 33 of its 41 references that are in the map.
- First posted on arXiv 1 Jul 2024.

## Placed in

- Offline RL and imitation learning from data / [Decision transformers and sequence modeling](https://mlsanity.com/offlinerl/decision-transformers-and-sequence-modeling/)

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Voyager: An Open-Ended Embodied Agent with Large Language Models](https://mlsanity.com/paper/2305.16291/) (2023)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [MineDojo: Building Open-Ended Embodied Agents with Internet-Scale Knowledge](https://mlsanity.com/paper/2206.08853/) (2022)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [FNet: Mixing Tokens with Fourier Transforms](https://mlsanity.com/paper/2105.03824/) (2021)
- [Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks](https://mlsanity.com/paper/2306.13831/) (2023)
- [The Impact of Positional Encoding on Length Generalization in Transformers](https://mlsanity.com/paper/2305.19466/) (2023)
- [StarCraft II: A New Challenge for Reinforcement Learning](https://mlsanity.com/paper/1708.04782/) (2017)
- [Visual Reinforcement Learning with Imagined Goals](https://mlsanity.com/paper/1807.04742/) (2018)
- [In-context Reinforcement Learning with Algorithm Distillation](https://mlsanity.com/paper/2210.14215/) (2022)
- [RvS: What is Essential for Offline RL via Supervised Learning?](https://mlsanity.com/paper/2112.10751/) (2021)
- [ViZDoom: A Doom-based AI Research Platform for Visual Reinforcement Learning](https://mlsanity.com/paper/1605.02097/) (2016)
- [Prompting Decision Transformer for Few-Shot Policy Generalization](https://mlsanity.com/paper/2206.13499/) (2022)
- [Sample Factory: Egocentric 3D Control from Pixels at 100000 FPS with Asynchronous Reinforcement Learning](https://mlsanity.com/paper/2006.11751/) (2020)
- [For SALE: State-Action Representation Learning for Deep Reinforcement Learning](https://mlsanity.com/paper/2306.02451/) (2023)
- [Elastic Decision Transformer](https://mlsanity.com/paper/2307.02484/) (2023)
- [Action Space Shaping in Deep Reinforcement Learning](https://mlsanity.com/paper/2004.00980/) (2020)
- [StARformer: Transformer with State-Action-Reward Representations for Visual Reinforcement Learning](https://mlsanity.com/paper/2110.06206/) (2021)
- [Structural Guidance for Transformer Language Models](https://mlsanity.com/paper/2108.00104/) (2021)
- [Unleashing the Power of Pre-trained Language Models for Offline Reinforcement Learning](https://mlsanity.com/paper/2310.20587/) (2023)
- [Chain-of-Thought Predictive Control](https://mlsanity.com/paper/2304.00776/) (2023)
- [What Improves the Generalization of Graph Transformers? A Theoretical Dive into the Self-attention and Positional Encoding](https://mlsanity.com/paper/2406.01977/) (2024)

## Abstract (arXiv)

Decision Transformers, in their vanilla form, struggle to perform on image-based environments with multi-discrete action spaces. Although enhanced Decision Transformer architectures have been developed to improve performance, these methods have not specifically addressed this problem of multi-discrete action spaces which hampers existing Decision Transformer architectures from learning good representations. To mitigate this, we propose Multi-State Action Tokenisation (M-SAT), an approach for tokenising actions in multi-discrete action spaces that enhances the model's performance in such environments. Our approach involves two key changes: disentangling actions to the individual action level and tokenising the actions with auxiliary state information. These two key changes also improve individual action level interpretability and visibility within the attention layers. We demonstrate the performance gains of M-SAT on challenging ViZDoom environments with multi-discrete action spaces and image-based state spaces, including the Deadly Corridor and My Way Home scenarios, where M-SAT outperforms the baseline Decision Transformer without any additional data or heavy computational overheads. Additionally, we find that removing positional encoding does not adversely affect M-SAT's performance and, in some cases, even improves it.