Title: Behavior Alignment via Reward Function Optimization
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement Learning

Behavior Alignment

Implicit Gradient

Bi-level Optimization
Abstract: Designing reward functions for efficiently guiding reinforcement learning (RL) agents toward specific behaviors is a complex task.
This is challenging since it requires the identification of reward structures that are not sparse and that avoid inadvertently inducing undesirable behaviors. Naively modifying the reward structure to offer denser and more frequent feedback can lead to unintended outcomes and promote behaviors that are not aligned with the designer's intended goal. Although potential-based reward shaping is often suggested as a remedy, we systematically investigate settings where deploying it often significantly impairs performance. To address these issues, we introduce a new framework that uses a bi-level objective to learn \emph{behavior alignment reward functions}. These functions integrate auxiliary rewards reflecting a designer's heuristics and domain knowledge with the environment's primary rewards. Our approach automatically determines the most effective way to blend these types of feedback, thereby enhancing robustness against heuristic reward misspecification. Remarkably, it can also adapt an agent's policy optimization process to mitigate suboptimalities resulting from limitations and biases inherent in the underlying RL algorithms. We evaluate our method's efficacy on a diverse set of tasks, from small-scale experiments to high-dimensional control challenges. We investigate heuristic auxiliary rewards of varying quality---some of which are beneficial and others detrimental to the learning process. Our results show that our framework offers a robust and principled way to integrate designer-specified heuristics. It not only addresses key shortcomings of existing approaches but also consistently leads to high-performing solutions, even when given misaligned or poorly-specified auxiliary reward functions.

Behavior Alignment via Reward Function Optimization

Abstract

Designing reward functions for efficiently guiding reinforcement learning (RL) agents toward specific behaviors is a complex task. This is challenging since it requires the identification of reward structures that are not sparse and that avoid inadvertently inducing undesirable behaviors. Naively modifying the reward structure to offer denser and more frequent feedback can lead to unintended outcomes and promote behaviors that are not aligned with the designer's intended goal. Although potential-based reward shaping is often suggested as a remedy, we systematically investigate settings where deploying it often significantly impairs performance. To address these issues, we introduce a new framework that uses a bi-level objective to learn behavior alignment reward functions . These functions integrate auxiliary rewards reflecting a designer's heuristics and domain knowledge with the environment's primary rewards. Our approach automatically determines the most effective way to blend these types of feedback, thereby enhancing robustness against heuristic reward misspecification. Remarkably, it can also adapt an agent's policy optimization process to mitigate suboptimalities resulting from limitations and biases inherent in the underlying RL algorithms. We evaluate our method's efficacy on a diverse set of tasks, from small-scale experiments to high-dimensional control challenges. We investigate heuristic auxiliary rewards of varying qualitysome of which are beneficial and others detrimental to the learning process. Our results show that our framework offers a robust and principled way to integrate designer-specified heuristics. It not only addresses key shortcomings of existing approaches but also consistently leads to high-performing solutions, even when given misaligned or poorly-specified auxiliary reward functions.

1 Introduction

In this paper, we investigate the challenge of enabling reinforcement learning (RL) practitioners, who may not be experts in the field, to incorporate domain knowledge through heuristic auxiliary reward functions. Our goal is to ensure that such auxiliary rewards not only induce behaviors that align with the designer's intentions but also allow for faster learning. RL practitioners typically model a given control problem by first designing simple reward functions that directly quantify whether (or how well) an agent completed a task. These could be, for instance, functions assigning a reward of +1 iff the agent reaches a specified goal state, and zero otherwise. However, optimizing a policy based on such a sparse reward function often proves challenging.

B

A

B

A

B

To address this issue, designers often introduce auxiliary reward functions that supplement the original rewards. Auxiliary rewards are heuristic guidelines aimed at facilitating and speeding up the learning process. One could, e.g., augment the previously described reward function (which gives a reward of +1 upon reaching a goal state) with an auxiliary reward accounting for the agent's distance to the goal. However, the effectiveness of using auxiliary reward functions largely depends on the problem's complexity and the designer's skill in crafting heuristics that, when combined with the original reward function, do not induce behaviors different than the ones originally intended [26, 27].

Existing methods like potential-based reward shaping [41] aim to incorporate domain knowledge without misaligning the behaviors induced by the resulting combined reward functions. However, as we discuss in Section 3, potential-based shaping has several limitations: (i) it is restricted to state-based functions; (ii) it amounts to a different initialization of the q -function; (iii) it does not alter policy gradients in expectation; and (iv) it can increase the variance in policy gradient methods.

To address these challenges, we introduce a scalable algorithm that empowers RL practitioners to specify potentially imperfect auxiliary reward functions. It ensures that the resulting optimization process will not inadvertently lead to unintended behaviors and that it will allow for faster learning. In particular, this paper addresses the following challenges:

(1) How to incorporate auxiliary reward information: We introduce a novel bi-level objective to analyze and automatically fine-tune designer-created auxiliary reward functions. It ensures they remain aligned with the original reward and do not induce behaviors different from those originally intended by the designer. Additionally, we formulate the problem to shape the optimization landscape, biasing our bi-level optimizer toward auxiliary reward functions that facilitate faster learning.

(2) How to use auxiliary reward to mitigate algorithmic biases: We show that our framework can automatically adjust how primary and auxiliary rewards are blended to mitigate limitations or biases inherent in the underlying RL algorithm (Section 4.1). For instance, many policy-gradient-based RL algorithms are subject to biases due to issues like discounting mismatch [59] or partial off-policy correction [51]. These biases can hinder the algorithm's ability to identify near-optimal policies.

(3) How to ensure scalability to high-dimensional problems: We introduce an algorithm that employs implicit gradients to automatically adjust primary and auxiliary rewards, ensuring that the combined reward function aligns with the designer's original expectations (see Figure 1). We evaluate our method's efficacy across a range of tasks, from small-scale to high-dimensional control settings (see Section 6). In these tasks, we experiment with auxiliary rewards of varying quality; some accelerate learning, while others can be detrimental to finding an optimal policy.

2 Notation

In this paper, we investigate sequential decision-making problems modeled as Markov decision processes (MDPs). An MDP is defined as a tuple ( S , A , p, r p , r aux , γ, d 0 ) , where S is the state set, A is the action set, p is the transition function, r p : S × A → R is the primary reward function, r aux :

S × A → R is an optional auxiliary reward function (possibly designed by a non-expert in machine learning, based on domain knowledge), and d 0 is the starting state distribution. Let π θ : S×A → [0 , 1] be any policy parameterized using θ ∈ Θ . For brevity, we will often use π θ and θ interchangeably. Let S t and A t be the random variables for the state and action observed at the time t . As in the standard RL setting, the performance J ( θ ) of a policy π θ is defined as the expected discounted return with respect to the (primary) reward function, r p ; i.e., J ( θ ) := E π [ ∑ T t =0 γ t r p ( S t , A t )] , where T +1 is the episode length. An optimal policy parameter θ ∗ is defined as θ ∗ ∈ arg max θ ∈ Θ J ( θ ) . A popular technique to search for θ ∗ is based on constructing sample estimates ˆ ∆( θ, r p ) of the ( γ -dropped) policy gradient, ∆( θ, r p ) , given an agent's interactions with the environment for one episode [58, 59]. Then, using ψ θ ( s, a ) as a shorthand for dln π θ ( s, a ) / d θ , these quantities are defined as follows:

3 Limitations of Potential Based Reward Shaping

When the objective function J ( θ ) is defined with respect to a sparse reward function r p (i.e., a reward function such that r p ( s, a ):= 0 for most s ∈ S and a ∈ A ), searching for θ ∗ is challenging [24]. A natural way to provide more frequent (i.e., denser) feedback to the agent, in the hope of facilitating learning, is to consider an alternate reward function, ˜ r naive := r p + r aux . However, as discussed earlier, r aux may be a designer-specified auxiliary reward function not perfectly aligned with the objective encoded in r p . In this case, using ˜ r naive may encourage undesired behavior. An alternative way to incorporate domain knowledge to facilitate learning was introduced by Ng et al. [41]. They proposed using a potential function , Φ : S → R (analogous to r aux ), to define new reward functions of the form ˜ r Φ ( S t , A t , S t +1 ):= r p ( S t , A t )+ γ Φ( S t +1 ) -Φ( S t ) . Importantly, they showed that optimal policies with respect to the objective E [ ∑ T t =0 γ t ˜ r Φ ( S t , A t , S t +1 )] are also optimal with respect to J ( θ ) .

