Unifying Counterfactual Data Augmentation and Architectural Inductive Biases in Offline Reinforcement Learning

Siyu Wang, Xiaocong Chen, Mingming Gong, Yong Li, Quan Z. Sheng, Lina Yao

Keywords: Causality, Offline Reinforcement Learning

Abstract: Transformer-based models have recently achieved strong results in offline reinforcement learning by casting decision-making as sequence modeling. However, when trained purely on fixed datasets, they are prone to causal confusion: reliance on spurious correlations that predict reward in the data but do not reflect the true causal mechanisms of the environment. This issue is exacerbated by the weak inductive bias of Transformers, whose global attention is not aligned with the Markovian and causal dependencies of decision processes. We introduce the Unified Causal Transformer (UCF), a framework that strengthens both the data and the model with causal consistency. On the data side, UCF employs a causal reward model to abduce exogenous factors and a counterfactual state generator to produce reward-preserving augmentations, yielding counterfactual trajectories that expose causal variability absent in observational data. On the model side, UCF integrates a causally structured hybrid architecture that combines disentangled convolutional encoders for local dynamics with supervised attention for global reasoning, guiding the model to allocate representational capacity according to true causal dependencies. We evaluate UCF on two distinct sequential decision-making tasks—robotic control and recommendation—and demonstrate consistent gains in robustness and generalization over Transformer-based baselines. These results highlight the importance of causal consistency in both data and architecture for reliable offline policy learning.

Primary Area: other topics in machine learning (i.e., none of the above)

Submission Number: 5443