Title: Interpretable Reward Redistribution in Reinforcement Learning: A Causal Approach
Venue: NeurIPS 2023
Accepted: yes
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

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/.

1 Introduction

Reinforcement Learning (RL) has achieved significant success in a variety of applications, such as autonomous driving [1, 2], robot [3, 4], games [5, 6, 7], financial trading [8, 9], and healthcare [10]. A challenge in real-world RL is the delay of reward feedback [11, 12, 13, 14]. In such a delayed reward setting, the learning process may suffer from instability due to the lack of proper guidance from the reward sequences [15, 16, 17]. Methods such as reward shaping [18, 19], curiosity-driven intrinsic reward [20, 21, 22, 23] and hindsight relabeling [24, 25, 26, 13] have been proposed to provide proxy rewards in RL.

How to compute the contribution of each state-action pair affected by a delayed reward and explain the reasons behind such contribution are equally important, as they can provide insights into the underlying dynamics 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 at each time step [15, 16, 17]. These methods allow for the

decomposition of the trajectory-wise return into Markovian rewards, preserving the policy invariance. However, previous methods construct reward redistribution by hand-designed rules [15, 16], or an uninterpretable model [17], 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, 28]. 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 action 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 a structural factored representation, relating to Factored MDP [32, 33], provides favorability to form and identify the Markovian reward function from the perspective of causality. Unlike previous approaches, GRDutilizes such a parsimonious graphical representation to discover how each dimension of state and action contributes to the Markovian reward. Moreover, 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 1 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 s · ,t and action a · ,t , Markovian rewards r t for t ∈ [1 , T ] , as well as the long-term return R . For sparse reward settings in RL, the Markovian rewards r t 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 t and evaluates the performance of the agent within the whole episode. A special case of delayed rewards is in episodic RL, where o 1: T -1 = 0 and o T = 0 . Theoretically, we prove that the underlying generative process, including the unknown causal structure, the Markovian reward function, and the dynamics function, are identifiable. GRD learns the Markovian reward function and dynamics function in a component-wise 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.

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 a compact representation to facilitate policy learning. Furthermore, empirical experiments on robot tasks from the MuJoCo environment demonstrate that our method outperforms the state-of-the-art methods and the provided visualization shows the interpretability of the redistributed rewards, indicating the usefulness of causal modeling for the sparse reward settings.

2 Related Work

Below we review the related work on reward redistribution and causality-facilitated RL methods.

Previous work redistributes rewards to deal with the delayed rewards in RL, such as reward shaping [18, 19] and curiosity-driven intrinsic reward [20, 21, 22, 23]. Return decomposition draws attention since RUDDER [15] rethinks the return-equivalent condition for reward shaping [34] to decompose the long-term return into proxy rewards for each time step through an LSTM-based long-term return predictor and manually designed assignment. Align-RUDDER [16] introduces the bioinformatical alignment method to align successful demonstration sequences, then manually scores new sequences according to the aligned demonstrations. However, in 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-equivalent 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 [17, 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, [27] learns factored representation and an individual change factor for different domains, and [31] extends it to cope with 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 is introduced to multi-agent task [41, 42], model-based RL [43], imitation learning [44] 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 a 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 [46] and return decomposition [15, 16, 17].

Markov Decision Process (MDP) is represented by a tuple ⟨S , A , R , γ, P ⟩ , where S , A , R denote the state space, action space, and reward function, separately. The state transition probability of the environment is denoted as P ( s t +1 | s t , a t ) . The goal of reinforcement learning is to find an optimal policy π : S → A to maximize the expected long-term return, i.e ., a discounted sum of the rewards with the predefined discounted factor γ and episode length T , J ( π ) = E s t ∼ P ( s t | s t -1 ,a t -1 ) ,a t ∼ π ( a t | s t ) [ ∑ T t =1 γ t -1 R ( s t , a t )] .

̸

Return decomposition is proposed to decompose the long-term return feedback into a sequence of Markovian rewards in RL with delayed rewards [15, 16] 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 o T at the end of each episode, i.e ., o 1: T -1 = 0 and o T = 0 , called episodic reward [17]. In general, the observed rewards o t are delayed and usually harm policy learning, since the contributions of the state-action pairs are not clear. Therefore, return decomposition is proposed to generate proxy rewards ˆ r t for each time step, which is expected to be non-delayed, dense, Markovian, and able to clarify the contributions of the state-action pairs to the delayed feedback. The work along this line shares a common assumption,

where R is the long-term return and o t , ˆ r t denote the observed delayed reward and the decomposed proxy reward for each time step, separately.

4 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 a 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 = ⟨S , A , g, P f , γ, G⟩ . S , A represent the sets of states s and actions a , respectively. G denotes a DBN that describes the causal relationship within the MDP environment, constructed over a finite number of random variables, { s 1 ,t , · · · , s | s | ,t , a 1 ,t , · · · , a | a | ,t , r t } T t =1 ∪ R where | s | and | a | are the dimension of s and a . r t is the unobservable Markovian reward that serves as the objective of return decomposition. The state transition P f can be represented as, P f ( s t +1 | s t , a t ) = ∏ | s | i =1 P ( s i,t +1 | pa ( s i,t +1 )) , where i is the index of dimension of s t . Here pa ( s i,t +1 ) denotes the causal parents of s i,t +1 and usually is a subset of the dimensions of s t and a t . We assume a given initial state distribution P ( s 1 ) . Similarly, we define the functions, g , which generates unobservable Markovian reward, r t = g ( pa ( r t )) , where pa ( r t ) is a subset of the dimensions of s t and a t . The trajectory-wise long-term return R is the causal effect of all the Markovian rewards. An example of the causal structure denoted by G is given in Figure 1. For simplicity, we reformalize the generative process as,

for i = (1 , 2 , · · · , | s | ) . ⊙ is the element-wise product. f and g stand for the dynamics function and reward function, separately. C ·→· are categorical masks organizing the causal structure in G . ϵ s,i,t and ϵ r,t are i.i.d random noises. Specially, C s → r ∈ { 0 , 1 } | s | and C a → r ∈ { 0 , 1 } | a | control if a specific dimension of state s t and action a t impact the Markovian reward r t , separately. For example, if there is an edge from s i,t to r t , then the C s → r i = 1 . Similarly, C s → s ∈ { 0 , 1 } | s |×| s | and C a → s ∈ { 0 , 1 } | a |×| s | indicate the causal relationship between s t , a t and s t +1 , separately. In particular, we assume that G is time-invariant, i.e ., f , g , C ·→· are invariant. In the system, s t , a t , and R are observable, while r t 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 s t , action a t , trajectory-wise long-term return R are observable while Markovian rewards r t are unobservable, and they form an MDP, as described in Eq. 2. Then, under the global Markov condition and faithfulness assumption, the reward function g and the Markovian rewards r t are identifiable, as well as the causal structure that is characterized by binary masks C ·→· and the transition dynamics f .

Details of the proof for Proposition 1 are deferred to Appendix B. Proposition 1 provides the foundation for us to identify the causal structure C ·→· and the unknown functions f , g in the generative process from the observed data. As a result, with the identified structures and functions in Eq. 2, 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 [17, 35], g in Eq. 2 is a general description to characterize the Markovian reward as the causal effect of a subset of the dimensions of state and action, 2) compared with [15, 16], treating long-term reward R as the causal effect of all the Markovian rewards in a sequence is an interpretation of return-equivalent assumption from the causal view, allowing a reasonable and flexible return decomposition.

