Bi-directional Recurrence Improves Transformer in Partially Observable Markov Decision Processes

TL;DR

This work tackles partial observability in reinforcement learning by integrating a bi-directional recurrence layer into a transformer-based Q-network (DBGFQN), replacing the conventional feed-forward network with a single BiGRU layer to capture both forward and backward temporal information. The approach combines multi-head self-attention with BiGRU memory, achieving about a $25\%$ reduction in parameter count relative to typical transformer baselines while enhancing sample efficiency. Across 23 POMDP environments, DBGFQN substantially outperforms both transformer-based and recurrence-based baselines, with average gains of $87.39\%$ over DTQN, $96.14\%$ over DGFQN, and $482.04\%$ over DRQN. These results demonstrate the practical value of bi-directional recurrence for partially observable tasks and suggest a route toward more compact, efficient architectures suitable for real-world deployment on limited hardware.

Abstract

In real-world reinforcement learning (RL) scenarios, agents often encounter partial observability, where incomplete or noisy information obscures the true state of the environment. Partially Observable Markov Decision Processes (POMDPs) are commonly used to model these environments, but effective performance requires memory mechanisms to utilise past observations. While recurrence networks have traditionally addressed this need, transformer-based models have recently shown improved sample efficiency in RL tasks. However, their application to POMDPs remains underdeveloped, and their real-world deployment is constrained due to the high parameter count. This work introduces a novel bi-recurrent model architecture that improves sample efficiency and reduces model parameter count in POMDP scenarios. The architecture replaces the multiple feed forward layers with a single layer of bi-directional recurrence unit to better capture and utilize sequential dependencies and contextual information. This approach improves the model's ability to handle partial observability and increases sample efficiency, enabling effective learning from comparatively fewer interactions. To evaluate the performance of the proposed model architecture, experiments were conducted on a total of 23 POMDP environments. The proposed model architecture outperforms existing transformer-based, attention-based, and recurrence-based methods by a margin ranging from 87.39% to 482.04% on average across the 23 POMDP environments.

Figures (8)

• Figure 1: Architecture of the proposed DBGFQN model.
• Figure 2: The effect of varying the number of FFN layers in the Transformer model.
• Figure 3: The effect of uni-directional recurrent layers in the Transformer model.
• Figure 4: The effect of bi-directional recurrent layers in the Transformer model.
• Figure 5: Comparitive analysis on POMDP environments 1-8.