Reward-Preserving Counterfactual State Editing for Offline Reinforcement Learning

Siyu Wang, Xiaocong Chen, Mingming Gong, Yong Li, Quan Z. Sheng, Lina Yao
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:125273-125296, 2026.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-wang26s, title = {Reward-Preserving Counterfactual State Editing for Offline Reinforcement Learning}, author = {Wang, Siyu and Chen, Xiaocong and Gong, Mingming and Li, Yong and Sheng, Quan Z. and Yao, Lina}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {125273--125296}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/wang26s/wang26s.pdf}, url = {https://proceedings.mlr.press/v306/wang26s.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Reward-Preserving Counterfactual State Editing for Offline Reinforcement Learning %A Siyu Wang %A Xiaocong Chen %A Mingming Gong %A Yong Li %A Quan Z. Sheng %A Lina Yao %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-wang26s %I PMLR %P 125273--125296 %U https://proceedings.mlr.press/v306/wang26s.html %V 306 %X 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.
APA
Wang, S., Chen, X., Gong, M., Li, Y., Sheng, Q.Z. & Yao, L.. (2026). Reward-Preserving Counterfactual State Editing for Offline Reinforcement Learning. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:125273-125296 Available from https://proceedings.mlr.press/v306/wang26s.html.

Related Material