SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alexander Meulemans\*1, Simon Schug\*1, Seijin Kobayashi\*1 Nathaniel D Daw2,3,4, Gregory Wayne2
$^{1}$ Department of Computer Science, ETH Zürich  
 $^{2}$ Google DeepMind
3Princeton Neuroscience Institute, Princeton University
$^{4}$ Department of Psychology, Princeton University {ameulema, sschug, seijink}@ethz.ch
To make reinforcement learning more sample efficient, we need better credit assignment methods that measure an action's influence on future rewards. Building upon Hindsight Credit Assignment (HCA) [1], we introduce Counterfactual Contribution Analysis (COCOA), a new family of model-based credit assignment algorithms. Our algorithms achieve precise credit assignment by measuring the contribution of actions upon obtaining subsequent rewards, by quantifying a counterfactual query: 'Would the agent still have reached this reward if it had taken another action?' We show that measuring contributions w.r.t. rewarding states, as is done in HCA, results in spurious estimates of contributions, causing HCA to degrade towards the high-variance REINFORCE estimator in many relevant environments. Instead, we measure contributions w.r.t. rewards or learned representations of the rewarding objects, resulting in gradient estimates with lower variance. We run experiments on a suite of problems specifically designed to evaluate long-term credit assignment capabilities. By using dynamic programming, we measure ground-truth policy gradients and show that the improved performance of our new model-based credit assignment methods is due to lower bias and variance compared to HCA and common baselines. Our results demonstrate how modeling action contributions towards rewarding outcomes can be leveraged for credit assignment, opening a new path towards sample-efficient reinforcement learning.[2]
Reinforcement learning (RL) faces two central challenges: exploration and credit assignment [2]. We need to explore to discover rewards and we need to reinforce the actions that are instrumental for obtaining these rewards. Here, we focus on the credit assignment problem and the intimately linked problem of estimating policy gradients. For long time horizons, obtaining the latter is notoriously difficult as it requires measuring how each action influences expected subsequent rewards. As the number of possible trajectories grows exponentially with time, future rewards come with a considerable variance stemming from stochasticity in the environment itself and from the stochasticity of interdependent future actions leading to vastly different returns [3-5].
Monte Carlo estimators such as REINFORCE [6] therefore suffer from high variance, even after variance reduction techniques like subtracting a baseline [6-9]. Similarly, in Temporal Difference methods such as Q-learning, this high variance in future rewards results in a high bias in the value estimates, requiring exponentially many updates to correct for it [5]. Thus, a common technique to
*Equal contribution; ordering determined by coin flip.
$^{2}$ Code available at https://github.com/seijin-kobayashi/cocoa
Figure 1: Counterfactual Contribution Analysis enables long-term credit assignment. (A) Given a sample trajectory that eventually results in a rewarding outcome, we estimate the policy gradient by considering the contribution of actions along the trajectory towards arriving at a rewarding outcome. In this example, we measure how much more likely the rewarding outcome with reward  $r_3$  is when following action  $a_1$  versus the counterfactual actions  $a_2$  and  $a_3$  in state  $s_2$ . This is quantified through the contribution coefficient  $w(s_2, a_1, r_3)$  which is used to update all possible action probabilities of the policy  $\pi(a \mid s_2)$ . (B) In the linear key-to-door environment increasing the distance between picking up the key and opening the door that leads to reward necessitates credit assignment over increasing time spans. COCOA consistently achieves good performance (left) compared to HCA and baselines which deteriorate when increasing the distance between an action and the resulting rewarding outcome. This is reflected in a higher signal-to-noise ratio of the policy gradient estimator of COCOA compared to baselines (right).
reduce variance and bias is to discount rewards that are far away in time resulting in a biased estimator which ignores long-term dependencies [10-13]. Impressive results have nevertheless been achieved in complex environments [14-16] at the cost of requiring billions of environment interactions, making these approaches sample inefficient.
Especially in settings where obtaining such large quantities of data is costly or simply not possible, model-based RL that aims to simulate the dynamics of the environment is a promising alternative. While learning such a world model is a difficult problem by itself, when successful it can be used to generate a large quantity of synthetic environment interactions. Typically, this synthetic data is combined with model-free methods to improve the action policy [17-19]. A notable exception to simply using world models to generate more data are the Stochastic Value Gradient method [20] and the closely related Dreamer algorithms [21-23]. These methods perform credit assignment by backpropagating policy gradients through the world model. Crucially, this approach only works for environments with a continuous state-action space, as otherwise sensitivities of the value with respect to past actions are undefined [20, 22, 24]. Intuitively, we cannot compute sensitivities of discrete choices such as a yes / no decision as the agent cannot decide 'yes' a little bit more or less.
Building upon Hindsight Credit Assignment (HCA) [1], we develop Counterfactual Contribution Analysis (COCOA), a family of algorithms that use models for credit assignment compatible with discrete actions. We measure the contribution of an action upon subsequent rewards by asking a counterfactual question: 'would the agent still have achieved the rewarding outcome, if it had taken another action?' (c.f. Fig. 1A). We show that measuring contributions towards achieving a future state, as is proposed in HCA, leads to spurious contributions that do not reflect a contribution towards a reward. This causes HCA to degrade towards the high-variance REINFORCE method in most environments. Instead, we propose to measure contributions directly on rewarding outcomes and we develop various new ways of learning these contributions from observations. The resulting algorithm differs from value-based methods in that it measures the contribution of an action to individual rewards, instead of estimating the full expected sum of rewards. This crucial difference allows our contribution analysis to disentangle different tasks and ignore uncontrollable environment influences, leading to a gradient estimator capable of long-term credit assignment (c.f. Fig. 1B). We introduce a new method for analyzing policy gradient estimators which uses dynamic programming to allow comparing to ground-truth policy gradients. We leverage this to perform a detailed bias-variance analysis of all proposed methods and baselines showing that our new model-based credit assignment algorithms achieve low variance and bias, translating into improved performance (c.f. Fig. 1C).
We consider an undiscounted Markov decision process (MDP) defined as the tuple  $(S, \mathcal{A}, p, p_r)$ , with  $S$  the state space,  $\mathcal{A}$  the action space,  $p(S_{t+1} \mid S_t, A_t)$  the state-transition distribution and  $p_r(R \mid S, A)$  the reward distribution with bounded reward values  $r$ . We use capital letters for random variables and lowercase letters for the values they take. The policy  $\pi(A \mid S)$ , parameterized by  $\theta$ , denotes the probability of taking action  $A$  at state  $S$ . We consider an undiscounted infinite-horizon setting with a zero-reward absorbing state  $s_\infty$  that the agent eventually reaches:  $\lim_{t \to \infty} p(S_t = s_\infty) = 1$ . Both the discounted and episodic RL settings are special cases of this setting. (c.f. App. B), and hence all theoretical results proposed in this work can be readily applied to both (c.f. App J).
We use  $\mathcal{T}(s,\pi)$  and  $\mathcal{T}(s,a,\pi)$  as the distribution over trajectories  $T = (S_{t},A_{t},R_{t})_{t\geq 0}$  starting from  $S_0 = s$  and  $(S_0,A_0) = (s,a)$  respectively, and define the return  $Z_{t} = \sum_{t = 0}^{\infty}R_{t}$ . The value function  $V^{\pi}(s) = \mathbb{E}_{T\sim \mathcal{T}(s,\pi)}[Z_t]$  and action value function  $Q^{\pi}(s,a) = \mathbb{E}_{T\sim \mathcal{T}(s,a,\pi)}[Z_t]$  are the expected return when starting from state  $s$ , or state  $s$  and action  $a$  respectively. Note that these infinite sums have finite values due to the absorbing zero-reward state (c.f. App. B).
The objective of reinforcement learning is to maximize the expected return  $V^{\pi}(s_0)$ , where we assume the agent starts from a fixed state  $s_0$ . Policy gradient algorithms optimize  $V^{\pi}(s_0)$  by repeatedly estimating its gradient  $\nabla_{\theta}V(s_0)$  w.r.t. the policy parameters. REINFORCE [6] (c.f. Tab. 1) is the canonical policy gradient estimator, however, it has a high variance resulting in poor parameter updates. Common techniques to reduce the variance are (i) subtracting a baseline, typically a value estimate, from the sum of future rewards [2, 25] (c.f. 'Advantage' in Tab. 1); (ii) replacing the sum of future rewards with a learned action value function  $Q$  [2, 3, 25, 26] (c.f. 'Q-critic' in Tab. 1); and (iii) using temporal discounting. Note that instead of using a discounted formulation of MDPs, we treat the discount factor as a variance reduction technique in the undiscounted problem [10, 11, 13] as this more accurately reflects its practical use [4, 27]. Rearranging the summations of REINFORCE with discounting lets us interpret temporal discounting as a credit assignment heuristic, where for each reward, past actions are reinforced proportional to their proximity in time.
$$
\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), \quad \gamma \in [ 0, 1 ]. \tag {1}
$$
Crucially, long-term dependencies between actions and rewards are exponentially suppressed, thereby reducing variance at the cost of disabling long-term credit assignment [4, 28]. The aim of this work is to replace the heuristic of time discounting by principled contribution coefficients quantifying how much an action contributed towards achieving a reward, and thereby introducing new policy gradient estimators with reduced variance, without jeopardizing long-term credit assignment.
HCA [1] makes an important step in this direction by introducing a new gradient estimator:
$$
\hat {\nabla} _ {\theta} ^ {\mathrm {H C A}} 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} \frac {p ^ {\pi} \left(A _ {t} = a \mid S _ {t} = s , S ^ {\prime} = S _ {t + k}\right)}{\pi (a \mid S _ {t})} R _ {t + k}\right) \tag {2}
$$
with  $r(s, a)$  a reward model, and the hindsight ratio  $\frac{p^{\pi}(a|S_t = s, S' = S_{t+k})}{\pi(a|S_t)}$  measuring how important action  $a$  was to reach the state  $S'$  at some point in the future. Although the hindsight ratio delivers precise credit assignment w.r.t. reaching future states, it has a failure mode of practical importance, creating the need for an updated theory of model-based credit assignment which we will detail in the next section.
3 Counterfactual Contribution Analysis
To formalize the 'contribution' of an action upon subsequent rewards, we generalize the theory of HCA [1] to measure contributions on rewarding outcomes instead of states. We introduce unbiased policy gradient estimators that use these contribution measures, and show that HCA suffers from high variance, making the generalization towards rewarding outcomes crucial for obtaining low-variance estimators. Finally, we show how we can estimate contributions using observational data.
3.1 Counterfactual contribution coefficients
To assess the contribution of actions towards rewarding outcomes, we propose to use counterfactual reasoning: 'how does taking action  $a$  influence the probability of obtaining a rewarding outcome, compared to taking alternative actions  $a'$ ?'.
Table 1: Comparison of policy gradient estimators.
Method	Policy gradient estimator (帽子Vπ(s0))
REINFORCE	Σt≥0∇θlog π(At|St) Σk≥0Rt+k
Advantage	Σt≥0∇θlog π(At|St) (Σk≥0Rt+k - V(ST))
Q-critic	Σt≥0 Σa∈A ∇θπ(a|St)Q(ST,a)
HCA-Return	Σt≥0 ∇θlog π(At|St) (1 - π(At|St)/pπ(At|St,Zt))Zt
TrajCV	Σt≥0 ∇θlog π(At|St) (Zt - Q(At,St) - Σt'>t (Q(St',At') - V(St')) + ...
... Σa∈A ∇θπ(a|St)Q(ST,a)
COCOA	Σt≥0 ∇θlog π(At|St) Rt + Σa∈A ∇θπ(a|St) Σk≥1w(St,a,Ut+k)Rt+k
HCA+	Σt≥0 ∇θlog π(At|St) Rt + Σa∈A ∇θπ(a|St) Σk≥1w(St,a,St+k)Rt+k
Definition 1 (Rewarding outcome). A rewarding outcome  $U' \sim p(U' \mid s', a', r')$  is a probabilistic encoding of the state-action-reward triplet.
If action  $a$  contributed towards the reward, the probability of obtaining the rewarding outcome following action  $a$  should be higher compared to taking alternative actions. We quantify the contribution of action  $a$  taken in state  $s$  upon rewarding outcome  $u'$  as
$$
w (s, a, u ^ {\prime}) = \frac {\sum_ {k \geq 1} p ^ {\pi} \left(U _ {t + k} = u ^ {\prime} \mid S _ {t} = s , A _ {t} = a\right)}{\sum_ {k \geq 1} p ^ {\pi} \left(U _ {t + k} = u ^ {\prime} \mid S _ {t} = s\right)} - 1 = \frac {p ^ {\pi} \left(A _ {t} = a \mid S _ {t} = s , U ^ {\prime} = u ^ {\prime}\right)}{\pi (a \mid s)} - 1 \tag {3}
$$
From a given state, we compare the probability of reaching the rewarding outcome  $u'$  at any subsequent point in time, given we take action  $a$  versus taking counterfactual actions according to the policy  $\pi$ , as  $p^{\pi}(U_{t+k} = u' \mid S_t = s) = \sum_{a'} \pi(a' \mid s)p^{\pi}(U_{t+k} = u' \mid S_t = s, A_t = a')$ . Subtracting this ratio by one results in an intuitive interpretation of the contribution coefficient  $w(s, a, u')$ : if the coefficient is positive/negative, performing action  $a$  results in a higher/lower probability of obtaining rewarding outcomes  $u'$ , compared to following the policy  $\pi$ . Using Bayes' rule, we can convert the counterfactual formulation of the contribution coefficients into an equivalent hindsight formulation (right-hand side of Eq. 3), where the hindsight distribution  $p^{\pi}(A_t = a \mid S_t = s, U' = u')$  reflects the probability of taking action  $a$  in state  $s$ , given that we encounter the rewarding outcome  $u'$  at some future point in time. We refer the reader to App. C for a full derivation.
Choice of rewarding outcome. For  $u' = s'$ , we recover state-based HCA [1]. In the following, we show that a better choice is to use  $u' = r'$ , or an encoding  $p(u' \mid s', a')$  of the underlying object that causes the reward. Both options lead to gradient estimators with lower variance (c.f. Section 3.3), while using the latter becomes crucial when different underlying rewarding objects have the same scalar reward (c.f. Section 4).
3.2 Policy gradient estimators
We now show how the contribution coefficients can be used to learn a policy. Building upon HCA [1], we propose the Counterfactual Contribution Analysis (COCOA) policy gradient estimator
$$
\hat {\nabla} _ {\theta} ^ {U} 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} w \left(S _ {t}, a, U _ {t + k}\right) R _ {t + k}. \tag {4}
$$
When comparing to the discounted policy gradient of Eq. 1, we see that the temporal discount factors are substituted by the contribution coefficients, replacing the time heuristic with fine-grained credit assignment. Importantly, the contribution coefficients enable us to evaluate all counterfactual actions instead of only the observed ones, further increasing the quality of the gradient estimator (c.f. Fig. 1A). The contribution coefficients allow for various different gradient estimators (c.f. App. C). For example, independent action samples can replace the sum over all actions, making it applicable to large action spaces. Here, we use the gradient estimator of Eq. 4, as our experiments consist of small action spaces where enumerating all actions is feasible. When  $U = S$ , the above estimator is almost equivalent to
<sup>3</sup>Note that Harutyunyan et al. [1] also introduce an estimator based on the return. Here, we focus on the state-based variant as only this variant uses contribution coefficients to compute the policy gradient (Eq. 4). The return-based variant instead uses the hindsight distribution as an action-dependent baseline as shown in Tab. 1. Importantly, the return-based estimator is biased in many relevant environments (c.f. Appendix L).
the state-based HCA estimator of Eq. 2, except that it does not need a learned reward model  $r(s,a)$ . We use the notation HCA+ to refer to this simplified version of the HCA estimator. Theorem 1 below shows that the COCOA gradient estimator is unbiased, as long as the encoding  $U$  is fully predictive of the reward, thereby generalizing the results of Harutyunyan et al. [1] to arbitrary rewarding outcome encodings.
Definition 2 (Fully predictive). A rewarding outcome  $U$  is fully predictive of the reward  $R$ , if the following conditional independence condition holds for all  $k \geq 0$ :  $p^{\pi}(R_k = r \mid 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$ .
Theorem 1. Assuming that  $U$  is fully predictive of the reward (c.f. Definition 2), the COCOA policy gradient estimator  $\hat{\nabla}_{\theta}^{U}V^{\pi}(s_{0})$  is unbiased, when using the ground-truth contribution coefficients of Eq. 3, that is
$$
\nabla_ {\theta} V ^ {\pi} (s _ {0}) = \mathbb {E} _ {T \sim \mathcal {T} (s _ {0}, \pi)} \hat {\nabla} _ {\theta} ^ {U} V ^ {\pi} (s _ {0}).
$$
3.3 Optimal rewarding outcome encoding for low-variance gradient estimators
Theorem 1 shows that the COCOA gradient estimators are unbiased for all rewarding outcome encodings  $U$  that are fully predictive of the reward. The difference between specific rewarding outcome encodings manifests itself in the variance of the resulting gradient estimator. Proposition 2 shows that for  $U' = S'$  as chosen by HCA there are many cases where the variance of the resulting policy gradient estimator degrades to the high-variance REINFORCE estimator [6]:
Proposition 2. In environments where each action sequence leads deterministically to a different state, we have that the  $HCA +$  estimator is equal to the REINFORCE estimator (c.f. Tab. 1).
In other words, when all previous actions can be perfectly decoded from a given state, they trivially all contribute to reaching this state. The proof of Proposition 2 follows immediately from observing that  $p^{\pi}(a \mid s, s') = 1$  for actions  $a$  along the observed trajectory, and zero otherwise. Substituting this expression into the contribution analysis gradient estimator (4) recovers REINFORCE. A more general issue underlies this special case: State representations need to contain detailed features to allow for a capable policy but the same level of detail is detrimental when assigning credit to actions for reaching a particular state since at some resolution almost every action will lead to a slightly different outcome. Measuring the contribution towards reaching a specific state ignores that the same rewarding outcome could be reached in slightly different states, hence overvaluing the importance of previous actions and resulting in spurious contributions. Many commonly used environments, such as pixel-based environments, continuous environments, and partially observable MDPs exhibit this property to a large extent due to their fine-grained state representations (c.f. App. G). Hence, our generalization of HCA to rewarding outcomes is a crucial step towards obtaining practical low-variance gradient estimators with model-based credit assignment.
Using rewards as rewarding outcomes yields lowest-variance estimators. The following Theorem 3 shows in a simplified setting that (i) the variance of the REINFORCE estimator is an upper bound on the variance of the COCOA estimator, and (ii) the variance of the COCOA estimator is smaller for rewarding outcome encodings  $U$  that contain less information about prior actions. We formalize this with the conditional independence relation of Definition 2 by replacing  $R$  with  $U'$ : encoding  $U$  contains less or equal information than encoding  $U'$ , if  $U'$  is fully predictive of  $U$ . Combined with Theorem 1 that states that an encoding  $U$  needs to be fully predictive of the reward  $R$ , we have that taking the reward  $R$  as our rewarding outcome encoding  $U$  results in the gradient estimator with the lowest variance of the COCOA family.
Theorem 3. Consider an MDP where only the states at a single (final) time step contain a reward, and where we optimize the policy only at a single (initial) time step. Furthermore, consider two rewarding outcome encodings  $U$  and  $U'$ , where  $S$  is fully predictive of  $U'$ ,  $U'$  fully predictive of  $U$ , and  $U$  fully predictive of  $R$ . Then, the following relation holds between the policy gradient estimators:
$$
\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} ^ {R E I N F} V ^ {\pi} \left(s _ {0}\right) \right]
$$
with  $\hat{\nabla}_{\theta}^{X}V^{\pi}(s_{0})$  the COCOA estimator (4) using  $U = X$ ,  $\mathbb{V}[Y]$  the covariance matrix of  $Y$  and  $A \preccurlyeq B$  indicating that  $B - A$  is positive semi-definite.
As Theorem 3 considers a simplified setting, we verify empirically whether the same arguments hold more generally. We construct a tree environment where we control the amount of information a
Figure 2: HCA suffers from spurious contributions which can be alleviated by using less informative rewarding outcome encodings. (A) and (C): Schematic of the tree environment where we parametrically adjust the amount of overlap between states by varying the amount of shared children of two neighboring nodes. We can decrease the information content of the rewarding outcome encoding  $u = f(s, a)$  by grouping state-action pairs that share the same reward value. (B) Normalized variance in dB using ground-truth coefficients and a random uniform policy (shaded region represents standard error over 10 random environments) comparing REINFORCE, HCA, COCOCA-reward and various degrees of intermediate grouping.
state contains about the previous actions by varying the overlap of the children of two neighbouring nodes (c.f. Fig 2), and assign a fixed random reward to each state-action pair. We compute the ground-truth contribution coefficients by leveraging dynamic programming (c.f. Section 4). Fig. 2B shows that the variance of HCA is as high as REINFORCE for zero state overlap, but improves when more states overlap, consistent with Proposition 2 and Theorem 3. To investigate the influence of the information content of  $U$  on the variance, we consider rewarding outcome encodings  $U$  with increasing information content, which we quantify with how many different values of  $u$  belong to the same reward  $r$ . Fig. 2B shows that by increasing the information content of  $U$ , we interpolate between the variance of COCOA with  $u = r$  and HCA+, consistent with Theorem 3.
Why do rewarding outcome encodings that contain more information than the reward lead to higher variance? To provide a better intuition on this question we use the following theorem: