Behavior Alignment via Reward Function Optimization
Authors: Dhawal Gupta
Yash Chandak
Emma Jordan
Philip S. Thomas
Bruno Castro da Silva
Venue: NeurIPS 2023 spotlight (2023)
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
Dhawal Gupta' University of Massachusetts
Yash Chandak"1 Stanford University
Scott M. Jordan' University of Alberta
Philip S. Thomas University of Massachusetts
Bruno Castro da Silva University of Massachusetts
Abstract
Designing reward functions for efficiently guiding reinforcement learning (RL) agents toward specific behaviors is complex task. This is challenging since it requires the identification of reward structures that are not sparse and that avoid in- advertently 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-basse 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 new framework that uses bi-level ob- jective to learn behavior alignment reward functions. These functions integrate auxiliary rewards reflecting a designer's ' heuristics and domain knowledge with the environment's 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, 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 s efficacy 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-specifie heuristics. It not only addresses key shorteomings of existing approaches but also consistently leads to high- tigh-perforning solutions, even when given misaligned poorly-specified auxiliary reward functions.
Introduction
In this paper, 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 ensure that such auxiliary rewards 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 specified goal state, and zero otherwise. However, optimizing policy based on such a sparse reward function often proves challenging.
Both authors contributed equally to this work. Work done while at the University of Massachusetts. Corresponding author: Dhawal Gupta (dgupta@cs umass, edu).
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Figure 1: Auxiliary rewards can be used to convey the agent how we (designers) think it should solve the problem. However, if not carefully designed, they can lead to policies that result in undesired behaviors. This figure provides visual illustration of toy example depicting how the proposed method works. The star represents the optimal policy, and the red dot represents the fixed point of policy optimization process under a "sub-optimal" heuristic; i.e. one that, when naively combined with Tp. induces behaviors different from those under the optimal policy for rp- (Left) Vector field of a policy optimization process converging to a sub-optimal policy. (Middle and Right) By changing the influence of auxiliary rewards, our method can dynamically correct the entire policy optimization process steering it towards a policy that results in the desired behavior.
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 's distance to the goal. However, the effectiveness of using auxiliary reward functions largely depends on the problem's complexity and the designer's '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 different initialization of the q-function; (ii) does not alter policy gradients in expectation; and (iv) it can increase the variance in policy gradient methods.
To address these challenges, introduce a scalable algorithm that empowers RL practitioners to specify potentially imperfect auxiliary reward functions. 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 level objective to analyze and automatically fine-tune designer- ereattdd auxiliary reward functions. 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 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 sarecy--rrdiet-baasee 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 identify near-optimal policies.
(3) How to ensure scalability to high-dimensional problems: We introduce an algorithm that employs implicit gradients 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 's efficacy across 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 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 tuple! (S, A.p. P. Tp: Tauxs 7, dg), where S is the state set, A is the action set, p is the transition function, Tp
:p:S :S A + R is the primary reward function, Faux

---

S X A is an optional auxiliary reward function (possibly designed by a non-expert in machine learning, based on domain knowledge), and do is the starting state distribution. Let Ty :8 + [0. 1] be any policy parameterized using 0 E O. For brevity, we will often use To and 0 interchangeably. Let St and A, be the random variables for the state and action observed at the time 1. As in the standard RL setting, the performance J(O) of policy Ty is defined as the expected discounted return with respect to the (primary) reward function, Tpi i.e., J(0) == E T ('rp(S,, A.)], where T + is the episode length. An optimal policy parameter 0" is defined as 0 E arg maX J(). A popular technique search for 0 is based on constructing sample estimates (66,rp) ry) of the (y-dropped) policy gradient, A rp), given an agent's 's interactions with the environment for one episode [58, 59]. Then, using (s a) as shorthand for d (T a)/d9, these quantities are defined as follows:
T
T
(0) Euo
[(0.rp)
and Á rppp =
Σ
bh (St, A1)
Σ o ( (s) (1) jmn
3 Limitations of Potential Based Reward Shaping
When the objective function J(0) is defined with respect to a sparse reward function Tp (i.e., a reward function such that rp(s, a)== for most E S and A), searching for 0* is challenging [24]. A natural way to provide more frequent (i.e., denser) feedback to the agent, in the hope of facilitating learning, is consider alternate reward function, i Favve r +rr However, as discussed earlier, Taux may be a designer-specified auxiliary reward function not perfectly aligned with the objective encoded in Ty- In this case, using 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 functionn, & S IR (analogous to Tzux)) to define new reward functions of the form Fo (St, A,, S:+1): rrp:)..-pp...... A -++1)-((S,)) Importantly, they showed that optimal policies with respect to the objective E E 7 (Se, Ac, S(+1)] are also optimal with respect to (0).
While potential-based reward shaping can partially alleviate some of the difficulties arising from sparse rewards, Wiewiora [64] showed that q-learning using Fo produces the exact same sequence of updates as q-learning using Tp but with a different initialization of q- q-values. In what follows, we establish 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.
Property 1. E[ fx)] (O)(O, and Var( f4)) can be higher than Var ((r()(0, rp))-
All proof's are deferred to Appendix A. The above points highlight some of the limitations of potential- based 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 T cannot depend on actions [41], which restricts the class of eligible auxiliary rewards Taux and heuristics functions that may be used. Finally, notice that & is designed independently of the agent's '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 inves- tigated 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 Taux to Tp may produce policies whose corresponding behaviors are misaligned with respect to the behaviors induced by Tp. In these cases, Tzux should be ignored during the search for an optimal policy. On the other hand, if rp and Taux 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 behavior alignment reward by combining Taux and rp using parameterized function. Our method is inspired by the optimal rewards framework by Singh et al. [52, 53]. Let 7 e [o, 1) be
3

---

a discount rate value parameterized by 9 E T.² rg 8 x A + R be behavior alignment reward: function of both rp and Taux. parameterized by o T, where T and are function classes. One example of a behavior alignment reward function is r a) = Sor (8,a) dp a) du (8, a), where Sa:S x -> TR and Φ: (01, 42, 0h3). Let A y be any based algorithm that outputs policy parameters. To mitigate possible divergence issues arising from certain policy optimization algorithms like DQN [60, IJ, we make the following simplifying assumption, which can generally be met with appropriate regularizers and step sizes:
Assumption 1. Given T and Tip the algorithm Al g Tip) converges to a fixed point f_ which we denote as ((o, :p) E Ô emphasize its indirect dependence on ó and ip through Alg, To. and Tip-
Given this assumption, we now specify the following bi-level objective:
4p" arg max J(0(), >y7pp where T
-Al((ra:))))))
Here, A, serves the regularization coefficient for the value of Yg. and Alg denotes given policy optimization algorithm. Let. as an example, Alg be an on-policy gradient algorithm that uses samples to estimate the gradient (0), as in (I). We can then define corresponding variant of A rp) that is compatible with our formulation and objective, and which uses both Te and 7pp as follows:
Don (0, O, ()) = E Σ Vo AD
Σ (o j=t
(3)
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: evaluates different parameterizations, denoted by é. for the behavior alignment reward function. These parameterizations influence the induced policy ((ú, sp), which is evaluated using the performance metric J. Recall that this metric quantifies the alignment of a policy with the designer's 's primary reward function, Tpp In essence, the outer-level objective seeks to optimize the behavior alignment reward function to produce policies that are effective according to Tp. This design adds robustness against any misspecification of the auxiliary rewards.' 3 In the inner-level optimization, by contrast, Alg identifies policy 0(Ø (p) that is optimal or near-optimal with respect to r (which combines r Taux through the behavior alignment reward). In the inner-level optimization, the algorithm Alg works to identify a policy ((Ø, sp) that is optimal or near-optimal in terms of r (which incorporates Taux via the behavior alignment reward). By employing a bi-level optimization structure, several benefits emerge. When Fzux well-crafted, re can exploit its detailed information to give Al frequent/dense reward feedback, thus aiding the search for an optimal 0". Conversely. if Taux leads to sub-optimal policies, then the influence of auxiliary rewards can be modulated decreased accordingly by the optimization process by adjusting rá. Consider, for example, case where the behavior alignment reward function is defined ( a) == for (s,a) Φ2r, a) + Φgraux (8, a). In an adversarial setting setting-where the designer-specified auxiliary reward Taux may lead to undesired behavior- the bi-level optimization process has the ability to set du O. This effectively allows the behavior alignment reward function T to exclude Taux 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 misaligned auxiliary reward, Taux- This design is particularly valuable because allows the objective (2) to leverage the potentially dense reward structure of Taux to provide frequent action evaluations when the auxiliary reward function is sell-specified. At the same time, the approach maintains robustness against possible misalignments.
The bi-level optimization problem (2) may have multiple optimal solutions for o-including the trivial solution where Taux is always ignored. The goal of regularizing the outer-level objective (in the form of the term Xy is to incorporate prior that adds preference for solutions, d²*, that provide useful and frequent evaluative feedback to the underlying RI. algorithm. In the next paragraphs, we discuss the need for such regularizer and motivate its mathematical form. First, recall that sparse rewards can pose challenges for policy optimization. An intuitive solution
Need for Outer Regularization:
"Our framework can be generalized to support state-action dependent discount rates, ?. 3 "Miispeccffication indicates that optimal policy for rp + Taux may not be optimal for rp alone.

---

to this problem could involve biasing the optimization process towards denser behavior alignment reward functions, e.g., by penalizing for sparsity of rg. 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., 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) 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 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.* 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 regularizer, 7/- This regularizer favors solutions that can generate policies with high performance (i.e., high expected return J with respect r,p) even when the discount factor 7 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..., 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, Ty-
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 Faux. Even with a well-specified Taux: RL algorithms often face design choices, such as the
trade-off, that can induce sub-optimal solutions. Below present examples to show how bias in the underlying RL algorithm may be mitigated by carefully optimizing rs and %p-
4.1.1 Bias in policy gradients: Recall that the popular "policy gradient" (0),r,) 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 (68,Tp) remain vastly popular in the
RL literature since they tend to be sample efficient [59]. Let (0, rp) denote the unbiased policy gradient, where (0, = E
T 7 Vo((S, A.) Σ "r, A,)]], We can show that with sufficiently expressive parameterization, optimized rs and 7op can effectively mimic the updates that would have resulted from using the unbiased gradient A. (0. rp). even if the underlying RI. algorithm uses the biased "gradient". Aun (f,, Φ, ()) as defined in (3). Detailed proofs are in Appendix A.
Property 2. There exists r& S x A + E and Tpp [0. 1) such that Don (0, o. x)) (0,rp)
4.1.2 Off off-policy learning without importance sampling: To increase sample efficiency when eval- uating a given policy Ty. it is often useful use off-policy data collected by a different policy, B. Nals.a) can
Under the assumption that Vs E S, Va E A, B(s.a] < co, importance ratios Pj == IIr_o [[s..))
xa(s.a)
be used to adjust the updates and account for the distribution shift between trajectories generated by and Ty. However, to avoid the high variance stemming from Pj. many methods tend drop most of the importance ratios and thus only partially correct for the distribution -wich can lead to bias [5]]. We can show (given a sufficiently expressive parameterization for the behavior aligmment reward function) that this type of bias can also be mitigated by carefully optimizing T and 7pp
Let denote the unbiased off-policy update with full-distribution correction as Auri(,, rp) = Eg T (Sr. Ad) EJ= P,r' A,)]. 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 ( d. sp) = E3 IT (h At) EJ= 26. (S),
Property 3. There exists ro 8 x A + R and 7 [0, 1) such that (0, +) = Amr(O, (p))))
Remark 1. Our method is capable of mitigating various types of algorithmic biases and imperfections in underlying RL algorithmss, without requiring any specialized learning rules. Additionally, thanks
E.g., if re se g", then its corresponding rewards instantly indicative of the optimal action any state.
5

