SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Then we denote the variable set in the system by  $V$ , with  $\boldsymbol{V} = \{\pmb{s}_{1,t},\dots,\pmb{s}_{|\pmb{s}|,t},\pmb{a}_{1,t},\dots,\pmb{a}_{|\pmb{a}|,t},r_t\}_{t=1}^T \cup R$  and the variables form a Bayesian network  $\mathcal{G}$ . Note, we assume that there are possible edges only from  $\pmb{s}_{i,t-1} \in \pmb{s}_{t-1}$  to  $\pmb{s}_{i',t} \in \pmb{s}_t$ , from  $\pmb{a}_{j,t-1} \in \pmb{a}_{t-1}$  to  $\pmb{s}_{i',t} \in \pmb{s}_t$ , from  $\pmb{s}_{i,t} \in \pmb{s}_t$  to  $r_t$ , and from  $\pmb{a}_{j,t} \in \pmb{a}_t$  to  $r_t$  in  $\mathcal{G}$ . In particular, the  $r_t$  are unobserved, while  $R = \sum_{t=1}^{T} \gamma^{t-1} o_t$  is observed. Thus, there are deterministic edges from each  $r_t$  to  $R$ .
Below we omit the  $\gamma$  for simplicity.
Proof 1 Given trajectory-wise long-term return  $R$ , the binary masks,  $C^{s \rightarrow r}$ ,  $C^{a \rightarrow 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,
$$
\begin{array}{l} R = \sum_ {t = 1} ^ {T} r _ {t} \\ = \sum_ {t = 1} ^ {T} \left[ W _ {g} ^ {T} \phi_ {g} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \epsilon_ {r, t} \right] \tag {11} \\ = W _ {g} ^ {T} \sum_ {t = 1} ^ {T} \phi_ {g} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \sum_ {t = 1} ^ {T} \epsilon_ {r, t}. \\ \end{array}
$$
For simplicity, we replace the components in Eq. 11 by,
$$
\begin{array}{r l} \zeta_ {g} (X) & = \sum_ {t = 1} ^ {T} \phi_ {g} (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}), \\ E _ {r} & = \sum_ {t = 1} ^ {T} \epsilon_ {r, t}, \end{array} \tag {12}
$$
where  $X \coloneqq [\pmb{s}_t, \pmb{a}_t]_{t=1}^T$  representing the concatenation of the covariates  $\pmb{s}_t$  and  $\pmb{a}_t$  from  $t = 1$  to  $T$ . Consequently, we derive the following equation,
$$
R = W _ {g} ^ {T} \zeta_ {g} (X) + E _ {r}. \tag {13}
$$
Then we can obtain a closed-form solution of  $W_{g}^{T}$  in Eq. 13 by modeling the dependencies between the covariates  $X_{\tau}$  and response variables  $R_{\tau}$ , 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,
$$
C \left(W _ {g}\right) = \frac {1}{2} \sum_ {X _ {\tau}, R _ {\tau} \sim \mathcal {D}} \left(R _ {\tau} - W _ {g} ^ {T} \zeta_ {g} \left(X _ {\tau}\right)\right) ^ {2} + \frac {1}{2} \lambda \| W _ {g} \| ^ {2}. \tag {14}
$$
where  $\tau$  represents trajectories consisting of state-action pairs  $X_{\tau}$  and long-term returns  $R_{\tau}$ , which are sampled from the replay buffer  $\mathcal{D}$ .  $\lambda$  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:
$$
\frac {\partial C \left(W _ {g}\right)}{\partial W _ {g}} = 0. \tag {15}
$$
Solving this equation will yield the closed-form solution for  $W_{g}^{T}$ , i.e.,
$$
W _ {g} = \left(\lambda I _ {d} + \zeta_ {g} \zeta_ {g} ^ {T}\right) ^ {- 1} \zeta_ {g} R = \zeta_ {g} \left(\zeta_ {g} ^ {T} \zeta_ {g} + \lambda I _ {n}\right) ^ {- 1} R \tag {16}
$$
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\rightarrow r}$ ,  $C^{a\rightarrow r}$  and Markovian reward function  $g$  and the rewards  $r_t$  are identifiable.
The binary masks,  $C^{s \to s}$ ,  $C^{a \to 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,
$$
\boldsymbol {s} _ {t + 1} = W _ {i, f} ^ {T} \phi_ {f} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \epsilon_ {s, i, t}. \tag {17}
$$
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  $\mathbf{s}_{t + 1}$ , both of which are continuous. The closed-form solution can be represented as:
$$
C \left(W _ {i, f}\right) = \frac {1}{2} \sum_ {\boldsymbol {s} _ {i, t}, \boldsymbol {s} _ {i, t + 1} \sim \mathcal {D}} \left(\boldsymbol {s} _ {i, t + 1} - W _ {i, f} ^ {T} \phi_ {i, f} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right)\right) ^ {2} + \frac {1}{2} \lambda \| W _ {i, f} \| ^ {2}. \tag {18}
$$
By taking derivatives of the cost function and setting them to zero, we can obtain the closed-form solution,
$$
\begin{array}{l} W _ {i, f} = \left(\lambda I _ {d} + \phi_ {i, f} \phi_ {i, f} ^ {T}\right) ^ {- 1} \phi_ {i, f} \boldsymbol {s} _ {i, t + 1} \tag {19} \\ = \phi_ {i, f} \left(\phi_ {i, f} ^ {T} \phi_ {i, f} + \lambda I _ {n}\right) ^ {- 1} \boldsymbol {s} _ {i, t + 1}. \\ \end{array}
$$
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\rightarrow s}$ ,  $C^{a\rightarrow 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} \in V$ ,  $V_{i}$  and  $V_{j}$  are not adjacent in the causal graph  $\mathcal{G}$  if and only if they are conditionally independent given some subset of  $\{V_{l} \mid l \neq i, l \neq 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 \rightarrow r}$ ,  $C^{a \rightarrow r}$ ,  $C^{s \rightarrow s}$ , and  $C^{a \rightarrow 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.
1: Initialize: Environment  $\mathcal{E}$  trajectory  $\tau \gets \emptyset$  buffer  $\mathcal{D} \gets \emptyset$   
2: Initialize: Generative Model  $\Phi_{\mathrm{m}} := [\phi_{\mathrm{cau}}, \phi_{\mathrm{dyn}}, \phi_{\mathrm{rew}}]$ ; Policy Model  $\Phi_{\pi}$   
3: for  $i = 1, 2, \dots, 3 \times 10^4$  do  
4:  $\tau \gets \emptyset$ , reset  $\mathcal{E}$   
5: for  $n_{\mathrm{step}} = 1, 2, \dots, 100$  do  
6: Sample data  $\langle s_t, a_t, o_t \rangle$  from  $\mathcal{E}$ , and store them to trajectory  $\tau$   
7:  
8: if  $\mathcal{E}$  done then  
9: store trajectory  $\tau = \{s_{1:T}, a_{1:T}, R\}$  to buffer  $\mathcal{D}$ , where  $R = \sum_{i=1}^{T} \gamma^{t-1} o_t$   
10:  $\tau \gets \emptyset$ , reset  $\mathcal{E}$   
11: end if  
12:  
for  $n_{\mathrm{batch}} = 1, 2, \dots, \text{train\_batches}$  do  
14: // Optimize Generative Model  $\Phi_{\mathrm{m}}$   
16: Sample  $D_1$  consisting of  $M$  trajectories from  $\mathcal{D}$ :  $D_1 = \{\langle s_t^m, a_t^m \rangle | t=1, R^m\} | M$   
17: Sample binary masks by Gumbel-Softmax from  $\phi_{\mathrm{cau}}: C^{s \rightarrow r}$  and  $C^{a \rightarrow r}$   
18: Optimize  $\phi_{\mathrm{rew}}$  with  $D_1: \phi_{\mathrm{rew}} \gets \phi_{\mathrm{rew}} - \alpha \nabla_{\phi_{\mathrm{rew}}} L_{\mathrm{rew}}$  (Eq. 4)  
19:  
Sample  $D_2$  consisting of  $N$  samples from  $\mathcal{D}$ :  $D_2 = \{s_{t_n}, a_{t_n}, s_{t_n+1}\} | N$   
20: Sample binary masks by Gumbel-Softmax from  $\phi_{\mathrm{cau}}: C^{s \rightarrow r}$ ,  $C^{a \rightarrow r}$ ,  $C^{s \rightarrow s}$  and  $\hat{C}^{a \rightarrow s}$   
21: Optimize  $\phi_{\mathrm{dyn}}$  with  $D_2$  (Using  $C^{s \rightarrow s}$  and  $C^{s \rightarrow s}$ ):  $\phi_{\mathrm{dyn}} \gets \phi_{\mathrm{dyn}} - \alpha \nabla_{\phi_{\mathrm{dyn}}} L_{\mathrm{dyn}}$  (Eq. 5)  
22: Optimize  $\phi_{\mathrm{cau}}: \phi_{\mathrm{cau}} \gets \phi_{\mathrm{cau}} - \alpha \nabla_{\phi_{\mathrm{cau}}} (L_{\mathrm{sp}} + L_{\mathrm{rew}} + L_{\mathrm{dyn}})$  (Eq. 6)  
23:  
Optimize  $\phi_{\mathrm{cau}}: \phi_{\mathrm{cau}} \gets \phi_{\mathrm{cau}} - \alpha \nabla_{\phi_{\mathrm{cau}}} (L_{\mathrm{sp}} + L_{\mathrm{rew}} + L_{\mathrm{dyn}})$  (Eq. 6)  
24: // Optimize Policy Model  $\Phi_{\pi}$   
25: Sample binary masks greedily from  $\phi_{\mathrm{cau}}: C^{s \rightarrow r}$ ,  $C^{a \rightarrow r}$ , and  $C^{s \rightarrow s}$   
26: Calculate  $C^{s \rightarrow r}$  based on  $C^{s \rightarrow r}$  and  $C^{s \rightarrow s}$   
27: Update  $D_2: D_2 \gets \{C^{s \rightarrow r} \odot s_{t_n}, a_{t_n}, s_{t_n+1}, \phi_{\mathrm{rew}}(s_{t_n}, a_{t_n}, C^{s \rightarrow r}, C^{a \rightarrow r})\} | N$   
28: Optimize  $\Phi_{\pi}: \Phi_{\pi} \gets \Phi_{\pi} - \alpha \nabla_{\Phi_{\pi}} J_{\pi}$  (Eq. 8)  
29: end for  
30: end for  
31: end for
C Details of Implementation
C.1Baselines
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  $\Phi_{\mathrm{m}}$  used in the MDP environment consists of three components:  $\phi_{\mathrm{cau}}$ ,  $\phi_{\mathrm{rew}}$ , and  $\phi_{\mathrm{dyn}}$ . We provide a detailed description of their model structures below.
Layer#	1	2	3
φrew	FC256	FC256	FC1
φdyn	FC256	FC256	FC9
φπ	FC256	FC256	FC2|a|
φv	FC256	FC256	FC1
Table 2: The network structures of  $\phi_{\mathrm{rew}}$ ,  $\phi_{\mathrm{dyn}}$ ,  $\phi_{\pi}$  and  $\phi_v$ .  $FC256$  denotes a fully-connected layer with an output size of 256. Each hidden layer is followed by an activation function, ReLU.  $|\alpha|$  is the number of dimensions of the action in a specific task.
$\phi_{\mathrm{cau}}$  for predicting the causal structure.  $\phi_{\mathrm{cau}}$  comprises a set of free parameters without input. We divide  $\phi_{\mathrm{cau}}$  into four parts, each corresponding to the binary masks in Equation 2. Specifically, we have
-  $\phi_{\mathrm{cau}}^{\pmb{s} \rightarrow \pmb{s}} \in \mathbb{R}^{|\pmb{s}| \times |\pmb{s}| \times 2}$  for  $C^{\pmb{s} \rightarrow \pmb{s}} \in \{0,1\}^{|\pmb{s}| \times |\pmb{s}|}$ ,
-  $\phi_{\mathrm{cau}}^{\pmb{a}\rightarrow \pmb{s}}\in \mathbb{R}^{|\pmb{a}|\times |\pmb{s}|\times 2}$  for  $C^{\alpha \to s}\in \mathbb{R}^{|\pmb {a}|\times |\pmb {s}|}$
-  $\phi_{\mathrm{cau}}^{s\rightarrow r}\in \mathbb{R}^{|s|\times 2}$  for  $C^{s\rightarrow r}\in \mathbb{R}^{|s|}$
$\phi_{\mathrm{cau}}^{a\to r}\in \mathbb{R}^{|a|\times 2}$  for  $C^{a\rightarrow r}\in \mathbb{R}^{|a|}$
Below we explain the shared workflows in  $\phi_{\mathrm{cau}}$  using the example of predicting the causal edge from the  $i$ -th dimension of state  $\pmb{s}_{i,t}$  to the  $j$ -th dimension of the next state  $\pmb{s}_{j,t+1}$ , by part of the free parameters,  $\phi_{\mathrm{cau},i,j}^{s \rightarrow s}$ .
For simplicity, we denote  $\phi_{\mathrm{cau},i,j}^{s\rightarrow s}$  as  $\psi$ . The shape of  $\psi$  is now easy to be determined. That is  $\psi \in \mathbb{R}^2$  and we write it as  $\psi = [\psi_0,\psi_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  $(\psi_0)$  or not existing  $(\psi_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:
$$
P \left(C _ {i, j} ^ {\mathbf {s} \rightarrow \mathbf {s}}\right) = \frac {\exp \left(\psi_ {0}\right)}{\exp \left(\psi_ {0}\right) + \exp \left(\psi_ {1}\right)}. \tag {20}
$$
One example of using the estimated mask is given as Figure 8.
During training: obtain  $C_{i,j}^{s \rightarrow s}$  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_{i,j}^{s \rightarrow s}$  by applying Gumbel-Softmax [48],
$$
C _ {i, j} ^ {\mathbf {s} \rightarrow \mathbf {s}} = \operatorname {G S} (\psi) \tag {21}
$$
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  $\phi_{\mathrm{cau}}$  for predicting a single causal edge provides insight into the overall implementation of the entire module  $\phi_{\mathrm{cau}}$  and can be applicable for predicting all the causal edges during the training phase. Therefore, we can obtain  $C^{a\rightarrow s}$  for optimizing  $\phi_{\mathrm{dyn}}$ ,  $C^{s\rightarrow r}$  and  $C^{a\rightarrow r}$  for optimizing  $\phi_{\mathrm{rew}}$  using similar procedures.
During inference: obtain  $C_{i,j}^{s \rightarrow s}$  by greedy selection. In the inference phase, including data sampling and policy learning, we get the prediction of  $C_{i,j}^{s \rightarrow s}$  through a greedy selection,
$$
C _ {i, j} ^ {\mathbf {s} \rightarrow \mathbf {s}} = \left\{\begin{array}{l}1, \psi_ {0} \geq \psi_ {1}\\0, \psi_ {0} <   \psi_ {1}.\end{array}\right. \tag {22}
$$
Such a greedy sampling avoids introducing randomness using the Gumble-Softmax sampling.
The above explanation of the workflow in  $\phi_{\mathrm{cau}}$  for predicting a single causal edge provides insight into the overall implementation of the entire module  $\phi_{\mathrm{cau}}$  and can be applicable for predicting all the causal edges during the inference phase. Therefore, we can obtain  $C^{s\rightarrow s}$ ,  $C^{s\rightarrow r}$  and  $C^{a\rightarrow r}$  using similar procedures to predict the Markovian reward and extract compact representation  $\pmb{s}^{\mathrm{min}}$ .
$\phi_{\mathrm{rew}}$  for predicting the Markovian rewards.  $\phi_{\mathrm{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
Figure 8: The illustration of using learnable masks to predict the next state.
written as,
$$
\hat {r} = \phi_ {\mathrm {r e w}} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}, \boldsymbol {C} ^ {\boldsymbol {s} \rightarrow r}, \boldsymbol {C} ^ {\boldsymbol {a} \rightarrow r}\right) = \operatorname {F C s} \left(\left[ \boldsymbol {C} ^ {\boldsymbol {s} \rightarrow r} \odot \boldsymbol {s} _ {t}, \boldsymbol {C} ^ {\boldsymbol {a} \rightarrow r} \odot \boldsymbol {a} _ {t} \right]\right), \tag {23}
$$
where  $[\cdot ,\cdot ]$  ,  $\odot$  denotes concatenation and element-wise multiply operations, respectively. FCs denote the stacked fully-connected network.  $C^{s\rightarrow r}$  and  $C^{s\rightarrow r}$  are derived from  $\phi_{\mathrm{cau}}$  by Gumbel-Softmax.
During inference, including policy learning and data sampling, the predicted Markovian reward is
$$
\hat {r} = \phi_ {\mathrm {r e w}} \left(\boldsymbol {s} _ {t}, \boldsymbol {a}\right) _ {t}, \boldsymbol {C} ^ {s \rightarrow r}, \boldsymbol {C} ^ {\boldsymbol {a} \rightarrow r}) = \operatorname {F C s} \left(\left[ \boldsymbol {C} ^ {s \rightarrow r} \odot \boldsymbol {s} _ {t}, \boldsymbol {C} ^ {\boldsymbol {a} \rightarrow r} \odot \boldsymbol {a} _ {t} \right]\right), \tag {24}
$$
where  $C^{s\rightarrow r}$  and  $C^{s\rightarrow r}$  are derived from  $\phi_{\mathrm{cau}}$  greedily by deterministic sampling.
$\phi_{\mathrm{dyn}}$  for modeling the environment dynamics. In our experiment, we do not directly utilize  $\phi_{\mathrm{dyn}}$  in policy learning. Instead, this module serves as a bridge to optimize  $\phi_{\mathrm{cau}}^{s\rightarrow s}$  and  $\phi_{\mathrm{cau}}^{a\rightarrow s}$ . Subsequently,  $\phi_{\mathrm{cau}}^{s\rightarrow s}$  can be utilized in the calculation of  $C^{s\rightarrow \pi}$ .
During training, we initially sample  $C^{s \to s}$  and  $C^{a \to s}$  using Gumbel-Softmax. The prediction for the  $i$ -th dimension of the next state can be represented as follows,
$$
\hat {\boldsymbol {s}} _ {i, t} = \operatorname {M D N} \left(\left[ \boldsymbol {C} _ {., i} ^ {\boldsymbol {s} \rightarrow \boldsymbol {s}} \odot \boldsymbol {s} _ {t}, \boldsymbol {C} _ {., i} ^ {\boldsymbol {a} \rightarrow \boldsymbol {s}} \odot \boldsymbol {a} _ {t} \right]\right), \tag {25}
$$
where  $[\cdot ,\cdot ]$  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  $\phi_{\mathrm{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  $\Phi_{\pi}$  comprises two components, the actor  $\phi_{\pi}$  and the critic  $\phi_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  $\mathcal{D}$ :
-  $D_{1} = \{\langle \pmb{s}_{t}^{m},\pmb{a}_{t}^{m}\rangle |_{t = 1}^{T},R^{m}\} |_{m = 1}^{M}$  consists of  $M$  trajectories. Provided with the trajectory-wise long-term returns  $R^{m}|_{m = 1}^{M}$ ,  $D_{1}$  is utilized to optimize  $\phi_{\mathrm{cau}}^{s\rightarrow r}$ ,  $\phi_{\mathrm{cau}}^{a\rightarrow r}$  and  $\phi_{\mathrm{rew}}$ , with  $L_{\mathrm{rew}}$ .
-  $D_{2} = \{s_{t_{n}}, a_{t_{n}}, s_{t_{n + 1}}\} \mid_{n = 1}^{N}$  consists of  $N$  state-action pairs.  $D_{2}$  are used for policy optimization and optimize the parts of causal structure,  $\phi_{\mathrm{cau}}^{\mathbf{s}\rightarrow \mathbf{s}}$  and  $\phi_{\mathrm{cau}}^{\mathbf{a}\rightarrow \mathbf{s}}$ ,  $\phi_{\mathrm{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.
Table 3: The table of the hyper-parameters used in the experiments for GRD.
Envs	λ1	λ2	λ3	λ4	λ5
Ant	10-5	0	10-7	10-8	10-8
HalfCheetah	10-5	10-5	10-5	10-6	10-5
Walker2d	10-5	10-5	10-6	10-6	10-7
Humanoid	10-5	10-8	10-5	10-7	10-8
Reacher	5 × 10-7	10-8	10-8	10-8	10-8
Swimmer	10-7	10-9	10-9	0	10-9
Hopper	10-6	10-6	10-6	10-7	10-6
HumanStandup	10-5	10-4	10-6	10-7	10-7
Table 4: The hyper-parameters.
hyperparameters	value	hyperparameters	value
epochs	3	optimizer	Adam
cycles	100	learning rate	3 × 10-4
iteration	100	N	256
train_batches	100	M	4
replay_buffer size	10^6	γ	1.00
evaluation episodes	10	Polyak-averaging coefficient	0.0005
C.5 Hyper-Parameters
The network is trained from scratch using the Adam optimizer, without any pre-training. The initial learning rate for both model estimation and policy learning is set to  $3 \times 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 \times 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 \times 10^{6}$  time steps. The warmup size of the buffer for training is set to  $1 \times 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  $\phi_{\mathrm{rew}}$ , we sample 4 episodes, each containing  $5 \times 10^{3}$  steps. As for policy learning and the optimization of  $\phi_{\mathrm{dyn}}$ , we use data from 256 time steps.  $\phi_{\mathrm{cau}}$  is trained together with  $\phi_{\mathrm{rew}}$  and  $\phi_{\mathrm{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\mathrm{GHz}$  and 5 terabytes of memory.
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.
Figure 9: Evaluation on the manipulation tasks of MetaWorld, Door Lock, Push Wall, Pick Place.
D.3 Visualization of Learned Causal Structure
We provide the visualization of learned causal structure in Swimmer, 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  $\pmb{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  $\pmb{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.