SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dhawal Gupta*  
University of Massachusetts
Yash Chandak*†  
Stanford University
Scott M. Jordan†  
University of Alberta
Philip S. Thomas  
University of Massachusetts
Bruno Castro da Silva  
University of Massachusetts
Designing reward functions for efficiently guiding reinforcement learning (RL) agents toward specific behaviors is a complex task. This is challenging since it requires the identification of reward structures that are not sparse and that avoid inadvertently inducing undesirable behaviors. Naively modifying the reward structure to offer denser and more frequent feedback can lead to unintended outcomes and promote behaviors that are not aligned with the designer's intended goal. Although potential-based reward shaping is often suggested as a remedy, we systematically investigate settings where deploying it often significantly impairs performance. To address these issues, we introduce a new framework that uses a bi-level objective to learn behavior alignment reward functions. These functions integrate auxiliary rewards reflecting a designer's heuristics and domain knowledge with the environment's primary rewards. Our approach automatically determines the most effective way to blend these types of feedback, thereby enhancing robustness against heuristic reward misspecification. Remarkably, it can also adapt an agent's policy optimization process to mitigate suboptimalities resulting from limitations and biases inherent in the underlying RL algorithms. We evaluate our method's efficacy on a diverse set of tasks, from small-scale experiments to high-dimensional control challenges. We investigate heuristic auxiliary rewards of varying 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.
In this paper, we investigate the challenge of enabling reinforcement learning (RL) practitioners, who may not be experts in the field, to incorporate domain knowledge through heuristic auxiliary reward functions. Our goal is to ensure that such auxiliary rewards not only induce behaviors that align with the designer's intentions but also allow for faster learning. RL practitioners typically model a given control problem by first designing simple reward functions that directly quantify whether (or how well) an agent completed a task. These could be, for instance, functions assigning a reward of  $+1$  iff the agent reaches a specified goal state, and zero otherwise. However, optimizing a policy based on such a sparse reward function often proves challenging.
*Both authors contributed equally to this work. † Work done while at the University of Massachusetts. Corresponding author: Dhawal Gupta (dgupta@cs.umass.edu).
Figure 1: Auxiliary rewards can be used to convey to 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 a visual illustration of a toy example depicting how the proposed method works. The star represents the optimal policy, and the red dot represents the fixed point of a policy optimization process under a "sub-optimal" heuristic; i.e., one that, when naively combined with  $r_p$ , induces behaviors different from those under the optimal policy for  $r_p$ . (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 distance to the goal. However, the effectiveness of using auxiliary reward functions largely depends on the problem's complexity and the designer's skill in crafting heuristics that, when combined with the original reward function, do not induce behaviors different than the ones originally intended [26, 27].
Existing methods like potential-based reward shaping [41] aim to incorporate domain knowledge without misaligning the behaviors induced by the resulting combined reward functions. However, as we discuss in Section 3, potential-based shaping has several limitations: (i) it is restricted to state-based functions; (ii) it amounts to a different initialization of the  $q$ -function; (iii) it does not alter policy gradients in expectation; and (iv) it can increase the variance in policy gradient methods.
To address these challenges, we introduce a scalable algorithm that empowers RL practitioners to specify potentially imperfect auxiliary reward functions. It ensures that the resulting optimization process will not inadvertently lead to unintended behaviors and that it will allow for faster learning. In particular, this paper addresses the following challenges:
(1) How to incorporate auxiliary reward information: We introduce a novel bi-level objective to analyze and automatically fine-tune designer-created auxiliary reward functions. It ensures they remain aligned with the original reward and do not induce behaviors different from those originally intended by the designer. Additionally, we formulate the problem to shape the optimization landscape, biasing our bi-level optimizer toward auxiliary reward functions that facilitate faster learning.
(2) How to use auxiliary reward to mitigate algorithmic biases: We show that our framework can automatically adjust how primary and auxiliary rewards are blended to mitigate limitations or biases inherent in the underlying RL algorithm (Section 4.1). For instance, many policy-gradient-based RL algorithms are subject to biases due to issues like discounting mismatch [59] or partial off-policy correction [51]. These biases can hinder the algorithm's ability to identify near-optimal policies.
(3) How to ensure scalability to high-dimensional problems: We introduce an algorithm that employs implicit gradients to automatically adjust primary and auxiliary rewards, ensuring that the combined reward function aligns with the designer's original expectations (see Figure 1). We evaluate our method's efficacy across a range of tasks, from small-scale to high-dimensional control settings (see Section 6). In these tasks, we experiment with auxiliary rewards of varying quality; some accelerate learning, while others can be detrimental to finding an optimal policy.
2 Notation
In this paper, we investigate sequential decision-making problems modeled as Markov decision processes (MDPs). An MDP is defined as a tuple  $(\mathcal{S},\mathcal{A},p,r_p,r_{\mathrm{aux}},\gamma ,d_0)$ , where  $\mathcal{S}$  is the state set,  $\mathcal{A}$  is the action set,  $p$  is the transition function,  $r_p:\mathcal{S}\times \mathcal{A}\to \mathbb{R}$  is the primary reward function,  $r_{\mathrm{aux}}$  :
$\mathcal{S} \times \mathcal{A} \to \mathbb{R}$  is an optional auxiliary reward function (possibly designed by a non-expert in machine learning, based on domain knowledge), and  $d_0$  is the starting state distribution. Let  $\pi_{\theta}: \mathcal{S} \times \mathcal{A} \to [0,1]$  be any policy parameterized using  $\theta \in \Theta$ . For brevity, we will often use  $\pi_{\theta}$  and  $\theta$  interchangeably. Let  $S_t$  and  $A_t$  be the random variables for the state and action observed at the time  $t$ . As in the standard RL setting, the performance  $J(\theta)$  of a policy  $\pi_{\theta}$  is defined as the expected discounted return with respect to the (primary) reward function,  $r_p$ ; i.e.,  $J(\theta) \coloneqq \mathbb{E}_{\pi}[\sum_{t=0}^{T} \gamma^t r_p(S_t, A_t)]$ , where  $T + 1$  is the episode length. An optimal policy parameter  $\theta^*$  is defined as  $\theta^* \in \arg \max_{\theta \in \Theta} J(\theta)$ . A popular technique to search for  $\theta^*$  is based on constructing sample estimates  $\hat{\Delta}(\theta, r_p)$  of the ( $\gamma$ -dropped) policy gradient,  $\Delta(\theta, r_p)$ , given an agent's interactions with the environment for one episode [58, 59]. Then, using  $\psi_{\theta}(s, a)$  as a shorthand for  $\mathrm{d}\ln \pi_{\theta}(s, a) / \mathrm{d}\theta$ , these quantities are defined as follows:
$$
\Delta (\theta , r _ {p}) = \mathbb {E} _ {\pi_ {\theta}} \left[ \hat {\Delta} (\theta , r _ {p}) \right] \quad \text {a n d} \quad \hat {\Delta} (\theta , r _ {p}) := \sum_ {t = 0} ^ {T} \psi_ {\theta} \left(S _ {t}, A _ {t}\right) \sum_ {j = t} ^ {T} \gamma^ {j - t} r _ {p} \left(S _ {j}, A _ {j}\right). \tag {1}
$$
3 Limitations of Potential Based Reward Shaping
When the objective function  $J(\theta)$  is defined with respect to a sparse reward function  $r_p$  (i.e., a reward function such that  $r_p(s, a) \coloneqq 0$  for most  $s \in S$  and  $a \in \mathcal{A}$ ), searching for  $\theta^*$  is challenging [24]. A natural way to provide more frequent (i.e., denser) feedback to the agent, in the hope of facilitating learning, is to consider an alternate reward function,  $\tilde{r}_{\mathrm{naive}} \coloneqq r_p + r_{\mathrm{aux}}$ . However, as discussed earlier,  $r_{\mathrm{aux}}$  may be a designer-specified auxiliary reward function not perfectly aligned with the objective encoded in  $r_p$ . In this case, using  $\tilde{r}_{\mathrm{naive}}$  may encourage undesired behavior. An alternative way to incorporate domain knowledge to facilitate learning was introduced by Ng et al. [41]. They proposed using a potential function,  $\Phi: S \to \mathbb{R}$  (analogous to  $r_{\mathrm{aux}}$ ), to define new reward functions of the form  $\tilde{r}_{\Phi}(S_t, A_t, S_{t+1}) \coloneqq r_p(S_t, A_t) + \gamma \Phi(S_{t+1}) - \Phi(S_t)$ . Importantly, they showed that optimal policies with respect to the objective  $\mathbb{E}\left[\sum_{t=0}^{T} \gamma^t \tilde{r}_{\Phi}(S_t, A_t, S_{t+1})\right]$  are also optimal with respect to  $J(\theta)$ .
While potential-based reward shaping can partially alleviate some of the difficulties arising from sparse rewards, Wiewiora [64] showed that  $q$ -learning using  $\tilde{r}_{\Phi}$  produces the exact same sequence of updates as  $q$ -learning using  $r_p$  but with a different initialization of  $q$ -values. In what follows, we establish a similar result: we show that performing potential-based reward shaping has no impact on expected policy gradient updates; and that it can, in fact, even increase the variance of the updates.
Property 1.  $\mathbb{E}[\hat{\Delta} (\theta ,\tilde{r}_{\Phi})] = \mathbb{E}[\hat{\Delta} (\theta ,r_p)]$  and  $\operatorname {Var}(\hat{\Delta} (\theta ,\tilde{r}_{\Phi}))$  can be higher than  $\operatorname {Var}(\hat{\Delta} (\theta ,r_p))$
All proofs 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  $\Phi$  cannot depend on actions [41], which restricts the class of eligible auxiliary rewards  $r_{\mathrm{aux}}$  and heuristics functions that may be used. Finally, notice that  $\Phi$  is designed independently of the agent's underlying learning algorithm. As we will show in the next sections, our method can autonomously discover auxiliary reward functions that not only facilitate learning but also help mitigate various types of algorithmic limitations and biases.
4 Behavior Alignment Reward Function
In this section, we introduce an objective function designed to tackle the primary challenge investigated in this paper: how to effectively leverage designer-specified auxiliary reward functions to rapidly induce behaviors envisioned by the designer. The key observation is that naively adding an auxiliary reward function  $r_{\mathrm{aux}}$  to  $r_p$  may produce policies whose corresponding behaviors are misaligned with respect to the behaviors induced by  $r_p$ . In these cases,  $r_{\mathrm{aux}}$  should be ignored during the search for an optimal policy. On the other hand, if  $r_p$  and  $r_{\mathrm{aux}}$  may be combined in a way that results in the desired behaviors, then combinations that produce frequent and informative feedback to the agent should be favored, as they are likely to facilitate faster learning.
To tackle the challenges discussed above, we employ a bi-level optimization procedure. This approach aims to create a behavior alignment reward by combining  $r_{\mathrm{aux}}$  and  $r_p$  using a parameterized function. Our method is inspired by the optimal rewards framework by Singh et al. [52, 53]. Let  $\gamma_{\varphi} \in [0,1)$  be
a discount rate value parameterized by  $\varphi \in \Gamma$ .<sup>2</sup> Let  $r_{\phi}: \mathcal{S} \times \mathcal{A} \to \mathbb{R}$  be a behavior alignment reward: a function of both  $r_p$  and  $r_{\mathrm{aux}}$ , parameterized by  $\phi \in \Upsilon$ , where  $\Upsilon$  and  $\Gamma$  are function classes. One example of a behavior alignment reward function is  $r_{\phi}(s, a) := f_{\phi_1}(s, a) + \phi_2 r_p(s, a) + \phi_3 r_{\mathrm{aux}}(s, a)$ , where  $f_{\phi}: \mathcal{S} \times \mathcal{A} \to \mathbb{R}$  and  $\phi := (\phi_1, \phi_2, \phi_3)$ . Let A1g be any gradient/semi-gradient/non-gradient-based algorithm that outputs policy parameters. To mitigate possible divergence issues arising from certain policy optimization algorithms like DQN [60, 1], we make the following simplifying assumption, which can generally be met with appropriate regularizers and step sizes:
Assumption 1. Given  $r_{\phi}$  and  $\gamma_{\varphi}$ , the algorithm  $\mathsf{Alg}(r_{\phi},\gamma_{\varphi})$  converges to a fixed point  $\theta$  which we denote as  $\theta (\phi ,\varphi)\in \Theta$  to emphasize its indirect dependence on  $\phi$  and  $\varphi$  through  $\mathsf{Alg}$ ,  $r_{\phi}$ , and  $\gamma_{\varphi}$ .
Given this assumption, we now specify the following bi-level objective:
$$
\phi^ {*}, \varphi^ {*} \in \underset {\phi \in \Upsilon , \varphi \in \Gamma} {\arg \max } J (\theta (\phi , \varphi)) - \lambda_ {\gamma} \gamma_ {\varphi}, \quad \text {w h e r e} \quad \theta (\phi , \varphi) := \operatorname {A l g} (r _ {\phi}, \gamma_ {\varphi}). \tag {2}
$$
Here,  $\lambda_{\gamma}$  serves as the regularization coefficient for the value of  $\gamma_{\varphi}$ , and  $\mathsf{A}1\mathsf{g}$  denotes a given policy optimization algorithm. Let, as an example,  $\mathsf{A}1\mathsf{g}$  be an on-policy gradient algorithm that uses samples to estimate the gradient  $\Delta (\theta ,r_p)$ , as in (1). We can then define a corresponding variant of  $\Delta (\theta ,r_p)$  that is compatible with our formulation and objective, and which uses both  $r_{\phi}$  and  $\gamma_{\varphi}$ , as follows:
$$
\Delta_ {\mathrm {o n}} (\theta , \phi , \varphi) := \mathbb {E} _ {\pi_ {\theta}} \left[ \sum_ {t = 0} ^ {T} \psi_ {\theta} \left(S _ {t}, A _ {t}\right) \sum_ {j = t} ^ {T} \gamma_ {\varphi} ^ {j - t} r _ {\phi} \left(S _ {j}, A _ {j}\right) \right]. \tag {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: it evaluates different parameterizations, denoted by  $\phi$ , for the behavior alignment reward function. These parameterizations influence the induced policy  $\theta(\phi, \varphi)$ , which is evaluated using the performance metric  $J$ . Recall that this metric quantifies the alignment of a policy with the designer's primary reward function,  $r_p$ . In essence, the outer-level objective seeks to optimize the behavior alignment reward function to produce policies that are effective according to  $r_p$ . This design adds robustness against any misspecification of the auxiliary rewards. In the inner-level optimization, by contrast,  $\mathrm{Alg}$  identifies a policy  $\theta(\phi, \varphi)$  that is optimal or near-optimal with respect to  $r_\phi$  (which combines  $r_{\mathrm{aux}}$  through the behavior alignment reward). In the inner-level optimization, the algorithm  $\mathrm{Alg}$  works to identify a policy  $\theta(\phi, \varphi)$  that is optimal or near-optimal in terms of  $r_\phi$  (which incorporates  $r_{\mathrm{aux}}$  via the behavior alignment reward). By employing a bi-level optimization structure, several benefits emerge. When  $r_{\mathrm{aux}}$  is well-crafted,  $r_\phi$  can exploit its detailed information to give  $\mathrm{Alg}$  frequent/dense reward feedback, thus aiding the search for an optimal  $\theta^*$ . Conversely, if  $r_{\mathrm{aux}}$  leads to sub-optimal policies, then the influence of auxiliary rewards can be modulated or decreased accordingly by the optimization process by adjusting  $r_\phi$ . Consider, for example, a case where the behavior alignment reward function is defined as  $r_\phi(s, a) := f_{\phi_1}(s, a) + \phi_2 r_p(s, a) + \phi_3 r_{\mathrm{aux}}(s, a)$ . In an adversarial setting—where the designer-specified auxiliary reward  $r_{\mathrm{aux}}$  may lead to undesired behavior—the bi-level optimization process has the ability to set  $\phi_3$  to 0. This effectively allows the behavior alignment reward function  $r_\phi$  to exclude  $r_{\mathrm{aux}}$  from consideration. Such a bi-level approach to optimizing the parameters of behavior alignment reward functions can act as a safeguard against the emergence of sub-optimal behaviors due to a misaligned auxiliary reward,  $r_{\mathrm{aux}}$ . This design is particularly valuable because it allows the objective in (2) to leverage the potentially dense reward structure of  $r_{\mathrm{aux}}$  to provide frequent action evaluations when the auxiliary reward function is well-specified. At the same time, the approach maintains robustness against possible misalignments.
Need for Outer Regularization: The bi-level optimization problem (2) may have multiple optimal solutions for  $\phi$  -including the trivial solution where  $r_{\mathrm{aux}}$  is always ignored. The goal of regularizing the outer-level objective (in the form of the term  $\lambda_{\gamma}\gamma_{\varphi}$ ) is to incorporate a prior that adds a preference for solutions,  $\phi^{*}$ , that provide useful and frequent evaluative feedback to the underlying RL algorithm. In the next paragraphs, we discuss the need for such a regularizer and motivate its mathematical form. First, recall that sparse rewards can pose challenges for policy optimization. An intuitive solution
$^{2}$ Our framework can be generalized to support state-action dependent discount rates,  $\gamma$ .
3"Misspecification" indicates that an optimal policy for  $r_p + r_{\mathrm{aux}}$  may not be optimal for  $r_p$  alone.
to this problem could involve biasing the optimization process towards denser behavior alignment reward functions, e.g., by penalizing for sparsity of  $r_{\phi}$ . Unfortunately, the distinction between sparse and dense rewards alone may not fully capture the nuances of what designers typically consider to be a "good" reward function. This is the case because a reward function can be dense and still may not be informative; e.g., a reward function that provides  $-1$  to the agent in every non-goal state is dense but fails to provide useful feedback regarding how to reach a goal state. A better characterization of how useful (or informative) a reward function is may be constructed in terms of how instructive and instantaneous the evaluation or feedback it generates is. We consider a reward function to be instructive if it produces rewards that are well-aligned with the designer's goals. A reward function is instantaneous if its corresponding rewards are dense, rather than sparse, and are more readily indicative of the optimal action at any given state. Reward functions that are both instructive and instantaneous can alleviate issues associated with settings with sparse rewards and long horizons. To bias our bi-level optimization objective towards this type of reward function, we introduce a regularizer,  $\gamma_{\varphi}$ . This regularizer favors solutions that can generate policies with high performance (i.e., high expected return  $J$  with respect to  $r_p$ ) even when the discount factor  $\gamma_{\varphi}$  is small. To see why, first notice that this regularizer encourages behavior alignment reward functions that provide more instantaneous feedback to the agent. This has to be the case; otherwise, it would be challenging to maximize long-term reward should the optimized alignment reward function be sparse. Second, the regularizer promotes instructive alignment reward functions—i.e., functions that facilitate learning policies that maximize  $J$ . This is equally crucial: effective policies under the metric  $J$  are the ones that align well with the designer's objectives as outlined in the original reward function,  $r_p$ .
4.1 Overcoming Imperfections of Policy Optimization Algorithms
The advantages of the bi-level formulation in (2) extend beyond robustness to sub-optimality from misspecified  $r_{\mathrm{aux}}$ . Even with a well-specified  $r_{\mathrm{aux}}$ , RL algorithms often face design choices, such as the bias-variance trade-off, that can induce sub-optimal solutions. Below we present examples to show how bias in the underlying RL algorithm may be mitigated by carefully optimizing  $r_{\phi}$  and  $\gamma_{\varphi}$ .
4.1.1 Bias in policy gradients: Recall that the popular "policy gradient"  $\Delta (\theta ,r_p)$  is not, in fact, the gradient of any function, and using it in gradient methods may result in biased and sub-optimal policies [44]. However, policy gradient methods based on  $\Delta (\theta ,r_p)$  remain vastly popular in the RL literature since they tend to be sample efficient [59]. Let  $\Delta_{\gamma}(\theta ,r_p)$  denote the unbiased policy gradient, where  $\Delta_{\gamma}(\theta ,r_p)\coloneqq \mathbb{E}[\sum_{t = 0}^{T}\gamma^{t}\psi_{\theta}(S_{t},A_{t})\sum_{j = t}^{T}\gamma^{j - t}r_{p}(S_{j},A_{j})]$ . We can show that with a sufficiently expressive parameterization, optimized  $r_\phi$  and  $\gamma_{\varphi}$  can effectively mimic the updates that would have resulted from using the unbiased gradient  $\Delta_{\gamma}(\theta ,r_p)$ , even if the underlying RL algorithm uses the biased "gradient",  $\Delta_{\mathrm{on}}(\theta ,\phi ,\varphi)$ , as defined in (3). Detailed proofs are in Appendix A.
Property 2. There exists  $r_{\phi}:\mathcal{S}\times \mathcal{A}\to \mathbb{R}$  and  $\gamma_{\varphi}\in [0,1)$  such that  $\Delta_{\mathrm{on}}(\theta ,\phi ,\varphi) = \Delta_{\gamma}(\theta ,r_p)$
4.1.2 Off-policy learning without importance sampling: To increase sample efficiency when evaluating a given policy  $\pi_{\theta}$ , it is often useful to use off-policy data collected by a different policy,  $\beta$ . Under the assumption that  $\forall s \in S, \forall a \in A, \frac{\pi_{\theta}(s,a)}{\beta(s,a)} < \infty$ , importance ratios  $\rho_{j} \coloneqq \prod_{k=0}^{j} \frac{\pi_{\theta}(s,a)}{\beta(s,a)}$  can be used to adjust the updates and account for the distribution shift between trajectories generated by  $\beta$  and  $\pi_{\theta}$ . However, to avoid the high variance stemming from  $\rho_{j}$ , many methods tend to drop most of the importance ratios and thus only partially correct for the distribution shift—which can lead to bias [51]. We can show (given a sufficiently expressive parameterization for the behavior alignment reward function) that this type of bias can also be mitigated by carefully optimizing  $r_{\phi}$  and  $\gamma_{\varphi}$ .