---

to the " regularization, it favors reward functions that lead to faster learning of high bigh-peeforming policies aligned with the designer's objectives, as outlined in the original reward function Tp-
BARFI: Implicitly Learning Behavior Alignment Rewards
Having introduced our bi-level objective and discussed the benefits of optimizing ra and 7,p an important question arises: Although ((1, sp) can be optimized using any policy learning algorithm, how efficiently identify the optimal 6- and 4p in equation (2)? Given the practical advantages of gradient-basse methods, one would naturally consider using them for optimizing 6 and 4 as well. However, key challenge in setting lies in computing d./ (0(ó, s))/dó and d. p))/dpp These computations require an analytical characterization the impact that r and 7ip have on the entire optimization process of the inner-level algorithm, Alg.
In addressing this challenge, initially focus Alg that employs policy gradients for updating Ty. Similar extensions for other update rules can be derived similarly. We start by re-writing the expression for clJ (0(6, /d) using the chain rule:
dj(((()))) dJ(0(á, p)) d0(), dó
(4)
d0(o. 1)) dø
(a)
where (a) is the policy gradient 0(d, s),, and (b) 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 RI. algo- rithms [50]. First, observe that when Alg converges to 0(0, sp), then it follows that
(5) Let Oʃ denote the partial derivative with respect to the immediate arguments of J. and d be the total derivative as before. That is, g(=)) == rg(r), then of (xr, g(æ) and df dr (Ix =
( (p), Ø, vp) = 0.
Bg g(z) + De (z)) Therefore, taking the total derivative of (5) with respect to é yields
(A(0((, d.P) @ p)
(2(0().)) 0, 00 4)
(6)
dø
dΦ
a0 ip)
06
By re-arranging terms in (6), we obtain the term (b) in (4). In particular,
00( v)
(2) (0()))))
(
(7)
as
Furthermore, by combining (7) and (4) we obtain the desired gradient expression for d:
aJ(O(e. (J((()))))
(A (O((0))))...)) ( up)
@A(O(á, v), c.
(8)
0Φ
(0 s))
Similarly, a gradient expression for can be derived; the full derivation is detailed in Appendix E. Using 6 as shorthand for (( sp). find that the terms A and H can be expressed as
A = Ep Σ wha (St, As)
Σ
Oral8,, A.)
H Ep
Bean (St, Ad)
Σ Enttro (ro(8)
06
When working with the equations above, assume the inverse of H 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 y in our bi-level optimization's outer loop. However, computing A and H directly can be impractical for high-dimensiioal 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) calculate (8) and (15) via Hessian-vector products [47], which are readily available in modern auto-diff libraries [46]. These methods climinate 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 ((d) complexity, where d 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 BARFT, an acronym for behavior alignment reward function's 's implicis optimization.? 5 BARFI designed to iteratively solve the bi-level optimization problem defined in (2). With policy regularization, the updates to r' and Nie incrementally modify (( p). 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.
Empirical Analyses
Our experiments serve multiple purposes and include detailed ablation studies. First, we demonstrate our bi-level objective's '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- reward shaping, showing how it can lead to suboptimal policies (Section 6.2). We then provide qualitative illustration of the behavior alignment reward function learned by BARFI (Section 6.3). Finally, we evaluate how well BARFI I scales to problems with high-dimensional, continuous action and state spaces (Section 6.4).
In the sections that follow, we examine range of methods and reward combinations for comparison:
Baseline RL methods: We consider baseline RL methods that employ a naive reward combination strategy: naive (s, a) rp(ss, a) + Tsux 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- Aator-Critic algorithms are used for optimization.
Potential crential-baseddae shaping: To assess how well potential- potential-based reward shaping performs, we introduce variants of the baseline methods. Specifically, we investigate the effectiveness of the reward function (8 0, 8') == (x,8, a) + 7/aaux (s') - Taux(s).
BARFI: We use REINFORCE
the underlying RL algorithm when implementing BARFI I and define (s a) == Φ (8, a) + ()())), a) + ( (s,
(s,a). a). Our implementation includes a warm-up period wherein the agent collects data for fixed number of episodes, using f Faivee 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], Cart- Pole [16], and HalfChectah-v4 [9]. These domains offer increasing levels of complexity and are intended to assess the algorithms' adaptability. Furthermore, we examine their performance under a variety of auxiliary reward functions, ranging from well-aligned to misaligned with respect to the designer's '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 '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 designer's '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 Taux that provides bonuses when the agent's 's actions align with a known effective policy in this domain. Second, we designed 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.c., well-known effective policy [18]) to
""BARFI" commonly refers to type of south-Asian sweet confectionery, typically pronounced as 'bur-fee'.
7

---

