SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The coefficients corresponding to zero-reward outcomes are multiplied with zero in the gradient estimator, and can hence be ignored. Now, we proceed by showing that  $\mathbb{E}[\hat{\nabla}_{\theta}^{\mathrm{REINFORCE}}V^{\pi}(A_0,R_T)\mid S_T,R_T] = \hat{\nabla}_{\theta}^{S}V^{\pi}(S_T,R_T),\mathbb{E}[\hat{\nabla}_{\theta}^{S}V^{\pi}(S_T,R_T)\mid U_T',R_T] = \hat{\nabla}_{\theta}^{U'}V^{\pi}(U_T',R_T),\mathbb{E}[\hat{\nabla}_{\theta}^{U'}V^{\pi}(U_T',R_T)\mid U_T,R_T] = \hat{\nabla}_{\theta}^{U}V^{\pi}(U_T,R_T)$ , and  $\mathbb{E}[\hat{\nabla}_{\theta}^{U}V^{\pi}(U_{T},R_{T})]$ $R_{T}] = \hat{\nabla}_{\theta}^{R}V^{\pi}(R_{T})$  , after which we can use the law of total variance to prove our theorem.
As  $S_{T}$  is fully predictive of  $R_{T}$ , the following conditional independence holds  $p^{\pi}(A_0 \mid S_0 = s, S_T, R_T) = p^{\pi}(A_0 \mid S_0 = s, S_T)$  (c.f. Definition 2). Hence, we have that
$$
\begin{array}{l} \mathbb {E} \left[ \hat {\nabla} _ {\theta} ^ {\text {R E I N F O R C E}} V ^ {\pi} \left(A _ {0}, R _ {T}\right) \mid S _ {T}, R _ {T} \right] (37) \\ = \sum_ {a \in \mathcal {A}} p \left(A _ {0} = a \mid S _ {0} = s, S _ {T}\right) \frac {\nabla_ {\theta} \pi (a \mid s)}{\pi (a \mid s)} R _ {T} = \hat {\nabla} _ {\theta} ^ {S} V ^ {\pi} \left(S _ {T}, R _ {T}\right) (38) \\ \end{array}
$$
where we used that  $\sum_{a}\nabla_{\theta}\pi (a\mid s) = \nabla_{\theta}\sum_{a}\pi (a\mid s) = 0$
Similarly, as  $S_T$  is fully predictive of  $U_T'$  (c.f. Definition 2), we have that
$$
\begin{array}{l} p \left(A \mid S, U _ {T} ^ {\prime}\right) = \sum_ {s _ {T} \in \mathcal {S}} p \left(S _ {T} = s _ {T} \mid S, U _ {T} ^ {\prime}\right) p \left(A \mid S, S _ {T} = s _ {T}, U _ {T} ^ {\prime}\right) (39) \\ = \sum_ {s _ {T} \in \mathcal {S}} p \left(S _ {T} = s _ {T} \mid S, U _ {T} ^ {\prime}\right) p \left(A \mid S, S _ {T} = s _ {T}\right) (40) \\ \end{array}
$$
Using the conditional independence relation  $p^{\pi}(S_T \mid S_0, U_T', R_T) = p^{\pi}(S_T \mid S_0, U_T')$ , following from d-separation in Fig. 6b, this leads us to
$$
\begin{array}{l} \mathbb {E} \left[ \hat {\nabla} _ {\theta} ^ {S} V ^ {\pi} \left(S _ {T}, R _ {T}\right) \mid U _ {T} ^ {\prime}, R _ {T} \right] (41) \\ = \sum_ {a \in \mathcal {A}} \sum_ {s _ {T} \in \mathcal {S}} p \left(S _ {T} = s _ {T} \mid S _ {0} = s, U _ {T} ^ {\prime}\right) p \left(A _ {0} = a \mid S = s, S _ {T} = s _ {T}\right) \frac {\nabla_ {\theta} \pi (a \mid s)}{\pi (a \mid s)} R _ {T} (42) \\ = \sum_ {a \in \mathcal {A}} p \left(A _ {0} = a \mid S = s, U _ {T} ^ {\prime}\right) \frac {\nabla_ {\theta} \pi (a \mid s)}{\pi (a \mid s)} R _ {T} (43) \\ = \hat {\nabla} _ {\theta} ^ {U ^ {\prime}} V ^ {\pi} \left(U _ {T} ^ {\prime}, R _ {T}\right) (44) \\ \end{array}
$$
Using the same derivation leveraging the fully predictive properties (c.f. Definition 2), we get
$$
\mathbb {E} \left[ \hat {\nabla} _ {\theta} ^ {U ^ {\prime}} V ^ {\pi} \left(U _ {T} ^ {\prime}, R _ {T}\right) \mid U _ {T}, R _ {T} \right] = \hat {\nabla} _ {\theta} ^ {U} V ^ {\pi} \left(U _ {T}, R _ {T}\right) \tag {45}
$$
$$
\mathbb {E} \left[ \hat {\nabla} _ {\theta} ^ {U} V ^ {\pi} \left(U _ {T}, R _ {T}\right) \mid R _ {T} \right] = \hat {\nabla} _ {\theta} ^ {R} V ^ {\pi} \left(R _ {T}\right) \tag {46}
$$
Now we can use the law of total variance, which states that  $\mathbb{V}[X] = \mathbb{E}[\mathbb{V}[X\mid Y]] + \mathbb{V}[\mathbb{E}[X\mid Y]]$ . Hence, we have that
$$
\begin{array}{l} \mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {\text {R E I N F O R C E}} V ^ {\pi} \right] = \mathbb {V} \left[ \mathbb {E} \left[ \hat {\nabla} _ {\theta} ^ {\text {R E I N F O R C E}} V ^ {\pi} \mid S _ {T} \right] \right] + \mathbb {E} \left[ \mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {\text {R E I N F O R C E}} V ^ {\pi} \mid S _ {T} \right] \right] (47) \\ = \mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {S} V ^ {\pi} \right] + \mathbb {E} \left[ \mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {\text {R E I N F O R C E}} V ^ {\pi} \mid S _ {T} \right] \right] \succcurlyeq \mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {S} V ^ {\pi} \right] (48) \\ \end{array}
$$
as  $\mathbb{E}[\mathbb{V}[\hat{\nabla}_{\theta}^{\mathrm{REINFORCE}}V^{\pi}\mid S_T]]$  is positive semi definite. Using the same construction for the other pairs, we arrive at
$$
\mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {R} V ^ {\pi} \left(s _ {0}\right) \right] \preccurlyeq \mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {U} V ^ {\pi} \left(s _ {0}\right) \right] \preccurlyeq \mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {U ^ {\prime}} V ^ {\pi} \left(s _ {0}\right) \right] \preccurlyeq \mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {S} V ^ {\pi} \left(s _ {0}\right) \right] \preccurlyeq \mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {\text {R E I N F O R C E}} V ^ {\pi} \left(s _ {0}\right) \right] \tag {49}
$$
thereby concluding the proof.
Additional empirical verification. To get more insight into how the information content of the rewarding-outcome encoding  $U$  relates to the variance of the COCOA gradient estimator, we repeat the experiment of Fig. 2, but now plot the variance as a function of the amount of information in the rewarding outcome encodings, for a fixed state overlap of 3. Fig. 7 shows that the variance of the resulting COCOA gradient estimators interpolate between COCOA-reward and HCA+, with more informative encodings leading to a higher variance. These empirical results show that the insights of Theorem 3 hold in this more general setting of estimating full policy gradients in an MDP with random rewards.
Proof. This proof follows a similar technique to the policy gradient theorem [26]. Let us first define the expected number of occurrences of  $u$  starting from state  $s$  as  $O^{\pi}(u,s) = \sum_{k\geq 1}p^{\pi}(U_k = u \mid S_0 = s)$ , and its action equivalent  $O^{\pi}(u,s,a) = \sum_{k\geq 1}p^{\pi}(U_k = u \mid S_0 = s, A_0 = a)$ . Expanding  $O^{\pi}(u,s)$  leads to
$$
\begin{array}{l} \nabla_ {\theta} O ^ {\pi} (u, s) = \nabla_ {\theta} \left[ \sum_ {a \in \mathcal {A}} \pi (a \mid s) \sum_ {s ^ {\prime} \in \mathcal {S}} p \left(s ^ {\prime} \mid s, a\right) \left(\delta \left(s ^ {\prime} = s\right) + O ^ {\pi} \left(u, s ^ {\prime}\right)\right) \right] (50) \\ = \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid s) O ^ {\pi} (u, s, a) + \sum_ {s ^ {\prime} \in \mathcal {S}} p \left(s ^ {\prime} \mid s\right) \nabla_ {\theta} O ^ {\pi} \left(u, s ^ {\prime}\right) (51) \\ \end{array}
$$
Figure 7: Less informative rewarding outcome encodings lead to gradient estimators with lower variance. Normalized variance in dB using ground-truth coefficients and a random uniform policy, for various gradient estimators on the tree environment (shaded region represents standard error over 10 random environments). To increase the information content of  $U$ , we increase the number of different encodings  $u$  corresponding to the same reward value (c.f.  $n_g$  in Section E.6), indicated by the x-axis. COCOA-u-group indicates the COCOA estimator with rewarding outcome encodings of increasing information content, whereas COCOA-reward and HCA+ have fixed rewarding outcome encodings of  $U = R$  and  $U = S$  respectively.
Now define  $\phi(u, s) = \sum_{a \in \mathcal{A}} \nabla_{\theta} \pi(a \mid s) O^{\pi}(u, s, a)$ , and  $p^{\pi}(S_l = s' \mid S_0 = s)$  as the probability of reaching  $s'$  starting from  $s$  in  $l$  steps. We have that  $\sum_{s'' \in \mathcal{S}} p^{\pi}(S_l = s'' \mid S_0 = s) p^{\pi}(S_1 = s' \mid S_0 = s'') = p^{\pi}(S_{l+1} = s' \mid S_0 = s)$ . Leveraging this relation and recursively applying the above equation leads to
$$
\begin{array}{l} \nabla_ {\theta} O ^ {\pi} (u, s) = \sum_ {s ^ {\prime} \in \mathcal {S}} \sum_ {l = 0} ^ {\infty} p ^ {\pi} \left(S _ {l} = s ^ {\prime} \mid S _ {0} = s\right) \phi (u, s ^ {\prime}) (52) \\ = \sum_ {s ^ {\prime} \in \mathcal {S}} \sum_ {l = 0} ^ {\infty} p ^ {\pi} \left(S _ {l} = s ^ {\prime} \mid S _ {0} = s\right) \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid s ^ {\prime}) O ^ {\pi} (u, s ^ {\prime}, a) (53) \\ \propto \mathbb {E} _ {S \sim \mathcal {T} (s, \pi)} \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid S) O ^ {\pi} (u, S, a) (54) \\ \end{array}
$$
where in the last step we normalized  $\sum_{l=0}^{\infty} p^{\pi}(S_l = s' \mid S_0 = s)$  with  $\sum_{s' \in S} \sum_{l=0}^{\infty} p^{\pi}(S_l = s' \mid S_0 = s)$ , resulting in the state distribution  $S \sim \mathcal{T}(s, \pi)$  where  $S$  is sampled from trajectories starting from  $s$  and following policy  $\pi$ .
Finally, we can rewrite the contribution coefficients as
$$
w (s, a, u) = \frac {O ^ {\pi} (u , s , a)}{O ^ {\pi} (u , s)} - 1 \tag {55}
$$
$$
\nabla_ {\theta} O ^ {\pi} (u, s) \propto \mathbb {E} _ {S \sim \mathcal {T} (s, \pi)} \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid S) w (S, a, u) O ^ {\pi} (u, S) \tag {56}
$$
where we used that  $\sum_{a\in \mathcal{A}}\nabla_{\theta}\pi (a\mid s) = 0$  , thereby concluding the proof.
D Learning the contribution coefficients
D.1 Proof Proposition 5
Proof. In short, as the logits  $l$  are a deterministic function of the state  $s$ , they do not provide any further information on the hindsight distribution and we have that  $p^{\pi}(A_0 = a \mid S_0 = s, L_0 = l, U' = u') = p^{\pi}(A_0 = a \mid S_0 = s, U' = u')$ .
We can arrive at this result by observing that  $p(A_0 = a \mid S_0 = s, L_0 = l) = p(A_0 = a \mid S_0 = s) = \pi(a \mid s)$ , as the policy logits encode the policy distribution. Using this, we have that
$$
\begin{array}{l} p ^ {\pi} \left(A _ {0} = a \mid S _ {0} = s, L _ {0} = l, U ^ {\prime} = u ^ {\prime}\right) (57) \\ = \frac {p ^ {\pi} \left(A _ {0} = a , U ^ {\prime} = u ^ {\prime} \mid S _ {0} = s , L _ {0} = l\right)}{p ^ {\pi} \left(U ^ {\prime} = u ^ {\prime} \mid S _ {0} = s , L _ {0} = l\right)} (58) \\ = \frac {p ^ {\pi} \left(U ^ {\prime} = u ^ {\prime} \mid S _ {0} = s , A _ {0} = a , L _ {0} = l\right) p \left(A _ {0} = a \mid S _ {0} = s , L _ {0} = l\right)}{\sum_ {a ^ {\prime} \in \mathcal {A}} p ^ {\pi} \left(A _ {0} = a ^ {\prime} , U ^ {\prime} = u ^ {\prime} \mid S _ {0} = s , L _ {0} = l\right)} (59) \\ = \frac {p ^ {\pi} \left(U ^ {\prime} = u ^ {\prime} \mid S _ {0} = s , A _ {0} = a\right) p \left(A _ {0} = a \mid S _ {0} = s\right)}{\sum_ {a ^ {\prime} \in \mathcal {A}} p ^ {\pi} \left(U ^ {\prime} = u ^ {\prime} \mid S _ {0} = s , A _ {0} = a ^ {\prime} ,\right) p \left(A _ {0} = a ^ {\prime} \mid S _ {0} = s\right)} (60) \\ = \frac {p ^ {\pi} \left(A _ {0} = a , U ^ {\prime} = u ^ {\prime} \mid S _ {0} = s\right)}{p ^ {\pi} \left(U ^ {\prime} = u ^ {\prime} \mid S _ {0} = s\right)} (61) \\ = p ^ {\pi} \left(A _ {0} = a \mid S _ {0} = s, U ^ {\prime} = u ^ {\prime}\right) (62) \\ \end{array}
$$
D.2 Learning contribution coefficients via contrastive classification.
Instead of learning the hindsight distribution, we can estimate the probability ratio  $p^{\pi}(A_t = a \mid S_t = s, U' = u') / \pi(a \mid s)$  directly, by leveraging contrastive classification. Proposition 7 shows that by training a binary classifier  $D(a, s, u')$  to distinguish actions sampled from  $p^{\pi}(A_t = a \mid S_t = s, R' = r')$  versus  $\pi(a \mid s)$ , we can directly estimate the probability ratio.
Proposition 7. Consider the contrastive loss
$$
L = \mathbb {E} _ {s, u ^ {\prime} \sim \mathcal {T} (s _ {0}, \pi)} \left[ \mathbb {E} _ {a \sim p ^ {\pi} (a | s, u ^ {\prime})} [ \log D (a, s, u ^ {\prime}) ] + \mathbb {E} _ {a \sim \pi (a | s)} [ \log (1 - D (a, s, u ^ {\prime})) ] \right], \tag {63}
$$
and  $D^{*}(s,a,u^{\prime})$  its minimizer. Then the following holds
$$
w (s, a, u ^ {\prime}) = \frac {D ^ {*} (a , s , u ^ {\prime})}{1 - D ^ {*} (a , s , u ^ {\prime})} - 1. \tag {64}
$$
Proof. Consider a fixed pair  $s, u'$ . We can obtain the discriminator  $D^{*}(a, s, u')$  that maximizes  $\mathbb{E}_{a \sim p^{\pi}(a|s, u')} [\log D(a, s, u')] + \mathbb{E}_{a \sim \pi(a|s)} [\log (1 - D(a, s, u'))]$  by taking the point-wise derivative of this objective and equating it to zero:
$$
\begin{array}{l} p ^ {\pi} (a \mid s, u ^ {\prime}) \frac {1}{D ^ {*} (a , s , u ^ {\prime})} - \pi (a \mid s) \frac {1}{1 - D ^ {*} (a , s , u ^ {\prime})} = 0 (65) \\ \Rightarrow \quad \frac {p ^ {\pi} (a \mid s , u ^ {\prime})}{\pi (a \mid s)} = \frac {D ^ {*} (a , s , u ^ {\prime})}{1 - D ^ {*} (a , s , u ^ {\prime})} (66) \\ \end{array}
$$
As for any  $(a,b) \in \mathbb{R}_{/0}^{2}$ , the function  $f(x) = a\log x + b\log (1 - x)$  achieves its global maximum on support  $x \in [0,1]$  at  $\frac{a}{a + b}$ , the above maximum is a global maximum. Repeating this argument for all  $(s,u')$  pairs concludes the proof.
We can approximate the training objective of  $D$  by sampling  $a^{(m)}, s^{(m)}, u'^{(m)}$  along the observed trajectories, while leveraging that we have access to the policy  $\pi$ , leading to the following loss
$$
L = \sum_ {m = 1} ^ {M} \left[ - \log D \left(a ^ {(m)}, s ^ {(m)}, u ^ {\prime (m)}\right) - \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s ^ {(m)}\right) \log \left(1 - D \left(a ^ {\prime}, s ^ {(m)}, u ^ {\prime (m)}\right)\right) \right] \tag {67}
$$
Numerical stability. Assuming  $D$  uses the sigmoid nonlinearity on its outputs, we can improve the numerical stability by computing the logarithm and sigmoid jointly. This results in
$$
\log \sigma (x) = - \log (1 + \exp (- x)) \tag {68}
$$
$$
\log (1 - \sigma (x)) = - x - \log (1 + \exp (- x)) \tag {69}
$$
The probability ratio  $D / (1 - D)$  can then be computed with
$$
\frac {\sigma (x)}{1 - \sigma (x)} = \exp [ - \log (1 + \exp (- x)) + x + \log (1 + \exp (- x)) ] = \exp (x) \tag {70}
$$
D.3 Successor representations.
If we take  $u'$  equal to the state  $s'$ , we can observe that the sum  $\sum_{k \geq 1} p^{\pi}(S_{t+k} = s' \mid s, a)$  used in the contribution coefficients (3) is equal to the successor representation  $M(s, a, s')$  introduced by Dayan [30]. Hence, we can leverage temporal differences to learn  $M(s, a, s')$ , either in a tabular setting [30], or in a deep feature learning setting [31]. Using the successor representations, we can construct the state-based contribution coefficients as  $w(s, a, s') = M(s, a, s') / [\sum_{\tilde{a} \in \mathcal{A}} \pi(\tilde{a} \mid s) M(s, \tilde{a}, s')] - 1$ .
For a rewarding outcome encoding  $U$  different from  $S$ , we can recombine the state-based successor representations to obtain the required contribution coefficients using the following derivation.
$$
\begin{array}{l} \sum_ {k \geq 0} p ^ {\pi} \left(U _ {k} = u ^ {\prime} \mid S _ {0} = s, A _ {0} = a\right) = \sum_ {k \geq 0} \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(U _ {k} = u ^ {\prime}, S _ {k} = s ^ {\prime} \mid S _ {0} = s, A _ {0} = a\right) (71) \\ = \sum_ {k \geq 0} \sum_ {s ^ {\prime} \in \mathcal {S}} p \left(U ^ {\prime} = u ^ {\prime} \mid S ^ {\prime} = s ^ {\prime}\right) p ^ {\pi} \left(S _ {k} = s ^ {\prime} \mid S _ {0} = s, A _ {0} = a\right) (72) \\ = \sum_ {s ^ {\prime} \in S} p \left(U ^ {\prime} = u ^ {\prime} \mid S ^ {\prime} = s ^ {\prime}\right) \sum_ {k \geq 0} p ^ {\pi} \left(S _ {k} = s ^ {\prime} \mid S _ {0} = s, A _ {0} = a\right) (73) \\ = \sum_ {s ^ {\prime} \in \mathcal {S}} p \left(U ^ {\prime} = u ^ {\prime} \mid S ^ {\prime} = s ^ {\prime}\right) M (s, a, s ^ {\prime}) (74) \\ \end{array}
$$
where in the second line we used that  $S$  is fully predictive of  $U$  (c.f. Definition 2). Note that  $p(U' \mid S')$  is policy independent, and can hence be approximated efficiently using offline data. We use this recombination of successor representations in our dynamic programming setup to compute the ground-truth contribution coefficients (c.f. App. E).
E Experimental details and additional results
Here, we provide additional details to all experiments performed in the manuscript and present additional results.
Algorithm 1 shows pseudocode for the COCOA family of algorithms.
E.2 Dynamic programming setup
In order to delineate the different policy gradient methods considered in this work, we develop a framework to compute ground-truth policy gradients and advantages as well as ground-truth contribution coefficients. We will first show how we can recursively compute a quantity closely related to the successor representation using dynamic programming which then allows us to obtain expected advantages and finally expected policy gradients.
E.2.1 Computing ground truth quantities
To compute ground truth quantities, we assume that the environment reaches an absorbing, terminal state  $s_{\infty}$  after  $T$  steps for all states  $s$ , i.e.  $p^{\pi}(S_T = s_{\infty}|s_0 = s) = 1$ . This assumption is satisfied by the linear key-to-door and tree environment we consider. As a helper quantity, we define the successor representation:
$$
M (s, a, s ^ {\prime}, T) = \sum_ {k = 1} ^ {T} p ^ {\pi} \left(S _ {k} = s ^ {\prime} \mid S _ {0} = s, A _ {0} = a\right) \tag {75}
$$
which captures the cumulative probability over all time steps of reaching state  $s'$  when starting from state  $S_0 = s$ , choosing initial action  $A_0 = a$  and following the policy  $\pi$  thereafter. We can compute  $M(s, a, s', T)$  recursively as
$$
M (s, a, s ^ {\prime}, t) = \sum_ {s ^ {\prime \prime} \in \mathcal {S}} p \left(S _ {1} = s ^ {\prime \prime} \mid S _ {0} = s, A _ {0} = a\right) \sum_ {a ^ {\prime \prime} \in \mathcal {A}} \pi \left(a ^ {\prime \prime} \mid s ^ {\prime \prime}\right) \left(\mathbb {1} _ {s ^ {\prime \prime} = s ^ {\prime}} + M \left(s ^ {\prime \prime}, a ^ {\prime \prime}, s ^ {\prime}, t - 1\right)\right) \tag {76}
$$
where  $\mathbb{1}$  is the indicator function and  $M(s,a,s^{\prime},0)$  is initialized to 0 everywhere. Then, the various quantities can be computed exactly as follows.
Algorithm 1 COCOA
Require: Initial  $\pi$ ,  $h$ , episode length  $T$ , number of episodes  $N$ , number of pretraining episodes  $M$ , batch size  $L$ ,  $K$  number of pretraining update steps.
1: for  $i = 1$  to  $M$  do  $\triangleright$  Collect random trajectories
2: Sample  $L$  trajectories  $\tau = \{\{S_t, A_t, R_t\}_{t=0}^T\}_{l=1}^L$  from a random policy
3: Add trajectories to the buffer
4: end for
5: for  $j = 1$  to  $K$  do  $\triangleright$  Pretrain reward features
6: Sample batch from buffer
7: Train reward features  $U = f(S, A)$  to predict rewards  $R$  via mean squared error
8: end for
9: for  $i = 1$  to  $N - M$  do
10: Sample  $L$  trajectories  $\tau = \{\{S_t, A_t, R_t\}_{t=0}^T\}_{l=1}^L$  from  $\pi$
1: for  $t = 1$  to  $T$  do  $\triangleright$  Update the contribution module
2: for  $k > t$  do
3: Train  $h(A_{t}|S_{t},U_{k},\pi_{t})$  via cross-entropy loss on all trajectories  $\tau$
4: end for
5: end for
6: for  $t = 1$  to  $T$  do  $\triangleright$  Update the policy
7: for  $a\in \mathcal{A}$  do
8:  $X_{a} = \nabla_{\theta}\pi (a\mid S_{t})\sum_{k\geq 1}w(S_{t},a,U_{t + k})R_{t + k}$
9: end for
20:  $\hat{\nabla}_{\theta}^{U}V^{\pi}(s_{0}) = \sum_{t\geq 0}\nabla_{\theta}\log \pi (A_{t}\mid S_{t})R_{t} + \sum_{a\in \mathcal{A}}X_{a}$
21: end for
22: end for
Contribution coefficients
$$
w (s, a, u ^ {\prime}) = \frac {M (s , a , u ^ {\prime} , T)}{\sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s\right) M \left(s , a ^ {\prime} , u ^ {\prime} , T\right)} - 1 \tag {77}
$$
$$
M (s, a, u ^ {\prime}, T) = \sum_ {s ^ {\prime} \in \mathcal {S}} M (s, a, s ^ {\prime}, T) \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s ^ {\prime}\right) \sum_ {r ^ {\prime}} p \left(r ^ {\prime} \mid s ^ {\prime}, a ^ {\prime}\right) \mathbb {1} _ {f \left(s ^ {\prime}, a ^ {\prime}, r ^ {\prime}\right) = u ^ {\prime}}, \tag {78}
$$
similar to the successor representations detailed in Section D. To discover the full state space  $S$  for an arbitrary environment, we perform a depth-first search through the environment transitions (which are deterministic in our setting), and record the encountered states.
$$
\begin{array}{l} V (s) = \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s\right) r \left(s, a ^ {\prime}\right) + \sum_ {k = 1} ^ {T} \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S _ {k} = s ^ {\prime} \mid S _ {0} = s, \pi\right) \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s ^ {\prime}\right) r \left(s ^ {\prime}, a ^ {\prime}\right) (79) \\ = \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} | s\right) r \left(s, a ^ {\prime}\right) + \sum_ {s ^ {\prime} \in \mathcal {S}} \sum_ {a \in \mathcal {A}} \pi (a \mid s) M \left(s, a, s ^ {\prime}, T\right) \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} | s ^ {\prime}\right) r \left(s ^ {\prime}, a ^ {\prime}\right) (80) \\ \end{array}
$$
where  $r(s', a') = \sum_{r'} r' p(r' \mid s', a')$
Action-value function
$$
\begin{array}{l} Q (s, a) = r (s, a) + \sum_ {k = 1} ^ {T} \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S _ {k} = s ^ {\prime} \mid S _ {0} = s, A _ {0} = a\right) \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s ^ {\prime}\right) r \left(s ^ {\prime}, a ^ {\prime}\right) (81) \\ = r (s, a) + \sum_ {s ^ {\prime} \in \mathcal {S}} M (s, a, s ^ {\prime}, T) \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s ^ {\prime}\right) r \left(s ^ {\prime}, a ^ {\prime}\right) (82) \\ \end{array}
$$
E.2.2 Computing expected advantages
As a next step we detail how to compute the true expected advantage given a potentially imperfect estimator of the contribution coefficient  $\hat{w}$ , the value function  $\hat{V}$  or the action-value function  $\hat{Q}$ .
$$
\begin{array}{l} \mathbb {E} _ {\pi} \left[ \hat {A} ^ {\pi} (s, a) \right] = r (s, a) - \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s\right) r \left(s, a ^ {\prime}\right) + \mathbb {E} _ {T \sim \mathcal {T} (s, \pi)} \left[ \sum_ {k = 1} ^ {T} \hat {w} (s, a, U _ {k}) R _ {k} \right] (83) \\ = r (s, a) - \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s\right) r \left(s, a ^ {\prime}\right) + \sum_ {u ^ {\prime} \in \mathcal {U}} \sum_ {k = 1} ^ {T} p ^ {\pi} \left(U _ {k} = u ^ {\prime} \mid S _ {0} = s\right) \hat {w} \left(s, a, u ^ {\prime}\right) r \left(u ^ {\prime}\right) (84) \\ = r (s, a) - \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s\right) r \left(s, a ^ {\prime}\right) + \sum_ {u ^ {\prime} \in \mathcal {U}} \hat {w} \left(s, a, u ^ {\prime}\right) r \left(u ^ {\prime}\right) \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s\right) M \left(s, a ^ {\prime}, u ^ {\prime}, T\right) (85) \\ \end{array}
$$
where  $r(u^{\prime}) = \sum_{r^{\prime}}r^{\prime}p(r^{\prime}\mid u^{\prime})$  and  $M(s,a^{\prime},u^{\prime},T)$  is defined as in E.2.1.
$$
\mathbb {E} _ {\pi} \left[ \hat {A} ^ {\pi} (s, a) \right] = Q (s, a) - \hat {V} (s) \tag {86}
$$
where  $Q$  is the ground truth action-value function obtained following E.2.1.
Q-critic
$$
\mathbb {E} _ {\pi} \left[ \hat {A} ^ {\pi} (s, a) \right] = \hat {Q} (s, a) - \sum_ {a ^ {\prime} \in \mathcal {A}} \pi \left(a ^ {\prime} \mid s\right) \hat {Q} (s, a ^ {\prime}) \tag {87}
$$
E.2.3 Computing expected policy gradients
Finally, we show how we can compute the expected policy gradient,  $\mathbb{E}_{\pi}\big[\hat{\nabla}_{\theta}V^{\pi}\big]$  of a potentially biased gradient estimator, given the expected advantage  $\mathbb{E}_{\pi}\big[\hat{A}^{\pi}(s,a)\big]$ .
$$
\begin{array}{l} \mathbb {E} _ {\pi} \left[ \hat {\nabla} _ {\theta} V ^ {\pi} \right] = \sum_ {k = 0} ^ {T} \mathbb {E} _ {S _ {k} \sim \mathcal {T} (s _ {0}, \pi)} \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a | s) \mathbb {E} _ {\pi} \left[ \hat {A} ^ {\pi} \left(S _ {k}, a\right) \right] (88) \\ = \sum_ {k = 0} ^ {T} \sum_ {s \in \mathcal {S}} p ^ {\pi} \left(S _ {k} = s \mid s _ {0}\right) \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid s) \mathbb {E} _ {\pi} \left[ \hat {A} ^ {\pi} (s, a) \right] (89) \\ \end{array}
$$
Using automatic differentiation to obtain  $\nabla_{\theta}\pi (a|s_0)$ , and the quantity  $M(s,a,s^{\prime},T)$  we defined above, we can then compute the expected policy gradient.