Policy Optimization

Generative Model Learning

Causal

Structure

...

Reward as

Guidence

CR: Compact Representation

CR as Input

...

5 Generative Return Decomposition

In this section, we propose a principled framework for reward redistribution from the causal perspective, named Generative Return Decomposition (GRD). Specifically, we first introduce how GRD recovers a 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, Φ m for the parameterized generative process and Φ π for the policy. Therefore, the corresponding loss function includes L m (Φ m ) and J π (Φ π ) , for causal generative model estimation and policy optimization, respectively. Hence, the overall objective can be formulated as

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 Φ m 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 τ = {⟨ s t , a t ⟩ | T t =1 , R } , where R is the discounted sum of observed delayed rewards o t , as given in Eq. 1.

Generative Model. The parameterized model Φ m consists of three parts, ϕ cau , ϕ dyn , ϕ rew, which are described below. (1) ϕ cau is used to identify the causal structure in G by predicting the values of binary masks in Eq. 2. It consists of four parts, ϕ s → s cau ∈ R | s |×| s |× 2 , ϕ a → s cau ∈ R | a |×| s |× 2 , ϕ s → r cau ∈ R | s |× 2 and ϕ a → r cau ∈ R | a |× 2 , which are used to predict C s → s , C a → s , C s → r and C a → r , separately. We take the prediction of C s → s with ϕ s → s cau as an example to explain the process of predicting causal structure. ϕ s → s cau characterizes | s | ×| s | i.i.d Bernoulli distributions for the existence of the edges organized in the matrix, C s → s . 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 an edge from i -th dimension of the state to j -th dimension of the next state in the causal graph as P ( C s → s i,j ) . Binary prediction of C s → s is sampled by applying element-wise gumbel-softmax [48] during training, while using greedy selection during

inference. Similarly, we can obtain P ( C a → s i,j ) , P ( C s → r i ) , P ( C a → r i ) , C a → s , C s → r and C a → r . (2) ϕ rew is constructed with fully-connected layers and approximates the reward function g in Eq. 2, which takes as input the state s t , action a t , as well as the predictions of C s → r and C a → r to obtain the prediction of Markovian rewards. (3) ϕ dyn is constructed on a Mixture Density Network [49], and is used to approximate the dynamics function f in Eq. 2, which takes as inputs the state s t , action a t , the predictions of causal structures, C s → s and C a → s . As an example, the predicted distribution for the i -th dimension of the next state is P ( s i,t +1 | s t , a t , C s → s · ,i , C a → s · ,i ; ϕ dyn ) . More details for Φ m can be found in the Appendix C.2.

Loss Terms. Accordingly, the loss term L m contains three components with L m (Φ m ) = L rew + L dyn + L reg. Considering the long-term return R as the causal effect of the Markovian rewards r t , we optimize ϕ rew , ϕ s → r cau and ϕ a → r cau by

where ˆ r t is predicted by ϕ rew , i.e ., ˆ r t = ϕ rew ( s t , a t , C s → r , C a → r ) . To optimize ϕ dyn, ϕ s → s cau and ϕ a → s cau , we minimize

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 D i ( x ) = log P ( x i ) , where P ( x i ) is the possibility that the edge x i exists. The regularizer term is,

̸

where ✶ ( · ) is the indicator function and hyper-parameters λ ( · ) are listed in Appendix C.5. 5.2 Policy Optimisation with Generative Models

Here we explain how the learned generative model aids policy optimization with delayed rewards.

