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.

Modern AI systems often learn from past examples. Imagine a robot learning to walk by watching recordings of previous attempts. But these systems can latch onto accidental patterns in the training data rather than the real causes of good decisions. For example, an AI might learn "speed up when a yellow object appears" simply because yellow taxis happened to appear in fast-driving clips, even though the color is irrelevant. When the system is later deployed and these accidental patterns no longer hold, performance can collapse. We developed CSET, a method that teaches AI to rely on what truly causes good outcomes rather than what merely coincides with them. It works in two complementary ways. First, it generates "what-if" variations of the original data where misleading features change but the meaningful outcome stays the same, exposing the AI to the difference between true causes and lucky coincidences. Second, it gently guides the AI's internal attention toward the genuinely informative signals. When we deliberately broke such deceptive patterns at deployment, the performance of competing AI methods dropped substantially, in some cases by more than half. Our method stayed close to its original level, suggesting that learning real causes helps AI systems hold up better when conditions change.