SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Remark 2. Notice that as with any optimization problem, issues of realizability and identifiability of the desired  $r_{\phi}$  must be taken into account. The examples provided in this section aim to highlight the capability of optimized behavior alignment reward functions. In particular, they not only improve and accelerate the learning process but are also capable of inducing updates capable of 'fixing' imperfections in the underlying RL algorithm.
B Extended Related Works
The bi-level objective draws inspiration from the seminal work of Singh et al. [52, 53] that provides an optimal-rewards framework for an agent. Prior works have built upon it to explore search techniques using evolutionary algorithms [43, 21], develop extensions for multi-agent setting [66, 22], and mitigate sub-optimality due to use of inaccurate models [55-57]. Our work also builds upon this direction and focuses on various aspects of leveraging auxiliary rewards  $r_{\mathrm{aux}}$ , while staying robust against its misspecification.
Apart from specifying auxiliary rewards  $r_{\mathrm{aux}}$  directly, other techniques for reward specification include linear temporal logic [35, 11, 65, 39] or reward machines [29-31] that allow exposing the reward functions as a white-box to the agent.
Recent works also explore  $\gamma$  that is state-action dependent [63, 48], or establishes connection between  $\gamma$  and value function regularization in TD learning Amit et al. [3]. These ideas are complementary to our proposed work and combining these with BARFI remains interesting directions for the future.
The concept of path-based meta-learning was initially popularized for few-shot task learning in supervised learning [15, 42]. Similar path-based approaches have been adopted in reinforcement
learning (RL) in various forms [28, 62, 69, 71]. Initially designed for stochastic gradient descent, these methods have been extended to other optimizers such as Adam [32] and RMSprop [25], treating them as differentiable counterparts [20].
C Algorithm
In this section we discuss the algorithm for the proposed method. As the proposed method does behavior alignment reward function's implicit optimization, we name it BARFI. Pseudo-code for BARFI is presented in Algorithm 5. We will first build on some preliminaries to understand the concepts.
C.1 Vector Jacobian Product
Let us assume that,  $x \in \mathbb{R}^d, y \in \mathbb{R}^m, f(x,y) \in \mathbb{R}$ . Then, we know that  $\partial f(x,y) / \partial x \in \mathbb{R}^d, \partial f(x,y) / \partial y \in \mathbb{R}^m, \partial^2 f(x,y) / \partial y\partial x \in \mathbb{R}^{d\times m}$ . Let us also assume that we have a vector  $v \in \mathbb{R}^d$ , and if we need to calculate the following, we can pull the derivative outside as shown:
$$
\underbrace{\underbrace{v}_{\mathbb{R}^{d}}\underbrace{\frac{\partial^{2}f(x,y)}{\partial y\partial x}}_{\mathbb{R}^{d\times m}}} _ {\mathbb{R}^{m}} = \underbrace{\frac{\partial}{\partial y}\left\langle\underbrace{v}_{\mathbb{R}^{d}},\underbrace{\frac{\partial f(x,y)}{\partial x}}_{\mathbb{R}^{d}}\right\rangle}_{\mathbb{R}^{1}}.
$$
As we can see, the vector Jacobian product can be broken down into differentiating a vector product but shifting the place of multiplication, in which case we assume that the gradient passes through  $v$  w.r.t.  $y$  and hence we don't ever have to deal with large multiplications. Also note that the outer partial w.r.t. can easily be handled by autodiff packages. A pseudo-code is show in Algorithm 1.
Algorithm 1: Jacobian Vector Product
1 Input:  $f(x,y)\in \mathbb{R}^1,x\in \mathbb{R}^d,y\in \mathbb{R}^m,v\in \mathbb{R}^d$    
2  $f^{\prime}\gets \mathrm{grad}(f(x,y),x)$    
3 jvp  $\leftarrow$  grad(f',y,grad_outputs=  $v$  ）   
4 Return: jvp
C.2 Neumann Series Approximation for Hessian Inverse
Recall, that for a given real number  $\beta \in \mathbb{R}$ , such that  $0 \leq \beta < 1$ , we know that the geometric series of this has a closed form solution, i.e.,
$$
\begin{array}{l} s = 1 + \beta^ {1} + \beta^ {2} + \beta^ {3} + \dots + \\ = \frac {1}{1 - \beta}. \\ \end{array}
$$
Similarly, given we have a value  $\alpha$  such that  $\beta = 1 - \alpha$ , we can write  $\alpha^{-1}$  as follows:
$$
\begin{array}{l} \frac {1}{1 - \beta} = 1 + \beta + \beta^ {2} + \beta^ {3} + \dots + \\ \frac {1}{1 - (1 - \alpha)} = 1 + (1 - \alpha) + (1 - \alpha) ^ {2} + (1 - \alpha) ^ {3} + \dots + \\ \alpha^ {- 1} = 1 + (1 - \alpha) + (1 - \alpha) ^ {2} + (1 - \alpha) ^ {3} + \dots + \\ \alpha^ {- 1} = \sum_ {i = 0} ^ {\infty} (1 - \alpha) ^ {i}. \\ \end{array}
$$
The same can be generalized for a matrix, i.e., given a matrix  $\mathbf{A} \in \mathbb{R}^{d \times d}$ , we can write  $\mathbf{A}^{-1}$  as follows:
$$
\mathbf {A} ^ {- 1} = \sum_ {i = 0} ^ {\infty} (\mathbf {I} - \mathbf {A}) ^ {i}.
$$
Note for the above to hold, matrix  $\mathbf{A}$ , where we represent  $\mathrm{eig}(\mathbf{A})$  as the eigenvalues of matrix  $\mathbf{A}$ , we should have the following condition to hold,  $0 < \mathrm{eig}(\mathbf{A}) < 1$ . Note here we would regularize  $\mathbf{A}$  to ensure that all eigenvalues are positive, and then we can always scale the matrix  $\mathbf{A}$ , by its biggest eigenvalue to ensure that the above condition holds. Let say  $\eta = 1 / \max \mathrm{eig}(\mathbf{A})$ . Then we can write the following:
$$
\begin{array}{l} \mathbf {A} ^ {- 1} = \frac {\eta}{\eta} \mathbf {A} ^ {- 1} \\ = \eta (\eta \mathbf {A}) ^ {- 1} \\ = \eta \sum_ {i = 0} ^ {\infty} (\mathbf {I} - \eta \mathbf {A}) ^ {i}. \\ \end{array}
$$
As  $\eta \mathbf{A}$  would always satisfy the above condition.
C.3 Neumann Approximation for Hessian Vector Product
Given we have seen how we can approximate the Inverse of a matrix without relying  $O(d^3)$  operations, through Neumann approximation, let's look what needs to be done for our updates. Recall that the update  $\phi, \varphi$  (8) and (15) were,
$$
\frac {\partial J (\theta (\phi , \varphi))}{\partial \phi} = - \underbrace {\frac {\partial J (\theta (\phi , \varphi))}{\partial \theta (\phi , \varphi)}} _ {v} \left(\underbrace {\frac {\partial \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \theta (\phi , \varphi)}} _ {\mathbf {H}}\right) ^ {- 1} \underbrace {\frac {\partial \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \phi}} _ {\mathbf {A}}
$$
$$
\frac {\partial \big (J (\theta (\phi , \varphi)) - \frac {1}{2} \| \gamma_ {\varphi} \| ^ {2} \big)}{\partial \varphi} = - \underbrace {\frac {\partial J (\theta (\phi , \varphi))}{\partial \theta (\phi , \varphi)}} _ {v} \left(\underbrace {\frac {\partial \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \theta (\phi , \varphi)}} _ {\mathbf {H}}\right) ^ {- 1} \underbrace {\frac {\partial \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \varphi}} _ {\mathbf {B}} - \frac {\partial \gamma_ {\varphi}}{\partial \varphi}.
$$
Let us look closely at the update for  $\phi$  and we can generalize the updates easily for the case of  $\varphi$ .
$$
\frac {\partial J (\theta (\phi , \varphi))}{\partial \phi} = - v \mathbf {H} ^ {- 1} \mathbf {A}
$$
We first look at how we can approximate the value of  $v\mathbf{H}^{-1}$  efficiently, as we can always make use of the Jacobian Vector product later to get  $(v\mathbf{H}^{-1})\mathbf{A}$ , as  $v\mathbf{H}^{-1}$  becomes a vector. Let us assume we wish to run the Neumann approximation up to  $n$  steps, i.e., we want to approximate  $\mathbf{H}^{-1}$  up to  $n$  order Neumann expansion,
$$
\eta \left(\eta \mathbf {H} ^ {- 1}\right) \approx \eta \sum_ {i = 0} ^ {n} (I - \eta \mathbf {H}) ^ {i} \tag {12}
$$
Here we are assuming that the outer optimization for update (1) is for the function  $J(\theta(\phi, \varphi))$  and the inner optimization which is represented by the update (3) is  $f(\theta(\phi, \varphi), \phi, \varphi)$ , i.e.,
$$
\Delta (\theta , r _ {p}) = \frac {\partial J (\theta (\phi , \varphi))}{\partial \theta}, \quad \Delta (\theta , \phi , \varphi) = \frac {\partial f (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \theta}.
$$
The most common form in which  $f(;;B)$  is usually defined is the following:
$$
f (\theta , \phi , \varphi ; B) := \frac {1}{| B |} \sum_ {\tau \in B} \left[ \sum_ {t = 0} ^ {T} \log (\pi_ {\theta} (S _ {t} ^ {\tau}, A _ {t} ^ {\tau})) \sum_ {j = t} ^ {T} \gamma_ {\varphi} ^ {j - t} r _ {\phi} (S _ {j} ^ {\tau}, A _ {j} ^ {\tau}) \right].
$$
Algorithm 2: Vector Hessian Inverse Product for (8) i.e.,  $v\mathbf{H}^{-1}$
1 Input:  $\theta, \phi, \varphi, J, f, n, \eta, \mathcal{D}_{\mathrm{off}}, \mathcal{D}_{\mathrm{on}}$   
2  $v \gets \operatorname{grad}(J(\theta; \mathcal{D}_{\mathrm{on}}), \theta)$   
3  $v' \gets \eta \times \operatorname{grad}(f(\theta, \phi, \varphi; \mathcal{D}_{\mathrm{off}}), \theta)$   
4 Let:  $v_0 \gets v, p_0 \gets v$   
5 for  $i \in [0, n)$  do  
6  $\begin{bmatrix} v_{i+1} & v_i - \operatorname{grad}(v', \theta, \operatorname{grad\_outputs} = v_i) \\ p_{i+1} & p_i + v_{i+1} \end{bmatrix}$   
7  
8 Return:  $\eta p_n$ ; // Approximation of  $vH^{-1}$  as in (12)
Similarly this can be defined for  $J$ , except making use of  $r_p$  and problem defined  $\gamma$ :
$$
J (\theta ; B) := \frac {1}{| B |} \sum_ {\tau \in B} \left[ \sum_ {t = 0} ^ {T} \log (\pi_ {\theta} (S _ {t} ^ {\tau}, A _ {t} ^ {\tau})) \sum_ {j = t} ^ {T} \gamma^ {j - t} r _ {p} (S _ {j} ^ {\tau}, A _ {j} ^ {\tau}) \right].
$$
Finally, once we have  $v\mathbf{H}^{-1}$ , we can use the Vector Jacobian Product to calculate  $(v\mathbf{H}^{-1})\mathbf{A}$  as described in Algorithm 3:
Algorithm 3: Update for  $\phi$ , i.e. (8) i.e.,  $v\mathbf{H}^{-1}\mathbf{A}$
1 Input:  $\theta, \phi, \varphi, J, f, n, \eta, \mathcal{D}_{\mathrm{off}}, \mathcal{D}_{\mathrm{on}}$   
2  $v \gets$  Algorithm 2 ( $\theta, \phi, \varphi, J, f, n, \eta, \mathcal{D}_{\mathrm{off}}, \mathcal{D}_{\mathrm{on}}$ )  
3  $v' \gets \operatorname{grad}(f(\theta, \phi, \varphi; \mathcal{D}_{\mathrm{off}}), \theta)$   
4  $\Delta_{\phi} \gets \operatorname{grad}(v', \phi, \operatorname{grad\_outputs} = v)$   
5 Return  $\Delta_{\phi}$
We can similarly derive updates for  $\varphi$ . Note we are not including the different forms of regularizers
Algorithm 4: Update for  $\phi$ , i.e. (15) i.e.,  $v\mathbf{H}^{-1}\mathbf{B}$
1 Input:  $\theta, \phi, \varphi, J, f, n, \eta, \mathcal{D}_{\mathrm{off}}, \mathcal{D}_{\mathrm{on}}$   
2  $v \gets$  Algorithm 2 ( $\theta, \phi, \varphi, J, f, n, \eta, \mathcal{D}_{\mathrm{off}}, \mathcal{D}_{\mathrm{on}}$ )  
3  $v' \gets \operatorname{grad}(f(\theta, \phi, \varphi), \mathcal{D}_{\mathrm{off}}), \theta)$   
4  $\Delta_{\varphi} \gets \operatorname{grad}(v', \varphi, \operatorname{grad\_outputs} = v)$   
5 Return  $\Delta_{\varphi}$
over here to reduce clutter, but adding them is simple.
C.4 Pseudo Code (Algorithm 5)
Lines 8-10 and 21-23 of Algorithm 5 represent the inner optimization process, and the outer optimization process if from lines 16-17. Lines 8-10 is the initial step of updates to converge to the current values of  $\phi, \varphi$ , and from there onwards after each update of outer optimization, we consequently update the policy in (21-23). The flow of the algorithm is shown in Figure 4.
As discussed in Section D, using regularizers in  $\Delta(\theta, \phi, \varphi)$  smoothens the objective  $J(\theta(\phi, \varphi))$  with respect to  $\phi$  and  $\varphi$ . This is helpful as gradual changes in  $r_{\phi}$  and  $\gamma_{\varphi}$  can result in gradually changes in the fixed point for the inner optimization. Therefore, for computational efficiency, we initialize the policy parameters from the fixed-point of the previous inner-optimization procedure such that the inner-optimization process may start close to the new fixed-point.
In lines 8-10, the inner optimization for the policy parameters  $\theta$  are performed till (approximate) convergence. Note that only trajectories from past interactions are used and no new-trajectories are sampled for the inner optimization.
Algorithm 5: BARFI: Behavior Alignment Reward Function's Implicit optimization
1 Input: \(J,f,\alpha_{\theta},\alpha_{\phi},\alpha_{\varphi},\eta ,n,\delta\) optim, \(E,N_{i},N_{0}\)   
2 Initialize: \(\pi_{\theta},r_{\phi},\gamma_{\varphi}\)   
3 Initialize: optim \(\theta \gets\) optim \((\alpha_{\theta})\) optim \(\phi \gets\) optim \((\alpha_{\phi})\) optim \(\varphi \gets\) optim \((\alpha_{\varphi})\)   
4 \(\mathcal{D}_{\mathrm{off}}\gets []\) # Collect a batch of data for warmup period   
5 for \(e\in [1,N_0)\) do Generate \(\tau_{e}\) using \(\pi_{\theta}\) Append \(\tau_{e}\) to \(\mathcal{D}_{\mathrm{off}}\) # Initial training steps using warmup data   
8 for \(i\in [0,N_i + N_0)\) do Sample a batch of trajectories \(B\) from \(\mathcal{D}_{\mathrm{off}}\) # Update policy \(\theta \leftarrow \theta +\mathrm{optim}_{\theta}(\mathrm{grad}(f(\theta ,\phi ,\varphi ;B),\theta))\) # Start reward alignment   
11 for \(e\in [N_0,E)\) do # Collect a batch of on-policy data \(\mathcal{D}_{\mathrm{on}}\gets []\) for \(j\in [0,\delta)\) do Generate trajectory \(\tau_{e + j}\) using \(\pi_{\theta}\) and append in \(\mathcal{D}_{\mathrm{on}}\) \(e\gets e + \delta\) # Update \(r_\phi\) and \(\gamma_{\varphi}\) \(\Delta_{\phi}\gets\) Algorithm 3(0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0.1\) \(\Delta_{\varphi}\gets\) Algorithm 4(0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, 1 \(\phi \gets \phi +\mathrm{optim}_{\phi}(\Delta_{\phi})\) \(\varphi \gets \varphi +\mathrm{optim}_{\varphi}(\Delta_{\varphi})\) \(\mathcal{D}_{\mathrm{off}}\gets \mathcal{D}_{\mathrm{off}} + \mathcal{D}_{\mathrm{on}}\) # Learn policy for new reward function, initializing from the last   
for \(i\in [0,N_i)\) do Sample a batch of trajectories \(B\) from \(\mathcal{D}_{\mathrm{off}}\) # Update policy \(\theta \leftarrow \theta +\mathrm{optim}_{\theta}(\mathrm{grad}(f(\theta ,\phi ,\varphi ;B),\theta))\)
Figure 4: Algorithm Flow: The change in different parameters
In Lines 13-14, a new batch  $\mathcal{D}_{\mathrm{on}}$  of data is sampled using the policy returned by the inner-optimization process. This data is used to compute  $\partial J(\theta (\phi ,\varphi)) / \partial \theta (\phi ,\varphi)$ . Existing data  $\mathcal{D}_{\mathrm{off}}$  that was used in the inner-optimization process is then used to compute  $\partial \theta (\phi ,\varphi) / \partial \phi$  and  $\partial \theta (\phi ,\varphi) / \partial \varphi$ . Using these in (8) and (15), the parameters for  $r_{\phi}$  and  $\gamma_{\varphi}$  are updated in Lines 16 and 17, respectively.
Finally, the new data  $\mathcal{D}_{\mathrm{on}}$  is merged into the existing data  $\mathcal{D}_{\mathrm{off}}$  and the entire process continues.
C.5 Note on Approximation
An important limitation of the methods discussed above is that  $\theta(\phi, \varphi)$  is considered such that  $\Delta(\theta(\phi, \varphi), \phi, \varphi) = 0$ , i.e., the  $\mathsf{Alg}$  is run to convergence. In practice, we only execute  $\mathsf{Alg}$  for a predetermined number of update steps that need not result in convergence to an optimum exactly. However, the impact of this approximation can be bounded by assuming convergence to an  $\epsilon$ -neighborhood of the optima [49]. Furthermore, due to smoothness in the functional space, slight
changes to  $\phi$  and  $\varphi$  should result in slight shifts in the optimum  $\theta(\phi, \varphi)$ . The continuity property allows for improvements in the optimization process: it suffices to initialize the parameters of each inner-loop optimization problem with the final parameters of the approximate fixed point solution,  $(\phi, \varphi)$ , identified in the previous iteration of the inner loop. The complete resulting algorithm is presented in the appendix as Algorithm 5.
C.6 Path-wise Bi-level Optimization
An alternative approach for computing the term (b) in (4) is possible. The formulation of BARFI described above, based on implicit bi-level optimization, is agnostic to the optimization path taken by  $\mathrm{Alg}$ . For the sake of completeness, let us also consider a version of BARFI that does take into account the path followed by the inner optimization loop. This is advantageous because it allows us to eliminate the need for the convergence criteria (5). We call this variant BARFI unrolled. The main difference, in this case, is that when computing the term (b) in (4), we now consider each inner update step until the point  $\theta (\phi ,\varphi)$  is reached—where the sequence of steps depends on the specific  $\mathrm{Alg}$  used for the inner updates. Details are deferred to Appendix C. Notice that this approach results in a path-wise optimization process that can be more demanding in terms of computation and memory. We further discuss this issue, and demonstrate the efficacy of this alternative approach, in the empirical analyses section (Section ??).
D Smoothing the objective
To understand why  $J(\theta(\phi, \varphi))$  might be ill-conditioned is to note that, often a small perturbation in the reward function doesn't necessarily lead to a change in the corresponding optimal policy. This can lead lack of gradient directions in the neighborhood of  $\phi, \varphi$  for gradient methods to be effective. This issue can be addressed by employing common regularization techniques like L2 regularization of the policy parameters or entropy regularization for the policy<sup>6</sup>. We discuss two ways to regularize the objective in the upcoming sections.
D.1 L2 Regularization
To understand how severely ill-conditioned  $J(\theta(\phi, \varphi))$  can be, notice that a small perturbation in the reward function often does not change the corresponding optimal policies or the outcome of a policy optimization algorithm A1g. Therefore, if the parameters of the behavior alignment reward are perturbed from  $\phi$  to  $\phi'$ , it may often be that  $J(\theta(\phi, \varphi)) = J(\theta(\phi', \varphi))$  and this limits any gradient based optimization for  $\phi$  as  $\partial J(\theta(\phi, \varphi)) / \partial \phi$  is 0. Similarly, minor perturbations in  $\varphi$  may result in no change in  $J(\theta(\phi, \varphi))$  either.
Fortunately, there exists a remarkably simple solution: incorporate regularization for the policy parameters  $\theta$  in objective for  $\mathrm{Alg}$  in the inner-level optimization. For example, the optimal policy for the following regularized objective  $\mathbb{E}_{\pi_{\theta}}[\sum_{t = 0}^{T}\gamma_{\varphi}^{t}r_{\phi}(S_{t},A_{t})] - \frac{\lambda}{2}\| \theta \|^{2}$  varies smoothly to trade-off between the regularization value of  $\theta$  and the magnitude of the performance characterized by  $(r_{\phi},\gamma_{\varphi})$  which changes with the values of  $r_{\phi}$  and  $\gamma_{\varphi}$ . See Figure 5 for an example with L2 regularization.
D.2 Entropy Regularized
In Section D.1, smoothing of  $J(\theta (\phi ,\varphi))$  was done by using L2 regularization on the policy parameters  $\theta$  in the inner-optimization process. However, alternate regularization methods can also be used. For example, in the following we present an alternate update rule for  $\theta$  based on entropy regularization,
$$
\Delta (\theta , \phi , \varphi) := \mathbb {E} _ {\mathcal {D}} \left[ \sum_ {t = 0} ^ {T} \psi_ {\theta} (S _ {t}, A _ {t}) \sum_ {j = t} ^ {T} \gamma_ {\varphi} ^ {j - t} \left(r _ {\phi} (S _ {j}, A _ {j}) - \lambda \ln \pi_ {\theta} (S _ {j}, A _ {j})\right) \right].
$$
This regularization is performed so as to avoid a noninvertible Hessian as we had discussed in Section 5
Figure 5: (Left) A bandit problem, where the data is collected from a policy  $\beta$  that samples action  $A$  mostly. (Middle) Each point on the 3D surface corresponds to the performance of  $\theta(\phi, 1)$  returned by an  $\mathbb{A}\mathbb{1}\mathbb{g}$  that uses the update rule  $\Delta_{\mathrm{off}}(\theta, \phi, 1)$  corresponding to the value of  $r_{\phi}$  for actions  $A$  and  $C$  in the bottom axes;  $r_{\theta}$  for action  $B$  is set to 0 to avoid another variable in a 3D plot. Notice that small perturbation in  $r_{\phi}$  may lead to no or sudden changes in  $J(\theta(\phi, 1))$ . (Right) Performance of  $\theta(\phi, 1)$  returned by an  $\mathbb{A}\mathbb{1}\mathbb{g}$  that uses the update rule  $\Delta_{\mathrm{off}}(\theta, \phi, 1) - \theta$  that incorporates gradient of the L2 regularizer. Vector fields in Figure 1 were also obtained from this setup.
Notice that new update rule for  $\phi$  and  $\varphi$  can be obtained from steps (4) to (15) with the following  $\mathbf{A}$ ,  $\mathbf{B}$ , and  $\mathbf{H}$  instead, where for shorthand  $\theta^{*} = \theta (\phi ,\varphi)$
$\mathbf{A} = \mathbb{E}_{\mathcal{D}}\left[\sum_{t = 0}^{T}\psi_{\theta^{*}}(S_{t},A_{t})\left(\sum_{j = t}^{T}\gamma_{\varphi}^{j - t}\frac{r_{\phi}(S_{j},A_{j})}{\partial\phi}\right)^{\top}\right],$
$\mathbf{B} = \mathbb{E}_{\mathcal{D}}\left[\sum_{t = 0}^{T}\psi_{\theta^{*}}(S_{t},A_{t})\left(\sum_{j = t}^{T}\frac{\partial\gamma_{\varphi}^{j - t}}{\partial\varphi}\left(r_{\phi}(S_{j},A_{j}) - \lambda \ln \pi_{\theta^{*}}(S_{j},A_{j})\right)\right)\right],$
$\mathbf{H} = \mathbb{E}_{\mathcal{D}}\left[\sum_{t = 0}^{T}\frac{\partial\psi_{\theta^*}(S_t,A_t)}{\partial\theta^*}\left(\sum_{j = t}^{T}\gamma_{\varphi}^{j - t}\left(r_{\phi}(S_j,A_j) - \lambda \ln \pi_{\theta^*}(S_j,A_j)\right)\right) - \lambda \psi_{\theta^*}(S_t,A_t)\left(\sum_{j = t}^{T}\gamma_{\varphi}^{j - t}\psi_{\theta^*}(S_j,A_j)^\top\right)\right].$
E Meta Learning via Implicit Gradient: Derivation
The general technique of implicit gradients [14, 34, 19] has been used in a vast range of applications, ranging from energy models [13, 36], differentiating through black-box solvers [61], few-shot learning [38, 49], model-based RL [50], differentiable convex optimization neural-networks layers [4, 2], to hyper-parameter optimization [37, 8, 12, 40]. In this work, we show how implicit gradients can also be useful to efficiently leverage auxiliary rewards  $r_a$  and overcome various sub-optimialities.
Taking total derivative in (5) with respect to  $\phi$
$$
\frac {d \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{d \phi} = \frac {\partial \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \phi} + \frac {\partial \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \theta (\phi , \varphi)} \frac {\partial \theta (\phi , \varphi)}{\partial \phi} = 0. \tag {13}
$$
Let us try to understand why the above is true, considering the finite difference approach for this derivative,
$$
\begin{array}{l} \frac {d \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{d \phi} = \lim  _ {\| d \phi \| \rightarrow 0} \frac {\Delta (\theta (\phi + d \phi , \varphi) , \phi + d \phi , \varphi) - \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{d \phi} \\ = \frac {0 - 0}{d \phi} = 0, \\ \end{array}
$$
$\Delta (\theta (\phi +d\phi ,\varphi),\phi +d\phi ,\varphi) = \Delta (\theta (\phi ,\varphi),\phi ,\varphi) = 0$  as  $\theta (\cdot ,\cdot)$  defines convergence to fixed point.
By re-arranging terms in (13) we obtain the term (b) in (4),
$$
\frac {\partial \theta (\phi , \varphi)}{\partial \phi} = - \left(\frac {\partial \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \theta (\phi , \varphi)}\right) ^ {- 1} \frac {\partial \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \phi}. \tag {14}
$$
By combining (14) with (4) we obtain the desired gradient expression for  $\phi$
$$
\frac {\partial J (\theta (\phi , \varphi))}{\partial \phi} = - \frac {\partial J (\theta (\phi , \varphi))}{\partial \theta (\phi , \varphi)} \left(\underbrace {\frac {\partial \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \theta (\phi , \varphi)}} _ {\mathbf {H}}\right) ^ {- 1} \underbrace {\frac {\partial \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\partial \phi}} _ {\mathbf {A}},
$$
and following similar steps, it can be observed that the gradient expression for  $\varphi$