SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Another important and relevant sub-field of reinforcement learning is work which studies non-stationary rewards. Perhaps most relevant, the setting of DMU can be seen as a special case of submodular planning and reward structures [31, 32]. [31] focus specifically on planning and not the form of diminishment we study, while [32] is a concurrent work which focuses on the general class of submodular
reward problems and introduces a policy-based, REINFORCE-like method which is necessarily on-policy. In contrast, we focus on a particular sub-class of problems especially relevant to natural behavior and introduce a family of approaches which exploit this reward structure, are value-based (and which can be used to modify the critic in policy-based methods), and are compatible with off-policy learning. Other important areas include convex [33] and constrained [34, 35] MDPs. In these cases, non-stationarity is introduced by way of a primal-dual formulation of distinct problem classes into min-max games.
E Further Experimental Details
E.1 Policy Evaluation
We perform policy evaluation for the policy shown in Fig. 4.1 on the  $6 \times 6$  gridworld shown. The discount factor  $\gamma$  was set to 0.9 for all experiments, which were run for  $H = 10$  steps per episode. The error metric was the mean squared error:
$$
Q _ {\text {e r r o r}} \triangleq \frac {1}{| \mathcal {S} | | \mathcal {A} |} \sum_ {s, a} \left(Q ^ {\pi} (s, a) - \hat {Q} (s, a)\right) ^ {2}, \tag {E.1}
$$
where  $Q^{\pi}$  is the ground truth  $Q$ -values and  $\hat{Q}$  is the estimate. Transitions are deterministic. For the dynamic programming result, we learned the  $\lambda \mathrm{R}$  using Eq. (4.3) for  $\lambda \in \{0.5, 1.0\}$  and then measured the resulting values by multiplying the resulting  $\lambda \mathrm{R}$  by the associated reward vector  $\mathbf{r} \in \{-1, 0, 1\}^{36}$ , which was  $-1$  in all wall states and  $+1$  at the reward state  $g$ . We compared the results to the ground truth values. Dynamic programming was run until the maximum Bellman error across state-action pairs reduced below 5e-2. For the tabular TD learning result, we ran the policy for three episodes starting from every available (non-wall) state in the environment, and learned the  $\lambda \mathrm{R}$  for  $\lambda \in \{0.5, 1.0\}$  as above, but using the online TD update:
$$
\Phi_ {\lambda} (s _ {t}, a _ {t}) \gets \Phi_ {\lambda} (s _ {t}, a _ {t}) + \alpha \delta_ {t},
$$
$$
\delta_ {t} = \mathbf {1} (s _ {t}) \odot (1 + \gamma \lambda \Phi_ {\lambda} (s _ {t + 1}, a _ {t + 1})) + \gamma (1 - \mathbf {1} (s _ {t})) \odot \Phi_ {\lambda} (s _ {t + 1}, a _ {t + 1}) - \Phi_ {\lambda} (s _ {t}, a _ {t}),
$$
where  $a_{t+1} \sim \pi(\cdot \mid s_{t+1})$ . The learned  $Q$ -values were then computed in the same way as the dynamic programming case and compared to the ground truth. For the  $\lambda \mathrm{F}$  result, we first learned Laplacian eigenfunction base features as described in [25] from a uniform exploration policy and normalized them to the range [0, 1]. We parameterized the base feature network as a 2-layer MLP with ReLU activations and 16 units in the hidden layer. We then used the base features to learn the  $\lambda \mathrm{Fs}$  as in the tabular case, but with the  $\lambda \mathrm{F}$  network parameterized as a three-layer MLP with 16 units in each of the hidden layers and ReLU activations. All networks were optimized using Adam with a learning rate of 3e-4. The tabular and neural network experiments were repeated for three random seeds, the former was run for 1,500 episodes and the latter for 2,000.
E.2 Policy Learning
We ran the experiments for Fig. 5.2 in a version of the TwoRooms environment from the NeuroNav benchmark [22] with reward modified to decay with a specified
Algorithm 1: Online Tabular  $Q_{\lambda}$ -Learning Update
1: Require: Current  $\lambda$ R-values  $\Phi_{\lambda}^{(t)} \in \mathbb{R}^{|S| \times |\mathcal{A}| \times |\mathcal{S}|}$ , current reward vector  $\mathbf{r}^{(t)}$ , observed  $(s_t, a_t, s_{t+1})$  tuple
2: Compute  $Q_{\lambda}$ -values:  $Q_{\lambda}^{(t)} \gets (\Phi_{\lambda}^{(t)})^{\top} \mathbf{r}^{(t)}$
3: Select greedy action:  $a_{t + 1} \gets \operatorname{argmax}_{a \in \mathcal{A}} Q_{\lambda}^{(t)}(s_{t + 1}, a)$
4: Update  $\Phi_{\lambda}$ :
$$
\begin{array}{l} \Phi_ {\lambda} ^ {(t + 1)} \left(s _ {t}, a _ {t}\right) \leftarrow \Phi_ {\lambda} ^ {(t)} \left(s _ {t}, a _ {t}\right) + \alpha \delta^ {(t)}, \quad \text {w h e r e} \\ \delta^ {(t)} = \mathbf {1} (s _ {t}) \odot (1 + \gamma \lambda \Phi_ {\lambda} ^ {(t)} (s _ {t + 1}, a _ {t + 1})) + \gamma (1 - \mathbf {1} (s _ {t})) \odot \Phi_ {\lambda} ^ {(t)} (s _ {t + 1}, a _ {t + 1}) - \Phi_ {\lambda} ^ {(t)} (s _ {t}, a _ {t}). \\ \end{array}
$$
5: Return updated  $\Phi_{\lambda}^{(t + 1)}$
$\lambda_{true} = 0.5$  and discount factor  $\gamma = 0.95$ . The initial rewards in the top right goal and the lower room goal locations were 5 and the top left goal had initial reward 10. The observations in the neural network experiment were one-hot state indicators. The tabular  $Q_{\lambda}$  experiments run the algorithm in Algorithm 1 for 500 episodes for  $\lambda \in \{0.0, 0.5, 1.0\}$ , with  $\lambda_{true}$  set to 0.5, repeated for three random seeds. Experiments used a constant step size  $\alpha = 0.1$ . There were five possible actions: up, right, down, left, and stay. The recurrent A2C agents were based on the implementation from the BSuite library [36] and were run for 7,500 episodes of maximum length  $H = 100$  with  $\gamma = 0.99$  using the Adam optimizer with learning rate 3e-4. The experiment was repeated for three random seeds. The RNN was an LSTM with 128 hidden units and three output heads: one for the policy, one for the value function, and one for the  $\lambda$ F. The base features were one-hot representations of the current state, 121-dimensional in this case.
E.3 Tabular GPI
The agent is assumed to be given or have previously acquired four policies  $\{\pi_0, \pi_1, \pi_2, \pi_3\}$  individually optimized to reach rewards located in each of the four rooms of the environment. There are three reward locations  $\{g_0, g_1, g_2\}$  scattered across the rooms, each with its own initial reward  $\bar{r} = [5, 10, 5]$  and all with  $\lambda = 0.5$ . At the beginning of each episode, an initial state  $s_0$  is sampled uniformly from the set of available states. An episode terminates either when the maximum reward remaining in any of the goal states is less than 0.1 or when the maximum number of steps  $H = 40$  is reached. Empty states carry a reward of 0, encountering a wall gives a reward of  $-1$ , and the discount factor is set to  $\gamma = 0.97$ .
For each of the four policies, we learn  $\lambda$ Rs with  $\lambda$  equal to 0, 0.5, and 1.0 using standard dynamic programming (Bellman error curves plotted in Fig. F.3), and record the returns obtained while performing GPE+GPI with each of these representations over the course of 50 episodes. Bellman error curves for the  $\lambda$ Rs are In the left panel of Fig. 5.3, we can indeed see that using the correct  $\lambda$  (0.5) nets the highest returns. Example trajectories for each of  $\lambda$ R are shown in the remaining panels.
E.4 Pixel-Based GPI
In this case, the base policies  $\Pi$  were identical to those used in the tabular GPI experiments. First, we collected a dataset consisting of 340 observation trajectories  $(o_0, o_1, \ldots, o_{H-1}) \in \mathcal{O}^H$  with  $H = 19$  from each policy, totalling 6,460 observations. Raw observations were  $128 \times 128 \times 3$  and were converted to grayscale. The previous seven observations were stacked and used to train a Laplacian eigenfunction base feature network in the same way as [25]. For observations less than seven steps from the start of an episode, the remaining frames were filled in as all black observations (i.e., zeros). The network consisted of four convolutional layers with  $32 \times 3$  filters with strides  $(2, 2, 2, 1)$ , each followed by a ReLU nonlinearity. This was then flattened and passed through a Layer Norm layer [37] and a tanh nonlinearity before three fully connected layers, the first two with 64 units each and ReLU nonlinearities and the final, output layer with 50 units. The output was  $L_2$ -normalized as in [25]. This network  $\phi: \mathcal{O}^7 \mapsto \mathbb{R}^D$  (with  $D = 50$ ) was trained on the stacked observations for 10 epochs using the Adam optimizer and learning rate 1e-4 with batch size  $B = 64$ . To perform policy evaluation, the resulting features, evaluated on the dataset of stacked observations were collected into their own dataset of  $(s_t, a_{t+1}, s_{t+1}, a_{t+1})$  tuples, where  $s_t \triangleq o_{t-6:t}$ . The "states" were normalized to be between 0 and 1, and a vector  $\mathbf{w}$  was fit to the actual associated rewards via linear regression on the complete dataset. The  $\lambda$ F network was then trained using a form of neural fitted Q-iteration [FQI; 38] modified for policy evaluation with  $\lambda$ Fs (Algorithm 2). The architecture for the  $\lambda$ F network was identical to the base feature network, with the exception that the hidden size of the fully connected layers was 128 and the output dimension was  $D|\mathcal{A}| = 250$ . FQI was run for  $K = 20$  outer loop iterations, with each inner loop supervised learning setting run for  $L = 100$  epochs on the current dataset. Supervised learning was done using Adam with learning rate 3e-4 and batch size  $B = 64$ . Given the trained networks, GPI proceeded as in the tabular case, i.e.,
$$
a _ {t} = \underset {a \in \mathcal {A}} {\operatorname {a r g m a x}} \max  _ {\pi \in \Pi} \mathbf {w} ^ {\top} \varphi_ {\theta} ^ {\pi} \left(s _ {t}, a\right). \tag {E.2}
$$
50 episodes were run from random starting locations for  $H = 50$  steps and the returns measured. Learning curves for the base features and for  $\lambda \mathrm{F}$  fitting are shown in Fig. E.1. The  $\lambda \mathrm{F}$  curve measures the mean squared error as in Eq. (E.1).
The feature visualizations were created by performing PCA to reduce the average  $\lambda$ F representations for observations at each state in the environment to 2D. Each point in the scatter plot represents the reduced representation on the  $xy$  plane, and is colored according to the  $\lambda$ -conditioned value of the underlying state.
E.5 Continuous Control
$\lambda$  -SAC See Appendix I for details.
E.6 Learning the  $\lambda \mathbf{O}$  with FB
Training the  $\lambda \mathrm{O}$  with the FB parameterization proceeds in much the same way as in [12], but adjusted for a different norm and non-Markovian environment. We
Figure E.1: Learning curves for  $\lambda \mathbf{F}$  policy evaluation. Results are averaged over three runs, with shading indicating one unit of standard error.
Algorithm 2: Fitted  $Q_{\lambda}$ -Iteration
1: Require: Dataset of base features  $\{\phi (s)\in \mathbb{R}^D\}_{s\in S}$  decay rate  $\lambda$  , discount factor  $\gamma$  , reward feature vector  $\mathbf{w}\in \mathbb{R}^{D}$  , batch size  $B$  , learning rate  $\alpha$    
2: Initialize  $\lambda F\varphi_{\theta}$  parameters  $\theta^{(1)}$  (we drop the subscript  $\lambda$  and superscript  $\pi$  for concision)   
3: for  $k = 1\dots ,K$  do   
4: // Stage 1: Construct dataset   
5:  $\mathcal{D}\gets \emptyset$    
6: for  $(s,a)\in S\times A$  do   
7: for  $(s^{\prime},a^{\prime})\in S\times A$  do   
8:  $\mathcal{D}\leftarrow \mathcal{D}\cup \left\{\left((s,a),\underbrace{\mathbf{w}^{\mathrm{T}}[\phi(s)\odot(1 + \lambda\gamma\bar{\varphi}_{\theta(k)}(s^{\prime},a^{\prime})) + \gamma(1 - \phi(s))\odot\bar{\varphi}_{\theta(k)}(s^{\prime},a^{\prime})]}_{\triangleq y(s,a)}\right)\right\}$    
9: end for   
10: end for   
11: // Stage 2: Supervised learning   
12: Randomly initialize  $\theta_0$    
13: for  $\ell = 1,\dots ,L$  do   
14: Randomly shuffle  $\mathcal{D}$    
15: for  $\{(s,a),y)\}_{b = 1}^{B}\in \mathcal{D}$  do   
16:  $\theta_{\ell}\gets \theta_{\ell -1} - \alpha \nabla_{\theta}\frac{1}{2B}\sum_{b = 1}^{B}\left(y_{b} - \mathbf{w}^{\mathrm{T}}\varphi_{\theta_{\ell -1}}(s_{b},a_{b})\right)^{2}$    
17: end for   
18: end for   
19:  $\theta^{(k + 1)}\gets \theta_L$    
20: end for
summarize the learning procedure in Algorithm 3. The loss function  $\mathcal{L}$  is derived in Appendix H, with the addition of the following regularizer:
$$
\left\| \mathbb {E} _ {s \sim \rho} B _ {\omega} (s) B _ {\omega} (s) ^ {\top} - I \right\| ^ {2}.
$$
This regularizer encourages  $B$  to be approximately orthonormal, which promotes identifiability of  $F_{\theta}$  and  $B_{\omega}$  [12].
F Additional Results
See surrounding sections.
1: Require: Probability distribution  $\nu$  over  $\mathbb{R}^d$ , randomly initialized networks  $F_{\theta}, B_{\omega}$ , learning rate  $\eta$ , mini-batch size  $B$ , number of episodes  $E$ , number of epochs  $M$ , number of time steps per episode  $T$ , number of gradient steps  $N$ , regularization coefficient  $\beta$ , Polyak coefficient  $\alpha$ , initial diminishing constant  $\lambda$ , discount factor  $\gamma$ , exploratory policy greediness  $\epsilon$ , temperature  $\tau$
2: // Stage 1: Unsupervised learning phase
3:  $\mathcal{D}\gets \emptyset$
4: for epoch  $m = 1,\dots ,M$  do
5: for episode  $i = 1\ldots ,E$  do
6: Sample  $z\sim \nu$
7: Observe initial state  $s_1$
8: for  $t = 1,\dots ,T$  do
9: Select  $a_{t}\epsilon$ -greedy with respect to  $F_{\theta}(s_t,a,z)^{\top}z$
10: Observe reward  $r_t(s_t)$  and next state  $s_{t+1}$ .
1:  $\mathcal{D}\gets \mathcal{D}\cup \{(s_t,a_t,r_t(s_t),s_{t + 1})\}$
2: end for
3: end for
4: for  $n = 1,\dots ,N$  do
5: Sample a minibatch  $\{(s_j, a_j, r_j(s_j), s_{j+1})\}_{j \in J} \subset \mathcal{D}$  of size  $|J| = B$
6: Sample a minibatch  $\{\tilde{s}_j\}_{j\in J}\subset \mathcal{D}$  of size  $|J| = B$
7: Sample a minibatch  $\{s_j^{\prime}\}_{j\in J}\stackrel {\mathrm{iid}}{\sim}\mu$  of size  $|J| = B$
8: Sample a minibatch  $\{z_j\}_{j\in J}\stackrel {\mathrm{iid}}{\sim}\nu$  of size  $|J| = B$
9: For every  $j\in J$  , set  $\pi_{z_j}(\cdot |s_{j + 1}) = \mathrm{softmax}\left(F_{\theta -}(s_{j + 1},\cdot ,z_j)^\top z_j / \tau\right)$
20:
$$
\begin{array}{l} \mathcal {L} (\theta , \omega) \leftarrow \frac {1}{2 B ^ {2}} \sum_ {j, k \in J ^ {2}} \left(F _ {\theta} (s _ {j}, a _ {j}, z _ {j}) ^ {\top} B _ {\omega} \left(s _ {k} ^ {\prime}\right) - \gamma \sum_ {a \in \mathcal {A}} \pi_ {z _ {j}} (a | s _ {j + 1}) F _ {\theta^ {-}} \left(s _ {j + 1}, a, z _ {j}\right) ^ {\top} B _ {\omega^ {-}} \left(s _ {k} ^ {\prime}\right)\right) ^ {2} \\ - \frac {1}{B} \sum_ {j \in J} F _ {\theta} \left(s _ {j}, a _ {j}, z _ {j}\right) ^ {\top} B _ {\omega} \left(s _ {j}\right) \\ + \frac {\gamma (1 - \lambda)}{B} \sum_ {j \in J} \mu \left(s _ {j}\right) F _ {\theta} \left(s _ {j}, a _ {j}, z _ {j}\right) ^ {\top} B _ {\omega} \left(s _ {j}\right) \sum_ {a \in \mathcal {A}} \pi_ {z _ {j}} \left(a \mid s _ {j + 1}\right) F _ {\theta^ {-}} \left(s _ {j + 1}, a, z _ {j}\right) ^ {\top} B _ {\omega^ {-}} \left(s _ {j}\right) \\ + \beta \left(\frac {1}{B ^ {2}} \sum_ {j, k \in J ^ {2}} B _ {\omega} (s _ {j}) ^ {\top} \bar {B} _ {\omega} (\tilde {s} _ {k}) \bar {B} _ {\omega} (s _ {j}) ^ {\top} \bar {B} _ {\omega} (\tilde {s} _ {k}) - \frac {1}{B} \sum_ {j \in J} B _ {w} (s _ {j}) ^ {\top} \bar {B} _ {\omega} (s _ {j})\right) \\ \end{array}
$$
21: Update  $\theta$  and  $\omega$  via one step of Adam on  $\mathcal{L}$
22: Sample a minibatch  $\{(s_j, r_j(s_j), s_{j+1}, r_{j+1}(s_{j+1}))\}_{j \in J}$  of size  $|J| = B$  from  $\mathcal{D}$
23:  $\mathcal{L}_{\lambda}(\lambda)\gets \frac{1}{2B}\sum_{j\in J}\mathbb{1}(s_{j + 1} = s_j)(r_{j + 1}(s_{j + 1}) - \lambda r_j(s_j))^2$
24: Update  $\lambda$  via one step of Adam on  $\mathcal{L}_{\lambda}$
25: end for
26:  $\theta^{-} \gets \alpha \theta^{-} + (1 - \alpha) \theta$
27:  $\omega^{-} \gets \alpha \omega^{-} + (1 - \alpha) \omega$
28: end for
29: // Stage 2: Exploitation phase for a single episode with initial reward  $r_0(s)$
30:  $z_{R}\gets \sum_{s\in \mathcal{S}}\mu (s)r_{0}(s)B_{\omega}(s)$
31: Observe initial state  $s_1$
32: for  $t = 1,\dots ,T$  do
33:  $a_{t}\gets \operatorname{argmax}_{a\in \mathcal{A}}F(s_{t},a,z_{R})^{\top}z_{R}$
34: Observe reward  $r_t(s)$  and next state  $s_{t+1}$ .
35:  $z_{R}\gets \sum_{s\in \mathcal{S}}\mu (s)r_{t}(s)B_{\omega}(s)$
36: end for
Algorithm 3:  $\lambda$  O FB Learning
Figure F.1: A simple grid and several policies.
Figure F.3: Convergence of dynamic programming on FourRooms with and without stochastic transitions.
Figure F.4: GPI with noisy transitions in FourRooms. To verify that performance was maintained even with stochastic transitions, we added a  $20\%$  probability that a given action would result in a random transition to neighboring state. Results are consistent with Fig. 5.3, indicating the having the correct value of  $\lambda$  produces better performance.
G Advantage of the Correct  $\lambda$
Importantly, for GPE using the  $\lambda \mathrm{R}$  to work in this setting, the agent must either learn or be provided with the updated reward vector  $\mathbf{r}_{\lambda}$  after each step/encounter with a rewarded state. This is because the  $\lambda \mathrm{R}$  is forward-looking in that it measures the (diminished) expected occupancies of states in the future without an explicit mechanism for remembering previous visits. For simplicity in this case, we provide this vector to the agent at each step—though if we view such a multitask agent as simply as a module carrying out the directives of a higher-level module or policy within a hierarchical framework as in, e.g., Feudal RL [39], the explicit provision of reward information is not unrealistic. Regardless, a natural question in this case is whether there is actually any value in using the  $\lambda \mathrm{R}$  with the correct value of  $\lambda$  in this
$\Phi_{1,0}^{\pi_2}$
$\Phi_{1,0}^{\pi_3}$
$\Phi_{1,0}^{\pi_1}$
$\Phi_{0.8}^{\pi_1}$
$\Phi_{0.8}^{\pi_2}$
$\Phi_{0.4}^{\pi_1}$
$\Phi_{0.4}^{\pi_2}$
$\Phi_{0.4}^{\pi_3}$
$\Phi_{0,0}^{\pi_1}$
$\Phi_{0,0}^{\pi_2}$
Φ0.0
Fn1
Figure F.2: Visualizing the SR, the  $\lambda$ R and the FR. We can see that the  $\Phi_1^\pi$  is equivalent to the SR and  $\Phi_0^\pi$  is equivalent to the FR, with intermediate values of  $\lambda$  providing a smooth transition between the two.
Fn2
Fn3
setting: If the agent is provided with the correct reward vector, then wouldn't policy evaluation work with any  $\lambda \mathbb{R}$ ?
To see that this is not the case, consider the three-state toy MDP shown in Figure Fig. G.1, where  $\bar{r}(s_1) = 10$ ,  $\bar{r}(s_2) = 6$ ,  $\bar{r}(s_0) = 0$ ,  $\lambda(s_1) = 0$ ,  $\lambda(s_2) = 1.0$ , and  $\gamma = 0.99$ . At time  $t = 0$ , the agent starts in  $s_0$ . Performing policy evaluation with  $\lambda(s_1) = \lambda(s_2) = 1$  (i.e., with the SR) would lead the agent to go left to  $s_1$ . However,
Figure G.1: A 3-state toy environment.
the reward would then disappear, and policy evaluation on the second step would lead it to then move right to  $s_0$  and then  $s_2$ , where it would stay for the remainder of the episode. In contrast, performing PI with the correct values of  $\lambda$  would lead the agent to go right to  $s_2$  and stay there. In the first two timesteps, the first policy nets a total reward of  $10 + 0 = 10$ , while the second policy nets  $6 + 5.94 = 11.94$ . (The remaining decisions are identical between the two policies.) This is a clear example of the benefit of having the correct  $\lambda$ , as incorrect value estimation leads to suboptimal decisions even when the correct reward vector/function is provided at each step.
H The  $\lambda$  Operator
To learn the  $\lambda$ O, we would like to define  $\Phi_{\lambda}^{\pi}(s_t, ds') \triangleq \varphi_{\lambda}^{\pi}(s_t, s') \mu(ds')$  for some base policy  $\mu$ . However, this would lead to a contradiction:
$$
\Phi_ {\lambda} ^ {\pi} (s, A \cup B) = \int_ {A} \varphi_ {\lambda} ^ {\pi} (s, d s ^ {\prime}) \mu (d s ^ {\prime}) + \int_ {B} \varphi_ {\lambda} ^ {\pi} (s, d s ^ {\prime}) \mu (d s ^ {\prime}) = \Phi_ {\lambda} ^ {\pi} (s, A) + \Phi_ {\lambda} ^ {\pi} (s, B)
$$
for all disjoint  $A, B$ , contradicting Lemma B.4.
For now, we describe how to learn the  $\lambda \mathrm{O}$  for discrete  $S$ , in which case we have  $\Phi_{\lambda}^{\pi}(s,s') = \varphi_{\lambda}^{\pi}(s,s')\mu(s')$ , i.e., by learning  $\varphi$  we learn a weighted version of  $\Phi$ . We define the following norm, inspired by Touati et al. [25]:
$$
\| \Phi_ {\lambda} ^ {\pi} \| _ {\rho} ^ {2} \triangleq \mathbb {E}  \left[ \left(\frac {\Phi_ {\lambda} ^ {\pi} (s , s ^ {\prime})}{\mu (s ^ {\prime})}\right) ^ {2} \right],
$$
where  $\mu$  is any density on  $S$ . In the case of finite  $S$ , we let  $\mu$  be the uniform density. We then minimize the Bellman error for  $\Phi_{\lambda}^{\pi}$  with respect to  $\| \cdot \|_{\rho, \mu}^{2}$  (dropping the sub/superscripts on  $\Phi$  and  $\varphi$  for clarity):
$$
\begin{array}{l} \mathcal {L} (\Phi) = \left\| \varphi \mu - (I \odot (\mathbf {1 1} ^ {\top} + \lambda \gamma P ^ {\pi} \varphi \mu) + \gamma (\mathbf {1 1} ^ {\top} + I) \odot P ^ {\pi} \varphi \mu) \right\| _ {\rho , \mu} ^ {2} \\ = \mathbb {E} _ {s _ {t} \sim \rho , s ^ {\prime} \sim \mu} \Big [ \big (\varphi (s _ {t}, s ^ {\prime}) - \frac {\mathbb {1} (s _ {t} = s ^ {\prime})}{\mu (s ^ {\prime})} \\ \left. \left. + \gamma (1 - \lambda) \frac {\mathbb {1} \left(s _ {t} = s ^ {\prime}\right)}{\mu \left(s ^ {\prime}\right)} \mathbb {E} _ {s _ {t + 1} \sim p ^ {\pi} (\cdot | s _ {t})} \bar {\Phi} \left(s _ {t + 1}, s ^ {\prime}\right) - \gamma \mathbb {E} _ {s _ {t + 1} \sim p ^ {\pi} (\cdot | s _ {t})} \bar {\varphi} \left(s _ {t + 1}, s ^ {\prime}\right)\right) ^ {2} \right] \\ \stackrel {\pm c} {=} \mathbb {E} _ {s _ {t}, s _ {t + 1} \sim \rho , s ^ {\prime} \sim \mu} \left[ \left(\varphi (s _ {t}, s ^ {\prime}) - \gamma \bar {\varphi} (s _ {t + 1}, s ^ {\prime})\right) ^ {2} \right] \\ - 2 \mathbb {E} _ {s _ {t}, s _ {t + 1} \sim \rho} \left[ \sum_ {s ^ {\prime}} \mu \left(s ^ {\prime}\right) \varphi \left(s _ {t}, s ^ {\prime}\right) \frac {\mathbb {1} \left(s _ {t} = s ^ {\prime}\right)}{\mu \left(s ^ {\prime}\right)} \right] \\ + 2 \gamma (1 - \lambda) \mathbb {E} _ {s _ {t}, s _ {t + 1} \sim \rho} \left[ \sum_ {s ^ {\prime}} \mu \left(s ^ {\prime}\right) \varphi \left(s _ {t}, s ^ {\prime}\right) \bar {\varphi} \left(s _ {t + 1}, s ^ {\prime}\right) \mu \left(s ^ {\prime}\right) \frac {\mathbb {1} \left(s _ {t} = s ^ {\prime}\right)}{\mu \left(s ^ {\prime}\right)} \right] \\ \stackrel {\pm c} {=} \mathbb {E} _ {s _ {t}, s _ {t + 1} \sim \rho , s ^ {\prime} \sim \mu} \left[ \left(\varphi \left(s _ {t}, s ^ {\prime}\right) - \gamma \bar {\varphi} \left(s _ {t + 1}, s ^ {\prime}\right)\right) ^ {2} \right] - 2 \mathbb {E} _ {s _ {t} \sim \rho} [ \varphi (s _ {t}, s _ {t}) ] \\ + 2 \gamma (1 - \lambda) \mathbb {E} _ {s _ {t}, s _ {t + 1} \sim \rho} [ \mu (s _ {t}) \varphi (s _ {t}, s _ {t}) \bar {\varphi} (s _ {t + 1}, s _ {t}) ], \\ \end{array}
$$
Note that we recover the SM loss when  $\lambda = 1$ . Also, an interesting interpretation is that when the agent can never return to its previous state (i.e.,  $\varphi(s_{t+1}, s_t) = 0$ ), then we also recover the SM loss, regardless of  $\lambda$ . In this way, the above loss appears to "correct" for repeated state visits so that the measure only reflects the first visit.