SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathbb {E} _ {\pi} \left[ \hat {\nabla} _ {\theta} ^ {\cdot} V ^ {\pi} \right] = \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a | s _ {0}) \mathbb {E} _ {\pi} \left[ \hat {A} ^ {\pi} \left(s _ {0}, a\right) \right] + \tag {90}
$$
$$
\sum_ {s \in \mathcal {S}} \sum_ {a _ {0} \in \mathcal {A}} \pi \left(a _ {0} \mid s _ {0}\right) M \left(s _ {0}, a _ {0}, s, T\right) \sum_ {a \in \mathcal {A}} \nabla_ {\theta} \pi (a | s) \mathbb {E} _ {\pi} \left[ \hat {A} ^ {\pi} (s, a) \right] \tag {91}
$$
Computing the ground-truth policy gradient. To compute the ground-truth policy gradient, we can apply the same strategy as for Eq. 90 and replace the expected (possibly biased) advantage  $\mathbb{E}_{\pi}[\hat{A}^{\pi}(s_0,a)]$  by the ground-truth advantage function, computed with the ground-truth action-value function (c.f. Section E.2.1).
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. A higher SNR indicates that we need fewer trajectories to estimate accurate policy gradients, hence reflecting better credit assignment. To obtain meaningful
scales, we normalize the bias and variance by the norm of the ground-truth policy gradient.
$$
S N R = \frac {\left\| \nabla_ {\theta} V ^ {\pi} \right\| ^ {2}}{\mathbb {E} _ {\pi} \left[ \left\| \hat {\nabla} _ {\theta} V ^ {\pi} - \nabla_ {\theta} V ^ {\pi} \right\| ^ {2} \right]} \tag {92}
$$
$$
\text {V a r i a n c e} = \frac {\mathbb {E} _ {\pi} \left[ \left\| \hat {\nabla} _ {\theta} V ^ {\pi} - \mathbb {E} _ {\pi} \left[ \hat {\nabla} _ {\theta} V ^ {\pi} \right] \right\| ^ {2} \right]}{\| \nabla_ {\theta} V ^ {\pi} \| ^ {2}} \tag {93}
$$
$$
\text {B i a s} = \frac {\left\| \mathbb {E} _ {\pi} \left[ \hat {\nabla} _ {\theta} V ^ {\pi} \right] - \nabla_ {\theta} V ^ {\pi} \right\| ^ {2}}{\left\| \nabla_ {\theta} V ^ {\pi} \right\| ^ {2}} \tag {94}
$$
We compute the full expectations  $\mathbb{E}_{\pi}[\hat{\nabla}_{\theta}V^{\pi}]$  and ground-truth policy gradient  $\nabla_{\theta}V^{\pi}$  by leveraging our dynamic programming setup, while for the expectation of the squared differences  $\mathbb{E}_{\pi}\left[\|\hat{\nabla}_{\theta}V^{\pi}-\nabla_{\theta}V^{\pi}\|^{2}\right]$  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, variance and bias of a partial policy gradient estimator considering only  $t = 0$  in the outer sum (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, we consider a shadow training setup in which the 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. 1B (right), Fig. 2B and Fig. 3C-D.
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 a 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 or the right of the agent and which can stochastically assume two different reward values. Finally, the agent reaches a door which it can open with the key to collect a treasure reward.
In our simulations we represent states using a nine-dimensional vector, encoding the relative position on the track as a 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-noise ratio (SNR) of the REINFORCE policy gradient [6] for solving the main task of picking up the key can be approximated by
$$
\mathrm {S N R} ^ {\text {R E I N F}} \approx \frac {\| \mathbb {E} \left[ \hat {\nabla} _ {\theta} ^ {\text {R E I N F}} V ^ {\pi} \right] \| ^ {2}}{C (\theta) \mathbb {V} \left[ \sum_ {t \in T 2} R _ {t} \right] + \operatorname {T r} \left[ \mathbb {V} \left[ \hat {\nabla} _ {\theta} ^ {\text {R E I N F}} V ^ {\pi} \mid \text {n o T 2} \right] \right]}. \tag {95}
$$
with  $\hat{\nabla}_{\theta}^{\mathrm{REINF}}V^{\pi}$  the REINFORCE estimator (c.f. Table 1),  $C(\theta)$  a reward-independent constant,  $T2$  the set of time steps corresponding to the distractor task, and  $\operatorname{Tr}\left[\mathbb{V}[\hat{\nabla}_{\theta}^{\mathrm{REINF}}V^{\pi} \mid \text{no T2}]\right]$  the trace of the covariance matrix of the REINFORCE estimator in an equivalent task setup without distractor rewards. Hence, we can adjust the difficulty of the task by increasing the number of distractor rewards and their variance. We perform experiments with environments of length  $L \in \{20, 40, \dots, 100\}$  choosing the reward values such that the total distractor reward remains approximately constant. Concretely, distractor rewards are sampled as  $r_{\text{distractor}} \sim \mathcal{U}\left(\left\{\frac{2}{L}, \frac{18}{L}\right\}\right)$  and the treasure leads to a deterministic reward of  $r_{\text{tresource}} = \frac{4}{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 can 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  $n_a$  determining the branching factor, and the state overlap  $o_s$ , defined as the number of overlapping children from two neighbouring nodes. The states are represented by two integers:  $i < d$  representing the current level in the tree, and  $j$  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 \gets i + 1$  and  $j \gets j(n_a - o_s) + a$ , where we represent the action by an integer  $0 \leq a < n_a$ . We assign each state-action pair a reward  $r \in \{-2, -1, 0, 1, 2\}$ , computed as
$$
r (s, a) = \left((\operatorname {i d x} (s) + a p + \text {s e e d}) \bmod n _ {r}\right) - n _ {r} / / 2 \tag {96}
$$
with the modulo operator mod, the number of reward values  $n_r = 5$ ,  $\mathrm{idx}(s)$  a unique integer index for the state  $s = (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  $n_g$  groups:
$$
u = \left(\operatorname {i d x} (s) + a p + \text {s e e d}\right) \mod \left(n _ {r} n _ {g}\right) - n _ {r} / / 2 \tag {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  $U$  corresponding to  $n_g = 4$  and  $n_g = 32$ .
E.7 Task interleaving environment setup
We simplify the task interleaving environment of Mesnard et al. [3] in the following way: the environment 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  $1 - \gamma$ . In each context, a set of  $2O$  objects is given, and a 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 a 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 interleaved 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 2 actions the agent can take: choose right, or choose left.
The state is defined by the following variables:
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 no 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 is a 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. If the agent is in an answer room, then regardless of the action, it receives a 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, it receives 0 reward. Then, the current context is removed from the set of open contexts.
If the agent is in a query room, then the agent is presented with 2 objects, one rewarding and one not, and must choose the side of the rewarding object. Regardless, the agent receives 0 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 a probability  $\frac{C_o}{B}$  (where  $C_o$  is the number of open contexts), the agent samples uniformly one of the open contexts. Otherwise, the agent samples uniformly one of the non-open contexts. The 2 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 a probability  $1 - \gamma$  of dying, in which case the agent receives the reward but is put into a terminal state at the next timestep.
Environment parameters For all our experiments, we choose  $C = 5, B = 3, O = 2$  and  $\gamma = 0.95$ .
Visualization of contribution coefficient magnitudes. For the heatmaps shown in Fig 4C, we computed each entry as  $\sum_{a\in \mathcal{A}}\pi (a|s_t)\cdot |w(s_t,a,r_{t + k})|$  where states  $s_t$  and rewards  $r_{t + k}$  are 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 action-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  $s$  and rewarding outcome encoding  $u'$  as inputs to perform poorly. We hypothesize that this is due to the policy dependence of the hindsight distribution creating a moving target during learning as the
policy changes. Leveraging Proposition 5, we therefore add the policy logits as an extra input to the hindsight network to 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 a logit for each possible action. The multiplicative interaction denoted by  $\otimes$  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 a gated version of the ReLU nonlinearity in the second layer which computes the difference between the split, rectified input effectively halving the output dimension:
$$
\operatorname {R e L U} - \mathrm {g} (x) = \operatorname {R e L U} \left(x _ {0: n / 2}\right) - \operatorname {R e L U} \left(x _ {n / 2: n}\right) \tag {98}
$$
with  $n$  the dimension of  $x$ . Gating in combination with the multiplicative interaction is a useful inductive bias for the hindsight model, since for actions which have zero contribution towards the rewarding outcome  $u$  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 multiplicative interaction.
E.8.2 Optimization
For training of all models we use the AdamW 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(λ) 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  $\frac{1}{\sqrt{n_{\mathrm{input}}}}$  where  $n_{\mathrm{input}}$  is the input dimension of the respective layer and.
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 a network to predict the rewards given a state-action pair, and take the penultimate layer as the feature  $u = f(s,a)$ . For the latter, there exist multiple approaches. When using a deterministic encoding  $u = f(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\mid S,A)$  we can impose an information bottleneck on the reward features  $U$ , enforcing the encoding to discard as much information about  $U$  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 to ensure positivity. A ReLU nonlinearity with straight through gradient estimator is then applied to the transformation. Finally, an action-independent readout weight transforms the activations to the prediction of the reward. The mask parameters are initialized to 1. Weights of the readout layer are initialized with a Gaussian distribution of mean 0 and std  $\frac{1}{\sqrt{d}}$  where  $d$  is the dimension of the input.
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 L1 regularization term to all weights up to the penultimate layer with a strength of  $\eta_{L_1} = 0.001$ . The readout weights are trained with standard L2 regularization with strength  $\eta_{L_2} = 0.03$ . For the linear key-to-door experiments, we choose  $(\eta_{L_1},\eta_{L_2}) = (0.001,0.03)$ . For the task interleaving environment, we choose  $(\eta_{L_1},\eta_{L_2}) = (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.
Table 2: The range of values swept over for each hyperparameter in a grid search for the linear key-to-door environment and task interleaving environment.
Hyperparameter	Range
lr_agent	{0.003, 0.001, 0.0003, 0.0001}
lr_hindsight	{0.01, 0.003, 0.001}
lr_value	{0.01, 0.003, 0.001}
lr_qvalue	{0.01, 0.003, 0.001}
lr_features	{0.01, 0.003, 0.001}
td_lambda_value	{0., 0.5, 0.9, 1.}
td_lambda_qvalue	{0., 0.5, 0.9, 1.}
entropy_reg	{0.3, 0.1, 0.03, 0.01}
Table 3: Hyperparameter values on the linear key-to-door environment. The best performing hyperparameters were identical across all environment lengths. COCOA-r stands for COCOA-return, COCOA-f for COCOA-feature.
Hyperparameter	COCOA-r	COCOA-f	HCA+	Q-critic	Advantage	REINFORCE	TrajCV	HCA-return
lr_agent	0.0003	0.0003	0.0003	0.0003	0.001	0.0003	0.003	0.0001
lr_hindsight	0.003	0.003	0.003	-	-	-	-	0.001
lr_value	-	-	-	-	0.001	-	-	-
lr_qvalue	-	-	-	0.003	-	-	0.01	-
lr_features	-	0.003	-	-	-	-	-	-
td_lambda_value	-	-	-	-	1.	-	-	-
td_lambda_qvalue	-	-	-	0.9	-	-	0.9	-
Pretraining. To learn the reward features, we collect the first  $B_{\mathrm{feature}}$  mini-batches of episodes in a buffer using a frozen random policy. We then sample triplets  $(s, a, r)$  from the buffer and train with full-batch gradient descent using the Adam optimizer over  $N_{\mathrm{feature}}$  steps. Once trained, the reward network is frozen, and the masked inputs, which are discretized using the element-wise threshold function  $\mathbb{1}_{x > 0.05}$ , are used to 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_{\mathrm{feature}}$  batches of episodes. For the linear key-to-door experiments, we chose  $(B_{\mathrm{feature}}, N_{\mathrm{feature}}) = (30, 20000)$ . For the task interleaving environment, we chose  $(B_{\mathrm{feature}}, N_{\mathrm{feature}}) = (90, 30000)$ .
E.8.4 Hyperparameters
Linear key-to-door setup (performance) For all our experiments on the linear key-to-door environment, 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 a 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\%$  of all seeds, given a 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 4: Hyperparameter values on the task interleaving environment.
Hyperparameter	COCOA-reward	COCOA-feature	HCA+	Q-critic	Advantage	REINFORCE	TrajCV
lr_agent	0.0003	0.0003	0.0001	0.0003	0.001	0.001	0.001
lr_hindsight	0.001	0.001	0.001	-	-	-	-
lr_value	-	-	-	-	0.001	-	-
lr_qvalue	-	-	-	0.003	-	-	0.001
lr_features	-	0.001	-	-	-	-	-
entropy_reg	0.01	0.01	0.01	0.01	0.01	0.01	0.01
td_lambda_value	-	-	-	-	1.	-	-
td_lambda_qvalue	-	-	-	0.9	-	-	0.9
Table 5: 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-aliasing)
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_reg. For the final values of entropy_reg 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, we used the best performing hyperparameters found in the performance setting. We kept a batch size of 8 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 following a similar selection procedure as 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 a 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, we 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, we found that a 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 a 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 a 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 in Table 4.