SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathcal {L} (\Phi) = \mathbb {E} _ {s _ {t}, a _ {t}, s _ {t + 1} \sim \rho , s ^ {\prime} \sim \mu} \left[ \left(F (s _ {t}, a _ {t}, z) ^ {\top} B (s ^ {\prime}) - \gamma \bar {F} (s _ {t + 1}, \pi_ {z} (s _ {t + 1}), z) ^ {\top} \bar {B} (s ^ {\prime})\right) ^ {2} \right] \\ - 2 \mathbb {E} _ {s _ {t}, a _ {t} \sim \rho} \left[ F (s _ {t}, a _ {t}, z) ^ {\top} B (s _ {t}) \right] \\ + 2 \gamma (1 - \lambda) \mathbb {E} _ {s _ {t}, a _ {t}, s _ {t + 1} \sim \rho} \left[ \mu \left(s _ {t}\right) F \left(s _ {t}, a _ {t}, z\right) ^ {\top} B \left(s _ {t}\right) \bar {F} \left(s _ {t + 1}, \pi_ {z} \left(s _ {t + 1}\right), z\right) ^ {\top} \bar {B} \left(s _ {t}\right) \right] \tag {H.1} \\ \end{array}
$$
Even though the  $\lambda O$  is not a measure, we can use the above loss to the continuous case, pretending as though we could take the Radon-Nikodym derivative  $\frac{\Phi(s,ds')}{\mu(ds')}$ .
H.1 Experimental Results with the FB Parameterization
To show that knowing the correct value of  $\lambda$  leads to improved performance, we trained  $\lambda_0$  with the FB parameterization on the FourRooms task of Fig. 5.3, but with each episode initialized at a random start state and with two random goal states. Average per-epoch reward is shown in Fig. H.2. We tested performance with  $\lambda_{\mathrm{true}}$ ,  $\lambda_{\mathrm{agent}} \in \{0.5, 1.0\}$ , where  $\lambda_{\mathrm{true}}$  denotes the true environment diminishing rate and  $\lambda_{\mathrm{agent}}$  denotes the diminishing rate that the agent uses. For the purpose of illustration, we do not allow the agent to learn  $\lambda$ . We see in Fig. H.2 that using the correct  $\lambda$  leads to significantly increased performance. In particular, the left plot shows that assuming  $\lambda = 1$ , i.e., using the SR, in a diminishing environment can lead to highly suboptimal performance.
Hyperparameters used are given in Table 1 (notation as in Algorithm 3).
Hyperparameter	Value
M	100
E	100
N	25
B	128
T	50
γ	0.99
α	0.95
η	0.001
τ	200
ε	1
Table 1:  $\lambda \mathrm{O} - \mathrm{{FB}}$  hyperparameters.
Figure H.1: Performance of the  $\lambda$ O-FB with two values of  $\lambda$ . Results averaged over six seeds and 10 episodes per seed. Error bars indicate standard error.
H.2  $\lambda O$  and the Marginal Value Theorem
To study whether the agent's behavior is similar to behavior predicted by the MVT, we use a very simple task with constant starting state and vary the distance between rewards (see Fig. H.1(a)). When an agent is in a reward state, we define an MVT-optimal leaving time as follows (similar to that of [8] but accounting for the non-stationarity of the reward).
Let  $R$  denote the average per-episode reward received by a trained agent,  $r(s_{t})$  denote the reward received at time  $t$  in a given episode,  $R_{t} = \sum_{u=0}^{t} r(s_{u})$  denote the total reward received until time  $t$  in the episode, and let  $T$  be episode length. Then, on average, the agent should leave its current reward state at time  $t$  if the next reward that it would receive by staying in  $s_{t}$ , i.e.,  $\lambda r(s_{t})$ , is less than
$$
\frac {R - R _ {t}}{T}.
$$
In other words, the agent should leave a reward state when its incoming reward falls below the diminished average per-step reward of the environment. We compute  $R$  by averaging reward received by a trained agent over many episodes.
Previous studies have trained agents that assume stationary reward to perform foraging tasks, even though the reward in these tasks is non-stationary. These agents can still achieve good performance and MVT-like behavior [8]. However, because
Figure H.2: Analysis of MVT-like behavior of  $\lambda 0$ -FB. a) Three environments with equal start state and structure but different distances between reward states. b) Difference between the agent's leave times and MVT-predicted leave times for  $\gamma = 0.99$ , with discounting taken into account. The agent on average behaves similar to the discounted MVT. c) Difference between the agent's leave times and MVT-predicted leave times for  $\gamma = 1.0$ , i.e., with no discounting taken into account. The agent on average behaves similar to the MVT. Results for (b) and (c) are averaged over three seeds. Error bars indicate standard error.
they target the standard RL objective
$$
\mathbb {E} _ {\pi} \left[ \sum_ {k = 0} ^ {\infty} \gamma^ {k} r (s _ {t + k}) \mid s _ {t} = s \right],
$$
which requires  $\gamma < 1$  for convergence, optimal behavior is recovered only with respect to the discounted MVT, in which  $R$  (and in our case,  $R_{t}$ ) weights rewards by powers of  $\gamma$  [8].
In Fig. H.1(b-c) we perform a similar analysis to that of [8] and show that, on average over multiple distances between rewards,  $\lambda$ O-FB performs similarly to the discounted MVT for  $\gamma = 0.99$  and the standard MVT for  $\gamma = 1.0$ . An advantage of the  $\lambda$ O is that it is finite for  $\gamma = 1.0$  provided that  $\lambda < 1$ . Hence, as opposed to previous work, we can recover the standard MVT without the need to adjust for discounting.
I SAC
Mitigating Value Overestimation One well-known challenge in deep RL is that the use of function approximation to compute values is prone to overestimation. Standard approaches to mitigate this issue typically do so by using two value functions and either taking the minimum  $\min_{i\in \{1,2\}}Q_i^\pi (s,a)$  to form the Bellman target for a given  $(s,a)$  pair [40] or combining them in other ways [41]. However, creating multiple networks is expensive in both computation and memory. Instead, we hypothesized that it might be possible to address this issue by using  $\lambda$ -based values. To test this idea, we modified the Soft Actor-Critic [SAC; 42] algorithm to compute  $\lambda$  Fs-based values by augmenting the soft value target  $\mathcal{T}_{soft}Q = r_t + \gamma \mathbb{E}V_{soft}(s_{t + 1})$ , where  $V_{soft}(s_{t + 1})$  is given by the expression
$$
\begin{array}{l} \mathbb {E} _ {a _ {t + 1} \sim \pi (\cdot | s _ {t + 1})} \left[ \bar {Q} \left(s _ {t + 1}, a _ {t + 1}\right) + (\lambda - 1) \mathbf {w} ^ {\top} \left(\phi \left(s _ {t}, a _ {t}\right) \odot \varphi_ {\lambda} \left(s _ {t + 1}, a _ {t + 1}\right)\right) \right. \\ \left. - \alpha \log \pi \left(a _ {t + 1} \mid s _ {t + 1}\right) \right] \\ \end{array}
$$
A derivation as well as pseudocode for the modified loss is provided in Appendix E.5. Observe that for  $\lambda = 1$ , we recover the standard SAC value target, corresponding to an assumed stationary reward. We apply this modified SAC algorithm, which we term  $\lambda$ -SAC to feature-based Mujoco continuous control tasks within OpenAI Gym [43]. We found that concatenating the raw state and action observations  $\tilde{\phi}_t = [s_t, a_t]$  and normalizing them to  $[0, 1]$  make effective regressors to the reward. That is, we compute base features as
$$
\phi_ {t} ^ {b} = \frac {\tilde {\phi} _ {t} ^ {b} - \min _ {b} \tilde {\phi} _ {t} ^ {b}}{\max _ {b} \tilde {\phi} _ {t} ^ {b} - \min _ {b} \tilde {\phi} _ {t} ^ {b}},
$$
where  $b$  indexes  $(s_t, a_t)$  within a batch. Let  $X \in [0,1]^{B \times D}$  be the concatenated matrix of features for a batch, where  $D = \dim(\mathcal{S}) + \dim(\mathcal{A})$ . Then,
$$
\mathbf {w} _ {t} = \left(X ^ {\intercal} X\right) ^ {- 1} X ^ {\intercal} \mathbf {r},
$$
where here  $\mathbf{r}$  denotes the vector of rewards from the batch. In addition to using a fixed  $\lambda$  value, ideally we'd like an agent to adaptively update  $\lambda$  to achieve the best balance of optimism and pessimism in its value estimates. Following [44], we frame this decision as a multi-armed bandit problem, discretizing  $\lambda$  into three possible values  $\{0, 0.5, 1.0\}$  representing the arms of the bandit. At the start of each episode, a random value of  $\lambda$  is sampled from these arms and used in the value function update. The probability of each arm is updated using the Exponentially Weighted Average Forecasting algorithm [45], which modulates the probabilities in proportion to a feedback score. As in [44], we use the difference in cumulative (undiscounted) reward between the current episode  $\ell$  and the previous one  $\ell - 1$  as this feedback signal:  $R_{\ell} - R_{\ell - 1}$ . That is, the probability of selecting a given value of  $\lambda$  increases if performance is improving and decreases if it's decreasing. We use identical settings for the bandit algorithm as in [44]. We call this variant  $\lambda$ -SAC.
We plot the results for SAC with two critics (as is standard), SAC with one critic, SAC with a single critic trained with  $\lambda$ F-based values (" $x$ -SAC" denotes SAC trained with a fixed  $\lambda = x$ ), and  $\lambda$ -SAC trained on the HalfCheetah-v2 and Hopper-v2 Mujoco environments. All experiments were repeated over eight random seeds. HalfCheetah-v2 was found by [44] to support "optimistic" value estimates in that even without pessimism to reduce overestimation it was possible to perform well. Consistent with this, we found that single-critic SAC matched the performance of standard SAC, as did 1-SAC (which amounts to training a standard value function with the auxiliary task of SF prediction). Fixing lower values of  $\lambda$  performed poorly, indicating that over-pessimism is harmful in this environment. However,  $\lambda$ -SAC eventually manages to learn to set  $\lambda = 1$  and matches the final performance of the best fixed algorithms. Similarly, in [44] it was observed that strong performance in Hopper-v2 was associated with pessimistic value estimates. Consistent with this,  $\lambda$ -SAC learns to select lower values of  $\lambda$ , again matching the performance of SAC while only requiring one critic and significantly reducing the required FLOPS Fig. I.2. We consider these results to be very preliminary, and hope to perform more experiments on other environments. We also believe  $\lambda$ -SAC could be improved by using the difference between the current episode's total reward and the average of the total rewards from previous episodes  $R_{\ell} - (\ell - 1)^{-1}\sum_{i=1}^{\ell-1} R_{i}$  as a more
stable feedback signal for the bandit. There is also non-stationarity in the base features due to the per-batch normalization, which could also likely be improved. Hyperparameters are described in Table 2.
Hyperparameter	Value
Collection Steps	1000
Random Action Steps	10000
Network Hidden Layers	256:256
Learning Rate	3 × 10-4
Optimizer	Adam
Replay Buffer Size	1 × 106
Action Limit	[-1,1]
Exponential Moving Avg. Parameters	5 × 10-3
(Critic Update:Environment Step) Ratio	1
(Policy Update:Environment Step) Ratio	1
Has Target Policy?	No
Expected Entropy Target	-dim(A)
Policy Log-Variance Limits	[-20,2]
Bandit Learning Rate*	0.1
λ Options*	{0,0.5,1.0}
Table 2: Hyperparameters for SAC Mujoco experiments. *Only applicable to  $\lambda$  -SAC
SAC (2 critics)
SAC
Figure I.1:  $\lambda$ -SAC (1 critic) matches the performance of SAC (2 critics) by adapting  $\lambda$  online.
1-SAC
0.5-SAC
#
一，
0-SAC
λ-SAC
Figure I.2:  $\lambda$ -SAC matches performance while saving in computational cost.
J Replenishing Rewards
We list below a few candidate reward replenishment schemes, which are visualized in Fig. J.1.
Time elapsed rewards
$$
r (s, t) = \lambda^ {n (s, t) / m (s, t)} \bar {r} (s),
$$
where  $m(s,t)$  is the time elapsed since the last visit to state  $s$ :
$$
m (s, t) \triangleq t - \max  \{j | s _ {t + j} = s, 0 \leq j \leq t - 1 \}.
$$
Due to the max term in  $m(s,t)$ , the corresponding representation
$$
\mathbb {E} _ {\boldsymbol \pi} \left[ \sum_ {k = 0} ^ {\infty} \gamma^ {k} \lambda^ {n (s, t) / m (s, t)} \mathbb {1} \big (s _ {t + k} = s ^ {\prime} \big) \Big | s _ {t} = s \right]
$$
does not admit a closed-form recursion. However, we empirically tested a version of this type of reward with  $Q_{\lambda}$ -learning in the TwoRooms environment, modified so that the exponent on  $\lambda$  is  $n(s,t) / (0.1m(s,t))$ . This was done so that reward replenishes at a slow rate, reducing the deviation from the standard diminishing setting. Episode length was capped at  $H = 100$ . All other settings are identical to the  $Q_{\lambda}$  experiment described in Appendix E. Results are presented in Fig. J.2 and a GIF is included in the supplementary material.
Eligibility trace rewards
$$
r (s, t) = \left(1 - (1 - \lambda_ {d}) \sum_ {j = 0} ^ {t - 1} \lambda_ {r} ^ {t - j} \mathbb {1} (s _ {t + j} = s)\right) \bar {r} (s),
$$
where  $\lambda_d, \lambda_r \in [0,1]$  are diminishment and replenishment constants, respectively. Denoting the corresponding representation by  $\Omega^{\pi}$ , i.e.,
$$
\Omega^ {\pi} (s, s ^ {\prime}) = \mathbb {E} \left[ \sum_ {k = 0} ^ {\infty} \gamma^ {k} \left(1 - (1 - \lambda_ {d}) \sum_ {j = 0} ^ {k} \lambda_ {r} ^ {k - j} \mathbb {1} (s _ {t + j} = s ^ {\prime})\right) \mathbb {1} (s _ {t + k} = s ^ {\prime}) \Bigg | s _ {t} = s \right],
$$
we obtain the following recursion:
$$
\begin{array}{l} \Omega^ {\pi} (s, s ^ {\prime}) = \mathbb {1} (s = s ^ {\prime}) (\lambda_ {d} - \gamma \lambda_ {r} (1 - \lambda_ {d}) \mathbb {E} _ {s _ {t + 1} \sim p ^ {\pi} (\cdot | s)} M _ {\gamma \lambda_ {r}} ^ {\pi} (s _ {t + 1}, s ^ {\prime})) \\ + \gamma \mathbb {E} _ {s _ {t + 1} \sim p ^ {\pi (\cdot | s)}} \Omega^ {\pi} \big (s _ {t + 1}, s ^ {\prime} \big), \\ \end{array}
$$
where  $M_{\gamma \lambda_r}^\pi$  denotes the successor representation of  $\pi$  with discount factor  $\gamma \lambda_r$ . This representation could be learned by alternating TD learning between  $\Omega^\pi$  and  $M_{\gamma \lambda_r}^\pi$ . We leave this to future work.
Total time rewards
$$
r (s, t) = \lambda_ {d} ^ {n (s, t)} \lambda_ {r} ^ {n (s, t) - t} \bar {r} (s),
$$
where  $\lambda_d, \lambda_r \in [0,1]$  are diminishment and replenishment constants, respectively. The corresponding representation is
$$
P ^ {\pi} (s, s ^ {\prime}) = \mathbb {E} \left[ \sum_ {k = 0} ^ {\infty} \gamma^ {k} \lambda_ {d} ^ {n _ {t} (s ^ {\prime}, k)} \lambda_ {r} ^ {k - n _ {t} (s ^ {\prime}, k)} \mathbb {1} (s _ {t + k} = s ^ {\prime}) \Bigg | s _ {t} = s \right],
$$
Figure J.1: Visualizing three different replenishment schemes. For all schemes,  $\bar{r}(s) = 1$  and visits to  $s$  are at  $t = 2,5$ . (Left) The time elapsed reward with  $\lambda = 0.5$ ; (Middle) The eligibility trace reward with  $\lambda_r = \lambda_d = 0.5$ ; (Right) The total time reward with  $\lambda_d = 0.5$ ,  $\lambda_r = 0.9$ .
Figure J.2: Performance on TwoRooms with replenishing rewards. Return is averaged over five runs, with shading indicating one unit of standard error.
which satisfies the following recursion:
$$
\begin{array}{l} P ^ {\pi} (s, s ^ {\prime}) = \mathbb {1} (s = s ^ {\prime}) + \gamma (\lambda_ {d} \mathbb {1} (s = s ^ {\prime}) + 1 \\ - \mathbb {1} (s = s ^ {\prime})) (\lambda_ {r} (1 - \mathbb {1} (s = s ^ {\prime})) + \mathbb {1} (s = s ^ {\prime})) \mathbb {E} _ {s _ {t + 1} \sim p ^ {\pi} (\cdot | s)} P ^ {\pi} (s _ {t + 1}, s ^ {\prime}). \\ \end{array}
$$
While neither the reward nor the representation are guaranteed to be finite,  $P^{\pi}$  could be learned via TD updates capped at a suitable upper bound.
K  $\lambda$  vs.  $\gamma$
We now briefly discuss the interaction between the temporal discount factor  $\gamma$  commonly used in RL and the diminishing utility rate  $\lambda$ . The key distinction between the two is that all rewards decay in value every time step with respect to  $\gamma$ , regardless of whether a state is visited or not. With  $\lambda$ , however, decay is specific to each state (or  $(s, a)$  pair) and only occurs when the agent visits that state. Thus,  $\gamma$  decays reward in a global manner which is independent of the agent's behavior, and  $\lambda$  decays reward in a local manner which dependent on the agent's behavior. In combination, they have the beneficial effect of accelerating convergence in dynamic programming (Fig. F.3). This indicates the potential for the use of higher discount factors in practice, as paired with a decay factor  $\lambda$ , similar (or faster) convergence rates could be observed even as agents are able to act with a longer effective temporal horizon.
L Compute Resources
The  $\lambda$ F-based experiments shown were run on a single NVIDIA GeForce GTX 1080 GPU. The recurrent A2C experiments took roughly 30 minutes, base feature learning for policy composition took approximately 45 minutes,  $\lambda$ F learning for policy composition took approximately 10 hours, and the SAC experiments took approximately 8 hours per run. The  $\lambda$ F training required roughly 30GB of memory due to the size of the dataset. All experiments in Section 6 and Appendix H were run on a single RTX5000 GPU and each training and evaluation run took about 30 minutes. All other experiments were run on a 2020 MacBook Air laptop 1.1 GHz Quad-Core Intel Core i5 CPU and took less than one hour to train.
M Broader Impact Statement
We consider this work to be primarily of a theoretical nature pertaining to sequential decision-making primarily in the context of natural intelligence. While it may have applications for more efficient training of artificial RL agents, it is hard to predict long-term societal impacts.