Would I have gotten that reward? Long-term credit assignment by counterfactual contribution analysis
Authors: Alexander Meulemans
Simon Schug
Seijin Kobayashi
Nathaniel Daw
Greg Wayne
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Reinforcement learning
Long-term credit assignment
contribution analysis
hindsight credit assignment
policy gradient methods

Abstract: 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), 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.

Would I have gotten that reward? Long-term credit assignment by counterfactual contribution analysis
Alexander Meulemans" Simon Schug"), Seijin Kobayashi"
Nathaniel D Daw Gregory Wayne?
Depparmeent of Computer Science, ETH Zürich
Google DeepMind
Princeton Neuroscience Institute, Princeton University
Department of Psychology, Princeton University (ameulema, ema, sschug, seijink]@ethhz
Abstract
To make reinforcement learning more sample efficient, we need better credit assign- ment 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 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 HCA, results in spurious estimates of contributions, causing HCA to degrade towards the high-variance REINFORCE. estimator 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 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- 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.?
Introduction
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 requires measuring how each action influences expected subsequent rewards. As the number of possible trajectories grows exponentially with time, future rewards come with 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 [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. "odee available httpa://githhr com/sei jineijiii---- jin-kobyasshi/cooaaa
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

A
B
COCOA-maard
COCOA-fettre
HGM
HGA-rmum
Eajectrry 2 = (bu, Hee ny)es
Average framiin of treasurre col collected
D-oritc
Advantage
REINFORCE
Trajov
0-
0.6
SNR dB
-20
Update the polcy to increase the probabillty of actions That cou ld have consiiuueed to reaching reward Fo.
04-
Avtee Ag] (e(ss, 19 7a
-40
0.-
-60
TO
100
100
Cradit assignment distance
Credit assignment distance
Figure 1: Counterfactual Contribution Analysis enables long-term credit assignment. (A) Given sample trajectory that eventually results in 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 r3 is when following action d1 versus the counterfactual actions dg and 03 in state $2. This is quantified through the contribution coefficient w($2, 01, 01, ra) which is used to update all possible action probabilities of the policy <( 82). (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-tor -ogii-to-nooie 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 a difficult problem by itself, when successful it can be used to generate 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 choices such as a yes / no decision as the agent cannot decide 'yes' a little bit more or less.
[20, 22, 24]. Intuitively, we cannot compute sensitivities of discrete
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. AA). We show that measuring contributions towards achieving future state, as is proposed in HCA, leads to spurious contributions that do not reflect a contribution towards 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 differencee 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. IB). 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 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 fc.f. Fig, 1C).
2

---

