SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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,  $\mathrm{HCA + }$
F.1 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 a learned reward model  $r(s,a)$ . We empirically verify that this simplification does not lead to a 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  $u$  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
Figure 9: 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 in terms of performance measured as the percentage of treasure reward collected.
Figure 10: COCOA-feature quickly adapts the policy to 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 the treasure, we can use this knowledge to avoid picking up the key and opening the door when the treasure becomes poisonous.
To illustrate these benefits, we 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 to 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-reward in contrast needs to encode the poisonous treasure as a new, previously unknown rewarding outcome. Since it 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.
F.3 Using returns instead of rewards as hindsight information
In Appendix L, we discussed that HCA-return is a biased estimator in many relevant environments, including our key-to-door and task-interleaving environments. This explains its worse performance compared to our COCOA algorithms. To isolate the difference of using returns instead of rewards as
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:
$$
\sum_ {t \geq 0} \sum_ {a} \nabla_ {\theta} \pi (a \mid S _ {t}) \left(\frac {p ^ {\pi} (a \mid S _ {t} , Z _ {t})}{\pi (a \mid S _ {t})} - 1\right) Z _ {t} \tag {99}
$$
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.
Figure 11 shows that the performance of HCA-return and Counterfactual return lags far behind the performance of COCOA-reward on the key-to-door environment. 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.
F.4 Investigation into the required accuracy of the hindsight models
To quantify the relation between the approximation quality of the hindsight models and value functions upon the SNR of the resulting policy gradient estimate, we introduce the following experiment. We start from the ground-truth hindsight models and value functions computed with our dynamic programming setup (c.f. E.2), 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.
Figure 12 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.
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 to contain detailed features to allow for a capable policy. The same level of detail however 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 a 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 a large extent due to their fine-grained state representations. Here, we take a closer look at how spurious contributions arise in continuous environments and Partially Observable MDPs (POMDPs).
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, as a function of the perturbance magnitude  $\log \sigma$ , where we average the results over 30 random seeds.
G.1 Spurious contributions in continuous state spaces
When using continuous state spaces,  $p^{\pi}(S_k = s' \mid s, a)$  represents a probability density function (PDF) instead of a probability. A ratio of PDFs  $p(X = x) / p(X = x')$  can be interpreted as a likelihood ratio of how likely a sample  $X$  will be close to  $x$  versus  $x'$ . Using PDFs, the contribution coefficients with  $U = S$  used by HCA result in
$$
w (s, a, s ^ {\prime}) = \frac {\sum_ {k \geq 1} p ^ {\pi} \left(S _ {t + k} = s ^ {\prime} \mid S _ {t} = s , A _ {t} = a\right)}{\sum_ {k \geq 1} p ^ {\pi} \left(S _ {t + k} = s ^ {\prime} \mid S _ {t} = s\right)} - 1 = \frac {p ^ {\pi} \left(A _ {t} = a \mid S _ {t} = s , S ^ {\prime} = s ^ {\prime}\right)}{\pi (a \mid s)} - 1 \tag {100}
$$
and can be interpreted as a likelihood ratio of encountering a state 'close' to  $s'$ , starting from  $(s, a)$  versus starting from  $s$  and following the policy. The contribution coefficient will be high if  $p^{\pi}(S' = s' \mid s, a)$  has a high probability density at  $s'$  conditioned on action  $a$ , compared to the other actions  $a' \neq a$ . Hence, the less  $p^{\pi}(S_k = s' \mid s, a)$  and  $p^{\pi}(S_k = s' \mid s, a')$  overlap around  $s'$ , the higher the contribution coefficient for action  $a$ . The variance of the distribution  $p^{\pi}(S_k = s' \mid 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 a sharp distribution  $p^{\pi}(S_k = s' \mid s, 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 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^{\pi}(S_k = s' \mid s, a)$ . For playing the violin, a slight change in 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 to a 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 in a high contribution for the action taken in the trajectory for both tasks. For playing violin, this could be a good approximation of what we intuitively think of as a '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.
G.2 Deterministic continuous reward functions can lead to excessive variance
Proposition 2 shows that HCA degrades to REINFORCE in environments where each action sequence leads to distinct states. In Section G.1, we discussed that continuous state spaces exhibit this property to a large extent due to their fine-grained state representation. If each environment state has a unique reward value, COCOA can suffer from a similar degradation to HCA. As the rewarding outcome encoding  $U$  needs to be fully predictive of the reward (c.f. Theorem 1), it needs to have a distinct encoding for each unique reward, and hence for each state.
With a continuous reward function this can become a significant issue, if the reward function is deterministic. Due to the fine-grained state representation in continuous state spaces, almost every action will lead to a (slightly) different state. When we have a 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  $p(R \mid S, A)$  is probabilistic, even for continuous rewards. If the variance of  $p(R \mid S, A)$  is bigger than zero, the variance of  $p(S, A \mid 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 to 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 environments 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 a deterministic reward function, we 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 \sim \mathcal{N}(R, \sigma)$ , with  $\mathcal{N}$  corresponding to a 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  $\sigma$ : a small sigma corresponds to a sharp distribution, akin to a fine discretization in the above binning strategy, and hence a low bias. Increasing  $\sigma$  leads to 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 is equivalent to adjusting the training scheme of the hindsight model  $h(a \mid s, u')$  by adding noise to the input  $u'$ . Here, we take  $U$  equal to the (deterministic) reward  $R$ , but add noise to it while training the hindsight model. Due to 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-variance trade-off for HCA and COCOA.
Smoothing hindsight states for HCA. 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
will likely result in sub-optimal bias-variance trade-offs, as it will be too coarse-grained in some regions causing large bias, whereas too fine-grained in other regions causing a low variance-reduction. Binning rewards in contrast does not suffer from this issue, as binning is directly performed in a 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 a discrete reward function with not too many distinct values compared to the state space, such that COCOA-reward is a low-variance gradient estimator. Proposition 8 shows that we can recombine the state-based hindsight distribution  $p^{\pi}(A_0 = a \mid S_0 = s, S' = s')$  into the reward-based hindsight distribution  $p^{\pi}(A_0 = a \mid S_0 = s, R' = r')$ , by leveraging a smoothing distribution  $\sum_{s' \in S} p^{\pi}(S' = s' \mid S_0 = s, R' = r')$ . When we consider a specific hindsight state  $s''$ , this means that we can obtain the low-variance COCOA-reward estimator, 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' \sim \mathcal{N}(s'', \sigma)$ , a more optimal binning strategy is to take the reward structure into account through  $p^{\pi}(S' \mid S_0 = s, R' = r(s'))$
Proposition 8.
$$
p ^ {\pi} \left(A _ {0} = a \mid S _ {0} = s, R ^ {\prime} = r ^ {\prime}\right) = \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S ^ {\prime} = s ^ {\prime} \mid S _ {0} = s, R ^ {\prime} = r ^ {\prime}\right) p ^ {\pi} \left(A _ {0} = a \mid S _ {0} = s, S ^ {\prime} = s ^ {\prime}\right) \tag {101}
$$
Proof. As  $A_0$  is d-separated from  $R'$  conditioned on  $S'$  and  $S_0$  (c.f. Fig. 6b) we can use the implied conditional independence to prove the proposition:
$$
\begin{array}{l} p ^ {\pi} \left(A _ {0} = a \mid S _ {0} = s, R ^ {\prime} = r ^ {\prime}\right) (102) \\ = \sum_ {s ^ {\prime} \in S} p ^ {\pi} \left(S ^ {\prime} = s ^ {\prime}, A _ {0} = a \mid S _ {0} = s, R ^ {\prime} = r ^ {\prime}\right) (103) \\ = \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S ^ {\prime} = s ^ {\prime} \mid S _ {0} = s, R ^ {\prime} = r ^ {\prime}\right) p ^ {\pi} \left(A _ {0} = a \mid S _ {0} = s, S ^ {\prime} = s ^ {\prime}, R ^ {\prime} = r ^ {\prime}\right) (104) \\ = \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S ^ {\prime} = s ^ {\prime} \mid S _ {0} = s, R ^ {\prime} = r ^ {\prime}\right) p ^ {\pi} \left(A _ {0} = a \mid S _ {0} = s, S ^ {\prime} = s ^ {\prime}\right) (105) \\ \end{array}
$$
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.
Gradient estimator. We can adjust the COCOA gradient estimator of Eq. 4 towards continuous action spaces by replacing the sum over  $a' \mathcal{A}$  by an integral over the action space  $A'$ .
$$
\hat {\nabla} _ {\theta} ^ {U} V ^ {\pi} \left(s _ {0}\right) = \sum_ {t \geq 0} \nabla_ {\theta} \log \pi \left(A _ {t} \mid S _ {t}\right) R _ {t} + \int_ {\mathcal {A}} \mathrm {d} a \nabla_ {\theta} \pi (a \mid S _ {t}) \sum_ {k \geq 1} w \left(S _ {t}, a, U _ {t + k}\right) R _ {t + k} \tag {106}
$$
In general, computing this integral is intractable. We can approximate the integral by standard numerical integration methods, introducing a 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
$$
\hat {\nabla} _ {\theta} V ^ {\pi} = \sum_ {t \geq 0} \nabla_ {\theta} \log \pi \left(A _ {t} \mid S _ {t}\right) R _ {t} + \frac {1}{M} \sum_ {m} \nabla_ {\theta} \log \pi \left(a ^ {m} \mid S _ {t}\right) \sum_ {k = 1} ^ {\infty} w \left(S _ {t}, a ^ {m}, U _ {t + k}\right) R _ {t + k} \tag {107}
$$
where we sample  $M$  actions independently from  $a^m \sim \pi(\cdot \mid S_t)$ . This gradient estimator is unbiased, but introduces extra variance through the sampling of actions.
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^{\pi}(A_0 = a \mid S_0 = s, S' = s')$  is equal to the Dirac delta function  $\delta(a = a_0)$ , with  $a_0$  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  $\mathcal{O}$ . Instead of directly observing Markov states, the agent now acquires an observation  $o_t$  at each time step. The probability of observing  $o$  in state  $s$  after performing action  $a$  is given by  $p_o(o \mid s, a)$ .
A popular strategy in deep RL methods to handle partial observability is learning an internal state  $x_{t}$  that summarizes the observation history  $h_{t} = \{o_{t' + 1}, a_{t'}, r_{t'}\}_{t' = 0}^{t - 1}$ , typically by leveraging recurrent neural networks [85-88]. This internal state is then used as input to a policy or value function. This strategy is intimately connected to estimating belief states [89]. The observation history  $h_{t}$  provides us with information on what the current underlying Markov state  $s_{t}$  is. We can formalize this by introducing the belief state  $b_{t}$ , which captures the sufficient statistics for the probability distribution over the Markov states  $s$  conditioned on the observation history  $h_{t}$ . Theoretically, such belief states can be computed by doing Bayesian probability calculus using the environment transition model, reward model and observation model:
$$
p _ {B} (s \mid b _ {t}) = p \left(S _ {t} = s \mid H _ {t} = h _ {t}\right) \tag {108}
$$
Seminal work has shown that a POMDP can be converted to an MDP, using the belief states  $b$  as new Markov states instead of the original Markov states  $s$  [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$ .