While potential-based reward shaping can partially alleviate some of the difficulties arising from sparse rewards, Wiewiora [64] showed that q -learning using ˜ r Φ produces the exact same sequence of updates as q -learning using r p but with a different initialization of q -values. In what follows, we establish a similar result: we show that performing potential-based reward shaping has no impact on expected policy gradient updates; and that it can, in fact, even increase the variance of the updates.

All proofs are deferred to Appendix A. The above points highlight some of the limitations of potentialbased shaping for policy gradients and q -learning-both of which form the backbone of the majority of model-free RL algorithms [58]. Furthermore, potential functions Φ cannot depend on actions [41], which restricts the class of eligible auxiliary rewards r aux and heuristics functions that may be used. Finally, notice that Φ is designed independently of the agent's underlying learning algorithm. As we will show in the next sections, our method can autonomously discover auxiliary reward functions that not only facilitate learning but also help mitigate various types of algorithmic limitations and biases.

4 Behavior Alignment Reward Function

In this section, we introduce an objective function designed to tackle the primary challenge investigated in this paper: how to effectively leverage designer-specified auxiliary reward functions to rapidly induce behaviors envisioned by the designer. The key observation is that naively adding an auxiliary reward function r aux to r p may produce policies whose corresponding behaviors are misaligned with respect to the behaviors induced by r p . In these cases, r aux should be ignored during the search for an optimal policy. On the other hand, if r p and r aux may be combined in a way that results in the desired behaviors, then combinations that produce frequent and informative feedback to the agent should be favored, as they are likely to facilitate faster learning.

To tackle the challenges discussed above, we employ a bi-level optimization procedure. This approach aims to create a behavior alignment reward by combining r aux and r p using a parameterized function. Our method is inspired by the optimal rewards framework by Singh et al. [52, 53]. Let γ ϕ ∈ [0 , 1) be

a discount rate value parameterized by ϕ ∈ Γ . 2 Let r φ : S ×A → R be a behavior alignment reward : a function of both r p and r aux , parameterized by φ ∈ Υ , where Υ and Γ are function classes. One example of a behavior alignment reward function is r φ ( s, a ):= f φ 1 ( s, a ) + φ 2 r p ( s, a ) + φ 3 r aux ( s, a ) , where f φ : S × A → R and φ :=( φ 1 , φ 2 , φ 3 ) . Let Alg be any gradient/semi-gradient/non-gradientbased algorithm that outputs policy parameters. To mitigate possible divergence issues arising from certain policy optimization algorithms like DQN [60, 1], we make the following simplifying assumption, which can generally be met with appropriate regularizers and step sizes:

Assumption 1. Given r φ and γ ϕ , the algorithm Alg ( r φ , γ ϕ ) converges to a fixed point θ -which we denote as θ ( φ, ϕ ) ∈ Θ to emphasize its indirect dependence on φ and ϕ through Alg , r φ , and γ ϕ .

Given this assumption, we now specify the following bi-level objective:

Here, λ γ serves as the regularization coefficient for the value of γ ϕ , and Alg denotes a given policy optimization algorithm. Let, as an example, Alg be an on-policy gradient algorithm that uses samples to estimate the gradient ∆( θ, r p ) , as in (1). We can then define a corresponding variant of ∆( θ, r p ) that is compatible with our formulation and objective, and which uses both r φ and γ ϕ , as follows:

Notice that the bi-level formulation in (2) is composed of three key components: outer and inner objectives, and an outer regularization term. In what follows, we discuss the need for these.

Need for Outer- and Inner-Level Objectives: The outer-level objective in Equation (2) serves a critical role: it evaluates different parameterizations, denoted by φ , for the behavior alignment reward function. These parameterizations influence the induced policy θ ( φ, ϕ ) , which is evaluated using the performance metric J . Recall that this metric quantifies the alignment of a policy with the designer's primary reward function, r p . In essence, the outer-level objective seeks to optimize the behavior alignment reward function to produce policies that are effective according to r p . This design adds robustness against any misspecification of the auxiliary rewards. 3 In the inner-level optimization, by contrast, Alg identifies a policy θ ( φ, ϕ ) that is optimal or near-optimal with respect to r φ (which combines r aux through the behavior alignment reward). In the inner-level optimization , the algorithm Alg works to identify a policy θ ( φ, ϕ ) that is optimal or near-optimal in terms of r φ (which incorporates r aux via the behavior alignment reward). By employing a bi-level optimization structure, several benefits emerge. When r aux is well-crafted, r φ can exploit its detailed information to give Alg frequent/dense reward feedback, thus aiding the search for an optimal θ ∗ . Conversely, if r aux leads to sub-optimal policies, then the influence of auxiliary rewards can be modulated or decreased accordingly by the optimization process by adjusting r φ . Consider, for example, a case where the behavior alignment reward function is defined as r φ ( s, a ) := f φ 1 ( s, a ) + φ 2 r p ( s, a ) + φ 3 r aux ( s, a ) . In an adversarial setting-where the designer-specified auxiliary reward r aux may lead to undesired behavior-the bi-level optimization process has the ability to set φ 3 to 0 . This effectively allows the behavior alignment reward function r φ to exclude r aux from consideration. Such a bi-level approach to optimizing the parameters of behavior alignment reward functions can act as a safeguard against the emergence of sub-optimal behaviors due to a misaligned auxiliary reward, r aux . This design is particularly valuable because it allows the objective in (2) to leverage the potentially dense reward structure of r aux to provide frequent action evaluations when the auxiliary reward function is well-specified. At the same time, the approach maintains robustness against possible misalignments.

Need for Outer Regularization: The bi-level optimization problem (2) may have multiple optimal solutions for φ -including the trivial solution where r aux is always ignored. The goal of regularizing the outer-level objective (in the form of the term λ γ γ ϕ ) is to incorporate a prior that adds a preference for solutions, φ ∗ , that provide useful and frequent evaluative feedback to the underlying RL algorithm. In the next paragraphs, we discuss the need for such a regularizer and motivate its mathematical form. First, recall that sparse rewards can pose challenges for policy optimization. An intuitive solution

to this problem could involve biasing the optimization process towards denser behavior alignment reward functions, e.g., by penalizing for sparsity of r φ . Unfortunately, the distinction between sparse and dense rewards alone may not fully capture the nuances of what designers typically consider to be a 'good' reward function. This is the case because a reward function can be dense and still may not be informative ; e.g., a reward function that provides -1 to the agent in every non-goal state is dense but fails to provide useful feedback regarding how to reach a goal state. A better characterization of how useful (or informative) a reward function is may be constructed in terms of how instructive and instantaneous the evaluation or feedback it generates is. We consider a reward function to be instructive if it produces rewards that are well-aligned with the designer's goals. A reward function is instantaneous if its corresponding rewards are dense, rather than sparse, and are more readily indicative of the optimal action at any given state. 4 Reward functions that are both instructive and instantaneous can alleviate issues associated with settings with sparse rewards and long horizons. To bias our bi-level optimization objective towards this type of reward function, we introduce a regularizer, γ ϕ . This regularizer favors solutions that can generate policies with high performance (i.e., high expected return J with respect to r p ) even when the discount factor γ ϕ is small . To see why, first notice that this regularizer encourages behavior alignment reward functions that provide more instantaneous feedback to the agent. This has to be the case; otherwise, it would be challenging to maximize long-term reward should the optimized alignment reward function be sparse. Second, the regularizer promotes instructive alignment reward functions-i.e., functions that facilitate learning policies that maximize J . This is equally crucial: effective policies under the metric J are the ones that align well with the designer's objectives as outlined in the original reward function, r p .