2 Background and notation
We consider an undiscounted Markov decision process (MDP) defined as the tuple (S, P. Pr)) with S the state space, A the action space, p(S,+1 S, A,) the state-transition distribution and Pr(R| 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 π <<A S), parameterized by 0, denotes the probability of taking action at state S. We consider an undiscounted infinite-horizon setting with rero-reward absorbing state Seo that the agent eventually reaches: lim, p(St sxc) = 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 Ts,, and ((8, a, IT) the distribution over trajectories T = (SS, A,, R2)020 starting from So = 8 and (So, Ao) (s, a) respectively, and define the return Z, = >i R. The value function (T()) ((..r) [Z,] and action value function QT a) = ErwT((s...)) [2,] 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 (so), where we assume the agent starts from fixed state su. Policy gradient algorithms optimize (*(ss) ) by repeatedly estimating its gradient Vo (so) w.r.t. the policy parameters. REINFORCE [6] (c.f. Tab. 1) 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 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 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 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 reinforced proportional their proximity in time.
A REINFORCE..
Σ 120
R
E
Age kgt
V log T (Ax [k), Y E [0.
(1)
Crucially, long-term dependencies between actions and rewards are exponentially suppressed, thereby reducing variance the cost of disabling long-term credit assignment [4, 28). The aim of this work to replace the heuristic of time discounting by principled contribution coefficients quantifying how much an action contributed towards achieving 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 new gradient estimator:
HC
ΣΣ Var(al Si) (r(S,, a) A
p" (A₁ S S, Setk), Re+k)
(2)
T (aa S:)
with r(s, a) reward model, and the hindsight ratio p" S5." measuring how important
Sz)
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 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.
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 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 influence the probability of obtaining a rewarding outcome, compared to taking alternative actions a'?'.
3

---

Table 1: Comparison of policy gradient estimators.
Method
Policy gradient estimator (VeV*(so))
REINFORCE
Σ 120 Ve log π (A, | S)) Ex20 Re+k
Advantage
- VV(s)) 120 Ve log π (A, S)) Ex20 R
Q-critic
Σ 120 EAEA Vo (,)Q(S;; a)
HCA-Return
Σ E1 Va logg((()) St) p" (Ar
*(A.(S.) Zz
TrajCV
Ei Vo log (A,1 Si) (Ze Q(A,, Sb)
Eusi(Q(Se, V((e)) +...
Enea Vor(a
()),()(.
COCOA HCA+
Σ 220 Vo log 7 (A, St) R2 + EaEA Vor(a St) S Exz w(ee a, U R
Eeo Ve log (A
S.) R2 DAEA Ver(a S)) EE> ES w(e,, a, S
",,'r'' .r' is probabilistic Definition 1 (Rewarding outcome). A rewarding outcome U p(U' |s',a', encoding of the actiin-rttwaadt triplet.
If action contributed towards the reward, the probability of obtaining the rewarding outcome follow- ing action a should be higher compared to taking alternative actions. We quantify the contribution of action a taken in state upon rewarding outcome ul
w(s a, v Σx21 p² (UL++ = u' St A==
D" (A. = S. s,U' w.)
(3)
Σne p" (U++: w/ S.=
From a given state, we compare the probability of reaching the rewarding outcome u' at any sub- sequent point in time, given we take action a versus taking counterfactual actions according to the policy π, as p" (U+++= ='' S, = 8) = E. T(a' | (Up"" = v' St = A, = d'). Subtracting this ratio by one results in an intuitive interpretation of the contribution coefficient w((s, a, u'): if the coefficient is positive/negaatvee, performing action a results in higher/lower probability of obtaining rewarding outcomes u', compared to following the policy л. Using Bayes' rule, we can convert the counterfactual formulation of the contribution coefficients into equivalent hindsight formulation (right-hand side of Eq. 3), where the hindsight distribution p" (A, = a S, = s, U' = u') reflects the probability of taking action in state 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.
(U+-
Choice of rewarding outcome. For u' 8', we recover state-based HCA [1].3 3 In the following, we show that a better choice is to use x' = r', or an encoding p(u') 8', 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
44 V = Σ Vo (A4
S.) + Σ Vor(a S.) Σ u((s, a, (
120
acA
#21
When comparing to the discounted policy gradient of Eq. 1, we see that the temporal discount factors 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. IA). 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 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
"Note that Harutyunyan et al. [1] also introduce an estimator based on the return. Here, we focus on the based variant as only this variant uses contribution coefficients to compute the policy gradient (Eq. 4). The return- baaud- 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 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 Rt if the following conditional independence condition holds for all k 2 0: p = r Sg = 8, Au = a. U2 u) = p" - u), where the right- side does not depend on the time k.
Theorem 1. Assuming that U is fully predictive of the reward (c.f. Desimition 2), the COCOA policy V gradient estimator V (8o) is unbiased, when using the ground-truth contribution coefficients of Eq. 3, that is
VeV* = ET~T(So. V' (su).
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 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-variane
REINFORCE estimator [6]:
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).
Proposition 2.
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 Propostion 2 follows immediately from observing that (a(aa s, s') 1 for actions 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 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 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 nowest-varianee estimators. The following Theo- rem 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 that states that an encoding U needs to be fully predictive of the reward R. we have that taking the reward R as 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' fully predictive of U, and U fully predictive of R. Then, the following relation holds between the policy gradient estimators:
V (su)] x V & V (**(s $ V( < V (su)
with VX (*o))) the COCOA estimator (4) using U= X, V(Y) the covariance matrix of Y and A & indicating that B A is positive semi-definite.
As Theorem considers simplified setting, we verify empirically whether the same arguments hold more generally. We construct tree environment where we control the amount of information
5

---

B
REINFORGE
(m)
HCAI
Ver dB
12-
COCOA roward
NO state overap
State averiap
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 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 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 better intuition this question we the following theorem:
Theorem 4. The policy gradient on the expected number of occurrences O (u' = Exz1 P" ("UU==
u' | So s) is proportional to
(uO x Eg"n((T..)) EASA Vor(a ("")((S" a, ('')'""','"))
Recall that the COCOA gradient estimator consists of individual terms that credit past actions state for current reward r encountered in u according EnEA Vaz(a s)w(s, a, y''r' (c.f. Eq. 4). Theorem 4 indicates that each such term aims to increase the average number of times we encounter u' in trajectory starting from s, proportional to the corresponding reward r'. If R', this update will correctly make all underlying states with the same reward r' more likely while decreasing the likeliness of all states for which u' * r'. Now consider the case where our rewarding outcome encoding contains bit more information, i.e. U = =f(R', AS') where A.' contains a little bit of information about the state. As 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 A.S' while decreasing the number of occurrences for unseen ones. As a single trajectory, we do not visit each possible AS', 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-trute -trth contribution coefficients, but need to learn them from observations. Following Harutyunyan et al. [1], we can approximate the hindsight distribution (A, 14400 a S, s, U = u'), now conditioned on rewarding outcome encodings instead of states, by training model hiaa s, u') on the supervised discriminative task of classifying the observed action a, given the current state and some future rewarding outcome u'. Note that if the model 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 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

---

600
0
COCON
COCOA-63tJ
HCA+
HGArrt.m
Q
Adaaraage
REINFORCE
Teye
Treasure reward colerted
Treasure reward colecied
C
Episode
Episode
100
CICCO fenture
HDA+
Advantage REINFORCE
TeaCVV
der
SNR diB
Vontance
-50-
-00
-120
COCOA
COCOA
HCA
HoA
Crontic
Advansage
RENFORCE ORCEE
RENI
-60
60
fanture
Biss dB
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-)vall functions, computed the same trajectories of a shared base policy. (D) Comparison the bias-variance trade-off incurred by different policy gradient estimators, computed in (C), normalized by the ground-truth policy gradient norm (scatter plot showing 30 seeds per method).
distribution, 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. pa(a u', ()) = p" (a s, u'), with deterministic function of representing the sufficient statistics of T T(a s).
As an alternative to learning the hindsight distribution, we can directly estimate the probability ratio (A, = a 5, = s, U' = (() 8) using a contrastive loss (c.f. App. D). Yet another path builds on the observation that the sums Ex>1 p" = w' 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 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.
Experimental analysis
To systematically investigate long-term credit assignment performance of COCOA compared to standard baselines, we design 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 to perform detailed bias 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-dood environment (c.f. Fig. 3A), a simplification of the key-to-door environment [3, 4, 32] to mi--imensional track. Here, the agent needs to pick up key in the first time step, after which engages in a distractor task of picking up apples with varying reward values. Finally, it can open door with the key and collect 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
7

---

among a growing number of varying distractor rewards [4]. We use the signal-to-ooise ratio, SNR= V 1 /EIIIV - - [I1" 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 dynamies 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 laking the penultimate network layer as U. We learn the contribution coefficients by approximating the hindsight distribution with a neural network classifier h( 5, =', that takes as input the current state 8, resulting policy logits I, 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-an control variate method that uses hindsight information in its baseline, and (ii) HCA-return [1], 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-rewart 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 oetion-value function leads to 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 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.
control variate method, TrajCV, pointing to potential benefit of the
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 fixed number of episodes
8

---

A
COCDAreward
Advantage REINFORCE
Corkec
CoCOA-feaure
A
HCA+
TraCVV
Conkent
HEA
G cnitic
Quera
COCOA
neward
R
Fraddion correct
0.8
Corkert
garly
0.7-
Canfer22
Sumple
Corkar Gr
0.5
HCA+
Epleode
Reward
Figure 4: COCOA improves performance by disentangling subtasks. (A) Schematic representa- tion of the task interleaving environment where colored borders indicate the context of a room. (B) Performance of COCOA and baselines with learned contribution coefficients value functions, mea- sured 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 al. [4] showed that the SNR of REINFORCE decreases inversely proportional growing distance between key and door. Figure IB shows that HCA and all baselines follow this trend, whereas the COCOA estimators maintain robust SNR.
We can explain the qualitative difference between COCOA and other methods by observing that the key-to-door task consists of two distinet subtasks: picking up the key 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 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 environ- ment of Mesnard et al. [3] (c.f. Fig. 4A, App. E.7). Here, the agent is faced with sequence of contextual bandit tasks, where the reward for a correct decision is given at a later point in
COcONrreward
Advantage
Treasure reward collected
0.8
Q-critic
REINFORCE
time, together with an observation 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 caused by stochastic dis- tractor rewards. Figure 4B shows that our COCOA algorithms outperform all baselines. The learned contribution coefficients of COCOA reward accurately capture that actions in one con- text only contribute to rewards in the same context as opposed to HCA that fails to disentangle the contributions (c.f. Fig- ure 4C).
0.4-
2%
4k
8k
10k
Episode
Figure 5: COCOA-features is ro- bust to reward aliasing. On a ver- sion of the linear key-to-door en- vironment where one of the dis- tractor reward values has the same magnitude as the treasure reward, COCOA-reward can no longer dis- tinguish between the distractor and treasure reward and as result its performance decreases. COCOA- feature is robust to this manipula- tion since it relies on learned fea- tures to distinguish rewarding ob- jects.
4.3 Learned credit assignment
features allow for disentangling aliased rewards
For COCOA-rewart 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 distin- guish between them and has to estimate common hindsight

---

probability, making it impossible to disentangle the tasks and hence rendering learning of the contri- bution 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. 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 -eaturee robust to reward aliasing, continuing to perform well on the task of picking up the treasure while performance of COCOA- rewaard noticeably suffers. Note that the performance of all methods has slightly decreased compared Fig. 3, as the magnitude of the treasure reward is now smaller relative to the variance of the distractor rewards, resulting in worse SNR for all methods.
Discussion
We present theory for model- based credit assignment compatible with discrete actions and show in 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, pa s, r²). 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 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 Stochastic Value Gradients [20] and Dreamer [21-23] with the crucial advance that it enables model-based credit assignment on discrete actions. Another possibility observations is to learn a generative model that can recombine inverse models based on state representations into reward contributions (c.f. App. H).
enable learning from non-rewarding
Related work has explored the credit assignment problem through the lens of transporting rewards or value estimates towards previous states bridge long-term dependencies [4, 5, 32, 35-4314, This approach is compatible with existing and well-established policy gradient estimators but determining how to redistribute rewards has relied heuristic contribution analyses, such 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 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 similar problem in the case where each environment state has unique reward value. In practice, this issue might occur in environments with continuous rewards. A potential remedy in this situation is 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
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 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. cowad-reward and COCOA-feature are promising first algorithms in this framework, opening the way towards sample-efficient reinforcement learning by model-based credit assignment.
10

---

6
Acknowledgements
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 Foun- dation (PZ00P3 _IpP33__1722... Seijin Kobayashi is supported by the Swiss National Science Foundation (CRSII5 173721). (ERS15_173721).. Simon Schug would like to kindly thank the TPU Research Cloud (TRC) program for providing access to Cloud TPUs from Google.
References
[1] Anna Harutyunyan, Will Dabney, Thomas Mesnard, Mohammad Gheshlaghi Azar, Bilal Piot, Nicolas Heess, Hado P van Hasselt, Gregory Wayne, Satinder Singh, Doina Precup. and Remi Munos. Hindsight Credit Assignment. In H. Wallach, H. Larochelle, A. Beygelzimer, F Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 12488- 12497. Curran Associates, Inc., 2019.
[2] Richard S. Sutton and Andrew G. Barto. Reinforcement Learning, second edition: An Introduc- tion. MIT Press, November 2018.
[3] Thomas Mesnard, Théophane Weber, Fabio Viola, Shantanu Thakoor, Alaa Saade, Anna Harutyunyan, Will Dabney, Tom Stepleton, Nicolas Heess, Arthur Guez, Marcus Hutter, Lars Buesing, and Rémi Munos. Counterfactual Credit Assignment in Model-Free Reinforcement Learning. In Proceedings of the 38 th International Conference Machine Learning, PMLR 139, 2021.
[4] Chia-Chun Hung, Timothy Lillicrap. Josh Abramson, Yan Wu, Mchdi Mirza, Federico Carnevale, Arun Ahuja, and Greg Wayne. Optimizing agent behavior over long time scales by transporting value. Nature Communications, 10(1):5223, November 2019. Number: Publisher: Nature Publishing Group.
[5] Jose A. Arjona-Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brandstetter, and Sepp Hochreiter. RUDDER: Return Decomposition for Delayed Rewards. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'textquotesingle Alché-Buc. E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 13566-13577. Curran Associates, Inc., 2019.
[6] Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine Learning, 8(3):229-22 256, May 1992.
[7] Evan Greensmith, Peter Bartlett, and Jonathan Baxter. Variance Reduction Techniques for Gradient Estimates in Reinforcement Learning. Journal of Machine Learning Research, 5(Nov): 1471-1530, 2004.
[8] Théophane Weber, Nicolas Heess, Lars Buesing, and David Silver. Credit Assignment Tech- niques in Stochastic Computation Graphs. In Proceedings of the Twenty-Secont International Conference on Artificial Intelligence and Statistics, pages 2650-26600 PMLR, April 2019. ISSN: 2640-3498.
[9] Jonathan Baxter and Peter L. Bartlett. Artificial Intelligence Research, 15:319-350, November 2001.
Gradiet Estimation. Journal of
[10] Philip Thomas. Bias in Natural Actor-Critie Algorithms. In Proceedings of Conference on Machine Learning, pages 441-448. PMIR, January 2014. ISSN: 1938-7228.
the 31st International
[11] Sham Kakade. Optimizing Average Reward Using Discounted Rewards. David Helmbold and Bob Williamson, editors, Computational Learning Theory, Lecture Notes in Computer Science, pages 605-615, Berlin, Heidelberg, 2001. Springer.
[12] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- Dimensional Continuous Control Using Generalized Advantage Estimation. In International Conference on Learning Representatiinns arXiv, 2016. 1.6v:15006.02338 [cs].
[13] Peter Marbach and John Tsitsiklis. Approximate Gradient Methods in Policy-Space Optimiza- tion of Markov Reward Processes. Discrete Event Dynamic Systems, page 38, 2003.
11

---

[14]
OpenAI, Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemystaw Dębiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, Rafal Józefowicz, Scott Gray, Catherine Olsson, Jakub Pachocki, Michael Petrov, Henrique P. d O. Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, Ilya Sutskever, Jie Tang, Filip Wolski, and Susan Zhang. Dota 2 with Large Scale Deep Reinforcement Learning, December 2019. arXiv: wwt.:1912..66800 [cs, stat].
[15] OpenAI, Ilge Akkaya, Marcin Andrychowicz. Maciek Chociej. Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas. Jonas Schneider, Nikolas Tezak, Jerry Tworek, Peter Welinder, Lilian Weng, Qiming Yuan, Wojeiech Zaremba, and Lei Zhang. Solving Rubik's Cube with Robot Hand, October 2019. arXiv:1910.0711 [cs, stat].
[16] Oriol Vinyals, Igor Babuschkin, Wojeiech M. Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H. Choi, Richard Powell, Timo Ewalds, Petko Georgiev, Junhyuk Oh, Dan Horgan, Manuel Kroiss, Ivo Danihelka, Aja Huang, Laurent Sifre, Trevor Cai, John P. Agapiou, Max Jaderberg, Alexander S. Vezhnevets, Rémi Leblond. Tobias Pohlen, Valentin Dalibard, David Budden, Yury Sulsky, James Molloy, Tom L. Paine. Caglar Gulcehre. Ziyu Wang, Tobias Pfaff, Yuhuai Wu, Roman Ring. Dani Yogatama, Dario Wünsch, Katrina McKin- ney, Oliver Smith, Tom Schaul, Timothy Lillicrap, Koray Kavukcuoglu, Demis Hassabis, Chris Apps, and David Silver. Grandmaster level in StarCraft II using multi-agent reinforcement learning. Nature, 575(7782):350-3 3544, November 2019. Number: 7782 Publisher: Nature Publishing Group.
[17] Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting, ACM SIGART Bulletin, 2(4): 163, July 1991.
[18] Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo. and Honglak Lee. Sample- Efficient Reinforcement Learning with Stochastic Ensemble Value Expansion. In 32nd Confer- ence on Neural Information Processing Systems. arXiv, 2018. arXiv:1807.0167. [cs, stat).
[19] Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-Ensemble Trust-Region Policy Optimization. 2018.
[20] Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning Continuous Control Policies by Stochastic Value Gradients. In Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015.
[21] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream Control: Learning Behaviors by Latent Imagination. In International Conference on Learning Represen- tations, March 2020. Number: arXiv:1912.0I. arXiv:1912.01660 [cs].
[22] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with Discrete World Models. In International Conference on Learning Representations, 2021. Number: arXiv:2010.02193 arXiv:2010.02193 [cs, stat].
[23] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering Diverse Domains through World Models, January 2023. arXiv:2301.044 [cs, stat].
[24] Mikael Henaff, William F. Whitney. and Yann LeCun. Model--asee Planning with Discrete and Continuous Actions, April 2018. arXiv:1705.077 10t:1000500177 [cs].
[25] Volodymyr Mnih Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous Methods for Deep Reinforce- ment Learning. In Proceedings of The 33rd International Conference on Machine Learning, pages 1928-199 1937. PMLR, June 2016. ISSN: 1938-7228.
[26] Richard S Sutton, David A McAllester, Satinder P Singh, and Yishay Mansour. Policy Gradient Methods for Reinforcement Learning with Function Approximation, In Advances in Neural Information Processing Systems 12, 1999.
[27] Chris Nota and Philip S Thomas. Is the Policy Gradient a Gradient? In Proc. of the 19th International Conference on Autonomous Agents and Multiagent Systems (AAMAS 2020).
page 9, 2020.
12

---

[28] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- Dimensional Continuous Control Using Generalized Advantage Estimation, October 2018. arXiv:1506.02436 [es].
[29] Vyacheslav Alipov, Riley Commons-Edlerr Nikita Putintsev, Pavel Kalinin, and Dmitry Vetrov. Towards Practical Credit Assignment for Deep Reinforcement Learning, February 2022.
arXiv:2106.04499 [es].
[30] Peter Dayan. Improving Generalization for Temporal Difference Learning: The Successor Representation. Neural Computation, 5(4):613-624, July 1993. Conference Name: Neural Computation.
[31] Tejas D. Kulkarni, Ardavan Saeedi, Simanta Gautam, and Samuel J. Gershman. Deep Successor Reinforcement Learning. June 2016. Number: arXiv:1606.02399 arXiv:1606.02349 [cs, stat].
[32] David Raposo, Sam Ritter, Adam Santoro, Greg Wayne, Theophane Weber, Matt Botvinick, Hado van Hasselt, and Francis Song. Synthetic Returns for Terg-Term Credit Assignment. arXiv:2102. 12425 [cs], February 2021. arXiv: 1202.12425.
[33] John W. Roberts and Russ Tedrake. Signal-to-Noise Ratio Analysis of Policy Gradient Algo- rithms. In D. Koller, D. Schuurmans, Y. Bengio, and L. Bottou, editors, Advances in Neural Information Processing Systems 21, pages 1361-1368. Curran Associates, Inc., 2009.
[34] Ching-An Cheng, Xinyan Yan, and Byron Boots. Trajectory-wise Control Variates for Variance Reduction in Policy Gradient Methods. In Proceedings of the Conference on Robot Learning, pages 1379-1322 1394. PMLR, May 2020. ISSN: 2640-3498.
[35] Zhizhou Ren, Ruihan Guo, Yuan Zhou, and Jian Peng. Learning Long-Term Reward Redistri- bution via Randomized Return Decomposition. January 2022.
[36] Yonathan Efroni, Nadav Merlis, and Shie Mannor. Reinforcement Learning with Trajectory Feedback. In The Thirty-Fifth AAAI Conference on Artificial Intelligenc. arXiv, March 2021.
[37] Minah Seo, Luiz Felipe Vecehietti, Sangkeum Lee, and Dongsoo Har. Rewards Prediction- Based Credit Assignment for Reinforcement Learning With Sparse Binary Rewards. IEEE ACCESS, 7:1 181118661187911 2019. Accepted: 1119-9924411:2215555 Publisher: IEEE-INST ELECTRICAL ELECTRONICS ENGINEERS INC.
[38] Markel Sanz Ausin, Hamoon Azizsoltani, Song Ju, Yeo Jin Kim, and Min Chi. InferNet for Delayed Reinforcement Tasks: Addressing the Temporal Credit Assignment Problem. In 2021 IEEE International Conference on Big Data (Big Data), pages 1337-1348, December 2021.
[39] Hamoon Azizsoltani, Yeo Jin Kim, Markel Sanz Ausin, Tiffany Barnes, and Min Chi. Unob- served Is Not Equal to Non-existent: Using Gaussian Processes to Infer Immediate Rewards Across Contexts. In Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence, pages 1974-1980, 2019.
[40] Vihang P. Patil, Markus Hofmarcher, Marius-Constanie Dinu, Matthias Dorfer, Patrick M. Blies, Johannes Brandstetter, Jose A. Arjona-Medina, and Sepp Hochreiter. Align-RUDDER: Learning From Few Demonstrations by Reward Redistribution. Proceedings of Machine Learning Research, 162, 2022. arXiv: 1009.14108.
[41] Johan Ferret, Raphaël Marinier, Matthieu Geist, and Olivier Pietquin. Self-Attentioaa Credit Assignment for Transfer in Reinforcement Learning. In Proceedings of the Twenty- Ninth International Joint Conference on Artificial Intelligence, pages 2655-2661, July 2020. arXiv:1907.0802: [cs].
[42] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Flements of causal inference: foun- dations and learning algorithms. Adaptive Computation and Machine Learning. November 2018.
[43] Lars Buesing, Theophane Weber, Yori Zwols, Sebastien Racaniere, Arthur Guez, Jean-Baptiste Lespiau, and Nicolas Heess. Woulda, Coulda, Shoulda: Counterfctuully--Guide Policy Search. In International Conference on Learning Representattons, 2019. arXiv: 1812.2....
13

---

[44] Ioana Bica, Ahmed M. Alaa, James Jordon, and Mihaela van der Schaar. Estimating Counter- factual Treatment Outcomes over Time Through Adversarially Balanced Representations. In International Conference on Learning Representations. arXiv, February 2020. arXiv:2002.04083 [es, stat].
[45] Pushi Zhang, Li Zhao, Guoqing Liu, Jiang Bian, Minlie Huang, Tao Qin, and Tan Liu. Independence-awwree Advantage Estimation. In Proceedings of the Thirtieth International Joint Conference Artificial Intelligence, pages 3349-3355, Montreal, Canada, August 2021. International Joint Conferences on Artificial Intelligence Organization.
[46] Kenny Young. Variance Reduced Advantage Estimation with Sidelta5 Hindsight Credit Assign- ment. arXiv:1911.08110 [cs], September 2020. arXiv: 1911.08362.
[47] Michel Ma and Bacon Pierre-Luc. Counterfactual Policy Evaluation and the Conditional Monte Carlo Method. In Offline Reinforcement Learning Workshop, NeurIPS, 2020.
[48] Paul Bratley, Bennett L. Fox, and Linus E. Schrage. A Guide to Simulation. Springer, New York, NY, 1987.
[49] J. M. Hammersley. Conditional Monte Carlo. Journal of the ACM, 3(2):73-76, April 1956.
[50] Dilip Arumugam, Peter Henderson, and Pierre-Luc Bacon. An Information-Thewe Theoretic Perspec- tive on Credit Assignment in Reinforcement Learning. arXiv:2103.062. [cs, math], March 2021. arviv: 2103.06224.
[51] Kenny Young. Hindsight Network Credit Assignment: Efficient Credit Assignment in Networks of Discrete Stochastic Units. Proceedings of the AAAI Conference on Artificial Intelligence, 36 (8):8919-892)) 8926, June 2022. Number: 8.
[52] Shixiang Gu. Timothy Lillicrap, Zoubin Ghahramani, Richard E. Turner, and Sergey Levine. Q-Prop: Sample- Efficient Policy Gradient with An Off-Policy Critic. 2017.
[53] Philip S. Thomas and Emma Brunskill. Policy Gradient Methods for Reinforcement Learning with Function Approximation and Action-Dependen Baselines, June 2017. arXiv: 10xx:1706.06645 [es].
[54] Hao Liu", Yihao Feng*. Yi Mao, Dengyong Zhou, Jian Peng, and Qiang Liu. Action-dependent Control Variates for Policy Optimization via Stein Identity. February 2022.
[55] Cathy Wu, Aravind Rajeswaran, Yan Duan, Vikash Kumar, Alexandre M. Bayen, Sham Kakade, Igor Mordatch, and Pieter Abbeel. Variance Reduction for Policy Gradient with Action- Dependent Factorized Baselines. February 2022.
[56] George Tucker, Surya Bhupatiraju, Shixiang Gu, Richard Turner, Zoubin Ghahramani, and Sergey Levine. The Mirage of Detion-Dep Baselines in Reinforcement Learning. In Proceedings of the 35th International Conference on Machine Learning, pages 5015-5024. PMLR, July 2018 ISSN: 2640-3498.
[57] Chris Nota, Philip Thomas, and Bruno C. Da Silva. Posterior Value Functions: Hindsight Baselines for Policy Gradient Methods. In Proceedings of the 38th International Conference on Machine Learning, pages 8238-8247. PMLR, July 2021. ISSN: 2640-3498.
[58] Arthur Guez, Fabio Viola, Theophane Weber, Lars Buesing, Steven Kapturowski, Doina Precup, David Silver, and Nicolas Heess. Value-driven Hindsight Modelling. Advances in Neural Information Processing Systems, 33, 2020.
[59] David Venuto, Elaine Lau, Doina Precup, and Ofir Nachum. Policy Gradients Incorporating the Future. January 2022.
[60] Jiawei Huang and Nan Jiang. From Importance Sampling to Doubly Robust Policy Gradient. In Proceedings of the 37th International Conference on Machine Learning, pages 4434 4443. PMLR, November 2020. ISSN: 2640-3498.
[61] Andrew Y. Ng. Daishi Harada, and Stuart Russell. Policy Invariance Under Reward Trans- formations: Theory and Application to Reward Shaping. In Proceedings of the Sixteenth International Conference on Machine Learning, ICMI. '99, pages 278-287, San Francisco, CA, USA, June 1999 Morgan Kaufmann Publishers Inc
14

---

[62] Jürgen Schmidhuber. Formal Theory of Creativity, Fun, and Intrinsic Motivation (1990-2010), IEEE Transactions on Autonomous Mental Development. 222, September 2010.
[63] Mare Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying Count- Bount-Based Exploration and Intrinsic Motivation. In Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
[64] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. December 2018.
[65] Ofir Marom and Benjamin Rosman. Belief Reward Shaping in Reinforcement Learning. Proceedings of the AAAI Conference on Artificial Intelligence, 32(1), April 2018. Number: 1.
[66] Farzan Memarian, Wonjoon Goo. Rudolf Lioutikov, Scott Niekum, and Ufuk Topcu. Self-Supervised Online Reward Shaping in Sparse-Reward Environments, July 2021. arXiv:2103.04529 [es].
[67] Halit Bener Suay, Tim Brys, Matthew E. Taylor, and Sonia Chernova, Learning from Demon- stration for Shaping through Inverse Reinforcement Learning. In Proceedings of the 2016 International Conference on Autonomous Agents & Multiagent Systems, AAMAS '16, pages 429 429,437, Richland, SC, May 2016. International Foundation for Autonomous Agents and Multiagent Systems.
[68] Yuchen Wu, Melissa Mozifian, and Florian Shkurti. Shaping Rewards for Reinforce- ment Learning with Imperfect Demonstrations using Generative Models, November 2020. 1xXiv:20110012999 [cs].
[69] John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient Estimation Using Stochastic Computation Graphs. In Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015.
[70] Michel Ma, Pierluca D'Oro, Yoshua Bengio, and Pierre-Luc Bacon. Long-Term Credit As- signment via Model-based Temporal Shortcuts. In Deep RL Workshop NeurIPS, October 2021.
[71] Nan Rosemary Ke, Anirudh Goyal ALIAS PARTH GOYAL, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio. Sparse Attentive Backtracking: Temporal Credit Assignment Through Reminding. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems
31, pages 7640-7651, Curran Associates, Inc., 2018.
[72] Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal Value Function Approximators. Proceedings of the 32nd International Conference on Machine Learning. pages 1312-1320. PMLR, June 2015. ISSN: 1938-7228.
[73] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight Experience Replay. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[74] Paulo Rauber, Avinash Ummadisingu, Filipe Mutz, and Jürgen Schmidhuber. Hindsight policy gradients. December 2018
[75] Anirudh Goyal, Philemon Brakel, William Fedus, Soumye Singhal, Timothy Lillicrap, Sergey Levine, Hugo Larochelle, and Yoshua Bengio. Recall Traces: Backtracking Models for Efficient Reinforcement Learning. In International Conference on Learning Representations. arXiv, January 2019. 180400380 [cs, stat].
[76] Juergen Schmidhuber. Reinforcement Learning Upside Down: Don't Predict Rewards Just Map Them to Actions, June 2020. arXiv:1912.02875 [cs].
[77] Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojeiech Jaskowski, and Jürgen Schmidhuber. Training Agents using Upside-Dowe -oole Reinforcement Learning, September 2021. arXiv:1912.02877 [cs].
15

---

[78] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision Transformer: Reinforcement Learning via Sequence Modeling. In Advances in Neural Information Processing Systems, volume 34, Curran Associates, Inc., 2021.
pages 15084-15097.
[79] Michael Janner, Qiyang Li, and Sergey Levine. Offline Reinforcement Learning as One Big Sequence Modeling Problem. In Advances in Neural Information Processing Systems, volume 34, pages 1273-1286. Curran Associates, Inc., 2021.
[80] Richard S. Sutton. Learning to predict by the methods of temporal differences. Machine Learning, 3(1):9-44, August 1988.
[81] Hado van Hasselt, Sephora Madjiheurem, Matteo Hessel, David Silver, André Barreto, and Diana Borsa. Expected Eligibility Traces In Association for the Advancement of Artificial Intelligence. arXiv, February 2021. arXiv:2007.0155 [cs, stat].
[82]
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojeiech Stokowiee, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020.
[83] Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep Variational Information Bottleneck. 2017.
[84] Naftali Tishby, Fernando C. Pereira, and William Bialek. The information bottleneck method, April 2000. arXivpphysies/0044055.
[85] Matthew Hausknecht and Peter Stone. Deep Recurrent Q-Learning for Partially Observ- able MDPs. In Association for the Advancement of Artificial Intelligence. arXiv, 2015.
1XXiv:1507006222 [cs].
[86] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning Latent Dynamies for Planning from Pixels. In Proceedings of the 36th International Conference on Machine Learning, pages 2555-2565. PMLR, May 2019.
ISSN: 2640-3498.
[87] Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping Belief States with Generative Environment Models for RL. In 33rd Conference on Neural Information Processing Systems (NeurIPS 2019). arXiv, June 2019. Number: LtXiv:1906.09233 arXiv:1906.09233 [cs, stat].
[88] Karol Gregor, George Papamakarios, Frederic Besse, Lars Buesing, and Theophane Weber. Temporal Difference Variational Auto-Encoder. 2019.
[89] Matthijs J Spaan. Partially Observable Markov Decision Processes. Reinforcement Learning. page 27.
[90] K. J. Astrom. Optimal Control of Markov decision processes with incomplete state estimation. J. Math. Anal. Applic., 10:174-205, 1965.
[91] Edward Tolman. Cognitive maps in rats and men. Psychological Review, 55:189-2088 2009 1948. Place: US Publisher: American Psychological Association.
[92] Alexander A. Alemi, Ben Poole, Ian Fischer, Joshua V. Dillon, Rif A. Saurous, and Kevin Murphy. Fixing Broken ELBO. In Proceedings of the 35 th International Conference on Machine Learning.. arXiv, February 2018. 1eXXx:1711.0000 [cs, stat].
[93] Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. July 2022.
16

---

[94] Diederik P. Kingma and Max Welling, Auto-Encoding Variational Bayes, arXiv, May 2014. Number: arXiv:1312.6114 14 11Xiv: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 Wanderman- 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.
17

---

Supplementary Materials
Table of Contents
Related work
B Undiscounted infinite-horizon MDPs
3
C Theorems, proofs and additional information for Section 3
4
C.I Contribution coefficients, hindsight distribution and graphical models
4
C.2 Proof Theorem
4
C.3 Different policy gradient estimators leveraging contribution coefficients.
6
C.4 Proof of Proposition 2
7
C.5 Proof Theorem
Proof Theorem
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
E.1
Algorithms
E.2 Dynamic programming sctup
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 enviromment setup
15
E.8
Training details
F
Additional results
F.1
HCA vs HCA+
19
E2
Learned credit assignment features allow for quick adaptation to a change of the reward function.
19
F3
Using returns instead of rewards hindsight information Investigation into the required accuracy of the hindsight models
20
F4
21
Contribution analysis in continuous spaces and POMDPs
21
Spurious contributions continuous state spaces
22
G.2 Deterministic continuous reward functions can lead to excessive variance
22
G.3 Smoothing can alleviate excess variance hy 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
Latent learning
26
Optimal rewarding outcome encodings 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

---

1.2
Extending COCOA to counterfactual interventions
29
Contribution analysis with temporal discounting
30
K Rootstrapping with COCOA
30
L HCA-retura 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
A
Related work
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-Lue [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 [5 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 of future rewards for long trajectories. A common technique to reduce the resulting variance in the policy gradients is to subtract baseline from the sum of future rewards [6 9]. To further reduce the variance, a line of work introduced laae-aactonnddeenddee 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
baselines [52-55]. However, Tucker et al. [56] argues
the baseline, accounting for a larger portion of the variance [1. 3, 34, 37-60g,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 Pradients.Mesnad 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 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 HCA that uses a return-conditioned hindsight distribution to construct 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 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 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 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 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 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 assigmment, while being compatible with discrete actions.
Incorporating hindsight information has 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 '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 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 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()) [80] considers eligibility traces to trade off bias and variance, crediting past actten 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 soo 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.c 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 Σk 7 RR with discount factor 7 E [0, 1], is equivalent to considering the above undiscounted MDP setting, but modifying the state transition probability function p(S(+1 S., A:) such that each sate-action pair has a fixed probability (1 - 7) 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 J.p. J.
The undiscounted infinite-horizon MDP with an absorbing state can also model episodic RL with fixed time horizon. We can include time as an extra feature the states S, and then have a probability of I to transition to the absorbing state when the agent reaches the final time step.
3

---

R
S
S
U
R
S,
A
K
A'
(a)
(b)
Figure 6: (a) Graphical model of the MDP. (b) Graphical model of the MDP, where 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, convert the classical graphical model of an MDP (c.f. Fig. 6a) into a graphical model that incorporates
S s, U' n²). we
the time A into separate node (c.f. Fig. 6b). Here, we rewrite p(U; u' S 8, A a), the probability distribution of a rewarding outcome U2 l time steps later, as p(U" = u' S s a, K = k). By giving K the geometric distribution p(K = k) (1 - 3) for some 3, we can rewrite the infinite sums used in the time-independent contribution coefficients as marginalization over K:
Σ PUB n' |S 8, a) lim 8++1
Σ p(Ux w' 1 ==
(5)
k21
lim
Σ
p(U = S= 8, == a, a.K = k)p(K = k)
k>1
(6)
lim
v
(7)
8-1
Note that this limit is finite, as for k cc, the probability of reaching absorbing state 8o0 and corresponding rewarding outcome Use goes to (and we take * u.o). Via Bayes rule, we then have that
a') = Ees1 P(U = u =
(8)
= S=
p(U' = = -1 p(A S U = n²)
(9)
-U'=uuS==)
T(a S)
will where assume x s) 0 and where dropped the limit of B + in the notation, always take this limit henceforth. Note that when (a s) = O, the hindsight distribution (a s, s,u") is also equal to zero, and hence the hagg--h of the above equation is undefined. However, the middle term using (U = v' S = s, = a) is still well-defines and hence we can use the contribution coefficients for actions where Taa s) 0.
C.2 Proof Theorem 1
We start with lemma showing that the contribution coefficients can be used to estimate the advantage A (s, a) = Q" (s,a)_V* by generalizing Theorem of Harutyunyan [1] towards rewarding outcomes.
Definition (Fully predictive, repeated from main text). rewarding outcome encoding U fully predictive of the reward R, if the following conditional independence condition holds: p" (Rtg

---

So A, = Uk u- = = U = u), where the right-hand side does not depend on the time k.
Lemma 6. For each state-action pair (s, a) with x(a 8) > O, and assuming that u = f(s, r, is fully predictive of the reward (c.f. Definition 3), we have that
as) = =((,,)-
ΣE(a [ + ETNT(s))
Σ w(ss a, RR
(10)
a'EA
with the advantage function A", and the reward function r(s, a) - E[R s.a).
Proof.
We rewrite the undiscounted state-action value function in the limit of a discounting factor + 1_ (the minus sign indicating that we approach I from the left):
Q(s, = lim (T,a(s,, $-1-
3h R
(11)
1
= =r(8,a) + lim
3 p" (R& -
(12)
A-1-
PER #21
as( + lim E ΣΣ
8 (R =r = uul 8, a,
(13)
8-1-
TER uEd F21
= =r(s,a) + lim E
p0p (R r U = uurp (Ux u s,a)
(14)
8-1-
TER uEl K21
=r(,,))+
lim Σ r(u)
p"pp(Up (Up+ = (.
(15)
8-1- wEll
Ex'I 8h 1" (Up n s,
=r(s,)) + lim
En(w)
r(u)
Σ
Su =u
(16)
B-1
Σxz1 3k (U
well
121
where we use the property that is fully predictive of the reward (c.f. Definition 2), and where
define r(s.a) a) & E(R s, a) and r(u) & E[R ui] Using the graphical model of Fig. 6b where
abstract time & (c.f. App. C.I), we have that 1 (R r' 8, = Ex 3 p" (RR =r s,a).
Leveraging Bayes rule, we get
PU u s.a)t(a lim
Σs>1 Sh p" (Up- = u 8, a)
(a 8, U' = =)
lim B-1
Σ②Z1 3 p" (Up. = u T(a 8)
P3(U"= uu
B-1-
(17)
Where we dropped the limit of 3 + 1- in the notation of pa | s, R = r), as we henceforth always consider this limit, Taking everything together, we have that
(s) o)) + ΣΣ
(U,
p" (a
T(a s) r(u)
U = u)
(18)
nEL 121
= rss,,) + ΣΣ
p (U, (w) + 1)r(u)
(19)
= =s(,,)) + ΣΣ Σ
(Ux = u
1)
Σ
p(R = Up = u)r
(20)
wEld 121
rER
0s( + ET~T(A,T)
(w(s. a, U2) + R)
(21)
Subtracting the value function, we get
AT A² = r(8,a) E a(a (s)r(s,a)) + -ETTT(N))
Σ
w(s a, RR
(22)
d'
A21
5

---

Now we are ready to prove Theorem 1.
Proof. Using the policy gradient theorem [26], we have
VoV* (80) (TTT(
ΣΣ> 220 REA
Vor(al S.)AA
(23)
VoT(U
St) Si)
r(St. a) + Σ
w(1,, a, &
(24)
120 DEA
121
[T=T(s0m
V To S.) +
Vor(a S:) Σ w(S,, a, Uetk) Retk
12:0
aEA
(25)
where we used that removing a baseline En'Es s)r(s, d') independent from the actions does not change the policy gradient, and replaced ET~T(en.st) [EnEA Vor(a St)r(Ss,, a)] by its sampled version Er~T(=o,m] log (AA) R.). Hence the policy gradient estimator of Eq. is unbiased.
C.3 Different policy gradient estimators leveraging contribution coefficients.
In this work, contribution coefficients of Eq. 3, as this estimator does not need a separate reward model, and works well for the small action spaces considered in experiments. However, can design other unbiased policy gradient estimators compatible with the same contribution coefficients.
the COCOA estimator of Eq. 4 to estimate policy gradients leveraging the
Harutyunyan et al. [1] introduced the HCA gradient estimator of Eq. 2, which can readily be extended to rewarding outcomes U:
VOVT V
ΣΣ 120
Va (r) Σ w
(26)
This gradient estimator uses a reward model r(s, a) to obtain the rewards corresponding to counter- factual actions, whereas the COCOA estimator (4) only uses the observed rewards Rg.
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.
VOVE
Vo (A, S. + 120
1 M
Σ Vo (a"" | Σ w(S a", U.-k (27)
where we sample M actions from a T(- S.). Importantly, for obtaining an unbiased policy gradient estimate, the actions at should be sampled independently from the actions used in the trajectory of the observed rewards Re+k- Hence, we cannot take the observed A, as sample a" 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 combined with the reward model as in Eq. 26.
One can show that the above policy gradient estimators are unbiased with proof akin the one of Theorem 1.
Comparison of COCOA to using time as a credit assignment heuristic. briefly the following discounted policy gradient:
In Section we discussed
A REINFORCE (so) = Σ, Volog S:)
k-t Rs
(28)
with discount factor 7 E [0, 1]. Note that we specifically put no discounting 7" in front of Vo log m( (A, St), which would required for being an unbiased gradient estimate the dis- counted expected total return, as the above formulation is most used in practice [12, 27). Rewriting

---

the summations reveals that this policy gradient uses time heuristic for credit assignment:
( (50)) = Σ 120 Rs. E
L log St).
(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:
T Vo VologT(A) ( + ΣΣ w(Sk, a, Vo Sk))
V"
Σ Rs 220
(30)
REA
C.4 Proof of Proposition 2
Proof. Proposition 2 assumes that each action sequence {Amm J leads to a unique state s'. Hence, all previous actions can be decoded perfectly from the state s', leading to p" (A, a S s') = 5(a = az), with 5 the indicator function and a the action taken in the trajectory that led to s'. Filling this into the COCOA gradient estimator leads
ö(a Ad)
V (S)O
Ve log S.
Σ
VoT(a| S.) Σ
1) Re+
(31)
Ta St)
VoT(Ar St)
B
(32)
<A, S.))
100
k20
Ve log π (Ac (e))
Ro4k
(33)
120
k20
where we used that E. Va 8) =
C.5
Proof Theorem 3
Proof. Theorem 3 considers the case where the environment only contains reward the final time
step t = T. and where we optimize the policy only on a single (initial) time step = 0. Then, the policy gradients are given by
V VAUT, RT) = Σ Vor(a s)w(s, UT) Rr
(34)
RFFINFORCE (Ao. (A RT) V log (Au RT
(35)
With U either S, R or Z, and the state at 0. As only the last time step can have reward by assumption, and the encoding 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) (Ao = a Su = s, Up = u)
(36)
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 E| (Ao, RT) | Rr] (Sr RT), EIVAV=(S RT) UG, Rr]
V/K (* Rr), V (TG,, Ur, Ryl V Rr). and E[ (TT, Rr) Rr] = V (Rr), after which we can use the law of total variance to prove our theorem.
As ST is fully predictive of Rr, the following conditional independence holds pd (Ao So 8, ST, RT) = (Ao So = s, ST) (c.f. Definition 2). Hence, we have that
L REINFORRC 7" (Ap, Rx) Sx, Rrl
(37)
VOT(a 5) RT (ST TTa
So = ST) Σ p(An =a 1 aEA
(38)
where we used that E. Vor(al | s) = Vo En a(a| 3) = 0.
7

---

Similarly, as ST is fully predictive of U. (c.f. Definition 2), we have that
(A1)) = Σ
p(Sr= = ST S S, Sr = ST, Uf)
(39)
STES
Σ
p(S = ST S, (A S, = ST)
(40)
Using the conditional independence relation y" (St So, U/, RT) p" (St So, U2). following from d-separation in Fig. 6b, this leads us to
E| E V RT) Uy: Rrl
(41)
Voz(a 8)
= Σ Σ pSS = ST = s.UDp =a s, = =ST)
:(a RT
(42)
aEAA STES
Σ p(An =a ss VoT(u Rr
(43)
<aa 8)
aEA
= - V" (* RT)
(44)
Using the same derivation leveraging the fully predictive properties (c.f. Definition 2), we get
E V (U", RT) Ur Rri = 4.. V"Ur, Rr)
(45)
E( Rr) RT] V.V" (RRT)
(46)
Now we can use the law of total variance, which states that V[x] = EMix | YL] + V|E(X Y]]. Hence, we have that
V 0
W = |[[ V REINFORCE Sr] + E(W WWWWWWWRRRRORRCCRCCC Sr]l
(47)
V REINFORCTY
Sr] > V
(48)
as [[ [ Sr] is positive semi definite. Using the same construction for the other pairs, we arrive at
VVVVV < V(V! < V VIV (= < V(VE < V V REINFORCE V (sn)
(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 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. 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.
C.6 Proof of Theorem 4
Proof. This proof follows 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² (u, 8) = >x21 p" (Ux =
So = s), and its action equivalent OF (u, a) = EE>1 D" (Uk- = u So = s, Ao = a). Expanding
O (u, 8) leads to
VOO 8) = Vo
Σ T(a 8)
Σ p(s' 18 ( o
(50)
L
e'ES
VoT(al s)00 8, + Epss 1
(51)
aEA
S'ES
8

---

COCOA-U-group
HCA+ REINFORCE
COCOA-reward
15
14-
Var dB
13-
12-
100
Number of groups per reward value
Figure 7: Less informative rewarding outcome encodings lead to gradient estimators with lower variance. Normalized variance in dB using ground-truth coefficients and random uniform policy, for various gradient estimators on the tree envrionment (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. ng in Section E.6), indicated by the x-axis. COCOA-u-gun 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 o(u, s) Eaed Vor(a s)O"(u, s, o), and p" (S) s So s) the probability of reaching s' starting from in I steps. We have that Ex p" (S: s" |So 8 s)p s' So s²) = p²((++11 = SS = s). Leveraging this relation and recursively applying the above equation leads to
8
VoO"(u, = Σ
(S = s'1 So = s))
(52)
E'ES 1=0
20
p (S 8' So = 8) Σ
(53)
E'ES 1=0
aad
x ES~T(a,))
Σ aEA
Vor(a (s) S.a)
(54)
where in the last step we normalized Ei= pd (S) 8 So 8) with Σ Exes Di 8
So = s), resulting in the state distribution S ~ T(s, T) where S is sampled from trajectories starting from s and following policy x.
Finally, we can rewrite the contribution coefficients as
O (n 0 (u,
w(s,a, a, u)
(55)
which leads to
voO x E (SNT(s,))
Voria (S)((S a, u)O S)
(56)
aEA
where we used that EaEA Var(a 8) = 0, thereby concluding the proof.
D Learning the contribution coefficients
D.1 Proof Proposition 5
Proof. In short, the logits are deterministic function of the state 8, they do not provide any further information on the hindsight distribution and we have that p" (An = a | So = 8, Lo = 1, U'
u') p" (Ao So s, U' = u').

---

We can arrive at this result by observing that p(.Au a So s. = = p(Ag = So =) T (a s), as the policy logits encode the policy distribution. Using this, we have that
p" (An== = a
s. -1, =u)
(57)
p"
=1)
(58)
p 1² = u' s =
p" (U' v' 1 So A, = Lo 1)p((A a So Lo 1)
(59)
En'eA (Ag a', U' u' Sq 8, Lg =1)
V" u' | S0 = s, A. a)p(Ao
So = 5)
(60)
EWEA 1- (U' = u/ Sq = 8, Ao = a', )p( Ao = d' So =)
(A = a, U" = S s)
(61)
P= = v' Su =)
= =An = a So = s, U' v¹)
(62)
D.2 Learning contribution coefficients via contrastive classification.
Instead of learning the hindsight distribution, we can estimate the probability ratio pa (A, = a St = 8, U' = w') wT) s) directly, by leveraging contrastive classification. Proposition 7 shows that by training a binary classifier D(a, 8, uu) to distinguish actions sampled from y" (AAe = a St 8, R' r') versus л (a 8), we directly estimate the probability ratio.
Proposition 7. Consider the contrastive loss
E (.wTT(()))))) [Fo-p (ale,ar) [log D(a, 8, u')] + Eawx(als) [log (1- D(a, w')]],
(63)
and (* a, w') its minimizer. Then the following holds
D (a, w(s,a, u') 1 (a, (,,")
1.
(64)
Proof. Consider a fixed pair 8,u'. We can obtain the discriminator (a 8, v/) that maximizes Easpp (als,w") [log D(a, s, EEanm(als) [log (1 ((a, s, w')) by taking the point-wise derivative
of this objective and equating it to zero:
1
p² (a s, 1')
&a()), T(a 8) 1- D* w()
(65)
822 p³ (aa
D (a,
(66)
1- (" 8, u!)
As for any (a, E RR the function f(x) a log r + b blog(1 - x)) achieves its global maximum on support x [0, at a+3. the above maximum a global maximum. Repeating this argument for all (s, u') pairs concludes the proof.
We can approximate the training objective of D by sampling amm)) s(m), m((m) along the observed trajectories, while leveraging that we have access to the policy n, leading to the following loss
log (ma()) sm)) x (m)))
Σ T(a' | s(m) log (1 D(a', s()) u (m))))
(67)
a'EA
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 (x) = log(1 (1 + exp(-x:) log(1 - (x)) -r - log(1 + exp(-x))
(68)
(69)
The probability ratio D/(1- - D) can then be computed with
(xx) = exp log(1 (1 + x)) + x + log(1 + exp(-x))] = exp(x)
(70)
l-a
10

---

D.3 Successor representations.
we take u' equal to the state s', we can observe that the sum Σx>1 DF = 5 8, a) used in the contribution coefficients (3) is equal to the successor representation M(s, a, 8') introduced by Dayan [30]. Hence, we can leverage temporal differences to learn M(s, 8'), either in 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, s') a(s, s')/LaaAA x( (s. a, s')]
For rewarding outcome encoding U different from S, we can recombine the state- successor representations to obtain the required contribution coefficients using the following derivation.
Σ p" (Uk- u' So 8, Ao a) k00
Σ
Σ
p" (Uk+ u', Sx 8' | Sa 8, Ao = a)
(71)
S'E
Σ
Σ
p(U" = ul 1S' = 5) = s | So = 3, Ao = a)
(72)
Σ
p(U" = u' 15' 15'=s') Σ p p"x(s = 1
So = s. a))
SEES
k20
(73)
D(U''u'' - 5' = s)) a,s')
(74)
S'ES
where in the second line we used that S is fully predictive of U (c.f. Definition 2). Note that
P(U' (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.
E.1
Algorithms
Algorithm shows pscudocode for the COCOA family of algorithms.
E.2 Dynamic programming setup
In order delineate the different policy gradient methods considered in this work, we develop framework to compute ground-truth policy gradients and advantages as well as ground- 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 & after T steps for all states s, i.c. (Sr = 8 So = s) = 1. This assumption is satisfied by the linear key-to-door and tree environment we consider. As helper quantity, define the successor representation:
M(s, a, 8,)
Σ p" -S = 5, An a) k=1
(75)
which captures the cumulative probability over all time steps of reaching state s' when starting from state So = 8, choosing initial action Ao = and following the policy π thereafter. We can compute
M(s, a, s', T) recursively as
M(s, a, 8,)= Σ ES
p(S = S = 8, Ao = a) Σ (1"""""""==""""""""""""""""""""" M(s", d" 1,8 1))
a"EA
(76)
where 1l is the indicator function and (s, a, 8', 0) is initialized to everywhere. Then, the various quantities can be computed exactly as follows.

---

Algorithm COCOA
Require: Initial т. h, episode length T, number of episodes N, number of pretraining episodes M, batch size L, K number of pretraining update steps.
1: for == 1 to M do
D Collect random trajectories {[S,, Ar, Re]T-ol-= from a random policy
2: Sample L trajectories
3: Add trajectories to the buffer
4: end for
5:
for j j= 1 to K do
> Pretrain reward features
6: Sample batch from buffer
7: Train reward features U = f(S to predict rewards R via mean squared error
8:
end for
9:
to N M do Sample trajectories T ((Sr; Ar, of from a
11:
for t = 1 to T do
Update the contribution module
12:
for > t do
13:
Train h(Ae|Se, Uk. Ty) via cross-entropy loss all trajectories
14:
end for
15:
end for
16:
for
= to T do
D Update the policy
17:
for E A do
18:
X. = (a S:) Ex21 w(St, a, R
19:
end for
20: V (8o) E1zo Vo Volog S. EaEs Xa
21: end for
22:
end for
Contribution coeflicients
M(s,a,'iTT -1
w(s,a,u)) EvEA s)M(s,
(77)
where
a() = Σ M(s,a, Σ T(a' 1 Eplr' of
s')
(78)
SES
d'EAA
similar to the successor representations detailed in Section D. To discover the full state space S for an arbitrary environment, we perform depth-first search through the environment transitions (which are deterministic in our setting), and record the encountered states.
Value function
T
(ss) Σ
a(d')s)r(s,d
Σ
Σ
(S 5 =S,T) Σ <(d((s))(((''')
(79)
a'EA
a'EA
Σ
x(d's))((s,''++
Σ ((s,s) Σ
(80)
a'EAA
a'EAA
where r'p(r' 1s
Action-value function
T
Q(s,a) =r(s,a) + Σ
Σ (S = s'10 = s, a) Σ
a(d'ls))rs''dd)
(81)
k=IvES
a'EA
=r(s,a)+ Σ
M(s, (T) Σ x(d'ls))((''a')
(82)
s'ES
a'EA
where r(s',a') -, Ep 7'p(1' "s'p''/''d'' ss',a')
12

---

E.2.2 Computing expected advantages
As next step detail how compute the true expected advantage given potentially imperfect estimator of the contribution coefficient úù the value function V or the action-vait value function Q.
COCOA
E [A*(s, 0() - Σ
v(a' ())((s,a) ETWT(E,))
Σ (0 a, Rs
(83)
d'EA
= = - Σ
T(a' 8))(s,a')+
+ ΣΣ
(UR = u' So - a)
d'EA
d'EUkE
(84)
=r =s(s, - Σ
T(a' s)r(s,a))+
E ((ss,,u)rru' Σ
T(a' 8)M d' u',T)
d'EA
weE
a'EA
(85)
where r(u) = En "p("' u') and ((s a', u', is defined as in E.2.1.
Advantage
(x) ((s.a) -s)))
(86)
where Q the ground truth action-value function obtained following E.2.1.
Q-critic
E
(ss Σ T(a' ((d(ss)((ss))
(87)
a'EAA
E.2.3 Computing expected policy gradients
Finally, we show how we can compute the expected policy gradient, E. [V of a potentially biased gradient estimator, given the expected advantage F [s"
E
V
E Esy T T (m
(88)
k=0
aed
T
p" (S = s|3))
Σ [o
(89)
K=
aEd
Using automatic differentiation obtain Voz(alsu). and the quantity ((s, a, s', T) we defined above, we can then compute the expected policy gradient.
F-
[V[VVV
V Voz (A4(
(90)
aEA
ΣΣ <(ao Mo) ag, s,
VoT [A'
(91)
SES 0
aEAA
Computing the ground-truth policy gradient. compute the ground-truth policy gradient, we can apply the same strategy as for Eq. 90 and replace the expected (possibly biased) advantage E, (so. a)] by the ground-truth advantage function, computed with the ground-truth action-value function (c.f. Section 1.2.
E.3 Bias, variance and SNR metrics
To analyze the quality of the policy gradient estimators, we use the signal-to-noise ratio (SNR), which we further subdivide into variance and bias. higher SNR indicates that we need fewer trajectories to estimate accurate policy gradients, hence reflecting better credit assignment. To obtain meaningful
13

---

scales, we normalize the bias and variance by the norm of the ground truth policy gradient.
1=12
SNR
(92)
E
V V
VoV=
E 1 V - E=18
Variance
(93)
VV
En V
V x V
Bias
(94)
We compute the full expectations E, V and ground-truth policy gradient VoV* by lever- aging our dynamic programming setup, while for the expectation of the squared differences E |V V=V=] we use Monte Carlo sampling with a sample size of 512. We report the metrics in Decibels in all figures.
Focusing on long-term credit assignment. As we are primarily interested in assessing the long- term credit assignment capabilities of the gradient estimators, we report the statistics of the policy gradient estimator corresponding to learning to pick up the key or not. Hence, we compare the SNR, in the outer sum
variance and bias of a partial policy gradient estimator considering only (corresponding to the state with the key) for all considered estimators (c.f. Table 1).
Shadow training. Policy gradients evaluated during training depend on the specific learning trajectory of the agent. Since all methods' policy gradient estimators contain noise, these trajectories are likely different for the different methods. As a result, it is difficult to directly compare the quality of the policy gradient estimators, since it depends on the specific data generated by intermediate policies during training. In order to allow for a controlled comparison between methods independent of the noise introduced by different trajectories, policy is trained with the Q-critic method using ground-truth action-value functions. We can then compute the policy gradients for the various estimators on the same shared data along this learning trajectory without using it to actually train the policy. We use this strategy to generate the results shown in Fig. IB (right), Fig. 2B and Fig. 3C-D.
consider a shadow training setup in which the
E.4 Linear key-to-door environment setup
We simplify the key-to-door environment previously considered by various papers [e.g. 3, 4, 32], to one-dimensional, linear track instead of the original two-dimensional grid world. This version still captures the difficulty of long-term credit assignment but reduces the computational burden allowing us to thoroughly analyze different policy gradient estimators with the aforementioned dynamic programming based setup. The environment is depicted in Fig. 3A. 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 which can either be to the left the right of the agent and which stochastically assume two different reward values. Finally, the agent reaches door which it can open with the key to collect treasure reward.
In our simulations we represent states using a nine-dimensional vector, encoding the relative position on the track floating number, a boolean for each item that could be present at the current location (empty, apple left, apple right, door, key, treasure) as well as boolean indicating whether the agent has the key and a boolean for whether the agent does not have the key.
There are four discrete actions the agent can pick at every time step: pick the key, pick to the left, pick to the right and open the door. Regardless of the chosen action, the agent will advance to the next position in the next time step. If it has correctly picked up the key in the first time step and opened the door in the penultimate time step, it will automatically pick up the treasure, not requiring an additional action.
Hung et al. [4] showed that the signal-to gegnal-tonoiise ratio (SNR) of the REINFORCE policy gradient [6] for solving the main task of picking up the key can be approximated by
|E V y*
|1
SNR RFINF
(95)
C(0))
Rd + Tr VIREINFF no T2
14

---

with V VE the REINFORCE estimator (c.f. Table 1), C(0) reward-indepenee constant, T2 the set of time steps corresponding to the distractor task, and Tr VIDREINFV" no T2] the trace of the covariance matrix of the REINFORCE estimator in an equivalent task setup without distractor rewards. Hence, we can adjust the difficultyy of the task by increasing the number of distractor rewards and their variance. We perform experiments with environments of length L E {20, 40, ...., 100) choosing the reward values such that the total distractor reward remains approximately constant. Concretely, distractor rewards are sampled as ddistractor ~ u(1 2, 18 and the treasure leads to a deterministic reward of Vreausure =
L.
E.5 Reward switching setup
While learning to get the treasure in the key-to-door environment requires long term credit assignment as the agent needs to learn to 1) pickup the key and 2) open the door, learning to stop picking up the treasure does not require long term credit assignment, since the agent simply learn to stop opening the door.
We therefore reuse the linear key-to-door environment of length L 40, with the single difference that we remove the requirement to open the door in order to get the treasure reward. The agent thus needs to perform similar credit assignment to both get the treasure and stop getting the treasure. When applying reward switching, we simply flip the sign of the treasure reward while keeping the distractor reward unchanged.
E.6 Tree environment setup
We parameterize the tree environment by its depth d, the number of actions na determining the branching factor, and the state overlap o,, defined the number of overlapping children from two neighbouring nodes. The states are represented by two integers: in the tree, and the position of the state within that level. The root node has state (i,j) (0, 0). and the state transitions are deterministic and given by + -i- i and j + jmm - 02) + a, where we represent the action by an integer I < Ra. We assign each state-action pair reward r E (- 1, 2}, computed as
representing the current level
r(s,a) ((idx(s) + ap + seed) mod n2,) - n
(96)
with the modulo operator mod, the number of reward values n, = 5, idx(s) unique integer index for the state = (i,j), a large prime number p, and an environment seed.
To introduce rewarding outcome encodings with varying information content, we group state-action pairs corresponding to the same reward value in ng groups:
(idx(s) + ap + seed) mod (n,r)) - nt,/12
(97)
Environment parameters for Figure 2. For the experiment of Fig. 2, we use 6 actions and a depth of 4. We plotted the variance of the COCOA estimator for encodings corresponding to ng = and ng = 32.
E.7 Task interleaving environment setup
We simplify the task interleaving environment of Mesnard et al. [3] in the following way: the enviromment is parameterized by the number of contexts C, the number of objects per context O the maximum number of open contexts B, and finally the dying probability - y. In each context, set of 20 objects is given, and subset of O objects is randomly chosen as the rewarding objects for that context. The task consists in learning, for all contexts, to chose the right objects when presented pair of objects, one rewarding and one non rewarding. While this is reminiscent of a contextual bandit setting, crucially, the agent receives the reward associated to its choice only at a later time, after potentially making new choices for other contexts, and receiving rewards from other contexts, resulting in an interlcaved bandit problem. The agent must thus learn to disentangle the reward associated to each context in order to perform proper credit assignment.
State and action At every timestep, there are actions the agent take: choose right, choose left.
The state is defined by the following variables:
15

---

Linear ReLVig
h(als,u)
Lineor ReLU
RoLUig
L Lincor
u
log π (a|s)
Figure 8: Schematic of the neural network architecture of the hindsight models
A set of open contexts, i.e. contexts where choices have been made but reward associated to those choices has been received
For each open context, a bit (or key) indicating whether the correct choice has been made or not
The context of the current task
A pair of objects currently being presented to the agent, if the current context query room. Each element of the pair is assigned to either the left side or the right side.
Transition and reward The transition dynamics work as follows. the agent in an answer room, then regardless of the action, receives context-dependent reward if it also has the key of the corresponding context (i.e. it has made the correct choice when last encountering query room corresponding to the context). Otherwise, receives 0 reward. Then, the current context removed from the set of open contexts.
If the agent is in query room, then the agent is presented with 2 objects, rewarding and one not, and must choose the side of the rewarding object. Regardless, the agent receives reward at this timestep, and the current context is added to the set of open contexts. Furthermore, if it did make the correct choice, a key associated to the current context is given.
Finally, the next context is sampled. With probability E (where C2 the number of open context), the agent samples uniformly one of the open contexts. Ótherwise, the agent samples uniformly one of the non-open contexts. The objects are then also sampled uniformly out of the O rewarding objects and O unrewarding objects. The rewarding object is placed either on the right or left side with uniform probability, and the unrewarding object placed on the other side.
Crucially, there is probability - 7 of dying, in which case the agent receives the reward but put into terminal state at the next timestep.
Environment parameters For all our experiments, choose C 5, B 3,0 and 0.95.
Visualization of contribution coefficient magnitudes. For the heatmaps shown in Fig 4C, we computed each entry Eae [ (woss, a, rt+*)) where states 8, and rewards r grouped by the context of their corresponding room. We average over all query-answer pairs grouped by contexts excluding query-query or answer-answer pairs and only consider answer rooms with non-zero rewards.
E.8 Training details E.8.1 Architecture
We use separate fully-connected ReLU networks to parameterize the policy, value function, and vation-value function. For the policy we use two hidden layers of size 64, for the value function and action-value function respectively we use a single hidden layer of size 256. For the hindsight model of both HCA and COCOA we found a simple multilayer perceptron with the state and rewarding outcome encoding as inputs to perform poorly. We hypothesize that this is due to the policy dependence of the hindsight distribution creating moving target during learning the
16

---

policy changes. Leveraging Proposition 5, we therefore add the policy logits as an extra input to the hindsight network ease tracking the changing policy. We found good performance using a simple hypernetwork, depicted in Fig. 8 that combines the policy logits with the state and hindsight object inputs through a multiplicative interaction, outputting logit for each possible action. The multiplicative interaction denoted by & consists of a matrix multiplication of a matrix output by the network with the policy logits, and can be interpreted as selecting a combination of policy logits to add to the output channel. In order to allow gating with both positive and negative values, we use gated version of the ReLU nonlinearity in the second layer which computes the difference between the split, rectified input effectively halving the output dimension:
ReLU-ggxx) = (ee.U())))/)) - ReI.U(Zn./..)
(98)
with n the dimension of x. Gating in combination with the multiplicative interaction is useful inductive bias for the hindsight model, since for actions which have zero contribution towards the rewarding outcome the hindsight distribution is equal to the policy. To increase the performance of our HCA+ baseline, we provide both the policy logits and one minus the policy logits to the multiplicatiie interaction.
E.8.2 Optimization
For training of all models we use the A optimizer with default parameters only adjusting the learning rates and clipping the global norm of the policy gradient. We use entropy regularization in combination with epsilon greedy to ensure sufficient exploration to discover the optimal policy, To estimate (action-) value functions we use TD(A) treating each À as a hyperparemeter. For all linear layers we use the default initialization of Haiku [82] where biases are initialized as zero and weights are sample from a truncated Gaussian with standard deviation dimension of the respective layer and.
where Tujputt is the input /uupe
E.8.3 Reward features
Learned reward features should both be fully predictive of the reward (c.f. Theorem 1), and contain as little additional information about the underlying state-action pair as possible (c.f. Theorem 3). We can achieve the former by training network to predict the rewards given state-action pair, and take the penultimate layer as the feature u = [(s, a). For the latter, there exist multiple approaches. When using deterministic encoding u = [(s, a), we can bin the features such that similar features predicting the same reward are grouped together. When using a stochastic encoding p(U S, A) we can impose an information bottleneck on the reward features U, enforcing the encoding to discard as much information about as possible [83, 84). We choose the deterministic encoding approach, as our Dynamic Programming routines require a deterministic encoding." We group the deterministic rewarding outcome encodings by discretizing the reward prediction network up to the penultimate layer.
Architecture. For the neural architecture of the reward prediction network we choose a linear model parameterized in the following way. The input is first multiplicatively transformed by an action-specific mask. The mask is parameterized by a vector of the same dimensionality as the input, and which is squared in order ensure positivity. A ReLU nonlinearity with straight through gradient estimator is then applied to the transformation. Finally, an action -idoonnnnndeddeett readout weight transforms the activations to the prediction of the reward. The mask parameters are initialized 1. Weights of the readout layer are initialized with Gaussian distribution of mean 0 and std 1 where d is the dimension of the input.
va
Loss function. We train the reward network on the mean squared error loss against the reward. To avoid spurious contributions (c.f. 3.3). we encourage the network to learn sparse features that discard information irrelevant to the prediction of the reward, by adding a LI regularization term to all weights up to the penultimate layer with a strength of π),, = 0.001. The readout weights are trained with standard 1.2 regularization with strength M, = 0.03. For the linear key-to-door experiments, we choose (M.,, m),,)) = (0.001, 0.03). For the task interleaving environment, we choose (ML,, ) = (0.05, 0.0003).
All weights regularization are treated as weight decay, with the decay applied after the gradient update.
'In principle, our dynamic programming routines can be extended to allow for probabilistic rewarding outcome encodings and probabilistic environment transitions, which we leave to future work.
17

---

Table 2: The range of values swept over for each hyperparameter in grid search for the linear key-to-door environment and task interleaving environment.
Hyperparameter
Range
1r 1r_agent
{0.003, 0.001, 0.0003, 0.0001}
1r 1r_hindsight
(0.01, 0.003, 0.001]
1r_value
{0.01, 0.003,00000000 0.001
1r-qvalue
{0.01, 0.003, 0.001
Ir lr_features
[0.01. 0.003, 0.001]
@i_lambda_valuu
0.5, 0.9, 00. 0.5, 0.9,
td td_lambda_qvalle
entropy_reg
(0.3, 0.03, 0.01}
Table 3: Hyperparameter values on the linear key-to-door environment. The best performing hyperparameters were identical accross all environment lengths. COCOA-r stands for COCOA-retum, COCOA-F for COCOA-fcaturre
Hyperparameter
COCOA-r COCOA-f HCA+ Q-critic Advantage REINFORCE TrajCV HCA-return
1r_ 1r_agent
0.0003
0.0003
0.0003
0.0003
0.001
0.0003
0.003
0.0001
1r rr_hinndighht
0.003
0.003
0.003
0.001
Ir_ r__value
0.001
1r 1r_qvalue
0.003
0.01
ir_features
0.003
td td_ 1illambda_vee vaalue
td. qvalue_
0.9
0.9
Pretraining. To learn the reward features, we collect the first Breatue mini-batches of episodes in buffer using frozen random policy. We then sample triplets (s, a)) from the buffer and train with full-batch gradient descent using the Adam optimizer over Nieature steps. Once trained, the reward network is frozen, and the masked inputs, which are discretized using the element-wise threshold function 1000. are used train the contribution coefficient as in other COCOA methods. To ensure a fair comparison, other methods are already allowed to train the policy on the first B batches of episodes. For the linear key-to-door experiments, we chose (Breature: Nieumure = (30, 20000). For the task interleaving environment, we chose (Bieature, Nieature) = (90, 30000).
E.8.4 Hyperparameters
Linear key-to-door setup (performance) For all experiments the linear key-to-door vironment, we chose a batch size of 8, while using a batch size of 512 to compute the average performance, SNR, bias and variance metrics. We followed 2-step selection procedure for selecting the hyperparameters: first, we retain the set of hyperparameters for which the environment can be solved for at least 90% all seeds, given large amount of training budget. An environment is considered to be solved for a given seed when the probability of picking up the treasure is above 90%. Then, out of all those hyperparameters, we select the one which maximizes the cumulative amount of treasures picked over 10000 training batches. We used 30 seeds for each set of hyperparameters to identify the best performing ones, then drew 30 fresh seeds for our evaluation. The range considered
Table Hyperparameter values on the task interleaving environment.
Hyperparameter
COcOA-reward COCOA-feature HCA+
Q-critic Advantage REINFORCE TrajCV
1r_ lr_agent
0.0003
0.0003
0.0001
0.0008
0.001
0.001
0.001
1r_ 1r_value 0.001
1r_ hindsight
0.001
0.001
0.001
1r_qvalue
0.003
0.001
ir_features 0.01
0.001
entropy_reg
0.01
0.01
0.01
0.01
0.01
0.01
_d_lambda__alue
1.
td xd_lambda_qvalue
0.9
0.9
18

---

Table The entropy regularization value selected for each environment length of the linear key-to- door environment. The values were obtained by linearly interpolating in log-log space between the best performing entropy regularization strength between environment length 20 and 100.
Environment length
20
40
60
80
100
100 (reward-aliase) -aliasinn
entropy_reg 0.03
0.0187
0.0142
0.0116
0.01
0.0062
for our hyperparameter search can be found in Table 2, and the selected hyperparameters in Table 3. Surprisingly, we found that for the environment length considered, the same hyperparameters were performing best, with the exception of entropy_rege For the final values of entropy for each environment length, we linearly interpolated in log-log space between the best performing values, 0.03 for length 20 and 0.01 for 100. The values can be found in Table 5.
Linear key-to-door setup (shadow training) For measuring the bias and variances of different methods in the shadow training setting, used the best performing hyperparameters found in the performance setting. We kept a batch size for the behavior policy and shadow training, while using a batch size of 512 during evaluation.
Reward switching setup For the reward switching experiment, we chose hyperparameters follow- ing similar selection procedure in the linear key-to-door setup, but in the simplified door-less environment of length 40, without any reward switching. We found very similar hyperparameters to work well despite the absence of door compared to the linear key-to-door setup. However, in order to ensure that noisy methods such as REINFORCE fully converged before the moment of switching the reward, needed to train the models for 60000 training batches before the switch. To stabilize the hindsight model during this long training period, we added coarse gradient norm clipping by 1.. Furthermore, found that slightly decreased learning rate of 0.001 for the Q-critic performed best. Once the best hyperparameters were found, we applied the reward switching to record the speed of adaptation for each algorithm. We kept batch size of 8 for the behavior policy and shadow training, while using a batch size of 512 during evaluation.
Task interleaving experiment We choose a batch size of 8 and train on 10000 batches, while using batch size of 512 to compute the average performance. We used 5 seeds for each set of hyperparameters to identify the best performing ones, then drew 5 fresh seeds for our evaluation. The selected hyperparameters can be found Table 4.
F
Additional results
We perform additional experiments to corroborate the findings of the main text. Specifically, we investigate how reward aliasing affects COCOA-reward and COCOA-feature and show that there is no significant difference in performance between HCA and our simplified version, HCA+.
F. HCA vs HCA+
Our policy gradient estimator for U = S presented in Eq. 4 differs slightly from Eq. 2, the version originally introduced by Harutyunyan et al. [1], as we remove the need for learned reward model r(s,a) a.. We empirically verify that this simplification does not lead to decrease in performance in Fig. 9. We run the longest and hence most difficult version of the linear key-to-door environment considered in our experiments and find no significant difference in performance between our simplified version (HCA+) and the original variant (HCA).
F.2 Learned credit assignment features allow for quick adaptation to a change of the reward function.
Disentangling rewarding outcomes comes with another benefit: When the reward value corresponding to a rewarding outcome changes, e.g. the treasure that was once considered a reward turns out to be poisonous, we only have to relearn the contribution coefficients corresponding to this rewarding outcome. This is in contrast to value-based methods for which potentially many state values are altered. Moreover, once we have access to credit assignment features that encode rewarding outcomes which generalize to the new setting, the contribution coefficients remain invariant to changes in reward contingencies. For example, if we remember that we need to open the door with
19

---

Treasure reward collected
HCA+ HCA
0.6
0.2
10k
Episode
Figure No performance difference between the original HCA method and our modified variant HCA+. In the linear key-to-door environment of length 103 both the original HCA method with an additional learned reward model and our simplified version HCA+ perform similarly terms of performance measured the percentage of treasure reward collected.
Treasure reward collected
0.8
Oracie
COCOA-reward
COCOA-lesture
critic
0.2
Advantage
newand Switch
58k
60k
Episode
Figure 10: COCOA-feature quickly adapts the policy a change of the reward function Percentage of treasure reward collected before and after the change in reward contingencies. COCOA- feature quickly adapts to the new setting, as its credit assignment mechanisms generalize to the new setting, whereas COCOA-reward, Q-critic and Advantage need to adjust their models before appropriate credit assignment can take place.
the key to get to treasure, we can use this knowledge to avoid picking up the key and opening door when the treasure becomes poisonous.
To illustrate these benefits, consider a variant of the linear key-to-door environment. where picking up the key always leads to obtaining the treasure reward, and where the treasure reward abruptly changes be negative after 60k episodes. In this case, COCOA-feature learns to encode the treasure as the same underlying rewarding object before and after the switch, and hence can readily reuse the learned contribution coefficients adapting almost instantly (c.f. Fig. 10). COCOA- contrast needs to encode the poisonous treasure new, previously unknown rewarding outcome. Since only needs to relearn the contribution coefficients for the disentangled subtask of avoiding poison it nevertheless adapts quickly (c.f. Fig. 10). Advantage and Q-critic are both faced with relearning their value functions following the reward switch. In particular, they lack the property of disentangling different rewarding outcomes and all states that eventually lead to the treasure need to be updated.
in
E33
Using returns instead of rewards as hindsight information
In Appendix L, we discussed that HCA-return is biased estimator in many relevant environments, including our key-to-door and iast-intrreeaving environments. This explains its worse performance compared to our COCOA algorithms. To isolate the difference of using returns instead of rewards as
20

---

CocDA-reward CocoA-feature HCA+
Advaniage
HCAvretum
REINFORCE
G-ortiic
TrajCV
Counterfactuallrouuu
Treasure reward collected
0.5
2k
4k
6k
8k
10k
Episode
Figure 11: Performance of COCOA and baselines on the main task of picking up the treasure. measured as the average fraction of treasure rewards collected including the Counterfactual return method.
hindsight information for constructing the contribution coefficients, we introduce the following new baseline:
ΣΣ VaT(al Si) ())(P(e)) St T( S.)
1)
(99)
Similar to HCA-return, this Counterfactual Return variant leverages a hindsight distribution con- ditioned on the return. Different from HCA-retum, we use this hindsight distribution to compute contribution coefficients that can evaluate all counterfactual actions. We prove that this estimator is unbiased.
Figure 11 shows that the performance of HCA-return and Counterfactual return lags far behind the performance of COCOA-reward the key-to-door environment. This is due to the high variance of HCA-return and Counterfactual Return, and the biasedness the former. As the return is combina- tion of all the rewards of trajectory, it cannot be used to disentangle rewarding outcomes, causing the variance of the distractor subtask to spill over to the subtask of picking the treasure.
F44 Investigation into the required accuracy of the hindsight models
To quantify the relation between the approximation quality hindsight models and value functions upon the SNR of the resulting policy gradient estimate, introduce the following experiment. We start from the ground-truth hindsight models and value functions computed with our dynamic programming setup fc.f. E.2), and introduce persistent bias into the output of the model by elementwise multiplying the output with (1 (1+ac), with zero-mean univariate Gaussian noise and a E (0.001, 0.003, 0.01, 0.03, 0.1, 0.3) scalar of which vary the magnitude in the experiment.
Figure 12 shows the SNR the COcOA-reward estimator and the baselines, as function of the perturbance magnitude log a, where we average the results over 30 random seeds. We see that sensitivity of the COCOA-reward estimator to model quality similar to that Q-critic. Furthermore, the SNR of COCOA-rewar remains of better quality compared to Advantage and HCA-state for a wide range of perturbations.
G Contribution analysis in continuous spaces and POMDPs
In Section 3.3 we showed that HCA can suffer from spurious contributions, as state representations need contain detailed features to allow for capable policy. The same level of detail however is detrimental when assigning credit to actions for reaching particular state, since at some resolution almost every action will lead to a slightly different outcome. Measuring the contribution towards specific state ignores that often the same reward could be obtained in a slightly different state, hence overvaluing the importance of past actions. Many commonly used environments, such as pixel-based environments, continuous environments, and partially observable MDPs exhibit this property to large extent due their fine-grained state representations, Here, we take closer look at how spurious contributions arise in continuous environments and Partially Observable MDPs (POMDPs).
21

---

Q-oritic
HCA+
20
COCOA reward
Advantage
SNR (dB)
-20
-40
0.001
Model Perturbation
Figure 12: The improved SNR of the COCOA estimator is robust to imperfect contribution coefficients. The SNR of the COCOA-reward estimator and the baselines, function of the perturbance magnitude log o, where we average the results over 30 random seeds.
G.1 Spurious contributions in continuous state spaces
When using continuous state spaces, p"(S, = s' s,a) represents a probability density function (PDF) instead of probability. A ratio of PDFs p(X x)/p(X x²) can be interpreted likelihood ratio of how likely sample X will be close to x versus x'. Using PDFs, the contribution coefficients with U used by HCA result in
w(s, a, 5') ERZIP (S4+k = 8 S: = 8, A
p a S, s, =
EAZI pr (St4k 8 S: 8))
π (a 88)
(100)
and can be interpreted as likelihood ratio of encountering state 'close' to s', starting from (s, a) versus starting from s and following the policy. The contribution coefficient will be high if (S' s' 8, s.a) has a high probability density at s' conditioned on action a, compared to the other actions a' + a. Hence, the less p" (Sk- = s' s, a) and pd (S, = s' s, a') overlap around s', the higher the contribution coefficient for action a. The variance of the distribution (S, s' s, a). and hence its room for overlap, is determined by how diffuse the environment transitions and policy are. For example, a very peaked policy and nearly deterministic environment transitions leads to sharp distribution p² (Sk- = s' 18 a).
The randomness of the policy and environment is a poor measure for the contribution of an action in obtaining a reward. For example, consider the musician's 's motor system, that learns to both: (i) play the violin and (ii) play the keyboard. We assume a precise control policy and near deterministic environment dynamics, resulting in peaked distributions p" (S), = s' s, a). For playing the violin, a slight change finger position significantly influences the pitch, hence the reward function sharply declines around the target state with perfect pitch. For playing the keyboard, the specific finger position matters
lesser extent, as long as the correct key is pressed, resulting in a relatively flat reward function w.r.t. the precise finger positioning. The state-based contribution coefficients of Eq. 100 result a high contribution violin, this could be good approximation of what we intuitively think of as 'high contribution whereas for the keyboard, this overvalues the importance of the past action in many cases. From this example, it is clear that measuring contributions towards rewarding states in continuous state spaces can lead to spurious contributions, as the contributions mainly depend on how diffuse the policy and environment dynamics are, and ignore that different reward structures can require vastly different contribution measures.
the action taken the trajectory for both tasks. For playing
G.2 Deterministic continuous reward functions can lead excessive variance
Proposition shows that HCA degrades REINFORCE in environments where each action sequence leads to distinct states. In Section G.I, we discussed that continuous state spaces exhibit this property to large extend due to their fine-grained state representation. If each environment state has a unique reward value, COCOA can suffer from similar degradation to HCA. As the rewarding outcome encoding needs to be fully predictive of the reward (c.f. Theorem 1), it needs have a distinct encoding for each unique reward, and hence for each state.
22

---

With continuous reward function this can become a significant issue, if the reward function is deterministic. Due to the fine-graud ane-grained state representation continuous state spaces, almost every action will lead to (slightly) different state. When we have deterministic, continuous reward function, two nearby but distinct states will often lead to nearby but distinct rewards. Hence, to a large extent, the reward contains the information of the underlying state, encoded in its infinite precision of a real value, resulting in COCOA degrading towards the high-variance HCA estimator.
The above problem does not occur when the reward function [R| S, A) is probabilistic, even for continuous rewards. If the variance of p(R S, A) bigger than zero, the variance of p(S. R) is also bigger than zero under mild assumptions. This means that it is not possible to perfectly decode the underlying state using a specific reward value r, and hence COCOA does not degrade towards HCA in this case. Intuitively, different nearby states can lead to the same sampled reward, removing the spurious contributions. In the following section, we will use this insight to alleviate spurious contributions, even for deterministic reward functions.
G.3 Smoothing can alleviate excess variance by trading variance for bias
When each state has a unique reward, an unbiased COCOA estimator degrades to the high-variance HCA estimator, since U needs to be fully predictive of the reward, and hence each state needs a unique rewarding outcome encoding. Here, we propose two ways forward overcome the excess variance of the COCOA estimator in this extreme setting that trade variance for bias.
Rewarding outcome binning. One intuitive strategy to avoid that each state has a unique rewarding outcome encoding U. is to group rewarding outcome encodings corresponding to nearby rewards together, resulting in discrete bins. As the rewarding outcome encodings now contain less details, the resulting COCOA estimator will have lower variance. However, as now the rewarding outcome encoding is not anymore fully predictive of the reward, the COCOA estimator will be biased. An intimately connected strategy is to change the reward function in the environment to a discretized reward function with several bins. Policy gradients in the new discretized environment will not be exactly equal to the policy gradients in the original environment, however for fine discretizations, we would not expect much bias. Similarly for binning the rewarding outcome encodings, when we group few nearby rewarding outcome encodings together, we expect a low bias, but also a low variance reduction. Increasing the amount of rewarding outcomes we group together we further lower the bias, at a cost of an increasing bias, hence creating a bias-variance trade-off.
Stochastic rewarding outcomes. We can generalize the above binning technique towards rewarding outcome encodings that bin rewards stochastically. In Section G.2, we discussed that when the reward function is probabilistic, the excessive variance problem is less pronounced, as different states can lead to the same reward. When dealing with deterministic reward function, can introduce stochasticity in the rewarding outcome encoding to leverage the same principle and reduce variance at the cost of increasing bias. For example, we can introduce the rewarding outcome encoding U ~ - NN(R, o). with /V corresponding to Gaussian distribution. As this rewarding outcome encoding is not fully predictive of the reward, it will introduce bias. We can control this bias- variance trade-off with the variance a: a small sigma corresponds to a sharp distribution, akin to a fine discretization in the above binning strategy, and hence a low bias. Increasing a leads more states that could lead to the same rewarding outcome encoding, hence lowering the variance at the cost of increasing the bias.
Implicit smoothing by noise perturbations or limited capacity networks. Interestingly, the above strategy of defining stochastic rewarding outcomes equivalent to adjusting the training scheme of the hindsight model h(a s, u') by adding noise to the input w'. Here, we take U equal to the (deterministic) reward R, but add noise to it while training the hindsight model. Due the noise, the hindsight model cannot perfectly decode the action a from its input, resulting in the same effects as explicitly using stochastic rewarding outcomes. Adding noise to the input of a neural network is a frequently used regularization technique. Hence, an interesting route to investigate is whether other regularization techniques on the hindsight model, such as limiting its capacity, can result in a bias- vasiinaiee trade-off for HCA and COCOA.
We can apply the same (stochastic) binning technique to hindsight states for HCA, creating a more coarse-grained state representation for backward credit assignment. However, the bias-variance trade-off is more difficult to control for states compared to rewards. The sensitivity of the reward function to the underlying state can be big in some regions of state-space, whereas small in others. Hence, a uniform (stochastic) binning of the state-space
Smoothing hindsight states for HCA.
23

---

will likely result in sub-optimal varranar trade-offs, as it will be too coarse -ersined in some regions causing large bias, whereas too fine-grained other regions causing low variance-reduction. Binning rewards in contrast does not suffer from this issue, as binning is directly performed in space relevant to the task.
Proposition 8 provides further insight on what the optimal smoothing or binning for HCA looks like. Consider the case where we have discrete reward function with not too many distinct values compared to the state space, such that COCOA-rewars is a low-variance gradient estimator. Proposition shows that we can recombine the state- base- hindsight distribution p" (Ao a So = s. = s') into the reward- based hindsight distribution p" (Ap leveraging a smoothing distribution Eves PS' = s' S = s, RR r'). When we consider obtain the low-variance COCOA-reward estimator,
| Sp s. R' r"), by
specific hindsight state s", this means that by considering all states S' that could have lead to the same reward r(s"), instead of only s". Hence, instead of using a uniform stochastic binning strategy with e.g. S Ns"" J). a more optimal binning strategy is to take the reward structure into account through p= So 8, R² r(s"))
Proposition 8.
| (Ao a Su 1 = Σ p" = So = =s, R' =r'pp (Au = =a So = =,, = s') SEES
(101)
Proof. As Ag is d-separated from R' conditioned on S' and So (c.f. Fig. 6b) we can use the implied
conditional independence to prove the proposition:
p" (Ao = a So s, R' = =r))
(102)
=
Σ p¹ (S' =s' An = aa So = = R =
(103)
E'S
Σ p= = =s | = = R, R' = r' )0'p" (Ao = a =5== =,
(104)
FES
Σ p= = ss S = s, R' (Ao = a = =s)=)))
(105)
d'ES
G.4 Continuous action spaces
In the previous section, we discussed how continuous state spaces can lead to spurious contributions resulting high variance for the HCA estimator. Here, we briefly discuss how COCOA can be applied to continuous action spaces, and how Proposition 2 translates to this setting.
We can adjust the COCOA gradient estimator of Eq. 4 towards continuous action spaces by replacing the sum over a' AA by an integral over the action space A'.
Gradient estimator.
V' (so) Vo log T (Ar S.) + / da Vor(a St) Σ w(St, a, Re4k
(106)
In general, computing this integral is intractable. We can approximate the integral by standard numerical integration methods, introducing bias due to the approximation. Alternatively. we introduced in App. C.3 another variant of the COCOA gradient estimator that samples independent actions A' from the policy instead of summing over the whole action space. This variant can readily be applied to continuous action spaces, resulting in
V V
Σ Vo log (At
+
Σ Vo log (aaT S.) Σ w(Sr, ((,S,,"""Ueea U++k) Re+ (107) m
M
120
where we sample M actions independently from ah T( St). but introduces extra variance through the sampling of actions.
This gradient estimator is unbiased,
24

---

Spurious contributions. Akin to discrete action spaces, HCA in continuous action spaces can suffer from spurious contributions when distinct action sequences lead to unique states. In this case, previous actions can be perfectly decoded from the hindsight state, and we have that the probability density function p" (Ag = a So = 8, S' s') is equal to the Dirac delta function 8(a ao), with ao the action taken in the trajectory that led to s'. Substituting this Dirac delta function into the policy gradient estimator of Eq. 106 results in the high- variance REINFORCE estimator. When we use the COCOA estimator of Eq. 107 using action sampling, HCA will even have a higher variance compared to REINFORCE.
G.5 Partially Observable MDPs
In many environments, agents do not have access to the complete state information s, but instead get observations with incomplete information. Partially observable Markov decision processes (POMDPs) formalize this case by augmenting MDPs with an observation space O. Instead of directly observing Markov states, the agent now acquires an observation on at each time step. The probability of observing o in state after performing action a is given by Po s.
A popular strategy in deep RL methods handle partial observability is learning an internal state de that summarizes the observation history h, +o ap, Te J='=0- typically by leveraging recurrent neural networks [85-88]. This internal state is then used as input to policy or value function. This strategy is intimately connected to estimating belief states [89]. The observation history h; provides us with information on what the current underlying Markov state 8t is. We can formalize this by introducing the belief state by. which captures the sufficient statisties for the probability distribution over the Markov states 8 conditioned on the observation history h, Theoretically, such belief states can be computed by doing Bayesian probability calculus using the environment transition model, reward model and observation model:
(
- - = He (0--ha) h2)
(108)
Seminal work has shown that a POMDP can be converted to an MDP, using the belief states as new Markov states instead of the original Markov states [90]. Now the conventional optimal control techniques can be used to solve the belief-state MDP, motivating the use of standard RL. methods designed for MDPs, combined with learning an internal state x.
As the internal state x summarizes the complete observation history h, past actions can be accurately decoded based on h, causing HCA to degrade to the hagi-variance REINFORCE estimator (c.f. Proposition 2). Here, the tension between forming good state representations for enabling capable policies and good representations for backward credit assignment is pronounced clearly. To enable optimal decision-making, good internal state x needs to encode which underlying Markov states the agent most likely occupies, as well as the corresponding uncertainty. To this end, the internal state needs to incorporate information about the full history. However, when using the same internal state for backward credit assignment, this leads to spurious contributions, as previous actions are encoded directly into the internal state.
HCA suffers from spurious contributions in POMDPs.
To make the spurious contributions more tangible, let us consider toy example where the state space S consists of three states v, y and 2. We assume the internal state represents a belief state b = [b', b2), which is sufficient statistic for the belief distribution:
p(S = x b) = p(S = y b) = b², r(S = = 11) b²
(109)
We assume that by is deterministically computed from the history hr, e.g. by an RNN. Now consider that at time step k. our belief state is be = (0.5, 0.25) and we get reward that resulted from the Markov state x. As the belief states are deterministically computed, the distribution p = b' | b, a) is a Dirac delta distribution. Now consider that the action does not influence the belicf distribution over the rewarding state x, but only changes the belief distribution over the non-rewarding states (e.g. B {0.5, 0.23] instead of B₁ (0.5, 0.25) when taking action a' instead of a). As the Dirac delta distributions (B' = 8' b, a) for different a do not overlap, we get a high contribution coefficient for the action ag that was taken in the actual trajectory that led to the belief state b', and a low contribution coefficient for all other actions, even though the actions did not influence the distribution over the rewarding state.
The spurious contributions originate from measuring contributions towards reaching a certain internal belief state, while ignoring that the same reward could be obtained in different belief states well.
25

---

Adapting Proposition 8 towards these internal belief states provides further insight on the difference between HCA and CocOA-reward:
(A = a X = x, Rt' = -r) = Σ
p= = x' X = x. RR ('= = xo= = x, x' = x')
X'EX
(110)
Here, we see that the contribution coefficients of COCOA-reward take into account that the same reward can be obtained while being in different internal belief states x' by averaging over them, while HCA only considers single internal state.
H
Learning contribution coefficients from non-rewarding observations
H.1 Latent learning
Building upon HCA [1], we learn the contribution coefficients (3) by approximating the hindsight distribution p" (a. 8, v'). The quality of this model is crucial for obtaining low-bias gradient estimates. However, its training data scarce, as it is restricted to learn from on-policy data, and rewarding observations in case the reward or rewarding object is used as encoding U. We can make the model that approximates the hindsight distribution less dependent on the policy by providing the policy logits input (c.f. Proposition 5). Enabling COCOA-reward or COCOA-feature to learn from non-rewarding states is more fundamental issue, as in general, the rewarding outcome encodings corresponding to zero rewards do not share any features with those corresponding to non-zero rewards.
Empowering COCOA with latent learning is an interesting direction to make the learning of contribu- tion coefficients more sample efficient. We refer to latent learning as learning useful representations of the environment structure without requiring task information, which we can then leverage for learning new tasks more quickly [91]. In our setting, this implies learning useful representations in the hindsight model without requiring rewarding observations, such that when new rewards are encountered, we can quickly learn the corresponding contribution coefficients, leveraging the existing representations.
H.2 Optimal rewarding outcome encodings for credit assignment.
Theorem shows that the less information a rewarding outcome encoding U contains, the lower the variance of the corresponding COCOA gradient estimator (4). Latent learning on the other hand considers the sample-efficienyy and corresponding bias of the learned contribution coefficients: when the hindsight model can learn useful representations with encodings U corresponding to zero rewards, it can leverage those representations to quickly learn the contribution coefficients for rewarding outcome encodings with non-zero rewards, requiring less training data to achieve a low bias.
These two requirements on the rewarding outcome encoding U are often in conflict. To obtain a low-variance gradient estimator, U should retain little information possible while still being predictive of the reward. To enable latent learning for sample-efficient learning of the contribution coefficients, the hindsight model needs pick up on recurring structure in the environment, requiring keeping as much information as possible in to uncover the structural regularities. Using the state as rewarding outcome encoding beneficial for enabling latent learning, as contains rich information on the environment structure, but results in spurious contributions and a resulting high variance. Using the rewards or rewarding object as rewarding outcome encoding removes spurious contributions resulting in low variance, but renders latent learning difficult.
A way out of these conflicting pressures for an optimal rewarding outcome encoding is rewards or rewarding objects as the optimal encoding for low variance, but extract these contribution coefficients from models that allow for sample-efficient, latent learning. One possible strategy is to learn probabilistic world models [20, 21, 88) which can be done using both non-rewarding and rewarding observations, and those to approximate the contribution coefficients of Eq. 3. Another strategy that we will explore in more depth is to learn hindsight models based on the state as rewarding outcome encoding enable latent learning, but then recombine those learned hindsight models to obtain contribution coefficients using the reward or rewarding object as U.
H.3 Counterfactual reasoning on rewarding states
HCA results in spurious contributions because it computes contributions towards reaching a precise rewarding state, while ignoring that the same reward could be obtained in other (nearby) states.
26

---

S
S'
Z
U
R
A
K
A'
Figure 13: Schematic of the graphical model used in variational information bottleneck approach.
Proposition 9 (a generalization of Proposition 8), shows that we can reduce the spurious contributions of the state-based contribution coefficients by leveraging counterfactual reasoning on rewarding states. Here, we obtain contribution coefficients for a certain rewarding outcome encoding (e.g. the rewarding object) by considering which other states s' could lead to the same rewarding outcome, and averaging over the corresponding coefficients w(s, a, 8').
Proposition Assuming S' is fully predictive O', we have that
w(s, al) (S' = So s, U' ='''"'=''(ssspp w')w(s,,,s)) S'ES
(111)
The proof follows the same technique as that of Proposition 8. Proposition 8) shows that is possible
learn state-based contribution coefficients, enabling latent learning. and obtain low-variance COCOA estimator by recombining the state-based contribution coefficients into coefficients with less spurious contributions, if we have access to generative model p" (S' = So 8, U' = w'). This model embodies the counterfactual reasoning rewarding states: 'which other states s are likely considering I am currently u' and visited state somewhere in the past'.
In general, learning this generative model is as difficult or more difficult than approximating the hindsight distribution p" (a u'), as it requires rewarding observations as training data. Hence, is not possible to directly use Proposition to combine latent learning with low-variance estimators. In the following section, we propose possible way forward circumvent this issue.
H.4 Learning credit assignment representations with an information bottleneck
Here, we outline strategy where learn latent representation Z that retains useful information the underlying states S, and crucially has a latent space structure such that p" (Z' = S 8, U' w') easy approximate. Then, leveraging Proposition (and replacing S' by Z') allows to learn a hindsight representation based on Z, enabling latent learning, while reducing the spurious contributions by counterfactual reasoning with p= 5, u').
To achieve this, we use the Variational Information Bottleneck approach [83, 92], closely related to the 3 Variational Autoencoder [93]. Fig. 13 shows the graphical model with the relations between the various variables and encoding: we learn a probabilistic encoding p(Z S, A; 6) parameterized by and assume that the latent variable Z is fully predictive of the rewarding outcome encoding U and reward R. We aim to maximize the mutual information I(Z'; S' A' S, U") under some information bottleneck. We condition S and U², end up later with decoder and prior variational model that we can combine with Proposition 9.
Following rate-distortion theory, Alemi et al. [92] consider the following tractable variational bounds on the mutual information
D-D<I A' |S,U') < Rate
(112)
with entropy H. distortion D and rate defined follows:
H = S..U A' S(U')
(113)
D = -Es.U' Es / de' p(z' ss a logg'' I ('.,)'))))))
(114)
Rate = Es. [Est.Ars. s. [DLL (p(z s s''''; |s,''))
(115)
where the decoder q(s', 8 (b) is variational approximation to p(s' a² x'), and the marginal g(2' 8, v) variational approximation to the true marginal p(=' 8, u'). The
distortion D quantifies how well we decode the state-action pair (s', a') from the encoding
27

---

z', by using a decoder a(s',a' s, =/:²/1)) 1)) parameterized by v. The distortion is reminiscent of an autoencoder loss, and hence encourages the encoding p(æ' s's a') to retain as much information as possible on the state-action pair ,'), The rate measures the average -I--vieeeeence between the encoder and variational marginal. In information theory, this rate measures the extra number of bits (or nats) required to encode samples from Z' with an optimal code designed for the variational marginal q(' s, :)) o).
We can use the rate to impose an information bottleneck on Z'. If we constrain the rate to Rate S with some positive constant, we restrict the amount of information Z' can encode about (S', AS", as p(z' s',, a²; 8) needs to remain close to the marginal g(z' s, u Φ), quantified by the KL-divergence. We can maximize the mutual information I S' A'' | S. U²) under the information bottleneck by minimizing the distortion under this rate constraint. Instead of imposing a fixed constraint on the rate, we can combine the rate and distortion into an unconstrained optimization problem by leveraging the Lagrange formulation
min D + RRat 0.0.0
(116)
Here, the 3 parameter determines the strength of the information bottleneck. This formulation equivalent to the Variational Autoencoder [93], and for 1 we recover the Variational Autoencoder [94].
To understand why this information bottleneck approach is useful to learn credit assignment represen- tations Z, we examine the rate in more detail. We can rewrite the rate as
Rate = Es [DKL (p(z' I s, ('la((2 s, u,
())] +
(117)
E [DKL (p(s,, a' s, 2l)''s a' | w'))
(118)
Hence, optimizing the information bottleneck objective (116) w.r.t. 6 fits the variational marginal q(z' s, u'; v'; to the true marginal p(=' s, u') induced by the encoder p(2' s a': 0). Proposition uses this true marginal to recombine coefficients based on Z into coefficients based on U. Hence, by optimizing the information bottleneck objective, we learn a model q(2' s. u'; Ø) that approximates the true marginal, which we then can use to obtain contribution coefficients with less spurious contributions by leveraging Proposition Furthermore, minimizing the rate w.r.t. 0 shapes the latent space of Z' such that the true marginal p(2' | 88 u') moves closer towards the variational marginal 4(21 8, ())
In summary, the outlined information bottleneck approach for learning a credit assignment representa- tion Z is a promising way forward merge the powers of latent learning with a low-variance COCOA gradient estimator. The distortion and rate terms in the information bottleneck objective of Eq. 116 represent a trade-off parameterized by 3. Minimizing the distortion results in detailed encoding Z' with high mutual information with the state-action pair (S', which can be leveraged for latent learning. Minimizing the rate shapes the latent space of Z' in such way that the true marginal p(a' v²) can be accurately approximated within the variational family of q(=' | | s, Φ), and fits the parameters & resulting in an accurate marginal model. We can then leverage the variational marginal q(=' | u'; Φ) to perform counterfactual reasoning on the rewarding state encodings z' according to Proposition 9, resulting a low-variance COCOA-estimant
I Contribution analysis and causality
I.1 Causal interpretation of COCOA
COCOA is closely connected to causality theory [42], where the contribution coefficients (3) cor- respond to performing Do- interventions on the causal graph to estimate their effect on future rewards. To formalize this connection with causality, we need to use a new set of tools beyond conditional probabilities, causation is in general not the same as correlation. We start with rep- resenting the MDP combined with the policy as directed acyclic graphical model (c.f. Fig. 6a in App. C. In causal reasoning, we have two different 'modes of operation'. On the one hand, we can use observational data, corresponding to 'observing' the states of the nodes in the graphical model, which is compatible with conditional probabilities measuring correlations. On the other hand, we can perform interventions on the graphical model, where we manually set a node, e.g. A, to a specific value a independent from its parents S,, and see what the influence on the probability distributions of other nodes in the graph, c.g. R4+1. These interventions are formalized with do-calculus [42],
28

---

R
S
A
S
N
N
y
Figure 14: Structural causal model (SCM) of the MDP. Squares represent deterministic functions and circles random variables.
where we denote an intervention of putting a node Vi equal to e as Do = =)), and can be used to investigate causal relations.
Using the graphical model of Fig. 6b that abstracts time, can use do- fo-iinerrventios to quantify the causal contribution of an action A, = taken in state S. = 8 upon reaching the rewarding outcome U u' in the future as,
(- = o' S = -; s, =a)
WD,(s, (ss,a,u)) ΣaEA क(Ō s)p" u' 1 = =, Do(A,
1.
(119)
As conditioning on S satisfies the backdoor criterion [42] for U' w.r.t A, the interventional distribution y" (U' u' St 8, Do(A = a)) is equal to the observational distribution
p" (U' u' St = s, A, a). Hence, the causal contribution coefficients of Eq. 119 equal to the contribution coefficients of Eq. 3 used by COCOA.
1,2 Extending COCOA to counterfactual interventions
Within causality theory, counterfuctual reasoning goes one step further than causal reasoning, by incorporating the hindsight knowledge of the external state of the world in its reasoning. Applied to COCOA, this more advanced counterfactual reasoning would evaluate the query: 'How does taking action a influence the probability of reaching rewarding outcome, compared to taking alternative actions a', given everything else remains the same'. To formalize the difference between causal and counterfactual reasoning, we need to convert the causal DAG of Figure 6a into structural causal model (SCM), as shown in Figure 14. The SCM expresses all conditional distributions as deterministic functions with independent noise variables N, akin the reparameterization trick [94]. In causal reasoning, we perform do-interventions on nodes, which is equivalent to cutting all incoming edges to node. To compute the resulting probabilities, we still use the prior distribution over the noise variables N. Counterfactual reasoning goes one step further. First, it infers the posterior probability p" ([N) T 7). with { the set of all noise variables. given the observed trajectory т. Then it performs a Do-intervention as in causal reasoning, However now, to compute the resulting probabilities on the nodes, it uses the posterior noise distribution combined with the modified graph.
One possible strategy to estimate contributions using the above counterfactual reasoning is to explicitly estimate the posterior noise distribution and combine it with forward dynamics models to obtain the counterfactual probability of reaching specific rewarding outcomes. Leveraging the work of Buesing et al. [43] is a promising starting point for this direction of future work. Alternatively, we avoid explicitly modeling the posterior noise distribution, by leveraging the hindsight distribution combined with the work of Mesnard et al. [3]. Here, we aim to learn a parameterized approximation the counterfactual hindsight distribution p (A, = a S; s, indicates the counterfactual distribution incorporating the noise posterior. Building upon the approach of Mesnard al. [3], we can learn h(a s', s, 4 (T)) to approximate the counterfactual hindsight distribution, with (T) a summary statistic trained to encapsulate the information of the posterior noise distribution P([N) T= 7). Mesnard et al. [3] show that such a summary statistic (F(T) can be used to amortize the posterior noise estimation if it satisfies the following two conditions: (i) it needs to provide useful information for predicting the counterfactual hindsight distribution and (ii) it needs to be independent from the action A,. We can achieve both characteristics by (i) training (a s',, (T)) on the hindsight action classification task and backpropagating the gradients to (T)., and (ii) training 4 on independence maximization loss (s),s), which is minimized iff
u'), where the 7-subscript
29

---

Table 6: Comparison of discounted policy gradient estimators
Method
Policy gradient estimator (VoV*
REINFORCE
E 1200 ) V'V log π (Ae St) Σ k20 2 Ri+K
Advantage
E 120 Vo log ग (Ae) St) E k20 * RR+* - V5.)
Q-critic
E120? 7 EAEA Va St (5,
COCOA
Ex V g log T (At
R.) + EoEA VaT(a S,) k21 y wn a, (1+k)
HCA+
E120 log ग (A
St) + DOEA VoT(a Si) Ex>1 w (Sc. a Ri4h
A, and D, are conditionally independent given S,. An example to minimize the KL divergence between =[a; 8t) and p(az 8t, (x(r)) where the latter be approximated by training classifier ((a st. (x))). Leveraging this approach to extend COCOA towards counterfactual interventions an exciting direction for future research.
Contribution analysis with temporal discounting
As discussed in App. B, we can implicitly incorporate discounting into the COCOA framework by adjusting the transition probabilities to have fixed probability of (1 7) of transitioning to the absorbing state Soe at each time step.
We can also readily incorporate explicit time discounting into the COCOA framework, which we discuss here. We consider now a discounted MDP defined as the tuple AS p. Pr, 7), with discount factor E [0. and (S,A,) p: Pr) defined in Section 2. The discounted value function VT( (s)
ET~T(a.)) Ee >" RR] and action value function (s, a) ET~T(s.a.x) E> 7 Re] are the expected discounted return when starting from state s, or state s and action a respectively. Table
shows the policy gradient estimators of V (s) for REINFORCE, Advantage and Q-critic.
In a discounted MDP, matters at which point time we reach rewarding outcome u, as the corresponding rewards are discounted. Hence, we adjust the contribution coefficients to
Wy (s, ass,a,u u') = Σς>1 ("U+
(U++k = v | s.
(120)
Σ121 (U1+: =u S.
(A St 3, = v¹)
(121)
x(a s)
Here, we define the discounted hindsight distribution p. (A = a S. 8. w') similarly to the undiscounted hindsight distribution explained in App. C.I, but now using different probability distribution on the time steps k: Pg(K k) = (1- -33) where we take B = We can readily extend Theorems 1-4 to the explicit discounted setting, by taking 3 = 7 instead of the limit of + and using the discounted COCOA policy gradient estimator shown in Table 6.
To approximate the discounted hindsight distribution p., (A, S. s, S² = u'), we need to incorporate the temporal discounting into the classification Pross-entropy loss:
L, E
ΣΣ
*CE(,, St, U(a)), Ai)) 120
(122)
with CE the cross-entropy loss, (. | Sr, Up+R) the classification model that approximates the discounted hindsight distribution, and ō(a Ar) a one-hot encoding of Aμ-
K
Bootstrapping with COCOA
Here, we show how COCOA can be combined with n-step returns, and we make a correction to Theorem of Harutyunyan al. [1] which considers n-step returns for HCA.
Consider the graphical model of Fig. 15a where we model time, K, as separate node in the graphical model (c.f. App. C. To model n-step returns, we now define the following prior probability on K,
30

---

S
S'
U
R
S
S
Z'
A
K
A'
A
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 f::
{ i 2 1 - -n
-
Pn. P [K = =k)=
(123)
else
Using Pã.g as the probability distribution induced by this graphical model, we have that
PT =u s,A a)
Σ [U =u,k = k =
(124)
Pnnsp(K = k)) (U" = =
S = 8, A a. k-
(125)
1-3 3 3(1 (n-1)
Σ ph u So 8, Ao = a)
(126)
We introduce the following contribution coefficients that will use for n-step returns
Wn,3(8, a.) = カリー 8tp (Us = (U== So = s, Ao Ao= = Σ 3h (U+ u' [So = 8)
(127)
(U ul S =, a)
-1 (A a ||SSS = s,U' = =d)
(128)
Pho (U' = vl -
x(a s)
p" (Sn 8 So Ao a) pd (Sn = So =
-1 nd (A a 5 = Sn = 8)) -1 (129)
Wi (8)
(a s)
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 7 1-
Theorem 10. Consider state s and action o for which it holds that Ta 8) > 0 and take 3 equal to the discount factor y E [0, 1]. Furthermore, assume that the rewarding outcome encoding f(s, a, r) is fully predictive of the reward (c.f. Definition 2). Then the advantage (r a) = Q" (s, a) (s*
equal
(s, = r(s,a) - 15 (s) ET(M.))
Σ
7 ww as, Un)R + (s,,,
with r(8, a) the reward model and (*)(() =. M(a s))
Proof. We start with the action-value function Q and will subtract the value V* to obtain the result on advantage function.
Q(s, a) ET~T(ANO.F)
7
& Rk
(130)
= +
Σ
p = r' s)))r'
(131)
MER k21
31

---

= + Σ n-11
(Rx= u' 1,,0)''
(132)
MER w'Ell k=1
7" Σ
1" (Sn = =8 (1(88)\=(((
(133)
d'ES
= +
ΣΣ
u' ssa''t
(134)
MER w'ell km
7 Σ
(S, = =s(ss))VV((s'
(135)
3'ES
n-1 Σ
=r(s,a)+ + Σr(u' r(u')
Epr(s, = =s1s,)) (s')
E'
kml
(136)
n-1 1
-k H p" (Up =u
= as( +
Σ r(u')
Σ
pm = u' |s)
(137)
Σk * uu
N E'
k=1
(Sn = 8' s,a) (T((s)
711 Σ
p" (S, = 8)
(138)
(S, s 8)
S'ES
n-1
=r(8,a) +
Σ r(u') Σ
(U = (s) v') +1)
(139)
W'Ell
k=1
"
Σ
(" s (8, 15)
(140)
3'ES
(141)
where we use that U' is fully predictive of the reward R', and define r(w') EVER U = =''rr' By subtracting the value function, we get
=
A(s,a) = r(s - (s)+ + Σ r(u')
Σ
= = (s),n,((,,a,)))
(142)
w'Ell k=1
+7" Σ (Sn = 8 (5,,
(143)
8'88
= =r(s,a) - r (s) + ET(a.T)
Σ
7 Wn, s(, a, ss, Ux) + ws
(144)
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- state-based contribution coefficients wn (,,ss, a, s') to bootstrap the value function into the n-step return, as the value function requires Markov state s' as input instead of a rewarding outcome encoding u'. Unfortunately, these state-based contribution coefficients will suffer from spurious contributions, akin to HCA, introducing 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 n-step return, while using rewarding-outcome contribution coefficients w(s, a, u²) instead of state-based contribution coefficients (,, 8')
Learning the contribution coefficients. We can learn the contribution coefficients Wa,., 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 discount 7 & 1, we need to take this discount factor into account in the training distribution or loss function (c.f. App. J).
32

---

Correction to Theorem 7 of Harulyunyan et al. [1]. Harutyunyan et al. [1] propose a theorem similar 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 different prior probability distribution on K':
8 (1
B)
if S k n-nn
HCA Bn.
(K = k) 8m-1
if = n
(145)
else
The graphical model combined with the distribution on defines the hindsight distribution Pa, (A = = =, S, ='). The second difference is the specific Q-value estimator Haru-
tyunyan al. [1] propose. They use the hindsight distribution PH, (A = a S = S' s') in front of the value function (c.f. Theorem 10), which considers that s' can be reached any time step HCA (k), whereas Theorem 10 uses wn (s, 8') which considers that s' is reached exactly at k ~ Pn. 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 simple fix. For the interested reader, we briefly explain the error. One indication of the problem is that for 3 - 1, all the probability mass of P (K k) is concentrated S S' s') considers only hindsight states s' encountered at time k 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, conditional independence relation is assumed that does not hold. The authors introduce a helper variable Z defined on the state space S, with conditional distribution
at k
hence the corresponding hindsight distribution PR.8.HCA (A
6(2 = 5') if f k 1 (z Z = = S' = =s) ======/--- (= =
(146)
8')
if k=
with the normalized discounted visit distribution d' = Ex
PS = S. 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" (Au =
=)
S. = S''=
(147)
- However, Fig. 15b shows that S' is collider on the path Ao - S' - K Z. Hence, by conditioning on S" we open this collider path, making Ao dependent on conditioned on So 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 Ag, beyond 5 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:
x (AA S.))
Σ Vo log T(A. 1 p (A. S.. 2,) 120
Z+
(148)
When comparing this estimator with COCOA-rewarts we see two important differences: (i) HCA- return uses 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 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
33

---

starting from state it holds that (a s, O. This restrictive assumption requires that for each observed state-action pair (st, at) along trajectory, all counterfactual returns Z resulting from counterfactual trajectory starting from (not including (z) result in p" (a, | 8t, 2) > O. This implies that all returns (or rewarding states) reachable from 8, should also be reachable from 0t).
Consider the following bandit setting as a simple example where the above assumption is not satisfied. The bandit has two arms, with reward of and -2, and a policy probability & and respectively. The advantage for both arms and -2 Applying eq. from Harutyunyan et al. results in (s, a1) = (1 - 3) and (s, a2) = -2(1 1/3) = -4/3. This shows that the needed assumptions for unbiased HCA-return estimator be violated even in simple bandit settings.
M Additional details
M.1
Author contributions
This paper was 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 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- bbane ground-truth methods, debugging, uning and running experiments, feedback to the manuscript, writing of supplementary materials.
Nathaniel Daw. Regular project meetings, conceptual input and feedback for method and ex- perimental design, connecting the project to existing literature, feedback to the manuscript and figures.
Gregory Wayne. Senior project supervision, conceptualising of the project idea, conceptual develop- ment 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, TPUV3-8 and 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 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 used wandb [96, MIT license] and for the generation of plots we used plotly [97, MIT license].
34


Reviews:
[{"review_id": "sp5yitjQez", "reviewer": "Reviewer_GAJb", "summary": "This paper presents family of algorithms namely Counterfactual Contribution Analysis (COCOA), which is based on measuring contribution of an action by asking counterfactual question. The new measure of contribution performs better than existing methods such as HCA in terms of lower variance. They define rewarding outcome as function of state, action and reward. They compute contribution coefficients which can be used to learn a policy with COCOA policy gradient estimator. ", "questions": "1. Could you please highlight how expensive it is to compute these contribution coefficients.\n2. Policy gradient methods have shown great performance in continuous state and continuous action MDPs. How can the coefficients be computed in such a setting? Or the proposed method is limited to tabular setting?\n3. In many real world applications, we only have sparse rewards, how does the performance of these estimators vary between dense and sparse reward scenarios?", "limitations": "-", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors show that the presented COCOA policy gradient estimator can result in lower variance than existing method such as HCA, which suffers from spurious contributions which can be alleviated by using less informative encodings.\n2. The authors demonstrate demonstrate the performance of COCOA compared to standard baselines in key-to-door environment. They further present improvement in sample efficiency due to favorable bias-variance trade off\n3. The way COCOA is learning the correct distanglement enables long-term credit assignment better than existing methods.", "weaknesses": "Compared to policy gradient methods like REINFORCE and Advantage, the COCOA estimators seem to be very computationally expensive.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review and instructive suggestions that have helped to improve the paper. Below we provide responses to your raised questions and concerns. \n\n## Computational cost\n\nWe provide the wall clock time of all methods after 3000 update steps on the linear key-to-door environment (averaged over three seeds) on our hardware (Nvidia RTX 3090) in the following table.\n\n| Method        | Wall clock time (seconds) |\n|---------------|---------------------------|\n| COCOA-reward  | 343.67                    |\n| COCOA-feature | 434.67                    |\n| HCA+          | 319.33                    |\n| Advantage     | 314.00                    |\n| Q-critic      | 319.00                    |\n| REINFORCE     | 330.33                    |\n| TrajCV        | 328.00                    |\n\nDuring training, value based methods such as Q-critic and Advantage have a complexity linear in the environment length, considering the amount of evaluations of the policy network and value network. The COCOA and HCA methods have a quadratic complexity in the environment length. Despite this, except for COCOA-feature, all methods have similar runtime, suggesting the computational bottleneck lies in the environment simulation in our case.\n\n## Continuous state and action spaces\nIn appendix F.1 and F.4 in the supplementary materials we expand upon this question of how COCOA can be applied to continuous state and action spaces. In short, the contribution coefficients of Equation 3 readily extend towards continuous state and action spaces, by using probability density functions $p(a\\mid s, u)$ and $p(u \\mid s,a)$ instead of probability distributions as in the discrete case. To construct a policy gradient estimator, we either need to replace the sum over actions in Equation 4 by an integral over the action space, or sample from it resulting in \n$$\\sum_{t\\geq0} \\nabla_\\theta \\log \\pi(A_t \\mid S_t) R_{t} + \\frac{1}{M} \\sum_m \\nabla_\\theta \\log \\pi(a^m\\mid S_t)\\sum_{k=1}^{\\infty} w(S_{t}, a^m, U_{t+k}) R_{t+k}$$\nwith $M$ the number of samples. We refer the reviewer to Appendix F for a more detailed discussion. \n\n## Difference between sparse and dense reward scenarios\nIn sparse reward settings, there are two main challenges: (i) efficient exploration of the environment to discover the sparse rewards and (ii) making efficient use of the sparse rewards to update the policy. These exploration and credit assignment challenges are mostly orthogonal, and in our work and experiments we focus on the credit assignment problem, assuming that the exploration challenge is solved to a sufficient extend. The COCOA methods particularly excel when long-term credit assignment is needed, as the contribution coefficients can reduce the variance of the policy gradient estimator without the need to resort to temporal discounting. The key-to-door environment investigates the credit assignment problem with sparse treasure rewards, with the added difficulty of a distractor task introducing noise. Hence, we expect that our key-to-door results translate to other sparse reward settings where the exploration problem is solved sufficiently. In a dense reward setting where no long-term credit assignment is needed, we expect that the Advantage and Q-net baselines will move closer towards the performance of our COCOA methods, as in this case weaker credit assignment techniques are likely sufficient in combination with temporal discounting."}, {"review_id": "PWLtwy5aI5", "reviewer": "Reviewer_r14e", "summary": "This paper proposes Counterfactual Contribution Analysis (COCOA) to improve credit assignment in the reinforcement learning problem by building upon Hindsight Credit Assignment.", "questions": "A space is missing between policy and [15-17] on line 41.", "limitations": "Seems okay to me.", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This paper is written in a clear manner and allows non-expert audiences to understand its difference and contribution in contrast to previous work.", "weaknesses": "I do not feel that I understand enough about this field to critique this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for his review and pointing us to the typo. "}, {"review_id": "iLIjctVqyT", "reviewer": "Reviewer_V1Ut", "summary": "The paper introduces a novel credit assignment algorithm for reinforcement learning known as COCOA (Counterfactual Contribution Analysis). The proposed method extends the concept of Hindsight Credit Assignment (HCA) and aims at making the learning process more sample efficient. By leveraging counterfactual reasoning, COCOA measures the contribution of actions upon obtaining subsequent reward, thereby achieving improved credit assignment. The paper provides an extensive theoretical analysis of COCOA while addressing its potential advantages and limitations.", "questions": "- How does COCOA compare to similar SOTA credit assignment methods?\n- What is the expected influence of imperfect modelling on the performance of COCOA?\n- When employing COCOA, what are best practices for modelling action contributions for ‘rewarding outcomes’ and how does this impact the learning process?\n", "limitations": "- Uncertain scalability to intricate environments.\n- Examination on how the choice of the dynamics model impacts COCOA's performance (arguably this is kept open in the discussion section).\n", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The proposal of a novel approach to credit assignment in RL.\n- Comprehensive theoretical derivations and analysis that support the argument.\n- Improvement in policy gradient estimators and progression in long-term credit assignment by distinguishing rewarding outcomes.\n", "weaknesses": "- The absence of empirical comparisons with related methods, particularly RUDDER which is cited in the paper.\n- Unclear scalability of the proposed method in more complex environments and restricted quality of the inverse dynamics model.\n- Limited empirical evaluation and ablation studies, i.e. examining the influence of the dynamics model on performance.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the encouraging review and for the useful feedback that we incorporated in a revised version of our paper. We reply to your questions point-by-point below.\n\n## 1. Comparison to SOTA credit assignment methods\nWe now include HCA-return in our experiments, as well as a new method Counterfactual Return that uses returns to construct contribution coefficients. Finally, we include Trajectory-wise CV (Cheng, Yan and Boots, 2019), a state-of-the-art control variate method that incorporates hindsight information to construct a near-optimal control variate to reduce the variance of the Monte Carlo estimator. The [the joint reply to all reviewers](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=ga5ZZkiTQt) and Answer 1 of [Rebuttal UxzB](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=P18iyPkGnV)  contains detailed information on the new baselines. While we acknowledge that it would be interesting to compare COCOA to the RUDDER method (Arjona-Medina et al. 2019) we consider it out of scope for our current work to integrate RUDDER with our dynamic programming set up.\n\n## 2. Influence of imperfect modelling on the performance of COCOA\n\nWe now include a new experiment investigating this interesting question. We refer to Answer 2 of [Rebuttal UxzB](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=P18iyPkGnV)  for further details. \n\n## 3. Best practices for modelling action contributions for ‘rewarding outcomes’\n\nWhen considering which rewarding outcome encoding $U$ to use, there are 2 main forces to take into consideration. First, section 3.3 argues that taking the reward as rewarding outcome encoding results in the lowest variance estimator, when we have access to the ground-truth contribution coefficients from an oracle. Second, when we need to learn the contribution coefficients, taking the rewards as rewarding outcome encodings is not always the best choice, as the rewards do not contain a lot of information about the underlying environment structure, which can make it harder to learn good hindsight distributions. Especially when multiple rewarding objects have the same scalar reward, this can pose problems to COCOA-reward, favoring COCOA-features that learns a rewarding object representation instead (c.f. Appendix E.7.1 in the supplementaries for an empirical evaluation with reward aliasing).\n\nFor environments with simple dynamics and without reward aliasing we recommend using COCOA-reward, as it is uses a straight-forward encoding $U=R$ and leads to good results in practice. \nIn environments where one expects the state space to contain useful information for learning the hindsight distribution, or where reward aliasing can occur, there can be a trade-off when choosing the information content of $U$. Indeed, a lower information content of $U$ leads to a lower variance, but can lead to a higher bias in the hindsight model as there is less information available to leverage structure of the environment dynamics for efficient learning.\n\nFrom a practitioners perspective, we can navigate this trade-off by considering the information content of the encoding $U$ as a hyperparameter that needs to be tuned for a certain environment. When learning the rewarding outcome features for COCOA-feature, we use a regularization term in the loss to control the information content of the encoding (c.f. appendix E.6.3), which can be tuned as a hyperparameter.\nIn Appendix G in the supplementary materials, we further expand upon the interesting question of navigating this trade-off. Our main aim here is to explore strategies that can leverage more state-information to learn the credit assignment models, without causing additional variance.\n\nWe highlight three interesting avenues:\n1. Learn a state-based hindsight distribution, as in HCA-state, which can leverage detailed state information, but use counterfactual reasoning on rewarding states to recombine the state-based hindsight distributions into a reward-based hindsight distribution, as in COCOA-reward (c.f. Appendix G.3). Intuitively, this approach evaluates the question “could I have gotten the same reward in different states as well?”\n2. Use a variational information bottleneck approach to explicitly control the information content of the encodings, which can also be combined with the previous point (c.f. Appendix G.4). \n3. Instead of learning backward hindsight distributions $p(a \\mid s, u)$, learn forward environment dynamics $p(s’ \\mid s,a)$ and $p(u \\mid s,a)$ and use this to compute the contribution coefficients using Eq. 3. Learning the forward environment dynamics can leverage the detailed state-information, hence bypassing the variance trade-off (c.f. Appendix G.2)\nWe kindly refer the reviewer to Appendix G for more information about these approaches.\n\n\n## 4. Uncertain scalability to intricate environments\n\nWe agree that investigating the scalability of COCOA is an important next step. In this work, however, we concentrated our efforts on a deep analysis of all methods by designing a dynamic programming toolbox, instead of dedicating our time towards investigating the scalability of our methods. This allows us to investigate the credit assignment capabilities of RL algorithms without confounding them with exploration difficulties, as is often the case with performance metrics. While we consider it out of the scope of our work to perform a detailed scalability analysis, we perform experiments on a new environment compatible with our dynamic programming toolbox, to show that the improved performance of the COCOA methods is not specific to the key-to-door environment. We refer to [the joint reply to all reviewers](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=ga5ZZkiTQt) for further details."}, {"review_id": "FkA3uHJuiL", "reviewer": "Reviewer_nRru", "summary": "The paper presents a policy gradient estimator using credit assignment measure on influence of specific actions toward future rewards. The paper offers theoretical analysis on resulting algorithm such as (1) policy gradient estimator is unbiased, (2) variance of the estimator in relation to some existing methods, and so on. The authors also include numerical studies using a simple but specifically designed example to highlight the benefits of the proposed algorithm.", "questions": "1. Rare but possible that reward-based approach may result in larger variance (e.g., S=(1,2) but r(s,a) \\in {-100, 0, 100)).\n2. Why did you use only state-conditional HCA? (i.e., not using return-conditional HCA)? \n3. In the last paragraph on p.5, it is argued that less informative encodings lead to lower variance. While this is intuitively true, less information also means less learning, which suggests that there should be a trade-off between lower variance and information loss in learning. How would you achieve a good balance? ", "limitations": "None", "rating": 4, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "1. With some assumptions, the paper provides analytical guarantees on a range of desirable aspects of the algorithm. ", "weaknesses": "1. Given HCA, contribution seems marginal. \n2. Claims are either straightforward or require strong assumptions (quarantees in extreme cases such as Prop 2, Theorem 3)\n3. Experiments too simple ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the useful feedback that helped us improve our work. We address the points you raised below.\n\n## 1. Contribution compared to HCA and importance of our theoretical results\nAlthough Harutyunyan et al. 2019 proposed HCA as a family of algorithms leveraging hindsight probabilities conditioned on any function of the trajectory, the authors only provided theoretical justifications and experimental results for HCA-state and HCA-return. Hence, our theoretical results for general rewarding outcomes combined with a rigorous empirical investigation of our newly proposed algorithms are important new contributions to the field of credit assignment in RL. More specifically:\n- Theorem 1 introduces a general condition (Definition 1) under which hindsight information COCOA and HCA estimators can be used to construct unbiased policy gradients. \n- Proposition 2 provides crucial insights on why HCA-state fails to perform well on many environments of interest, relating it to the high-variance REINFORCE estimator and pointing to a more general issue: state representations need to contain detailed features to enable powerful policies, whereas the same level of detail leads to high variance in HCA-state. These insights combined with our experimental analysis show that HCA-state suffers due to high variance even in environments that significantly violate the strict assumptions of Proposition 2.\n- Theorem 3 is to the best of our knowledge the first theoretical result analyzing the variance of the HCA and COCOA policy gradient estimators, providing important insights on how the information content of hindsight objects influence the variance of the estimator, and suggesting that COCOA-reward results in the lowest-variance estimator. Although Theorem 3 considers only a limited setting, the relative ordering between the methods in terms of variance as predicted by Theorem 3 seems to translate to the general setting of dense rewards and a full policy gradient (c.f. Figure 2 and 3, and Figure 6 in Appendix C). \n- We introduce COCOA-reward and COCOA-features, two practical new algorithms resulting in low-variance gradient estimators with strong performance. \n- We developed a dynamic programming setup which makes it possible to investigate the credit assignment capabilities of RL algorithms without confounding them with exploration difficulties, as is often the case with performance metrics.\n- Specifically, this allows us for the first time to empirically investigate the performance of the ground-truth HCA and COCOA variants (i.e. without approximating models for the hindsight distribution) and further allows us to investigate the bias and variance of the approximate HCA and COCOA methods.\n\n\n## 2. Extending COCOA to more environments\nWe agree that investigating the scalability of COCOA is an important next step. In this work, however, we concentrated our efforts on a deep analysis of all methods by designing a dynamic programming toolbox, instead of dedicating our time towards investigating the scalability of our methods. This allows us to investigate the credit assignment capabilities of RL algorithms without confounding them with exploration difficulties, as is often the case with performance metrics. While we consider it out of the scope of our work to perform a detailed scalability analysis, we include a new environment compatible with our dynamic programming toolbox, to show that the improved performance of the COCOA methods is not specific to the key-to-door environment (c.f. Figure 2 in the attached pdf).\n\n## 3. Comparison with HCA-return\nWe have included new experimental and theoretical results, comparing COCOA-reward to HCA-return. Please refer to the [joint reply for details](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=ga5ZZkiTQt).\n\n## 4. Tradeoff between variance and information loss in learning\nWe thank the reviewer for raising this excellent point. Section 3.3 uses *ground-truth* contribution coefficients, and in this case less informative encodings indeed lead to lower variance. However, when *learning* the contribution coefficients, it could be that less informative encodings lead to slower learning and hence a bigger bias in the contribution coefficients, as fewer observations contain useful information for learning the hindsight distributions. We did not observe this issue in our key-to-door environment, as HCA-state, which has more informative encodings, has a higher bias on average compared to COCOA-reward in Fig. 3. However, in environments where there is more structure present that can be learned from non-rewarding observations, we would expect a trade-off between lower variance and higher bias due to information loss in learning. \n\nFrom a practitioners perspective, we can navigate this trade-off by considering the information content of the encoding $U$ as a hyperparameter that needs to be tuned for a certain environment. When learning the rewarding outcome features for COCOA-feature, we use a regularization term in the loss to control the information content of the encoding (c.f. appendix E.6.3), which can be tuned as a hyperparameter.\nIn Appendix G in the supplementary materials, we further expand upon the interesting question of navigating this trade-off. Our main aim here is to explore strategies that can leverage more state-information to learn the credit assignment models, without causing additional variance. We kindly refer the reviewer to Appendix G for more information about these approaches.\n\n> Rare but possible that reward-based approach may ...\n\nThis is an interesting consideration but even in this case, COCOA-reward still has smaller-or-equal variance compared to HCA-state. We assume the one-step MDP depicted in Figure 3a of the accompanying PDF: the agent has 2 actions, which give different transition probabilities to $s_1$ and $s_2$ respectively. This case exactly satisfies the assumptions of Theorem 3, independent of the reward structure."}, {"review_id": "UvXVw4HDdn", "reviewer": "Reviewer_Embf", "summary": "This paper introduces COCOA, a new family of hindsight credit assignment methods that build on HCA, which uses hindsight importance weights for the policy gradient estimator to reduce its variance. HCA uses state-conditioned importance weights, i.e., the ratio of $p(a_t|s_t, s_{t+k})$ to $\\pi(a_t|s_t)$. COCOA generalizes HCA by replacing $s_{t+k}$ with $u_{t+k} = f(s_{t+k}, a_{t+k}, r_{t+k})$ with a general function $f$. Specifically, the authors suggest using either $f(s, a, r) = r$ or $f(s, a, r) = g(s, a)$ where $g(s, a)$ is a function that can fully determine the reward. These variants lead to a lower variance compared to vanilla HCA by conditioning only on information that is related to the reward function. They evaluate COCOA on synthetic domains, showing that COCOA outperforms HCA with improved credit assignment.\n", "questions": "- How is COCOA-reward different from HCA-return?\n- Given that COCOA-reward shows the lowest variance (Theorem 3), when is COCOA-feature preferred to COCOA-reward?\n- Is COCOA-reward always *optimal* in terms of variance, meaning that there are no importance sampling-based gradient estimators with lower variance?\n- What would happen if the reward function is a constant function (e.g., $r(s, a) = 1$) but there are terminal signals? Can COCOA-reward directly handle this MDP (without transforming the MDP -- e.g., by creating an additional absorbing state with a zero reward)?", "limitations": "The authors adequately address the limitations of COCOA in Section 5.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors generalize HCA in an important direction, and the theoretical benefits of COCOA are significant and clear.\n- The paper contains informative discussions about various aspects of hindsight gradient estimators, which provides further insights into hindsight credit assignment in general.\n- The paper is well-written and easy to follow. ", "weaknesses": "- The paper lacks discussion regarding return-conditioned HCA. While the authors argue that (state-conditioned) HCA suffers from high variance, the original HCA paper also proposes a return-conditioned variant (Sec 3.2 and Theorem 5 in Appendix B), which can enjoy similar benefits to COCOA. Also, when referring to HCA, the authors only mention its state-conditioned variant throughout the paper. Discussing and comparing with the return-conditioned variant of HCA would be more relevant given its similarity to COCOA-reward.\n- The novelty of the method is a bit limited, as the original HCA paper also proposes a general version of HCA that is conditioned on any function of the future trajectory (Sec 3 in the HCA paper). Nonetheless, the paper further discusses sufficient conditions of this function to make the policy gradient estimator unbiased.\n- The experiments are conducted only on toy, synthetic environments. It is questionable how COCOA works in more realistic domains.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments and the useful feedback that helped us improve our work. We address the points you raised one-by-one below.\n\n## 1. Comparison to HCA-return\n\nWe agree that a dedicated comparison to and discussion of the return-conditioned variant of HCA was missing and we added it now to the manuscript. Notably, we added HCA-return as a baseline to our experiments (see Figure 1 in attached PDF) and discuss its theoretical properties (see joint reply).\n\n\n> The novelty of the method is a bit limited, as the original HCA paper also proposes a general version of HCA that is conditioned on any function of the future trajectory (Sec 3 in the HCA paper). Nonetheless, the paper further discusses sufficient conditions of this function to make the policy gradient estimator unbiased\n## 2. Contributions compared to HCA\nAlthough Harutyunyan et al. 2019 proposed HCA as a family of algorithms leveraging hindsight probabilities conditioned on any function of the trajectory, the authors only provided theoretical justifications and experimental results for HCA-state and HCA-return. Hence, our theoretical results for general rewarding outcomes combined with a rigorous empirical investigation of our newly proposed algorithms are important new contributions to the field of credit assignment in RL. In [Answer 1 of rebuttal nRru](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=YPsUl3Q9i1) we provide further details on the contributions of our work.\n\n## 3. Scalability of COCOA\n\nWe agree that investigating the scalability of COCOA is an important next step. In this work, however, we concentrated our efforts on a deep analysis of all methods by designing a dynamic programming toolbox, instead of dedicating our time towards investigating the scalability of our methods. This allows us to investigate the credit assignment capabilities of RL algorithms without confounding them with exploration difficulties, as is often the case with performance metrics. While we consider it out of the scope of our work to perform a detailed scalability analysis, we perform experiments on a new environment compatible with our dynamic programming toolbox, to show that the improved performance of the COCOA methods is not specific to the key-to-door environment. We refer to [the joint reply to all reviewers](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=ga5ZZkiTQt) for further details.\n\n\n## 4. Comparison with HCA-return\nIn the [general post for all reviewers](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=ga5ZZkiTQt) we provide a detailed theoretical and empirical analysis of HCA-return, comparing it to COCOA-reward.\n\n\n## 5. When is COCOA-feature preferred to COCOA-reward?\n\n<!-- The reviewer points out correctly that COCOA-reward shows the lowest variance, raising the question when COCOA-feature would be preferred over COCOA-reward. -->\nWhen an oracle provides access to the ground-truth contribution coefficients, COCOA-reward is better-or-equal compared to COCOA-features in MDPs. However, when we need to learn the contribution coefficients, it is often beneficial to use COCOA-features, especially when different rewarding objects have the same scalar reward, but different features. In this case, COCOA-feature can disentangle the rewarding objects, whereas COCOA-reward cannot (c.f. Section 4 in the main text and Fig. 8 in Appendix E). Furthermore, when the environment has many structural regularities, it can help to have more informative rewarding outcome encodings $U$ to leverage latent learning, on which we elaborate in Appendix G. Finally, in a non-MDP setting where the reward function but not the environment dynamics can change throughout training, COCOA-feature allows for quick adaptation (c.f. Fig. 4).\n\n## 6. Is COCOA-reward optimal?\nWhile our theoretical and empirical results suggest that the COCOA-reward estimator with ground-truth contribution coefficients has the lowest variance when considering the COCOA and HCA family of methods, we currently cannot make any conclusive claims on the optimality of COCOA-reward. The main two reasons are that (i) Theorem 3 only applies in a limited setting, hence cannot be used to theoretically show the optimality of COCOA-reward in the more general case, and (ii) our work focuses on gradient estimators that use rewarding outcome encodings based on state-action pairs, leaving it to future work to analyse the variance of other hindsight functions of future trajectories, such as the return.\n\n> What would happen if the reward function is a constant function (e.g., $r(s,a) = 1 $) but there are terminal signals? ...\n\nWe interpret ‘terminal signals’ as state features indicating the end of a trajectory. If this interpretation is not what the reviewer intended, we kindly ask for a clarification of the question.\n\nCOCOA-reward estimates the correct policy gradient in this extreme setting. To gain an intuition, consider the MDP illustrated in Fig 3b in the attached pdf, where each state has a constant reward of 1. Following action 1 results in a longer trajectory compared to action 2, and hence more rewards $r=1$ are encountered after taking action 1. Consequently, the probability $p^\\pi(A_0=a_1 \\mid S_0=s, R=1) > p^\\pi(A_0=a_2 \\mid S_0=s, R=1)$, as $\\sum_k p^\\pi(R_k=1 \\mid A_0=a_1, S_0=s) > \\sum_k p^\\pi(R_k=1 \\mid A_0=a_2, S_0=s)$, and hence COCOA-reward reinforces action $a_1$ more than action $a_2$, independently of which trajectory is followed (see also Appendix C.1). When learning the hindsight function $p^\\pi(A_0 \\mid S_0=s, R=1)$ with a classifier, there will be more triplets $(s,a_1, r=1)$ in the dataset compared to $(s,a_2, r=1)$, hence the learned hindsight function will approximate the true distribution."}, {"review_id": "H1hLikySvl", "reviewer": "Reviewer_5aMg", "summary": "This paper proposes Counterfactual Contribution Analysis (COCOA), an RL credit assignment approach inspired by the Hindsight Credit Assignment (HCA) family of algorithms. \n\nThe paper notes that in some instances, the previously proposed State-HCA approach can degrade to as high a variance as REINFORCE due to using the future state for obtaining hindsight weights. \n\nInstead of estimating hindsight contributions of actions based on future states (as in State-HCA), this paper introduces the idea of a rewarding outcome which is a function of a future state, action, and reward (s’, a’, r’) tuple. The paper proposes two choices for the rewarding outcome– the reward itself or the reward predictive features of the state-action pair (COCOA-reward and COCOA-feature). The paper provides intuitions on how the COCOA approaches should lead to lower-variance policy gradients compared to State-HCA due to reduced encoding of state information. \n\nExperiments in a key-to-door environment shows that the COCOA approaches perform much better than State-HCA, REINFORCE (with and without baseline), and all-actions policy gradient on Q-values.\n", "questions": "- In Figure 4, why does the COCOA policy not always collect the treasure before the reward switch? Is it due to an exploratory policy? Furthermore, why does the REINFORCE with baseline (Advantage) get much closer to always collecting the treasure in that period?", "limitations": "Negative societal impact is not directly applicable. Limitations adequately discussed in the Discussion section.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "**S1.** Hindsight-based approaches are a promising solution to the credit assignment problem in RL, which is of great interest to the research community. \n\n**S2.** While identifying the high-variance failure case of State-HCA may not be a novel contribution (see W2), the proposed solution to use rewards or reward predictive features instead of states is novel. The solution idea is simple, well-motivated with examples (e.g., Figure 2 and text below proposition 2), and fits the HCA family of approaches well. Focusing on features that can be useful for HCA is an interesting direction and opens a way to improve previous HCA approaches.\n\n**S3.** The paper is well-written and easy to follow.\n", "weaknesses": "**W1.** The paper's claims would be more accurate with the clarification that HCA refers to the State-HCA algorithm from the HCA paper. The claim that this approach generalizes HCA (e.g., line 103) is inaccurate. HCA generally proposes any function of the future trajectory to estimate contributions, and future states are simply one instance in the proposed family.\n\n**W2.** The observation that State-HCA-based approaches can suffer from high-variance or spurious contributions (even in comparison to Monte-Carlo approaches) is not a novel finding, as it has also been noted in previous work [1,2]. Could the authors clarify the similarities and differences to the observation made in the papers above? \n\n**W3.** A crucial weakness is that it is hard to judge the applicability of the proposed approach based on evaluation in a single environment (linear key-to-door). While linear key-to-door is a challenging problem from the perspective of credit assignment, an evaluation in Atari games (along the lines of Deep HCA [3]) could be more informative in guiding intuitions of how this approach scales to complex settings.\n\n**W4.** The paper also needs to include evaluation against other hindsight baselines, for instance, return-conditioned HCA, which was concretely proposed in the HCA paper. An even more interesting comparison would be to Counterfactual Credit Assignment [4], which learns features for hindsight contribution. \n\n\nOverall, I appreciate the direction the authors took to address the limitation of State-HCA. Should these weaknesses be adequately clarified/addressed, I would happily increase my score.\n\n\n—------------------—------------------—------------------—------------------—------------------\n\n### References\n\n[1] Young, K. (2019). Variance Reduced Advantage Estimation with $\\delta$ Hindsight Credit Assignment. arXiv preprint arXiv:1911.08362.\n\n[2] Zhang, P., Zhao, L., Liu, G., Bian, J., Huang, M., Qin, T., & Liu, T. Y. (2019). Independence-aware advantage estimation.\n\n[3] Alipov, V., Simmons-Edler, R., Putintsev, N., Kalinin, P., & Vetrov, D. (2021). Towards practical credit assignment for deep reinforcement learning. arXiv preprint arXiv:2106.04499.\n\n[4] Mesnard, T., Weber, T., Viola, F., Thakoor, S., Saade, A., Harutyunyan, A., ... & Munos, R. (2020). Counterfactual credit assignment in model-free reinforcement learning. arXiv preprint arXiv:2011.09464.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your encouraging review and the useful feedback that helped us improve our work.\n\n## 1. Adjusting the claims of COCOA\nWe acknowledge the foundational work of Harutyunyan et al. 2019 in proposing a general class of credit assignment algorithms using hindsight probabilities conditioned on any function of the future trajectory. The main goal of our work is to extend the -theory- of HCA, as the authors only showed theoretical results for the unbiasedness of State-HCA and Return-HCA, leaving open the question of which other functions of the future trajectory lead to unbiased policy gradient estimators, and what impact these functions have on the variance of the estimators. We will clarify this main focus in the manuscript.\n\n## 2. Comparison to Young 2020 and Zhang 2019\n1. **Scope of the analysis**: Both works focus their analyses on a specific constructed setting with constant Monte Carlo (MC) returns, no matter which state trajectories are followed. Proposition 2 and its accompanying explanations demonstrate that the high variance of HCA is not just limited to this constructed setting but becomes problematic in a wider range of environments. Interestingly, our work highlights the inherent conflict between obtaining detailed state-representations that enable powerful policies and the resultant high variance in HCA.\n2. **Estimator focus**: A significant difference lies in which estimator we analyse. While both Young 2020 and Zhang et al 2019 analyse the value estimator, our work studies the policy gradient estimator. This distinction is crucial, as it can lead to different conclusions. Even though a value estimate can be of higher variance compared to the MC value estimate, when combining these value estimates for all actions into the policy gradient estimator, the resulting gradient estimate can be of equal or lower variance compared to its MC equivalent. To illustrate, consider the MDP of Figure 3c in the attached pdf. Here, the return is constant and hence the MC estimator of the Q-value has zero variance, whereas the HCA Q-value estimator does not. However, when comparing the REINFORCE policy gradient estimator with HCA, HCA still demonstrates a reduction in variance, as it accurately estimates that the penultimate state is accessible through both actions in state A, while REINFORCE solely credits the observed action in state A for the entire return. It's worth noting that in scenarios with negative rewards, one can construct toy examples where rewards nullify each other resulting in a constant zero return independent of the state trajectory, which could result in HCA exhibiting higher variance compared to REINFORCE.\n4. **Solutions to the variance problem**: Both Young 2020 and Zhang et al 2019 propose solutions to the variance issue that continue to utilize the state-based contribution coefficients. This approach does not address the core issue causing the high variance. Given that state representations in realistic environments are often detailed, aligning with Proposition 2, the solutions proposed by these works may not significantly reduce the policy gradient variance when compared to their MC counterparts, in contrast to our low-variance COCOA methods.\n\nIn conclusion, while we acknowledge the important work of Young 2020 and Zhang et al 2019, our paper extends and deepens the understanding of the variance issues in the HCA policy gradient estimator and its relationship with state representations. These insights have enabled us to introduce new COCOA estimators with markedly reduced variance.\n\n\n## 3. Extending COCOA to more environments\nWe agree that investigating the scalability of COCOA is an important next step. In this work, however, we concentrated our efforts on a deep analysis of all methods by designing a dynamic programming toolbox, instead of dedicating our time towards investigating the scalability of our methods. This allows us to investigate the credit assignment capabilities of RL algorithms without confounding them with exploration difficulties, as is often the case with performance metrics. While we consider it out of the scope of our work to perform a detailed scalability analysis, we include a new environment compatible with our dynamic programming toolbox, to show that the improved performance of the COCOA methods is not specific to the key-to-door environment (c.f. Figure 2 in the attached pdf)\n\n## 4. Comparison to new baselines\n\nWe now include HCA-return in our experiments, as well as a new method Counterfactual Return that uses returns to construct contribution coefficients (see general rebuttal). Finally, we include Trajectory-wise CV (Cheng, Yan and Boots, 2019), a state-of-the-art method that incorporates hindsight information to construct a near-optimal control variate (see [Answer 1 Rebuttal UxzB](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=P18iyPkGnV)). While we acknowledge that comparing COCOA to Counterfactual Credit Assignment could lead to interesting insights, the source code of this method is not publicly available and we consider it out of scope for our work to reproduce this method consisting of several intrically connected models and losses. We have contacted the lead authors of Mesnard et al. 2021 to request access to their code, to be able to run this baseline. \n\n## 5. Clarification Figure 4\nThe entropy regularization that we use to ensure sufficient exploration has a different effect on the various methods. We perform an independent hyperparameter tuning for all methods and baselines, including the entropy regularization, as detailed in Appendix E.6.3 in the supplementary materials. This led to a hyperparameter configuration for COCOA-reward that solves the environment fast above 90%, but does not solve it 100% due to the entropy regularizer. The tuned configuration for Advantage solves the environment closer to a 100% after 60k episodes, due to a different effect of the entropy regularizer on the exploration behavior. "}, {"review_id": "mpPDIltAGI", "reviewer": "Reviewer_UxzB", "summary": "This paper lays out a goal of addressing a weakness in HCA that HCA can confuse the contributions of actions to reaching a state, increasing the variance of gradient estimation. Instead, the COCOA method proposed in this paper generalizes HCA to use any feature that is predictive of rewards, such as the reward or reward-predicting features. Thus, it allows for selecting an appropriate rewarding outcome to disentangle the contributions of each action from the observed rewards. The paper derives a new unbiased gradient estimator (given the exact hindsight model). Experiments are then conducted to show that when these models are known, they can reduce variance in cases not handled by standard control variates. ", "questions": "From Theorem 3 it wasn't clear if one should expect the same trend to hold in other cases. Would we expect the same ordering of variances in more settings? ", "limitations": "The authors do mention that further work is needed on developing methods to learn the hindsight models so they do acknowledge this limitation. But a pilot study on the same environment is really warranted in this work. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The biggest strength of this paper is thorough experimentation to illustrate when COCOA reduces variance when the hindsight models are known. \n\nThe experiments and environment clearly show the COCOA method's impact in the cases: when distractor rewards are encountered during the episode and when the distractor rewards are of the same magnitude as the main reward. These experiments show clear evidence that selecting the appropriate disentangling factor is essential to reducing variance. \n\nThe experiments to show the SNR as a function of credit assignment distance highlights how each gradient estimator degrades as the episode length and distance between action and rewarding stimulus increases. However, this should be evaluated at longer distances as the slopes and shapes of the curves are difficult to observe in this range. ", "weaknesses": "The biggest weakness of this paper is that there is no investigation into how accurate the estimates of hindsight models need to be to reduce variance. Past variance reduction techniques for policy gradients have proven useless in practice (Tucker et al. 2018). So it is crucial to understand the practical limitations when developing a new method. \n\n\nThe comparison to q critic, when q is known, is odd. Because if q is known, then the natural gradient will be known (Kakade, 2001). It should be made precise what sources of variance these experiments are measuring. \n\nThe bias can be much higher for COCOA methods than other methods (Figure 3 (D)), and it needs to be clearly discussed since this method is supposed to be unbiased. \n \nA more relevant baseline instead of the Advantage method would be the doubly robust control variate for policy gradient methods (Huang and Jiang, 2020). This has a significant variance reduction compared to standard baselines in policy gradient. The comparison would also better highlight the limits of information between the two approaches. \n\nThe success/failure of an algorithm depends greatly on hyperparameter choices. It needs to be clear how each method's hyperparameters were chosen. Since this paper is trying to argue that COCOA has better sample efficiency, it should be made clear in what exact context this claim is being made. Furthermore, if one wants to say that COCOA has better information utilization through measuring performance, then the hyperparameters for the baseline methods need to be tuned to their absolute maximum. As the paper currently stands, the experiment design is insufficient to make this argument, even if it is likely true. \n\nThough not poor, the writing in the paper has room for improvement. One of the best ways for this paper to have a bigger impact is to make it written so that most RL people can easily understand the lessons. So I recommend that the authors spend time revising the writing for a future or camera-ready version of the paper. \n\n___minor comments__\nwhat are the decibel units? This is nonstandard in ML, so it is worth specifying. \n\nTable1: In the Advantage method, $v^\\pi$ should just be $v$ because the baseline can be any state-dependent function\n\nAlthough common in the ML community $\\nabla_\\theta$ is a mathematical symbol for a directional derivative. However, it is used in this work as a partial derivative. Further, since the gradient is a function that takes a derivative with respect to all inputs and partial derivates are the quantities of interest, it makes mathematical sense to use mathematical symbols for partial differentiation, not gradients. \n\nLine 100: \"it has a failure mode of practical importance\" —say what this is. \n\nSection 3.1 Define reward outcomes first. It is unclear what they mean, making it hard to interpret. It also needs to be made clear when or why one would use different choices of f. Giving a few motivating examples would help here. \n\nLine 133: Eq 1. Check the JMLR guidelines for standard practice for referencing equations. It should be just (1) not Eq, and equation blocks should only be numbered if they are referenced. \n\nLine 146: missing for all k?\n\nDefinition 1. Remark on definition 1 to clarify what is needed to satisfy this property in practice. \n\nTheorem 1: It is important to remark that the gradient estimator is unbiased only if the counterfactual term is known precisely, which is unlikely in practice. This would be like saying the DDPG is unbiased, which is misleading because it is only true if Q is known. \n\nTheorem 4. \"equals\" should be changed because the expression is not a statement of equality but proportionality. \n\nLine 246: \"SNR indicates… fewer trajectories…\" This is not universally true. It can require an exponential number of trajectories to get a gradient estimate, and lower variance estimates will not help this. \n\n\n\nREFERENCES\nHuang, Jiawei, and Nan Jiang. \"From importance sampling to doubly robust policy gradient.\" International Conference on Machine Learning. PMLR, 2020.\n\nKakade, Sham M. \"A natural policy gradient.\" Advances in neural information processing systems 14 (2001).\n\nTucker, George, et al. \"The mirage of action-dependent baselines in reinforcement learning.\" International conference on machine learning. PMLR, 2018.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comprehensive remarks and for the useful feedback that we incorporated in a revised version of our paper. We address your main concerns below.\n\n## 1. SOTA Control Variate baseline.\n\nWe thank the reviewer for suggesting an extra baseline that leverages hindsight information for a control variate (CV). Doubly-Robust CV requires learning an environment dynamics model to contruct a cheap estimate of the Q-value gradient, which we consider out-of-scope for our work. Instead, we investigate TrajCV (Cheng, Yan and Boots, 2019), a related SOTA CV method that also leverages hindsight information but does not require an environment model. \n\nFigure 1 in the attached pdf shows the results of the TrajCV method on our linear key-to-door environment. We observe that it improves upon the advantage baseline but still falls short of the performance of the COCOA methods with ground-truth contribution coefficients. The main reason for the performance gap between the COCOA methods and TrajCV is that the control variate of TrajCV can only reduce the variance caused by the stochastic policy and not from stochastic environment dynamics, whereas COCOA can reduce both. In the linear key-to-door environment, a large part of the variance is due to the stochastic distractor rewards. The value functions used in TrajCV necessarily need to approximate the average returns to use as control variates, whereas COCOA can ignore the stochastic distractor rewards completely when applicable, by multiplying them with contribution coefficients equal to zero. Interestingly, this points 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.\n\n## 2. Investigation into the needed accuracy of the hindsight models\nWe include a new experiment quantifying the relation between the approximation quality of the hindsight models and value functions upon the SNR of the resulting policy gradient estimate. We start from the ground-truth hindsight models and value functions computed with our dynamic programming setup, and introduce a persistent bias into the output of the model by elementwise multiplying the output with $(1 + \\sigma \\epsilon)$, with $\\epsilon$ zero-mean univariate Gaussian noise and $\\sigma \\in \\{0.001, 0.003, 0.01, 0.03, 0.1,0.3\\}$ a scalar of which we vary the magnitude in the experiment. \n\nFigure 1a in the attached PDF shows the SNR of the COCOA-reward estimator and the baselines, as a function of the perturbance magnitude $\\log\\sigma$, where we average the results over 30 random seeds. We see that the sensitivity of the COCOA-reward estimator to its model quality is similar to that of Q-critic. Furthermore, the SNR of COCOA-reward remains of better quality compared to Advantage and HCA-state for a wide range of perturbations.\n\n## 3. Clarification on hyperparameter tuning\nAppendix E.6.4 and Table 2 provide details on our rigorous hyperparameter tuning. We emphasize that we used the same extensive hyperparameter search for all methods and baselines, and hence we are confident that our empirical results support the main claims we made in our paper.\n\n## 4. Generality of Theorem 3\nTheorem 3 considers only a limited setting. To extend these theoretical results, we performed an extensive empirical evaluation on random MDPs with dense rewards considering the full policy gradient estimates. Figure 2 and Figure 6 in Appendix C.5 indicate that the relative ordering of Theorem 3 holds in more general settings as well. This is further corroborated by our results on the key-to-door environment in Figure 3. It is an open question whether it is possible to formally extend Theorem 3 to a more general setting. \n\n## 5. Q-critic baseline\nOur main motivation to include the Q-critic when the Q-value is learned, is to have a baseline that considers all counterfactual actions for a given trajectory, similar to the COCOA estimators. The ground-truth Q-critic is an interesting upper bound for the other methods. Compared to the ground-truth policy gradient, the only source of variance in the ground-truth Q-critic is the sampled trajectory states $S_t$ where the policy is evaluated. In Figure 3, we see that the performance of ground-truth COCOA-reward matches the Q-critic upper bound. We note that in the bias-variance and SNR plot of Figure 3, we consider the policy gradient on the first decision moment of picking up the key, as this determines the performance on the key-to-door task. In this case, the ground-truth Q-critic is equal to the ground-truth policy gradient, as the first state is always the same. \n\n## 6. Biasedness of COCOA with learned coefficients\nTheorem 1 considers only the case with ground-truth contribution coefficients. When we approximate the contribution coefficients, this can lead to a biased COCOA estimator. Theorem 1 nonetheless is an important result, as it was not clear beforehand which rewarding outcome encodings $U$ lead to a an unbiased policy gradient estimator *with ground-truth contribution coefficients*. We thank the reviewer for pointing out this possible confusion, and we will clarify it in the main text by adjusting the text of Theorem 1 and discussing the potential biasedness in section 3.4.\n## 7. Remaining comments\n- We will run the experiments of Figure 1B for longer environment lengths for the final version of our paper. \n- We thank the reviewer for the detailed suggestions to improve the clarity of our manuscript and will incorporate them in our final version."}]

Author rebuttal:
We thank all reviewers for their constructive comments and useful suggestions that we believe have helped us to significantly improve our paper. Here we summarize the changes we did addressing the main concerns of the reviewers, linking to the corresponding detailed responses. 

## Main changes

1. Reviewers nRru, Embf and 5aMg asked for a comparison with HCA-return. We now include HCA-return in our experimental results, showing that it performs similarly to HCA-state and significantly lags behind our COCOA methods and other baselines. Furthermore, we provide new theoretical results showing that HCA-return with ground-truth hindsight distributions is a biased estimator in many relevant environments. See below for a detailed discussion.
2. Reviewers 5aMg, UxzB, V1Ut, nRru and Embf asked to introduce extra baselines to compare COCOA against state-of-the-art credit assignment methods. We included three new baselines: (i) Trajectory-wise Control Variate (Cheng, Yan and Boots, 2019; TrajCV), a SOTA control variate method to reduce the variance of MC estimators, (ii) HCA-return as explained above, and (iii) Counterfactual Return, a new baseline investigating the difference between taking rewards or returns as hindsight information for constructing contribution coefficients. See below and [Answer 1 Rebuttal UxzB](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=P18iyPkGnV) for details.
3. Reviewers UxzB and V1Ut asked for additional experiments investigating the influence of the hindsight model accuracy upon the gradient estimator quality. We now include new experimental results investigating this question, showing that COCOA-reward remains performant for significantly biased models, compared to the other baselines. See Answer 2 in [Rebuttal UxzB](https://openreview.net/forum?id=yvqqkOn9Pi&noteId=P18iyPkGnV) for details.
4. Reviewers 5aMg, V1Ut, nRru and Embf raised the concern that our experiments are focused around a single environment. We have added an additional environment based on Mesnard et al. (2021), see below for details. 


## Comparison with HCA-return

### Difference between HCA-return and COCOA-reward
The first main differences between HCA-return and COCOA-reward are that (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:
$\sum_{t\geq 0} \nabla_\theta \log \pi(A_t \mid S_t)\Big(1 - \frac{\pi(A_t \mid S_t)}{p^\pi(A_t \mid S_t, Z_t)}\Big)Z_t$

### HCA-return is a biased estimator in many relevant environments.
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 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.

### Using returns instead of rewards as hindsight information
To isolate the difference of using returns instead of rewards as hindsight information for constructing the contribution coefficients, we introduce the following new baseline:
$\sum_{t\geq 0} \sum_{a} \nabla_\theta \pi(a \mid S_t)\Big(\frac{p^\pi(a \mid S_t, Z_t)}{\pi(a \mid S_t)} -1\Big)Z_t$. 
Similar to HCA-return, this *Counterfactual Return* variant leverages a hindsight distribution conditioned on the return. Different from HCA-return, we use this hindsight distribution to compute contribution coefficients that can evaluate all counterfactual actions. We prove that this estimator is unbiased.


### Empirical comparison HCA-return and Counterfactual Return

Figure 1 in the attached pdf shows that the performance of HCA-return and Counterfactual return lags far behind the performance of COCOA-reward, both in a setting with ground-truth hindsight distributions and learned hindsight distributions. This is due to the high variance of HCA-return and Counterfactual Return, and the biasedness of the former. As the return is a combination of all the rewards of a trajectory, it cannot be used to disentangle rewarding outcomes, causing the variance of the distractor subtask to spill over to the subtask of picking up the treasure.

## New environment
To show that the improved performance of the COCOA methods is not specific to the key-to-door environment, we include a stylized version of the Task Interleaving environment of Mesnard et al. 2021. Strong credit assignment is needed to solve this task, due to the delayed rewards. 
Figure 2 in the attached pdf describes the environment and shows the results for our main methods and baselines. We will include the remaining baselines, ground-truth performance and bias/variance analyses in the final manuscript.

Meta-review: Based on the reviews provided by the reviewers, I recommend accepting the paper. The proposed Counterfactual Contribution Analysis (COCOA) approach addresses limitations of existing methods, particularly the State-Hindsight Credit Assignment (HCA) variant, by introducing rewarding outcomes and disentangling action contributions using various features. The theoretical analysis and experiments demonstrate its potential in reducing variance and improving credit assignment in reinforcement learning. While some weaknesses and limitations are acknowledged, the paper's contributions are technically sound, and it shows promise in improving sample efficiency and addressing issues with credit assignment. Therefore, I recommend accepting the paper as it offers a novel approach with moderate-to-high impact in the field.
Decision: Accept (spotlight)