Interpretable Reward Redistribution in Reinforcement Learning: A Causal Approach
Authors: Yudi Zhang
Yali Du
Biwei Huang
Ziyan Wang
Jun Wang
Meng Fang
Mykola Pechenizkiy
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement learning
sparse reward
return decomposition
causal modeling

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/](https://reedzyd.github.io/GenerativeReturnDecomposition/).

Interpretable Reward Redistribution in Reinforcement Learning: A Causal Approach
Yudi Zhang'
Yali Du Biwei Huang' Ziyan Wang² Jun Wang*
Meng Fang Mykola Pechenizkiy'
Eindhoven University of Technology 'King's College London
3 University of California San Diego *University College London
University of Liverpool
{y zhang5, pec ml, Lyali.du, ziyan. wang) ac. uk bih0078uced. edu, jun. wangdcs ucl.ac uk, Meng.Fang@ll wnggglllerooo...ooom ac. uk
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 solution to re-assign credits for each time step from observed sequences. While the majority of current approaches construct the reward redistribution an uninterpretable manner, we propose to explicitly model the contributions of state and action from causal perspective, resulting in an interpretable reward redistribu- tion and preserving policy invariance. In this paper, start by studying the role of causal generative models in reward redistribution by characterizing the generation of Markovian rewards and trajectory-we long-terna return and further propose framework, called Generative Return Decomposition (GRD), for policy optimiza- tion 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 representa- tion to train policy over the most favorable subspace of the state space of the agent. Theoretically,
show that the unobservable Markovian reward function is identi- fiable, as well as the underlying causal structure and causal models. Experimental results show that our method outperforms ttate-of-the-ur oh-ther methods and the provided visualization further demonstrates the interpretability of our method. The project page is located at http://eeddy..gguuu.o.oeeeeeette potprraiieeeeeummmoooooosssooooi
Introduction
Reinforcement Learning (RL) has achieved significant success in a variety of applications, such as
6 71, financial trading [8] 9 and healtheare [I0]. 14] In such delayed reward setting, the learning process may suffer from instability due the lack of proper guidance
autonomous driving [ 2) robot [3] 4). games 6
A challenge in real-world RL is the delay of reward feedback [12 13
from the reward sequences 16 I7]. Methods such as reward shaping [8 19
curiosity-driven
intrinsic reward [20 21 22 23 and hindsight relabeling [24 25 26
have been proposed to
provide proxy rewards in RL.
How compute the contribution of each state-action pair affected by delayed reward and explain the reasons behind such contribution are equally important, as they can provide insights into the underlying dynamies and the decision process, guiding the development of more effective RL algorithms. Recent return decomposition methods employ the return equivalent hypothesis to produce proxy rewards for the state-action pair each time step 15 16 [7]. These methods allow for the
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

decomposition of the trajectory-wise return into Markovian rewards, preserving the policy invariance. However, previous methods construct reward redistribution by denddddieeedd rules [ [6]. or an uninterpretable model [7], where the explanation of how the state and action contribute to
the proxy reward is unclear. On the other hand, recently, causal modeling has shown promise in
environmental model estimation and follow-up policy learning for RL [27 (8). Works along this
direction investigate and characterize the significant causes with respect to their expected outcomes, which help improve the efficiency of exploration by making use of structural constraints [29] 30 31], resulting in an impressive performance improvement. Therefore, causal modeling serves as a natural tool to investigate the contributions of states and actions toward the Markovian rewards.
In this paper, we propose a novel algorithm for reward redistribution with causal treatment, called Generative Return Decomposition (GRD). GRD tackles reward redistribution with causal modeling that enjoys the following advantages. First, instead of a flat representation, GRD specifies each state and ac- tion as a combination of values of several constituent variables relevant to the problem and accounts for the causal relationships among variables in the system. Such structural factored representation, relating Factored MDP (3] provides favorability to form and identify the Markovian reward function from the perspective of causality. Unlike previous approaches, GRD utilizes such a parsimonious graphical represen- tation to discover how each dimension of state and action contributes to the Markovian reward. More- over, within the generative process of the MDP envi- ronment, we can naturally explain and model the observed delayed return as a causal effect of the unobserved Markovian reward sequence, which provides insights into the underlying dynamics of the environment, preserves the policy invariance as well. Figure I shows the framework of GRD, involving the causal relationship among environmental variables. The nodes denote different variables in the MDP environment, i.e., all dimensions of state 8,, and action a. Markovian rewards Tt for & E [1, T). as well as the long-term return R. For sparse reward settings in RL., the Markovian rewards r, are unobservable, which are represented by nodes with blue filling. While considering the return-equivalent assumption in return decomposition, we can observe the trajectory-wise long-term return, R, which equals the discounted sum of delayed reward o, and evaluates the performance of the agent within the whole episode. A special case of delayed rewards is in episodic RL, where 01.7-1 = and or * O. Theoretically, we prove that the underlying generative process, including the unknown causal structure, the Markovian reward function, and the dynamics function, GRD learns the Markovian reward function and dynamics function in a component wiss way to recover the underlying causal generative process, resulting in an explainable and principled reward redistribution. Furthermore, we identify a minimal sufficient representation for policy training from the learned parameters, consisting of the dimensions of the state that have an impact on the generated Markovian rewards, both directly and indirectly. Such a compact representation has a matching causal structure with the generation of the Markovian reward, aiding in the effectiveness and stability of policy learning.
$1
22
Map
$3,018
01413
dgg-1
dg.
T+-1
Figure I:
A graphical illustration of causal structure in Generative Return Decomposi- tion. See the main text for the interpretation.
identifiable.
We summarize the main contributions of this paper as follows. First, we reformulate the reward redistribution by introducing a graphical representation to describe the causal relationship among the dimensions of state, action, the Markovian reward, and the long-term return. The causal structure over the environmental variables, the unknown Markovian reward function, and the dynamics function are identifiable, which is supported by theoretical proof. Second, we propose GRD to learn the underlying causal generative process. Based on the learned model, we construct interpretable reward redistribution and identify compact representation facilitate policy learning. Furthermore, empirical experiments on robot tasks from the MuJoCo environment demonstrate that our method outperforms the state-of-the-ar methods and the provided visualization shows the interpretability of the redistributed rewards, indicating the usefulness of causal modeling for the sparse reward settings.
2

---

2 Related Work
Below we review the related work on reward redistribution and causality-facilica RL methods.
Previous work redistributes rewards to deal with the delayed rewards in RL, such reward shap-
ing 18 19] and curiosity-driven intrinsic reward [20] 21 22 23 Return decomposition draws
attention since RUDDER [15) rethinks the returnequivalent condition for reward shaping [4 to decompose the long-term return into proxy rewards for each time step through an LST'M-hased long-term return predictor and manually designed assignment. Align-RUDDER [T6] introduces the bioinformatical alignment method to align successful demonstration sequences, then manually scores new sequences according to the aligned demonstrations. However, many reinforcement learning (RL) tasks, obtaining a sufficient number of successful demonstrations can be challenging. [35] and 36 explore to improve RUDDER by the replacement of expressive language models [37 and the continuous modern Hopfield network [38] to decompose delayed rewards. Apart from them, RRD 17 proposes an upper bound for the common return-equee -eqiivalent assumption to serve as a surrogate optimization objective, bridging the return decomposition and uniform reward redistribution in IRCR (39 However, those methods lack the explanation of how the contributions derives [7 35 or applying unflexible manual design [15 16 to decompose the long-term returns into dense proxy rewards for the collected sequences. By contrast, we study the role of the causal model, investigate the relationships within the generation of the Markovian reward and exploit them to guide interpretable return decomposition and efficient policy learning.
Plenty of work explores solving diverse RL problems with causal treatment. Most conduct research on the transfer ability of RL agents. For instance, (7 learns factored representation and an individual change factor for different domains, and [3] extends it to cope with non non-stationary changes. More recently, 28 40 remove unnecessary dependencies between states and actions variables in the causal dynamics model to improve the generalizing capability in the unseen state. Also, causal modeling introduced to multi-agent task [4 42,, model-based RL, 43), imitation learning 4 and so on. However, most of them do not consider the cases of MDP with observed delayed and sparse rewards explicitly. As an exception, (45) distinguishes the impact of the actions from one-time step and future time step to the delayed feedback by policy gradient estimator, but still suffering from very delayed sparse rewards. Compared with the previous work, we investigate the causes for the generation of sequences of Markovian rewards to address very delayed rewards, even episodic delayed ones.
3 Preliminaries
In this section, we review the Markov Decision Process 446 and return decomposition [15 16
Markov Decision Process (MDP) represented by a tuple (S. R. 7, P). where S, A, R de- note the state space, action space, and reward function, separately. The state transition proba- bility of the environment is denoted as P(8++1 I 8t, at). The goal of reinforcement learning is to find an optimal policy : S - A to maximize the expected long-term return, i.e,, discounted sum of the rewards with the predefined discounted factor y and episode length T,
((x) = E 1.94 11.0-1)).0... (a 7 RRs,, a.)].
Return decomposition is proposed to decompose the long-term return feedback into a sequence of Markovian rewards in RL with delayed rewards [5 [6] while maintaining the invariance of optimal policy [15]. In the case of delayed reward setting, the agent can only observe some sparse delayed rewards. An extreme case is that the agent can only get an episodic non-zero reward or at the end of each episode, i.e., 01:T-1 0 and or * O, called episodic reward [17. In general, the observed rewards or are delayed and usually harm policy learning, since the contributions of the state-action pairs not clear. Therefore, return decomposition is proposed to generate proxy rewards ft for each time step, which is expected to be non-delayed, dense, Markovian, and able to clarify the contributions of the sate-action pairs to the delayed feedback. The work along this line shares a common assumption,
R ΣT T
0: 7
(I)
where is the long-term return and ot, cg denote the observed delayed reward and the decomposed proxy reward for each time step, separately.
3

---

Causal Reformulation of Reward Redistribution
As a foundation, below we introduce a Dynamic Bayesian Network (DBN) 47 to reformulate reward redistribution and characterize the underlying generative process, leading to natural interpretation of the explicit contribution of each dimension of state and action towards the Markovian rewards.
4.1 Causal Modeling
We describe the MDP process with a trajectory-wise long-term return, P= AS,, S, represent the sets of states and actions a, respectively. denotes a DBN that describes the causal relationship within the MDP environment, constructed over a finite number of random variables, {81.t.-... 8|al,t; 0,,,,, aJa),,, ! R where (s| and |a| are the dimension of and a. rı is the unobservable Markovian reward that serves as the objective of return decomposition, The state
transition Py can be represented as, Py(sr+1 Se, az) I P(si.++ pa(ss,c+1)), where is the index of dimension of 8t. Here pa( denotes the causal parents of 81,++1 and usually is subset of the dimensions of 8; and aę. We assume given initial state distribution P(s1). Similarly, we define the functions, g. which generates unobservable Markovian reward, rt (pa(r))). where pa(r.) is subset of the dimensions of s, and a, The trajectory wise long-tenn return R is the causal effect of all the Markovian rewards. An example of the causal structure denoted by G is given in Figure For simplicity, reformalize the generative process as,
8i = S 4 8t. C ② at, (s,it)
Ft = g(C @ St, Ca O az Gr,))
== ?
for = (1,2..... |s)). O is the element- product. f and g stand for the dynamics function and reward function, separately. C' are categorical masks organizing the causal structure in g. is and 6r,t i.i.d random noises. Specially, C {0, 1)|a and Ca-tr E {) 10 |a| control if a specific dimension of state and action a₂ impact the Markovian reward Tt, separately. For example, if there is edge from $;, Tı, then the C or 1. Similarly, C {0, 1] [e] x|a]
and Ca-+ E (0, 1] |a|x indicate the causal relationship between 8t, a₂ and 81+1, separately. In particular, assume that is time-invariant, i.e., f. g. C invariant. In the system, $t, at- and R are observable, while Tt are unobservable and there are no unobserved confounders and instantaneous causal effects.
4.2 Identifiability of Generative Process
Below we give the identifiability result of learning the latent causal structure and unknown functions in the above causal generative model.
Proposition 1 (Identifiability) Suppose the state 8t, action as trajectory-wise long-term return R are observable while Markovian rewards Tt are unobservable, and they form an MDP, as described in Eq. [2] Then, under the global Markov condition and faithfulness assumption, the reward function and the Markovian rewards T are identifiabl,, as well as the causal structure that is characterized by binary masks C and the transition dynamies f.
Details of the proof for Proposition I are deferred to Appendix B Proposition | provides the foundation for identify the causal structure - and the unknown functions f. in the generative process from the observed data. As result, with the identified structures and functions in [ we can naturally construct an interpretable return decomposition. Here we clarify the difference between our work and the previous work for a better understanding: 1) compared with [7 35), in Eq. 2 2 a general description to characterize the Markovian reward as the causal effect of a subset treating long-term reward R as the causal effect of all the Markovian rewards in a sequence is an interpretation of eeturn-equivalent assumption from the causal view, allowing reasonable and fiexible return decomposition.
of the dimensions of state and action, 2) compared with [13 16

---

Learning
Pollcy Optimization
Lepors Leea
Ldyo, Line
CR Input
Reward as Guidence
Tdys
cau
d
e-ir
d
C
C
C
② dp-1
dyya
Causal Structure
Cra-e
Ch
C
C
Pveer
++1
Ct
Orew
ca 41
Fs-3
Lywr
R
O4
R
CR: Compact Representation
Figure 2: The framework the proposed Generative Return Decomposition (GRD). Deau, Prew. Qdyn in generative model Pm are marked yellow, blue and green, while policy model o, is marked as orange. The observable variables, state 8t, action a₁, and the delayed reward ot, are marked as gray. The mediate results, binary masks, C outputs of policy, the predicted Markovian rewards A, and the compact representation shth are denoted as purple squares. The policy model 4 takes as input smin and its training is supervised by the predicted Markovian reward ft. The dotted lines represent the supervision signals, ie., losses and predicted Markovian rewards.
5 Generative Return Decomposition
In this section, propose a principled framework for reward redistribution from the causal perspec- tive, named Generative Return Decomposition (GRD). Specifically, we first introduce how GRD recovers generative model within the MDP environment in Sec. [5.1] and then show how GRD deploys the learned parameters of the generative model to facilitate efficient policy learning in Sec. 5.2] The GRD consists of two parts, 4n for the parameterized generative process and D for the policy. Therefore, the corresponding loss function includes Lm (fm) and J= (PP),, for causal generative model estimation and policy optimization, respectively. Hence, the overall objective can be formulated as
mine... .. L( 47) = L,n (..) ) + (
In the following subsections, we will present each component of the objective function.
5.1 Generative Model Estimation
In this subsection, we present how our proposed GRD recovers the underlying generative process. This includes the identification of binary masks (C*))) as well as the estimation of unknown functions (f and g). An overview of our method is illustrated in Figure 2 The parameterized model on that incorporates the structures and functions in Eq. 2 is used to approximate the causal generative process. The optimization is carried out by minimizing L'm over the replay buffer D, which consists of trajectories, with each trajectory т {(81, ar) I. where R is the discounted sum of observed
delayed rewards o,, given in Eq.
Generative Model. The parameterized model 4.m consists of three parts, deau Pdyns Orew. which are described below. (1) dcau is used to identify the causal structure in g by predicting the values of binary aa
R/a x|s|x2 & B |R|s|x and d R a which are used predict C Ca-s C and Ca-tr separately. We take the prediction of c with Φā-* as an example to explain the process of predicting causal structure. da-* characterizes |s] [s] i.i.d Bernoulli distributions for the existence of the edges organized the matrix, C Each Bernoulli distribution is denoted by a two-element vector, where each element corresponds to the unnormalized probability of classifying the edge as existing or not, respectively. We denote the probability of the existence of edge from i-th dimension of the state to j-th dimension of the next state in the causal graph as P(C Binary prediction of C is sampled by applying element-wist gumbel-softmax [88 during training, while using greedy selection during
masks in Eq. 2 It consists of four parts, & E R sx/x|s x2
5

---

inference. Similarly, we can obtain P(C P(C PIC;""), Ca C and Ca Prew is constructed with fully-connected layers and approximates the reward function g Eq. 2 which takes as input the state 8;, action a,, as well as the predictions of C and Ca-tr to obtain the prediction of Markovian rewards. (3) Odyn is constructed on a Mixture Density Network [49] and
(2)
is used to approximate the dynamics function f in Eq. 2 which takes as imputs the state 8,, action ae. the predictions causal structures, C and C As an example, the predicted distribution for the i-th dimension of the next state is P(s 8t; Ot, C C Odyn). More details for @m can be found in the Appendix C.2
Loss Terms, Accordingly, the loss term Lm contains three components with Lm (m)) Lrew Layn + Lreg. Considering the long-term return I as the causal effect of the Markovian rewards Tr, we optimize Prew, D and da by
T Σ 7
Trew (Drew: Peau Pemm
ETWDII R
Σ
E
o
1111
(4)
where Pe is predicted by drew, i.e., ft = (row ac, Ce+r, C.-t)). To optimize Odyn 6aa and da we minimize
18
& 0 -amt))= E, Bei Σ log P (s1,+11 st, ar, C (om))) (5)
Additionally, we minimize the following cross-entropy terms to regulate the sparsity of learned causal structure to avoid trial solutions. It is achieved by force to optimize the parameters towards the direction of the nonexistence of the causal edge. Let (xx = log P(æ:), where P(x;) is the possibility that the edge x; exists. The regularizer term is,
Inee ( (Drw))
A1 E, Di(C*TT) + A2 E, D1(Ca+T)+
As Eift (Co
quic-to-reaard
+ + (C Ens (C
(6)
(excluding self connections)
(self connections)
where 1(.) is the indicator function and prrpee-prraeetes A(.) are listed in Appendix C.5
5.2 Policy Optimisation with Generative Models
Here explain how the learned generative model aids policy optimization with delayed rewards.
Compact Representation as Inputs. Inspired by (30, we identify minimal and sufficient state set, gmin for policy learning as the input of the policy model, Φx called compact representation. contains the states' dimensions that directly or indirectly impact the reward. To be more specific, it includes each dimension of state 8;, 8t, which either 1) has edge to the reward i.e., C Ce-or = 1, or 2) has an edge to another state component in the next time-step, 83,411 E 8t, i.e., C = 1, such that the same component at time / the compact representation, i.e., 83,1 E spiln We first select C and C greedily from the i.i.d Bernoulli distributions characterized by de and d=-"" (as illustrated in Appendix C.2) and organize the state dimensions which exist in smin as C We define C := {i, VC =T3 to denote the dimensions of 8, which have impact on the Markovian reward ry directly. Then we have,
v... - (C V.. V.
where |C*| denotes the size of C and V denotes element wise logical or operation. Then, the compact representation which serves as the input of policy is defined as, sqmin C 8t. In this way, GRD considers the dimensions contributing to the current reward provider (the learned Markovian reward function) to choose action, thus leading to efficient policy learning.
Markovian Rewards as Guidance. With the predicted Markovian rewards, we adopt soft actor-critic (SAC) (50) for policy optimization. Specifically, given and a,, we replace the observed delayed reward or by Az = Vrew (8t; at, C Ca*+"), where C and Ca- are selected greedily from de-r and øa-r, illustrated in Appendix The policy model 4, contains two parts, a critic Φ

---

and an actor o due to the applied SAC algorithm. SAC alternates between policy evaluation step. which trains critic o, (sftin at) to estimate QT ax) Ex, i>i= -f aa) 81 8, 01 = a] using the Bellman backup operator, and a policy improvement step, which trains an actor on (spmi) by minimizing the expected KL-divergence,
-++) = EAAD [Dki. (P=]l (Q" V*))].
(8)
where V is the value function of state-action pairs (50] More detailed implementation of Φx can be founded in Appendix C.3
We optimize the generative model and the policy model simultaneously, cutting the need for collecting additional diverse data for causal discovery. It is worth noticing that SAC can be replaced by any
policy learning backbones [51 52 53 54]. Besides, for those model- based ones [53] [54], our learned generative model serves as a natural alternative to the used environment model.
6 Experiments
In this section, we begin by performing experiments on a collection of MuJoCo locomotion benchmark tasks to highlight the remarkable performance of our methods. We then conduct an ablation study to demonstrate the effectiveness of the compact representation. Finally, we investigate the interpretability of our GRD by visualizing learned causal structure and decomposed rewards.
6.1 Setup
Below we introduce the tasks, baselines, and evaluation metric in our experiments.
Tasks. We evaluate method eight widely used classical robot control tasks in the MuJoCo environment [55], including Half- Walf-Cheetah, Ant, Walker2d, Humanoid, Swimmer, Hopper, Humanoid- Standup, and Reacher tasks. All robots are limited to observing only one episode reward at the end of each episode, which is equivalent to the accumulative Markov reward that describes their performance throughout the episode. The maximum episode length is shared among all tasks and set to be 1000. In these tasks, the contributions across dimensions in the state to the Markovian reward can differ significantly due the agents' different objectives, even in Humanoid and HumanoidStandup, where the robots hold the same dynamic physics. A common feature is that the control costs are designed to penalize the robots if they take overly large actions. Therefore, the grounded Markovian reward functions consider different dimensions of state and all dimensions of action with varying weights.
Baselines. We compare our GRD with RRD (biased) [17]. RRD (unbiased) 17.. and IRCR [39], whose implementation details are provided in Appendix C.1 We use the same srper-parameters for policy learning RRD for a fair comparison. While RUDDER and RUiDDR are also popular methods for return decomposition, RRD and IRCR demonstrate superior performance compared to them. Moreover, Align-RUDDEE requires successful trajectories to form scoring rule for state-action pairs, which are unavailable in MuloCo. Therefore, we do not compare our performance with theirs.
Evaluation Metric. We report the average accumulative reward across 5 seeds with random initial- ization to demonstrate the performance of evaluated methods. Intuitively, method that carns higher accumulative rewards in evaluation indicates better reward redistribution.
6.2
Main Results
Figure B provides an overview of the performance comparison among the full version of our method, GRD, an ablation version without compact representation, GRD w/o CR, and the baseline methods. The ablation version GRD w/o CR uses all state dimensions instead of the compact representation as the input of policy model. For better visualization, we scale the time step axis to highlight the different converging speeds for each task. Our method GRD outperforms the baseline methods in all tasks, as shown by higher average accumulative rewards and faster convergence, especially in the tasks with high-dimensional states. For instance, HalfCheetah and Ani, GRD obtains the highest average accumulative reward of 1.5 x 10d and 6.5 x 103, separately, while the baseline methods achieve the best value of 1.25 x 10 and 103, respectively. In the task where the agents are equipped with high-dimension states, such as HumanoidStandup of 376-dimension states and Ant of
7

---

Ant-v2
HalfCheetah-v2
Walker2d-v2
Humanoid-v2
2001
pl0
Average Accumulative Reward
1.0
2.5
5.0
150
320
203
500
50
100
150
285
40
Number of frames (xtof)
Number of frames (x10%)
Number of frames oct(P)
Number of frames bctof)
Reacherv2
Swimmer-v2
Hopperv2
HumanoidStanduu-v2
xips
20
25
-0.4
2.0
-0.0
10-
-1.00
Number frames Ox16f) Numbar (xio")
Number
Number
RRD-Biased
RRD-Unbiased
IRCR
GRD w/o CR
Figure 3: Learning curves on suite of MuJoCo benchmark tasks with episodic rewards, based on 5 independent runs with random initialization. The shaded region indicates the standard deviation and the curves were smoothed by averaging the 10 most recent evaluation points using exponential moving average. An evaluation point was established every 10a time steps.
Table 1: The mean and the standard variance of average accumulative reward and sparsity rate S regarding diverse A at different time step in Swimmer.
A1 /1
le5 87 41(0777
2c5
3e5
0
153+ 43(0.80)
151 32(0.80)
55(088)
170 36(0.77)
159# 25(0.78)
3e-6
(82=203(0.72)
201=109(0.71) 130=103(0.56)
217=108(0.72)
217=100X0.68)
229 95(0.52)
220=102(0.50)
5e-6
245=104(0.69)
261#100(0.74)
286 77(0.77)
272 78(0.75)
262+ 94(0.66)
5e-5
118=100(0:46) 109=107(0.41) 123+103(0.34)
141# 96(0.25)
1524 93(0.19)
158 90(0.17)
11l-dimension states, GRD gains significant improvement. Taking HumanoidStandup as an example, GRD achieves higher average accumulative reward 2.3 10 the end of the training and always demonstrates better performance at a certain time step during the training. The possible explanation is that GRD can quickly learn the grounded causal structure and filter out the nearly useless dimensions for approximating the real reward function. The visualization of the learned causal structure and the related discussion be found in Sec. 6.4
6.3 Ablation Study
In this section, we first demonstrate that the compact representation improves policy learning and then investigate the potential impact of varying causal structures policy learning.
Compact representation for policy learning. According to Figure 3 although the GRD w/o CR has access to more information (it takes all the dimensions of the state as the input of policy model), GRD earns higher accumulative rewards on all the tasks. This is because only the dimensions tied to the predicted Markovian rewards reflect the agents' goals. That is, the agent is supervised by the learned reward function while choosing an action based on the associated state dimensions, ensuring the input of the policy model aligns with its supervision signals. This constrains the policy model to be optimized over a smaller state space, leading to a more succinct and efficient training process.
The impact of learned causal structure. Define sparsity of causal structure of the state to reward, S Elel C We control the value of ê to obtain different sparsities of learned causal struc- ture. The average accumulative reward and the sparsity of causal structure during the training process in Swimmer are presented in Table With the increase of X (like from 5e-6 6 to 5e-5), the causal structure generally gets more sparse (the sparsity S decreases), leading to less policy improvement. The most possible explanation is, GRD does not take enough states to predict Markovian rewards and thus can not reflect the true goal the agent, leading to misguided policy learning. By contrast, we set a relatively low X; to form a denser structure, GRD may consider redundant dimensions that harm policy learning. Therefore, a reasonable causal structure for the reward function can improve both the convergence speed and the performance of policy training.
8

---

1= le4
(1()))))
1e4
1e4
== 1e4
U (1,54)
(E [1.44]
(E (1.8)
(LE (1.A)))) (188)
5e5
0= 5e5
8 5e5
Sye4s e (1,54))
of (1.54)
if (1.54)
(f [1.8]
1e6
(1 (1.8))
1e6
c 1e6
t 1e6
(1.5A)
Fiz E (1.54)
Ag (6 (1.8)
(b) Causal Structure from 51.+1
(c) Causal Structure from S, to (d) Causal Structure
from Bgg tov;
Syz-1 E (1554)
(a) Causal Structure from Sja=1
0D
04
9.b
0.5
1.0
The visualization of learned causal structure for Ant when t E [1e4, 5c5, 1e6). The color indicates the probability of the existence of causal edges, whereas darker colors represent higher probabilities. There are 111 dimensions in the state variables, but only the first 27 ones are used. (a) The learned causal structure among the first 27 dimensions of the state variable 8, the first 54 dimensions of the next state variable 8141- Due to the limited space, we only visualize the structure t 1e4 and 1e6. (b) The learned causal structure among all dimensions of the action variable a, to the first 54 dimensions of the next state variable 8:+1- (c) The learned causal structure among the first 54 dimensions of the state variable 8; the Markovian reward variable ry- (d) The learned causal structure among all dimensions of action variable a, to the Markovian reward variable rt-
Figure 4:
Average Accumulative Reward
x104
SAC
DDPG
2101
TD3
1.50
1.50
1.50
1.00
1.00
1.
0.75
0.75
0.75
0.50
0.50
0.50
0.55
0.5
0.00
50
100
150
200
250
300
50
100
150
200
250
300
50
100
150
200
250 J00
Number of frames (x10")
Number of frames (x10")
Number of frames (x10*)
RRD-Bias
RRD-Unblas
IRCR
None
GRD
Figure 5: Evaluation with different RL backbones. SAC, DDPG and TD3. "None" is training with the observed delayed rewards.
Robustness. To evaluate the robustness of GRD, we provide results under noisy states in Ant. A significant characteristic of Ant that only the first 28 dimensions of state are used. Therefore, the noise on the other dimensions (28 ~ 111) should not impact learned robust policy. To verify this, during policy evaluation, we
Average Accumulative Reward
Gaussian
163
GRD
RRD-Unblas
IRCR
introduce the independent Gaussian noises with the mean of 0 and
standard deviation of 0 into those insignificant dimensions
(28 ~ 111). As shown in Figure 6 GRD is unaffected by the injected noises, while the performance of the baseline methods decreases. That is because the insignificant dimensions are not in the compact representation, which serves as the policy input of GRD. Therefore, method learns more robust policy with the usage of compact representation.
0.0
0.2
0.4
0.6
0.0
1.0
Standard Deviation of Noise
Figure 6: Evaluation with Gaus-
sian Noise in the State.
Results over different RL backbones. We provide the results of training with DDPG [66] and TD3 [57 F As the experimental result shows, on the tasks HalfCheetah, GRD consis- tently outperforms the baseline methods, RRD-Bias, RRD-Unbias, and IRCR, which are modified to run based on the same policy optimization algorithm, DDPG and TD3. We also provide results of "None", which utilizes the observed delayed reward for policy learning directly.

---

Consistent improvement. We provide results to showcase the consistent improvement of our method. 1) On different RL backbones; we provide results to demonstrate the consistent improvement of our method over different RI backbones (DDPG [56] and TD3 [57]) as shown in Figure | 2) On manipulation tasks: we provide the results on manipulation tasks MetaWorld Appendix D
[88]
Please refer to
6.4 Visualization
We visualize the learned causal structure and redistributed rewards in Ant, where the robot observes 11 Ill-dimension states and 84 dimensions are not used [55] The actions are 8-dimensional.
Causal Structure. As Figure 4 the existence of each element in the causal structure becomes more deterministic (indicated by
Episode
Episode
ground
bedistnbutted rowards
Reward
Reward
the more distinguishable color) and further filters out more ele- ments with the training. At the le4 time step, GRD regards 75% elements in Ca-tr (0. 1) 70% dimensions in Ca-+ E
groundd truin
redistr bused rewants
100
250
250
300
750
850
900
1000
Time Steps
Steps
17
Figure 7: The visualization of decomposed rewards (blue solid lines) and the grounded rewards (red dotted lines).
(0, 1) 3 8x376 98% dimensions in C E {0, 1) 376x376 be 0 with high probabilities, indicat- ing that the causal edge not exists. The learned structure is nearly consistent with the fact that there should not be an edge from the unused dimensions of the state variable to the other variables. Although some redundant edges have been learned here (as shown in Figure (a). (a), the grids for the causal structure from some dimensions of state to the 28 ~ 54 dimensions of the next state are of dark color), these edges do not affect our policy learning, since there is no edge from the 28 ~ 54 dimensions of s, to Ft. i..., these dimensions do not exist in the identified compact representation. Additionally, according to Figure d), the edges from different dimensions of to always exist. It corresponds to the reward design in Ant since the robots are expected to act with the lowest control cost, defined as Tt,cast Díe af. Therefore, the learned causal structure captures some causal characteristics of the environment, which can explain the generation of the Markovian rewards.
Decomposed Rewards. We visualize the decomposed rewards and the ground truth rewards to demonstrate the accuracy of predicting Markovian reward by GRD. According to Figure 7 the redistributed rewards (blue lines) consistently align with the ground truth (red lines), indicating that our method indeed distinguishes the state-action pairs with less contribution the long-term return from those with more contributions. More cases are provided in Appendix D.2
7 Conclusion
In this paper, propose Generative Return Decomposition (GRD) to address reward redistribution for the RL tasks with delayed rewards in an interpretable and principled way. GRD reformulates reward redistribution from causal view and recovers the generative process within the MDPs with rrajectory-wiie return, supported by theoretical evidence. Furthermore, GRD forms a compact representation for efficient policy learning. Experimental results show that, in GRD, not only does the explainable reward redistribution produce a qualified supervision signal for policy learning, but also the identified compact representation promotes policy optimization.
Limitation. The limitations of our work arise from the made assumption. We presumed a stationary reward function, limiting method's use in the case of dynamic, non-stationary reward function. The assumed static causal graph and used causal discovery methods might not cater to situations with confounders. Besides, an important direction for future work is to extend our proposed method in the case of partial observable MDP and to discuss the identifiability in such scenarios.
Broader Impacts. Our motivation is to strive to make decisions that are both understood and trusted by humans. By enhancing the transparency and credibility algorithms, we aim to harmonize human- Al collaboration to enable more reliable and responsible decision-making making in various fields. Our GRD algorithm notably advances this by offering a causal view of reward generation and identifying a compact representation for policy learning.
10

---

Acknowledgments
We thank the reviewers for the constructive comments and questions, which improved the quality of our paper. Part of this work used the Dutch national e-infrastructure with the support of the SURF Cooperative using grant no. EINF-4550. Yali Du thanks the support by the EPSRC grant EP/Y003187/1.
References
[1] B Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey, IEEE Transactions on Intelligent Transportation Systems, 2021.
[2] Jianyu Chen, Bodi Yuan, and Masayoshi Tomizuka. Model-free deep reinforcement learning for urban autonomous driving. In 20/9 TEEE Intelligent Transportation Systems Conference (ITSC), pages 2765- 2771. IEEE, 2019.
[3] Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep rein- forcement learning from human preferences. Advances in neural information processing systems, 30, 2017.
[4] Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured representations for model-based reinforcement learning. In International conference on machine learning, pages 7444 7453. PMLR, 2019.
[5] David Silver, Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap. Karen Simonyan, and Demis Hassabis. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419): 140 1144, 2018.
[6] Lei Han, Peng Sun, Yali Du, Jiechao Xiong, Qing Wang, Xinghai Sun, Han Liu, and Tong Zhang. Grid-wise control multi-agent reinforcement learning in video game ai. In International Conference on Machine Learning (JCML), pages 2576-2585 PMLR, 2019.
[7] Yali Du, Lei Han, Meng Fang, Tianhong Dai, Ji Liu, and Dacheng Tao. Liir: learning individual intrinsic reward in multi-agent reinforcement learning. Proceedings of the 33rd International Conference on Neural Information Processing Systems, pages 4403-4414, 2019.
[8] Zihao Zhang, Stefan Zohren, and Stephen Roberts. Deep reinforcement learning for trading. The Journal of Financial Data Science, 2(2):25 40, 2020.
[9] Hongyang Yang, Xiao- Yang Liu, Shan Zhong, and Anwar Walid. Deep reinforcement learning for automated stock trading: An ensemble strategy. In Proceedings of First ACM International Conference on Al in Finance, pages 1 8, 2020.
[10] Chao Yu, Jiming Liu, Shamim Nemati, and Guosheng Yin. Reinforcement learning in healthcare: A survey. ACM Computing Surveys (CSUR), 55(1):1-36, 2021.
[11] Nan Rosemary Ke, Anirudh Goyal ALIAS PARTH GOYAL, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio. Sparse attentive backtracking: Temporal credit assignment through reminding. Advances neural information processing systems, 31, 2018.
[12] Meng Fang, Cheng Zhou, Bei Shi, Boqing Gong, Jia Xu, and Tong Zhang. Dher: Hindsight experience replay for dynamic goals. International Conference on Learning Representations, 2018.
[13] Meng Fang, Tianyi Zhou, Yali Du, Lei Han, and Zhengyou Zhang. Curriculum-guided hindsight experience replay. Advances in neural information processing systems, 32, 2019.
[14] Beining Han, Zhizhou Ren, Zuofan Wu, Yuan Zhou, and Jian Peng. Off-polic reinforcement learning with delayed rewards. In International Conference on Machine Learning, pages 8280-8303. PMLR, 2022.
[15] Jose A Arjona-Meding Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brandstetter, and Sepp Hochreiter. Rudder: Return decomposition for delayed rewards. Advances in Neural Information Processing Systems, 32, 2019.
[16] Vihang P. Patil, Markus Hofmarcher, Marius Conussatnnnnnnnn Dinu, Matthias Dorfer, Patrick M. Blies, Jo- hannes Brandstetter. Jose A. Arjona-Medina Medina, and Sepp Hochreiter. Align-rudder: Learning from few demonstrations reward redistribution. In International Conference Machine Learning. PMLR, 2022.
11

---

[17] Zhizhou Ren, Ruihan Guo, Yuan Zhou, and Jian Peng. Learning long-term reward redistribution via randomized return decomposition. In International Conference on Learning Representations, 2022.
[18] Yujing Hu, Weixun Wang. Hangtian Jia, Yixiang Wang. Yingfeng Chen, Jianye Hao, Feng Wu, and Changjie Fan. Learning utilize shaping rewards: A new approach of reward shaping. Advances in Neural Information Processing Systems, 33:1,5911-1599 2020.
[19] Pradyumma Tambwekar, Murtaza Dhuliawala, Lara Martin, Animesh Mehta, Brent Harrison, and Mark Riedl. Controllable neural story plot generation via reward shaping. In Proceedings of the Sixteenth International Joint Conference on Artificial Intelligence, pages
2019.
[20] Sai Rajeswar, Cyril Ibrahim, Nitin Surya, Florian Golemo, David Vazquez, Aaron Courville, and Pedro O Pinheiro. Haptics-based curiosity for sparsc-reward tasks. In Conference Robot Learning, pages 395-405. PMLR, 2022.
[21] Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Danell. Curiosity-driver exploration by sulf-supervised prediction. In International Conference on Machine Learning, pages 2778 -7777 PMLR, 2017.
[22] Lulu Zheng, Jiarui Chen, Jianhao Wang, Jiamin He, Yujing Hu, Yingfeng Chen, Changjie Fan. Yang Gao, and Chongjie Zhang. Episodic multi-agent reinforcement learning with curiosity-driven exploration. Advancex in Neural Information Processing Systems, 34:3757-3769, 2021.
[23] Cédric Colas, Tristan Karch, Nicolas Lair, Jean-Michel Dussoux, Clément Moulin-Frier, Peter Dominey, and Pierre-Y Oudeyer. Language as cognitive tool to imagine goals in curiosity driven exploration. Advances in Neural Information Processing Systems, 33:3761-3774 2020.
[24] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. Advances in Neural Information Processing Systems, 30, 2017.
[25] Ben Eysenbach, Xinyang Geng, Sergey Lovine, and Russ R Salakhutdinov. Rewriting history with inverse rl: Hindsight inference for policy improvement. Advances in neural information processing systems. 33:14783- 14795, 2020.
[26] Meng Fang. Cheng Zhou, Bei Shi, Boqing Gong, Jia Xu, and Tong Zhang. Dher: Hindsight experience replay for dynamic goals. In International Conference on Learning Representations, 2019.
[27] Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane. and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. In International Conference on Learning Representations, 2021.
[28] Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based counterfactual data augmentation. Advances in Neural Information Processing Systems, 2022.
[29] Carlos Guestrin, Daphne Koller, and Ronald Parr: Multiagent planning with factored mdps. Advances in Neural Information Processing Systems, 14, 2001.
[30] Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hernández-Lobato, Clark Glymour, Bernhard Schölkopf, and Kun Zhang. Action-sufficient state representation learning for control with structural constraints. In International Conference on Machine Learning, pages 9260-9279. PMIR, 2022.
[31] Fan Feng, Biwei Huang, Kun Zhang, and Sara Magliacane. Factored adaptation for non-stationary reinforcement learning. Advances in Neural Information Processing Systems, 2022.
[32] Assaf Hallak, François Schnitzler, Timothy Mann, and Shie Mannor, Off-policy model-based learning under unknown factored dynamics. In International Conference on Machine Learning, pages 711-719. PMIR, 2015.
[33] Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In Thomas Dean, editor, Proceedings of the Sixteenth International Joint Conference on Artificial Intelligence, pages 740 747. Morgan Kaufmann, 1999.
[34] Andrew Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In Proceedings of the Sixteenth International Conference on Machine Learning, pages 278-287, 1999.
[35] Yang Liu, Yunan Luo, Yuanyi Zhong, Xi Chen, Qiang Liu, and Jian Peng. Sequence modeling of temporal credit assignment for episodic reinforcement learning. arXiv preprint 1rXiv:1905.1340 2019.
12

---

[36] Michael Widrich, Markus Hofmarcher, Vihang Prakash Patil, Angela Bitto-Nemling. and Sepp Hochreiter. Modern hopfield networks for return decomposition for delayed rewards. In Deep RL Workshop NeurIPS 2021, 2021.
[37] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, Illia Polosukhin. Attention is all you need. Advances Neural Information Processing Systems. 30, 2017.
[38] Hubert Ramsauer, Bernhard Schäff, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hoptield networks is all you need. arXiv preprint arXivv2008.022) 2020.
[39] Tanmay Gangwani, Yuan Zhou, and Jian Peng. Learning guidance rewards with trajectory- smoothing. Advances in Neural Information Processing Systems, 33, 2020.
[40] Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke 7hu, and Peter Stone. Causal dynamics learning for task- independent state abstraction. International Conference on Machine Learning, pages 23151-23180, 2022.
[41] St John Grimbly, Jonathan Shock, and Arnu Pretorius. Causal multi-agent reinforcement learning: Review and open problems. In Cooperative Al Workshop, Advances in Neural Information Processing Systems, 2021.
[42] Natasha Jaques, Angeliki Lazaridou, Edward Hughes, Caglar Gulcehre. Pedro Ortega, DJ Strouse, Joel Z Leibo, and Nando De Freitas. Social influence intrinsic motivation for multi-agent deep reinforcement learning. International Conference on Machine Learning, pages 1040-3049 PMLR, 2019.
[43] Junzhe Zhang and Elias Bareinboim. Markov decision processes with unobserved confounders: A causal approach. Technical report. Technical report, Technical Report R-23, Purdue AI Lab, 2016.
[44] Junzhe Zhang, Daniel Kumor, and Elias Bareinboim. Causal imitation learning with unobserved con-
founders. Advances in Neural Information Processing Systems, 33:12263 -2274,, 2020.
[45] Thomas Mesnard, Theophane Weber, Fabio Viola, Shantanu Thakoor. Alaa Saade, Anna Harutyunyan. Will Dabney, Thomas Stepleton, Nicolas Heess, Arthur Guez, et al. Counterfactual credit assignment model-free reinforcement learning. In International Conference on Machine Learning, pages 7654-7664. PMLR, 2021.
[46] Richard Sutton and Andrew Barto. Reinforcement learning: An introduction. MIT press, 2018
[47] Kevin Patrick Murphy. Dynamic bayesian networks: representation, inference and learning. University of California, Berkeley, 2002.
[48] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizztion with gumbel-softmax. In International Conference on Learning Representations, 2016.
[49] Douglas A Reynolds al. Gaussian mixture models. Encyclopedia of biometrics, (41(659-)63), 2009.
[50] Tuomas Haamoja, Aurick Zhou, Pieter Abbecl, and Sergey Levine. Soft actor-critie; Off Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, pages 1861-1870. PMLR, 2018
[51] Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. CoRR, abs/1602.01 2016.
[52] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, abs/ 00//100770634777 00346,637777 2017.
[53] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in neurul information processing systems, 32, 2019.
[54] Ignasi Clavera, Jonas Rothfuss, John Schulman, Yasuhiro Fujita, Tamim Asfour, and Pieter Abbeel. Model- based reinforcement learning via meta-policy optimization. In Conference on Robot Learning, pages 617-629. PMLR, 2018.
[55] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model- control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 2012.
IEEE,
13

---

[56] David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Determin- istic policy gradient algorithms. In International conference machine learning, pages 387-395. Pmnr, 2014.
[57] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. International conference on machine learning, pages 1587-1596. PMLR, 2018.
[58] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-worlds benchmark and evaluation for multi-tast and meta reinforcement learning. In Conference on robot learning, pages 1094-1100 PMLR, 2020.
[59] Judea Pearl. Causality: Models, reasoning, inference, 2000.
[60] Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, prediction, and search. MIT press, 2000.
[61] Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning, volume 2. MIT press Cambridge, MA, 2006.
14

---

A Background on Causal Inference
directed acyelic graph (DAG), (V, EE), can be deployed represent a graphical criterion carrying out set of conditions on the paths, where V and E' denote the set of nodes and the set of directed edges, separately.
Definition 1 (d-separation (59) A set of nodes Z S V blocks the path p if and only if (1) contains chain such that the middle node in Z. or (2) p contains collider
i + j or fork +
such that the middle node sets of nodes. If and only if the set blocks all paths from one node to one node in d-separate X from Y, denoting as (X l Y 2).
in Z and such that descendant of m is in Z. Let X. Y and Z be disjunct
considered to
Definition 2 (Global Markov Condition [59 60D If. for any partition (X, Y. Z). X is d-separated from Y given Z. ie.. X La Y Z. Then the distribution P over V satisfies the global Markov condition graph G. and can factorizes as, P(X, given Z, writing as X 11 Y Z.
P(X ZY Z). That is, conditionally independent of Y
Definition 3 (Faithfulness Assumption 59 60) The variables, which are not entailed by the Markov Condi- fion, are not independent of each other.
Under the above assumptions, can apply d-separation criterion to understand the conditional indepen- dencies from given DAG G. That is, for any disjoint subset of nodes X,Y Y. (X 11 Y Z) and X La are necessary and sufficient condition of each other.
B Details of Theoretical Analysis
Proposition 1 (Identifiability) Suppose the state 8s, action a., trajectory-wise long-term return R are observ- able while Markovian rewards T. are unobservable, and they form an MDP, as described in Eq 2 Then under the global Markov condition and foithfulness assumption, the reward function 9 and the Markovian rewards rt identifiale,, well the causal structure that is characterized by binary masks - and C - and the transition dynamics f-
Below is the proof of Proposition ! We begin by clarifying the assumptions we made and then provide the mathernaticl proof.
Assumption We assume that, ta,t, and 6,, in Eq. 2 are i.i.d additive noise. From the weight-space view of Gaussian Process 61l, cquivalently, the causal models for 81.41 and r be represented as follows, respectively,
8i,t+1 = fi((e, ar) + te,iee Wh ax) fasi.t
(9)
Ft g(s as) + cr.t w (sc,, as) + Cr.ts
(10)
where E [1, [s]], and Φ and Φ, denote basis function sets.
Then we denote the variable set in the system by V, with V= fs1....... a 01,1, Gjal.t, rx re U R and the variables form Bayesian network G. Note, we assume that there are possible edges only from St.t- 8 to S E from a, E a:-1 sy s.. from 8,, rt, and from aj,a a, to r in G. In particular, the are unobserved, while R = ΣT, " " is observed. Thus, there are deterministic edges from each r, R.
Below omit the 7 for simplicity.
Given trajectory-wist wise long-term return R, the binary masks, Ca-77 C and Markovian reward function g and the rewards T. are identifiable, Following the above assumption, we first rewrite the function to calculate trajectory-wise long-term return in Eq. 2 as.
Proof
rt
WTo (8t, ar) +en
(1I)
T
T (at, ax) + Σ
W
Σ
Exa
[=]
1=1
15

---

For simplicity, we replace the compoments in Eq. 17 by.
(x(X) = (5 ac);
(12)
E
Er.ts
where X -st,,a) am) representing the concatenation of the covariates 8, and as from 1 T. Conse- quently, we derive the following equation.
R = + Er,
(13)
Then we can obtain a closed-form solution of W in Eq. X, and response variables R., where both continuous. One classical approach to finding such solution involves minimizing the quadratic and incorporating weight-decay regularizer prevent overfitting. Specifically, we define cost fanction as.
13 by modeling the dependencies between the covariates
C(W,)
Σ (R, - -W76666)) + NIWW1. 2 RR
1
(14)
where T represents trajectories consisting state-action pairs X and long-term returns R.. which are sampled from the replay buffer D. is the weight-decay regularization parameter. To find the closed-form solution, we differentiate the cost function with respect to W, and set the derivative to zero:
ac(W,) 9) 0.
(15)
aw,
Solving this equation will yield the closed-form solution for W.,. i.e.
Wg = (Aaa Co Co AI.)
(16)
Therefore, W,, which indicates the causal structure and strength of the edge, can be identified the observed data. summary, given trajectory-wise long-term return the binary masks, Ca-ar C and Markovian reward function g and rewards Tı are identifiable.
The binary masks, Ca -a and the transition dynamics are identifiable a similar manner, based on the assumption and Eq. we can rewrite Eq. 9 t0.
Sc+1 W (St, at) + to.itt.
(17)
To obtain cloxed-form solution for Wi, 1 in Equation 17 covariates X₁ and the response variables 8141, both of which are continuous. closed-form solution can be represented as:
model the dependencies between the
C(W..) =
Σ RLLIN
(8L141- WT az)) A1IWWLI1
(18)
By taking derivatives of the cost function and setting them zero, obtain the closed solution.
We, (Ala + D , ) -1 -00,,88,41
(19)
= (0. + AI.) 88,441-
Therefore, W., can be identified from the observed data. This conclusion applies all dimensions of the state. As a result, the f. which indicates the parent nodes of the i-dimension of the state, as well as the strength of the causal edge, are identifiable. In summary, the binary masks, =* Ca-+ and the transition dynamics f are identifiable.
Considering the Markov condition and faithfulness assumption, can conclude that for any pair of variables V.. V3 V. V. and V, not adjacent the causal graph if and only if they conditionally independent given some subset of (Vi + i, + i}. Additionally, since there are no instantaneous causal relationships and the direction of causality can be determined if an edge exists, the binary structural masks Ca O C and c defined over the set V are identifiable with conditional independence relationships [30]. Consequently, the functions f and in Equation 2 are also identifiabl.
16

---

Algorithm Learning the generative process and policy jointly.
I: Initialize: Environment &, trajectory + butfer D <- 0
Initialize: Generative Model 2m = [0ams Oayo: Prew] Policy Model 2..
3: for i = ........ ...,,
10 do
4: -- 0, resct &
5: for Metep 1.2... 100 do
6: Sample data (st, ai, 0n) from &, and store them trajectory
8:
if E done then store trajectory T {81.2, 01:r, RG to buffer D, where R = E/= 1 700000 т + 0, reset 3
9:
10:
end if
12:
13:
for Thbathh = 1,2,.... train trrin_batches batches do
14:
15:
// Optimize Generative Model 4m Sample consisting of trajectories from D: D2 ((s ai [[- Hm)
17:
Sample binary masks by Gumbel-Softmax from demi C and Ca- Optimize &rew with Di: Crew +- Orew dn
18:
dnow Leew (Eg.
19:
20:
Sample D2 consisting of N samples from D: D2 {sta atn Stat1] | l Sample binary masks by Gumbel-Softme from Oeam: C Ca- C and Ca Optimize Payn with D2 (Using C and C Odyn & Odyn Oayn Ldyn (Eq.
21:
22:
23:
24:
Optimize Deai Pcan <- Pcan a den (Tup + Irew Taya) (Eq. 6
25:
26:
// Optimize Policy Model d
27:
Sample binary masks greedily from demi Ca Ca-r and C Calculate Ca based on C and Cams
28:
29:
Update D2: D2 + {Ca-* O stn: Atwy +1 Prew (St,, Or, Car Ca-tr)) i Optimize oa 2.r 4. a av Jx (Eq,
30:
31:
32:
end for
33:
end for
34:
end for
C Details of Implementation
C.1 Baselines
We compare our method against the following baselines,
RRD (biased). This baseline utilizes surrogate objective called randomized return decomposition loss for reducing the consumption of estimating the Markovian reward function. applies Monte-Cate Carlo sampling to get biased estimation of the Mcan Square Error (MSE) between the observed episodic reward and the sum of Markovian reward predictions in sequence. We keep the same setting and hyper-parameters with its official implementation reproduce the results, in which the policy module is optimized by soft actor-critic (SAC) algorithm [50].
RRD (unbiased). This variant of RRD (biased) provides unbiased estimation of MSE by sampling short sequeeeeee.. offers a computationally efficient approach optimize MSE. According to [77. RRD (biased) and RRD (unbiased) achieve state-of-the- tf-th---t performance in episodic MuJoCo tasks.
This baseline performs non- parameeettiic uniform reward redistrihution. At each time step, the proxy reward set to the normalized value of the trajectory return. IRCR is a simple and efficient approach, and except for RRD. it achieves fate-of-the-ar performance in the literature. The implementation is from RRD I7
C.2 Detailed Generative Model
The parametric generative model 2.. used in the MDP environment consists of three components: Ouan Drow- and dua. We provide detailed description of their model structures below.
17

---

Layer#
1
Prew
FC256
FC256
FCI
Odyn
FC256
FC256
FC9
FC256
FC256
FC2 a
FC256
FC256
FCI
Table 2: The network structures of Frew. Odyns Øn and Φ.. FC256 denotes a fully-connected layer with output size of 256. Each hidden layer is followed by activation function, ReLU. |a] is the number dimensions of the action in a specific task.
Dean for predicting the causal structure. des comprises a set of free parameters without input. divide deau into four parts, each corresponding to the binary masks in Equation 2 Specifically, we have
Rlaxxeexx2 C E {0, 13
p
cma
a-
E Rialxidxx2 for Ca-s E Rla/xlel
O E Rlaix2
for C E R
6a E Rlalx2 for Ca-tr Rlal
Below explain the shared workflows in den using the example of predicting the causal edge from the i-th dimension state ai, the j-th dimension of the next state 8,,441, by part of the free parameters, Pem.i.j.
For simplicity, we denote dom.... as tb. The shape of is now easy to be determined. That is 1 E R and we write it [e'u, 1/1.. With this 2-element vector, element corresponds the unnormalized probability of classifying the edge existing (to) or not existing (11)) respectively. Therefore, the probability of the causal edge existing from the i-th dimension of state s,, to the j-th dimension of the next state 8j,t+ can he calculated as:
characterize Berouuli distribution. where each
exp(1io)
P(CT)')) exp(10) + exp(V1))
(20)
One example of using the estimated mask given Figure
During training: obtain C²,² through Gumbel- Softmax sampling. In the training phase, it is crucial to main- tain the gradient flow for backpropagation. To achieve this, sample the binary values of Oty'* by applying Gumbel-Softnn
(21) where GS denotes the Gumbel-Softmax sampling, which allows obtain binary discrete samples from the Bernoulli distribution. Applying Gumbel Softmax sampling allows us to randomly sample from the Bernoulli distribution in a stochastic manner rather than simply selecting the class with the highest probability. This introduces some randomness, enabling the model to explore the balance and uncertainty between different classifications more flexibly.
CT,² = GS(x)
The above explanation of the workflow in dem for predicting a single causal edge provides insight into the overall implementation of the entire module dew can applicable predicting causal edges during the training phase. Therefore, we can obtain Ca-+ for optimizing duyn: C and C for optimizing Prew. using similar procedures.
During inference: obtain C by greedy selection. In the inference phase, including data sampling and policy learning, we get the prediction of C through greedy selection,
s/ 4/2
C
(22)
0, t/u 4/2-
Such greedy sampling avoids introducing randomness using the Gumble-Softmax sampling.
The above explanation of the workflow in dcm for predicting single causal edge provides insight into the overall implementation of the entire module Pum and can he applicable for predicting all the causal edges during the inference phase. Therefore, we can obtain =* Caa-p and Ca using similar procedures to predict the Markovian reward and extract compact representation min
érew for predicting the Markovian rewards. &rew is a stacked, fully colle-comnected network, and the details of the network structure provided in Table 2 During training, the prediction of Markovian reward can be
18

---

St
C5-2
$1,+1
St
C5
$2,t+1
St
C555
3,6+1
Learnable Parameters
Ca
2
a1
ddyn
at
Element-wise Product
Figure 8: The illustration of using learnable masks to predict the next state.
written as,
== (Ac, Ci. C (,-,,, ---" FCs([C*** c 81, Ca a.l):
(23)
where [.. O denotes concatenation and element-wise multiply operations, respectively. denotes stacked mully-connected network. C and C are derived from dem by
During inference, including policy learning and data sampling, the predicted Markovian reward is
f == @rew (#t, a)2, C** "*,,,,"""" CP*") = FCs([C*** O 8s; Ca @ at)).
(24)
where
C and C are derived from Ocn greedily by deterministic sampling.
ddn for modeling the environment dynamics. In our experiment, we not directly utilize duyn in policy learning. Instead, this module serves as bridge to optimize $5 and da-s, Suhsequently, de can be utilized in the calculation of C
During training, initially sample C and Ca-+ using Gumbel-Softmax. The prediction for the i-th dimension of the next state be represented follows,
$. MDN( о S1, C C aa)),
(25)
where [+ denotes concatenation and MDN denotes Mixture Density Network, which outputs the means, variances, and probabilities for No.. Gaussian cores. The parameters of MDN are shared across the predictions
of different dimensions of the next state. We set Neau 3 in our experiments. More details about daya can be found in Table 2
C.3 Detailed Policy Model
Considering the specific requirements the employed RL algorithm, Soft Actor- (SAC), our Policy Model Φ comprises two components, the actor dix and the critic &o. Detailed network structures for both components can be found in Table 2
C.4 Training Process
We follow line of joint learning data sampled by optimal or sub-optimal policies [16]. During each mini-batch training iteration, we sample two sets of data separately from the replay buffer D:
[17,, which avoids learning return decomposition model advance using
D = ((""", at P-1, RnD [ consists of trajectories. Provided with the rajectory-wiss long-term returns RT mm=, D2 utilized optimize pe-rr oa-2r and drew, with T.ow-
D2 = (st,, as,, St,+1 +1} - consists of N state-action pairs. D2 are used for policy optimization and optimize the parts of causal structure, &e and da-ts duyn. With such Dz, GRD breaks the temporal cues in the training data learn the policy and dynamics function.
Please refer to Algorithm for detailed training process.
19

---

Table 3: The table of the hyper-parameterss used in the experiments for GRD.
Envs
A2
A
Ant
10-5
0
10-
10-8
10
8
HalfCheetah
10 10-5
10-5
10
10
5
10-6
10
Walker2d
10-b
10
10
10-0
10-7
Humanoid
10-5
10-3
10
10-7
10
Reacher
5 x 10-1
10-8
10-8
10-8
10
Swimmer
10-T
10-9
10-9
10 10-9
Hopper
10 6
10 6
10-
10 6
10-
10 6
HumanStandup
10-5
10
10-6
10
10
Table 4: The hyper-parameters.
hyperparameters
value
hyperparameters
value
epochs
optimizer
Adam
cycles
100
learning rate
x 10-4
iteration
100
N
256
train_ batches
100
M
replay buffer size
10
1.00
evaluation episodes
Polyak-averaging coefficient
0.0005
C.5 Hyper-Parameters
The network is trained from scratch using the Adam optimizer, without any pre-training. The initial learning rate for both model estimation and policy learning is set to x 10 The hyperparameters for policy learning are shared across all tasks, with discount factor of 1.00 and Polyak-averaging coefficient of 5 x 10 The target entropy is
the negative value of the dimension of the robot action. To facilitate training, replay buffer with size of x 10th time steps. The warmup size the buffer training set to 10 The model trained for 3 epochs, with each epoch consisting of 100 training cycles. In cach cycle, we repeat the process of data collection and model training for 100 iterations. During each iteration, we collect data from 100 time steps of interaction with the MuloCo simulation, which is then stored in the replay buffer. For training the grew. sample episodes, each containing 5 10 steps. As for policy learning and the optimization of duyn. we use data from 256 time steps. dem trained together with svew and dayn- Validation is performed after every cycle, and the average metric is computed based on 10 test rollouts. The hyperparameters for learning the GRD model can be found in Table [] All experiments were conducted on an HPC system equipped with 128 Intel Xeon processors operating at clock speed of 2.2 GHz and terabytes of memory.
utilize
D Additional Results
D.1 Results over manipulation tasks
We provide the comparison our method with RRD the three tasks of MetaWorld, as shown in Figure 9 The evaluation metric is the success ratc. As shown in the results, compared with RRD, GRD achieves comparable or better performance.
Door Lock: The agent must lock the door by rotating the lock clockwise.
Push Wall: The agent is required to bypass wall and push puck goal. The puck and goal positions are randomized.
Pick Place: The agent needs to pick and place a puck to a goal. The puck and goal positions are randomized.
D.2 Visualization of Decomposed Rewards
As shown in Figure 10 we visualize the redistributed rewards in Ant by GRD, as well as the grounded rewards provided by the environment.
20

---

Door Lock
Push Wall
Pick Place
1.0
1.0
1.0
Success Rate
0.8
0.8
0.8
0.6
0.5
0.6
0.4
0.4
0.4
0.2
0.0
50
160
150
200
0.0g
0.0
Number of frames (x10+)
Number of frames (x10*)
Number of frames (x10*)
RRD-Bias
RRD-Unbias
GRD
Figure 9: Evaluation on the manipulation tasks of MetaWorld, Door Lock, Push Wall, Pick Place.
D.3 Visualization of Learned Causal Structure
We provide the visualization of learned causal structure in Swimmer, shown in 1T Since the grounded causal structure is not accessible, verify the reasonability of the learned causal structure by observations:
All the edges from different dimensions shares the same characteristic with Ani that the edges from different dimensions of a to r always exist, corresponding with the reward design of penalizing the agent if it takes actions that are too large, measured by the sum of the values of all the action dimensions.
to always exist, as shown Figure (d): Swinmer
According to Figure (b), the first dimension of action (Torque applied on the first rotor) has an impact on the last three dimensions of state (angular velocity of front tip, angular velocity of first rotor, second rotor), which is corresponding with the intuitive that the first dimension of action should impact the part that connects first rotor. We get similar observation for the second action dimension from Figure 11
We observe that all the state dimensions are learned be connected to the reward; the possible explanation that the swimmer robot, any changes of the front tip, rotors will impact the position the robot, potentially influencing the reward.
21

---

Episode
Episode 2
Reward
Reward
ground truth redistributed rewerds
ground truth redistributed rewnerts
Steps
Steps
Episode 3
Episode 4
Reward
Reward
ground truth redist riuute rewards
ground bruth redistributed rewards
50
250
300
35o
50
250
300
350
Time Steps
Time Steps
Episode 5
Episode 6
Reward
Reward
ground redistributed rewerds
ground redistributed rewerdss
600
650
700
eon
oin
10o0
600
650
700
850
eon
osn
1000
Time Steps
Time Steps
Episode 7
Episode 8
Reward
Reward
truth
truth redistributed rewards
600
650
700
750
800
850
900
850
1000
600
650
700
750
800
85o
900
550
2000
Time Steps
Time Steps
Figure 10: The visualization of redistributed rewards and grounded rewards in Ant. The results produced by the GRD model trained after 1 x 106 steps. The redistributed rewards are shown in red, and the grounded rewards shown blue.
T
1
4
6
7
aL
$/
514
(b) Causal Structure from arz Sjees
(d) Causal Structure from ai to
0
3
5
0.2
Si,
S1141
(a) Causul Structure from
(c) Causal Structure from Sig
(e) Color Bar
Si Sj,411
Figure 11: Learned causal structure in Swimmer-v2.
22


Reviews:
[{"review_id": "sZjntEeKHz", "reviewer": "Reviewer_Xv6U", "summary": "The paper addresses a major challenge in reinforcement learning: identifying which state-action pairs contribute to delayed future rewards. They propose a solution called \"Return Decomposition\" that redistributes rewards from observed sequences while maintaining policy invariance. Unlike other methods, their approach explicitly models state and action contributions from a causal perspective, making it interpretable.\n\nThe authors introduce a framework called \"Generative Return Decomposition (GRD)\" for optimizing policies in scenarios with delayed rewards. GRD identifies unobservable Markovian rewards and causal relationships in the generative process. Using this causal generative model, GRD creates a compact representation to train policies efficiently.\n\nThe paper proves the identifiability of the unobservable Markovian reward function and the underlying causal structure and causal models. Experimental results show that their method outperforms existing techniques, and visualizations demonstrate its interpretability. The source code for their approach is publicly available.", "questions": "I have no knowledge about reinforcement learning, and I don't understand why the system assigned me to review papers on this topic. Please disregard my review comments.", "limitations": "I have no knowledge about reinforcement learning, and I don't understand why the system assigned me to review papers on this topic. Please disregard my review comments.", "rating": 5, "confidence": 1, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "  - The authors provide theoretical proof for the identifiability of the unobservable Markovian reward function and the underlying causal structure. This solidifies the theoretical foundation and robustness of the model.\n  - The GRD method outperforms state-of-the-art methods in experimental results across a range of tasks. This demonstrates its practical effectiveness and application potential.\n  - Visualization of the learned causal structure and decomposed rewards contributes to the interpretability aspect, a valued characteristic in contemporary machine learning.", "weaknesses": "I have no knowledge about reinforcement learning, and I don't understand why the system assigned me to review papers on this topic. Please disregard my review comments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer Xv6U\n\nWe appreciate you for reviewing our paper! Thank you for your positive support!"}, {"review_id": "psuK92iisM", "reviewer": "Reviewer_6iFi", "summary": "The paper introduces a new algorithm called Generative Return Decomposition (GRD) for return decomposition with causal treatment. GRD addresses the problem by modeling causal relationships among variables, providing advantages over flat representations. It specifies each state and action as a combination of constituent variables and considers causal relationships within the system. The algorithm utilizes a factored representation similar to Factored MDP, enabling the formation and identification of the Markovian reward function based on causality. Unlike previous approaches, GRD uses a graphical representation to determine the contribution of each dimension of state and action to the Markovian reward. It also explains and models the observed delayed return as a causal effect of the unobserved Markovian reward sequence. The framework of GRD visualizes the causal relationships among environmental variables. The paper proves the identifiability of the underlying generative process and introduces a component-wise learning approach for recovering the causal generative process and redistributing rewards. The learned parameters provide a minimal sufficient representation for policy training, aiding in the effectiveness and stability of policy learning. The main contributions of the paper include the reformulation of return decomposition with a graphical representation, the introduction of GRD for learning the causal generative process, and empirical experiments demonstrating the method's superiority over state-of-the-art approaches in robot tasks with sparse rewards.", "questions": "1 - How do you decide which trajectory to store in the memory for training? \n\n2 - Why is the dynamics model needed? The reward redistribution should be possible without the dynamics model.", "limitations": "Yes, the limitations have been addressed. ", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "1 - Interpretability: Having interpretable reward redistribution is an advantage over non-interpretable methods. This can be used to diagnose the reason for failures policy optimization. \n\n2 - Reduces the state dimensionality: A very nice side effect of learning causal masks using a dynamics models is that a policy can be learned using very few features of the state. This leads to simpler policies, which could be more robust. \n", "weaknesses": "1 - Writing: The paper needs a lot of work in explaining the method. Especially section 4 and section 5.1. A figure showing how the causal masks are applied would be a good idea. I am willing to improve my score, if the method explanation is improved. \n\n2 - Experiments: The experiments include only Mujoco tasks. It would be interesting to see how the method behaves on delayed reward Atari environments like Bowling. \n\nMissing Related work: \n[1] Modern hopfield networks for return decomposition for delayed rewards", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer 6iFi\n\nThank you for your positive support for our paper! Below we provide a point-wise response to your concerns.\n\n**Weakness 1:** Writing: The paper needs a lot of work in explaining the method. Especially section 4 and section 5.1. A figure showing how the causal masks are applied would be a good idea. I am willing to improve my score, if the method explanation is improved.\n> **Reply 1:** Thank you for your comments. We will revise Sec 4 and Sec 5.1 in the future version to include more clear and detailed explanations, including a figure for the overall framework. Additionally, please refer to Figure 2 in the attached PDF, which illustrates how the causal masks are applied.\n\n\n**Weakness 2:** Experiments: The experiments include only Mujoco tasks. It would be interesting to see how the method behaves on delayed reward Atari environments like Bowling.\n> **Reply 2:** Thank you for your suggestion. Our method focuses on state-based environments, which Atari does not satisfy. However, to further showcase the applicability of GRD, we provide additional experimental results:\n> 1) on three tasks from Meta-World, *pick-place-v2*, *push-wall-v2* and *door-lock-v2* (Figure 4) to demonstrate better performance of GRD, compared with the baseline methods. \n> 2) with different RL backbones (Figure 2), *TD3*, *DDPG*, to show the consistency improvement of GRD on the *HalfCheetah-v2*. \n> 3) against the Gaussian noise with different standard deviations in the insignificant dimensions of state in *Ant-v2* to demonstrate the more robust performance compared with baselines. During the evaluation, the noise is inserted in the $28\\sim 111$ dimensions, which are not used in *Ant-v2*. The performance of GRD is not affected due to correctly identifying the compact representation for policy learning.\n\n**Weakness 3:** Missing Related work: [1] Modern hopfield networks for return decomposition for delayed rewards.\n> **Reply 3:** Thank you for pointing this out. Hopfield-RUDDER improves RUDDER by the replacement of LSTM with a continuous modern Hopfield network, and further employs history compression to facilitate the detection and storage of the key events. However, it still follows the line of RUDDER and shares the drawback of lacking interpretability. We will incorporate the suggested paper in the future version.\n\n\n\n**Q1:** How do you decide which trajectory to store in the memory for training?\n> **A1:** We collect trajectories using the on-training policy and subsequently store them in the buffer. During model training, we uniformly sample data from this buffer.\n\n**Q2:** Why is the dynamics model needed? The reward redistribution should be possible without the dynamics model.\n> **A2:** It is true that reward redistribution doesn't necessitate a dynamics model. We construct an ablation version of GRD, GRD wo CR. It only relies on the learned causal structure and reward function to train policy, and does not perform as well as GRD. Incorporating the dynamics model benefits policy learning by helping identify a compact representation: through learning the dynamics model, we gain insights into the causal relationships that dictate the generation of the next state, *i.e.*, how the $\\boldsymbol{s} _{t}$ determine the next state $\\boldsymbol{s} _{t+1}$ . Such causal knowledge helps identify the minimal sufficient dimensions of the state for policy learning ($\\boldsymbol{\\tilde{c}} ^{\\boldsymbol{s}\\rightarrow \\pi}$ in Eq. 7), which is called compact representation in our paper. Compact representation is associated with the learned causal structure for generating Markovian reward and varies with the learning of the reward function. Since the supervision signals for policy learning are generated by the learned reward function, the policy learning consistently be optimized within the smallest state space, resulting in a more efficient training procedure.\n\n\n"}, {"review_id": "qYbc1bJsjR", "reviewer": "Reviewer_ynzt", "summary": "Delayed reward in reinforcement learning is the major challenge in reinforcement learning. The return distribution technique is the direct way to resolve this issue while preserving policy. The existing works redistribute the returns in an uninterpretable manner. In this regard, this paper proposes a GRD which generates the Markovian rewards in delayed reward scenarios. GRD first checks the casual relations of state and actions and from a compact representation using causal generative model. The experiment results show that GRD outperforms the baselines and helps visualization. ", "questions": "* I could not understand the experiment results in l345-352.", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* Experiment results seem promising.\n", "weaknesses": "* Explanation on line 345-352 is not sufficient and hard to understand. This experiment section is very important since authors insist that GRD give the interpretable structure of reward. \n* The methods have to construct casual inference which is only possible when all the states are exactly defined. If the number of states explodes, the parameters to learn casual structure would explode. If states and actions are given, we can construct the casual structure without learning with parameters.\n\nMinor\n* Equations are too messy and hard to understand. Use under bracket in equation 6. \n* Notations are not familiar. It is hard to understand C^{\\cdot -> \\cdot}, d^a , d^s (?). The authors have to redefine all the variables step-by-step to improve the presentation of this paper.\n* The arrow size is not consistent in Figure 2. The arrow to \\hat{r}_3 is narrower. Also, arrow directions which come from R are weird. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer ynzt\n\n\nThanks for your constructive feedback. We provide a point-to-point response below.\n\n**Weakness 1:** Explanation on line 345-352 is not sufficient and hard to understand. This experiment section is very important since authors insist that GRD give the interpretable structure of reward.\n\n> **Reply 1:** We appreciate your observation, but there seems to be a misunderstanding. The experimental result (Line 345 - 352) aims to demonstrate the accuracy of predicting Markovian reward by our learned model, rather than the interpretable structure of the learned reward function (presented in L327 - 344).  We visualize the comparison of decomposed rewards and the ground truth rewards to demonstrate the accuracy of Markovian reward prediction by GRD. As shown in Figure 5, the blue lines, representing redistributed rewards, consistently align with the corresponding red lines, which are the ground truth. This visualization demonstrates that GRD indeed distinguishes the state action pairs with less contribution to the long-term return (episodic reward) from those with more contributions.  We will revise this section in the future version to provide a clearer explanation.\n\n**Weakness 2:** The methods have to construct causal inference which is only possible when all the states are exactly defined. If the number of states explodes, the parameters to learn causal structure would explode. If states and actions are given, we can construct the causal structure without learning with parameters.\n> **Reply 2:** Thanks for the question. This work focuses mainly on state-based tasks. 1) For the setting where states are not defined such as image-like input, our work can be applied by employed in the latent state space, which requires learning a latent vector representation, as AdaRL [1] does.  2) For a large number of dimensions of state, this would increase the complexity of the environment dynamics and thus requires more efficient causal discovery methods. 3) As an advantage, by learning the causal structure, we can constrain the optimization of the models over a small subspace of state and action, resulting in a lower requirement of parameters of the neural network. 4) Yes, it is possible to construct causal structures without learning the models, but it requires additional diverse data for causal discovery. Overall, the main focus of this work is to address the delayed rewards by interpretable return decomposition, therefore we consider the basic causal discovery algorithm.\n\n\n**Weakness 3:** Equations are too messy and hard to understand. Use under bracket in equation 6.\n\n> **Reply 3:** Eq. 6 is introduced to regulate the sparsity of learned causal structure to avoid trivial solutions. It is achieved by optimizing the parameters towards the direction of the nonexistence of the causal edge. Below, we revise Eq. 6 and hope to provide better delineation and understanding. \n> \n> Let $D_i(\\boldsymbol{x})=\\log P(\\boldsymbol{x} _i)$, where $P(\\boldsymbol{x}_i)$ is the possibility that the edge $\\boldsymbol{x} _i$ exists, . Minimizing $D _i(\\boldsymbol{s})$ prevents the causal edge from existing. Then Eq. 6 is,\n> \\begin{array}{ll}\nL _{\\text{sp}}(\\phi _{\\text{cau}}) = \n\\underbrace{\\lambda _1 \\sum _i D _i(\\boldsymbol{c}^{\\boldsymbol s\\rightarrow r})} _{\\text{state-to-reward}} +\n\\underbrace{\\lambda_2 \\sum _i D _i(\\boldsymbol{c}^{\\boldsymbol a\\rightarrow r})} _{\\text{action-to-reward}} + \\\\\n\\underbrace{\\lambda _3 \\sum _{j \\ne i} D _{i, j}(\\boldsymbol{C}^{\\boldsymbol s\\rightarrow \\boldsymbol s})} _{\\text{state-to-state (excluding self-connections)}} +\n\\underbrace{\\lambda _4 \\sum _{j = i} D _{i, j}(\\boldsymbol{C}^{\\boldsymbol s\\rightarrow \\boldsymbol s})} _{\\text{state-to-state (self-connections)}} +\n\\underbrace{\\lambda_5 \\sum _{j,i} D _{i, j}(\\boldsymbol{C}^{\\boldsymbol a\\rightarrow \\boldsymbol s})} _{\\text{action-to-state}}.\n\\end{array}\n> \n> These five terms are responsible for the sparsity of the causal structures of state-to-reward， action-to-reward， state-to-state (excluding self-connections)， state-to-state (self connections) and action-to-state, separately. Here self-connection represents the causal edge from $\\boldsymbol{s} _{i, t}$ to $\\boldsymbol{s} _{i, t+1}$.  \n\n**Weakness 4:** Notations are not familiar. It is hard to understand $C^{\\cdot \\rightarrow \\cdot}$, $d^a$ , $d^s$ (?). The authors have to redefine all the variables step-by-step to improve the presentation of this paper.\n\n> **Reply 4:** We genuinely appreciate your feedback on this matter, and we will revisit and clarify the symbols in our paper. Here we explain the mentioned notations:\n> - $\\boldsymbol{C}^{\\cdot \\rightarrow \\cdot}$ denotes all the causal masks, *i.e.*, $\\boldsymbol{C}^{\\cdot \\rightarrow \\cdot}:=[\\boldsymbol{C}^{\\boldsymbol{s}\\rightarrow \\boldsymbol{s}}, \\boldsymbol{C}^{\\boldsymbol{a} \\rightarrow \\boldsymbol{s}}, \\boldsymbol{c}^{\\boldsymbol{s} \\rightarrow r}, \\boldsymbol{c}^{\\boldsymbol{a} \\rightarrow r}]$;\n> - $d^\\boldsymbol{s}$ denotes the number of dimensions of state;\n> - $d^\\boldsymbol{a}$ denotes the number of dimensions of action.\n\n\n**Weakness 5:** The arrow size is not consistent in Figure 2. The arrow to $\\hat{r}_3$ is narrower. Also, arrow directions which come from R are weird.\n\n> **Reply 5:** Thank you for pointing this out. We will revise and ensure consistency in the size and direction of the arrows in Figure 2.\n\n\n**Reference**\n\n[1] Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. In International Conference on Learning Representations, 2021.\n"}, {"review_id": "JTj0D1hGb9", "reviewer": "Reviewer_EmbG", "summary": "This study introduces a novel approach, termed Generative Return Decomposition (GRD), to address a key challenge in reinforcement learning: identifying the state-action pairs that contribute to future, delayed rewards. While many methods redistribute rewards in a non-transparent manner, GRD offers a clear return decomposition by explicitly modeling the contributions of states and actions from a causal perspective.\n\nGRD works by first recognizing unobservable Markovian rewards and causal relations in the data generation process. Then, it leverages these to create a compact representation for policy training over the agent's most favorable state-space subset. The researchers provide theoretical proof of the identifiability of the Markovian reward function and underlying causal structure and models. Experimental data also reveal GRD's superior performance and interpretability compared to other methods. However, some limitations exist due to the assumptions made, such as the stationary nature of the reward function, which may not be applicable in dynamic or online RL scenarios.", "questions": "See Weaknesses.", "limitations": "See Weaknesses.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper excels in presenting Generative Return Decomposition (GRD), an innovative method that improves interpretability in reinforcement learning. GRD successfully addresses the identification of impactful state-action pairs for future rewards. Its effectiveness is supported by both theoretical evidence and practical experiments, demonstrating its superior performance over other existing methods. Moreover, the paper demonstrates a high degree of clarity and coherence, enabling smooth comprehension. Additionally, the explicit description of assumptions contributes to a more profound comprehension of the inherent strengths and weaknesses of the study.", "weaknesses": "- The quality of the text in Figure 1 could be improved by removing the shadow around the text. The same applies to Figure 2.\n- Minor typos and errors that need to be edited for the next version of the paper. E.g. in line 172: \"provide\" -> \"provides\"\n- In Section 5, only one policy is considered, which is SAC. How about having the experiments run based on another policy optimization algorithm? What would be the differences in performance and results?\n- Regarding the last paragraph of Section 5, there are two possible scenarios to train the agent.\n    - First, the generative model is learned while the policy is being updated, in an end-to-end paradigm.\n    - Second, the generative model is first trained (and stays fixed thereafter), then the policy begins to be optimized.\n\n    In either case, how would that affect the policy training and performance? And how the insights from the GRD interpretations would be changed?\n\n- In the experiment section, the visualizations for the learned causal structure are only provided for Ant. Please provide the same type of analysis for other environments.\n\n- Considering Figure 4, having GRD, how the agent's robustness and generalizability would be affected? For example, consider the case where there are some anomalies injected into the agent and environment interaction, more specifically changing some values from the state-space. If such anomalies target the features that are less important to the agent, then its performance should not be affected that much, right? If so, could you provide some results in this regard?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer EmbG\n\nWe thank the reviewer for the comments. Below please see our responses as well as clarifications.\n\n**Weakness 1:** The quality of the text in Figure 1 could be improved by removing the shadow around the text. The same applies to Figure 2.\n> **Reply 1:** Thank you for your suggestion. We will revise the figures in the future version\n\n**Weakness 2:** Minor typos and errors that need to be edited for the next version of the paper. E.g. in line 172: \"provide\" -> \"provides\"\n> **Replay 2** Thank you for the comments. We will revise in the future version.\n\n**Weakness 3:** In Section 5, only one policy is considered, which is SAC. How about having the experiments run based on another policy optimization algorithm? What would be the differences in performance and results? \n> **Reply 3:** We provide results training with DDPG and TD3 in Figure 3 in the attached PDF. As the experimental result shows, on the tasks of *HalfCheetah*, GRD consistently outperforms the baseline methods, *RRD-Bias*, *RRD-Unbias*, and *IRCR*, which are modified to run based on the same policy optimization algorithm, DDPG and TD3. We also provide results of *None*, which utilizes the observed delayed reward for policy learning directly.\n\n**Weakness 4:** Regarding the last paragraph of Section 5, there are two possible scenarios to train the agent. \nFirst, the generative model is learned while the policy is being updated, in an end-to-end paradigm.\nSecond, the generative model is first trained (and stays fixed thereafter), then the policy begins to be optimized.\nIn either case, how would that affect the policy training and performance? And how the insights from the GRD interpretations would be changed?\n> **Reply 4:** Although these two lines are possible for learning the reward model, we follow the first line and enjoy the following advantages: \n> 1) We can use the on-training policy to collect data rather than the random policy, which would collect data without diversity.\n> 2) Apart from that, it is more data-efficient to optimize simultaneously since it avoids collecting data for learning the generative model and the policy separately.\n> \n> As for the insight of GRD, it will not be changed, since GRD models the generative process of the environment, which is totally the same for the two different lines.\n\n**Weakness 5:** In the experiment section, the visualizations for the learned causal structure are only provided for Ant. Please provide the same type of analysis for other environments.\n> **Reply 5:** Thank you for pointing that out. We provide the visualization of learned causal structure in *Swimmer-v2*, as shown in Figure 5 in the attached PDF, and more results for other environments will be presented in the future version.\n> Since the grounded causal structure is not accessible, we verify the reasonability of the learned causal structure by some observations:\n> - All the edges from different dimensions of $\\boldsymbol a$ to $r$ always exist, as shown in Figure 5 (d): *Swimmer-v2* shares the same characteristic that the edges from different dimensions of $\\boldsymbol a$ to $r$ always exist, corresponding with the reward design of penalizing the swimmer if it takes actions that are too large, measured by the sum of the values of all the action dimensions.\n> - According to Figure 5 (b), the first dimension of action (Torque applied on the first rotor) has an impact on the last three dimensions of state (angular velocity of front tip, angular velocity of first rotor, second rotor), which is corresponding with the intuitive that the part that connects to first rotor should be impacted by the first dimension of action. We can get a similar observation for the second action dimension from Figure 5.\n> - We can observe that all the state dimensions are learned to be connected to the reward; the possible explanation is that in the swimmer robot, any changes of the front tip, or two rotors will impact the position of the robot, potentially influencing the reward.\n\n\n**Weakness 6:** Considering Figure 4, having GRD, how the agent's robustness and generalizability would be affected? For example, consider the case where there are some anomalies injected into the agent and environment interaction, more specifically changing some values from the state-space. If such anomalies target the features that are less important to the agent, then its performance should not be affected that much, right? If so, could you provide some results in this regard?\n> **Reply 6:** Sure, we provide the results with certain noisy states at the insignificant dimensions of the state, as shown in Figure 1 in the attached PDF. To mimic the anomalies, we introduced independent Gaussian noise (mean: $0$, std: $0\\sim 1$) to the dimensions ranging from $28$ to $111$ while evaluating the policy. According to Figure 1, unlike the baseline methods, GRD is unaffected by the injected noises, demonstrating that GRD is more robust than others. That is because the insignificant dimensions are not in the compact representation, which serves as the input of policy. \n"}, {"review_id": "KMnFTtn4on", "reviewer": "Reviewer_24n7", "summary": "This paper proposes a novel algorithm for return decomposition with causal treatment. To do reward redistribution, GRD uses factored representations to model the Markovian reward function and dynamics function. ", "questions": "1 The technical contribution of the paper seems to be limited. The major contribution is the learned generative model. However, the proposed framework is similar to [27], which also uses factored MDP, learns a mask on the factors, and optimizes the model in a generative way. Is it correct to say GRD equals [27] adds markovian reward functions?\n\n2 As for the experiments in Fig.3, GRD does not seem to outperform existing methods by a large margin. Except Halfcheetah, swimmer and humanoid stand up, GRD seems to be rather close to the baselines. Maybe it will be better to further test on some environment with more obvious performance gap.\n\n3 It is interesting to see some horizontal lines in Fig. 4(a), where some dimensions have definite contribution to dimensions #27:54. I think it will be helpful to provide illustrations to correspond this to the mujoco ground truth dynamics, which proves that GRD indeed learns the true and interpretable casual structures.\n", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The writing is clear and easy to follow. It is interesting to see the visualization in section 6.4, especially Fig. 4.", "weaknesses": "The technical contribution is somehow limited and the stronger experiments are expected.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer 24n7\n\nThank you for your positive support and constructive comments. We provide our point-wise response below.\n\n**Weakness 1:** The technical contribution is somehow limited and the stronger experiments are expected.\n> **Reply 1:** Thank you for your comments. \n> As complemented by Reviewer 6iFi, 1) \"Interpretability: Having interpretable reward redistribution is an advantage over non-interpretable methods. This can be used to diagnose the reason for failures policy optimization.\" 2) \"Reduces the state dimensionality: A very nice side effect of learning causal masks using a dynamics models is that a policy can be learned using very few features of the state. This leads to simpler policies, which could be more robust.\", our contribution of providing an interpretable solution for delayed rewards is recognized by other reviewers, and we also reclarify our technical contribution in the Q1. As for the experiments, we provide additional experimental results on other RL training backbones (Figure 3 in attached PDF), on the tasks from Metaworld environment (Figure 4 in attached PDF) to demonstrate the state-of-the-art performance, as well as against noisy state to demonstrate the robustness of GRD. Please refer to the attached PDF.\n\n**Q1:** The technical contribution of the paper seems to be limited. The major contribution is the learned generative model. However, the proposed framework is similar to [27], which also uses factored MDP, learns a mask on the factors, and optimizes the model in a generative way. Is it correct to say GRD equals [27] adds markovian reward functions? \n> **A1:** Thanks for your questions. Our work differs from [27] in three ways. 1)Different task: [27] aims to identify the change factors across different domains and does not address the delayed rewards that we are interested in.  2) Identifiability: For GRD, without the given long-term return, the reward function and the related causal structure are not identifiable. For [27], the reward function is identifiable given the Markovian rewards, which are not observable in our setting. 3) Estimation method: due to the unobservable Markovian reward, we treat the long-term return as the causal effect of all the Markovian rewards within an episode and use the corresponding loss (Eq 4). Additionally, different from [27], we treat the existence of the causal edges as variables, leading to different losses for the minimal-edge assumption (likelihood in GRD and L1 loss in [27]) and different model structures (Gumbel-softmax for sampling in GRD). Overall, although we both utilize DBN, factored representation and model the causal structure as the binary masks, which are the common line for causal modeling [1][2], our work emphasizes the importance of learning an interpretable reward function and is the first to introduce causality into return decomposition.  \n\n\n\n**Q2:** As for the experiments in Fig.3, GRD does not seem to outperform existing methods by a large margin. Except Halfcheetah, swimmer and humanoid stand up, GRD seems to be rather close to the baselines. Maybe it will be better to further test on some environment with more obvious performance gap. \n> **A2:** Thank you for the observation. We've extended our experiments to include the Meta-World environment, further showcasing our method's performance. Given the time constraints, we do not compare GRD with IRCR as it does not perform as well as RRD.  Besides this merit, GRD demonstrates 1) interpretability (our primary objective) and 2) robustness compared to the baseline methods. As an example of 2), upon introducing Gaussian noise to the $28\\sim 111$ dimensions of states in Ant-v2, the performance of GRD does not decrease. Please refer to Figure 1 in the attached PDF for the results.\n\n**Q3:** It is interesting to see some horizontal lines in Fig. 4(a), where some dimensions have definite contribution to dimensions #27:54. I think it will be helpful to provide illustrations to correspond this to the mujoco ground truth dynamics, which proves that GRD indeed learns the true and interpretable causal structures.\n> **A3:** While the corresponding real causal structure would be helpful to verify the interpretability of our method, it is inaccessible. Therefore, we provide some evidence to underscore the reliability of the learned causal structure:\n> - there should not be an edge from the unused dimensions of the state variable to the other variables. (L334-L336)\n> - the edges from different dimensions of $\\boldsymbol a$ to $r$ always exist, corresponding with the reward design of penalising the robot if it takes actions that are too large, measured by the sum of the values of all the action dimensions. (L340-342)\n> \n> Apart from above, we explain why the learned redundant edges (horizontal lines) do not impact the policy learning in L336-340: there is no edge from the  $28\\sim 111$ dimensions of the next state to the reward, i.e., these dimensions do not exist in the identified compact representation, thus having no influence on policy learning.\n> \n\n**Reference**\n\n[1] Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. In International Conference on Learning Representations, 2021.\n\n[2] Huang, B., Lu, C., Leqi, L., Hernández-Lobato, J. M., Glymour, C., Schölkopf, B., & Zhang, K. (2022, June). Action-sufficient state representation learning for control with structural constraints. In International Conference on Machine Learning (pp. 9260-9279). PMLR."}]

Author rebuttal:
# Attached PDF

Thank you to all the reviewers for your invaluable insights and thoughtful feedback. Your expertise has greatly helped us refine and improve our work.

Here we provide five figures in the attached PDF as supplementary of our response:

- Figure 1: Evaluation with Gaussian noise.
- Figure 2: The illustration of using learnable masks to predict the next state.
- Figure 3: Learning curves on *HalfCheetch-v2* with different training backbone, DDPG and TD3.
- Figure 4: Learning curves on three tasks from *MetaWorld*, Door Lock, Push Wall, and Pick Place.
- Figure 5: Learned causal structure in *Swimmer-v2*.


Meta-review: The reviews and discussion offered cautious to enough positive support. We encourage the authors to use the feedback to further improve the work into the final version.
Decision: Accept (poster)