4.1 Overcoming Imperfections of Policy Optimization Algorithms

The advantages of the bi-level formulation in (2) extend beyond robustness to sub-optimality from misspecified r aux . Even with a well-specified r aux , RL algorithms often face design choices, such as the bias-variance trade-off, that can induce sub-optimal solutions. Below we present examples to show how bias in the underlying RL algorithm may be mitigated by carefully optimizing r φ and γ ϕ .

4.1.1 Bias in policy gradients: Recall that the popular 'policy gradient' ∆( θ, r p ) is not, in fact, the gradient of any function, and using it in gradient methods may result in biased and sub-optimal policies [44]. However, policy gradient methods based on ∆( θ, r p ) remain vastly popular in the RL literature since they tend to be sample efficient [59]. Let ∆ γ ( θ, r p ) denote the unbiased policy gradient, where ∆ γ ( θ, r p ) := E [ ∑ T t =0 γ t ψ θ ( S t , A t ) ∑ T j = t γ j -t r p ( S j , A j )] . We can show that with a sufficiently expressive parameterization, optimized r φ and γ ϕ can effectively mimic the updates that would have resulted from using the unbiased gradient ∆ γ ( θ, r p ) , even if the underlying RL algorithm uses the biased 'gradient', ∆ on ( θ, φ, ϕ ) , as defined in (3). Detailed proofs are in Appendix A.

Property 2. There exists r φ : S × A → R and γ ϕ ∈ [0 , 1) such that ∆ on ( θ, φ, ϕ ) = ∆ γ ( θ, r p ) .

4.1.2 Off-policy learning without importance sampling: To increase sample efficiency when evaluating a given policy π θ , it is often useful to use off-policy data collected by a different policy, β . Under the assumption that ∀ s ∈ S , ∀ a ∈ A , π θ ( s,a ) β ( s,a ) < ∞ , importance ratios ρ j := ∏ j k =0 π θ ( s,a ) β ( s,a ) can be used to adjust the updates and account for the distribution shift between trajectories generated by β and π θ . However, to avoid the high variance stemming from ρ j , many methods tend to drop most of the importance ratios and thus only partially correct for the distribution shift--which can lead to bias [51]. We can show (given a sufficiently expressive parameterization for the behavior alignment reward function) that this type of bias can also be mitigated by carefully optimizing r φ and γ ϕ .