Table 1: Summary of the performance of various reward combination methods and types of Taux
CartPole
MountainCar
Mothod for Reward Combination
Well-aligned Paux
Misaligned Faux
Well-aligned Paux (wet. energy policy)
Partially-alinned Taux (wirt. high velocity policy)
BARFT I (our method) Fraive (naive reward combination)
487 2 + 9.4
475.5 + 15.5
0.99 + 0.0
0.90 + 0.1
498.9
1.0
9.04 + 0.2
0.99 0.0
0.63 + 0.1
(potential-based shaping)
8.98
0.2
500
+ 0.0
0.00 0.0
0.00 0.0
BARFI's performance compared to two baselines that use: Fnavve and rq. respectively. CartPole uses an retioondepenneen Taux function that either rewards agents when actions align with a known effective policy (well- whell-aligned Taux) or with a moolyyperforming policy (misaligned T Taux). MountainCar uses either an action-dependent function aligned with energy- poarging policy [18] a partially-alignn function incentivizing higher velocities. BARFI consistently achieves near-optimal performance across scenarios, even if given poorly speciffeddmisslignedd 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 partially-alineed 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 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 Section 6.2.
The key takeaway from the experimental results in Table is that BARFI consistently performs well across various domains and under different types of auxiliary rewards. Specifically, when designer-speciffed feedback is appropriate and can assist in accelerating learning, BARFI I efficiently exploits it to produce high-performing policies. Conversely, if auxiliary rewards are misaligned with the designer's 's intended goals, BARFT 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, 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 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 no prescribed way for designing potential shaping when Taux is action-dependent, we use direct extension of the original formulation [41] by considering (a s', s',a') ;= ry a) yr (s', aa) - Taux(s. (s, a). Furthermore, we designed an auxiliary reward function, Taux: that is well aligned: it provides positive reward signals of fixed magnitude both for (s, a) and (s', a²) whenever the agent's s actions coincide with the optimal policy. Notice, however, that if y < 1, the resultant value from "Tx (s', a') T a) is negative. Such negative component may deter the agent from selecting actions that are otherwise optimal, depending on how Taux and rp differ in magnitude. Conversely, potential-based shaping can also occasionally perform well under misaligned rewards. In these cases, the shaping function may yield positive value whenever the agent selects an optimal action, which could induce well- well-performing behaviors.
8

---

Reward Canter State
GridWorld
gl((s)
GridWorld
CartPole
Gool
100
100
100
Raistoraa Patantial
030
+00
230
Narm Pwrldd
Actor Oriti Potertia
130-
Actor Craid
Releforce
$0
500
1815
1000
2050
205
400
100
1002
Episodes
Eplxodea
Figure 2: (Left) Misaligned Taux (+50) for entering the center state and rp (+100) at the goal state. (Center Left) Optimized weight (s(s) depicting how the learned reward function penalizes the agent for entering the center state. (Center Right) Performances in the GridWorld domain under a misspecified Faux- (Right) Performances in the CartPole domain under a Taux providing positive feedback for mimicking a known ineffective policy.
6.3 What does no 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 the agent when it makes progress towards the goal. However, such intermediate positive feedhack 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 [center left]), thereby avoiding poorly- performing behaviors that focus on revisiting irrelevant central states. Similarly, when Cartpole operates under misspecified Yaux (Figure 2 [right]), BARFI capable of rapidly adapting (after a warm-up period) and effectively disregarding misleading auxiliary reward signals. These results highlight once again BARFI I 's robustness when faced with reward misspecification.
6.4 Scalability to High-Dimensional Continuous Control
One might wonder whether computing implicit gradients for & and yp would be feasible in high- dimensional problems, due to the computational cost of inverting Hessians. To address this concern, we leverage Neumann series approximation with Hessian-vector products (See Appendix C) and con- duct further experiments, as shown in Figure 3. These experiments focus on evaluating the scalability of BARFI in control problems with high-dimensiioal state spaces and continuous actions- that often rely on neural networks for both the policy and critic function approximators. For more comprehensive evaluation, we also introduced an alternative method named BARFI I unrolled Unlike BARFI, I. which uses implicit bi-level optimization, BARFI unrolles employs path-wise bi-level optimization. It maintains a complete record of the optimization path to determine updates for 6 and pp. Further details regarding this alternative method can be found Appendix C.6.
We conducted experiments on the HalfCheetah-v4 domain and investigated, in particular, a reward function comprising two components with varying weights. This empirical analysis was designed to help us understand how different weight assignments to each reward component could influence the learning process. Specifically, in HalfCheetah-v44, the agent receives positive reward rp proportional to how much it moved forward. It also incurs a small negative reward (concretely. an auxiliary reward, Taux a) == clJa||2, I|2. known as a control cost) for the torque applied to its joints. A hyperparameter determines the balance between these rewards. The naive combination of such primary and auxiliary rewards is defined Fraive (s, a) rp,s, a) + Taux (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 unroll led 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 (4) increases, due to growing computational and memory requirements (Figure 3 [right]). Although we do not recommend BARFI unrol led, we include its results for completeness. Additional ablation studies on (a) the effect of the inner optimization step; (b) Neumann approximations: (c) decay of y; and (d) returns based on re. are provided in Appendix H.

---

Mujoco
Mujoco
Compute/Memory
2500
2500
Compute Time (Hri))
2000
2000
Barfi Unrolled
(PBW) Amww
Return
1500
Borti
2000
Barl (.Stppsl
3000
Barhi Uncolled
Critie
Barth Stree
Bort Stops)
300
Reinforce
Actor Critic Potential
Unrotled
G0OD
Episodes
Episodes
Number of Inner Updates
Figure 3: Results for MuJoCo environment. (Left) Auxiliary reward is defined to be -c|al where e is a positive hyperparameter and is the continuous high- -ihennmennional action vector. (Middle) Similar setting as before, but uses amplified variant of the auxiliary reward: -4c||a| [.. It is worth highlighting that under alternative reward weighting schemes, hoth variants of our (bchavior-aligned) bi-level optimization methods demonstrate successful learning. Learning curves correspond to mean return over 15 trials, and the shaded regions correspond to one standard error. (Right) Required compute and memory for BARFI anrolled compared to BARF I, function of the number of inner-optimization updates. This figure also showcases BARFI's characteristics under various orders of Neumann approximation.
7 Related work
This paper focuses primarily on how to efficiently leverage auxiliary rewards Taux- Notice, however, that in the absence of Taux: the resulting learned behavior alignment rewards ra 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 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 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 techniques, most rely on path-wise meta meta-gradien gradd-ggddents..
mentioning that among the abooe-mentiones 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, I, by contrast, exploits implicit gradients to alleviate these issues by directly characterizing the fixed point of Alg induced by learned behavior alignment rewards.
Finally, 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 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 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- may not be experts in the field-to incorporate domain knowledge through heuristic auxiliary reward functions. Our framework allows for more expressive reward functions be learned while ensuring they remain aligned with a designer's 's original intentions. BARFI can also identify reward functions that foster faster learning while mitigating various limitations and biases underlying RL algorithms. We empirically show that BARFI is effective in training agents in sparse-reware scenarios where (possibly poorly-specified auxiliary reward information is available. If the provided auxiliary rewards are determined to be misaligned with the designer's s intended goals, BARFI I autonomously adapts and effectively disincentivizes their use as needed. This adaptability results in a reliable pathway to identifying tigh-performin policies. The conceptual insights offered by this work provide RL practitioners with a structured way to design more robust and mas-tto-op reward functions. We believe this will contribute to making RL more accessible to broader audience.
10

---

Acknowledgement and Funding Disclosures
We thank Andy Barto for invaluable discussions and insightful feedback on an earlier version of this manuscript, which significantly improved the quality of our work.
This work is partially supported by the National Science Foundation under grant no. CCF-2018372 and by gift from the Berkeley Existential Risk Initiative.
References
[1] Joshua Achiam, Ethan Knight, and Pieter Abbeel. Towards characterizing divergence in deep q-learming, arXiv preprint arXiv:1903.0889/0 2019.
[2] Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and Zico Kolter. Differentiable convex optimization layers, arXiv preprint arXiv:1910.124333 2019.
[3] Ron Amit, Ron Meir, and Kamil Ciosek. Discount factor as a regularizer in reinforcement learning. In International Conference Machine Learning. 2020.
[4] Brandon Amos and J Zico Kolter. Optnet:
Differentiable optimization as layer in neural networks. International Conference on Machine Learning. 2017.
[5] Stuart Armstrong, Jan Leike, Laurent Orseau, and Shane Legg. Pitfalls of learning reward function online. arXiv preprint arXiv:2004.1366 2020.
[6] Karl Åström and Tore Hägglund. Pid control. IEEE Control Systems Magazine, 2006.
[7] Sarah Bechtle, Artem Molchanov, Yevgen Chebotar, Edward Grefenstette, Ludovic Righetti, Gaurav Sukhatme, and Franziska Meier. Meta-learning via learned loss. arXiv preprint arXiv:1906.05377 05374, 2019.
[8] Yoshua Bengio. Gradient-basce optimization of hyperparameters. Neural computation, 2000.
[9] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint
2016.
[10) 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.02555 2023.
[11] Xu Chu Dennis Ding, Stephen L Smith, Calin Belta, and Daniela Rus. LTL control uncertain environments with probabilistic satisfaction guarantees. IFAC Proceedings Volumes, 2011.
[12] Chuong B Do, Chuan-Sheng Foo, and Andrew Y Ng. Efficient multiple hyperparameter learning for log-linear models. In Advances Neural Information Processing Systems, 2007.
[13] Justin Domke. Generic methods for optimization-based modeling. In Artificial Intelligence and Statistics, 2012.
[14] Asen L Dontchev and R Tyrrell Rockafellar. Implicit functions and solution mappings, volume 543. Springer, 2009,
[15] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In Proceedings of the 34th International Conference Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, 2017.
[16] Razvan V Florian. Correct equations for the dynamics of the cart-pole system. Center for Cognitive and Neural Studies (Coneural), Romania, 2007.
[17] Zhao- Zhao-Yang Fu, De-Chuan Zhan, Xin-Chun Li, and Yi-Xing Lu. Automatic successive reinforce- ment learning with multiple auxiliary rewards. In IJCAI, 2019.

---

[18] 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- 9-1 2020, 2020.
[19] 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 1607.::600.05..00 2016.
[20] Edward Grefenstette, Brandon Amos, Denis Yarats, Artem Molchanov, Franziska Meier, and Kyunghyun Cho. higher: A pytorch meta-learning library. 2020.
[21] Ricardo Grunitzki, Bruno C da Silva, and Ana LCC. Bazzan. A flexible approach for designing optimal reward functions. In Proceedings of the 16th Conference on Autonomous Agents and MultiAgent Systems, 2017.
[22] Ricardo Grunitzki, Bruno C da Silva, and LC Ana Bazzan. Towards designing optimal re- ward functions in multi-agent reinforcement learning problems. In 2018 International Joint Conference on Neural Networks (IJCNN), 2018.
[23] 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: 170400055, 2016.
[24] Joshua Hare. Dealing with sparse rewards in reinforcement learning. arXiv preprint arXiv:191000928// 2019.
[25] Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Cited on, 2012.
[26] Mark K Ho, Michael Littman, Fiery Cushman, and Joseph Austerweil. Teaching with rewards and punishments: Reinforcement or communication? In CogSci, 2015.
[27] Mark K Ho. Fiery Cushman, Michael L. Littman, and Joseph I. Austerweil. People teach with rewards and punishments as communication, not reinforcements. Journal of Experimental Psychology: General, 2019.
[28] 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.
[29] Rodrigo Toro Icarte, Toryn Klassen, Richard Valenzano, and Sheila Mcllraith. Using reward machines for high-level task specification and decomposition in reinforcement learning. In International Conference on Machine Learning. 2018.
[30] Rodrigo Toro Icarte, Ethan Waldie, Toryn Klassen, Rick Valenzano, Margarita Castro, and Sheila Mcllraith. Learning reward machines for partially observable reinforcement learning. Advances in Neural Information Processing Systems, 2019.
[31] Rodrigo Toro Icarte, Toryn Q Klassen, Richard Valenzano, and Sheila A Mcllraith. Reward machines: Exploiting reward function structure reinforcement learning. arXiv preprint arXiv:2010.03950, 2020.
[32] 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,201 2015.
[33] Louis Kirsch, Sjoerd van Steenkiste, and Jürgen Schmidhuber. Improving generalization in meta reinforcement learning using learned objectives. arXiv preprint 1rXiv:1910.04098, 2019.
[34] Steven G Krantz and Harold R Parks. The implicit function theorem: history, theory, and applications. Springer Science & Business Media, 2012.
[35] Hadas Przss, Georgios Fainekos, and George Pappas. Temporalllogic-baseed reactive mission and motion planning. IEEE Transactions Robotics, 2009.
12

---

[36] Karl Kunisch and Thomas Pock. A bilevel optimization approach for parameter learning in variational models. SIAM Journal on Imaging Sciences, 2013.
[37] Jan Larsen, Lars Kai Hansen, Claus Svarer, and 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.
[38] 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.
[39] Michael I. Littman, Ufuk Topeu, Jie Fu, Charles Isbell, Min Wen, and James MacGlashan. Environmenttindeeendenn task specifications via GLTL. arXiv preprint 1rXiv:/704.043: 2017.
[40] Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In Infernational Conference Artificial Intelligence and Statistics, 2020.
[41] Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transforma- tions: Theory and application to reward shaping. In ICML. 1999.
[42] Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. CoRR, 2018.
[43] Scott Niekum, Andrew Barto, and Lee Spector. Genetic programming for reward function search. IEFE Transactions on Autonomous Mental Development, 2010.
[44] Chris Nota and Philip S. Thomas. Is the policy gradient a gradient? arXiv preprint arXiv: 07073, 2020.
[45] Junhyuk Oh, Matteo Hessel, Wojeiech M Czarnecki, Zhongwen Xu, Hado van Hasselt, Satinder Singh, and David Silver. Discovering reinforcement learning algorithms. arXiv preprint N0Ki::200..08794, 2020.
[46] 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.
[47] Barak A Pearlmutter. Fast exact multiplication by the hessian. Neural computation, 1994.
[48] Silviu Pitis. Rethinking the discount factor in reinforcement learning: A decision theoretic approach. Proceedings of the AAAI Conference on Artificial Intelligence, 2019.
[49] Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with
implicit gradients. In Advances in Neural Information Processing Systems, 2019.
[50] Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar. A game theoretic framework for model based reinforcement learning. arXiv preprint arXiv:2004.07804 2020.
[51] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1007.06777.0634 06347, 2017.
[52] 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.
[53] Satinder Singh, Richard I. Lewis, Andrew G Barto, and Jonathan Sorg. Intrinsically motivated reinforcement learning: An evolutionary perspective. IEEE Transactions on Autonomous Mental Development, 2010.
[54] Jonathan Sorg, Richard I. Lewis, and Satinder Singh. Reward design via online gradient ascent. Advances in Neural Information Processing Systems, 2010.
13

---

[55] Jonathan Sorg, Satinder P Singh, and Richard L Lewis. Internal rewards mitigate agent boundedness. In Proceedings of the 27th international conference on machine learning (ICML- 10), 2010.
[56] 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.
[57] Jonathan Daniel Sorg. The Optimal Reward Problem: Designing Effective Reward for Bounded Agents. PhD thesis, University of Michigan, 2011.
[58] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[59] Philip Thomas. Bias natural actor-critic algorithms. International Conference on Machine Learning, 2014.
[60] John N Tsitsiklis and Benjamin Van Roy. An analysis of temporal -iffreeeedf learning with function approximation. IEEE Transactions on Automatic Control, 1997.
[61] Marin Vlastelica, Anselm Paulus, Vít Musil, Georg Martius, and Michal Rolínek. Differentiation of blackbox combinatorial solvers. arXiv preprint 19X1v:11 2019.
[62] Yufei Wang, Qiwei Ye, and Tin Liu. Beyond exponentially discounted sum: Automatic learning of return function. arXiv preprint arXiv: 1905. 2019.
[63] Martha White. Unifying task specification in reinforcement learning. In International Confer- ence on Machine Learning. 2017.
[64] Eric Wiewiora. Potential-based shaping and q-value initialization are equivalent. Journal of Artificial Intelligence Research, 2003.
[65] Eric M Wolff, Ufuk Topcu, and Richard M Murray. Robust control of uncertain markov decision processes with temporal logic specifications. In 2012 IEEE 51 IEEE Conference on Decision and Control (CDC), 2012.
[66] David Wolpert and Kagan Tumer. Optimal reward functions in distributed reinforcement learning. In Intelligent agent technology: Research and development. World Scientific, 2001.
[67] Yuhuai Wu, Mengye Ren, Renjie Liao, and Roger Grosse. Understanding short-horizon bias in stochastic meta-optimization. arXiv preprint 1xXiv:803..0000 2018.
[68] Zhongwen Xu, Hado van Hasselt, and David Silver. deta-nrrddent reinforcement learning. In Advances in Neural Information Processing Systems, 2018.
[69] 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.
[70] Zeyu Zheng, Junhyuk Oh, and Satinder Singh. On learning intrinsic rewards for policy gradient methods. In Advances in Neural Information Processing Systems, 2018.
[71] Zeyu Zheng, Junhyuk Oh, Matteo Hessel, Zhongwen Xu, Manuel Kroiss, Hado Van Hasselt, David Silver, and Satinder Singh. What can learned intrinsic rewards capture? International Conference Machine Learning, 2020.
[72] Haosheng Zou, Tongzheng Ren, Dong Yan, Hang Su, and Jun Zhu. Reward shaping via meta-learning. arXiv preprint 1rX1v:1901.09330, 2019.
14

---

Behavior Alignment via Reward Function Optimization (Supplemental Material)
Table 2: Notations
Symbol
Description
Parameters for policy π Parameters for reward function
Parameters for learned
No, ne: The
Functional form of policy, reward and y with their respective parameters
00, Cld. Oke
Step sizes for the respective parameters Regularization for policy, reward and y function
Dg, Xg, Xo
Number of on-policy samples collected between subsequent updates to Ø, Yp
Neumann Approximator Eigen value scaling factor
n
Number of loops used in Neumann Approximation
optim
Any standard optimizer like Adam, RMSprop, SGD, which takes input
as gradients and outputs the appropriate update Total Number of episodes to sample from the environment Number of updates to be performed for updating the æ by Alg
E
N
No
Number of initial updates be peformed
Sample of a trajectory from full episode
A Proofs for Theoretical Results
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 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 denoted as
T
(s. = Σ
Pr(S: s, a)
(9)
1=0
(s Σ
Pr(S, 8, A, a)
(10)
We normalize it that it is distribution as follows:
(s dr (s
d' (s) =
The expected update performed by the biased policy gradient update is same when using the primary reward and reward modified with ontential-based shaping, i.e., A(0. 7) = A(0,r Further, the variance of the update when using potential-ba based reward shaping can be higher than the variance of the update performed using the primary reward, i.e., Var (A(O.F) Z Var (A(),rp))
Property 1.
15

---

Proof. Part 1: Equality of the expected update
(10, F) BT
(S Ad)
Σ
BT
(S Ad)
Σ
(r A3) + +(S)))
j=1
T
T
T
En
(S Ad)
Σ
"p Aj)
+ E
Σ
(S AA)
Σ
-
5=t
j=t
T
= A(0,rp) + ExD
Σ
Vi (St, Ad)
Σ
q 20 (y00(())) -
/=1
(a)
(p) + E
WoOSS A.) (7 ((T+1))))))))) - +(Si))
T
(b) A(0,rp) + E
(S,, c c-
A00,T)) E..
Σ
(7 x--t1c ((S)))EEo (th (S., S.)
4 A (0 rp),
where (a) holds because on the expansion of future return, intermediate potential values cancel out, (b) holds because Sr+1 the terminal state and potential function is defined be a fixed constant for any terminal state [41], (c) holds from the law of total expectation, and (d) holds because,
In no (St,
@mo(S a)
a 00
E,, (St, A.)
Σ (S
To a) =0. 0.
08
09
REA
In the stochastic setting, i.e., when using sample average estimates instead of the true expectation, -T-++1 c o(S.) is analogous dependee baseline for the sum of discounted future primary rewards. It may reduce or increase the variance of A(8, rp), depending on this baseline's 's co-variance with _ 7" "r= (S).
Note: As we encountered the potential at the terminal state to be c, as it is constant, we will use the value of e 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 Ag at the starting state So and then transitions to the terminal state. Hence, the stochastic update (0, F) can be written as:
Á F) (So Ap) Au) =vo (o)(T, Ao) <(())),
wherein, we assume that Φ for terminal states is 0 and similarly, Á(0, vo As)r,(So, Ao). For the purpose of this proof we will consider the case wherein we have scalar i.e., 0 IR such that, (...,.) e R.
16