Compact Representation as Inputs. Inspired by [30], we identify a minimal and sufficient state set, s min , for policy learning as the input of the policy model, Φ π , called compact representation . It contains the states' dimensions that directly or indirectly impact the reward. To be more specific, it includes each dimension of state s i,t ∈ s t , which either 1) has an edge to the reward r t i.e ., C s → r i = 1 , or 2) has an edge to another state component in the next time-step, s j,t +1 ∈ s t , i.e ., C s → s i,j = 1 , such that the same component at time t is in the compact representation, i.e ., s j,t ∈ s min t . We first select C s → s and C s → r greedily from the i.i.d Bernoulli distributions characterized by ϕ s → s cau and ϕ s → r cau (as illustrated in Appendix C.2) and organize the state dimensions which exist in s min t as C s → π . We define C := { i, ∀ C s → r i = 1 } to denote the dimensions of s t which have impact on the Markovian reward r t directly. Then we have,

where | C | denotes the size of C and ∨ denotes element-wise logical or operation. Then, the compact representation which serves as the input of policy is defined as, s min t = C s → π ⊙ s t . 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 s t and a t , we replace the observed delayed reward o t by ˆ r t = ϕ rew ( s t , a t , C s → r , C a → r ) , where C s → r and C a → r are selected greedily from ϕ s → r cau and ϕ a → r cau , as illustrated in Appendix C.2. The policy model Φ π contains two parts, a critic ϕ v

and an actor ϕ π due to the applied SAC algorithm. SAC alternates between a policy evaluation step, which trains a critic ϕ v ( s min t , a t ) to estimate Q π ( s min t , a t ) = E π [ ∑ T t =1 γ t -1 ˆ r ( s min t , a t ) | s 1 = s , a 1 = a ] using the Bellman backup operator, and a policy improvement step, which trains an actor ϕ π ( s min t ) by minimizing the expected KL-divergence,

where V π is the value function of state-action pairs [50]. More detailed implementation of Φ π 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 our method on eight widely used classical robot control tasks in the MuJoCo environment [55], including Half-Cheetah , Ant , Walker2d , Humanoid , Swimmer , Hopper , HumanoidStandup , 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 to 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 hyper-parameters for policy learning as RRD for a fair comparison. While RUDDER and Align-RUDDER are also popular methods for return decomposition, RRD and IRCR demonstrate superior performance compared to them. Moreover, Align-RUDDER requires successful trajectories to form a scoring rule for state-action pairs, which are unavailable in MuJoCo. Therefore, we do not compare our performance with theirs.

Evaluation Metric. We report the average accumulative reward across 5 seeds with random initialization to demonstrate the performance of evaluated methods. Intuitively, a method that earns higher accumulative rewards in evaluation indicates better reward redistribution.

6.2 Main Results

Figure 3 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, in HalfCheetah and Ant , GRD obtains the highest average accumulative reward of 1 . 5 × 10 4 and 6 . 5 × 10 3 , separately, while the baseline methods achieve the best value of 1 . 25 × 10 4 and 6 × 10 3 , respectively. In the task where the agents are equipped with high-dimension states, such as HumanoidStandup of 376 -dimension states and Ant of

0

20

40

60

80

100

Number of frames

(×10 4

)

0

1

2

3

×

10 2

Swimmer-v2

0

50

100

150

200

Number of frames

(×10 4

)

0

2

4

6

×

10 3

0

100

200

300

Number of frames

(×10 4

)

0.0

0.5

1.0

1.5

×

10 4

HalfCheetah-v2

0

50

100

150

Number of frames

(×10 4

)

0

1

2

3

×

10 3

Hopper-v2

0

50

100

150

200

Number of frames

(×10 4

)

1.0

0.8

0.6

0.4

0.2

0.0

×

10 3

Reacher-v2

0

100

200

300

Number of frames

(×10 4

)

0.5

1.0

1.5

2.0

2.5

3.0

×

10 5

HumanoidStandup-v2

0

100

200

300

Number of frames

(×10 4

)

0

2

4

6

×

10 3

0

20

40

60

80

Number of frames

(×10 4

)

0

1

2

3

4

5

×

10 3

Humanoid-v2

Average Accumulative Reward

RRD-Biased

RRD-Unbiased

IRCR

GRD w/o CR

GRD