Let us denote the unbiased off-policy update with full-distribution correction as ∆ off ( θ, r p ) := E β [ ∑ T t =0 γ t ψ θ ( S t , A t ) ∑ T j = t ρ j γ 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 ∆ off ( θ, φ, ϕ ) := E β [ ∑ T t =0 ψ θ ( S t , A t ) ∑ T j = t γ j -t ϕ r φ ( S j , A j )] .

Property 3. There exists r φ : S × A → R and γ ϕ ∈ [0 , 1) such that ∆ off ( θ, φ, ϕ ) = ∆ off ( θ, 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

to the γ ϕ ∗ 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 φ and γ ϕ , an important question arises: Although θ ( φ, ϕ ) can be optimized using any policy learning algorithm, how can we efficiently identify the optimal φ ∗ and ϕ ∗ in equation (2)? Given the practical advantages of gradient-based methods, one would naturally consider using them for optimizing φ ∗ and ϕ ∗ as well. However, a key challenge in our setting lies in computing d J ( θ ( φ, ϕ )) / d φ and d J ( θ ( φ, ϕ )) / d ϕ . These computations require an analytical characterization of the impact that r φ and γ ϕ have on the entire optimization process of the inner-level algorithm, Alg .

In addressing this challenge, we initially focus on an Alg that employs policy gradients for updating π θ . Similar extensions for other update rules can be derived similarly. We start by re-writing the expression for d J ( θ ( φ, ϕ )) / d φ using the chain rule:

where (a) is the policy gradient at θ ( φ, ϕ ) , and (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 Alg converges to θ ( φ, ϕ ) , then it follows that

Let ∂f denote the partial derivative with respect to the immediate arguments of f , and d f be the total derivative as before. That is, if f ( x, g ( x )) := xg ( x ) , then ∂f ∂x ( x, g ( x )) = g ( x ) and d f d x ( x, g ( x )) = g ( x ) + x ∂g ∂x ( x ) . Therefore, taking the total derivative of (5) with respect to φ yields

By re-arranging terms in (6), we obtain the term (b) in (4). In particular,

Furthermore, by combining (7) and (4) we obtain the desired gradient expression for φ :

Similarly, a gradient expression for ϕ can be derived; the full derivation is detailed in Appendix E. Using θ ∗ as shorthand for θ ( φ, ϕ ) , we find that the terms A and H can be expressed as

When working with the equations above, we assume the inverse of H -1 exists. To mitigate the risk of ill-conditioning, we discuss regularization strategies for Alg in Appendix D. Notice that equations (8) and (15) are the key elements needed to calculate the updates to φ and ϕ in our bi-level optimization's outer loop. However, computing A and 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 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 H or 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 φ and γ ϕ incrementally modify θ ( φ, ϕ ) . 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: ˜ r naive ( s, a ) := r p ( s, a ) + r 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 ˜ r Φ ( s, a, s ′ ) := r p ( s, a ) + γr aux ( s ′ ) -r aux ( s ) .

BARFI : We use REINFORCE as the underlying RL algorithm when implementing BARFI and define r φ ( s, a ) := φ 1 ( s, a ) + φ 2 ( s ) r p ( s, a ) + φ 3 ( s ) r aux ( s, a ) . Our implementation includes a warm-up period wherein the agent collects data for a fixed number of episodes, using ˜ r naive, prior to performing the first updates to φ and ϕ (See Appendix 5 for the complete algorithm).

We evaluate each algorithm across four distinct environments: GridWorld, MountainCar [58], CartPole [16], and HalfCheetahv4 [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 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

BARFI 's performance compared to two baselines that use: ˜ r naive and ˜ r Φ , respectively. CartPole uses an action-dependent r aux function that either rewards agents when actions align with a known effective policy ( well-aligned r aux ) or with a poorly-performing policy ( misaligned r 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 aux is action-dependent, we use a direct extension of the original formulation [41] by considering ˜ r Φ ( s, a, s ′ , a ′ ) := r p ( s, a ) + γr aux ( s ′ , a ′ ) -r aux ( s, a ) . Furthermore, we designed an auxiliary reward function, r 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 γ < 1 , the resultant value from γr aux ( s ′ , a ′ ) -r aux ( s, a ) is negative . Such a negative component may deter the agent from selecting actions that are otherwise optimal, depending on how r 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.

s

0

Goal

+100

+50

100

0

-20

GridWorld

Reward Center State

Learned

ϕ

3

(

s

)

CartPole

GridWorld

Episodes

Episodes

Actor Critic Potential

Reinforce Potential

Reinforce

Barfi

Actor Critic

Return

{

Warm up Period

100

500

80

400

60

300

40

200

20

100

0

0

0

500

1000

1500

2000

0

200

400

600

800

1000

Reinforce

Actor Critic

Episodes

Episodes

6.3 What does r φ 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 poorlyperforming behaviors that focus on revisiting irrelevant central states. Similarly, when Cartpole operates under a misspecified r 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 φ and ϕ would be feasible in highdimensional 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 φ and ϕ . Further details regarding this alternative method can be found in Appendix C.6.

We conducted experiments on the HalfCheetahv4 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 HalfCheetahv4 , 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 aux ( s, a ) := 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 ˜ r naive ( s, a ) = r p ( s, a ) + r 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 γ ; and (d) returns based on r φ , are provided in Appendix H.

Episodes

Episodes

Return

Mujoco

Mujoco

Actor Critic Potential

Reinforce

Barfi

Actor Critic

Barfi Unrolled

Barfi

Barfi Unrolled

Memory (MBs)

Number of Inner Updates

Compute/Memory

Barfi Unrolled

Barfi (5 Steps)

Barfi (3 Steps)

Barfi (1 Step)

2500

2500

2000

2000

1500

1500

1000

1000

500

500

0

0

0

2000

4000

6000

8000

10000

0

2000

4000

6000

8000

10000

7 Related work

This paper focuses primarily on how to efficiently leverage auxiliary rewards r aux . Notice, however, that in the absence of r aux , the resulting learned behavior alignment rewards r φ 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).

8 Conclusion and Future Work

In this paper, we introduced BARFI , a novel framework that empowers RL practitioners-who may not be experts in the field-to incorporate domain knowledge through heuristic auxiliary reward functions. Our framework allows for more expressive reward functions to be learned while ensuring they remain aligned with a designer's original intentions. BARFI can also identify reward functions that foster faster learning while mitigating various limitations and biases in underlying RL algorithms. We empirically show that BARFI is effective in training agents in sparse-reward scenarios where (possibly poorly-specified) auxiliary reward information is available. If the provided auxiliary rewards are determined to be misaligned with the designer's intended goals, BARFI autonomously adapts and effectively disincentivizes their use as needed. This adaptability results in a reliable pathway to identifying high-performing policies. The conceptual insights offered by this work provide RL practitioners with a structured way to design more robust and easy-to-optimize reward functions. We believe this will contribute to making RL more accessible to a broader audience.

References

Joshua Achiam, Ethan Knight, and Pieter Abbeel. Towards characterizing divergence in deep q-learning. arXiv preprint arXiv:1903.08894 , 2019.

Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and Zico Kolter. Differentiable convex optimization layers. arXiv preprint arXiv:1910.12430 , 2019.

Ron Amit, Ron Meir, and Kamil Ciosek. Discount factor as a regularizer in reinforcement learning. In International Conference on Machine Learning , 2020.

Brandon Amos and J Zico Kolter. Optnet: Differentiable optimization as a layer in neural networks. In International Conference on Machine Learning , 2017.

Stuart Armstrong, Jan Leike, Laurent Orseau, and Shane Legg. Pitfalls of learning a reward function online. arXiv preprint arXiv:2004.13654 , 2020.

Karl J Åström and Tore Hägglund. Pid control. IEEE Control Systems Magazine , 2006.

Sarah Bechtle, Artem Molchanov, Yevgen Chebotar, Edward Grefenstette, Ludovic Righetti, Gaurav Sukhatme, and Franziska Meier. Meta-learning via learned loss. arXiv preprint arXiv:1906.05374 , 2019.

Yoshua Bengio. Gradient-based optimization of hyperparameters. Neural computation , 2000.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540 , 2016.

Souradip Chakraborty, Amrit Singh Bedi, Alec Koppel, Dinesh Manocha, Huazheng Wang, Furong Huang, and Mengdi Wang. Aligning agent policy with externalities: Reward design via bilevel rl. arXiv preprint arXiv:2308.02585 , 2023.

Xu Chu Dennis Ding, Stephen L Smith, Calin Belta, and Daniela Rus. LTL control in uncertain environments with probabilistic satisfaction guarantees. IFAC Proceedings Volumes , 2011.

Chuong B Do, Chuan-Sheng Foo, and Andrew Y Ng. Efficient multiple hyperparameter learning for log-linear models. In Advances in Neural Information Processing Systems , 2007.

Justin Domke. Generic methods for optimization-based modeling. In Artificial Intelligence and Statistics , 2012.

Asen L Dontchev and R Tyrrell Rockafellar. Implicit functions and solution mappings , volume 543. Springer, 2009.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In , Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017 , 2017.

Razvan V Florian. Correct equations for the dynamics of the cart-pole system. Center for Cognitive and Neural Studies (Coneural), Romania , 2007.

Zhao-Yang Fu, De-Chuan Zhan, Xin-Chun Li, and Yi-Xing Lu. Automatic successive reinforcement learning with multiple auxiliary rewards. In IJCAI , 2019.

Sina Ghiassian, Banafsheh Rafiee, Yat Long Lo, and Adam White. Improving performance in reinforcement learning by breaking generalization in neural networks. In , Proceedings of the 19th International Conference on Autonomous Agents and Multiagent Systems, AAMAS '20, Auckland, New Zealand, May 9-13, 2020 , 2020.

Stephen Gould, Basura Fernando, Anoop Cherian, Peter Anderson, Rodrigo Santa Cruz, and Edison Guo. On differentiating parameterized argmin and argmax problems with application to bi-level optimization. arXiv preprint arXiv:1607.05447 , 2016.

Edward Grefenstette, Brandon Amos, Denis Yarats, Artem Molchanov, Franziska Meier, and Kyunghyun Cho. higher: A pytorch meta-learning library. 2020.

Ricardo Grunitzki, Bruno C da Silva, and Ana LC Bazzan. A flexible approach for designing optimal reward functions. In Proceedings of the 16th Conference on Autonomous Agents and MultiAgent Systems , 2017.

Ricardo Grunitzki, Bruno C da Silva, and LC Ana Bazzan. Towards designing optimal reward functions in multi-agent reinforcement learning problems. In 2018 International Joint Conference on Neural Networks (IJCNN) , 2018.

Xiaoxiao Guo, Satinder Singh, Richard Lewis, and Honglak Lee. Deep learning for reward design to improve monte carlo tree search in atari games. arXiv preprint arXiv:1604.07095 , 2016.

Joshua Hare. Dealing with sparse rewards in reinforcement learning. arXiv preprint arXiv:1910.09281 , 2019.

Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Cited on , 2012.

Mark K Ho, Michael L Littman, Fiery Cushman, and Joseph L Austerweil. Teaching with rewards and punishments: Reinforcement or communication? In CogSci , 2015.

Mark K Ho, Fiery Cushman, Michael L Littman, and Joseph L Austerweil. People teach with rewards and punishments as communication, not reinforcements. Journal of Experimental Psychology: General , 2019.

Yujing Hu, Weixun Wang, Hangtian Jia, Yixiang Wang, Yingfeng Chen, Jianye Hao, Feng Wu, and Changjie Fan. Learning to utilize shaping rewards: A new approach of reward shaping. Advances in Neural Information Processing Systems , 2020.

Rodrigo Toro Icarte, Toryn Klassen, Richard Valenzano, and Sheila McIlraith. Using reward machines for high-level task specification and decomposition in reinforcement learning. In International Conference on Machine Learning , 2018.

Rodrigo Toro Icarte, Ethan Waldie, Toryn Klassen, Rick Valenzano, Margarita Castro, and Sheila McIlraith. Learning reward machines for partially observable reinforcement learning. Advances in Neural Information Processing Systems , 2019.

Rodrigo Toro Icarte, Toryn Q Klassen, Richard Valenzano, and Sheila A McIlraith. Reward machines: Exploiting reward function structure in reinforcement learning. arXiv preprint arXiv:2010.03950 , 2020.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In , 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings , 2015.

Louis Kirsch, Sjoerd van Steenkiste, and Jürgen Schmidhuber. Improving generalization in meta reinforcement learning using learned objectives. arXiv preprint arXiv:1910.04098 , 2019.

Steven G Krantz and Harold R Parks. The implicit function theorem: history, theory, and applications . Springer Science & Business Media, 2012.

Hadas Kress-Gazit, Georgios E Fainekos, and George J Pappas. Temporal-logic-based reactive mission and motion planning. IEEE Transactions on Robotics , 2009.

Karl Kunisch and Thomas Pock. A bilevel optimization approach for parameter learning in variational models. SIAM Journal on Imaging Sciences , 2013.

Jan Larsen, Lars Kai Hansen, Claus Svarer, and M Ohlsson. Design and regularization of neural networks: the optimal use of a validation set. In Neural Networks for Signal Processing VI. Proceedings of the 1996 IEEE Signal Processing Society Workshop , 1996.

Kwonjoon Lee, Subhransu Maji, Avinash Ravichandran, and Stefano Soatto. Meta-learning with differentiable convex optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2019.

Michael L Littman, Ufuk Topcu, Jie Fu, Charles Isbell, Min Wen, and James MacGlashan. Environment-independent task specifications via GLTL. arXiv preprint arXiv:1704.04341 , 2017.

Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics , 2020.

Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In ICML , 1999.

Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. CoRR , 2018.

Scott Niekum, Andrew G Barto, and Lee Spector. Genetic programming for reward function search. IEEE Transactions on Autonomous Mental Development , 2010.

Chris Nota and Philip S. Thomas. Is the policy gradient a gradient? arXiv preprint arXiv:1906.07073 , 2020.

Junhyuk Oh, Matteo Hessel, Wojciech M Czarnecki, Zhongwen Xu, Hado van Hasselt, Satinder Singh, and David Silver. Discovering reinforcement learning algorithms. arXiv preprint arXiv:2007.08794 , 2020.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32 . Curran Associates, Inc., 2019.

Barak A Pearlmutter. Fast exact multiplication by the hessian. Neural computation , 1994.

Silviu Pitis. Rethinking the discount factor in reinforcement learning: A decision theoretic approach. In Proceedings of the AAAI Conference on Artificial Intelligence , 2019.

Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. In Advances in Neural Information Processing Systems , 2019.

Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar. A game theoretic framework for model based reinforcement learning. arXiv preprint arXiv:2004.07804 , 2020.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Satinder Singh, Richard L Lewis, and Andrew G Barto. Where do rewards come from. In Proceedings of the Annual Conference of the Cognitive Science Society , 2009.

Satinder Singh, Richard L Lewis, Andrew G Barto, and Jonathan Sorg. Intrinsically motivated reinforcement learning: An evolutionary perspective. IEEE Transactions on Autonomous Mental Development , 2010.

Jonathan Sorg, Richard L Lewis, and Satinder Singh. Reward design via online gradient ascent. Advances in Neural Information Processing Systems , 2010.

Jonathan Sorg, Satinder P Singh, and Richard L Lewis. Internal rewards mitigate agent boundedness. In Proceedings of the 27th international conference on machine learning (ICML10) , 2010.

Jonathan Sorg, Satinder P Singh, and Richard L Lewis. Optimal rewards versus leaf-evaluation heuristics in planning agents. In Proceedings of the AAAI Conference on Artificial Intelligence , 2011.

Jonathan Daniel Sorg. The Optimal Reward Problem: Designing Effective Reward for Bounded Agents. PhD thesis, University of Michigan, 2011.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Philip Thomas. Bias in natural actor-critic algorithms. In International Conference on Machine Learning , 2014.

John N Tsitsiklis and Benjamin Van Roy. An analysis of temporal-difference learning with function approximation. IEEE Transactions on Automatic Control , 1997.

Marin Vlastelica, Anselm Paulus, Vít Musil, Georg Martius, and Michal Rolínek. Differentiation of blackbox combinatorial solvers. arXiv preprint arXiv:1912.02175 , 2019.

Yufei Wang, Qiwei Ye, and Tie-Yan Liu. Beyond exponentially discounted sum: Automatic learning of return function. arXiv preprint arXiv:1905.11591 , 2019.

Martha White. Unifying task specification in reinforcement learning. In International Conference on Machine Learning , 2017.

Eric Wiewiora. Potential-based shaping and q-value initialization are equivalent. Journal of Artificial Intelligence Research , 2003.

Eric M Wolff, Ufuk Topcu, and Richard M Murray. Robust control of uncertain markov decision processes with temporal logic specifications. In 2012 IEEE 51st IEEE Conference on Decision and Control (CDC) , 2012.

David H Wolpert and Kagan Tumer. Optimal reward functions in distributed reinforcement learning. In Intelligent agent technology: Research and development . World Scientific, 2001.

Yuhuai Wu, Mengye Ren, Renjie Liao, and Roger Grosse. Understanding short-horizon bias in stochastic meta-optimization. arXiv preprint arXiv:1803.02021 , 2018.

Zhongwen Xu, Hado P van Hasselt, and David Silver. Meta-gradient reinforcement learning. In Advances in Neural Information Processing Systems , 2018.

Zhongwen Xu, Hado P van Hasselt, Matteo Hessel, Junhyuk Oh, Satinder Singh, and David Silver. Meta-gradient reinforcement learning with an objective discovered online. Advances in Neural Information Processing Systems , 2020.

Zeyu Zheng, Junhyuk Oh, and Satinder Singh. On learning intrinsic rewards for policy gradient methods. In Advances in Neural Information Processing Systems , 2018.

Zeyu Zheng, Junhyuk Oh, Matteo Hessel, Zhongwen Xu, Manuel Kroiss, Hado Van Hasselt, David Silver, and Satinder Singh. What can learned intrinsic rewards capture? In International Conference on Machine Learning , 2020.

Haosheng Zou, Tongzheng Ren, Dong Yan, Hang Su, and Jun Zhu. Reward shaping via meta-learning. arXiv preprint arXiv:1901.09330 , 2019.

Behavior Alignment via Reward Function Optimization (Supplemental Material)

A Proofs for Theoretical Results

In this section, we provide proofs for Property 1 , Property 2 , and Property 3 . For the purpose of these proofs, we introduce some additional notation. To have a unified MDP notation for goal-based and time-based tasks, we first consider that in the the time-based task, time is a part of the state such that Markovian dynamics is ensured.

The (un-normalized) discounted and (un-normalized) undiscounted visitation probability is denoted as

We can normalize it so that it is a distribution as follows :

Property 1. The expected update performed by the biased policy gradient update is same when using the primary reward and reward modified with potential-based shaping, i.e., ∆( θ, ˜ r ) = ∆( θ, r p ) . Further, the variance of the update when using potential-based reward shaping can be higher than the variance of the update performed using the primary reward, i.e., Var ( ˆ ∆( θ, ˜ r ) ) ≥ Var ( ˆ ∆( θ, r p ) ) .

Proof. Part 1: Equality of the expected update

where (a) holds because on the expansion of future return, intermediate potential values cancel out, (b) holds because S T +1 is the terminal state and potential function is defined to be a fixed constant c for any terminal state [41], (c) holds from the law of total expectation, and (d) holds because,

In the stochastic setting, i.e., when using sample average estimates instead of the true expectation, γ T -t +1 c -φ ( S t ) is analogous to a state-dependent baseline for the sum of discounted future primary rewards. It may reduce or increase the variance of ∆( θ, r p ) , depending on this baseline's co-variance with ∑ T j = t γ j -t r p ( S j , A j ) .

Note: As we encountered the potential at the terminal state to be c , as it is a constant, we will use the value of c = 0 in accordance with [41].

Part 2: Variance characterization

For this result that discusses the possibility of the variance being higher when using potential-based reward shaping, we demonstrate the result using a simple example. We will consider the single-step case wherein an episode lasts for one time step. That is, the agent takes an action A 0 at the starting state S 0 and then transitions to the terminal state. Hence, the stochastic update ˆ ∆( θ, ˜ r ) can be written as:

wherein, we assume that Φ for terminal states is 0 and similarly, ˆ ∆( θ, r p ) = ψ θ ( S 0 , A 0 ) r p ( S 0 , A 0 ) . For the purpose of this proof we will consider the case wherein we have a scalar θ , i.e., θ ∈ R , such that, ψ θ ( ., . ) ∈ R .

Hence, Var ( ˆ ∆( θ, ˜ r ) ) can be written as:

Therefore,

Subsequently, variance of ˆ ∆( θ, ˜ r ) will be higher than that of ˆ ∆( θ, r p ) if E [ ψ θ ( S 0 , A 0 ) 2 Φ( S 0 ) 2 ] -2 E [ ψ θ ( S 0 , A 0 ) 2 Φ( S 0 ) r p ( S 0 , A 0 ) ] > 0 .

Example: Let us look at an example where the above condition can be true. Let us consider an MDP with a single state and a single-step horizon. In that case, we can consider the variance of the update to the policy at the said state, i.e.,

where s is the fixed state. Hence, the variance of the potential-based method might be more than the variance from using only the primary reward when

glyph[negationslash]

Further, let us consider the case where Φ( s ) = 0 , because otherwise the variance of the update for those states would be same, and Φ( s ) > 0 .

We can see that the above condition can be satisfied by choosing a potential function that might be overly optimistic about the average reward of the state s , i.e. any Φ( s ) , s.t. Φ( s ) > 2 r p ( s, a ) ∀ a would lead to an increase in variance. A common place where this might be true is the use of an optimal value function (as hinted by [41]) as a baseline for a bad/mediocre policy initially.

Property 2. There exists r φ : S × A → R and γ ϕ ∈ [0 , 1) such that ∆ on ( θ, φ, ϕ ) = ∆ γ ( θ, r p ) .

Proof. Recall the definition of ∆ γ ( θ, r p ) from Section 4.1:

Using the law of total expectation,

Notice from (9) and (10) that for any ( s, a ) pair, if d π θ γ ( s, a ) > 0 , then ¯ d π θ ( s, a ) > 0 since γ ≥ 0 . Therefore, dividing and multiplying by ¯ d π θ ( s, a ) leads to:

Now, notice that if γ ϕ = 0 and r φ ( s, a ) = q π θ ( s, a ) d π θ γ ( s,a ) ¯ d π θ ( s,a ) , for all s ∈ S and a ∈ A , then

Property 3. There exists r φ : S × A → R and γ ϕ ∈ [0 , 1) such that ∆ off ( θ, φ, ϕ ) = ∆ off ( θ, r p ) .

Proof. This proof follows a similar technique as the proof for Property 2 . Recall the definition of ∆ off ( θ, r p ) :

Now using the law of total expectations,

where the last line follows similar to (11). Now, notice that for any ( s, a ) pair, the assumption that π θ ( s, a ) /β ( s, a ) < ∞ for all s ∈ S , a ∈ A , implies d π θ γ ( s, a ) /d β γ ( s, a ) < ∞ . Further, if d β γ ( s, a ) > 0 it has to be that d β ( s, a ) > 0 as well. Therefore, d π θ γ ( s, a ) /d β ( s, a ) < ∞ as well. Multiplying and dividing by d β ( s, a ) results in:

Finally, notice that if γ ϕ = 0 and r φ ( s, a ) = q π θ ( s, a ) d π θ γ ( s,a ) ¯ d β ( s,a ) for all s ∈ S and a ∈ A ,

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 aux , while staying robust against its misspecification.

Apart from specifying auxiliary rewards r 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 γ that is state-action dependent [63, 48], or establishes connection between γ 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 ∈ R d , y ∈ R m , f ( x, y ) ∈ R . Then, we know that ∂f ( x, y ) /∂x ∈ R d , ∂f ( x, y ) /∂y ∈ R m , ∂ 2 f ( x, y ) /∂y∂x ∈ R d × m . Let us also assume that we have a vector v ∈ R d , and if we need to calculate the following, we can pull the derivative outside as shown:

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 ) ∈ R 1 , x ∈ R d , y ∈ R m , v ∈ R d

2 f ′ ← grad ( f ( x, y ) , x )

3 jvp ← grad ( f ′ , y, grad_outputs = v )

4 Return: jvp

C.2 Neumann Series Approximation for Hessian Inverse

Recall, that for a given real number β ∈ R , such that 0 ≤ β < 1 , we know that the geomertric series of this has a closed form solution, i.e.,

Similarly, given we have a value α such that β = 1 -α , we can write α -1 as follows:

The same can be generalized for a matrix, i.e., given a matrix A ∈ R d × d , we can write A -1 as follows:

Note for the above to hold, matrix A , where we represent eig ( A ) as the eigenvalues of matrix A , we should have the following condition to hold, 0 < eig ( A ) < 1 . Note here we would regularize A to ensure that all eigenvalues are positive, and then we can always scale the matrix A , by its biggest eigenvalue to ensure that the above condition holds. Let say η = 1 / max eig ( A ) . Then we can write the following:

As η 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, lets look what needs to be done for our updates. Recall that the update φ, ϕ (8) and (15) were,

and

Let us look closely at the update for φ and we can generalize the updates easily for the case of ϕ .

We first look at how we can approximate the value of v H -1 efficiently, as we can always make use of the Jacobian Vector product later to get ( v H -1 ) A , as v 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 H -1 up to n order Neumann expansion,

Here we are assuming that the outer optimization for update (1) is for the function J ( θ ( φ, ϕ )) and the inner optimization which is represented by the update (3) is f ( θ ( φ, ϕ ) , φ, ϕ ) , i.e.,

The most common form in which f (; B ) is usually defined is the following:

Similarly this can be defined for J , except making use of r p and problem defined γ :

Finally, once we have v H -1 , we can use the Vector Jacobian Product to calculate ( v H -1 ) A as described in Algorithm 3:

Algorithm 3: Update for φ , i.e. (8) i.e., v H -1 A

We can similarly derive updates for ϕ . Note we are not including the different forms of regularizers

Algorithm 4: Update for φ , i.e. (15) i.e., v H -1 B 1 Input: θ, φ, ϕ, J, f, n, η, D off , D on 2 v ← Algorithm 2 ( θ, φ, ϕ, J, f, n, η, D off , D on) 3 v ′ ← grad ( f ( θ, φ, ϕ ) , D off ) , θ ) 4 ∆ ϕ ← grad ( v ′ , ϕ, grad_outputs = v ) 5 Return ∆ ϕ

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 φ, ϕ , and from there onwards after each update of outer optimization, we consequently update the policy in (21-23). The flow of the algorithm is show in Figure 4.

As discussed in Section D, using regularizers in ∆( θ, φ, ϕ ) smoothens the objective J ( θ ( φ, ϕ )) with respect to φ and ϕ . This is helpful as gradual changes in r φ an γ ϕ 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 θ 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

Line 8-10

Line 16-17

Line 16-17

Line 21-23

Line 21-23

In Lines 13-14, a new batch D on of data is sampled using the policy returned by the inner-optimization process. This data is used to compute ∂J ( θ ( φ, ϕ )) /∂θ ( φ, ϕ ) . Existing data D off that was used in the inner-optimization process is then used to compute ∂θ ( φ, ϕ ) /∂φ and ∂θ ( φ, ϕ ) /∂ϕ . Using these in (8) and (15), the parameters for r φ and γ ϕ are updated in Lines 16 and 17 , respectively.

Finally, the new data D on is merged into the existing data D off and the entire process continues.

C.5 Note on Approximation

An important limitation of the methods discussed above is that θ ( φ, ϕ ) is considered such that ∆( θ ( φ, ϕ ) , φ, ϕ ) = 0 , i.e., the Alg is run to convergence. In practice, we only execute 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 glyph[epsilon1] -neighborhood of the optima [49]. Furthermore, due to smoothness in the functional space, slight

changes to φ and ϕ should result in slight shifts in the optimum θ ( φ, ϕ ) . 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, ( φ, ϕ ) , 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 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 θ ( φ, ϕ ) is reached-where the sequence of steps depends on the specific 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 ( θ ( φ, ϕ )) 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 φ, ϕ 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 6 . We discuss two ways to regularize the objective in the upcoming sections.

D.1 L2 Regularization

To understand how severely ill-conditioned J ( θ ( φ, ϕ )) 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 Alg . Therefore, if the parameters of the behavior alignment reward are perturbed from φ to φ ′ , it may often be that J ( θ ( φ, ϕ )) = J ( θ ( φ ′ , ϕ )) and this limits any gradient based optimization for φ as ∂J ( θ ( φ, ϕ )) /∂φ is 0 . Similarly, minor perturbations in ϕ may result in no change in J ( θ ( φ, ϕ )) either.

Fortunately, there exists a remarkably simple solution: incorporate regularization for the policy parameters θ in objective for Alg in the inner-level optimization. For example, the optimal policy for the following regularized objective E π θ [ ∑ T t =0 γ t ϕ r φ ( S t , A t )] -λ 2 ‖ θ ‖ 2 varies smoothly to trade-off between the regularization value of θ and the magnitude of the performance characterized by ( r φ , γ ϕ ) , which changes with the values of r φ and γ ϕ . See Figure 5 for an example with L2 regularization.

D.2 Entropy Regularized

In Section D.1, smoothing of J ( θ ( φ, ϕ )) was done by using L2 regularization on the policy parameters θ 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 θ based on entropy regularization,

A

B

C

r

[

A

]

3

2

1

0

1

2

r

[

C

]

3

2

1

0

1

2

J

(

(

, 1))

1.25

1.50

1.75

2.00

2.25

2.50

2.75

r

[

A

]

3

2

1

0

1

2

r

[

C

]

3

2

1

0

1

2

J

(

(

, 1))

1.6

1.8

2.0

2.2

2.4

2.6

2.8

Notice that new update rule for φ and ϕ can be obtained from steps (4) to (15) with the following A , B , and H instead, where for shorthand θ ∗ = θ ( φ, ϕ ) ,

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-optimalities.

Taking total derivative in (5) with respect to φ ,

Let us try to understand why the above is true, considering the finite difference approach for this derivative,

∆( θ ( φ + dφ, ϕ ) , φ + dφ, ϕ ) = ∆( θ ( φ, ϕ ) , φ, ϕ ) = 0 , as θ ( · , · ) defines convergence to fixed point. By re-arranging terms in (13) we obtain the term (b) in (4),

By combining (14) with (4) we obtain the desired gradient expression for φ ,

and following similar steps, it can be observed that the gradient expression for ϕ ,

where using θ ∗ as a shorthand for θ ( φ, ϕ ) the terms A , B and H can be expressed as

These provide the necessary expressions for updating φ and ϕ in the outer loop. As A involves an outer product and 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 ( 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 Alg to θ ( φ, ϕ ) and then updating r φ and γ ϕ . Importantly, being based on (3) for sample efficiency, 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 φ and γ ϕ changes the policy resulting from Alg gradually. Therefore, for compute efficiency, we start 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 HalfCheetahv4 from Mujoco (MJ) suite of OpenAI Gym [9].

For each environment, we define two auxiliarly reward functions. For GridWorld, we define the functions: r 1 aux , GW := -( s -s goal ) 2 , which provides the negative L2 squared distance from the goal position, and r 1 aux , GW := 50 × 1 s ∈S 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 1 aux , MC ( s, a ) := | v | , encourages a higher absolute velocity of the car, and the second, r 1 aux , MC ( s, a ) := 1 sign(v) = a , encourages the direction of motion to increase the magnitude of the velocity (also knows 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, PD ∗ : S → 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 1 aux , CP ( s, a ) := 5 × 1 PD ∗ ( s )= a -(1 -1 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 1 aux , GW ( s, a ) := -r 1 aux , CP ( 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 1 aux , GW , r 1 aux , GW , r 1 aux , MC ). Reward functions that encode a guess of an optimal policy (i.e., r 1 aux , MC , r 1 aux , CP ) and reward functions that change the optimal policy (i.e., r 1 aux , GW , r 1 aux , CP ). We also have rewards that only depend on states (i.e., r 1 aux , GW , r 1 aux , GW , r 1 aux , MC ) as well as ones that depend on both state and actions (i.e., r 1 aux , MC , r 1 aux , CP , r 1 aux , CP ). 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 HalfCheetahv4 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 aux ( s, a ) := 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 ˜ r ( s, a ) = r p ( s, a ) + r 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 1 aux , MJ ( s, a ) := -c | a | 2 2 , uses the default setting, while the second weighting, denoted as r 1 aux , MJ ( s, a ) := -4 c | 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 γ 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 σ = 0 . 1 .

General Details: All the outer returns are evaluated without any discounting, whereas all the inner optimizations were initialized with γ ϕ = 0 . 99 . Hence to do this we made ϕ a single bias unit, initialized to 4 . 6 , and passed through a sigmoid (i.e., σ (4 . 6) = 0 . 99 ).

For GW, CP and MC r φ is defined as below

Wherein φ 1 , φ 2 , φ 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

Wherein φ 1 is initialized to zero and φ 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

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 α θ ∈ { 5 ×

10 -3 , 2 . 5 × 10 -3 , 1 × 10 -3 , 5 × 10 -4 , 2 . 5 × 10 -4 , 1 × 10 -4 , 7 . 5 × 10 -5 } , we use similar ranges for α φ , α ϕ (which tend to be larger). For λ θ and λ φ , we sweeped from [0 , 0 . 25 , 0 . 5 , 1 . 0] and for λ γ we sweeped 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 BARFI across different methods as we can see from the tables above. δ ∈ [1 , 3 , 5] , n ∈ [1 , 3 , 5] , N i ∈ [1 , 3 , 6 , 9 , 12 , 15] , η ∈ [1 × 10 -3 , 5 × 10 -4 , 1 × 10 -4 ] , N 0 ∈ [30 , 50 , 100 , 150] , buffer ∈ [25 , 50 , 100 , 1000] . α for Tilecoding was adopted from [58] and hence similar ranges were sweeped 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 . 60 GHz. 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.

H Extra Results & Ablations

H.1 Experiment on partially misspecified r 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 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.

r

aux

(

s

)

0

10

20

30

40

50

60

70

(a)

r

aux

(

s

)

3(

s

)

0.30

0.25

0.20

0.15

0.10

0.05

0.00

0.05

0.10

(b) Learned

φ

3

(

s

)

, the weighting on

r

aux

(

s

)

r

(

s

)

10 1

10 0

0

10 0

10 1

(c)

The

net

reward

induced

i.e.,

r

p

(

s

) +

φ

3

(

s

)

r

aux

(

s

)

Barfi

Naive Addition

(d) Learning performance

Learning Curve

40

Return

20

0

0

200

400

600

800

1000

Episodes

MountainCar

Mujoco

CartPole

GridWorld

GridWorld

MountainCar

Mujoco

CartPole

Episodes

Aux Return

Aux Return

Episodes

Episodes

Episodes

(a)

(b)

(c)

(d)

20

800

100

800

80

600

20

0

600

400

60

40

20

200

400

-60

40

200

0

0

0

-80

-200

0

500

1000

1500

2000

0

200

400

600

800

1000

0

100

200

300

400

500

0

2000

4000

6000

8000

10000

600-

30

500

1500

200

400

20

1000

0

300

10

200

500

-200

100

0

0

0

400

0

500

1000

1500

2000

0

200

400

600

800

1000

0

500

1000

1500

20

0

2000

4000

6000

8000

10000

H.2 Return based on learned r φ and γ ϕ

Figure 7 and Figure 8 summarize the achievable return based on r φ and the γ 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 φ , but is also able to successively decay γ 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 η = 0 . 001 , 0 . 0005 , n = 5 works for the approximation.

MountainCar

MountainCar

GridWorld (Bad)

GridWorld

GridWorld

CartPole

CartPole (Good)

CartPole

Episodes

Episodes

Episodes

(a)

(b)

(c)

γ

γ

1.0

0.990

1.0

0.9

0.8

0.989

0.6

0.7

0.988

0.4

0.6

0.987

0.2

0.5

0.986

0.0

0

500

2000

0

200

400

600

800

1000

0

100

200

300

400

500

1.0

066°0

1.0

0.9

0.988

0.8

0.6

0.986

0.4

0.7

0.984

0.2

0.6

0.982

0.0

0

500

1000

1500

2000

0

200

400

600

800

1000

0

500

1000

1500

2000

2

6

2

4

2

2

2 0

2 2

2 4

91

92

93

94

95

96

97

98

Returns

Returns

Returns

2

6

2

4

2

2

2 0

2 2

2 4

91

92

93

94

95

96

97

98

2

9

2

7

2

5

2

3

2

1

2 1

94

95

96

97

98

2

12

2

10

2

8

2

6

2

4

20

40

60

80

100

2

12

2

10

2

8

2

6

2

4

75

80

85

90

95

2

12

2

10

2

8

2

6

2

4

92

93

94

95

96

97

98

99

2 0

2 1

2 2

94.0

94.5

95.0

95.5

96.0

96.5

97.0

97.5

98.0

2 0

2 1

2 2

2 3

93

94

95

96

97

98

2

17

2

15

2

13

2

11

2

9

2

7

86

88

90

92

94

96

98

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on how to better combine primary rewards that are typically sparse with auxiliary rewards designed by human heuristics to maximize the primary rewards. \n\nThe authors propose a bi-level optimization procedure, where the upper level optimizes a behavior alignment reward to guide agent to maximize the primary rewards and meanwhile minimizes the discount value to force the learned rewards to encode long-horizon information, and the lower level can be any policy searching algorithm. They provide both intuitive and theoretical analysis to show that the formulated procedure is not only robust to misspecified auxiliary rewards, but also overcomes the imperfection of policy optimization algorithms, such as biased policy gradients caused by omitted discounted term and no importance sampling in off-policy learning. \n\nBesides the theoretical analysis, the authors also provide two methods to empirically solve this untractable bi-level problem since the optimization of upper level problem requires characterizing the impact of learned rewards and discount value on the entire optimization process. One is a gradient based method that makes use of the implicit gradient denoted as Implicit Bi-level optimization and the other takes into account the complete path followed by the inner optimization loop denoted as Path-wise Bi-level optimization.\n\nThe authors also conduct experiments on some case studies to show the robustness and extent to Mujoco tasks to show the scalability to high-dimensional continuous control.\n\nSTRENGTHS:\n1. This paper presents a good motivation to motivate why this work is of great significance and why previous works like potential based methods fail (higher variance and only state-based potential function).\n\n2. This paper dives deep to analyze the advantages of the proposed bi-level objective (Section 3), which is helpful for readers to catch the main idea of this paper.\n\n3. The analysis in Section 3.1 seems novel to me. The proposed bi-level objective can not only address auxiliary reward imperfection but also the imperfection of policy optimization algorithms, which is novel and valuable.\n\nWEAKNESSES:\n1. The resulting practical algorithms seem too complex to conveniently apply for higher-dimensional tasks like vision-based tasks.\n\n2. The proposed method introduce a lot of additional hyperparameters as shown in Table 2.3.4.5, which may be hard to tune. For instance, in Table 3, $\\alpha_\\theta$ is tuned to 0.015625, which does not seem like a common value and may be carefully tuned.\n\n3. The experimental results are not so convincing. In figure 3, Barfi only wins 1 out of 4 tasks. Barfi also cannot outperforms baselines in figure 2. 5. \n\n\n4. The authors claim that Barfi requires a warm-up stage, after that, Barfi can fastly adapt to find good policies. However, the reasons why Barfi requires warm-up and how does Barfi performs without such warm-up remains unclear.\n\n\n5. The empirical analysis part is hard to follow, it would be better if the authors could separate some large paragraphs and simplify the abundant detailed task descriptions but focuses on more in-depth analysis.\n\n\n\nThere is also a recent work that focuses on correcting the imperfect rewards using a bi-level objective in the offline RL setting [1]\n\n[1] Li, J., Hu, X., Xu, H., Liu, J., Zhan, X., Jia, Q. S., & Zhang, Y. Q. Mind the gap: Offline policy optimization for imperfect rewards. ICLR2023.\n\nQUESTIONS:\nPlease refer to weaknesses for details.\n\nLIMITATIONS:\nPlease refer to weaknesses for details.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel framework that employs a bi-level objective for learning a “behavior alignment reward” to find the exact reward specification that elicits the desired behavior in RL. Specifically, the reward function combines auxiliary rewards, defined by a designer’s heuristics, with primary rewards provided by the environment. The experiment shows the performance of the proposed method. However, I have some concerns about this paper. My detailed comments are as follows.\n\nSTRENGTHS:\n1. In order to find the exact reward specification, the authors propose Barfi to dynamically adapts the agent’s policy optimization procedure and address other sub-optimalities introduced by algorithmic design choices.\n2. The proposed method allows for the separation of the desired behavior specification from the optimization process with a learned reward function.\n3. The authors show how the use of implicit gradients can enhance the efficiency of optimizing the reward function compared to traditional path-wise gradient-based meta optimization methods\n\nWEAKNESSES:\n1. Figure 1 explains the effect brought by the proposed method, but does not reflect why the proposed method can change the direction of the gradient and dynamically correct the entire policy optimization process by changing the behavior alignment rewards. It is suggested to visualize the principle of the method in this figure.\n2. As for the experiment results, I found that only a few results show that Barfi outperforms other methods, which makes me confused about the effectiveness of the method. Please specify the advantages of the proposed method.\n3. The authors compared the proposed method with potential-based reward shaping. Is there any other kind of reward shaping method? It is recommended to compare the proposed method with more relevant work to demonstrate its effectiveness.\n4. Some visualized results about the behavior of the agent with the proposed method are required to support the effectiveness.\n5. In Figure 9, only two curves in each sub-figure. Please explain where are other curves.\n6. This paper provides a detailed theoretical description of the proposed method. However, it is not easy to follow. It would be better to represent the method with a flowchart.\n\nQUESTIONS:\nSee the weaknesses.\n\nLIMITATIONS:\nSee the weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

comments: SUMMARY:
The authors propose a method (Barfi) that makes use of auxiliary rewards (provided by the designer), in addition to the rewards already defined for the environment. The method is able to use auxiliary rewards if they are helpful for the task, and avoid the auxiliary signal if they seem to be misspecified. The paper seems well written, for example, it motivates the need for such an algorithm, explains the objective function well, why it's needed, etc. and also seems to have decent experiments.

QUESTIONS:
On Line 91 -  The work [1] that shows that potential based reward shaping is not suited to Q learning is only for the tabular case, right? Would the results hold for gradient based deep RL algorithms?

It would be interesting to see how Barfi works when the given auxiliary reward is sometimes helpful (for a part of the state-action space) and sometimes misspecified (for another part of the state-action space). Assuming a parameterization like on line 126, would it be possible to show what happens with a Gridworld environment when the auxiliary reward is designed in this way? Does this need a parameterization where $\phi_i$ also depend on $s,a$?

Suggestions related to language and typos:
1. On Line 157 - "due its" -> "due to its"
2. On Line 127 - $\theta_3$ should be $\phi_3$?

References:
1. Potential-based shaping and q-value initialization are equivalent, Wiewiora (2003)

EDIT (16 Aug 2023): updated score from 7->8.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nRL agents often have difficulty learning from sparse rewards, and thus, it is commonplace to introduce an additional auxiliary reward to guide the learning process. However, these heuristic auxiliary rewards can introduce new side effects and biases into learning. In the paper, the authors develop a novel framework for combining the auxiliary reward and true reward in a way that is aligned with the true reward (even given the *learning* process, in contrast to e.g. shaping theorem which says whether the new reward has the same opt policy but does not take into account the potential imperfections in the learning process). In particular, they formulate a bi-level optimization problem in which the outer problem is to learn a combination of the auxiliary and base reward that does well wrt to the true reward given a specific algorithm that solves the inner problem of computing a poilicy given the combined reward.\n\nSTRENGTHS:\nThis paper is original, a high-quality, principled solution to a problem with high practical significance. The paper is also very well-written and clear.  Every choice is justified and the authors do a good job of contextualizing their work (e.g. in contrast with potential based reward-shaping).\n\nThe experiments are meaningful, well-organized and conveyed. I appreciated that the experiments used a range of realistic auxiliary reward functions including heuristic-based reward functions, reward functions that encode a guess of an optimal policy, ones that change the optimal policy, ones that depend on only states, and ones that depend on both states and actions. The authors report results on robustness, scalability, to high dimensions, and even compute/memory comparisons. In addition, there are several ablation studies in the appendix.\n\nWEAKNESSES:\nThere is probably more related work to cite in terms of learning from auxiliary rewards, e.g. Inverse Reward Design https://arxiv.org/abs/1711.02827\n\nQUESTIONS:\nn/a\n\nLIMITATIONS:\nThe authors address some limitations and extensions in their conclusion.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}