SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\frac{\partial\big(J(\theta(\phi,\varphi)) - \frac{1}{2}\|\gamma_{\varphi}\|^{2}\big)}{\partial\varphi} = -\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\varphi}}_{\mathbf{B}} - \frac{\partial\gamma_{\varphi}}{\partial\varphi},
$$
where using  $\theta^{*}$  as a shorthand for  $\theta (\phi ,\varphi)$  the terms  $\mathbf{A}$ ,  $\mathbf{B}$  and  $\mathbf{H}$  can be expressed as
$$
\begin{array}{l} \mathbf {A} = \mathbb {E} _ {\mathcal {D}} \left[ \sum_ {t = 0} ^ {T} \psi_ {\theta^ {*}} \left(S _ {t}, A _ {t}\right) \left(\sum_ {j = t} ^ {T} \gamma_ {\varphi} ^ {j - t} \frac {\partial r _ {\phi} \left(S _ {j} , A _ {j}\right)}{\partial \phi}\right) ^ {\top} \right], \quad \mathbf {B} = \mathbb {E} _ {\mathcal {D}} \left[ \sum_ {t = 0} ^ {T} \psi_ {\theta^ {*}} \left(S _ {t}, A _ {t}\right) \left(\sum_ {j = t} ^ {T} \frac {\partial \gamma_ {\varphi} ^ {j - t}}{\partial \varphi} r _ {\phi} \left(S _ {j}, A _ {j}\right)\right) \right], \tag {15} \\ \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} r _ {\phi} (S _ {j}, A _ {j})\right) \right] - \lambda . \\ \end{array}
$$
These provide the necessary expressions for updating  $\phi$  and  $\varphi$  in the outer loop. As  $\mathbf{A}$  involves an outer product and  $\mathbf{H}$  involves second derivatives, computing them exactly might not be practical when dealing with high-dimensions. Standard approximation techniques like conjugate-gradients or Neumann series can thus be used to make it more tractable [40]. In our experiments, we made use of the Neumann approximation to the Hessian Inverse vector product  $(\mathbf{AH}^{-1})$ , which requires the same magnitude of resources as the baseline policy gradient methods that we build on top off.
Algorithm: Being based on implicit gradients, we call our method BARFI, shorthand for behavior alignment reward function's implicit optimization. Overall, BARFI iteratively solves the bi-level optimization specified in (2) by alternating between using (3) till approximate converge of  $\mathrm{Alg}$  to  $\theta (\phi ,\varphi)$  and then updating  $r_{\phi}$  and  $\gamma_{\varphi}$ . Importantly, being based on (3) for sample efficiency,  $\mathrm{Alg}$  leverages only the past samples and does not sample any new trajectories for the inner level optimization. Further, due to policy regularization which smoothens the objective as discussed in D, updates in  $r_{\phi}$  and  $\gamma_{\varphi}$  changes the policy resulting from  $\mathrm{Alg}$  gradually. Therefore, for compute efficiency, we start  $\mathrm{Alg}$  from the policy obtained from the previous inner optimization, such that it is in proximity of the new fixed point. This allows BARFI to be both sample and compute efficient while solving the bi-level optimization iteratively online. Pseudo-code for BARFI and more details on the approximation techniques can be found in Appendix C.
F Environment & Reward Details
The first environment is a GridWorld (GW), where the start state is in the bottom left corner and a goal state is in the top right corner. The agent receives an  $r_p$  of  $+100$  on reaching the goal followed by termination of the episode. The second environment is MountainCar (MC) [58], wherein we make use of the sparse reward variant, wherein the agent receives a  $+1$  reward on reaching on top of the hill and 0 otherwise. The third environment is CartPole (CP) [16]. Finally, to assess the scalability we pick HalfCheetah-v4 from Mujoco (MJ) suite of OpenAI Gym [9].
For each environment, we define two auxiarily reward functions. For GridWorld, we define the functions:  $r_{\mathrm{aux,GW}}^{1} \coloneqq -(s - s_{\mathrm{goal}})^{2}$ , which provides the negative L2 squared distance from the goal position, and  $r_{\mathrm{aux,GW}}^{1} \coloneqq 50 \times \mathbf{1}_{s \in S_{\mathrm{center}}}$ , which provides an additional bonus of +50 to the agent along the desired path to the goal state (i.e. the center states). In MountainCar the state is composed of two components: the position  $x$ , and velocity  $v$ . The first auxiliary reward function,  $r_{\mathrm{aux,MC}}^{1}(s,a) \coloneqq |\mathbf{v}|$ , encourages a higher absolute velocity of the car, and the second,  $r_{\mathrm{aux,MC}}^{1}(s,a) \coloneqq \mathbf{1}_{\mathrm{sign(v)} = a}$ , encourages the direction of motion to increase the magnitude of the velocity (also known as the energy pumping policy [18]). For CartPole, we consider a way to reuse knowledge from a hand crafted policy. CartPole can be solved using a Proportional Derivate (PD) controller [6], hence we tune a PD controller,  $\mathsf{PD}^* : \mathcal{S} \to \mathcal{A}$ , to solve CartPole for the max possible return. We design two auxiliary
reward functions which make use of this PD controller. The first,  $r_{\mathrm{aux,CP}}^{1}(s,a) \coloneqq 5 \times \mathbf{1}_{\mathrm{PD}^{*}(s) = a} - (1 - 1_{\mathrm{PD}^{*}(s) = a})$ , encourages the agent to match the action of the optimal PD controller, and penalizes it for not matching. The second auxiliary reward function,  $r_{\mathrm{aux,GW}}^{1}(s,a) \coloneqq -r_{\mathrm{aux,CP}}^{1}(s,a)$ , encourages the agent to do the opposite. In the case of Mujoco, the reward function provided by the environment is itself composed of multiple different functions. We explain the same and the respective auxiliary functions for this case later.
We have considered several forms of information encoded as auxiliary rewards for these experiments. We have heuristic-based reward functions (i.e.,  $r_{\mathrm{aux,GW}}^{1}, r_{\mathrm{aux,GW}}^{1}, r_{\mathrm{aux,MC}}^{1}$ ). Reward functions that encode a guess of an optimal policy (i.e.,  $r_{\mathrm{aux,MC}}^{1}, r_{\mathrm{aux,CP}}^{1}$ ) and reward functions that change the optimal policy (i.e.,  $r_{\mathrm{aux,GW}}^{1}, r_{\mathrm{aux,CP}}^{1}$ ). We also have rewards that only depend on states (i.e.,  $r_{\mathrm{aux,GW}}^{1}, r_{\mathrm{aux,GW}}^{1}, r_{\mathrm{aux,MC}}^{1}$ ) as well as ones that depend on both state and actions (i.e.,  $r_{\mathrm{aux,MC}}^{1}, r_{\mathrm{aux,CP}}^{1}, r_{\mathrm{aux,CP}}^{1}$ ). Therefore, we can test if BARFI can overcome misspecified auxiliary reward functions and does not hurt performance when well-specified.
Mujoco Environment In this experiment, we investigate the scalability of BARFI in learning control policies for high-dimensional state spaces with continuous action spaces. In HalfCheetah-v4 the agent's task is to move forward, and it receives a reward based on its forward movement (denoted as  $r_p$ ). Additionally, there is a small cost associated with the magnitude of torque required for action execution (denoted as  $r_{\mathrm{aux}}(s,a) \coloneqq c|a|_2^2$ ). The weighting between the main reward and the control cost is pre-defined as  $c$  for this environment, and we form the reward as  $\tilde{r}(s,a) = r_p(s,a) + r_{\mathrm{aux}}(s,a)$ . We explore how an arbitrary weighting choice can cause the agent to fail in learning, while BARFI is still able to adapt and learn the appropriate weighting, remaining robust to possible misspecification. We consider two different weightings for the control cost: the first weighting, denoted as  $r_{\mathrm{aux},\mathrm{MJ}}^1(s,a) \coloneqq -c|a|_2^2$ , uses the default setting, while the second weighting, denoted as  $r_{\mathrm{aux},\mathrm{MJ}}^1(s,a) \coloneqq -4c|a|_2^2$ , employs a scaled variant of the first weighting. Additionally, we implement the path-wise bi-level optimization variant i.e., BARFI Unrolled. In these experiments, we keep the value of  $\gamma$  fixed to isolate the agent's capability to adapt and recover from an arbitrary reward weighting. We will also measure what computational and performance tradeoffs we might have to make between using the implicit version i.e., BARFI against, the path-wise version i.e., BARFI Unrolled.
G Details for the Empirical Results
G.1 Implementation Details
In this section we will briefly describe the implementation details around the different environments that were used.
GridWorld (GW): In the case of GridWorld we made use of the Fourier basis (of Order  $= 3$ ) over the raw coordinates of agent position in the GridWorld. Details about this could be found in the src/utils/Basis.py file.
MountainCar (MC): For this environment, to reduce the limitation because of the function approximator we used TileCoding [58], which offers a suitable representation for the MountainCar problem. We used 4 Tiles and Tilings of 5.
CartPole (CP): For CartPole also make use of Fourier Basis of (Order = 3), with linear function approximator on top of that.
MuJoco (MJ): For this we made use of a neural network with 1 hidden layer of 32 nodes and ReLU activation as the function approximator over the raw observations. The output of the policy is continuous actions, hence we used a Gaussian representation, where the policy outputs the mean of the multivariate Gaussian and we used a fixed diagonal standard deviation, fixed to  $\sigma = 0.1$ .
General Details: All the outer returns are evaluated without any discounting, whereas all the inner optimizations were initialized with  $\gamma_{\varphi} = 0.99$ . Hence to do this we made  $\varphi$  a single bias unit, initialized to 4.6, and passed through a sigmoid (i.e.,  $\sigma(4.6) = 0.99$ ).
For GW, CP and MC  $r_{\phi}$  is defined as below
$$
r _ {\phi} (s, a) = \phi_ {1} (s) + \phi_ {2} (s) r _ {p} + \phi_ {3} (s) r _ {a}
$$
Wherein  $\phi_1, \phi_2, \phi_3$  are scalar outputs of a 3-headed function, in this case simply a linear layer over the states inputs.
Whereas in the case of MJ, we have
$$
r _ {\phi} (s, a) = \phi_ {1} + r _ {p} + \phi_ {3} r _ {a}
$$
Wherein  $\phi_{1}$  is initialized to zero and  $\phi_3$  is 1.0 act like bias units.
Gradient normalization was used for all the cases where neural nets were involved (i.e., MJ), and also for MJ we modified the Baseline (REINFORCE) update to subtract the running average of the performance as a baseline to get acceptable performance for the baseline method.
G.2 Hyper-parameter Selection
As different make use of different function approximators the range of hyper-params can vary we talk about all the above over here.
Best-performing Parameters for different methods and environments are listed where
Table 3: Hyper-parameters for GridWorld
Hyper Parameter	BARFI Value	REINFORCE Value	Actor-Critic Value
αθ	1 × 10-3	1 × 10-3	1 × 10-3
αφ	5 × 10-3	-	-
αφ	5 × 10-3	-	-
optim	RMSprop	RMSprop	RMSprop
λθ	0.25	0.25	0.25
λφ	0.0625	-	-
λφ	4.0	-	-
Buffer	1000	-	-
Batch Size	1	1	1
η	0.0005	-	-
δ	3	-	-
n	5	-	-
N0	150	-	-
Ni	15	-	-
Table 4: Hyper-parameters for MountainCar
Hyper Parameter	BARFI Value	REINFORCE Value	Actor-Critic Value
αθ	0.015625	0.125	0.03125
αφ	0.0625	-	-
αφ	0.0625	-	-
optim	RMSprop	RMSprop	RMSprop
λθ	0.0	0.0	0.25
λφ	0.0	-	-
λφ	0.25	-	-
Buffer	50	-	-
Batch Size	1	1	1
η	0.001	-	-
δ	3	-	-
n	5	-	-
N0	50	-	-
Ni	15	-	-
Hyperparameter Sweep : Here we list the details about how we swept the values for different hyper-params. We used PyTorch [46] for all our implementations. We usually used an optimizer between RMSProp or Adam with default parameters as provided in Pytorch. For  $\alpha_{\theta} \in \{5 \times$
Table 5: Hyper-parameters for CartPole
Hyper Parameter	BARFI Value	REINFORCE Value	Actor-Critic Value
αθ	1 × 10-3	1 × 10-3	5 × 10-4
αφ	1 × 10-3	-	-
αφ	5 × 10-3	-	-
optim	RMSprop	RMSprop	RMSprop
λθ	1.0	1.0	0.0
λφ	0.0	-	-
λφ	4.0	-	-
Buffer	10000	-	-
Batch Size	1	1	1
η	0.0005	-	-
δ	3	-	-
n	5	-	-
N0	150	-	-
Ni	15	-	-
Table 6: Hyper-parameters for MuJoco
Hyper Parameter	BARFI Value	REINFORCE Value	Actor-Critic Value
αθ	7.5 × 10-5	5 × 10-4	2.5 × 10-4
αφ	2.5 × 10-3	-	-
αφ	0.0	-	-
optim	Adam	Adam	Adam
λφ	0.0625	-	-
λφ	0.0	-	-
Buffer	50	-	-
Batch Size	1	1	1
η	0.0005	-	-
δ	3	-	-
n	5	-	-
N0	30	-	-
Ni	15	-	-
$10^{-3}, 2.5 \times 10^{-3}, 1 \times 10^{-3}, 5 \times 10^{-4}, 2.5 \times 10^{-4}, 1 \times 10^{-4}, 7.5 \times 10^{-5}\}$ , we use similar ranges for  $\alpha_{\phi}, \alpha_{\varphi}$  (which tend to be larger). For  $\lambda_{\theta}$  and  $\lambda_{\phi}$ , we swepted from  $[0, 0.25, 0.5, 1.0]$  and for  $\lambda_{\gamma}$  we swepted from  $[0, 0.25, 1.0, 4.0, 16.0]$ . We simply list ranges for different values and later we present sensitivity curves showing that these values are usually robust for BARF I across different methods as we can see from the tables above.  $\delta \in [1, 3, 5]$ ,  $n \in [1, 3, 5]$ ,  $N_{i} \in [1, 3, 6, 9, 12, 15]$ ,  $\eta \in [1 \times 10^{-3}, 5 \times 10^{-4}, 1 \times 10^{-4}]$ ,  $N_{0} \in [30, 50, 100, 150]$ , buffer  $\in [25, 50, 100, 1000]$ .  $\alpha$  for Tilecoding was adopted from [58] and hence similar ranges were swepted in that case. Most sweeps were done with around 10 seeds, and later the parameter ranges were reduced and performed with more seeds.
G.3 Compute
The computer is used for a cluster where the CPU class is Intel Xeon Gold 6240 CPU @2.60GHz. The total compute required for GW was around 3 CPU years $^7$ , CP also required around 3 CPU years, and MC required around 4 CPU years. For MJ we needed around 5-6 CPU years. In total we utilized around 15-16 CPU years, where we needed around 1 GB of memory per thread.
1 CPU year := Compute equal to running a CPU thread for a year.
H Extra Results & Ablations
H.1 Experiment on partially misspecified  $r_{\mathrm{aux}}$
In these set of experiments, we consider the case where auxiliary reward signals could be helpful only in a few—possibly arbitrary—state-action pairs. In general, we anticipate that solutions in this scenario would be such that assigns weightings, allowing the agent to avoid regions where might be misspecified. Meanwhile, the agent would still make use of the places where is well specified and useful.
We consider another  $r_{\mathrm{aux}}$  in the GridWorld domain in which the auxiliary reward is misspecified for a subset of states near the starting position. Meanwhile, it is still well-specified for states near the goal (Figure 6 (a)). Figures 6 (b) and (c) illustrate the learned and the weighting on, showcasing the agent's ability to depict the expected behavior described above.
(a)  $r_{\mathrm{aux}}(s)$
(b) Learned  $\phi_3(s)$ , the weighting on  $r_{\mathrm{aux}}(s)$
(c) The net reward induced i.e.,  $r_p(s) + \phi_3(s)r_{\mathrm{aux}}(s)$
(d) Learning performance
Figure 6: 40 random seeds were used to generate the plots. The starting state is at the bottom left and the goal state is at the top right corner. The primary reward  $r_p$  is  $+100$  when the agent reaches the goal and 0 otherwise. (a) A state-dependent  $r_{\mathrm{aux}}$  function that is partially misspecified (in the blue region  $r_{\mathrm{aux}}$  provides a value equal to the Manhattan distance, thereby incentivizing the agent to stay near the start), and partially well specified (in the red region, it is the negative Manhattan distance and encourages movement towards the goal). (b) The state-dependent weighting  $\phi_3(s)$  learned by BARFI negates the positive value from  $r_{\mathrm{aux}}$  near the start state. (c) The effective reward function  $r_p(s) + \phi_3(s)r_{\mathrm{aux}}(s)$  learned by BARFI. (d) Learning curves for BARFI, and the baseline that uses a naive addition  $(r_p(s) + r_{\mathrm{aux}}(s))$  of the above auxiliary reward.
Figure 7: Return induced by learned reward functions: This figure illustrates the aux return collected by agent based on the learned  $r_{\phi}$ , the curves are chosen based on best-performing curves on  $r_p$ , and averaged over 20 runs (except 40 for GW). (a) Top  $-r_{aux,GW}^{1}$ , Bottom  $-r_{aux,GW}^{2}$ , (b) Top  $-r_{aux,CP}^{1}$ , Bottom  $-r_{aux,CP}^{2}$ , (c) Top  $-r_{aux,MC}^{2}$ , Bottom  $-r_{aux,MC}^{1}$ , (d) Top  $-r_{aux,MJ}^{1}$ , Bottom  $-r_{aux,MJ}^{2}$ .
H.2 Return based on learned  $r_{\phi}$  and  $\gamma_{\varphi}$
Figure 7 and Figure 8 summarize the achievable return based on  $r_{\phi}$  and the  $\gamma$  learned by the agent across different domains and reward specification. We observe that REINFORCE often optimizes the naive combination of reward for sure, but that doesn't really lead to a good performance on  $r_p$ , whereas BARFI does achieve an appropriate return on  $r_{\phi}$ , but is also able to successively decay  $\gamma$  as the learning progress across different domains. Particularly notice Figure 7 (a) Bottom, where REINFORCE does optimize aux return a lot, but actually fails to solve the problem, as it simply learns to loop around the center state.
H.3 Ablations
Figure 9 represents the ablation of BARFI on GridWorld with the misspecified reward for its different params. We can see that usually having  $\eta = 0.001$ ,  $0.0005$ ,  $n = 5$  works for the approximation.
Figure 8: Learned discounting  $\gamma_{\varphi}$ : This figure illustrates the learned  $\gamma_{\varphi}$  for BARFI and normal  $\gamma$  for other methods, the curves are chosen based on best-performing curves on  $r_p$ , and averaged over 20 runs (except 40 for GW). (a) Top  $r_{aux,GW}^{1}$ , Bottom  $r_{aux,GW}^{2}$ , (b) Top  $r_{aux,CP}^{1}$ , Bottom  $r_{aux,CP}^{2}$ , (c) Top  $r_{aux,MC}^{1}$ , Bottom  $r_{aux,MC}^{1}$ . Mujoco is not included as the  $\gamma$  was not learned in that case. We can observe that the agents start to learn to decay  $\gamma$  at the appropriate pace. Note that the curves for methods other than BARFI and BARFI Unrolled are overlapping as the baselines don't change the value of  $\gamma$ .
Figure 9: Sensitivity Curves: The set of graphs representing the sensitivity of different hyper-params keeping all the other params fixed. The sensitivity is for BARFI in GW with  $r_{\mathrm{aux},\mathrm{GW}}^2$ , i.e., the misspecified reward. We choose the best-performing parameters and vary each parameter to see its influence. The curves are obtained for 50 runs (seeds) in each case, and error bars are standard errors. We can notice that  $\alpha_{\theta}$  and  $\alpha_{\phi}$  can have a large influence, and tend to stay around similar values.  $\lambda_{\theta,\phi,\varphi}$  tends to help but doesn't really influence a lot in terms of its magnitude, except larger values of  $\lambda_{\varphi}$  seem to do better. Smaller values of  $\eta$  seems to work fine, hence something around  $5\times 10^{-4}$ ,  $1\times 10^{-3}$  usually should suffice.  $n, \delta$  can be chosen to around 5 and 3, and usually workout fine. We also defined  $N_{i} = 5\times \delta$  in this case.