---

Hence, Var Var ((0. can be written as:
Var (A
= A R) - F)
[Vo (So, Ao)(Tp(S0 Ao) ((S6)))* Ao) Ao) (())))]
- B (So, An) Ao) (S))2 rr Ao))] (So. Ao) Ao))]
= [o| (So, Ao) (00 Ab)) +
EE cis (So Ao) To Ao) Blua a (6,, Ao) (r) Ao))
va((A(0.)))
Therefore,
Var Var(A (ar()(()F)) Var (ar()())r))) (So, (o) (5( An))]
Subsequently, variance of Á(0, F) will be higher than that of (00,rp) if Ao) 2E (S (S))rr(S Ao) >> 0.
Example: Let us look at an example where the above condition can be true. Let us consider an MDP with single state and single-step horizon. In that case, we can consider the variance of the update to the policy at the said state, i.e.,
Var (A(0.)) Var (A0.(((). (ss) (v00ss A)2] (s, (rp(rss A))]
where 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
4 (20 (0(s)
vo A) (rp(s, A)) > 0
(s)) (v A)? > Ez (r, A))]
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.
+ ss) ( [wo (s, A)
[ ())
Φ {s) (s A) 2E, wo (s. (s,,
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 8, i.e. any ((s), s.t. +fs) > (s, aa 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 bad/mediocre policy initially.
Property 2. There exists T S A - R and Tpp [0. 1) such that (f, 0. (0. rp).
Proof. Recall the definition of (0, rp) from Section 4.1:
(0,rp) = EEgg
Σ
o dh (S, A.)
Σ
2 'r, (S))
1=0
==
17

---

Using the law of total expectation,
T
(0,rrpp) ENT
Σ
/' v/o (St, AL)
Σ
rp(s) Az) St, Ar
Exs
2/ VoooSe A1)9 (S,
T
E
Pr(S, = S, =070)
EES to
T
Σ
Σ
y
Pr(S: = 8, =
sES,eEA
1=
Σ sES
vals,
(11)
Notice from (9) and (10) that for any (s, a) pair, if dij (s, a) 0, then do a) Therefore, dividing and multiplying by d (s, a) leads to:
since y 2
(s.,,))00sss))0" a) (s, =ES,aEA
(,(0,pp) = Σ
dao (s.
a) (s,a) jun (8,
Σ
Σ
Pr(S: s, = a
EES
Σ (S, Ar)a Th (St, Ar) 1 (St, Ac)
E
t S1 Ac)
d,'s (a.a) for all E S and E A, then 7 (a.a)
Now, notice that 7/p and rs(s 0) =0
(0, = (0,0p)))
Property There exists T ro: S x A + R and 7pp E [0, 1) such that (0 6. = Anr(O, rp).
Proof. This proof follows a similar technique as the proof for Property 2. Recall the definition of Aott (0.rp):
T
T Σ 2
Aut((rr = Eg
(S Ai)
j=t
Eg
P'pie (Sc,
==1
18

---

Now using the law of total expectations,
A0n Eg
Σ
3 Py-oppS A,) Se. A
E
(S Ar)
Σ (p A1) Sc. Ae
E
7' Nov (St, Ar)aa (S) As)
vols a a)4² as) (s)
sES,
where the last line follows similar to (11). Now, notice that for any (s, a) pair, the assumption that (s aa/aa)//33)) (s a) oo for all E S,a A, implies (s, a) (s, a) < cc. Further, (s, a) > 0 it has to be that d' (8.a) > 0 as well. Therefore, d:o (s, a) /d (8, a) < co as well.
Multiplying and dividing by d (s, a) results in:
d' (1) (s,a)
(0, = Σ
(s,
(s,
#ES,aEA
T
Pr(S: = 8, A a; (s, (s, (s,
Σ
SES aEA t=0
T
1 (St, A.) Ar)
= Eg
SalS Ar)a (Sr,
d (S A,)
d (a.a) for all E and E A,
Finally, notice that 7op = and ro(s, a) = =q (8,
P(.,0)
(0 x) (0.
Remark Notice that as with any optimization problem, issues of realizability and identifiability of the desired T must be taken into account. The examples provided in this section aim to highlight the capability of optimized behavior alignment reward functions. In particular, they not only improve and accelerate the learning process but are also capable of inducing updates capable of 'fixing' imperfections in the underlying RL algorithm.
B
Extended Related Works
The bi-level objective draws inspiration from the seminal work of Singh et al. [52, 53] that provides an optimal-rewarde rewards framework for an agent. Prior works have built upon 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 various aspects of leveraging auxiliary rewards Taux: while staying robust against its misspecification.
Apart from specifying auxiliary rewards Taux directly, other techniques for reward specification include linear temporal logic [35, 11, 65, 39] reward machines [29-31] that allow exposing the reward functions as a white-box to the agent.
Recent works also explore y that is state-action dependent [63, 48), or establishes connection between y 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
19

---

learning (RL) in various forms [28, 62, 69, 71]. Initially designed for stochastic gradient descent, these methods have been extended other optimizers such as Adam [32] and RMSprop [25], treating them as differentiable counterparts [20].
C
Algorithm
In this section diseuss the algorithm for the proposed method. As the proposed method does behavior alignment reward function's s implicit optimization, we name BARFI. Pseudo-code for BARFI is presented in Algorithm 5. We will first build some preliminaries to understand the concepts.
C.1 Vector Jacobian Product
Let us assume that, I E N E IR" f(r", E R. Then, we know that 0f w)/0x E R 0z, v)/Oy E R", "e v)/yy0 E Raxm. Let us also assume that we have a vector v E IRd, and if need calculate the following, we can pull the derivative outside as shown:
8 f(x, v) Bydr
dy
Rixx
R1
B
m
As we can see, the vector Jacobian product can be broken down into differentiating vector product but shifting the place of multiplication, in which case we assume that the gradient passes through w.r.t y and hence we don't 't ever have deal with large multiplications. Also note that the outer partial w.r.t can easily be handled by aut odiff packages. A pseudo-code is show in Algorithm 1.
Algorithm 1: Jacobian Vector Product
Input: f(x, v) E RR x R E R" e Rd f' t (zad(f(2.y)).)) jvp t gf grad = u)
Return: jvp
C.2 Neumann Series Approximation for Hessian Inverse
Recall, that for a given real number R, such that 0 S 1, we know that the geomertric series of this has a closed form solution, i.e..
1 B + B
3 +...
1- B
Similarly, given we have a value a such that B a, we can write -1 follows:
1- 1 8 82 33 ...
=1+ (1 a) + (1 a)² + (1 - a) + +...
(1
a = = (1 a) + (1 a)² + (1 - a) 3 +...
a = Σ (1- - a)
1=0
The same can be generalized for matrix, i.e., given matrix A E Idddd we can write A-1 as follows:
--
Σ I - AI-A).
2=0
20

---

Note for the above to hold, matrix A, where we represent e (ig(A) as the eigenvalues of matrix A, we should have the following condition hold, 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 1/ max e (AA). Then we can write the following:
A
?
A
-1
"
M((A)~
-1
Σ II (I- nA-
As n would always satisfy the above condition.
C.3 Neumann Approximation for Hessian Vector Product
Given have seen how approximate the Inverse of matrix without relying (dd) operations, through Neumann approximation, lets look what needs to be done for our updates. Recall that the update o, ip (8) and (15) were.
-1
aJ(0( (J((())))) 06
(A( 06
00(0,)) (p) a0(6, sp)
and
a(J(((O, p)) 133)
0J(0(6,))) +))
( 0.
O s), ())
0%
0+
00(0,
00(0, sp)
O+p
0+
H
B
Let us look closely at the update for ái and we can generalize the updates easily for the case of yp.
( v))
-
aΦ
We first look at how we can approximate the value of vH efficiently, as we can always make use of the Jacobian Vector product later to get (vH-1) A, as vH becomes vector. Let us assume wish to run the Neumann approximation up to n steps, i.e., we want to approximate H up to n order Neumann expansion,
=
Σ i=0
HH
(12)
Here we are assuming that the outer optimization for update (1) for the function J(O(s), sp) and the inner optimization which is represented by the update (3) is f(((), sp), o)
A0,, rp) = (J((())))) v))
(0) 0f ( k))
00
00
The most common form in which f(; B)) is usually defined is the following:
f(0.
Σ
log (To (ST A7))
Σ
p r (ST.A7)
TEB 8
t=
f=1
21

---

Algorithm 2: Vector Hessian Inverse Product for (8) i.e., vH
Input: 0, 0, 9, J, n. Dott. Don
f grad(J(0; ,, 0)
o' < y grad(f Φ, (p; Don). 0)
Let: to t- v. +
for ie [0, do
Di+1 < f/ grad(v', g,g.0,0grrddduuuuttt =v)
Pet1 + pi #+11
Return: 'Pp
// Approximation of cH in (12)
Similarly this can be defined for J, except making use of rp and problem defined ?
T
J(0: B)
log( (S7 AD)
Σ
r-'
13)
B
TEB
j=t
Finally, once we have vH we can use the Vector Jacobian Product to calculate (uH as described in Algorithm 3:
Algorithm 3: Update for o, i.e. (8) i... vH AA
Input: o. p. J. n. 17, Donl. Dan
f Algorithm 2 (8, Ø, p. J. f. 17, Dou, D
v' < (f, 0, p: Dot),
A t grad(v', Ø. grad outputs = v)
Return As
can similarly derive updates for yp Note we are not including the different forms of regularizers
Algorithm 4: Update for d. i.e. (15) i... wH 'B
Input: 8, o, 4, J. n. 17, Don Den
< Algorithm 2 (0, Φ. f0, J. n, 17, Don, Doa)
o' < grad(f(f. Φ, p), Doa),
Ap grad(v', ip, grad outputs = v)
Return Ag
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 converge to the current values of Φ s., 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 A00, Φ, (p) smoothens the objective J (0(6, :p)) with respect and sp. This is helpful gradual changes in T 7;p result in gradually changes in the fixed point for the inner optimization. Therefore, for computational efficiency, we initialize the policy parameters from the fixed-poine 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 0 are performed till (approximate) convergence. Note that only trajectories from past interactions are used and no new-trajectories are sampled for the inner optimization.
22

---

Algorithm 5: BARFI: Behavior Alignment Reward Function's Implicit optimization
Input: J,J,00, (00,00,0000 agp ?) n, optim, F., N₂, No,
Initialize: (T, Thp
Initialize: optimg f- optoimim optim(a)); opt im + optim(as), optim, + optim(a.)
Doll +
Collect a batch of data for warmup period
for e E [1, No)
Generate using Tg
Append T. Dan
# Initial training steps using warmup data
f E [0, NN; + No) do
Sample batch of trajectories B from Doff
Update policy
- +-0 optimgggrad(((0) B),
Start reward alignment
for e [No; F.) do
Collect batch of on-policy data
12
Don <-
13
for jE 0,,6) do
14
Generate trajectory Te+s using To and append in Don
15
Update rg and % A + Algorithm 3(0, o, p0, J, f. n. Douts D Ap + Algorithm 4(0, Φ. 40, J. f, n, 17, Dotl Dea)
16
17
18
Φ < Φ + optimm + optooiiimm.. optim
19
20
Don + Don + Don
Learn policy for new reward function, initializing from the last
21
for i E [0, N₂) do
Sample batch of trajectories B from Doff
# Update policy
0 + opt optimyggradd(()) B),
Line 16-17
Line 16-17
(do, sP)
(Φ), PP))
(d2, (22)
Oo
01
0₂
03
Line 8-10
Line 21-23
Line 21-23
Figure 4: Algorithm Flow: The change in different parameters
In Lines 13-14, new batch Don of data is sampled using the policy returned by the inner-optimization process. This data is used to compute aJ (0(6, :)) (0((), sp) Existing data Dair that used in the inner-optimization process is then used to compute @0(ú, sp) /8/ and (0 /0/p Using these in (8) and (15), the parameters for T and Yip are updated in Lines 16 and 17. respectively.
Finally, the new data Don is merged into the existing data Donf and the entire process continues.
C.5 Note on Approximation
An important limitation of the methods discussed above is that 8(Ø, (p) is considered such that (0(Ø, ()) 0, s) = 0, i.e.. the Alg is run to convergence. In practice, we only execute Alg for predetermined number of update steps that need not result in convergence to optimum exactly. However, the impact of this approximation can be bounded by assuming convergence to an 6- neighborhood of the optima [49]. Furthermore, due to smoothness in the functional space, slight
23

---

changes to & and & should result in slight shifts in the optimum ( p)). The continuity property allows for improvements in the optimization process: suffices to initialize the parameters of each inner-loop optimization problem with the final parameters of the approximate fixed point solution, (ø, vp), identified the previous iteration of the inner loop. The complete resulting algorithm 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 BARFI described above, based 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 unrol led. The main difference, in this case, is that when computing the term (b) in (4), we now consider each the sequence of steps depends on the
inner update step until the point ((, (p) is reached- specific Al g used for the inerr updates. Details are deferred to Appendix C. Notice that this approach results in path path-wise optimization process that can be more demanding in terms of computation and memory. We further discuss this issue, and demonstrate the efficacy this alternative approach, in the empirical analyses section (Section ??).
D Smoothing the objective
To understand why [[b( ()) might be ill-conditioned is to note that, often a small perturbation in the reward function doesn't II necessarily lead to a change in the corresponding optimal policy. This can lead lack of gradient directions in the neighborhood of go for gradient methods to be effective. This issue can he addressed by employing common regularization techniques like 1.2 regularization of the policy parameters entropy regularization for the policy We discuss two ways regularize the objective in the upcoming sections.
D.1 L2 Regularization
To understand how severely ill-conditioned J(O(o, up) can be, notice that small perturbation in the reward function often does not change the corresponding optimal policies the outcome of policy optimization algorithm Alg. Therefore, the parameters the behavior alignment reward are perturbed from & to ø', it may often be that [(((,, v)) (0(o,, , sp)) and this limits any gradient based optimization for ó as ( f)) is 0. Similarly, minor perturbations in may result in no change in ((( :p)) either.
Fortunately, there exists remarkably simple solution: incorporate regularization for the policy parameters 0 in objective for Alg in the inner-level optimization. For example, the optimal policy for the following regularized objective E,,, EE (S, Ad) à |l@l varies smoothly to trade-off between the regularization value of 9 and the magnitude of the performance characterized by (ra. Tp), which changes with the values of na and Yp. See Figure 5 for an example with L2 regularization.
D.2 Entropy Regularized
In Section D.1, smoothing of J(O(d, (p) done by using 1.2 regularization the policy parameters 0 in the nuner-optimization process. However, alternate regularization methods can also be used. For example, in the following we present an alternate update rule for 0 based entropy regularization,
T Σ
T
A0, d. (p) = ED
(St, A,)
Σ j=t
2/
1-t (ro(S), A3) AIn No (S,, Az))
TTii regularization is performed avoid noninvertible Hessian we had discussed in Section 5

---

275
2.50
225
2.4 m
8
130
+
125
18
3 2 ici
2
Jcl
Figure 5: (Left)) A bandit problem, where the data is collected from policy / that samples action A mostly. (Middle) Each point on the 3D surface corresponds to the performance of 0(ó. 1) returned by an Alg that uses the update rule A 9, ø. 1) corresponding to the value of T. for actions A and C' in the bottom axes; ro for action B is set to avoid another variable in 3D plot. Notice that small perturbation in T may lead to no sudden changes in [( 1)). (Right) Performance of (( 1) returned by an Alg that uses the update rule Aor Φ. 1) that incorporates gradient of the L2 regularizer. Vector fields in Figure were also obtained from this setup.
Notice that new update rule for & and &o can be obtained from steps (4) to (15) with the following A, B, and H instead, where for shorthand 0" =(Ø 9).
ra (S)
A = ED
1o (Se, As)
Σ 1=
7
O
B = Ep
Σ Vo- (St, As)
(ro(S), A,)
Aln (S), A,)) ()))
Op
Σ Orbo- (Se; Ar)
T Aetio (St, Ad)
I Vo- (t, (8,,,)))))
Ep
Σ 9/p
(re(S),, Aln (s,, (s)
00
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 RI [50], differentiable convex optimization neural-networks layers [4, 2]. to hyper-parameter optimization [37, 8, 12, 40). In this work, show how implicit gradients can also be useful to efficiently leverage auxiliary rewards ra and overcome various sub-optimalities.
Taking total derivative in (5) with respect to á
d o). 0A(0(0, 6)) dΦ
@A(0(), à, v) 00(0. v)
(13)
dæ
(0(.i))
Let us try to understand why the above true, considering the finite difference approach for this derivative,
dA(0(0, s),
( + deb s). & + do, v) - ((0(s) p),
lim [dø] +0
dø
dē
0-0 0
do
A ( + do, p) + dØ, :p) = ( p), Φ. (p) = 0, as ((:.)) defines convergence to fixed point.
By re-arranging terms in (13) we obtain the term (b) in (4),
00(0. pp 00
(a0(0(6) v). 02 ()) -0 00(0,))
(14)
25

---

By combining (14) with (4) we obtain the desired gradient expression for ø,
aJ p)) aJ(B(0, s))
@0(Φ (p) ( s0,, 0, sp)
a ( s)), o, sp
()(
dΦ
H
and following similar steps, can be observed that the gradient expression for yp, a(J(O(á, +)) 701"1) aJ(O():))) a0(0; p)
8 p).
(O( 07.
0Φ
00(0, .p)
dip
dip
where using 8" as a shorthand for 0(0, up) the terms A, B and H can be expressed as
= EDD
Σ Vh- (Se: Ad)
Σ %
Or (S),
B ED
Σ 1/ (St, Ac)
Σ
T
Ori-
ro(S,, A1)
dep
(15)
H Ep
Odho (Se, Ax)
Σ
ro
-A.
00*
These provide the necessary expressions for updating & and oo 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- sonjugate-gradients Neumann series can thus be used to make it more tractable [40]. In our experiments, we made use of the Neumann approximation the Hessian Inverse vector product (AH- which requires the same magnitude of resources as the baseline policy gradient methods that we build on top off.
Being based on implicit gradients, call method BARFI, shorthand for behavior alignment reward function's 's implicit optimization. Overall, BARFI I iteratively solves the bi-level optimization specified in (2) by alternating between using (3) till approximate converge of Alg to 06, sp) and then updating Alg leverages only the past samples and does sample any new trajectories for the inner level optimization. Further, due to policy regularization which smoothens the objective as discussed in D, updates in ra and Tip 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 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 I and more details on the approximation techniques can be found in Appendix C.
Algorithm:
and 7/p- Importantly, being based on (3) for sample efficiency,
F Environment & Reward Details
The first environment is Grid (GW), where the start state is in the bottom left corner and goal state in the top right corner. The agent receives an rp of +100 on reaching the goal followed by termination of the episode. The second environment is MountainCan (MC) [58], wherein make use of the sparse reward variant, wherein the agent receives a +1 reward on reaching on top
of the hill and ( otherwise. The third environment is CartPole (CP) [16]. Finally, to assess the scalability we pick HalfCheetah-v4 from Mujoco (MJ) suite of OpenAl Gym [9].
For each environment, we define two auxiliarly reward functions. For GridWorld, we define the functions: y1 reux.e9 == (ss )g )2 which provides the negative L2 squared distance from the goal position, and r 3ux.39 == 50 x 1.S..... which provides an additional bonus of +50 to the agent along the desired path the goal state (i.e. the center states). In MountainCar the state is composed of two components: the position ær and velocity v. The first auxiliary reward function, r1 rlux. (8, a) == [v],
encourages a higher absolute velocity of the car, and the second, T.ux.Mc (s, a) == 1 Lign (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 hand crafted policy. CartPole can be solved using a Proportional Derivate (PD) controller [6], hence we tune PD controller, PD* S + A. to solve CartPole for the max possible return. We design two auxiliary
26

---

reward functions which make use of this PD controller. The first, T Taxx,cp a) == 5 x 1p (s)=a (1 (1- (_=Ap==<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, T Taux,ox a) == -T -p (s,a), encourages the agent to do the opposite. 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 exper- iments. We have heuristic-based reward functions (i.e., r aux I auxx, T axxx Reward func- tions that encode guess of an optimal policy (i.e., r Mux r1 (axxcc) and reward functions that change the optimal policy (i.e., r Faux,ca> r1 aaxx.c).. We also have rewards that only depend on states (i.e., Taux, Taux,cxx ca F Taax,ow: Meux,Me) as well as ones that depend on both state and actions (i.e., can test if BARFI can overcome misspecified auxiliary
rh Yaux,ec 71 "aax,ce cp: T ce). Therefore, reward functions and does not hurt performance when well-specified.
Mujoco Environment In this experiment, we investigate the scalability of BARF I in learning control policies for high-ding -igh-dimessioal state spaces with continuous action spaces. In HalfCheetah-v4 the agent's 's task is to move forward, and it receives reward based on its forward movement (denoted as rp). Additionally, there is a small cost associated with the magnitude of torque required for clal3). The weighting between the main reward and
action execution (denoted as (s, a) the control cost is pre-defined as c for this environment, and we form the reward as F(ss a) = rp(s,a) + Taux (s, a). We explore how an arbitrary weighting choice can cause the agent f'il in learning, while BARFI I 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 auxx, (s, a) == -cla|2?, uses the default setting, while the second weighting, denoted rh ruux,xu (s, a) == - 4c/aaB. employs scaled variant of the first weighting, Additionally, we implement the path-wise bi-level optimization variant i.c., BARFI Unrolled. In these experiments, keep the value of y fixed to isolate the agent's 's capability 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 approx- imator used TileCoding [58), which offers 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 Gaussian representation, where the policy outputs the mean of the multivariate Gaussian and we used 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 7 0.99. Hence to do this made yp single bias unit, initialized to 4.6, and passed through sigmoid (i.c., (4.6) = 0.99).
For GW, CP and MC Te is defined as below
(s a) = (s)) + (s)r, (s)r.
27

---

Wherein @1. 02; Oy are scalar outputs of 3-headed function, in this case simply a linear layer over the states inputs.
Whereas the case of MJ, we have
(5 01 + ry Φ38a
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 baseline to get acceptable performance for the baseline method.
G.2 Slyper-parameter Selection
As different make use of different function approximators the range of hyper- ppper-params can vary we talk about all the above over here.
Best- pest-performing Parameters for different methods and environments listed where
Table Hyper-parameters for GridWorld
Hyper Parameter
BARFI Value
REINFORCE Value Actor-Critie Value
ag
10-3
10-3
1 10-3
as
10
ap
5 x 10
optim
RMSprop
RMSprop 0.25
RMSprop 0.25
A
0.25
X
0.0625
X
4.0
Buffer
1000
Batch Size
1 0.0005
3
n
5
No
150
N
15
Table 4: Hyper-parameters for MountainCar
Hyper Parameter
BARFI Value
REINFORCE Value
Actor-Critic Value
0g
0.015625
0.125
0.03125
0.0625
Oy
0.0625
optim
RMSprop
RMSprop
RMSprop
A
0.0
0.0
0.25
A
0.0
A
0.25
Buffer
50
Batch Size
1
0.001
3
5
No
50
N;
15
Hyperparameter Sweep: Here we list the details about how we swept the values for different hyper-params. We used PyTorch [46] for all our implementations. We usually used an optimizer between RMSProp Adam with default parameters provided in Pytorch. For as E (5 x
28

---

Table 5: Hyper-parameters for CartPole
Hyper Parameter
BA AFRFI Value
REINFORCE Value
Actor-Critic Value
ag
10-3
x 10-3
5 x 10
as
10
ap
5 X 10-3
optim
RMSprop
RMSprop 1.0
RMSprop
A
1.0
0.0
0.0
X
4.0
Buffer
10000
Batch Size
1 0.0005
3
n
5
No
150
N;
15
Table 6: Hyper-parameters for Muloco
Hyper Parameter
BARFI Value
REINFORCE Value
Actor-Critic Value
ag
7.5 x 10-5
5 x 10
2.5 x 10
2.5 10-3
0.0
optim
Adam
Adam
Adam
A
0.0625
X
0.0
Buffer
50
Batch Size
1
1
0.0005
3
5
No
30
N:
15
10-3 105 10-3, x 10-3, 5 x 10-1 10-4, 10-4, 100-,,, 10-5) we use similar ranges for Olg. Otp (which tend to be larger). For Xo and As, we sweeped from |0, 0.25, 0.5, 1.0) and for X we sweeped from 0.,0.25,14 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 I across different methods as we can see from the tables above.
[1, 3, 5]. n [1, 5]. N; E [1, 3, 6, 9, 12, 15), [1 10-3,5 5,5 x 10-1, 1 x 10 10 No E [30, 50, 100, 150). buffer E [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 cluster where the CPU class is Intel Xcon Gold 6240 CPU @2. 22.6OGHz. The total compute required for GW was around 3 CPU years?, CP also required around 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 GB of memory per thread.
71 CPU year Compute equal to running CPU thread for year.
29

---

Extra Results & Ablations
H.1 Experiment on partially misspecified Taux
In these set of experiments, we consider the case where auxiliary reward signals could be helpful only in a few- new-possibly arbitrary- state-action pairs. In general, we anticipate that solutions 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 Faux in the GridWorld domain in which the auxiliary reward is misspecified for subset of states near the starting position. Mcanwhile, it is still well-specified for states near the goal (Figure 6 (a)). Figures (b) and illustrate the learned and the weighting on, showcasing the agent's ss ability to depict the expected behavior described above.
(S)
Φ3(s)
70
-0.0
-0.0
30
-0.15
20
10
-0.5
-030
(a) Faux (s)
(b) Learned (s(s), the weighting on ra)x(s)
(s(s)
Learning Curve
101
40
Return
20
200
400
600
800
1000
Fpisodes
Barfi
Naive Addition
(d) Learning performance
(c) The net reward induced i.e.. (s()))..n())
rp(s)
Figure 6: 40 random seeds were used to generate the plots. The starting state is at the bottom left and the goal state is at the top right corner. The primary reward rp is +100 when the agent reaches the goal and 0 otherwise. blue region Taux provides value equal to the Manhattan distance, thereby incentivizing the agent to stay near the start), and partially well specified (in the red region, it is the negative Manhattan distance and encourages movement towards the goal). (b) The state-dependent weighting @s(s) learned by BARFI negates the positive value from Faux near the start state. (c) The effective reward function rp (s)) + (s)(s)aaux(6 learned by BARFI. (d) Learning curves for BARFI, and the baseline that uses a naive addition (r) + Taux (s)x(()) of the above auxiliary reward.
(a) A state-dependent Teux function that is partially misspecified (in the
30

---

GridWorld
CartPole
MountainCar
Mujoco
20
600
200
1940
2000
222
(03)
2000 C0CO
10000
GrrdWorld
CartPole
MountainCar
1922
2001
420
20
333
-000
2000
2000
4000
6000 2000 Episodes
Episodes
(4)
(6)
Actor Critic
Actor Crilic Potential
Barfi Unrolled
REINFORCE
REINFORCE Potertlel
This figure illustrates the aux return collected by agent based on the learned rig the curves are chosen based on best-performing curves on
Figure 7:
Return induced by learned reward functions:
Tp, and averaged over 20 runs (except 40 for GW). (a) Top Taxx,c>> Bottom r2 Taux,ck* (b) Top
"P Bottom - aux,ce (c) Top - aux,,c- Bottom - aux, (d) Top copp-raxe -Taxx Bottom aux.
H.2 Return based learned ns and 7ip
Figure and Figure 8 summarize the achievable return based on 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 good performance rp- whereas BARFI does achieve an appropriate return on ras but is also able to successively decay ? as the learning progress across different domains. Particularly notice Figure (a) Bottom, where REINFORCE does optimize aux return a lot, but actually fails solve the problem, 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.
31

---

GridWorld
CartPole
MountainCar
1.0
0.990
0.9
0.989
7
0.8
0.7
0.088
0.4
0.2
500
1000
1000
GridWorld
CartPole
MountainCar
0.990
0.9
0.088
0.8
0.5
0.806
0.7
0.4
5
500
1000
1500
2000
5
200
400
6:00
800
1000
500
3008
1500
2006
Epieodes
Epieodes
Epieodes
Actor Critic
Actor Critic Potential
Barfi
Barfi Unrolled
NFORCE
REINFORCE Potential
Figure 8: Learned discounting Yit This figure illustrates the learned for BARFI and normal y for other methods, the curves chosen based on best-performing curves Tp, and averaged over 20 runs (except 40 for GW). (a) Top Toux,ew- Bottom 22 (b) Top - ""ax,cc=* Bottom 22 -uaux.ce (c) Top- 2 -2 Bottom y1 aux.Mc- Mujoco is not included the was learned in that case. We can observe that the agents start to learn to decay at the appropriate pace. Note that the curves for methods other than BARFI and BARFI Unrolled are overlapping as the baselines don't change the value of 7.
32

---

98
57-
Returns
97
96-
96
96
95
94
98-
92
92
91.
A6
Xup
100
90
Returns
97
98
80
94
80
93-
92
2.4
0.9
Ok
Olip
94
97
96-
94
88
83
86
25
8
25
22 2:57
2.15
2-30
2-10
2.9
Figure 9: Sensitivity Curves: The of graphs representing the sensitivity of different hyper- params keeping all the other params fixed. The sensitivity is for 3ARFI in GW with r cas le.. the misspecified reward. We choose the best-performing parameters and vary each parameter to its influence. The curves are obtained for 50 runs (seeds) in each case, and error bars are standard errors. We can notice that og and 0g values. Ao.e.p tends to help but doesn't t really influence a lot terms of its magnitude, except larger values of X seem to do better. Smaller values of seems work fine, hence something around 5 10-4, x 10-3 usually should suffice. n, 8 can be chosen around 5 and 3, and usually workout fine. We also defined N₂ 5 x this case.
have large influence, and tend to stay around similar
33


Reviews:
[{"review_id": "K868FICrHu", "reviewer": "Reviewer_8Eto", "summary": "RL 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.", "questions": "n/a", "limitations": "The authors address some limitations and extensions in their conclusion.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This 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.", "weaknesses": "There 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and for acknowledging the originality and quality of our proposed method, including its formal properties, as well as the strengths of our experiments. We are glad you recognized the practical significance of the problem and the novelty of our work. Your positive feedback on the paper are encouraging and greatly appreciated.\n\n\n### Possibly Relevant Related Work\nThank you for the suggestion to include further discussion on methods for learning from proxy rewards, such as the paper by Hadfield-Menell et al. (2020). This approach is related and relevant to our goals. It investigates the problem of inferring a designer's \"true\" underlying objective/reward based on a proxy reward function, an intended decision problem (e.g., an MDP), and a set of possible reward functions. While this is not precisely our objective, their paper sheds light on possible new and improved methodologies for incorporating information into reward functions, as well as possible connections between our work and inverse reward design techniques. In light of your suggestion, we will extend our literature review and add further discussion on methods for reward inference, as well as on reward function identification under various settings, such as when given demonstrations or proxy rewards.\n\n\n### References\nHadfield-Menell, D., Milli, S., Abbeel, P., Russell, S., & Dragan, A. (2020). Inverse Reward Design (arXiv:1711.02827). arXiv. https://doi.org/10.48550/arXiv.1711.02827\n"}, {"review_id": "ddAV5slyt7", "reviewer": "Reviewer_BidQ", "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?\n\nIt 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$?\n\nSuggestions related to language and typos:\n1. On Line 157 - \"due its\" -> \"due to its\"\n2. On Line 127 - $\\theta_3$ should be $\\phi_3$?\n\nReferences:\n1. Potential-based shaping and q-value initialization are equivalent, Wiewiora (2003)\n\nEDIT (16 Aug 2023): updated score from 7->8.\n", "limitations": " ", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": " ", "weaknesses": " ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's valuable and positive feedback and are encouraged by their appreciation of the quality, relevance, and soundness of our proposed method.\n\nBelow, we provide individual responses addressing the reviewer's comments.\n\n### On using potential-based reward shaping in deep RL settings\nThe reviewer is correct: the seminal paper by Ng, Harada, and Russell (1999) discusses the properties of reward shaping only in the tabular setting. Indeed, applying potential-based shaping to problems with function approximation may—as the reviewer suggests—lead to outcomes different than those expected in the tabular case. As an example, the suggested initialization value technique by Wiewiora (2003) is known not to directly translate to the function approximation setting. Having said that, from a practical perspective, the differences between the outcomes of applying reward shaping in the tabular case and in the value function approximate settings appear to be less significant when rich function approximators, like the ones used currently in deep RL. Regarding the properties of shaping in more general settings, we emphasize that while the existing literature on this issue does not provide definite answers about the properties of shaping for the case of *arbitrary* gradient-based deep RL algorithms, in our paper we show  (Property 1) that for policy gradient algorithms such as REINFORCE, the expected update doesn’t change, but can potentially increase the variance. We do appreciate the reviewer's comments and questions regarding a discussion on whether—and to what point—potential-based reward shaping can be applied in deep RL settings. We will update the paper to better discuss these points.\n\n\n### BARFI's performance under partially misspecified auxiliary reward functions\nThis is a great point. The reviewer is correct that the more general setting where auxiliary reward signals could be helpful *only* in a few—possibly arbitrary—state-action pair is indeed an interesting property to investigate.\n\nIn general, we anticipate that solutions in this scenario would be such that $\\phi_3(s,a)$  assigns weightings, allowing the agent to avoid regions where $r_\\texttt{aux}$ might be misspecified. Meanwhile, the agent would still make use of the places where $r_\\texttt{aux}$ is well specified and useful. \n\nBARFI's mathematical formulation is expressive enough to facilitate solutions that tackle the challenging setting brought up by the reviewer. Thus, it doesn't require any modification in terms of parameterization and optimization approach to handle this scenario.\n\nTo complement this hypothesis and further investigate the setting suggested by the reviewer, we performed additional experiments using a modified grid-world environment. The results are available in the attached PDF. Specifically, we updated the auxiliary reward function in this domain (as depicted in Figure 1(a)) so that $r_\\texttt{aux}$ is misspecified for a subset of states near the starting position. Meanwhile, it is still well-specified for states near the goal. Figures 1(b) and (c) illustrate the learned $r_\\phi$ and the weighting on $r_\\texttt{aux}$, showcasing the agent's ability to depict the expected behavior described above. \n\n### Feedback on the presentation\nWe thank the reviewer for their attention to detail and for the suggestions on how to improve the presentation. We will address them in an updated version of the paper.\n\n\n\n\n### References\nNg, A. Y., Harada, D., & Russell, S. (1999). Policy invariance under reward transformations: Theory and application to reward shaping. In ICML (Vol. 99, pp. 278-287).\n\nWiewiora, E. (2003). Potential-based shaping and Q-value initialization are equivalent. Journal of Artificial Intelligence Research, 19, 205-208.\n"}, {"review_id": "FfS6f2zRBl", "reviewer": "Reviewer_GcHf", "summary": "This 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.", "questions": "See the weaknesses.", "limitations": "See the weaknesses.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. 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", "weaknesses": "1. 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's valuable feedback and acknowledgment of BARFI's general capabilities of dealing with reward misspecification, addressing various potential sub-optimalities introduced by algorithmic design choices, and novel use of implicit gradient techniques to more efficiently optimize reward functions compared to traditional path-wise meta-optimization methods.\n\n\n### Clarification on Figure 1 \nThank you for your suggestion of an alternative, clearer way of discussing and interpreting Figure 1 to provide further intuition on why and how BARFI changes the direction of the gradients. Extending the paper with this discussion will indeed help better depict BARFI’s dynamic policy optimization procedure. We will update this figure in the final version of the paper by adding arrows indicating how the gradient directions may change over time. \nIn terms of additional discussion, we will further clarify that our method can modify the direction of gradients for the inner optimization by adjusting the learned reward function. We hope this explanation, combined with a complementary figure/diagram, helps clarify the mechanism.\n\n### Experimental Results\nThank you for highlighting the need to clarify the specific properties of BARFI investigated in Figures 2 and 3. These figures empirically support our primary claims about BARFI's robustness against potential inaccuracies in auxiliary reward functions, a unique feature of our method described in the Abstract (lines 13--15). In the Introduction (lines 64--66), we further outline our focus on robustness against potential inaccuracies in auxiliary rewards, rather than solely on the goal of accelerating learning.\n\nFigures 2 and 3 effectively showcase our method's success with respect to this goal. For example, in the top left of Figure 3, various methods, including BARFI, perform well with well-defined auxiliary rewards. However, when rewards are misspecified, as in the top right of Figure 3, conventional RL methods (such as REINFORCE and Actor-Critic) often fail, leading to *significantly* suboptimal solutions. BARFI, by contrast, consistently achieves near-optimal performance, thus demonstrating that its performance is not sensitive to whether (or how well) the auxiliary rewards were specified. \n\nWhen discussing these figures, we will better emphasize their main takeaway message: that predicting the success of an arbitrary RL method under particular (possibly misspecified) auxiliary reward functions is challenging and often infeasible. BARFI remains consistently reliable, regardless of auxiliary reward quality. This robustness is particularly vital in high-stakes applications such as medical uses of RL, where poor performance can lead to significantly risky outcomes.\n\nThe experiments in Figures 2 and 3, therefore, serve to illustrate BARFI's robustness and capability of consistently achieving near-optimal performance, regardless of the provided auxiliary reward. This distinguishes it from other techniques, including those based on reward shaping. We will update the paper to better emphasize and discuss these points, as well as how they strongly support BARFI's unique strengths.\n\n### Alternative Reward Shaping Approaches\nThe reviewer raises the concern regarding the existence of any alternative well-known and principled way of performing reward shaping, other than via potential-based reward shaping algorithms. In that case, the reviewer argues that we should compare BARFI with them. To the best of our knowledge, however, and after performing a further literature review, no other principled reward shaping methods exist (that are capable of ensuring that the optimal policy will remain unchanged) other than potential-based reward shaping techniques. Kindly let us know if you are aware of any alternative principled techniques of this nature. We will further discuss this point in the paper.\n\n### Visualization of Results\nPlease let us know if we have misunderstood your suggestion, but Figures 1 and 4 were specifically designed to offer insight into the behavior of BARFI's optimization process, as well as the type of reward function that it autonomously identifies to ensure robustness. \n\nFurthermore, to address the reviewer's suggestion of providing further ways of visualizing the solutions produced by BARFI, we have (in response to a comment by reviewer BidQ) performed additional experiments with a novel type of auxiliary reward function: one that is only helpful in a few—possibly arbitrary—state-action pairs, while being misspecified in other. This allows us to directly visualize how the reward functions learned by BARFI adapt to circumvent regions of the state space where auxiliary rewards are misspecified while prioritizing areas where auxiliary rewards are helpful. The results of this experiment are shown in Figures 1(b), (c), and (d), of the attached PDF file.\n\n### Interpreting Figure 9 \nWe appreciate your attention to Figure 9. The absence of distinct curves is due to the overlapping of many curves, as all algorithms (except for BARFI) use a constant value of $\\gamma$, making them appear indistinguishable. Though it may seem that only two colors are shown, the figure includes colors for all algorithms. We will clarify in the paper that the curves overlap due to the nature of the experiment. \n\n### Method Flowchart\nWe appreciate your feedback regarding augmenting the paper with additional, alternative ways of describing our method. Figure 6 of the appendix already provides a simplified flowchart of how our method works. We will address your suggestion by including more details in Figure 6 (as well as an additional corresponding discussion) in order to produce a more thorough flowchart representation of our algorithm.\n\n\n### References:\nZou, H., Ren, T., Yan, D., Su, H., & Zhu, J. (2019). Reward shaping via meta-learning. arXiv preprint arXiv:1901.09330.\n\n"}, {"review_id": "i4aCBMQTOH", "reviewer": "Reviewer_tPX3", "summary": "This 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.", "questions": "Please refer to weaknesses for details.", "limitations": "Please refer to weaknesses for details.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and for acknowledging the originality and quality of our proposed method. We are glad to hear your positive comments on the practical significance of the problem we address and the novelty of our work. Your positive feedback is encouraging and greatly appreciated.\n\n\n### High-Dimensional Tasks\nThank you for bringing this up. BARFI can, indeed, be seamlessly adapted to vision-based tasks. We will emphasize in the paper that our large-scale experiments on MuJoco employed deep neural networks (multi-layer perceptrons) to model both policy and reward functions. We will clarify, thus, that BARFI can be seamlessly adapted to vision-based tasks by simply substituting these MLPs with ConvNets.\n\n\n### Hyperparameters\nExcellent point. Please notice that the value of $\\alpha_\\theta$ we reported is simply $2^{-6}$, which can be expressed as $0.015625$ in standard decimal notation. This choice is not arbitrary; rather, it was obtained through a standard hyper-parameter search procedure described in Chapter 3 of Sutton and Barto. This procedure involves exploring candidate step sizes in the form of $2^{-x}$, where $x \\in {2, 4, 6, 8}$. Among these options, $2^{-6}$ yielded the best results. We will clarify this process in the main text to ensure readers understand that hyperparameters were fine-tuned using established techniques. Furthermore, we will further emphasize that we also *do* provide sensitivity analysis curves for various hyperparameters, including $\\alpha_\\theta$ (please see Figure 10 of the appendix).\n\n\n### Experimental Results\nThank you for highlighting the need to clarify the specific properties of BARFI investigated in Figures 2 and 3. These figures empirically support our primary claims about BARFI's robustness against potential inaccuracies in auxiliary reward functions, a unique feature of our method described in the Abstract (lines 13--15). In the Introduction (lines 64--66), we further outline our focus on robustness against potential inaccuracies in auxiliary rewards, rather than solely on the goal of accelerating learning.\n\nFigures 2 and 3 effectively showcase our method's success with respect to this goal. For example, in the top left of Figure 3, various methods, including BARFI, perform well with well-defined auxiliary rewards. However, when rewards are misspecified, as in the top right of Figure 3, conventional RL methods (such as REINFORCE and Actor-Critic) often fail, leading to *significantly* suboptimal solutions. BARFI, by contrast, consistently achieves near-optimal performance, thus demonstrating that its performance is not sensitive to whether (or how well) the auxiliary rewards were specified. \n\nWhen discussing these figures, we will better emphasize their main takeaway message: that predicting the success of an arbitrary RL method under particular (possibly misspecified) auxiliary reward functions is challenging and often infeasible. BARFI remains consistently reliable, regardless of auxiliary reward quality. This robustness is particularly vital in high-stakes applications such as medical uses of RL, where poor performance can lead to significantly risky outcomes.\n\nThe experiments in Figures 2 and 3, therefore, serve to illustrate BARFI's robustness and capability of consistently achieving near-optimal performance, regardless of the provided auxiliary reward. This distinguishes it from other techniques, including those based on reward shaping. We will update the paper to better emphasize and discuss these points, as well as how they strongly support BARFI's unique strengths.\n\n### BARFI Warmup Stage\nWe appreciate your observation and will further improve the discussion on BARFI's warmup stage. In Equations 2 and 3, the bi-level objective is defined as the point of convergence with respect to the inner-level objective. When randomly initializing the parameters of the initial policy considered by the algorithm, the warmup phase aids in achieving approximate convergence to the inner-level objective. This is advantageous since subsequent inner-level optimization steps benefit from (a) being able to initialize the policy search process from the previous convergence point; and (b) adapting to gradual changes to the inner-level objective. Empirically, we have observed that this allows BARFI to rapidly identify near-optimal policies in the inner objective. Omitting the warmup phase precludes BARFI from efficiently optimizing the inner-level objective, thereby potentially leading to lower performance. We will update the paper with a more thorough discussion of the properties and requirements of having a warmup phase.\n\n### Presentation of Empirical Results\nYour feedback on how to better organize the discussion of our empirical results is greatly appreciated. We will follow your suggestion to break down larger paragraphs into shorter ones, and we will move some details about the task descriptions to the footnotes or the appendix.\n\n### Recent Work\nThank you for directing us to this potentially relevant paper. It addresses a different problem than ours but also uses similar tools from bi-level optimization. We will expand our Related Work section to discuss this relevant related technique.\n"}]

Author rebuttal:
### Additional figures/experiments (as suggested by reviewer BidQ)



Meta-review: This paper tackles the problem of reward shaping, highlighting issues with the traditional potential-based reward shaping, and proposing an alternative based on bi-level optimization.  Reviewers found the work to be high-quality, the topic interesting and significant, and the approach principled, so I am happy to recommend acceptance.
Decision: Accept (spotlight)