111 -dimension states, GRD gains significant improvement. Taking HumanoidStandup as an example, GRDachieves a higher average accumulative reward of 2 . 3 × 10 5 at 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 can 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 on 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 = 1 | s | ∑ | s | i =1 C s → r i . We control the value of λ 1 to obtain different sparsities of learned causal structure. The average accumulative reward and the sparsity of causal structure during the training process in Swimmer are presented in Table 1. With the increase of λ 1 (like from 5 e -6 to 5 e -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 of the agent, leading to misguided policy learning. By contrast, if we set a relatively low λ 1 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.

𝑠

!,#

to

𝑠

$,#%&

𝑎

!,#

to

𝑠

$,#%&

𝑠

!,#

to

𝑟

#

from

𝑎

!,#

to

𝑟

#

𝑡 = 1𝑒4

0.0

0.2

0.4

0.6

0.8

1.0

𝑠

!,#

(

𝑖 ∈

1,27

)

𝑠

$,#%&

(

𝑗 ∈

1,54

)

1            6            11          16          21          26          31          36           41         46           51

𝑡 = 1𝑒6

𝑠

$,#%&

(

𝑗 ∈

1,54

)

𝑠

!,#

(

𝑖 ∈

1,27

)

1            6            11         16           21         26

1            6            11          16          21          26          31          36           41         46           51

1            6            11         16           21         26

𝑡 = 5𝑒5

1              6            11         16           21           26          31           36          41          46           51

𝑎

!,#

(

𝑖 ∈

1,8

)

1    3   5    7

𝑠

$,#%&

(

𝑗 ∈

1,54

)

𝑎

!,#

(

𝑖 ∈

1,8

)

1              6            11         16           21           26          31           36          41          46           51

1    3   5    7

𝑠

$,#%&

(

𝑗 ∈

1,54

)

𝑡 = 1𝑒4

𝑎

!,#

(

𝑖 ∈

1,8

)

1    3   5    7

1              6            11         16           21           26          31           36          41          46           51

𝑠

$,#%&

(

𝑗 ∈

1,54

)

𝑡 = 1𝑒6

𝑟

!

1            6            11

16           21           26

31          36          41          46          51

𝑠

!,#

(

𝑖 ∈

1,54

)

𝑡 = 5𝑒5

𝑟

!

1            6            11

16           21          26

31           36          41

46          51

𝑠

!,#

(

𝑖 ∈

1,54

)

𝑡 = 1𝑒4

𝑟

!

1            6            11

16           21           26          31           36

41          46

51

𝑠

!,#

(

𝑖 ∈

1,54

)

𝑡 = 1𝑒6

1         3         5         7

𝑟

!

𝑎

!,#

(

𝑖 ∈

1,8

)

𝑡 = 5𝑒5

𝑎

!,#

(

𝑖 ∈

1,8

)

𝑟

!

1         3         5         7

𝑡 = 1𝑒4

𝑟

!

1         3         5         7

𝑎

!,#

(

𝑖 ∈

1,8

)

𝑡 = 1𝑒6

口

0

50

100

150

200

250

300

Number of frames

(×10 4

)

0.00

0.25

0.50

0.75

1.00

1.25

1.50

×

10 4

SAC

0

50

100

150

200

250

300

Number of frames

(×10 4

)

0.00

0.25

0.50

0.75

1.00

1.25

1.50

×

10 4

TD3

0

50

100

150

200

250

300

Number of frames

(×10 4

)

0.00

0.25

0.50

0.75

1.00

1.25

1.50

×

10 4

DDPG

Average Accumulative Reward

RRD-Bias

RRD-Unbias

IRCR

None

GRD

Robustness. To evaluate the robustness of GRD, we provide results under noisy states in Ant . Asignificant characteristic of Ant is that only the first 28 dimensions of state are used. Therefore, the noise on the other dimensions ( 28 ∼ 111 ) should not impact a learned robust policy. To verify this, during policy evaluation, we introduce the independent Gaussian noises with the mean of 0 and standard deviation of 0 ∼ 1 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, our method learns a more robust policy with the usage of compact representation.

0.0

0.2

0.4

0.6

0.8

1.0

Standard Deviation of Noise

0

2

4

6

Average Accumulative Reward

1e3

GRD

RRD-Bias

RRD-Unbias

IRCR

Results over different RL backbones. We provide the results of training with DDPG [56] and TD3 [57] in Figure 5. 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.

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 RL backbones (DDPG [56] and TD3 [57]) as shown in Figure 5; 2) On manipulation tasks: we provide the results on manipulation tasks in MetaWorld [58]. Please refer to Appendix D.

6.4 Visualization

We visualize the learned causal structure and redistributed rewards in Ant , where the robot observes 111 -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 the more distinguishable color) and further filters out more elements with the training. At the 1 e 4 time step, GRD regards 75% elements in C s → r ∈ { 0 , 1 } 376 , 70% dimensions in C a → s ∈

Episodel

Episode2

8

8

6

6

redistributedrewards

Reward

4

Reward

4

2

2

0

0

redistributedrewards

0

50

100

150

200

250

300

700

750

800

850

900

950

1000

Time Steps

Time Steps

{ 0 , 1 } 8 × 376 , 98% dimensions in C s → s ∈ { 0 , 1 } 376 × 376 to be 0 with high probabilities, indicating 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 4 (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 t to r t , i.e ., these dimensions do not exist in the identified compact representation. Additionally, according to Figure 4(d), the edges from different dimensions of a to r always exist. It corresponds to the reward design in Ant since the robots are expected to act with the lowest control cost, defined as r t,cost = ∑ | a | i =1 a 2 i . 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 to the long-term return from those with more contributions. More cases are provided in Appendix D.2.

7 Conclusion

In this paper, we 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 a causal view and recovers the generative process within the MDPs with trajectory-wise 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 our method's use in the case of a 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 of algorithms, we aim to harmonize humanAI collaboration to enable more reliable and responsible decision-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.

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 to the support by the EPSRC grant EP/Y003187/1.

References

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.

Jianyu Chen, Bodi Yuan, and Masayoshi Tomizuka. Model-free deep reinforcement learning for urban autonomous driving. In 2019 IEEE Intelligent Transportation Systems Conference (ITSC) , pages 27652771. IEEE, 2019.

Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing systems , 30, 2017.

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.

David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis 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):1140-1144, 2018.

Lei Han, Peng Sun, Yali Du, Jiechao Xiong, Qing Wang, Xinghai Sun, Han Liu, and Tong Zhang. Grid-wise control for multi-agent reinforcement learning in video game ai. In International Conference on Machine Learning (ICML) , pages 2576-2585. PMLR, 2019.

Yali Du, Lei Han, Meng Fang, Tianhong Dai, Ji Liu, and Dacheng Tao. Liir: learning individual intrinsic reward in multi-agent reinforcement learning. In Proceedings of the 33rd International Conference on Neural Information Processing Systems , pages 4403-4414, 2019.

Zihao Zhang, Stefan Zohren, and Stephen Roberts. Deep reinforcement learning for trading. The Journal of Financial Data Science , 2(2):25-40, 2020.

Hongyang Yang, Xiao-Yang Liu, Shan Zhong, and Anwar Walid. Deep reinforcement learning for automated stock trading: An ensemble strategy. In Proceedings of the First ACM International Conference on AI in Finance , pages 1-8, 2020.

