Interpretable Reward Redistribution in Reinforcement Learning: A Causal Approach
Yudi Zhang, Yali Du, Biwei Huang, Ziyan Wang, Jun Wang, Meng Fang, Mykola Pechenizkiy
Abstract
A major challenge in reinforcement learning is to determine which state-action pairs are responsible for future rewards that are delayed. Reward redistribution serves as a solution to re-assign credits for each time step from observed sequences. While the majority of current approaches construct the reward redistribution in an uninterpretable manner, we propose to explicitly model the contributions of state and action from a causal perspective, resulting in an interpretable reward redistribution and preserving policy invariance. In this paper, we start by studying the role of causal generative models in reward redistribution by characterizing the generation of Markovian rewards and trajectory-wise long-term return and further propose a framework, called Generative Return Decomposition (GRD), for policy optimization in delayed reward scenarios. Specifically, GRD first identifies the unobservable Markovian rewards and causal relations in the generative process. Then, GRD makes use of the identified causal generative model to form a compact representation to train policy over the most favorable subspace of the state space of the agent. Theoretically, we show that the unobservable Markovian reward function is identifiable, as well as the underlying causal structure and causal models. Experimental results show that our method outperforms state-of-the-art methods and the provided visualization further demonstrates the interpretability of our method. The project page is located at https://reedzyd.github.io/GenerativeReturnDecomposition/ .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 78103d9b-9863-4ba5-88a4-a3867360b73cCited by top-tier papers8
- From Causal to Concept-Based Representation LearningGoutham Rajendran, Simon Buchholz, Bryon Aragam, Bernhard Schölkopf et al.NeurIPS 2024 · 37 citations
- Latent Reward: LLM-Empowered Credit Assignment in Episodic Reinforcement LearningYun Qu, Yuhang Jiang, Boyuan Wang, Yixiu Mao et al.AAAI 2025 · 29 citations
- STAS: Spatial-Temporal Return Decomposition for Solving Sparse Rewards Problems in Multi-agent Reinforcement LearningSirui Chen, Zhaowei Zhang, Yaodong Yang, Yali DuAAAI 2024 · 11 citations
- Causal Information Prioritization for Efficient Reinforcement LearningHongye Cao, Fan Feng, Tianpei Yang, Jing Huo et al.ICLR 2025 · 1 citation
- Self-evolving LLM agents with in-distribution OptimizationYudi Zhang, Meng Fang, Zhenfang Chen, Mykola PechenizkiyICML 2026 · 1 citation
Related papers
- Learning Long-Term Reward Redistribution via Randomized Return DecompositionZhizhou Ren, Ruihan Guo, Yuan Zhou, Jian PengICLR 2022 · 45 citations
- Generalizing Goal-Conditioned Reinforcement Learning with Variational Causal ReasoningWenhao Ding, Haohong Lin, Bo Li, Ding ZhaoNeurIPS 2022 · 59 citations
- Episodic Return Decomposition by Difference of Implicitly Assigned Sub-trajectory RewardHaoxin Lin, Hongqiu Wu, Jiaji Zhang, Yihao Sun et al.AAAI 2024 · 3 citations
- Temporally Disentangled Representation LearningWeiran Yao, Guangyi Chen, Kun ZhangNeurIPS 2022 · 84 citations
- Reinforcement Learning of Causal Variables Using Mediation AnalysisTue Herlau, Rasmus LarsenAAAI 2022 · 8 citations
