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

TL;DR

This paper introduces Multi-State Action Tokenisation (M-SAT), a novel method for improving Decision Transformers in multi-discrete action spaces by disentangling actions and incorporating auxiliary state information, leading to better performance and interpretability.

Contribution

The paper proposes M-SAT, a new tokenisation approach that enhances Decision Transformers for multi-discrete action spaces by disentangling actions and adding auxiliary state info.

Findings

M-SAT outperforms baseline Decision Transformers in ViZDoom environments.

Removing positional encoding can improve M-SAT performance.

M-SAT enhances interpretability of individual actions within attention layers.

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.