Chao Yu, Jiming Liu, Shamim Nemati, and Guosheng Yin. Reinforcement learning in healthcare: A survey. ACM Computing Surveys (CSUR) , 55(1):1-36, 2021.

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 in neural information processing systems , 31, 2018.

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 , 2018.

Meng Fang, Tianyi Zhou, Yali Du, Lei Han, and Zhengyou Zhang. Curriculum-guided hindsight experience replay. Advances in neural information processing systems , 32, 2019.

Beining Han, Zhizhou Ren, Zuofan Wu, Yuan Zhou, and Jian Peng. Off-policy reinforcement learning with delayed rewards. In International Conference on Machine Learning , pages 8280-8303. PMLR, 2022.

Jose A Arjona-Medina, 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.

Vihang P. Patil, Markus Hofmarcher, Marius-Constantin Dinu, Matthias Dorfer, Patrick M. Blies, Johannes Brandstetter, Jose A. Arjona-Medina, and Sepp Hochreiter. Align-rudder: Learning from few demonstrations by reward redistribution. In International Conference on Machine Learning . PMLR, 2022.

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.

Yujing Hu, Weixun Wang, Hangtian Jia, Yixiang Wang, Yingfeng Chen, Jianye Hao, Feng Wu, and Changjie Fan. Learning to utilize shaping rewards: A new approach of reward shaping. Advances in Neural Information Processing Systems , 33:15931-15941, 2020.

Pradyumna Tambwekar, Murtaza Dhuliawala, Lara J Martin, Animesh Mehta, Brent Harrison, and Mark O Riedl. Controllable neural story plot generation via reward shaping. In Proceedings of the Sixteenth International Joint Conference on Artificial Intelligence , pages 5982-5988, 2019.

Sai Rajeswar, Cyril Ibrahim, Nitin Surya, Florian Golemo, David Vazquez, Aaron Courville, and Pedro O Pinheiro. Haptics-based curiosity for sparse-reward tasks. In Conference on Robot Learning , pages 395-405. PMLR, 2022.

Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International Conference on Machine Learning , pages 2778-2787. PMLR, 2017.

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. Advances in Neural Information Processing Systems , 34:3757-3769, 2021.

Cédric Colas, Tristan Karch, Nicolas Lair, Jean-Michel Dussoux, Clément Moulin-Frier, Peter Dominey, and Pierre-Yves Oudeyer. Language as a cognitive tool to imagine goals in curiosity driven exploration. Advances in Neural Information Processing Systems , 33:3761-3774, 2020.

Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAI Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. Advances in Neural Information Processing Systems , 30, 2017.

Ben Eysenbach, Xinyang Geng, Sergey Levine, and Russ R Salakhutdinov. Rewriting history with inverse rl: Hindsight inference for policy improvement. Advances in neural information processing systems , 33:14783-14795, 2020.

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.

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.

Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based counterfactual data augmentation. Advances in Neural Information Processing Systems , 2022.

Carlos Guestrin, Daphne Koller, and Ronald Parr. Multiagent planning with factored mdps. Advances in Neural Information Processing Systems , 14, 2001.

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. PMLR, 2022.

Fan Feng, Biwei Huang, Kun Zhang, and Sara Magliacane. Factored adaptation for non-stationary reinforcement learning. Advances in Neural Information Processing Systems , 2022.

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. PMLR, 2015.

Michael J. 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.

Andrew Y Ng, Daishi Harada, and Stuart J 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.

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 arXiv:1905.13420 , 2019.

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.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems , 30, 2017.

Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlovi´ c, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.02217 , 2020.

Tanmay Gangwani, Yuan Zhou, and Jian Peng. Learning guidance rewards with trajectory-space smoothing. Advances in Neural Information Processing Systems , 33, 2020.

Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for taskindependent state abstraction. International Conference on Machine Learning , pages 23151-23180, 2022.

St John Grimbly, Jonathan Shock, and Arnu Pretorius. Causal multi-agent reinforcement learning: Review and open problems. In Cooperative AI Workshop, Advances in Neural Information Processing Systems , 2021.

Natasha Jaques, Angeliki Lazaridou, Edward Hughes, Caglar Gulcehre, Pedro Ortega, DJ Strouse, Joel Z Leibo, and Nando De Freitas. Social influence as intrinsic motivation for multi-agent deep reinforcement learning. In International Conference on Machine Learning , pages 3040-3049. PMLR, 2019.

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.

Junzhe Zhang, Daniel Kumor, and Elias Bareinboim. Causal imitation learning with unobserved confounders. Advances in Neural Information Processing Systems , 33:12263-12274, 2020.

Thomas Mesnard, Theophane Weber, Fabio Viola, Shantanu Thakoor, Alaa Saade, Anna Harutyunyan, Will Dabney, Thomas S Stepleton, Nicolas Heess, Arthur Guez, et al. Counterfactual credit assignment in model-free reinforcement learning. In International Conference on Machine Learning , pages 7654-7664. PMLR, 2021.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Kevin Patrick Murphy. Dynamic bayesian networks: representation, inference and learning . University of California, Berkeley, 2002.

Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations , 2016.

Douglas A Reynolds et al. Gaussian mixture models. Encyclopedia of biometrics , 741(659-663), 2009.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning , pages 1861-1870. PMLR, 2018.

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.01783, 2016.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR , abs/1707.06347, 2017.

Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in neural information processing systems , 32, 2019.

