Biasing the Future: Gaussian Attention for Sequential Decision-Making

ICLR 2026 Conference Submission22545 Authors

Keywords: Decision Transformer, Offline reinforcementlearning

Abstract: Transformers have emerged as powerful sequence models for offline reinforcement learning (RL), but their reliance on purely self-attention mechanisms can limit their ability to capture fine-grained local dependencies and Markovian dynamics present in many RL datasets. In this work, we introduce a modified Decision Transformer architecture that incorporates a Gaussian-biased masked causal attention mechanism. By augmenting attention scores with a distance-aware bias, the model adaptively emphasizes temporally local relationships while still retaining the ability to capture long-range dependencies through self-attention. Experimental results on benchmark offline RL tasks show that our Gaussian-biased Decision Transformer achieves achieves state-of-the-art performance and notable gains over the standard DT, particularly in environments with strong Markovian structure. This demonstrates the importance of explicitly encoding locality into attention mechanisms for sequential decision-making.

Primary Area: reinforcement learning

Submission Number: 22545