SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Let us denote the unbiased off-policy update with full-distribution correction as  $\Delta_{\mathrm{off}}(\theta, r_p) \coloneqq \mathbb{E}_{\beta}[\sum_{t=0}^{T} \gamma^t \psi_\theta(S_t, A_t) \sum_{j=t}^{T} \rho_j \gamma^j r_p(S_j, A_j)]$ . Now consider an extreme scenario where off-policy evaluation is attempted without any correction for distribution shift. In this situation, and with a slight abuse of notation, we define  $\Delta_{\mathrm{off}}(\theta, \phi, \varphi) \coloneqq \mathbb{E}_{\beta}[\sum_{t=0}^{T} \psi_\theta(S_t, A_t) \sum_{j=t}^{T} \gamma_{\varphi}^{j-t} r_\phi(S_j, A_j)]$ .
Property 3. There exists  $r_{\phi} : S \times \mathcal{A} \to \mathbb{R}$  and  $\gamma_{\varphi} \in [0,1)$  such that  $\Delta_{\mathrm{off}}(\theta, \phi, \varphi) = \Delta_{\mathrm{off}}(\theta, r_p)$ .
Remark 1. Our method is capable of mitigating various types of algorithmic biases and imperfections in underlying RL algorithms, without requiring any specialized learning rules. Additionally, thanks
<sup>4</sup>E.g., if  $r_{\phi} \approx q^{*}$ , then its corresponding rewards are instantly indicative of the optimal action at any state.
to the  $\gamma_{\varphi^*}$  regularization, it favors reward functions that lead to faster learning of high-performing policies aligned with the designer's objectives, as outlined in the original reward function  $r_p$ .
5 BARFI: Implicitly Learning Behavior Alignment Rewards
Having introduced our bi-level objective and discussed the benefits of optimizing  $r_{\phi}$  and  $\gamma_{\varphi}$ , an important question arises: Although  $\theta (\phi ,\varphi)$  can be optimized using any policy learning algorithm, how can we efficiently identify the optimal  $\phi^{*}$  and  $\varphi^{*}$  in equation (2)? Given the practical advantages of gradient-based methods, one would naturally consider using them for optimizing  $\phi^{*}$  and  $\varphi^{*}$  as well. However, a key challenge in our setting lies in computing  $\mathrm{d}J(\theta (\phi ,\varphi)) / \mathrm{d}\phi$  and  $\mathrm{d}J(\theta (\phi ,\varphi)) / \mathrm{d}\varphi$ . These computations require an analytical characterization of the impact that  $r_{\phi}$  and  $\gamma_{\varphi}$  have on the entire optimization process of the inner-level algorithm, A1g.
In addressing this challenge, we initially focus on an  $\mathsf{Alg}$  that employs policy gradients for updating  $\pi_{\theta}$ . Similar extensions for other update rules can be derived similarly. We start by re-writing the expression for  $\mathrm{d}J(\theta (\phi ,\varphi)) / \mathrm{d}\phi$  using the chain rule:
$$
\frac {\mathrm {d} J (\theta (\phi , \varphi))}{\mathrm {d} \phi} = \underbrace {\frac {\mathrm {d} J (\theta (\phi , \varphi))}{\mathrm {d} \theta (\phi , \varphi)}} _ {(a)} \underbrace {\frac {\mathrm {d} \theta (\phi , \varphi)}{\mathrm {d} \phi}} _ {(b)}, \tag {4}
$$
where  $(\pmb{a})$  is the policy gradient at  $\theta (\phi ,\varphi)$ , and  $(\pmb{b})$  can be computed via implicit bi-level optimization, as discussed below.
Implicit Bi-Level Optimization: We compute (4) by leveraging implicit gradients [14, 34, 19], an approach previously employed, e.g., in few-shot learning [38, 49] and model-based RL algorithms [50]. First, observe that when  $\mathsf{A}\mathsf{1}\mathsf{g}$  converges to  $\theta (\phi ,\varphi)$ , then it follows that
$$
\Delta (\theta (\phi , \varphi), \phi , \varphi) = 0. \tag {5}
$$
Let  $\partial f$  denote the partial derivative with respect to the immediate arguments of  $f$ , and  $\mathrm{d}f$  be the total derivative as before. That is, if  $f(x,g(x))\coloneqq xg(x)$ , then  $\frac{\partial f}{\partial x} (x,g(x)) = g(x)$  and  $\frac{\mathrm{d}f}{\mathrm{d}x} (x,g(x)) = g(x) + x\frac{\partial g}{\partial x} (x)$ . Therefore, taking the total derivative of (5) with respect to  $\phi$  yields
$$
\frac {\mathrm {d} \Delta (\theta (\phi , \varphi) , \phi , \varphi)}{\mathrm {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 {6}
$$
By re-arranging terms in (6), we obtain the term  $(\pmb{b})$  in (4). In particular,
$$
\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 {7}
$$
Furthermore, by combining (7) and (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}}. \tag {8}
$$
Similarly, a gradient expression for  $\varphi$  can be derived; the full derivation is detailed in Appendix E. Using  $\theta^{*}$  as shorthand for  $\theta (\phi ,\varphi)$ , we find that the terms  $\mathbf{A}$  and  $\mathbf{H}$  can be expressed as
$$
\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 {\partial r _ {\phi} (S _ {j} , A _ {j})}{\partial \phi}\right) ^ {\top} \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} r _ {\phi} (S _ {j}, A _ {j})\right) \right].
$$
When working with the equations above, we assume the inverse of  $\mathbf{H}^{-1}$  exists. To mitigate the risk of ill-conditioning, we discuss regularization strategies for  $\mathsf{A}\mathsf{1}\mathsf{g}$  in Appendix D. Notice that equations (8) and (15) are the key elements needed to calculate the updates to  $\phi$  and  $\varphi$  in our bi-level optimization's outer loop. However, computing  $\mathbf{A}$  and  $\mathbf{H}$  directly can be impractical for high-dimensional problems due to the need for outer products and second derivatives. To address this, we employ two strategies: (1) We approximate  $\mathbf{H}^{-1}$  using the Neumann series [40], and (2) we calculate (8) and (15) via Hessian-vector products [47], which are readily available in modern auto-diff libraries [46]. These methods eliminate the need for explicit storage or computation of  $\mathbf{H}$  or  $\mathbf{A}$ .
The mathematical approach outlined above results in an algorithm with linear compute and memory footprint, having  $O(d)$  complexity, where  $d$  is the number of parameters for both the policy and the reward function. Details can be found in Appendix C. We refer to our method as BARFI, an acronym for behavior alignment reward function's implicit optimization.5 BARFI is designed to iteratively solve the bi-level optimization problem defined in (2). With policy regularization, the updates to  $r_{\phi}$  and  $\gamma_{\varphi}$  incrementally modify  $\theta(\phi, \varphi)$ . This enables us to initialize Alg using the fixed point achieved in the previous inner optimization step, further reducing the time for subsequent inner optimizations.
6 Empirical Analyses
Our experiments serve multiple purposes and include detailed ablation studies. First, we demonstrate our bi-level objective's efficacy in discovering behavior alignment reward functions that facilitate learning high-performing policies. We focus especially on its robustness in situations where designers provide poorly specified or misaligned auxiliary rewards that could disrupt the learning process (Section 6.1). Second, we present a detailed analysis of the limitations of potential-based reward shaping, showing how it can lead to suboptimal policies (Section 6.2). We then provide a qualitative illustration of the behavior alignment reward function learned by BARFI (Section 6.3). Finally, we evaluate how well BARFI scales to problems with high-dimensional, continuous action and state spaces (Section 6.4).
In the sections that follow, we examine a range of methods and reward combinations for comparison:
- Baseline RL methods: We consider baseline RL methods that employ a naive reward combination strategy:  $\tilde{r}_{\mathrm{naive}}(s,a) \coloneqq r_p(s,a) + r_{\mathrm{aux}}(s,a)$ . In this case, the auxiliary reward from the designer is simply added to the original reward without checks for alignment. Both the REINFORCE and Actor-Critic algorithms are used for optimization.
- Potential-based shaping: To assess how well potential-based reward shaping performs, we introduce variants of the baseline methods. Specifically, we investigate the effectiveness of the reward function  $\tilde{r}_{\Phi}(s,a,s') \coloneqq r_p(s,a) + \gamma r_{\mathrm{aux}}(s') - r_{\mathrm{aux}}(s)$ .
- BARFI: We use REINFORCE as the underlying RL algorithm when implementing BARFI and define  $r_{\phi}(s,a) \coloneqq \phi_1(s,a) + \phi_2(s)r_p(s,a) + \phi_3(s)r_{\mathrm{aux}}(s,a)$ . Our implementation includes a warm-up period wherein the agent collects data for a fixed number of episodes, using  $\tilde{r}_{\mathrm{naive}}$ , prior to performing the first updates to  $\phi$  and  $\varphi$  (See Appendix 5 for the complete algorithm).
We evaluate each algorithm across four distinct environments: GridWorld, MountainCar [58], Cart-Pole [16], and HalfCheetah-v4 [9]. These domains offer increasing levels of complexity and are intended to assess the algorithms' adaptability. Furthermore, we examine their performance under a variety of auxiliary reward functions, ranging from well-aligned to misaligned with respect to the designer's intended objectives.
In our experiments, we investigate different types of auxiliary reward functions for each environment: some are action-dependent, while others are designed to reward actions aligned with either effective or ineffective known policies. These functions, therefore, vary in their potential to either foster rapid learning or inadvertently mislead the agent away from the designer's primary objectives, hindering the efficiency of the learning process. Comprehensive details of each environment and their corresponding auxiliary reward functions can be found in Appendix F.
6.1 BARFI's Robustness to Misaligned Auxiliary Reward Functions
In this section, we evaluate the performance of various methods for reward combination, particularly in scenarios where auxiliary reward functions can either be well-aligned with a designer's intended goals or be misaligned or poorly specified, thus inadvertently hindering efficient learning. We introduce two types of auxiliary reward functions for CartPole. First, we used domain knowledge to design an  $r_{\mathrm{aux}}$  that provides bonuses when the agent's actions align with a known effective policy in this domain. Second, we designed an adversarial example where the auxiliary reward function rewards actions that are consistent with a particularly poorly performing policy. For MountainCar, we first leveraged knowledge about an energy pumping policy (i.e., a well-known effective policy [18]) to
$5_{\text{BARFI}}$  commonly refers to a type of south-Asian sweet confectionery, typically pronounced as 'bur-fee'.
Table 1: Summary of the performance of various reward combination methods and types of  ${r}_{\text{aux }}$
Method for Reward Combination	CartPole	MountainCar
Well-aligned raux	Misaligned raux	Well-aligned raux (w.r.t. energy policy)	Partially-aligned raux (w.r.t. high velocity policy)
BARFI (our method)	487.2 ± 9.4	475.5 ± 15.5	0.99 ± 0.0	0.90 ± 0.1
rnaive (naive reward combination)	498.9 ± 1.0	9.04 ± 0.2	0.99 ± 0.0	0.63 ± 0.1
rΦ (potential-based shaping)	8.98 ± 0.2	500 ± 0.0	0.00 ± 0.0	0.00 ± 0.0
BARFI's performance compared to two baselines that use:  $\tilde{r}_{\mathrm{naive}}$  and  $\tilde{r}_{\Phi}$ , respectively. CartPole uses an action-dependent  $r_{\mathrm{aux}}$  function that either rewards agents when actions align with a known effective policy (well-aligned  $r_{\mathrm{aux}}$ ) or with a poorly-performing policy (misaligned  $r_{\mathrm{aux}}$ ). MountainCar uses either an action-dependent function aligned with an energy-pumping policy [18] or a partially-aligned function incentivizing higher velocities. BARFI consistently achieves near-optimal performance across scenarios, even if given poorly specified/misaligned auxiliary rewards. Competitors, by contrast, often induce suboptimal policies. Performances significantly below the optimal are shown in red, above.
craft an auxiliary reward function that provides bonuses for actions in line with such a control strategy. We also experimented with a partially-aligned auxiliary function that rewards high velocity—a factor not particularly indicative of high performance.
Table 1 summarizes the results across different reward functions and combination methods. The results suggest that if auxiliary rewards provide positive feedback when agents' actions align with effective policies, naive combination methods perform well. In such cases, auxiliary rewards effectively "nudge" agents towards emulating expert actions. However, our experimental results also indicate that all baseline methods are susceptible to poor performance when auxiliary rewards are not well aligned with the designer's goals. We provide more discussion for potential-based shaping in Section 6.2.
The key takeaway from the experimental results in Table 1 is that BARFI consistently performs well across various domains and under different types of auxiliary rewards. Specifically, when designer-specified feedback is appropriate and can assist in accelerating learning, BARFI efficiently exploits it to produce high-performing policies. Conversely, if auxiliary rewards are misaligned with the designer's intended goals, BARFI is capable of adapting and effectively dismissing "misleading rewards". This adaptability ensures that high-performing policies can be reliably identified. Other methods, by contrast, succeed only in some of these scenarios. Importantly, the unpredictability of whether a given auxiliary reward function will aid or hinder learning makes such alternative methods less reliable, as they may fail to learn effective policies.
6.2 Pitfalls of potential-based reward shaping
We now turn our attention to the (possibly negative) influence of action-dependent auxiliary rewards, particularly when used in combination with potential-based reward shaping. Our results in Table 1 reveal a key limitation: potential-based shaping struggles to learn efficient policies even when auxiliary rewards are well-aligned with effective strategies. This shortcoming is attributable to the action-dependent nature of the auxiliary rewards, which compromises the potential shaping technique's guarantee of policy optimality.
As there is no prescribed way for designing potential shaping when  $r_{\mathrm{aux}}$  is action-dependent, we use a direct extension of the original formulation [41] by considering  $\tilde{r}_{\Phi}(s,a,s',a') \coloneqq r_p(s,a) + \gamma r_{\mathrm{aux}}(s',a') - r_{\mathrm{aux}}(s,a)$ . Furthermore, we designed an auxiliary reward function,  $r_{\mathrm{aux}}$ , that is well aligned: it provides positive reward signals of fixed magnitude both for  $(s,a)$  and  $(s',a')$  whenever the agent's actions coincide with the optimal policy. Notice, however, that if  $\gamma < 1$ , the resultant value from  $\gamma r_{\mathrm{aux}}(s',a') - r_{\mathrm{aux}}(s,a)$  is negative. Such a negative component may deter the agent from selecting actions that are otherwise optimal, depending on how  $r_{\mathrm{aux}}$  and  $r_p$  differ in magnitude. Conversely, potential-based shaping can also occasionally perform well under misaligned rewards. In these cases, the shaping function may yield a positive value whenever the agent selects an optimal action, which could induce well-performing behaviors.
Figure 2: (Left) Misaligned  $r_{\mathrm{aux}}(+50)$  for entering the center state and  $r_p(+100)$  at the goal state. (Center Left) Optimized weight  $\phi_3(s)$  depicting how the learned reward function penalizes the agent for entering the center state. (Center Right) Performances in the GridWorld domain under a misspecified  $r_{\mathrm{aux}}$ . (Right) Performances in the CartPole domain under a  $r_{\mathrm{aux}}$  providing positive feedback for mimicking a known ineffective policy.
6.3 What does  $r_{\phi}$  learn?
We now investigate BARFI's performance and robustness in the GridWorld when operating under misspecified auxiliary reward functions. Consider the reward function depicted in Figure 2 [left]. This reward function provides the agent with a bonus for visiting the state at the center, akin to providing intermediate feedback to the agent when it makes progress towards the goal. However, such intermediate positive feedback can lead to behaviors where the agent repeatedly cycles around the middle state (i.e., behaviors that are misaligned with the original objective of reaching the goal state at the top right corner of the grid). Importantly, BARFI is capable of autonomously realizing that it should disregard such misleading incentives (Figure 2 [center left]), thereby avoiding poorly-performing behaviors that focus on revisiting irrelevant central states. Similarly, when Cartpole operates under a misspecified  $r_{\mathrm{aux}}$  (Figure 2 [right]), BARFI is capable of rapidly adapting (after a warm-up period) and effectively disregarding misleading auxiliary reward signals. These results highlight once again BARFI's robustness when faced with reward misspecification.
6.4 Scalability to High-Dimensional Continuous Control
One might wonder whether computing implicit gradients for  $\phi$  and  $\varphi$  would be feasible in high-dimensional problems, due to the computational cost of inverting Hessians. To address this concern, we leverage Neumann series approximation with Hessian-vector products (See Appendix C) and conduct further experiments, as shown in Figure 3. These experiments focus on evaluating the scalability of BARFI in control problems with high-dimensional state spaces and continuous actions—scenarios that often rely on neural networks for both the policy and critic function approximators. For a more comprehensive evaluation, we also introduced an alternative method named BARFI unrolled. Unlike BARFI, which uses implicit bi-level optimization, BARFI unrolled employs path-wise bi-level optimization. It maintains a complete record of the optimization path to determine updates for  $\phi$  and  $\varphi$ . Further details regarding this alternative method can be found in Appendix C.6.
We conducted experiments on the HalfCheetah-v4 domain and investigated, in particular, a reward function comprising two components with varying weights. This empirical analysis was designed to help us understand how different weight assignments to each reward component could influence the learning process. Specifically, in HalfCheetah-v4, the agent receives a positive reward  $r_p$  proportional to how much it moved forward. It also incurs a small negative reward (concretely, an auxiliary reward,  $r_{\mathrm{aux}}(s,a)\coloneqq c\| a\| _2^2$ , known as a control cost) for the torque applied to its joints. A hyperparameter  $c$  determines the balance between these rewards. The naive combination of such primary and auxiliary rewards is defined as  $\tilde{r}_{\mathrm{naive}}(s,a) = r_p(s,a) + r_{\mathrm{aux}}(s,a)$ . Figure 3 [left] shows that the baselines and both variants of BARFI appear to learn effectively. With alternative reward weighting schemes, however, only BARFI and BARFI unrolled show learning progress, as seen in Figure 3 [middle]. It is worth noting that path-wise bi-level optimization can become impractical as the number of update steps in (4) increases, due to growing computational and memory requirements (Figure 3 [right]). Although we do not recommend BARFI unrolled, we include its results for completeness. Additional ablation studies on (a) the effect of the inner optimization step; (b) Neumann approximations; (c) decay of  $\gamma$ ; and (d) returns based on  $r_{\phi}$ , are provided in Appendix H.
Figure 3: Results for MuJoCo environment. (Left) Auxiliary reward is defined to be  $-c\|a\|_2^2$ , where  $c$  is a positive hyperparameter and  $a$  is the continuous high-dimensional action vector. (Middle) Similar setting as before, but uses an amplified variant of the auxiliary reward:  $-4c\|a\|_2^2$ . It is worth highlighting that even under alternative reward weighting schemes, both variants of our (behavior-aligned) bi-level optimization methods demonstrate successful learning. Learning curves correspond to mean return over 15 trials, and the shaded regions correspond to one standard error. (Right) Required compute and memory for BARFI unrolled, compared to BARFI, as a function of the number of inner-optimization updates. This figure also showcases BARFI's characteristics under various orders of Neumann approximation.
7 Related work
This paper focuses primarily on how to efficiently leverage auxiliary rewards  $r_{\mathrm{aux}}$ . Notice, however, that in the absence of  $r_{\mathrm{aux}}$ , the resulting learned behavior alignment rewards  $r_{\phi}$  may be interpreted as intrinsic rewards [70, 71]. Furthermore, several prior works have investigated meta-learning techniques, which are methods akin to the bi-level optimization procedures used in our work. Such prior works have employed meta-learning in various settings, including automatically inferring the effective return of trajectories [68, 62, 7, 71], parameters of potential functions [72, 28, 17], targets for TD learning [69], rewards for planning [54, 23], and even fully specified reinforcement learning update rules [33, 45]. Additionally, various other relevant considerations to effectively learning rewards online have been discussed by Armstrong et al. [5]. Our work complements these efforts by focusing on the reward alignment problem, specifically in settings where auxiliary information is available. An extended discussion on related works can be found in Appendix B. It is worth mentioning that among the above-mentioned techniques, most rely on path-wise meta-gradients. As discussed in Section 6.4, this approach can be disadvantageous as it often performs only one or a few inner-optimization steps, which limits its ability to fully characterize the result of the inner optimization [67]. Further, it requires caching intermediate steps, which increases computational and memory costs. BARFI, by contrast, exploits implicit gradients to alleviate these issues by directly characterizing the fixed point of Alg induced by learned behavior alignment rewards.
Finally, it is also important to highlight that a concurrent work on reward alignment using bi-level optimization was made publicly available after our manuscript was submitted for peer-reviewing at NeurIPS [10]. While our work analyses drawbacks of potential-based shaping and establishes different forms of correction that can be performed via bi-level optimization, this concurrent work provides complementary analyses on the convergence rates of bi-level optimization, as well as a discussion on its potential applications to Reinforcement Learning from Human Feedback (RLHF).