Ignasi Clavera, Jonas Rothfuss, John Schulman, Yasuhiro Fujita, Tamim Asfour, and Pieter Abbeel. Modelbased reinforcement learning via meta-policy optimization. In Conference on Robot Learning , pages 617-629. PMLR, 2018.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems , pages 5026-5033. IEEE, 2012.

David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning , pages 387-395. Pmlr, 2014.

Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning , pages 1587-1596. PMLR, 2018.

Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning , pages 1094-1100. PMLR, 2020.

Judea Pearl. Causality: Models, reasoning, and inference, 2000.

Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, prediction, and search . MIT press, 2000.

Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning , volume 2. MIT press Cambridge, MA, 2006.

A Background on Causal Inference

A directed acyclic graph (DAG), G = ( V , E ) , can be deployed to represent a graphical criterion carrying out a 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 ⊆ V blocks the path p if and only if (1) p contains a chain i → m → j or a fork i ← m → j such that the middle node m is in Z , or (2) p contains a collider i → m ← j such that the middle node m is not in Z and such that no descendant of m is in Z . Let X , Y and Z be disjunct sets of nodes. If and only if the set Z blocks all paths from one node in X to one node in Y , Z is considered to d-separate X from Y , denoting as ( X ⊥ d Y | Z ) .

Definition 2 (Global Markov Condition [59, 60]) If, for any partition ( X , Y , Z ) , X is d-separated from Y given Z , i.e ., X ⊥ d Y | Z . Then the distribution P over V satisfies the global Markov condition on graph G , and can be factorizes as, P ( X , Y | Z ) = P ( X | Z ) P ( Y | Z ) . That is, X is conditionally independent of Y given Z , writing as X ⊥ ⊥ Y | Z .

Definition 3 (Faithfulness Assumption [59, 60]) The variables, which are not entailed by the Markov Condition, are not independent of each other.

Under the above assumptions, we can apply d-separation as a criterion to understand the conditional independencies from a given DAG G . That is, for any disjoint subset of nodes X , Y , Z ⊆ V , ( X ⊥ ⊥ Y | Z ) and X ⊥ d Y | Z are the necessary and sufficient condition of each other.

B Details of Theoretical Analysis

Proposition 1 (Identifiability) Suppose the state s t , action a t , trajectory-wise long-term return R are observable while Markovian rewards r t are unobservable, and they form an MDP, as described in Eq. 2. Then under the global Markov condition and faithfulness assumption, the reward function g and the Markovian rewards r t are identifiable, as well as the causal structure that is characterized by binary masks C ·→· and C ·→· and the transition dynamics f .

Below is the proof of Proposition 1. We begin by clarifying the assumptions we made and then provide the mathematical proof.

Assumption We assume that, ϵ s,i,t and ϵ r,t in Eq. 2 are i.i.d additive noise. From the weight-space view of Gaussian Process [61], equivalently, the causal models for s i,t +1 and r t can be represented as follows, respectively,

where ∀ i ∈ [1 , | s | ] , and ϕ f and ϕ g denote basis function sets.

Then we denote the variable set in the system by V , with V = { s 1 ,t , . . . , s | s | ,t , a 1 ,t , . . . , a | a | ,t , r t } T t =1 ∪ R , and the variables form a Bayesian network G . Note, we assume that there are possible edges only from s i,t -1 ∈ s t -1 to s i ′ ,t ∈ s t , from a j,t -1 ∈ a t -1 to s i ′ ,t ∈ s t , from s i,t ∈ s t to r t , and from a j,t ∈ a t to r t in G . In particular, the r t are unobserved, while R = ∑ T t =1 γ t -1 o t is observed. Thus, there are deterministic edges from each r t to R .

Below we omit the γ for simplicity.

Proof 1 Given trajectory-wise long-term return R , the binary masks, C s → r , C a → r and Markovian reward function g and the rewards r t are identifiable. Following the above assumption, we first rewrite the function to calculate trajectory-wise long-term return in Eq. 2 as,

For simplicity, we replace the components in Eq. 11 by,

where X := [ s t , a t ] T t =1 representing the concatenation of the covariates s t and a t from t = 1 to T . Consequently, we derive the following equation,

Then we can obtain a closed-form solution of W T g in Eq. 13 by modeling the dependencies between the covariates X τ and response variables R τ , where both are continuous. One classical approach to finding such a solution involves minimizing the quadratic cost and incorporating a weight-decay regularizer to prevent overfitting. Specifically, we define the cost function as,

where τ represents trajectories consisting of 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 g and set the derivative to zero:

Solving this equation will yield the closed-form solution for W T g , i.e .,

Therefore, W g , which indicates the causal structure and strength of the edge, can be identified from the observed data. In summary, given trajectory-wise long-term return R , the binary masks, C s → r , C a → r and Markovian reward function g and the rewards r t are identifiable.

The binary masks, C s → s , C a → s and the transition dynamics f are identifiable In a similar manner, based on the assumption and Eq. 2, we can rewrite Eq. 9 to,

To obtain a closed-form solution for Wi,f T in Equation 17, we can model the dependencies between the covariates X t and the response variables s t +1 , both of which are continuous. The closed-form solution can be represented as:

By taking derivatives of the cost function and setting them to zero, we can obtain the closed-form solution,

Therefore, W i,f can be identified from the observed data. This conclusion applies to 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, C s → s , C a → s and the transition dynamics f are identifiable.

̸

̸

Considering the Markov condition and faithfulness assumption, we can conclude that for any pair of variables V i , V j ∈ V , V i and V j are not adjacent in the causal graph G if and only if they are conditionally independent given some subset of { V l | l = i, l = j } . Additionally, since there are no instantaneous causal relationships and the direction of causality can be determined if an edge exists, the binary structural masks C s → r , C a → r , C s → s , and C a → s defined over the set V are identifiable with conditional independence relationships [30]. Consequently, the functions f and g in Equation 2 are also identifiable.

