SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Theorem 4. The policy gradient on the expected number of occurrences  $O^{\pi}(u', s) = \sum_{k \geq 1} p^{\pi}(U_k = u' \mid S_0 = s)$  is proportional to
$$
\nabla_ {\theta} O ^ {\pi} (u ^ {\prime}, s) \propto \mathbb {E} _ {S ^ {\prime \prime} \sim \mathcal {T} (\pi , s)} \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a \mid S ^ {\prime \prime}) w (S ^ {\prime \prime}, a, u ^ {\prime}) O ^ {\pi} (u ^ {\prime}, S ^ {\prime \prime})
$$
Recall that the COCOA gradient estimator consists of individual terms that credit past actions  $a$  at state  $s$  for a current reward  $r$  encountered in  $u$  according to  $\sum_{a\in \mathcal{A}}\nabla_{\theta}\pi (a\mid s)w(s,a,u')r'$  (c.f. Eq. 4). Theorem 4 indicates that each such term aims to increase the average number of times we encounter  $u^{\prime}$  in a trajectory starting from  $s$ , proportional to the corresponding reward  $r^\prime$ . If  $U^{\prime} = R^{\prime}$ , this update will correctly make all underlying states with the same reward  $r^\prime$  more likely while decreasing the likelihood of all states for which  $u^{\prime}\neq r^{\prime}$ . Now consider the case where our rewarding outcome encoding contains a bit more information, i.e.  $U^{\prime} = f(R^{\prime},\Delta S^{\prime})$  where  $\Delta S^{\prime}$  contains a little bit of information about the state. As a result the update will distinguish some states even if they yield the same reward and increase the number of occurrences only of states containing the encountered  $\Delta S^{\prime}$  while decreasing the number of occurrences for unseen ones. As in a single trajectory, we do not visit each possible  $\Delta S^{\prime}$ , this adds variance. The less information an encoding  $U$  has, the more underlying states it groups together, and hence the less rewarding outcomes are 'forgotten' in the gradient estimator, leading to lower variance.
3.4 Learning the contribution coefficients
In practice, we do not have access to the ground-truth contribution coefficients, but need to learn them from observations. Following Harutyunyan et al. [1], we can approximate the hindsight distribution  $p^{\pi}(A_t = a \mid S_t = s, U' = u')$ , now conditioned on rewarding outcome encodings instead of states, by training a model  $h(a \mid s, u')$  on the supervised discriminative task of classifying the observed action  $a_t$  given the current state  $s_t$  and some future rewarding outcome  $u'$ . Note that if the model  $h$  does not approximate the hindsight distribution perfectly, the COCOA gradient estimator (4) can be biased (c.f. Section 4). A central difficulty in approximating the hindsight distribution is that it is policy dependent, and hence changes during training. Proposition 5 shows that we can provide the policy logits as an extra input to the hindsight network without altering the learned hindsight
Figure 3: COCOA enhances policy gradient estimates and sample efficiency whereas HCA fails to improve over baselines. (A) Schematic representation of the linear key-to-door environment. (B) Performance of COCOA and baselines on the main task of picking up the treasure, measured as the average fraction of treasure rewards collected. (Left) ground-truth policy gradient estimators computed using dynamic programming, (right) learning the contribution coefficients or (action-value function using neural networks. Shaded regions are the standard error (30 seeds). (C) Violin plot of the signal-to-noise ratio (SNR) in Decibels for the various policy gradient estimators with learned coefficients and (action-)value functions, computed on the same trajectories of a shared base policy. (D) Comparison of the bias-variance trade-off incurred by different policy gradient estimators, computed as in (C), normalized by the ground-truth policy gradient norm (scatter plot showing 30 seeds per method).
distribution, to make the parameters of the hindsight network less policy-dependent. This observation justifies and generalizes the strategy of adding the policy logits to the hindsight model output, as proposed by Alipov et al. [29].
Proposition 5.  $p^{\pi}(a \mid s, u', l) = p^{\pi}(a \mid s, u')$ , with  $l$  a deterministic function of  $s$ , representing the sufficient statistics of  $\pi(a \mid s)$ .
As an alternative to learning the hindsight distribution, we can directly estimate the probability ratio  $p^{\pi}(A_t = a \mid S_t = s, U' = u') / \pi(a \mid s)$  using a contrastive loss (c.f. App. D). Yet another path builds on the observation that the sums  $\sum_{k \geq 1} p^{\pi}(U_{t+k} = u' \mid s, a)$  are akin to Successor Representations and can be learned via temporal difference updates [30, 31] (c.f. App. D). We experimented both with the hindsight classification and the contrastive loss and found the former to work best in our experiments. We leverage the Successor Representation to obtain ground truth contribution coefficients via dynamic programming for the purpose of analyzing our algorithms.
4 Experimental analysis
To systematically investigate long-term credit assignment performance of COCOA compared to standard baselines, we design an environment which pinpoints the core credit assignment problem and leverage dynamic programming to compute ground-truth policy gradients, contribution coefficients, and value functions (c.f. App E.2). This enables us to perform detailed bias-variance analyses and to disentangle the theoretical optimal performance of the various gradient estimators from the approximation quality of learned contribution coefficients and (action-)value functions.
We consider the linear key-to-door environment (c.f. Fig. 3A), a simplification of the key-to-door environment [3, 4, 32] to a one-dimensional track. Here, the agent needs to pick up a key in the first time step, after which it engages in a distractor task of picking up apples with varying reward values. Finally, it can open a door with the key and collect a treasure. This setting allows us to parametrically increase the difficulty of long-term credit assignment by increasing the distance between the key and door, making it harder to pick up the learning signal of the treasure reward
among a growing number of varying distractor rewards [4]. We use the signal-to-noise ratio,  $\mathrm{SNR} = \|\nabla_{\theta} V^{\pi}\|^{2} / \mathbb{E}[\|\hat{\nabla}_{\theta} V^{\pi} - \nabla_{\theta} V^{\pi}\|^{2}]$ , to quantify the quality of the different policy gradient estimators; a higher SNR indicates that we need fewer trajectories to estimate accurate policy gradients [33].
Previously, we showed that taking the reward as rewarding outcome encoding results in the lowest-variance policy gradients when using ground-truth contribution coefficients. In this section, we will argue that when learning the contribution coefficients, it is beneficial to use an encoding  $u$  of the underlying rewarding object since this allows to distinguish different rewarding objects when they have the same scalar reward value and allows for quick adaptation when the reward function but not the environment dynamics changes.
We study two variants of COCOA, COCOA-reward which uses the reward identity for  $U$ , and COCOA-feature which acquires features of rewarding objects by learning a reward model  $r(s,a)$  and taking the penultimate network layer as  $U$ . We learn the contribution coefficients by approximating the hindsight distribution with a neural network classifier  $h(a\mid s,u',l)$  that takes as input the current state  $s$ , resulting policy logs  $l$ , and rewarding outcome  $u'$ , and predicts the current action  $a$  (c.f. App. E for all experimental details). As HCA+ (c.f. Tab. 1) performs equally or better compared to HCA [1] in our experiments (c.f. App. F), we compare to HCA+ and several other baselines: (i) three classical policy gradient estimators, REINFORCE, Advantage and Q-critic, (ii) TrajCV [34], a state-of-the-art control variate method that uses hindsight information in its baseline, and (iii) HCA-return [1], a different HCA variant that uses the hindsight distribution conditioned on the return as an action-dependent baseline (c.f. Tab. 1).
4.1 COCOA improves sample-efficiency due to favorable bias-variance trade-off.
To investigate the quality of the policy gradient estimators of COCOA, we consider the linear key-to-door environment with a distance of 100 between key and door. Our dynamic programming setup allows us to disentangle the performance of the estimators independent of the approximation quality of learned models by using ground truth contribution coefficients and (action-)value functions. The left panel of figure 3B reveals that in this ground truth setting, COCOA-reward almost immediately solves the task performing as well as the theoretically optimal Q-critic with a perfect action-value function. This is in contrast to HCA and HCA-return which perform barely better than REINFORCE, all failing to learn to consistently pick up the key in the given number of episodes. This result translates to the setting of learning the underlying models using neural networks. COCOA-reward and -feature outperform competing policy gradient estimators in terms of sample efficiency while HCA only improves over REINFORCE. Notably, having to learn the full action-value function leads to a less sample-efficient policy gradient estimator for the Q-critic.
In Figure 3C and D we leverage dynamic programming to compare to the ground truth policy gradient. This analysis reveals that improved performance of COCOA is reflected in a higher SNR compared to other estimators due to its favorable bias-variance trade-off. Fig 12 in App. F indicates that COCOA maintains a superior SNR, even when using significantly biased contribution coefficients. As predicted by our theory in Section 3.3, HCA significantly underperforms compared to baselines due to its high variance caused by spurious contributions. In particular, the Markov state representation of the linear key-to-door environment contains the information of whether the key has been picked up. As a result, HCA always credits picking up the key or not, even for distractor rewards. These spurious contributions bury the useful learning signal of the treasure reward in noisy distractor rewards. HCA-return performs poorly as it is a biased gradient estimator, even when using the ground-truth hindsight distribution (c.f. Appendix F and L). Interestingly, the variance of COCOA is significantly lower compared to a state-of-the-art control variate method, TrajCV, pointing to a potential benefit of the multiplicative interactions between contribution coefficients and rewards in the COCOA estimators, compared to the additive interaction of the control variates: the value functions used in TrajCV need to approximate the full average returns, whereas COCOA can ignore rewards from the distractor subtask, by multiplying them with a contribution coefficient of zero.
4.2 COCOA enables long-term credit assignment by disentangling rewarding outcomes.
The linear key-to-door environment allows us to parametrically increase the difficulty of the long-term credit assignment problem by increasing the distance between the key and door and thereby increasing the variance due to the distractor task [4]. Figure 1B reveals that as this distance increases, performance measured as the average fraction of treasure collected over a fixed number of episodes
Figure 4: COCOA improves performance by disentangling subtasks. (A) Schematic representation of the task interleaving environment where colored borders indicate the context of a room. (B) Performance of COCOA and baselines with learned contribution coefficients or value functions, measured as the fraction of correct choices. (C) Visualization of the contribution coefficient magnitudes of each query room on reward rooms for COCOA (top) and HCA (bottom), c.f. App. E.7.
drops for all baselines including HCA but remains relatively stable for the COCOA estimators. Hung et al. [4] showed that the SNR of REINFORCE decreases inversely proportional to growing distance between key and door. Figure 1B shows that HCA and all baselines follow this trend, whereas the COCOA estimators maintain a robust SNR.
We can explain the qualitative difference between COCOA and other methods by observing that the key-to-door task consists of two distinct subtasks: picking up the key to get the treasure, and collecting apples. COCOA can quickly learn that actions relevant for one task, do not influence rewarding outcomes in the other task, and hence output a contribution coefficient equal to zero for those combinations. Value functions in contrast estimate the expected sum of future rewards, thereby mixing the rewards of both tasks. When increasing the variance of the return in the distractor task by increasing the number of stochastic distractor rewards, estimating the value functions becomes harder, whereas estimating the contribution coefficients between state-action pairs and distinct rewarding objects remains of equal difficulty, showcasing the power of disentangling rewarding outcomes.
To further showcase the power of disentangling subtasks, we consider a simplified version of the task interleaving environment of Mesnard et al. [3] (c.f. Fig. 4A, App. E.7). Here, the agent is faced with a sequence of contextual bandit tasks, where the reward for a correct decision is given at a later point in time, together with an observation of the relevant context. The main credit assignment difficulty is to relate the reward and contextual observation to the correct previous contextual bandit task. Note that the variance in the return is now caused by the stochastic policy, imperfectly solving future tasks, and by stochastic state transitions, in contrast to the linear key-to-door environment where the variance is caused by stochastic distractor rewards. Figure 4B shows that our COCOA algorithms outperform all baselines. The learned contribution coefficients of COCOA reward accurately capture that actions in one context only contribute to rewards in the same context as opposed to HCA that fails to disentangle the contributions (c.f. Figure 4C).
4.3 Learned credit assignment features allow for disentangling aliased rewards
For COCOA-reward we use the scalar reward value to identify rewarding outcomes in the hindsight distribution, i.e.  $U = R$ . In cases where multiple rewarding outcomes yield an identical scalar reward value, the hindsight distribution cannot distinguish between them and has to estimate a common hindsight
Figure 5: COCOA-features is robust to reward aliasing. On a version of the linear key-to-door environment where one of the distractor reward values has the same magnitude as the treasure reward, COCOA-reward can no longer distinguish between the distractor and treasure reward and as a result its performance decreases. COCOA-feature is robust to this manipulation since it relies on learned features to distinguish rewarding objects.
probability, making it impossible to disentangle the tasks and hence rendering learning of the contribution coefficients potentially more difficult. In contrast, COCOA-feature learns hindsight features of rewarding objects that are predictive of the reward. Even when multiple rewarding objects lead to an identical scalar reward value their corresponding features are likely different, allowing COCOA-feature to disentangle the rewarding outcomes.
In Fig. 5, we test this reward aliasing setting experimentally and slightly modify the linear key-to-door environment by giving the treasure reward the same value as one of the two possible values of the stochastic distractor rewards. As expected, COCOA-feature is robust to reward aliasing, continuing to perform well on the task of picking up the treasure while performance of COCOA-reward noticeably suffers. Note that the performance of all methods has slightly decreased compared to Fig. 3, as the magnitude of the treasure reward is now smaller relative to the variance of the distractor rewards, resulting in a worse SNR for all methods.
We present a theory for model-based credit assignment compatible with discrete actions and show in a comprehensive theoretical and experimental analysis that this yields a powerful policy gradient estimator, enabling long-term credit assignment by disentangling rewarding outcomes.
Building upon HCA [1], we focus on amortizing the estimation of the contribution coefficients in an inverse dynamics model,  $p^{\pi}(a \mid s, u')$ . The quality of this model is crucial for obtaining low-bias gradient estimates, but it is restricted to learn from on-policy data, and rewarding observations in case of  $u = r$ . Scaling these inverse models to complex environments will potentially exacerbate this tension, especially in sparse reward settings. A promising avenue for future work is to leverage forward dynamics models and directly estimate contribution coefficients from synthetic trajectories. While learning a forward model is a difficult problem in its own, its policy independence increases the data available for learning it. This would result in an algorithm close in spirit to Stochastic Value Gradients [20] and Dreamer [21-23] with the crucial advance that it enables model-based credit assignment on discrete actions. Another possibility to enable learning from non-rewarding observations is to learn a generative model that can recombine inverse models based on state representations into reward contributions (c.f. App. H).
Related work has explored the credit assignment problem through the lens of transporting rewards or value estimates towards previous states to bridge long-term dependencies [4, 5, 32, 35-41]. This approach is compatible with existing and well-established policy gradient estimators but determining how to redistribute rewards has relied 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.
While we have demonstrated that contribution coefficients with respect to states as employed by HCA suffer from spurious contributions, any reward feature encoding that is fully predictive of the reward can in principle suffer from a similar problem in the case where each environment state has a unique reward value. In practice, this issue might occur in environments with continuous rewards. A potential remedy in this situation is to assume that the underlying reward distribution is stochastic, smoothing the contribution coefficients as now multiple states could have led to the same reward. This lowers the variance of the gradient estimator as we elaborate in App. G.
Finally, we note that our contribution coefficients are closely connected to causality theory [42] where the contribution coefficients correspond to performing Do-interventions on the causal graph to estimate their effect on future rewards (c.f. App I). Within causality theory, counterfactual reasoning goes a step further by inferring the external, uncontrollable environment influences and considering the consequences of counterfactual actions given that all external influences remain the same [3, 20, 42-44]. Extending COCOA towards this more advanced counterfactual setting by building upon recent work [3, 43] is an exciting direction for future research (c.f. App. I).
Concluding remarks. By overcoming the failure mode of spurious contributions in HCA, we have presented here a comprehensive theory on how to leverage model information for credit assignment, compatible with discrete action spaces. COCOA-reward and COCOA-feature are promising first algorithms in this framework, opening the way towards sample-efficient reinforcement learning by model-based credit assignment.
We thank Angelika Steger, Yassir Akram, Ida Momennejad, Blake Richards, Matt Botvinick and Joel Veness for discussions and feedback. Simon Schug is supported by the Swiss National Science Foundation (PZ00P3_186027). Seijin Kobayashi is supported by the Swiss National Science Foundation (CRSII5_173721). Simon Schug would like to kindly thank the TPU Research Cloud (TRC) program for providing access to Cloud TPUs from Google.
[94] Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. arXiv, May 2014. Number: arXiv:1312.6114 arXiv:1312.6114 [cs, stat].
[95] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wunderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[96] Lukas Biewald. Experiment Tracking with Weights and Biases, 2020.
[97] Plotly Technologies Inc. Collaborative data science, 2015. Place: Montreal, QC Publisher: Plotly Technologies Inc.
A Related work 2
B Undiscounted infinite-horizon MDPs 3
C Theorems, proofs and additional information for Section 3 4
C.1 Contribution coefficients, hindsight distribution and graphical models 4
C.2 Proof Theorem 1 4
C.3 Different policy gradient estimators leveraging contribution coefficients. 6
C.4 Proof of Proposition 2 7
C.5 Proof Theorem 3. 7
C.6 Proof of Theorem 4 8
D Learning the contribution coefficients 9
D.1 Proof Proposition 5 9
D.2 Learning contribution coefficients via contrastive classification. 10
D.3 Successor representations. 11
E Experimental details and additional results 11
E.1 Algorithms 11
E.2 Dynamic programming setup 11
E.3 Bias, variance and SNR metrics 13
E.4 Linear key-to-door environment setup 14
E.5 Reward switching setup 15
E.6 Tree environment setup 15
E.7 Task interleaving environment setup 15
E.8 Training details 16
F Additional results 19
F.1 HCA vs HCA+ 19
F.2 Learned credit assignment features allow for quick adaptation to a change of the reward function. 19
F.3 Using returns instead of rewards as hindsight information 20
F.4 Investigation into the required accuracy of the hindsight models 21
G Contribution analysis in continuous spaces and POMDPs 21
G.1 Spurious contributions in continuous state spaces 22
G.2 Deterministic continuous reward functions can lead to excessive variance 22
G.3 Smoothing can alleviate excess variance by trading variance for bias 23
G.4 Continuous action spaces 24
G.5 Partially Observable MDPs 25
H Learning contribution coefficients from non-rewarding observations 26
H.1 Latent learning 26
H.2 Optimal rewarding outcome encodings for credit assignment. 26
H.3 Counterfactual reasoning on rewarding states 26
H.4 Learning credit assignment representations with an information bottleneck 27
I Contribution analysis and causality 28
I.1 Causal interpretation of COCOA 28
I.2 Extending COCOA to counterfactual interventions 29
J Contribution analysis with temporal discounting 30
K Bootstrapping with COCOA 30
L HCA-return is a biased estimator in many relevant environments 33
L.1 HCA-return 33
L.2 HCA-return can be biased 33
M Additional details 34
M.1 Author contributions 34
M.2 Compute resources 34
M.3 Software and libraries 34