SlowGuess/NeurIPS_2023_OCR huggingface dataset:
HCA suffers from spurious contributions in POMDPs. 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 high-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, a 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.
To make the spurious contributions more tangible, let us consider a toy example where the state space  $S$  consists of three states  $x$ ,  $y$  and  $z$ . We assume the internal state represents a belief state  $b = \{b^{1}, b^{2}\}$ , which is a sufficient statistic for the belief distribution:
$$
p (S = x \mid b) = b ^ {1}, \quad p (S = y \mid b) = b ^ {2}, \quad p (S = z \mid b) = 1 - b ^ {1} - b ^ {2} \tag {109}
$$
We assume that  $b_{t}$  is deterministically computed from the history  $h_{t}$ , e.g. by an RNN. Now consider that at time step  $k$ , our belief state is  $b_{k} = \{0.5, 0.25\}$  and we get a reward that resulted from the Markov state  $x$ . As the belief states are deterministically computed, the distribution  $p^{\pi}(B' = b' \mid b, a)$  is a Dirac delta distribution. Now consider that the action  $a$  does not influence the belief distribution over the rewarding state  $x$ , but only changes the belief distribution over the non-rewarding states (e.g.  $B_{k} = \{0.5, 0.23\}$  instead of  $B_{k} = \{0.5, 0.25\}$  when taking action  $a'$  instead of  $a$ ). As the Dirac delta distributions  $p^{\pi}(B' = b' \mid b, a)$  for different  $a$  do not overlap, we get a high contribution coefficient for the action  $a_{0}$  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 as well.
Adapting Proposition 8 towards these internal belief states provides further insight on the difference between HCA and COCOA-reward:
$$
p ^ {\pi} \left(A _ {0} = a \mid X _ {0} = x, R ^ {\prime} = r ^ {\prime}\right) = \sum_ {X ^ {\prime} \in \mathcal {X}} p ^ {\pi} \left(X ^ {\prime} = x ^ {\prime} \mid X _ {0} = x, R ^ {\prime} = r ^ {\prime}\right) p ^ {\pi} \left(A _ {0} = a \mid X _ {0} = x, X ^ {\prime} = x ^ {\prime}\right) \tag {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 a 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^{\pi}(a \mid s, u')$ . The quality of this model is crucial for obtaining low-bias gradient estimates. However, its training data is 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 as input (c.f. Proposition 5). Enabling COCOA-reward or COCOA-feature to learn from non-rewarding states is a 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 contribution 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 3 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-efficiency 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 as little information as possible while still being predictive of the reward. To enable latent learning for sample-efficient learning of the contribution coefficients, the hindsight model needs to pick up on recurring structure in the environment, requiring keeping as much information as possible in  $U$  to uncover the structural regularities. Using the state as rewarding outcome encoding is beneficial for enabling latent learning, as it 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 to use 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 use 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 to 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.
Figure 13: Schematic of the graphical model used in our 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, s')$ .
Proposition 9. Assuming  $S'$  is fully predictive of  $U'$ , we have that
$$
w (s, a, u ^ {\prime}) = \sum_ {s ^ {\prime} \in \mathcal {S}} p ^ {\pi} \left(S ^ {\prime} = s ^ {\prime} \mid S _ {0} = s, U ^ {\prime} = u ^ {\prime}\right) w (s, a, s ^ {\prime}) \tag {111}
$$
The proof follows the same technique as that of Proposition 8. Proposition 8) shows that it is possible to learn state-based contribution coefficients, enabling latent learning, and obtain a low-variance COCOA estimator by recombining the state-based contribution coefficients into coefficients with less spurious contributions, if we have access to the generative model  $p^{\pi}(S' = s' \mid S_0 = s, U' = u')$ . This model embodies the counterfactual reasoning on rewarding states: 'which other states  $s'$  are likely considering I am currently in  $u'$  and visited state  $s$  somewhere in the past'.
In general, learning this generative model is as difficult or more difficult than approximating the hindsight distribution  $p^{\pi}(a \mid s, u')$ , as it requires rewarding observations as training data. Hence, it is not possible to directly use Proposition 9 to combine latent learning with low-variance estimators. In the following section, we propose a possible way forward to circumvent this issue.
H.4 Learning credit assignment representations with an information bottleneck
Here, we outline a strategy where we learn a latent representation  $Z$  that retains useful information on the underlying states  $S$ , and crucially has a latent space structure such that  $p^{\pi}(Z' = z' \mid S = s, U' = u')$  is easy to approximate. Then, leveraging Proposition 9 (and replacing  $S'$  by  $Z'$ ) allows us to learn a hindsight representation based on  $Z$ , enabling latent learning, while reducing the spurious contributions by counterfactual reasoning with  $p^{\pi}(z' \mid s, u')$ .
To achieve this, we use the Variational Information Bottleneck approach [83, 92], closely related to the  $\beta$  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\mid S,A;\theta)$  parameterized by  $\theta$ , and we 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  $\mathcal{I}(Z';S',A'\mid S,U')$  under some information bottleneck. We condition on  $S$  and  $U'$ , to end up later with a 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
$$
H - D \leq \mathcal {I} \left(Z ^ {\prime}; S ^ {\prime}, A ^ {\prime} \mid S, U ^ {\prime}\right) \leq \text {R a t e} \tag {112}
$$
with entropy  $H$ , distortion  $D$  and rate defined as follows:
$$
H = - \mathbb {E} _ {S ^ {\prime}, A ^ {\prime}, S, U ^ {\prime}} \left[ \log p ^ {\pi} \left(S ^ {\prime}, A ^ {\prime} \mid S, U ^ {\prime}\right) \right] \tag {113}
$$
$$
D = - \mathbb {E} _ {S, U ^ {\prime}} \left[ \mathbb {E} _ {S ^ {\prime}, A ^ {\prime} | S, U ^ {\prime}} \left[ \int \mathrm {d} z ^ {\prime} p \left(z ^ {\prime} \mid s ^ {\prime}, a ^ {\prime}; \theta\right) \log q \left(s ^ {\prime}, a ^ {\prime} \mid s, z ^ {\prime}; \psi\right) \right] \right] \tag {114}
$$
$$
\operatorname {R a t e} = \mathbb {E} _ {S, U ^ {\prime}} \left[ \mathbb {E} _ {S ^ {\prime}, A ^ {\prime} | S, U ^ {\prime}} \left[ D _ {\mathrm {K L}} \left(p \left(z ^ {\prime} \mid s ^ {\prime}, a ^ {\prime}; \theta\right) \| q \left(z ^ {\prime} \mid s, u ^ {\prime}; \phi\right)\right) \right] \right] \tag {115}
$$
where the decoder  $q(s', a' \mid s, z'; \psi)$  is a variational approximation to  $p(s', a' \mid s, z')$ , and the marginal  $q(z' \mid s, u'; \phi)$  is a variational approximation to the true marginal  $p(z' \mid s, u')$ . The distortion  $D$  quantifies how well we can decode the state-action pair  $(s', a')$  from the encoding
$z^{\prime}$ , by using a decoder  $q(s^{\prime},a^{\prime}\mid s,z^{\prime};\psi)$  parameterized by  $\psi$ . The distortion is reminiscent of an autoencoder loss, and hence encourages the encoding  $p(z^{\prime}\mid s^{\prime},a^{\prime})$  to retain as much information as possible on the state-action pair  $(s^{\prime},a^{\prime})$ . The rate measures the average KL-divergence 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^{\prime}$  with an optimal code designed for the variational marginal  $q(z^{\prime}\mid s,u^{\prime};\phi)$ .
We can use the rate to impose an information bottleneck on  $Z'$ . If we constrain the rate to  $\text{Rate} \leq a$  with  $a$  some positive constant, we restrict the amount of information  $Z'$  can encode about  $(S', A')$ , as  $p(z' \mid s', a'; \theta)$  needs to remain close to the marginal  $q(z' \mid s, u'; \phi)$ , quantified by the KL-divergence. We can maximize the mutual information  $\mathcal{I}(Z'; S', A' \mid 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  _ {\theta , \phi , \psi} D + \beta \text {R a t e} \tag {116}
$$
Here, the  $\beta$  parameter determines the strength of the information bottleneck. This formulation is equivalent to the  $\beta$  Variational Autoencoder [93], and for  $\beta = 1$  we recover the Variational Autoencoder [94].
To understand why this information bottleneck approach is useful to learn credit assignment representations  $Z$ , we examine the rate in more detail. We can rewrite the rate as
$$
\operatorname {R a t e} = \mathbb {E} _ {S, U ^ {\prime}} \left[ D _ {\mathrm {K L}} \left(p \left(z ^ {\prime} \mid s, u ^ {\prime}\right) \| q \left(z ^ {\prime} \mid s, u ^ {\prime}; \phi\right)\right) \right] + \tag {117}
$$
$$
\mathbb {E} _ {S, U ^ {\prime}, Z ^ {\prime}} \left[ D _ {\mathrm {K L}} \left(p \left(s ^ {\prime}, a ^ {\prime} \mid s, z ^ {\prime}\right) \| p \left(s ^ {\prime}, a ^ {\prime} \mid s, u ^ {\prime}\right)\right) \right] \tag {118}
$$
Hence, optimizing the information bottleneck objective (116) w.r.t.  $\phi$  fits the variational marginal  $q(z^{\prime}\mid s,u^{\prime};\phi)$  to the true marginal  $p(z^{\prime}\mid s,u^{\prime})$  induced by the encoder  $p(z^{\prime}\mid s^{\prime},a^{\prime};\theta)$ . Proposition 9 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(z^{\prime}\mid s,u^{\prime};\phi)$  that approximates the true marginal, which we then can use to obtain contribution coefficients with less spurious contributions by leveraging Proposition 9. Furthermore, minimizing the rate w.r.t.  $\theta$  shapes the latent space of  $Z^{\prime}$  such that the true marginal  $p(z^{\prime}\mid s,u^{\prime})$  moves closer towards the variational marginal  $q(z^{\prime}\mid s,u^{\prime};\phi)$ .
In summary, the outlined information bottleneck approach for learning a credit assignment representation  $Z$  is a promising way forward to 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  $\beta$ . Minimizing the distortion results in a detailed encoding  $Z'$  with high mutual information with the state-action pair  $(S', A')$ , which can be leveraged for latent learning. Minimizing the rate shapes the latent space of  $Z'$  in such way that the true marginal  $p(z' \mid s, u')$  can be accurately approximated within the variational family of  $q(z' \mid s, u'; \phi)$ , and fits the parameters  $\phi$  resulting in an accurate marginal model. We can then leverage the variational marginal  $q(z' \mid s, u'; \phi)$  to perform counterfactual reasoning on the rewarding state encodings  $Z'$  according to Proposition 9, resulting in a low-variance COCOA-estimator.
I Contribution analysis and causality
I.1 Causal interpretation of COCOA
COCOA is closely connected to causality theory [42], where the contribution coefficients (3) correspond 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, as causation is in general not the same as correlation. We start with representing the MDP combined with the policy as a directed acyclic graphical model (c.f. Fig. 6a in App. C.1). 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_{t}$  to a specific value  $a$  independent from its parents  $S_{t}$ , and see what the influence is on the probability distributions of other nodes in the graph, e.g.  $R_{t + 1}$ . These interventions are formalized with do-calculus [42],
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  $V_{i}$  equal to  $v$  as  $\mathrm{Do}(V_i = v)$ , and can be used to investigate causal relations.
Using the graphical model of Fig. 6b that abstracts time, we can use do-interventions to quantify the causal contribution of an action  $A_{t} = a$  taken in state  $S_{t} = s$  upon reaching the rewarding outcome  $U^{\prime} = u^{\prime}$  in the future as,
$$
w _ {\mathrm {D o}} (s, a, u ^ {\prime}) = \frac {p ^ {\pi} \left(U ^ {\prime} = u ^ {\prime} \mid S _ {t} = s , \operatorname {D o} \left(A _ {t} = a\right)\right)}{\sum_ {\tilde {a} \in \mathcal {A}} \pi (\tilde {a} \mid s) p ^ {\pi} \left(U ^ {\prime} = u ^ {\prime} \mid S _ {t} = s , \operatorname {D o} \left(A _ {t} = \tilde {a}\right)\right)} - 1. \tag {119}
$$
As conditioning on  $S$  satisfies the backdoor criterion [42] for  $U'$  w.r.t.  $A$ , the interventional distribution  $p^{\pi}\big(U' = u' \mid S_t = s, \mathrm{Do}(A_t = a)\big)$  is equal to the observational distribution  $p^{\pi}\big(U' = u' \mid S_t = s, A_t = a\big)$ . Hence, the causal contribution coefficients of Eq. 119 are equal to the contribution coefficients of Eq. 3 used by COCOA.
I.2 Extending COCOA to counterfactual interventions
Within causality theory, counterfactual 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 a 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 a structural causal model (SCM), as shown in Figure 14. The SCM expresses all conditional distributions as deterministic functions with independent noise variables  $N$ , akin to the reparameterization trick [94]. In causal reasoning, we perform do-interventions on nodes, which is equivalent to cutting all incoming edges to a 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^{\pi}(\{N\} \mid T = \tau)$ , with  $\{N\}$  the set of all noise variables, given the observed trajectory  $\tau$ . 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 can 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  $h$  to the counterfactual hindsight distribution  $p_{\tau}^{\pi}\big(A_t = a \mid S_t = s, U' = u'\big)$ , where the  $\tau$ -subscript indicates the counterfactual distribution incorporating the noise posterior. Building upon the approach of Mesnard et al. [3], we can learn  $h(a \mid s', s, \Phi_t(\tau))$  to approximate the counterfactual hindsight distribution, with  $\Phi_t(\tau)$  a summary statistic trained to encapsulate the information of the posterior noise distribution  $p(\{N\} \mid T = \tau)$ . Mesnard et al. [3] show that such a summary statistic  $\Phi_t(\tau)$  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_t$ . We can achieve both characteristics by (i) training  $h(a \mid s', s, \Phi_t(\tau))$  on the hindsight action classification task and backpropagating the gradients to  $\Phi_t(\tau)$ , and (ii) training  $\Phi_t$  on an independence maximization loss  $\mathcal{L}_{\mathrm{IM}}(s)$ , which is minimized iff
Table 6: Comparison of discounted policy gradient estimators
Method	Policy gradient estimator (帽子Vπ(s0))
REINFORCE	Σt≥0 γt∇θ log π(At | St) Σk≥0 γkRt+k
Advantage	Σt≥0 γt∇θ log π(At | St) (Σk≥0 γkRt+k - Vγπ(St))
Q-critic	Σt≥0 γt Σa∈A ∇θπ(a | St)Qγπ(St, a)
COCOA	Σt≥0 γt∇θ log π(At | St) Rt + Σa∈A ∇θπ(a | St) Σk≥1 γkwγ(St, a, Ut+k)Rt+k
HCA+	Σt≥0 γt∇θ log π(At | St) Rt + Σa∈A ∇θπ(a | St) Σk≥1 γkwγ(St, a, St+k)Rt+k
$A_{t}$  and  $\Phi_t$  are conditionally independent given  $S_{t}$ . An example is to minimize the KL divergence between  $\pi (a_{t} \mid s_{t})$  and  $p(a_{t} \mid s_{t}, \Phi_{t}(\tau))$  where the latter can be approximated by training a classifier  $q(a_{t} \mid s_{t}, \Phi_{t}(\tau))$ . Leveraging this approach to extend COCOA towards counterfactual interventions is an exciting direction for future research.
J 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 a fixed probability of  $(1 - \gamma)$  of transitioning to the absorbing state  $s_{\infty}$  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  $(\mathcal{S},\mathcal{A},p,p_r,\gamma)$ , with discount factor  $\gamma \in [0,1]$ , and  $(\mathcal{S},\mathcal{A},p,p_r)$  as defined in Section 2. The discounted value function  $V_{\gamma}^{\pi}(s) = \mathbb{E}_{T\sim \mathcal{T}(s,\pi)}[\sum_{t = 0}^{\infty}\gamma^{t}R_{t}]$  and action value function  $Q_{\gamma}^{\pi}(s,a) = \mathbb{E}_{T\sim \mathcal{T}(s,a,\pi)}[\sum_{t = 0}^{\infty}\gamma^{t}R_{t}]$  are the expected discounted return when starting from state  $s$ , or state  $s$  and action  $a$  respectively. Table 6 shows the policy gradient estimators of  $V_{\gamma}^{\pi}(s)$  for REINFORCE, Advantage and Q-critic.
In a discounted MDP, it matters at which point in time we reach a rewarding outcome  $u$ , as the corresponding rewards are discounted. Hence, we adjust the contribution coefficients to