Multi-State-Action Tokenisation in Decision Transformers for Multi-Discrete Action Spaces
This addresses a specific bottleneck for researchers using Decision Transformers in reinforcement learning with complex action spaces, though it is an incremental improvement focused on tokenisation.
The paper tackles the problem of Decision Transformers struggling with multi-discrete action spaces in image-based environments by proposing Multi-State Action Tokenisation (M-SAT), which outperforms baseline Decision Transformers on ViZDoom scenarios like Deadly Corridor and My Way Home without extra data or computational costs.
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.