Algorithm 1 Learning the generative process and policy jointly.

C Details of Implementation

C.1 Baselines

We compare our method against the following baselines,

RRD(biased). This baseline utilizes a surrogate objective called randomized return decomposition loss for reducing the consumption of estimating the Markovian reward function. It applies Monte-Carlo sampling to get a biased estimation of the Mean Square Error (MSE) between the observed episodic reward and the sum of Markovian reward predictions in a sequence. We keep the same setting and hyper-parameters with its official implementation to 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 an unbiased estimation of MSE by sampling short sub-sequences. It offers a computationally efficient approach to optimize MSE. According to [17], RRD (biased) and RRD (unbiased) achieve state-of-the-art performance in episodic MuJoCo tasks.

This baseline performs non-parametric uniform reward redistribution. At each time step, the proxy reward is set to the normalized value of the trajectory return. IRCR is a simple and efficient approach, and except for RRD, it achieves state-of-the-art performance in the literature. The implementation is from RRD [17].

C.2 Detailed Generative Model

The parametric generative model Φ m used in the MDP environment consists of three components: ϕ cau , ϕ rew , and ϕ dyn. We provide a detailed description of their model structures below.

ϕ cau for predicting the causal structure. ϕ cau comprises a set of free parameters without input. We divide ϕ cau into four parts, each corresponding to the binary masks in Equation 2. Specifically, we have

ϕ s → s cau ∈ R | s |×| s |× 2 for C s → s ∈ { 0 , 1 } | s |×| s | ,

ϕ a → s cau ∈ R | a |×| s |× 2 for C a → s ∈ R | a |×| s | ,

ϕ s → r cau ∈ R | s |× 2 for C s → r ∈ R | s | ,

ϕ a → r cau ∈ R | a |× 2 for C a → r ∈ R | a | .

Below we explain the shared workflows in ϕ cau using the example of predicting the causal edge from the i -th dimension of state s i,t to the j -th dimension of the next state s j,t +1 , by part of the free parameters, ϕ s → s cau ,i,j .

For simplicity, we denote ϕ s → s cau ,i,j as ψ . The shape of ψ is now easy to be determined. That is ψ ∈ R 2 and we write it as ψ = [ ψ 0 , ψ 1 ] . With this 2-element vector, we can characterize a Bernoulli distribution, where each element corresponds to the unnormalized probability of classifying the edge as existing ( ψ 0 ) or not existing ( ψ 1 ), respectively. Therefore, the probability of the causal edge existing from the i -th dimension of state s i,t to the j -th dimension of the next state s j,t +1 can be calculated as:

One example of using the estimated mask is given as Figure 8.

During training: obtain C s → s i,j through Gumbel-Softmax sampling. In the training phase, it is crucial to maintain the gradient flow for backpropagation. To achieve this, we sample the binary values of C s → s i,j by applying Gumbel-Softmax [48],

where GS denotes the Gumbel-Softmax sampling, which allows us to 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.

The above explanation of the workflow in ϕ cau for predicting a single causal edge provides insight into the overall implementation of the entire module ϕ cau and can be applicable for predicting all the causal edges during the training phase. Therefore, we can obtain C a → s for optimizing ϕ dyn, C s → r and C a → r for optimizing ϕ rew , using similar procedures.

During inference: obtain C s → s i,j by greedy selection. In the inference phase, including data sampling and policy learning, we get the prediction of C s → s i,j through a greedy selection,

Such a greedy sampling avoids introducing randomness using the Gumble-Softmax sampling.

The above explanation of the workflow in ϕ cau for predicting a single causal edge provides insight into the overall implementation of the entire module ϕ cau and can be applicable for predicting all the causal edges during the inference phase. Therefore, we can obtain C s → s , C s → r and C a → r using similar procedures to predict the Markovian reward and extract compact representation s min .

ϕ rew for predicting the Markovian rewards. ϕ rew is a stacked, fully-connected network, and the details of the network structure are provided in Table 2. During training, the prediction of Markovian reward can be

written as,

⊙

𝒔௧

𝒂௧

𝑪ȉ,ଵ

𝑪ȉ,ଶ

𝑪ȉ,ଷ

𝑪ȉ,ଵ

𝒂→𝒔

𝑪ȉ,ଶ

𝒂→𝒔

𝑪ȉ,ଷ

𝒂→𝒔

𝑪 𝒔→𝒔

𝑪 𝒂→𝒔

𝒔ଵ,௧ାଵ

𝒔ଶ,௧ାଵ

𝒔ଷ,௧ାଵ

Learnable Parameters

𝜙 dyn

⊙

⊙

⊙

⊙

⊙

𝒔௧

𝒔௧

𝒂௧

𝒂௧

: 0

: 1

Element-wise Product

⊙

where [ · , · ] , ⊙ denotes concatenation and element-wise multiply operations, respectively. FCs denotes the stacked fully-connected network. C s → r and C s → r are derived from ϕ cau by Gumbel-Softmax.

During inference, including policy learning and data sampling, the predicted Markovian reward is

where C s → r and C s → r are derived from ϕ cau greedily by deterministic sampling.

ϕ dyn for modeling the environment dynamics. In our experiment, we do not directly utilize ϕ dyn in policy learning. Instead, this module serves as a bridge to optimize ϕ s → s cau and ϕ a → s cau . Subsequently, ϕ s → s cau can be utilized in the calculation of C s → π .

During training, we initially sample C s → s and C a → s using Gumbel-Softmax. The prediction for the i -th dimension of the next state can be represented as follows,

