SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yudi Zhang $^{1}$  Yali Du $^{2}$  Biwei Huang $^{3}$  Ziyan Wang $^{2}$  Jun Wang $^{4}$  Meng Fang $^{5,1}$  Mykola Pechenizkiy $^{1}$
$^{1}$ Eindhoven University of Technology  $^{2}$ King's College London
<sup>3</sup>University of California San Diego <sup>4</sup>University College London
$^{5}$ University of Liverpool
{y.zhang5,m.pechenizkiy}@tue.nl,{yali.du,ziyan.wang}@kcl.ac.uk
bih007@ucsd.edu, jun.wang@cs.ucl.ac.uk, Meng.Fang@liverpool.ac.uk
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/.
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, GRD utilizes 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
Figure 1: A graphical illustration of causal structure in Generative Return Decomposition. See the main text for the interpretation.
environment, 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_{\cdot,t}$  and action  $a_{\cdot,t}$ , Markovian rewards  $r_t$  for  $t \in [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 \neq 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.
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.
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  $\langle S, \mathcal{A}, \mathcal{R}, \gamma, P \rangle$ , where  $\mathcal{S}$ ,  $\mathcal{A}$ ,  $\mathcal{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} \mid s_t, a_t)$ . The goal of reinforcement learning is to find an optimal policy  $\pi: \mathcal{S} \to \mathcal{A}$  to maximize the expected long-term return, i.e., a discounted sum of the rewards with the predefined discounted factor  $\gamma$  and episode length  $T$ ,  $J(\pi) = \mathbb{E}_{s_t \sim P(s_t | s_{t-1}, a_{t-1}), a_t \sim \pi(a_t | s_t)}[\sum_{t=1}^T \gamma^{t-1} \mathcal{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} \neq 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  $\hat{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,
$$
R = \sum_ {t = 1} ^ {T} \gamma^ {t - 1} o _ {t} = \sum_ {t = 1} ^ {T} \gamma^ {t - 1} \hat {r} _ {t}, \tag {1}
$$
where  $R$  is the long-term return and  $o_{t}, \hat{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,  $\mathcal{P} = \langle \mathcal{S},\mathcal{A},g,P_f,\gamma ,\mathcal{G}\rangle$  .  $\mathcal{S},\mathcal{A}$  represent the sets of states  $s$  and actions  $\pmb{a}$ , respectively.  $\mathcal{G}$  denotes a DBN that describes the causal relationship within the MDP environment, constructed over a finite number of random variables,  $\{s_{1,t},\dots ,s_{|s|,t},a_{1,t},\dots ,a_{|a|,t},r_t\}_{t = 1}^T\cup R$  where  $|\boldsymbol {s}|$  and  $|\pmb {a}|$  are the dimension of  $s$  and  $\pmb{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}\mid s_{t},a_{t}) = \prod_{i = 1}^{|\boldsymbol{s}|}P(s_{i,t + 1}\mid \mathrm{pa}(s_{i,t + 1}))$  , where  $i$  is the index of dimension of  $s_t$ . Here  $\mathrm{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  $\pmb{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(\mathrm{pa}(r_t))$  , where  $\mathrm{pa}(r_t)$  is a subset of the dimensions of  $s_t$  and  $\pmb{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  $\mathcal{G}$  is given in Figure 1. For simplicity, we reformulate the generative process as,
$$
\left\{\begin{array}{l}s _ {i, t + 1} = f \left(C _ {., i} ^ {s \rightarrow s} \odot \boldsymbol {s} _ {t}, C _ {., i} ^ {\boldsymbol {a} \rightarrow \boldsymbol {s}} \odot \boldsymbol {a} _ {t}, \epsilon_ {s, i, t}\right)\\r _ {t} = g \left(\boldsymbol {C} ^ {\boldsymbol {s} \rightarrow r} \odot \boldsymbol {s} _ {t}, \boldsymbol {C} ^ {\boldsymbol {a} \rightarrow r} \odot \boldsymbol {a} _ {t}, \epsilon_ {r, t}\right)\\R = \sum_ {t = 1} ^ {T} \gamma^ {t - 1} r _ {t}\end{array}\right. \tag {2}
$$
for  $i = (1,2,\dots ,|s|)$ .  $\odot$  is the element-wise product.  $f$  and  $g$  stand for the dynamics function and reward function, separately.  $C^{\rightarrow}$  are categorical masks organizing the causal structure in  $\mathcal{G}$ .  $\epsilon_{s,i,t}$  and  $\epsilon_{r,t}$  are i.i.d random noises. Specially,  $C^{s\rightarrow r}\in \{0,1\}^{|s|}$  and  $C^{a\rightarrow r}\in \{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_i^{s\rightarrow r} = 1$ . Similarly,  $C^{s\rightarrow s}\in \{0,1\}^{|s|\times |s|}$  and  $C^{a\rightarrow s}\in \{0,1\}^{|a|\times |s|}$  indicate the causal relationship between  $s_t$ ,  $a_{t}$  and  $s_{t + 1}$ , separately. In particular, we assume that  $\mathcal{G}$  is time-invariant, i.e.,  $f,g,C^{\rightarrow}$  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^{\rightarrow}$  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^{\rightarrow}$  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.
Figure 2: The framework of the proposed Generative Return Decomposition (GRD).  $\phi_{\mathrm{cau}}$ ,  $\phi_{\mathrm{rew}}$ ,  $\phi_{\mathrm{dyn}}$  in generative model  $\Phi_{\mathrm{m}}$  are marked as yellow, blue and green, while policy model  $\Phi_{\pi}$  is marked as orange. The observable variables, state  $s_t$ , action  $a_t$ , and the delayed reward  $o_t$ , are marked as gray. The mediate results, binary masks,  $C^{\rightarrow}$ , outputs of policy, the predicted Markovian rewards  $\hat{r}_t$  and the compact representation  $s_t^{\mathrm{min}}$  are denoted as purple squares. The policy model  $\Phi_{\pi}$  takes as input  $s_t^{\mathrm{min}}$  and its training is supervised by the predicted Markovian reward  $\hat{r}_t$ . The dotted lines represent the supervision signals, i.e., losses and predicted Markovian rewards.
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,  $\Phi_{\mathrm{m}}$  for the parameterized generative process and  $\Phi_{\pi}$  for the policy. Therefore, the corresponding loss function includes  $L_{m}(\Phi_{m})$  and  $J_{\pi}(\Phi_{\pi})$ , for causal generative model estimation and policy optimization, respectively. Hence, the overall objective can be formulated as
$$
\min  _ {\Phi_ {m}, \Phi_ {\pi}} L \left(\Phi_ {m}, \Phi_ {\pi}\right) = L _ {m} \left(\Phi_ {m}\right) + J _ {\pi} \left(\Phi_ {\pi}\right). \tag {3}
$$
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^{\rightarrow})$ , as well as the estimation of unknown functions  $(f$  and  $g)$ . An overview of our method is illustrated in Figure 2. The parameterized model  $\Phi_{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_{\mathrm{m}}$  over the replay buffer  $\mathcal{D}$ , which consists of trajectories, with each trajectory  $\tau = \{\langle \boldsymbol{s}_t,\boldsymbol{a}_t\rangle |_{t = 1}^T,R\}$ , where  $R$  is the discounted sum of observed delayed rewards  $o_t$ , as given in Eq. 1.
Generative Model. The parameterized model  $\Phi_{\mathrm{m}}$  consists of three parts,  $\phi_{\mathrm{cau}}$ ,  $\phi_{\mathrm{dyn}}$ ,  $\phi_{\mathrm{rew}}$ , which are described below. (1)  $\phi_{\mathrm{cau}}$  is used to identify the causal structure in  $\mathcal{G}$  by predicting the values of binary masks in Eq. 2. It consists of four parts,  $\phi_{\mathrm{cau}}^{s \rightarrow s} \in \mathbb{R}^{|s| \times |s| \times 2}, \phi_{\mathrm{cau}}^{a \rightarrow s} \in \mathbb{R}^{|a| \times |s| \times 2}, \phi_{\mathrm{cau}}^{s \rightarrow r} \in \mathbb{R}^{|s| \times 2}$  and  $\phi_{\mathrm{cau}}^{a \rightarrow r} \in \mathbb{R}^{|a| \times 2}$ , which are used to predict  $C^{s \rightarrow s}$ ,  $C^{a \rightarrow s}$ ,  $C^{s \rightarrow r}$  and  $C^{a \rightarrow r}$ , separately. We take the prediction of  $C^{s \rightarrow s}$  with  $\phi_{\mathrm{cau}}^{s \rightarrow s}$  as an example to explain the process of predicting causal structure.  $\phi_{\mathrm{cau}}^{s \rightarrow s}$  characterizes  $|s| \times |s|$  i.i.d Bernoulli distributions for the existence of the edges organized in the matrix,  $C^{s \rightarrow 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_{i,j}^{s \rightarrow s})$ . Binary prediction of  $C^{s \rightarrow s}$  is sampled by applying element-wise gumbel-softmax [48] during training, while using greedy selection during
inference. Similarly, we can obtain  $P(C_{i,j}^{a\rightarrow s}), P(C_i^{s\rightarrow r}), P(C_i^{a\rightarrow r}), C^{a\rightarrow s}, C^{s\rightarrow r}$  and  $C^{a\rightarrow r}$ . (2)  $\phi_{\mathrm{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\rightarrow r}$  and  $C^{a\rightarrow r}$  to obtain the prediction of Markovian rewards. (3)  $\phi_{\mathrm{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\rightarrow s}$  and  $C^{a\rightarrow s}$ . As an example, the predicted distribution for the  $i$ -th dimension of the next state is  $P(s_{i,t + 1}\mid s_t,a_t,C_{.,i}^{s\rightarrow s},C_{.,i}^{a\rightarrow s};\phi_{\mathrm{dyn}})$ . More details for  $\Phi_{\mathrm{m}}$  can be found in the Appendix C.2.
Loss Terms. Accordingly, the loss term  $L_{\mathrm{m}}$  contains three components with  $L_{\mathrm{m}}(\Phi_{\mathrm{m}}) = L_{\mathrm{rew}} + L_{\mathrm{dyn}} + L_{\mathrm{reg}}$ . Considering the long-term return  $R$  as the causal effect of the Markovian rewards  $r_t$ , we optimize  $\phi_{\mathrm{rew}}, \phi_{\mathrm{cau}}^{s \to r}$  and  $\phi_{\mathrm{cau}}^{a \to r}$  by
$$
L _ {\mathrm {r e w}} \left(\phi_ {\mathrm {r e w}}, \phi_ {\mathrm {c a u}} ^ {s \rightarrow r}, \phi_ {\mathrm {c a u}} ^ {a \rightarrow r}\right) = \mathbb {E} _ {\tau \sim \mathcal {D}} [ \| R - \sum_ {t = 1} ^ {T} \gamma^ {t - 1} \hat {r} _ {t} \| ^ {2} ] = \mathbb {E} _ {\tau \sim \mathcal {D}} [ \| \sum_ {t = 1} ^ {T} \gamma^ {t - 1} o _ {t} - \sum_ {t = 1} ^ {T} \gamma^ {t - 1} \hat {r} _ {t} \| ^ {2} ], \tag {4}
$$
where  $\hat{r}_t$  is predicted by  $\phi_{\mathrm{rew}}$ , i.e.,  $\hat{r}_t = \phi_{\mathrm{rew}}(s_t, a_t, C^{s \rightarrow r}, C^{a \rightarrow r})$ . To optimize  $\phi_{\mathrm{dyn}}$ ,  $\phi_{\mathrm{cau}}^{s \rightarrow s}$  and  $\phi_{\mathrm{cau}}^{a \rightarrow s}$ , we minimize