SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
L _ {\mathrm {d y n}} \left(\phi_ {\mathrm {d y n}}, \phi_ {\mathrm {c a u}} ^ {s \rightarrow s}, \phi_ {\mathrm {c a u}} ^ {a \rightarrow s}\right) = \mathbb {E} _ {\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}, \boldsymbol {s} _ {t + 1} \sim \mathcal {D}} \left[ - \sum_ {i = 1} ^ {| s |} \log P \left(\boldsymbol {s} _ {i, t + 1} \mid \boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}, \boldsymbol {C} _ {\cdot , i} ^ {s \rightarrow s}, \boldsymbol {C} _ {\cdot , i} ^ {a \rightarrow s}; \phi_ {\mathrm {d y n}}\right)\right]. \tag {5}
$$
Additionally, we minimize the following cross-entropy terms to regulate the sparsity of learned causal structure to avoid trial solutions. It is achieved by force to optimize the parameters towards the direction of the nonexistence of the causal edge. Let  $D_{i}(\pmb{x}) = \log P(\pmb{x}_{i})$ , where  $P(\pmb{x}_i)$  is the possibility that the edge  $\pmb{x}_i$  exists. The regularizer term is,
$$
\begin{array}{l} L _ {\text {r e g}} \left(\phi_ {\text {c a u}}\right) = \underbrace {\lambda_ {1} \sum_ {i} D _ {i} \left(\boldsymbol {C} ^ {\boldsymbol {s} \rightarrow r}\right)} + \underbrace {\lambda_ {2} \sum_ {i} D _ {i} \left(\boldsymbol {C} ^ {\boldsymbol {a} \rightarrow r}\right)} + \\ \underbrace {\lambda_ {3} \sum_ {j \neq i} ^ {\text {s t a t e - t o - r e w a r d}} D _ {i , j} \left(\boldsymbol {C} ^ {\boldsymbol {s} \rightarrow \boldsymbol {s}}\right)} _ {\text {t e - t o - s t a t e (e x c l u d i n g s e l f - c o n n e c t i o n s)}} + \underbrace {\lambda_ {4} \sum_ {j = i} ^ {\text {a c t i o n - t o - r e w a r d}}} _ {\text {s t a t e - t o - s t a t e (s e l f - c o n n e c t i o n s)}} + \underbrace {\lambda_ {5} \sum_ {i , j} D _ {i , j} \left(\boldsymbol {C} ^ {\boldsymbol {s} \rightarrow \boldsymbol {s}}\right)} _ {\text {a c t i o n - t o - s t a t e}}. \tag {6} \\ \end{array}
$$
where  $\mathbb{1}(\cdot)$  is the indicator function and hyper-parameters  $\lambda_{(\cdot)}$  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^{\mathrm{min}}$ , for policy learning as the input of the policy model,  $\Phi_{\pi}$ , 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} \in s_t$ , which either 1) has an edge to the reward  $r_t$  i.e.,  $C_i^{s \to r} = 1$ , or 2) has an edge to another state component in the next time-step,  $s_{j,t+1} \in s_t$ , i.e.,  $C_{i,j}^{s \to s} = 1$ , such that the same component at time  $t$  is in the compact representation, i.e.,  $s_{j,t} \in s_t^{\mathrm{min}}$ . We first select  $C^{s \to s}$  and  $C^{s \to r}$  greedily from the i.i.d Bernoulli distributions characterized by  $\phi_{\mathrm{cau}}^{s \to s}$  and  $\phi_{\mathrm{cau}}^{s \to r}$  (as illustrated in Appendix C.2) and organize the state dimensions which exist in  $s_t^{\mathrm{min}}$  as  $C^{s \to \pi}$ . We define  $\mathbb{C} := \{i, \forall C_i^{s \to r} = 1\}$  to denote the dimensions of  $s_t$  which have impact on the Markovian reward  $r_t$  directly. Then we have,
$$
C ^ {s \rightarrow \pi} = \left\{C ^ {s \rightarrow r} \right\} \vee \left\{C _ {\cdot , \mathbb {C} _ {1}} ^ {s \rightarrow s} \vee \dots \vee C _ {\cdot , \mathbb {C} _ {| C |}} ^ {s \rightarrow s} \right\}, \tag {7}
$$
where  $|C|$  denotes the size of  $\mathbb{C}$  and  $\vee$  denotes element-wise logical or operation. Then, the compact representation which serves as the input of policy is defined as,  $s_t^{\min} = C^{s \to \pi} \odot 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  $\hat{r}_t = \phi_{\mathrm{rew}}(s_t, a_t, C^{s \to r}, C^{a \to r})$ , where  $C^{s \to r}$  and  $C^{a \to r}$  are selected greedily from  $\phi_{\mathrm{cau}}^{s \to r}$  and  $\phi_{\mathrm{cau}}^{a \to r}$ , as illustrated in Appendix C.2. The policy model  $\Phi_\pi$  contains two parts, a critic  $\phi_v$
and an actor  $\phi_{\pi}$  due to the applied SAC algorithm. SAC alternates between a policy evaluation step, which trains a critic  $\phi_v(s_t^{\min}, \pmb{a}_t)$  to estimate  $Q^{\pi}(s_t^{\min}, \pmb{a}_t) = \mathbb{E}_{\pi}[\sum_{t=1}^{T} \gamma^{t-1} \hat{r}(s_t^{\min}, \pmb{a}_t) \mid s_1 = \pmb{s}, \pmb{a}_1 = \pmb{a}]$  using the Bellman backup operator, and a policy improvement step, which trains an actor  $\phi_{\pi}(s_t^{\min})$  by minimizing the expected KL-divergence,
$$
J _ {\pi} \left(\Phi_ {\pi}\right) = \mathbb {E} _ {s _ {t} \sim \mathcal {D}} \left[ D _ {\mathrm {K L}} \left(\Phi_ {\pi} \| \exp \left(Q ^ {\pi} - V ^ {\pi}\right)\right) \right], \tag {8}
$$
where  $V^{\pi}$  is the value function of state-action pairs [50]. More detailed implementation of  $\Phi_{\pi}$  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.
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.
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, Humanoid-Standup, and Reacher tasks. All robots are limited to observing only one episode reward at the end of each episode, which is equivalent to the accumulative Markov reward that describes their performance throughout the episode. The maximum episode length is shared among all tasks and set to be 1000. In these tasks, the contributions across dimensions in the state to the Markovian reward can differ significantly due 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.
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 \times 10^{4}$  and  $6.5 \times 10^{3}$ , separately, while the baseline methods achieve the best value of  $1.25 \times 10^{4}$  and  $6 \times 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
Figure 3: Learning curves on a suite of MuJoCo benchmark tasks with episodic rewards, based on 5 independent runs with random initialization. The shaded region indicates the standard deviation and the curves were smoothed by averaging the 10 most recent evaluation points using an exponential moving average. An evaluation point was established every  $10^{4}$  time steps.
Table 1: The mean and the standard variance of average accumulative reward and sparsity rate  $S$  regarding diverse  ${\lambda }_{1}$  at different time step  $t$  in Swimmer.
λ1 / t	1e5	2e5	3e5	5e5	8e5	1e6
0	87 ± 41(0.77)	153± 43(0.80)	151± 32(0.80)	131± 55(0.80)	170± 36(0.77)	159± 25(0.78)
3e-6	182±103(0.72)	201±109(0.71)	217±108(0.72)	217±100(0.68)	229± 95(0.52)	220±102(0.50)
5e-6	130±103(0.56)	245±104(0.69)	261±100(0.74)	286± 77(0.77)	272± 78(0.75)	262± 94(0.66)
5e-5	118±100(0.46)	109±107(0.41)	123±103(0.34)	141± 96(0.25)	152± 93(0.19)	158± 90(0.17)
111-dimension states, GRD gains significant improvement. Taking Humanoid Standup as an example, GRD achieves a higher average accumulative reward of  $2.3 \times 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.
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 = \frac{1}{|s|}\sum_{i=1}^{|s|}C_i^{s\rightarrow r}$ . We control the value of  $\lambda_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  $\lambda_1$  (like from  $5e - 6$  to  $5e - 5$ ), the causal structure generally gets more sparse (the sparsity  $S$  decreases), leading to less policy improvement. The most possible explanation is, GRD does not take enough states to predict Markovian rewards and thus can not reflect the true goal of the agent, leading to misguided policy learning. By contrast, if we set a relatively low  $\lambda_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.
(a) Causal Structure from  $s_{i,t}$  to  $s_{j,t + 1}$
(b) Causal Structure from  $a_{i,t}$  to  $s_{j,t + 1}$
(c) Causal Structure from  $s_{i,t}$  to  $r_t$
(d) Causal Structure from  $a_{i,t}$  to  $r_t$
Figure 4: The visualization of learned causal structure for Ant when  $t \in [1e4, 5e5, 1e6]$ . The color indicates the probability of the existence of causal edges, whereas darker colors represent higher probabilities. There are 111 dimensions in the state variables, but only the first 27 ones are used. (a) The learned causal structure among the first 27 dimensions of the state variable  $s_t$  to the first 54 dimensions of the next state variable  $s_{t+1}$ . Due to the limited space, we only visualize the structure at  $t = 1e4$  and  $t = 1e6$ . (b) The learned causal structure among all dimensions of the action variable  $a_t$  to the first 54 dimensions of the next state variable  $s_{t+1}$ . (c) The learned causal structure among the first 54 dimensions of the state variable  $s_t$  to the Markovian reward variable  $r_t$ . (d) The learned causal structure among all dimensions of action variable  $a_t$  to the Markovian reward variable  $r_t$ .
Figure 5: Evaluation with different RL backbones, SAC, DDPG and TD3. "None" is training with the observed delayed rewards.
Robustness. To evaluate the robustness of GRD, we provide results under noisy states in Ant. A significant characteristic of Ant is that only the first 28 dimensions of state are used. Therefore, the noise on the other dimensions  $(28\sim 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\sim 1$  into those insignificant dimensions  $(28\sim 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.
Figure 6: Evaluation with Gaussian Noise in the State.
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  $1e4$  time step, GRD regards  $75\%$  elements in  $C^{s \rightarrow r} \in \{0,1\}^{376}$ ,  $70\%$  dimensions in  $C^{a \rightarrow s} \in$
Figure 7: The visualization of decomposed rewards (blue solid lines) and the grounded rewards (red dotted lines).
$\{0,1\}^{8 \times 376}$ ,  $98\%$  dimensions in  $C^{s \rightarrow s} \in \{0,1\}^{376 \times 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 \sim 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 \sim 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  $\mathbf{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} = \sum_{i=1}^{|a|} a_i^2$ . 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.
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 human-AI 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.
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.
[56] 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.
[57] 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.
[58] 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.
[59] Judea Pearl. Causality: Models, reasoning, and inference, 2000.
[60] Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, prediction, and search. MIT press, 2000.
[61] Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning, volume 2. MIT press Cambridge, MA, 2006.
A Background on Causal Inference
A directed acyclic graph (DAG),  $\mathcal{G} = (\mathbf{V},\mathbf{E})$ , can be deployed to represent a graphical criterion carrying out a set of conditions on the paths, where  $\mathbf{V}$  and  $\mathbf{E}$  denote the set of nodes and the set of directed edges, separately.
Definition 1 (d-separation [59]) A set of nodes  $Z \subseteq V$  blocks the path  $p$  if and only if (1)  $p$  contains a chain  $i \to m \to j$  or a fork  $i \gets m \to j$  such that the middle node  $m$  is in  $Z$ , or (2)  $p$  contains a collider  $i \to m \gets 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 disjoint 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 \perp_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 \perp_d Y \mid Z$ . Then the distribution  $P$  over  $V$  satisfies the global Markov condition on graph  $G$ , and can be factorizes as,  $P(X, Y \mid Z) = P(X \mid Z)P(Y \mid Z)$ . That is,  $X$  is conditionally independent of  $Y$  given  $Z$ , writing as  $X \perp \perp Y \mid 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  $\mathcal{G}$ . That is, for any disjoint subset of nodes  $X, Y, Z \subseteq V$ ,  $(X \perp Y \mid Z)$  and  $X \perp_d Y \mid Z$  are the necessary and sufficient condition of each other.
B Details of Theoretical Analysis
Proposition 1 (Identifiability) Suppose the state  $\mathbf{s}_t$ , action  $\mathbf{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  $C^{\rightarrow}$  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,  $\epsilon_{s,i,t}$  and  $\epsilon_{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,
$$
\boldsymbol {s} _ {i, t + 1} = f _ {i} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \epsilon_ {s, i, t} = W _ {i, f} ^ {T} \phi_ {f} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \epsilon_ {s, i, t}, \tag {9}
$$
$$
r _ {t} = g \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \epsilon_ {r, t} = W _ {g} ^ {T} \phi_ {g} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \epsilon_ {r, t}, \tag {10}
$$
where  $\forall i\in [1,|\pmb {s}|]$  , and  $\phi_f$  and  $\phi_g$  denote basis function sets.