where [ · , · ] denotes concatenation and MDN denotes the Mixture Density Network, which outputs the means, variances, and probabilities for N Gau Gaussian cores. The parameters of MDN are shared across the predictions of different dimensions of the next state. We set N cau = 3 in our experiments. More details about ϕ dyn can be found in Table 2.

C.3 Detailed Policy Model

Considering the specific requirements of the employed RL algorithm, Soft Actor-Critic (SAC), our Policy Model Φ π comprises two components, the actor ϕ π and the critic ϕ v . Detailed network structures for both components can be found in Table 2.

C.4 Training Process

We follow the line of joint learning in [17], which avoids learning a return decomposition model in advance using 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 :

D 1 = {⟨ s t m , a t m ⟩ | T t =1 , R m } | M m =1 consists of M trajectories. Provided with the trajectory-wise long-term returns R m | M m =1 , D 1 is utilized to optimize ϕ s → r cau , ϕ a → r cau and ϕ rew, with L rew .

D 2 = { s t n , a t n , s t n +1 } | N n =1 consists of N state-action pairs. D 2 are used for policy optimization and optimize the parts of causal structure, ϕ s → s cau and ϕ a → s cau , ϕ dyn. With such a D 2 , GRD breaks the temporal cues in the training data to learn the policy and dynamics function.

Please refer to Algorithm 1 for a detailed training process.

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 3 × 10 -4 . The hyperparameters for policy learning are shared across all tasks, with a discount factor of 1 . 00 and a Polyak-averaging coefficient of 5 × 10 -4 . The target entropy is set to the negative value of the dimension of the robot action. To facilitate training, we utilize a replay buffer with a size of 1 × 10 6 time steps. The warmup size of the buffer for training is set to 1 × 10 4 . The model is trained for 3 epochs, with each epoch consisting of 100 training cycles. In each 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 MuJoCo simulation, which is then stored in the replay buffer. For training the ϕ rew, we sample 4 episodes, each containing 5 × 10 3 steps. As for policy learning and the optimization of ϕ dyn, we use data from 256 time steps. ϕ cau is trained together with ϕ rew and ϕ dyn. 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 3. All experiments were conducted on an HPC system equipped with 128 Intel Xeon processors operating at a clock speed of 2.2 GHz and 5 terabytes of memory.

D Additional Results

D.1 Results over manipulation tasks

We provide the comparison of our method with RRD in the three tasks of MetaWorld, as shown in Figure 9. The evaluation metric is the success rate. 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 a wall and push a puck to a 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.

0

25

50

75

100

125

150

Number of frames

(×10 4

)

0.0

0.2

0.4

0.6

0.8

1.0

Push Wall

0

50

100

150

200

Number of frames

(×10 4

)

0.0

0.2

0.4

0.6

0.8

1.0

Door Lock

0

50

100

150

200

Number of frames

(×10 4

)

0.0

0.2

0.4

0.6

0.8

1.0

Pick Place

Success Rate

RRD-Bias

RRD-Unbias

GRD

D.3 Visualization of Learned Causal Structure

We provide the visualization of learned causal structure in Swimmer , as shown in 11. Since the grounded causal structure is not accessible, we verify the reasonability of the learned causal structure by some observations:

All the edges from different dimensions of a to r always exist, as shown in Figure 11 (d): Swimmer shares the same characteristic with Ant 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.

According to Figure 11 (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 to first rotor. We can get a similar observation for the second action dimension from Figure 11.

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.

Episode 1

Episode 2

8

8

6

Reward

4

Reward

4

2

01

.....

ground truth

2

redistributedrewards

-2

redistributedrewards

0

50

100

150

200

250

300

350

400

0

50

100

150

200

250

300

350

400

Time Steps

Time Steps

Episode3

Episode 4

Reward

4

Reward

ground truth

-2

redistributedrewards

-2

redistributedrewards

0

50

100

150200

Time Steps

250

300

350

400

0

50

100

150200

Time Steps

250

300350

400

Episode5

Episode 6

8

6

Reward

Reward

of

0

.....

ground truth

-2

redistributedrewards

2

redistributedrewards

600

650

700

750

800

850

900

950

1000

600

650

700

750

800

850

900

950

1000

Time Steps

Time Steps

Episode7

Episode8

8

8

6

Reward

4

Reward

2

0

2

redistributedrewards

-2

redistributedrewards

600

650

700

750800

850

900

950

1000

600

650

700

750

800

850

900950

1000

Time Steps

Time Steps

𝑠 ௜,௧

௝,௧ାଵ

௧

௜,௧

௧

௜,௧

௜,௧

to

௝,௧ାଵ

௝,௧ାଵ

𝑎௜,௧

௜,௧

to

௝,௧ାଵ

௜,௧

to

௧

from

௜,௧

to

௧

0.0

0.2

0.4

0.6

0.8

1.0

0

1

1

2

0

12345

67

3

4

5

6

01234567

0

1234567

0

1

0

1

2

3

4

5

6

7

Peer reviews:
comments: 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.

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.

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?

2 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.

3 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.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\nSee Weaknesses.\n\nLIMITATIONS:\nSee Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

comments: 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.

STRENGTHS:
* Experiment results seem promising.

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. 
* 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.

Minor
* Equations are too messy and hard to understand. Use under bracket in equation 6. 
* 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.
* 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.

QUESTIONS:
* I could not understand the experiment results in l345-352.

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1 - 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\nWEAKNESSES:\n1 - 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\n\nQUESTIONS:\n1 - 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.\n\nLIMITATIONS:\nYes, the limitations have been addressed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-23"}