SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Our work builds upon Hindsight Credit Assignment (HCA) [1] which has sparked a number of follow-up studies. We generalize the theory of HCA towards estimating contributions upon rewarding outcomes instead of rewarding states and show through a detailed variance analysis that HCA suffers from spurious contributions leading to high variance, while using rewards or rewarding objects as rewarding outcome encodings leads to low-variance gradient estimators. Follow-up work on HCA discusses its potential high variance in constructed toy settings, and reduces the variance of the HCA advantage estimates by combining it with Monte Carlo estimates [45] or using temporal difference errors instead of rewards [46]. Alipov et al. [29] leverages the latter approach to scale up HCA towards more complex environments. However, all of the above approaches still suffer from spurious contributions, a significant source of variance in the HCA gradient estimator. In addition, recent studies have theoretically reinterpreted the original HCA formulation from different angles: Ma and Pierre-Luc [47] create the link to Conditional Monte Carlo methods [48, 49] and Arumugam et al. [50] provide an information theoretic perspective on credit assignment. Moreover, Young [51] applies HCA to the problem of estimating gradients in neural networks with stochastic, discrete units.
Long-term credit assignment in RL is hard due to the high variance of the sum of future rewards for long trajectories. A common technique to reduce the resulting variance in the policy gradients is to subtract a baseline from the sum of future rewards [6-9]. To further reduce the variance, a line of work introduced state-action-dependent baselines [52-55]. However, Tucker et al. [56] argues that these methods offer only a small benefit over the conventional state-dependent baselines, as the current action often only accounts for a minor fraction of the total variance. More recent work proposes improved baselines by incorporating hindsight information about the future trajectory into the baseline, accounting for a larger portion of the variance [1, 3, 34, 57-60]. Cheng et al. [34] includes Q-value estimates of future state-action pairs into the baseline, and Huang and Jiang [60] goes one step further by also including cheap estimates of the future Q-value gradients.Mesnard et al. [3] learn a summary metric of the uncontrollable, external environment influences in the future trajectory, and provide this hindsight information as an extra input to the value baseline. Nota et al. [57] consider partially observable MDPs, and leverage the future trajectory to more accurately infer the current underlying Markov state, thereby providing a better value baseline. Harutyunyan et al. [1] propose return-HCA, a different variant of HCA that uses a return-conditioned hindsight distribution to construct a baseline, instead of using state-based hindsight distributions for estimating contributions. Finally, Guez et al. [58] and Venuto et al. [59] learn a summary representation of the full future trajectory and provide it as input to the value function, while imposing an information bottleneck to prevent the value function from overly relying on this hindsight information.
Environments with sparse rewards and long delays between actions and corresponding rewards put a high importance on long-term credit assignment. A popular strategy to circumvent the sparse and delayed reward setting is to introduce reward shaping [61-68]. These approaches add auxiliary rewards to the sparse reward function, aiming to guide the learning of the policy with dense rewards. A recent line of work introduces a reward-shaping strategy specifically designed for long-term credit assignment, where rewards are decomposed and distributed over previous state-action pairs that were instrumental in achieving that reward [4, 5, 32, 35-41]. To determine how to redistribute rewards,
these approaches rely on heuristic contribution analyses, such as via the access of memory states [4], linear decompositions of rewards [32, 35-39] or learned sequence models [5, 40, 41]. Leveraging our unbiased contribution analysis framework to reach more optimal reward transport is a promising direction for future research.
When we have access to a (learned) differentiable world model of the environment, we can achieve precise credit assignment by leveraging path-wise derivatives, i.e. backpropagating value gradients through the world model [20-24, 69]. For stochastic world models, we need access to the noise variables to compute the path-wise derivatives. The Dreamer algorithms [21-23] approach this by computing the value gradients on simulated trajectories, where the noise is known. The Stochastic Value Gradient (SVG) method [20] instead infers the noise variables on real observed trajectories. To enable backpropagating gradients over long time spans, Ma et al. [70] equip the learned recurrent world models of SVG with an attention mechanism, allowing the authors to leverage Sparse Attentive Backtracking [71] to transmit gradients through skip connections. Buesing et al. [43] leverages the insights from SVG in partially observable MDPs, using the inferred noise variables to estimate the effect of counterfactual policies on the expected return. Importantly, the path-wise derivatives leveraged by the above model-based credit assignment methods are not compatible with discrete action spaces, as sensitivities w.r.t. discrete actions are undefined. In contrast, COCOA can leverage model-based information for credit assignment, while being compatible with discrete actions.
Incorporating hindsight information has a wide variety of applications in RL. Goal-conditioned policies [72] use a goal state as additional input to the policy network or value function, thereby generalizing it to arbitrary goals. Hindsight Experience Replay [73] leverages hindsight reasoning to learn almost as much from undesired outcomes as from desired ones, as at hindsight, we can consider every final, possibly undesired state as a 'goal' state and update the value functions or policy network [74] accordingly. Goyal et al. [75] train an inverse environment model to simulate alternative past trajectories leading to the same rewarding state, hence leveraging hindsight reasoning to create a variety of highly rewarding trajectories. A recent line of work frames RL as a supervised sequence prediction problem, learning a policy conditioned on goal states or future returns [76-79]. These models are trained on past trajectories or offline data, where we have in hindsight access to states and returns, considering them as targets for the learned policy.
Finally, Temporal Difference (TD) learning [80] has a rich history of leveraging proximity in time as a proxy for credit assignment [2]. TD( $\lambda$ ) [80] considers eligibility traces to trade off bias and variance, crediting past state-action pairs in the trajectory for the current reward proportional to how close they are in time. van Hasselt et al. [81] estimate expected eligibility traces, taking into account that the same rewarding state can be reached from various previous states. Hence, not only the past state-actions on the trajectory are credited and updated, but also counterfactual ones that lead to the same rewarding state. Extending the insights from COCOA towards temporal difference methods is an exciting direction for future research.
B Undiscounted infinite-horizon MDPs
In this work, we consider an undiscounted MDP with a finite state space  $S$ , bounded rewards and an infinite horizon. To ensure that the expected return and value functions remain finite, we require some standard regularity conditions [2]. We assume the MDP contains an absorbing state  $s_{\infty}$  that transitions only to itself and has zero reward. Moreover, we assume proper transition dynamics, meaning that an agent following any policy will eventually end up in the absorbing state  $s_{\infty}$  with probability one as time goes to infinity.
The discounted infinite horizon MDP formulation is a special case of this setting, with a specific class of transition dynamics. An explicit discounting of future rewards  $\sum_{k\geq o}\gamma^k R_k$  with discount factor  $\gamma \in [0,1]$ , is equivalent to considering the above undiscounted MDP setting, but modifying the state transition probability function  $p(S_{t + 1}\mid S_t,A_t)$  such that each state-action pair has a fixed probability  $(1 - \gamma)$  of transitioning to the absorbing state [2]. Hence, all the results considered in this work can be readily applied to the discounted MDP setting, by modifying the environment transitions as outlined above. We can also explicitly incorporate temporal discounting in the policy gradients and contribution coefficients, which we outline in App. J.
The undiscounted infinite-horizon MDP with an absorbing state can also model episodic RL with a fixed time horizon. We can include time as an extra feature in the states  $S$ , and then have a probability of 1 to transition to the absorbing state when the agent reaches the final time step.
Figure 6: (a) Graphical model of the MDP. (b) Graphical model of the MDP, where we abstracted time.
C Theorems, proofs and additional information for Section 3
C.1 Contribution coefficients, hindsight distribution and graphical models
Here, we provide more information on the derivation of the contribution coefficients of Eq. 3.
Abstracting time. In an undiscounted environment, it does not matter at which point in time the agent achieves the rewarding outcome  $u'$ . Hence, the contribution coefficients (3) sum over all future time steps, to reflect that the rewarding outcome  $u'$  can be encountered at any future time, and to incorporate the possibility of encountering  $u'$  multiple times. Note that when using temporal discounting, we can adjust the contribution coefficients accordingly (c.f. App J).
Hindsight distribution. To obtain the hindsight distribution  $p(A_{t} = a \mid S_{t} = s, U' = u')$ , we convert the classical graphical model of an MDP (c.f. Fig. 6a) into a graphical model that incorporates the time  $k$  into a separate node (c.f. Fig. 6b). Here, we rewrite  $p(U_{k} = u' \mid S = s, A = a)$ , the probability distribution of a rewarding outcome  $U_{k}$ $k$  time steps later, as  $p(U' = u' \mid S = s, A = a, K = k)$ . By giving  $K$  the geometric distribution  $p(K = k) = (1 - \beta)\beta^{k-1}$  for some  $\beta$ , we can rewrite the infinite sums used in the time-independent contribution coefficients as a marginalization over  $K$ :
$$
\begin{array}{l} \sum_ {k \geq 1} p \left(U _ {k} = u ^ {\prime} \mid S = s, A = a\right) = \lim  _ {\beta \rightarrow 1} \sum_ {k \geq 1} p \left(U _ {k} = u ^ {\prime} \mid S = s, A = a\right) \beta^ {k - 1} (5) \\ = \lim  _ {\beta \rightarrow 1} \frac {1}{1 - \beta} \sum_ {k \geq 1} p \left(U ^ {\prime} = u ^ {\prime} \mid S = s, A = a, K = k\right) p (K = k) (6) \\ = \lim  _ {\beta \rightarrow 1} \frac {1}{1 - \beta} p \left(U ^ {\prime} = u ^ {\prime} \mid S = s, A = a\right) (7) \\ \end{array}
$$
Note that this limit is finite, as for  $k \to \infty$ , the probability of reaching an absorbing state  $s_{\infty}$  and corresponding rewarding outcome  $u_{\infty}$  goes to 1 (and we take  $u' \neq u_{\infty}$ ). Via Bayes rule, we then have that
$$
\begin{array}{l} w (s, a, u ^ {\prime}) = \frac {\sum_ {k \geq 1} p \left(U _ {k} = u ^ {\prime} \mid S = s , A = a\right)}{\sum_ {k \geq 1} p \left(U _ {k} = u ^ {\prime} \mid S = s\right)} - 1 (8) \\ = \frac {p \left(U ^ {\prime} = u ^ {\prime} \mid S = s , A = a\right)}{p \left(U ^ {\prime} = u \mid S = s\right)} - 1 = \frac {p (A = a \mid S = s , U ^ {\prime} = u ^ {\prime})}{\pi (a \mid s)} - 1 (9) \\ \end{array}
$$
where we assume  $\pi(a \mid s) > 0$  and where we dropped the limit of  $\beta \to 1$  in the notation, as we will always take this limit henceforth. Note that when  $\pi(a \mid s) = 0$ , the hindsight distribution  $p^{\pi}(a \mid s, u')$  is also equal to zero, and hence the right-hand-side of the above equation is undefined. However, the middle term using  $p^{\pi}(U = u' \mid S = s, A = a)$  is still well-defined, and hence we can use the contribution coefficients even for actions where  $\pi(a \mid s) = 0$ .
C.2 Proof Theorem 1
We start with a lemma showing that the contribution coefficients can be used to estimate the advantage  $A^{\pi}(s,a) = Q^{\pi}(s,a) - V^{\pi}(s)$ , by generalizing Theorem 1 of Harutyunyan et al. [1] towards rewarding outcomes.
Definition 3 (Fully predictive, repeated from main text). A rewarding outcome encoding  $U$  is fully predictive of the reward  $R$ , if the following conditional independence condition holds:  $p^{\pi}(R_k = r|$
$S_0 = s, A_0 = a, U_k = u) = p^\pi (R = r\mid U = u)$ , where the right-hand side does not depend on the time  $k$ .
Lemma 6. For each state-action pair  $(s, a)$  with  $\pi(a \mid s) > 0$ , and assuming that  $u = f(s, a, r)$  is fully predictive of the reward (c.f. Definition 3), we have that
$$
A ^ {\pi} (s, a) = r (s, a) - \sum_ {a ^ {\prime} \in \mathcal {A}} \pi (a \mid s) r (s, a) + \mathbb {E} _ {T \sim \mathcal {T} (s, \pi)} \left[ \sum_ {k \geq 1} w (s, a, U _ {k}) R _ {k} \right] \tag {10}
$$
with the advantage function  $A^{\pi}$ , and the reward function  $r(s,a)\triangleq \mathbb{E}[R\mid s,a]$ .
Proof. We rewrite the undiscounted state-action value function in the limit of a discounting factor  $\beta \rightarrow 1_{-}$  (the minus sign indicating that we approach 1 from the left):
$$
\begin{array}{l} Q (s, a) = \lim  _ {\beta \rightarrow 1 -} \mathbb {E} _ {T \sim \mathcal {T} (s, a, \pi)} \left[ \sum_ {k \geq 1} \beta^ {k} R _ {k} \right] (11) \\ = r (s, a) + \lim  _ {\beta \rightarrow 1 _ {-}} \sum_ {r \in \mathcal {R}} \sum_ {k \geq 1} \beta^ {k} p ^ {\pi} \left(R _ {k} = r \mid s, a\right) r (12) \\ = r (s, a) + \lim  _ {\beta \rightarrow 1 -} \sum_ {r \in \mathcal {R}} \sum_ {u \in \mathcal {U}} \sum_ {k \geq 1} \beta^ {k} p ^ {\pi} \left(R _ {k} = r, U _ {k} = u \mid s, a\right) r (13) \\ = r (s, a) + \lim  _ {\beta \rightarrow 1 _ {-}} \sum_ {r \in \mathcal {R}} \sum_ {u \in \mathcal {U}} \sum_ {k \geq 1} \beta^ {k} p ^ {\pi} (R = r \mid U = u) r p ^ {\pi} \left(U _ {k} = u \mid s, a\right) (14) \\ = r (s, a) + \lim  _ {\beta \rightarrow 1 -} \sum_ {u \in \mathcal {U}} r (u) \sum_ {k \geq 1} \beta^ {k} p ^ {\pi} \left(U _ {k} = u \mid s, a\right) (15) \\ = r (s, a) + \lim  _ {\beta \rightarrow 1 -} \sum_ {u \in \mathcal {U}} r (u) \sum_ {k \geq 1} \beta^ {k} p ^ {\pi} \left(U _ {k} = u \mid s\right) \frac {\sum_ {k ^ {\prime} \geq 1} \beta^ {k ^ {\prime}} p ^ {\pi} \left(U _ {k ^ {\prime}} = u \mid s , a\right)}{\sum_ {k ^ {\prime} \geq 1} \beta^ {k ^ {\prime}} p ^ {\pi} \left(U _ {k ^ {\prime}} = u \mid s\right)} (16) \\ \end{array}
$$
where we use the property that  $u$  is fully predictive of the reward (c.f. Definition 2), and where we define  $r(s,a) \triangleq \mathbb{E}[R \mid s,a]$  and  $r(u) \triangleq \mathbb{E}[R \mid u]$ . Using the graphical model of Fig. 6b where we abstract time  $k$  (c.f. App. C.1), we have that  $\frac{1}{1 - \beta} p_{\beta}^{\pi}(R = r' \mid s,a) = \sum_{k=0} \beta^{k} p^{\pi}(R_{k} = r \mid s,a)$ . Leveraging Bayes rule, we get
$$
p ^ {\pi} (a \mid s, U ^ {\prime} = u) = \lim  _ {\beta \rightarrow 1 _ {-}} \frac {p _ {\beta} ^ {\pi} \left(U ^ {\prime} = u \mid s , a\right) \pi (a \mid s)}{p _ {\beta} ^ {\pi} \left(U ^ {\prime} = u \mid s\right)} = \lim  _ {\beta \rightarrow 1 _ {-}} \frac {\sum_ {k \geq 1} \beta^ {k} p ^ {\pi} \left(U _ {k} = u \mid s , a\right)}{\sum_ {k \geq 1} \beta^ {k} p ^ {\pi} \left(U _ {k} = u \mid s\right)} \pi (a \mid s) \tag {17}
$$
Where we dropped the limit of  $\beta \to 1_{-}$  in the notation of  $p^\pi (a\mid s,R = r)$ , as we henceforth always consider this limit. Taking everything together, we have that
$$
\begin{array}{l} Q ^ {\pi} (s, a) = r (s, a) + \sum_ {u \in \mathcal {U}} \sum_ {k \geq 1} p ^ {\pi} \left(U _ {k} = u \mid s\right) \frac {p ^ {\pi} (a \mid s , U ^ {\prime} = u)}{\pi (a \mid s)} r (u) (18) \\ = r (s, a) + \sum_ {u \in \mathcal {U}} \sum_ {k \geq 1} p ^ {\pi} \left(U _ {k} = u \mid s\right) \left(w (s, a, u) + 1\right) r (u) (19) \\ = r (s, a) + \sum_ {u \in \mathcal {U}} \sum_ {k \geq 1} p ^ {\pi} \left(U _ {k} = u \mid s\right) \left(w (s, a, u) + 1\right) \sum_ {r \in \mathcal {R}} p \left(R _ {k} = r \mid U _ {k} = u\right) r (20) \\ = r (s, a) + \mathbb {E} _ {T \sim \mathcal {T} (s, \pi)} \left[ \sum_ {k \geq 1} \left(w (s, a, U _ {k}) + 1\right) R _ {k} \right] (21) \\ \end{array}
$$
Subtracting the value function, we get
$$
A ^ {\pi} (s, a) = 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 \geq 1} w (s, a, U _ {k}) R _ {k} \right] \tag {22}
$$
Now we are ready to prove Theorem 1.
Proof. Using the policy gradient theorem [26], we have
$$
\begin{array}{l} \nabla_ {\theta} V ^ {\pi} \left(s _ {0}\right) = \mathbb {E} _ {T \sim \mathcal {T} \left(s _ {0}, \pi\right)} \left[ \sum_ {t \geq 0} \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid S _ {t}) A ^ {\pi} \left(S _ {t}, a\right) \right] (23) \\ = \mathbb {E} _ {T \sim \mathcal {T} \left(s _ {0}, \pi\right)} \left[ \sum_ {t \geq 0} \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid S _ {t}) \left[ r \left(S _ {t}, a\right) + \sum_ {k \geq 1} w \left(S _ {t}, a, U _ {t + k}\right) R _ {t + k} \right] \right] (24) \\ = \mathbb {E} _ {T \sim \mathcal {T} \left(s _ {0}, \pi\right)} \left[ \sum_ {t \geq 0} \nabla_ {\theta} \log \pi \left(A _ {t} \mid S _ {t}\right) R _ {t} + \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid S _ {t}) \sum_ {k \geq 1} w \left(S _ {t}, a, U _ {t + k}\right) R _ {t + k} \right] (25) \\ \end{array}
$$
where we used that removing a baseline  $\sum_{a' \in \mathcal{A}} \pi(a' \mid s)r(s, a')$  independent from the actions does not change the policy gradient, and replaced  $\mathbb{E}_{T \sim \mathcal{T}(s_0, \pi)}[\sum_{a \in \mathcal{A}} \nabla_\theta \pi(a \mid S_t)r(S_t, a)]$  by its sampled version  $\mathbb{E}_{T \sim \mathcal{T}(s_0, \pi)}[\nabla_\theta \log \pi(A_t \mid S_t)R_t]$ . Hence the policy gradient estimator of Eq. 4 is unbiased.
C.3 Different policy gradient estimators leveraging contribution coefficients.
In this work, we use the COCOA estimator of Eq. 4 to estimate policy gradients leveraging the contribution coefficients of Eq. 3, as this estimator does not need a separate reward model, and it works well for the small action spaces we considered in our experiments. However, one can design other unbiased policy gradient estimators compatible with the same contribution coefficients.
Harutyunyan et al. [1] introduced the HCA gradient estimator of Eq. 2, which can readily be extended to rewarding outcomes  $U$ :
$$
\hat {\nabla} _ {\theta} V ^ {\pi} = \sum_ {t \geq 0} \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid S _ {t}) \left(r \left(S _ {t}, a\right) + \sum_ {k \geq 1} w \left(S _ {t}, a, U _ {t + k}\right) R _ {t + k}\right) \tag {26}
$$
This gradient estimator uses a reward model  $r(s, a)$  to obtain the rewards corresponding to counterfactual actions, whereas the COCOA estimator (4) only uses the observed rewards  $R_{k}$ .
For large action spaces, it might become computationally intractable to sum over all possible actions. In this case, we can sample independent actions from the policy, instead of summing over all actions, leading to the following policy gradient.
$$
\hat {\nabla} _ {\theta} V ^ {\pi} = \sum_ {t \geq 0} \nabla_ {\theta} \log \pi \left(A _ {t} \mid S _ {t}\right) R _ {t} + \frac {1}{M} \sum_ {m} \nabla_ {\theta} \log \pi \left(a ^ {m} \mid S _ {t}\right) \sum_ {k = 1} ^ {\infty} w \left(S _ {t}, a ^ {m}, U _ {t + k}\right) R _ {t + k} \tag {27}
$$
where we sample  $M$  actions from  $a^m \sim \pi(\cdot \mid S_t)$ . Importantly, for obtaining an unbiased policy gradient estimate, the actions  $a^m$  should be sampled independently from the actions used in the trajectory of the observed rewards  $R_{t+k}$ . Hence, we cannot take the observed  $A_t$  as a sample  $a^m$ , but need to instead sample independently from the policy. This policy gradient estimator can be used for continuous action spaces (c.f. App. G, and can also be combined with the reward model as in Eq. 26.
One can show that the above policy gradient estimators are unbiased with a proof akin to the one of Theorem 1.
Comparison of COCOA to using time as a credit assignment heuristic. In Section 2 we discussed briefly the following discounted policy gradient:
$$
\hat {\nabla} _ {\theta} ^ {\text {R E I N F O R C E}, \gamma} V ^ {\pi} \left(s _ {0}\right) = \sum_ {t \geq 0} \nabla_ {\theta} \log \pi \left(A _ {t} \mid S _ {t}\right) \sum_ {k \geq t} \gamma^ {k - t} R _ {k} \tag {28}
$$
with discount factor  $\gamma \in [0,1]$ . Note that we specifically put no discounting  $\gamma^t$  in front of  $\nabla_{\theta}\log \pi (A_t\mid S_t)$ , which would be required for being an unbiased gradient estimate of the discounted expected total return, as the above formulation is most used in practice [12, 27]. Rewriting
the summations reveals that this policy gradient uses time as a heuristic for credit assignment:
$$
\hat {\nabla} _ {\theta} ^ {\text {R E I N F O R C E}, \gamma} V ^ {\pi} \left(s _ {0}\right) = \sum_ {t \geq 0} R _ {t} \sum_ {k \leq t} \gamma^ {t - k} \nabla_ {\theta} \log \pi \left(A _ {k} \mid S _ {k}\right). \tag {29}
$$
We can rewrite the COCOA gradient estimator (4) with the same reordering of summation, showcasing that it leverages the contribution coefficient for providing precise credit to past actions, instead of using the time discounting heuristic:
$$
\hat {\nabla} _ {\theta} ^ {U} V ^ {\pi} \left(s _ {0}\right) = \sum_ {t \geq 0} R _ {t} \left[ \nabla_ {\theta} \log \pi \left(A _ {t} \mid S _ {t}\right) + \sum_ {k <   t} \sum_ {a \in \mathcal {A}} w \left(S _ {k}, a, U _ {t}\right) \nabla_ {\theta} \pi (a \mid S _ {k}) \right] \tag {30}
$$
C.4 Proof of Proposition 2
Proof. Proposition 2 assumes that each action sequence  $\{A_m\}_{m = t}^{t + k}$  leads to a unique state  $s^\prime$ . Hence, all previous actions can be decoded perfectly from the state  $s^\prime$ , leading to  $p^{\pi}(A_{t} = a\mid S_{t},S^{\prime} = s^{\prime}) = \delta (a = a_{t})$ , with  $\delta$  the indicator function and  $a_{t}$  the action taken in the trajectory that led to  $s^\prime$ . Filling this into the COCOA gradient estimator leads to
$$
\begin{array}{l} \hat {\nabla} _ {\theta} ^ {S} V ^ {\pi} \left(s _ {0}\right) = \sum_ {t \geq 0} \nabla_ {\theta} \log \pi \left(A _ {t} \mid S _ {t}\right) R _ {t} + \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid S _ {t}) \sum_ {k \geq 1} \left(\frac {\delta \left(a = A _ {t}\right)}{\pi \left(a \mid S _ {t}\right)} - 1\right) R _ {t + k} (31) \\ = \sum_ {t \geq 0} \frac {\nabla_ {\theta} \pi \left(A _ {t} \mid S _ {t}\right)}{\pi \left(A _ {t} \mid S _ {t}\right)} \sum_ {k \geq 0} R _ {t + k} (32) \\ = \sum_ {t \geq 0} \nabla_ {\theta} \log \pi \left(A _ {t} \mid S _ {t}\right) \sum_ {k \geq 0} R _ {t + k} (33) \\ \end{array}
$$
where we used that  $\sum_{a}\nabla_{\theta}\pi (a\mid s) = 0$
C.5 Proof Theorem 3
Proof. Theorem 3 considers the case where the environment only contains a reward at the final time step  $t = T$ , and where we optimize the policy only on a single (initial) time step  $t = 0$ . Then, the policy gradients are given by
$$
\hat {\nabla} _ {\theta} ^ {U} V ^ {\pi} \left(U _ {T}, R _ {T}\right) = \sum_ {a} \nabla_ {\theta} \pi (a \mid s) w (s, a, U _ {T}) R _ {T} \tag {34}
$$
$$
\hat {\nabla} _ {\theta} ^ {\text {R E I N F O R C E}} V ^ {\pi} \left(A _ {0}, R _ {T}\right) = \nabla_ {\theta} \log \pi \left(A _ {0} \mid s\right) R _ {T} \tag {35}
$$
With  $U$  either  $S, R$  or  $Z$ , and  $s$  the state at  $t = 0$ . As only the last time step can have a reward by assumption, and the encoding  $U$  needs to retain the predictive information of the reward, the contribution coefficients for  $u$  corresponding to a nonzero reward are equal to
$$
w (s, a, u) = \frac {p ^ {\pi} \left(A _ {0} = a \mid S _ {0} = s , U _ {T} = u\right)}{\pi (a \mid s)} - 1 \tag {36}
$$