SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} w _ {\gamma} (s, a, u ^ {\prime}) = \frac {\sum_ {k \geq 1} \gamma^ {k} p ^ {\pi} \left(U _ {t + k} = u ^ {\prime} \mid S _ {t} = s , A _ {t} = a\right)}{\sum_ {k \geq 1} \gamma^ {k} p ^ {\pi} \left(U _ {t + k} = u ^ {\prime} \mid S _ {t} = s\right)} - 1 (120) \\ = \frac {p _ {\gamma} ^ {\pi} \left(A _ {t} = a \mid S _ {t} = s , U ^ {\prime} = u ^ {\prime}\right)}{\pi (a \mid s)} - 1 (121) \\ \end{array}
$$
Here, we define the discounted hindsight distribution  $p_{\gamma}^{\pi}(A_t = a \mid S_t = s, U' = u')$  similarly to the undiscounted hindsight distribution explained in App. C.1, but now using a different probability distribution on the time steps  $k$ :  $p_{\beta}(K = k) = (1 - \beta)\beta^{k-1}$ , where we take  $\beta = \gamma$ . We can readily extend Theorems 1-4 to the explicit discounted setting, by taking  $\beta = \gamma$  instead of the limit of  $\beta \to 1$ , and using the discounted COCOA policy gradient estimator shown in Table 6.
To approximate the discounted hindsight distribution  $p_{\gamma}^{\pi}(A_t = a \mid S_t = s, U' = u')$ , we need to incorporate the temporal discounting into the classification cross-entropy loss:
$$
L _ {\gamma} = \mathbb {E} _ {\pi} \left[ \sum_ {t \geq 0} \sum_ {k \geq 1} \gamma^ {k} C E \left(h _ {\gamma} (\cdot \mid S _ {t}, U _ {t + k}), \delta (a = A _ {t})\right) \right], \tag {122}
$$
with  $CE$  the cross-entropy loss,  $h_{\gamma}(\cdot \mid S_t,U_{t + k})$  the classification model that approximates the discounted hindsight distribution, and  $\delta (a = A_{t})$  a one-hot encoding of  $A_{t}$ .
K Bootstrapping with COCOA
Here, we show how COCOA can be combined with n-step returns, and we make a correction to Theorem 7 of Harutyunyan et al. [1] which considers n-step returns for HCA.
Consider the graphical model of Fig. 15a where we model time,  $K$ , as a separate node in the graphical model (c.f. App. C.1). To model  $n$ -step returns, we now define the following prior probability on  $K$
(b)
Figure 15: (a) Graphical model where we abstract time. (b) Graphical model implicitly used in the proof of Theorem 7 in Harutyunyan et al. [1].
parameterized by  $\beta$
$$
p _ {n, \beta} (K = k) = \left\{ \begin{array}{l l} \frac {\beta^ {k}}{z} & \text {i f} 1 \leq k \leq n - 1 \\ 0 & \text {e l s e} \end{array} \quad z = \beta \frac {1 - \beta^ {n - 1}}{1 - \beta} \right. \tag {123}
$$
Using  $p_{n,\beta}^{\pi}$  as the probability distribution induced by this graphical model, we have that
$$
\begin{array}{l} p _ {n, \beta} ^ {\pi} (U ^ {\prime} = u \mid S = s, A = a) = \sum_ {k} p _ {n, \beta} ^ {\pi} (U ^ {\prime} = u, K = k \mid S = s, A = a) (124) \\ = \sum_ {k} p _ {n, \beta} (K = k) p _ {n, \beta} ^ {\pi} \left(U ^ {\prime} = u \mid S = s, A = a, K = k\right) (125) \\ = \frac {1 - \beta}{\beta \left(1 - \beta^ {n - 1}\right)} \sum_ {k = 1} ^ {n - 1} \beta^ {k} p ^ {\pi} \left(U _ {k} = u \mid S _ {0} = s, A _ {0} = a\right) (126) \\ \end{array}
$$
We introduce the following contribution coefficients that we will use for n-step returns
$$
\begin{array}{l} w _ {n, \beta} (s, a, u ^ {\prime}) = \frac {\sum_ {k = 1} ^ {n - 1} \beta^ {k} p ^ {\pi} \left(U _ {k} = u ^ {\prime} \mid S _ {0} = s , A _ {0} = a\right)}{\sum_ {k = 1} ^ {n - 1} \beta^ {k} p ^ {\pi} \left(U _ {k} = u ^ {\prime} \mid S _ {0} = s\right)} (127) \\ = \frac {p _ {n , \beta} ^ {\pi} \left(U ^ {\prime} = u ^ {\prime} \mid S = s , A = a\right)}{p _ {n , \beta} ^ {\pi} \left(U ^ {\prime} = u ^ {\prime} \mid S = s\right)} - 1 = \frac {p _ {n , \beta} ^ {\pi} \left(A = a \mid S = s , U ^ {\prime} = u ^ {\prime}\right)}{\pi (a \mid s)} - 1 (128) \\ \end{array}
$$
$$
w _ {n} (s, a, s ^ {\prime}) = \frac {p ^ {\pi} \left(S _ {n} = s ^ {\prime} \mid S _ {0} = s , A _ {0} = a\right)}{p ^ {\pi} \left(S _ {n} = s ^ {\prime} \mid S _ {0} = s\right)} - 1 = \frac {p ^ {\pi} \left(A = a \mid S = s , S _ {n} = s ^ {\prime}\right)}{\pi (a \mid s)} - 1 \tag {129}
$$
Now we are ready to prove the n-step return theorem for the discounted MDP setting. We can recover the undiscounted setting by taking the limit of  $\gamma \rightarrow 1_{-}$ .
Theorem 10. Consider state  $s$  and action  $a$  for which it holds that  $\pi(a \mid s) > 0$  and take  $\beta$  equal to the discount factor  $\gamma \in [0,1]$ . Furthermore, assume that the rewarding outcome encoding  $u = f(s,a,r)$  is fully predictive of the reward (c.f. Definition 2). Then the advantage  $A^{\pi}(s,a) = Q^{\pi}(s,a) - V^{\pi}(s)$  is equal to
$$
A ^ {\pi} (s, a) = r (s, a) - r ^ {\pi} (s) + \mathbb {E} _ {\mathcal {T} (s, \pi)} \left[ \sum_ {k = 1} ^ {n - 1} \gamma^ {k} w _ {n, \beta} (s, a, U _ {k}) R _ {k} + \gamma^ {n} w _ {n} (s, a, S _ {n}) V ^ {\pi} (S _ {n}) \right]
$$
with  $r(s,a)$  the reward model and  $r^{\pi}(s) = \sum_{a}\pi (a\mid s)r(s,a)$
Proof. We start with the action-value function  $Q^{\pi}$ , and will subtract the value  $V^{\pi}$  to obtain the result on the advantage function.
$$
\begin{array}{l} Q (s, a) = \mathbb {E} _ {T \sim \mathcal {T} (s, a, \pi)} \left[ \sum_ {k \geq 1} \gamma^ {k} R _ {k} \right] (130) \\ = r (s, a) + \sum_ {r ^ {\prime} \in \mathcal {R}} \sum_ {k \geq 1} \gamma^ {k} p ^ {\pi} \left(R _ {k} = r ^ {\prime} \mid s, a\right) r ^ {\prime} (131) \\ \end{array}
$$
$$
\begin{array}{l} = r (s, a) + \sum_ {r ^ {\prime} \in \mathcal {R}} \sum_ {u ^ {\prime} \in \mathcal {U}} \sum_ {k = 1} ^ {n - 1} \gamma^ {k} p ^ {\pi} \left(R _ {k} = r ^ {\prime}, U _ {k} = u ^ {\prime} \mid s, a\right) r ^ {\prime} + (132) \\ \gamma^ {n} \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S _ {n} = s ^ {\prime} \mid s, a\right) V ^ {\pi} \left(s ^ {\prime}\right) (133) \\ = r (s, a) + \sum_ {r ^ {\prime} \in \mathcal {R}} \sum_ {u ^ {\prime} \in \mathcal {U}} \sum_ {k = 1} ^ {n - 1} \gamma^ {k} p ^ {\pi} \left(R ^ {\prime} = r ^ {\prime} \mid U ^ {\prime} = u ^ {\prime}\right) p ^ {\pi} \left(U _ {k} = u ^ {\prime} \mid s, a\right) r ^ {\prime} + (134) \\ \gamma^ {n} \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S _ {n} = s ^ {\prime} \mid s, a\right) V ^ {\pi} \left(s ^ {\prime}\right) (135) \\ = r (s, a) + \sum_ {u ^ {\prime} \in \mathcal {U}} r \left(u ^ {\prime}\right) \sum_ {k = 1} ^ {n - 1} \gamma^ {k} p ^ {\pi} \left(U _ {k} = u ^ {\prime} \mid s, a\right) + \gamma^ {n} \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S _ {n} = s ^ {\prime} \mid s, a\right) V ^ {\pi} \left(s ^ {\prime}\right) (136) \\ = r (s, a) + \sum_ {u ^ {\prime} \in \mathcal {U}} r \left(u ^ {\prime}\right) \sum_ {k = 1} ^ {n - 1} \gamma^ {k} p ^ {\pi} \left(U _ {k} = u ^ {\prime} \mid s\right) \frac {\sum_ {k ^ {\prime} = 1} ^ {n - 1} \gamma^ {k ^ {\prime}} p ^ {\pi} \left(U _ {k ^ {\prime}} = u \mid s , a\right)}{\sum_ {k ^ {\prime} = 1} ^ {n - 1} \gamma^ {k ^ {\prime}} p ^ {\pi} \left(U _ {k ^ {\prime}} = u \mid s\right)} + (137) \\ \gamma^ {n} \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S _ {n} = s ^ {\prime} \mid s\right) \frac {p ^ {\pi} \left(S _ {n} = s ^ {\prime} \mid s , a\right)}{p ^ {\pi} \left(S _ {n} = s ^ {\prime} \mid s\right)} V ^ {\pi} \left(s ^ {\prime}\right) (138) \\ = r (s, a) + \sum_ {u ^ {\prime} \in \mathcal {U}} r \left(u ^ {\prime}\right) \sum_ {k = 1} ^ {n - 1} \gamma^ {k} p ^ {\pi} \left(U _ {k} = u ^ {\prime} \mid s\right) \left(w _ {n, \beta} (s, a, u ^ {\prime}) + 1\right) (139) \\ + \gamma^ {n} \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S _ {n} = s ^ {\prime} \mid s\right) \left(w _ {n} \left(s, a, s ^ {\prime}\right) + 1\right) V ^ {\pi} \left(s ^ {\prime}\right) (140) \\ \end{array}
$$
where we use that  $U'$  is fully predictive of the reward  $R'$ , and define  $r(u') = \sum_{r' \in \mathcal{R}} p(R' = r'|U' = u') r'$ . By subtracting the value function, we get
$$
\begin{array}{l} A (s, a) = r (s, a) - r ^ {\pi} (s) + \sum_ {u ^ {\prime} \in \mathcal {U}} r \left(u ^ {\prime}\right) \sum_ {k = 1} ^ {n - 1} \gamma^ {k} p ^ {\pi} \left(U _ {k} = u ^ {\prime} \mid s\right) w _ {n, \beta} (s, a, u ^ {\prime}) (142) \\ + \gamma^ {n} \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S _ {n} = s ^ {\prime} \mid s\right) w _ {n} \left(s, a, s ^ {\prime}\right) V ^ {\pi} \left(s ^ {\prime}\right) (143) \\ = r (s, a) - r ^ {\pi} (s) + \mathbb {E} _ {\mathcal {T} (s, \pi)} \left[ \sum_ {k = 1} ^ {n - 1} \gamma^ {k} w _ {n, \beta} (s, a, U _ {k}) R _ {k} + \gamma^ {n} w _ {n} (s, a, S _ {n}) V ^ {\pi} \left(S _ {n}\right) \right] (144) \\ \end{array}
$$
Finally, we can sample from this advantage function to obtain an n-step COCOA gradient estimator, akin to Theorem 1.
Note that we require to learn the state-based contribution coefficients  $w_{n}(s,a,s^{\prime})$  to bootstrap the value function into the n-step return, as the value function requires a Markov state  $s^{\prime}$  as input instead of a rewarding outcome encoding  $u^{\prime}$ . Unfortunately, these state-based contribution coefficients will suffer from spurious contributions, akin to HCA, introducing a significant amount of variance into the n-step COCOA gradient estimator. We leave it to future research to investigate whether we can incorporate value functions into an n-step return, while using rewarding-outcome contribution coefficients  $w(s,a,u^{\prime})$  instead of state-based contribution coefficients  $w_{n}(s,a,s^{\prime})$ .
Learning the contribution coefficients. We can learn the contribution coefficients  $w_{\beta, n}(s, a, u')$  with the same strategies as described in Section 3, but now with training data from  $n$ -step trajectories instead of complete trajectories. If we use a discount  $\gamma \neq 1$ , we need to take this discount factor into account in the training distribution or loss function (c.f. App. J).
Correction to Theorem 7 of Harutyunyan et al. [1]. Harutyunyan et al. [1] propose a theorem similar to Theorem 10, with two important differences. The first one concerns the distribution on  $K$  in the graphical model of Fig. 15a. Harutyunyan et al. [1] implicitly use this graphical model, but with a different prior probability distribution on  $K$ :
$$
p _ {n, \beta} ^ {H C A} (K = k) = \left\{ \begin{array}{l l} \beta^ {k - 1} (1 - \beta) & \text {i f} 1 \leq k \leq n - 1 \\ \beta^ {n - 1} & \text {i f} k = n \\ 0 & \text {e l s e} \end{array} \right. \tag {145}
$$
The graphical model combined with the distribution on  $K$  defines the hindsight distribution  $p_{n,\beta ,HCA}^{\pi}(A = a\mid S = s,S' = s')$ . The second difference is the specific  $Q$ -value estimator Harutyunyan et al. [1] propose. They use the hindsight distribution  $p_{n,\beta ,HCA}^{\pi}(A = a\mid S = s,S' = s')$  in front of the value function (c.f. Theorem 10), which considers that  $s^\prime$  can be reached at any time step  $k\sim p_{n,\beta}^{HCA}(k)$ , whereas Theorem 10 uses  $w_{n}(s,a,s^{\prime})$  which considers that  $s^\prime$  is reached exactly at time step  $k = n$ .
To the best of our knowledge, there is an error in the proposed proof of Theorem 7 by Harutyunyan et al. [1] for which we could not find a simple fix. For the interested reader, we briefly explain the error. One indication of the problem is that for  $\beta \to 1$ , all the probability mass of  $p_{n,\beta}^{HCA}(K = k)$  is concentrated at  $k = n$ , hence the corresponding hindsight distribution  $p_{n,\beta,HCA}^{\pi}(A = a \mid S = s, S' = s')$  considers only hindsight states  $s'$  encountered at time  $k = n$ . While this is not a mathematical error, it does not correspond to the intuition of a 'time independent hindsight distribution' the authors provide. In the proof itself, a conditional independence relation is assumed that does not hold. The authors introduce a helper variable  $Z$  defined on the state space  $S$ , with a conditional distribution
$$
\mu_ {k} (Z = z \mid S ^ {\prime} = s ^ {\prime}) = \left\{ \begin{array}{l l} \delta \left(z = s ^ {\prime}\right) & \text {i f} 1 \leq k \leq n - 1 \\ \tilde {d} ^ {\pi} (z \mid s ^ {\prime}) & \text {i f} k = n \end{array} \right. \tag {146}
$$
with the normalized discounted visit distribution  $\tilde{d}^{\pi}(z\mid s^{\prime}) = (1 - \gamma)\sum_{k}\gamma^{k}p^{\pi}(S_{k} = z\mid S_{0} = s)$ . We can model this setting as the graphical model visualized in Fig. 15b. In the proof (last line on page 15 in the supplementary materials of Harutyunyan et al. [1]), the following conditional independence is used:
$$
p ^ {\pi} \left(A _ {0} = a \mid S _ {0} = s, S ^ {\prime} = s ^ {\prime}, Z = z\right) = p ^ {\pi} \left(A _ {0} = a \mid S _ {0} = s, S ^ {\prime} = s ^ {\prime}\right) \tag {147}
$$
However, Fig. 15b shows that  $S'$  is a collider on the path  $A_0 \to S' \gets K \to Z$ . Hence, by conditioning on  $S'$  we open this collider path, making  $A_0$  dependent on  $Z$  conditioned on  $S_0$  and  $S'$ , thereby invalidating the assumed conditional independence. For example, if  $Z$  is different from  $S'$ , we know that  $K = n$  (c.f. Eq. 146), hence  $Z$  can contain information about action  $A_0$ , beyond  $S'$ , as  $S'$  ignores at which point in time  $s'$  is encountered.
L HCA-return is a biased estimator in many relevant environments
L.1 HCA-return
Besides HCA-state, Harutyunyan et al. [1] introduced HCA-return, a policy gradient estimator that leverages the hindsight distribution conditioned on the return:
$$
\sum_ {t \geq 0} \nabla_ {\theta} \log \pi \left(A _ {t} \mid S _ {t}\right) \left(1 - \frac {\pi \left(A _ {t} \mid S _ {t}\right)}{p ^ {\pi} \left(A _ {t} \mid S _ {t} , Z _ {t}\right)}\right) Z _ {t} \tag {148}
$$
When comparing this estimator with COCOA-reward, we see two important differences: (i) HCA-return uses a hindsight function conditioned on the return instead of individual rewards, and (ii) HCA-return leverages the hindsight function as an action-dependent baseline for a Monte Carlo policy gradient estimate, instead of using it for contribution coefficients to evaluate counterfactual actions. Importantly, the latter difference causes the HCA-return estimator to be biased in many environments of relevance, even when using the ground-truth hindsight distribution.
L.2 HCA-return can be biased
An important drawback of HCA-return is that it can be biased, even when using the ground-truth hindsight distribution. Theorem 2 of Harutyunyan et al. 2019, considering the unbiasedness HCA-return, is valid under the assumption that for any possible random return  $Z$  for all possible trajectories
starting from state  $s$ , it holds that  $p^{\pi}(a \mid s, z) > 0$ . This restrictive assumption requires that for each observed state-action pair  $(s_t, a_t)$  along a trajectory, all counterfactual returns  $Z$  resulting from a counterfactual trajectory starting from  $s_t$  (not including  $a_t$ ) result in  $p^{\pi}(a_t \mid s_t, Z) > 0$ . This implies that all returns (or rewarding states) reachable from  $s_t$  should also be reachable from  $(s_t, a_t)$ .
Consider the following bandit setting as a simple example where the above assumption is not satisfied. The bandit has two arms, with a reward of 1 and  $-2$ , and a policy probability of  $\frac{2}{3}$  and  $\frac{1}{3}$  respectively. The advantage for both arms is 1 and  $-2$ . Applying eq. 6 from Harutyunyan et al. results in  $A^{\pi}(s,a_1) = (1 - \frac{2}{3}) = \frac{1}{3}$  and  $A^{\pi}(s,a_2) = -2(1 - 1/3) = -4/3$ . This shows that the needed assumptions for an unbiased HCA-return estimator can be violated even in simple bandit settings.
M Additional details
M.1 Author contributions
This paper was a collaborative effort of all shared first authors working closely together. To do this fact better justice we give an idea of individual contributions in the following.
Alexander Meulemans*. Original idea, conceptualizing the theory and proving the theorems, conceptual development of the algorithms, experiment design, implementation of main method and environments, debugging, neural network architecture design, running experiments, connecting the project to existing literature, writing of manuscript, first draft and supplementary materials, feedback to the figures.
Simon Schug*. Conceptual development of the algorithms, experiment design, implementation of main method, baselines and environments, neural network architecture design, debugging, tuning and running experiments, writing of manuscript, creation of figures, writing of supplementary materials.
Seijin Kobayashi*. Conceptual development of the algorithms, experiment design, implementation of environments, baselines, main method and Dynamic Programming-based ground-truth methods, debugging, tuning and running experiments, feedback to the manuscript, writing of supplementary materials.
Nathaniel Daw. Regular project meetings, conceptual input and feedback for method and experimental design, connecting the project to existing literature, feedback to the manuscript and figures.
Gregory Wayne. Senior project supervision, conceptualising of the project idea, conceptual development of the algorithms, regular project meetings, technical and conceptual feedback for method and experimental design, connecting the project to existing literature, feedback to the manuscript and figures.
M.2 Compute resources
We used Linux workstations with Nvidia RTX 2080 and Nvidia RTX 3090 GPUs for development and conducted hyperparameter searches and experiments using 5 TPUv2-8, 5 TPUv3-8 and 1 Linux server with 8 Nvidia RTX 3090 GPUs over the course of 9 months. All of the final experiments presented take less than a few hours to complete using a single Nvidia RTX 3090 GPU. In total, we spent an estimated amount of 2 GPU months.
M.3 Software and libraries
For the results produced in this paper we relied on free and open-source software. We implemented our experiments in Python using JAX [95, Apache License 2.0] and the Deepmind Jax Ecosystem [82, Apache License 2.0]. For experiment tracking we used wandb [96, MIT license] and for the generation of plots we used plotly [97, MIT license].