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

TL;DR

This work addresses the challenge of applying Decision Transformers to image-based environments with multi-discrete action spaces by introducing Multi-State Action Tokenisation (M-SAT). M-SAT expands each multi-discrete action into multiple action tokens and injects preceding state information as auxiliary embeddings, enhancing the transformer’s attention-based reasoning over state-action and action-action relationships. Empirical results on ViZDoom Deadly Corridor and My Way Home show that M-SAT improves performance and reduces variance relative to the baseline, with notable interpretability benefits as attention heads attend to individual action tokens. The approach requires no pretraining or heavy architectural changes and demonstrates improved sample efficiency and robustness, suggesting broader applicability to RL tasks with structured action spaces. Future work includes deeper exploration of positional encoding effects and extending evaluations to additional environments and state spaces.

Abstract

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.

Figures (15)

• Figure 1: Tokenisation of the RTG, state and multi-discrete actions in the M-SAT Decision Transformer, with the proposed multi-token actions
• Figure 2: Sample of image-based states from My Way Home (Left) with multiple rooms and corridors leading to a green vest goal state and Deadly Corridor (Right) with attacking adversaries and a green vest goal state
• Figure 3: Evaluation results over 50 runs and 5 seeds for DC (left) and MWH (right). See Appendix \ref{['sec:appendix-dt-implementation-details']} for training details
• Figure 4: Simple map outline of ViZDoom's Deadly Corridor scenario showing the locations of the agent, enemies and the goal.
• Figure 5: Sample Factory architecture overview - Source: petrenko2020sf