Reward-Preserving Counterfactual State Editing for Offline Reinforcement Learning

Transformer sequence models such as Decision Transformer can learn strong offline policies from logged trajectories, but they often suffer from causal confusion: reliance on spurious correlations that predict reward in the data but do not reflect the true causal mechanisms of the environment. We propose CSET (Counterfactual State Editing Transformer), which improves robustness in strictly offline reinforcement learning without learning environment transition dynamics. On the data side, CSET fits a causal reward model as a conditional variational autoencoder and uses a counterfactual state generator to propose minimally edited observations whose predicted reward matches the factual reward, under a normalized move-band constraint and an acceptance gate that enforce plausibility and reward consistency; augmentation replaces only the observation token to avoid synthetic successor transitions. On the model side, CSET uses a causally structured hybrid transformer: modality-specific convolutional encoders process return-to-go, state, and action streams, and a final attention block is softly supervised so action prediction focuses on its direct causal parents. Experiments on D4RL locomotion, AntMaze, and offline recommendation benchmarks show consistent gains within the DT family, and CSET remains substantially more robust than strong value-based and DT baselines under injected spurious distractors.