Heuristic-Guided Reinforcement Learning

We focus on discounted infinite-horizon Markov Decision Processes (MDPs) for ease of exposition. The technique proposed here can be extended to other MDP settings.³³3The results here can be readily applied to finite-horizon MDPs; for other infinite-horizon MDPs, we need further, e.g., mixing assumptions for limits to exist. A discounted infinite-horizon MDP is denoted as a 5-tuple $\mathcal{M}=(\SS,\mathcal{A},P,r,\gamma)$, where $\SS$ is the state space, $\mathcal{A}$ is the action space, $P(s^{\prime}|s,a)$ is the transition dynamics, $r(s,a)$ is the reward function, and $\gamma\in[0,1)$ is the discount factor. Without loss of generality, we assume $r:\SS\times\mathcal{A}\to[0,1]$. We allow the state and action spaces $\SS$ and $\mathcal{A}$ to be either discrete or continuous. Let $\Delta(\cdot)$ denote the space of probability distributions. A decision-making policy $\pi$ is a conditional distribution $\pi:\SS\to\Delta(\mathcal{A})$, which can be deterministic. We define some shorthand for writing expectations: For a state distribution $d\in\Delta(\SS)$ and a function ${V:\SS\to\mathbb{R}}$, we define $V(d)\coloneqq\mathbb{E}_{s\sim d}[V(s)]$; similarly, for a policy $\pi$ and a function $Q:\SS\times\mathcal{A}\to\mathbb{R}$, we define ${Q(s,\pi)\coloneqq\mathbb{E}_{a\sim\pi(\cdot|s)}[Q(s,a)]}$. Lastly, we define $\mathbb{E}_{s^{\prime}|s,a}\coloneqq\mathbb{E}_{s^{\prime}\sim P(\cdot|s,a)}$.

Central to solving MDPs are the concepts of value functions and average distributions. For a policy $\pi$, we define its state value function $V^{\pi}$ as $V^{\pi}(s)\coloneqq\mathbb{E}_{\rho_{s}^{\pi}}\left[\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})\right],$ where $\rho_{s}^{\pi}$ denotes the trajectory distribution of $s_{0},a_{0},s_{1},\dots$ induced by running $\pi$ starting from $s_{0}=s$. We define the state-action value function (or the Q-function) as $Q^{\pi}(s,a)\coloneqq r(s,a)+\gamma\mathbb{E}_{s^{\prime}|s,a}[V^{\pi}(s^{\prime})]$. We denote the optimal policy as $\pi^{*}$ and its state value function as $V^{*}\coloneqq V^{\pi^{*}}$. Under the assumption that rewards are in $[0,1]$, we have $V^{\pi}(s),Q^{\pi}(s,a)\in[0,\frac{1}{1-\gamma}]$ for all $\pi$, $s\in\SS$, and $a\in\mathcal{A}$. We denote the initial state distribution of interest as $d_{0}\in\Delta(\SS)$ and the state distribution of policy $\pi$ at time $t$ as $d_{t}^{\pi}$, with $d_{0}^{\pi}=d_{0}$. Given $d_{0}$, we define the average state distribution of a policy $\pi$ as $d^{\pi}\coloneqq(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}d_{t}^{\pi}$. With a slight abuse of notation, we also write $d^{\pi}(s,a)\coloneqq d^{\pi}(s)\pi(a|s)$.

We consider RL with prior knowledge expressed in the form of a heuristic value function. The goal is to find a policy $\pi$ that has high return through interactions with an unknown MDP $\mathcal{M}$, i.e., $\max_{\pi}V^{\pi}(d_{0})$. While the agent here does not fully know $\mathcal{M}$, we suppose that, before interactions start the agent is provided with a heuristic $h:\SS\to\mathbb{R}$ which the agent can query throughout learning.

The heuristic $h$ represents a prior guess of the optimal value function $V^{*}$ of $\mathcal{M}$. Common sources of heuristics are domain knowledge as typically employed in planning, and logged data collected by exploratory or by expert behavioral policies. In the latter, a heuristic guess of $V^{*}$ can be computed from the data by offline RL algorithms. For instance, when we have trajectories of an expert behavioral policy, Monte-Carlo regression estimate of the observed returns may be a good guess of $V^{*}$.

Using heuristics to solve MDP problems has been popular in planning and control, but its usage is rather limited in RL. The closest provable technique in RL is PBRS [29], where the reward is modified into $\overline{r}(s,a)\coloneqq r(s,a)+\gamma\mathbb{E}_{s^{\prime}|s,a}[h(s^{\prime})]-h(s)$. It can be shown that this transformation does not introduce bias into the policy ordering, and therefore solving the new MDP $\overline{\mathcal{M}}\coloneqq(\SS,\mathcal{A},P,\overline{r},\gamma)$ would yield the same optimal policy $\pi^{*}$ of $\mathcal{M}$.

Conceptually when the heuristic is the optimal value function $h=V^{*}$, the agent should be able to find the optimal policy $\pi^{*}$ of $\mathcal{M}$ by acting myopically, as $V^{*}$ already contains all necessary long-term information for good decision making. However, running an RL algorithm with the PBRS reward (i.e. solving $\overline{\mathcal{M}}\coloneqq(\SS,\mathcal{A},P,\overline{r},\gamma)$) does not take advantage of this shortcut. To make learning efficient, we need to also let the base RL algorithm know that acting greedily (i.e., using a smaller discount) with the shaped reward can yield good policies. An intuitive idea is to run the RL algorithm to maximize $\overline{V}^{\pi}_{\lambda}(d_{0})$, where $\overline{V}^{\pi}_{\lambda}$ denotes the value function of $\pi$ in an MDP $\overline{\mathcal{M}}_{\lambda}\coloneqq(\SS,\mathcal{A},P,\overline{r},\lambda\gamma)$ for some $\lambda\in[0,1]$. However this does not always work. For example, when $\lambda=0$, $\max_{\pi}\overline{V}^{\pi}_{\lambda}(d_{0})$ only optimizes for the initial states $d_{0}$, but obviously the agent is going to encounter other states in $\mathcal{M}$. We next propose a provably correct version, HuRL, to leverage this short-horizon insight.

HuRL takes a reduction-based approach to realize the idea of heuristic guidance. As summarized in Algorithm 1, HuRL takes a heuristic $h:\SS\to\mathbb{R}$ and a base RL algorithm $\mathcal{L}$ as input, and outputs an approximately optimal policy for the original MDP $\mathcal{M}$. During training, HuRL iteratively runs the base algorithm $\mathcal{L}$ to collect data from the MDP $\mathcal{M}$ and then uses the heuristic $h$ to modify the agent’s collected experiences. Namely, in iteration $n$, the agent interacts with the original MDP $\mathcal{M}$ and saves the raw transition tuples⁴⁴4If $\mathcal{L}$ learns only with trajectories, we transform each tuple and assemble them to get the modified trajectory. $\mathcal{D}_{n}=\{(s,a,r,s^{\prime})\}$ (line 2). HuRL then defines a reshaped MDP $\widetilde{\mathcal{M}}_{n}\coloneqq(\SS,\mathcal{A},P,\widetilde{r}_{n},\widetilde{\gamma}_{n})$ (line 3) by changing the rewards and lowering the discount factor:

where $\lambda_{n}\in[0,1]$ is the mixing coefficient. The new discount $\widetilde{\gamma}_{n}$ effectively gives $\widetilde{\mathcal{M}}_{n}$ a shorter horizon than $\mathcal{M}$’s, while the heuristic $h$ is blended into the new reward in (2) to account for the missing long-term information. We call $\widetilde{\gamma}_{n}=\lambda_{n}\gamma$ in (2) the guidance discount to be consistent with prior literature [20], which can be viewed in terms of our framework as using a zero heuristic. In the last step (line 4), HuRL calls the base algorithm $\mathcal{L}$ to perform updates with respect to the reshaped MDP $\widetilde{\mathcal{M}}_{n}$. This is realized by 1) setting the discount factor used in $\mathcal{L}$ to $\widetilde{\gamma}_{n}$, and 2) setting the sampled reward to $r+(\gamma-\widetilde{\gamma}_{n})h(s^{\prime})$ for every transition tuple $(s,a,r,s^{\prime})$ collected from $\mathcal{M}$. We remark that the base algorithm $\mathcal{L}$ in line 2 always collects trajectories of lengths proportional to the original discount $\gamma$, while internally the optimization is done with a lower discount $\widetilde{\gamma}_{n}$ in line 4.

Over the course of training, HuRL repeats the above steps with a sequence of increasing mixing coefficients $\{\lambda_{n}\}$. From (2) we see that as the agent interacts with the environment, the effects of the heuristic in MDP reshaping decrease and the effective horizon of the reshaped MDP increases.

We can think of HuRL as introducing a horizon-based regularization for RL, where the regularization center is defined by the heuristic and its strength diminishes as the mixing coefficient increases. As the agent collects more experiences, HuRL gradually removes the effects of regularization and the agent eventually optimizes for the original MDP.

HuRL’s regularization is designed to reduce learning variance, similar to the role of regularization in supervised learning. Unlike the typical weight decay imposed on function approximators (such as the agent’s policy or value networks), our proposed regularization leverages the structure of MDPs to regulate the complexity of the MDP the agent faces, which scales with the MDP’s discount factor (or, equivalently, the horizon). When the guidance discount $\widetilde{\gamma}_{n}$ is lower than the original discount $\gamma$ (i.e. $\lambda_{n}<1$), the reshaped MDP $\widetilde{\mathcal{M}}_{n}$ given by (2) has a shorter horizon and requires fewer samples to solve. However, the reduced complexity comes at the cost of bias, because the agent is now incentivized toward maximizing the performance with respect to the heuristic rather than the original long-term returns of $\mathcal{M}$. In the extreme case of $\lambda_{n}=0$, HuRL would solve a zero-horizon contextual bandit problem with contexts (i.e. states) sampled from $d^{\pi}$ of $\mathcal{M}$.

We illustrate this idea in a chain MDP environment in Fig. 1. The optimal policy $\pi^{*}$ for this MDP’s original $\gamma=0.9$ always picks action $\rightarrow$, as shown in Fig. 1(b)-(1), giving the optimal value $V^{*}$ in Fig. 1(a)-(2). Suppose we used a smaller guidance discount $\widetilde{\gamma}=0.5\gamma$ to accelerate learning. This is equivalent to HuRL with a zero heuristic $h=0$ and $\lambda=0.5$. Solving this reshaped MDP yields a policy $\widetilde{\pi}^{*}$ that acts very myopically in the original MDP, as shown in Fig. 1(b)-(2); the value function of $\widetilde{\pi}^{*}$ in the original MDP is visualized in Fig. 1(a)-(4).

Now, suppose we use Fig. 1(a)-(4) as a heuristic in HuRL instead of $h=0$. This is a bad choice of heuristic (Bad $h$) as it introduces a large bias with respect to $V^{*}$ (cf. Fig. 1(a)-(2)). On the other hand, we can roll out a random policy in the original MDP and use its value function as the heuristic (Good $h$), shown in Fig. 1(a)-(3). Though the random policy has an even lower return at the initial state $s=3$, it gives a better heuristic because this heuristic shares the same trend as $V^{*}$ in Fig. 1(a)-(1). HuRL run with Good $h$ and Bad $h$ yields policies in Fig. 1(b)-(3,4), and the quality of the resulting solutions in the original MDP, $V_{\lambda}^{\widetilde{\pi}^{*}}(d_{0})$, is reported in Fig. 1(c) for different $\lambda$. Observe that HuRL with a good heuristic can achieve $V^{*}(d_{0})$ with a much smaller horizon $\lambda\leq 0.5$. Using a bad $h$ does not lead to $\pi^{*}$ at all when $\lambda=0.5$ (Fig. 1(b)-(4)) but is guaranteed to do so when $\lambda$ converges to $1$. (Fig. 1(b)-(5)).

Our main result is a performance decomposition, which characterizes how a heuristic $h$ and suboptimality in solving the reshaped MDP $\widetilde{\mathcal{M}}$ relate to performance in the original MDP $\mathcal{M}$.

The theorem shows that suboptimality of a policy $\pi$ in the original MDP $\mathcal{M}$ can be decomposed into 1) a bias term due to solving a reshaped MDP $\widetilde{\mathcal{M}}$ instead of the original MDP $\mathcal{M}$, and 2) a regret term (i.e. the learning variance) due to $\pi$ being suboptimal in the reshaped MDP $\widetilde{\mathcal{M}}$. Moreover, it shows that heuristics are equivalent up to constant offsets. In other words, only the relative ordering between states that a heuristic induces matters in learning, not the absolute values.

Balancing the two terms trades off bias and variance in learning. Using a smaller $\lambda$ replaces the long-term information with the heuristic and make the horizon of the reshaped MDP $\widetilde{\mathcal{M}}$ shorter. Therefore, given a finite interaction budget, the regret term in (3) can be more easily minimized, though the bias term in (4) can potentially be large if the heuristic is bad. On the contrary, with $\lambda=1$, the bias is completely removed, as the agent solves the original MDP $\mathcal{M}$ directly.

The regret term in (3) characterizes the performance gap due to $\pi$ being suboptimal in the reshaped MDP $\widetilde{\mathcal{M}}$, because $\mathrm{Regret}(h,\lambda,\widetilde{\pi}^{*})=0$ for any $h$ and $\lambda$. For learning, we need the base RL algorithm $\mathcal{L}$ to find a policy $\pi$ such that the regret term in (3) is small. By the definition in (3), the base RL algorithm $\mathcal{L}$ is required not only to find a policy $\pi$ such that $\widetilde{V}^{*}(s)-\widetilde{V}^{\pi}(s)$ is small for states from $d_{0}$, but also for states $\pi$ visits when rolling out in the original MDP $\mathcal{M}$. In other words, it is insufficient for the base RL algorithm to only optimize for $\widetilde{V}^{\pi}(d_{0})$ (the performance in the reshaped MDP with respect to the initial state distribution; see Section 2.2). For example, suppose $\lambda=0$ and $d_{0}$ concentrates on a single state $s_{0}$. Then maximizing $\widetilde{V}^{\pi}(d_{0})$ alone would only optimize $\pi(\cdot|s_{0})$ and the policy $\pi$ need not know how to act in other parts of the state space.

To use HuRL, we need the base algorithm to learn a policy $\pi$ that has small action gaps in the reshaped MDP $\widetilde{\mathcal{M}}$ but along trajectories in the original MDP $\mathcal{M}$, as we show below. This property is satisfied by off-policy RL algorithms such as Q-learning [34].

Another way to comprehend the regret term is through studying its dependency on $\lambda$. When $\lambda=1$, $\mathrm{Regret}(h,0,\pi)=V^{*}(d_{0})-V^{\pi}(d_{0})$, which is identical to the policy regret in $\mathcal{M}$ for a fixed initial distribution $d_{0}$. On the other hand, when $\lambda=0$, $\mathrm{Regret}(h,0,\pi)=\max_{\pi^{\prime}}\frac{1}{1-\gamma}\mathbb{E}_{s\sim d^{\pi}}[\widetilde{r}(s,\pi^{\prime})-\widetilde{r}(s,\pi)]$, which is the regret of a non-stationary contextual bandit problem where the context distribution is $d^{\pi}$ (the average state distribution of $\pi$). In general, for $\lambda\in(0,1)$, the regret notion mixes a short-horizon non-stationary problem and a long-horizon stationary problem.

One natural question is whether the reshaped MDP $\widetilde{\mathcal{M}}$ has a more complicated and larger value landscape than the original MDP $\mathcal{M}$, because these characteristics may affect the regret rate of a base algorithm. We show that $\widetilde{\mathcal{M}}$ preserves the value bounds and linearity of the original MDP.

On the contrary, the MDP $\overline{\mathcal{M}}_{\lambda}\coloneqq(\SS,\mathcal{A},P,\overline{r},\lambda\gamma)$ in Section 2.2 does not have these properties. We can show that $\overline{\mathcal{M}}_{\lambda}$ is equivalent to $\widetilde{\mathcal{M}}$ up to a PBRS transformation (i.e., $\bar{r}(s,a)=\tilde{r}(s,a)+\tilde{\gamma}\mathbb{E}_{s^{\prime}|s,a}[h(s^{\prime})]-h(s)$). Thus, HuRL incorporates guidance discount into PBRS with nicer properties.

The bias term in (4) characterizes suboptimality due to using a heuristic $h$ in place of long-term state values in $\mathcal{M}$. What is the best heuristic in this case? From the definition of the bias term in (4), we see that the ideal heuristic is the optimal value $V^{*}$, as $\mathrm{Bias}(V^{*},\lambda,\pi)=0$ for all $\lambda\in[0,1]$. By continuity, we can expect that if $h$ deviates from $V^{*}$ a little, then the bias is small.

To better understand how the heuristic $h$ affects the bias, we derive an upper bound on the bias by replacing the first term in (4) with an upper bound that depends only on $\pi^{*}$.

In Proposition 4.3, the term $(1-\lambda)\gamma\mathcal{C}(\pi^{*},V^{*}-h,\lambda\gamma)$ is the underestimation error of the heuristic $h$ under the states visited by the optimal policy $\pi^{*}$ in the original MDP $\mathcal{M}$. Therefore, to minimize the first term in the bias, we would want the heuristic $h$ to be large along the paths that $\pi^{*}$ generates.

However, Proposition 4.3 also discourages the heuristic from being arbitrarily large, because the second term in the bias in (4) (or, equivalently, the second term in Proposition 4.3) incentivizes the heuristic to underestimate the optimal value of the reshaped MDP $\widetilde{V}^{*}$. More precisely, the second term requires the heuristic to obey some form of spatial consistency. A quick intuition is the observation that if $h(s)=V^{\pi^{\prime}}(s)$ for some $\pi^{\prime}$ or $h(s)=0$, then $h(s)\leq\widetilde{V}^{*}(s)$ for all $s\in\SS$. More generally, we show that if the heuristic is improvable with respect to the original MDP $\mathcal{M}$ (i.e. the heuristic value is lower than that of the max of Bellman backup), then $h(s)\leq\widetilde{V}^{*}(s)$. By Proposition 4.3, learning with an improvable heuristic in HuRL has a much smaller bias.

While Corollary 4.1 shows that HuRL can handle an imperfect heuristic, this result is not ideal. The corollary depends on the $\ell_{\infty}$ approximation error, which can be difficult to control in large state spaces. Here we provide a more refined sufficient condition of good heuristics. We show that the concept of pessimism in the face of uncertainty provides a finer mechanism for controlling the approximation error of a heuristic and would allow us to remove the $\ell_{\infty}$-type error. This result is useful for constructing heuristics from data that does not have sufficient support.

From Proposition 4.3 we see that the source of the $\ell_{\infty}$ error is the second term in the bias upper bound, as it depends on the states that the agent’s policy visits which can change during learning. To remove this dependency, we can use improvable heuristics (see Proposition 4.4), as they satisfy $h(s)\leq\widetilde{V}^{*}(s)$. Below we show that Bellman-consistent pessimism yields improvable heuristics.

The Bellman-consistent pessimism in Proposition 4.5 essentially says that $h$ is pessimistic with respect to the Bellman backup. This condition has been used as the foundation for designing pessimistic off-policy RL algorithms, such as pessimistic value iteration [30] and algorithms based on pessimistic absorbing MDPs [31]. In other words, these pessimistic algorithms can be used to construct good heuristics with small bias in Proposition 4.3 from offline data. With such a heuristic, the bias upper bound would be simply $\mathrm{Bias}(h,\lambda,\pi)\leq(1-\lambda)\gamma\mathcal{C}(\pi^{*},V^{*}-h,\lambda\gamma)$. Therefore, as long as enough batch data are sampled from a distribution that covers states that $\pi^{*}$ visits, these pessimistic algorithms can construct good heuristics with nearly zero bias for HuRL with high probability.

We consider four MuJoCo environments with dense rewards (Hopper-v2, HalfCheetah-v2, Humanoid-v2, and Swimmer-v2) and a sparse reward version of Reacher-v2 (denoted as Sparse-Reacher-v2)⁶⁶6The reward is zero at the goal and $-1$ otherwise.. We design the experiments to simulate two learning scenarios. First, we use Sparse-Reacher-v2 to simulate the setting where an engineered heuristic based on domain knowledge is available; since this is a goal reaching task, we designed a heuristic $h(s)=r(s,a)-100\|e(s)-g(s)\|$, where $e(s)$ and $g(s)$ denote the robot’s end-effector position and the goal position, respectively. Second, we use the dense reward environments to model scenarios where a batch of data collected by multiple behavioral policies is available before learning, and a heuristic is constructed by an offline policy evaluation algorithm from the batch data (see Appendix C.1 for details). In brief, we generated these behavioral policies by running SAC from scratch and saved the intermediate policies generated in training. We then use least-squares regression to fit a neural network to predict empirical Monte-Carlo returns of the trajectories in the sampled batch of data. We also use behavior cloning (BC) to warm-start all RL agents based on the same batch dataset in the dense reward experiments.

The base RL algorithm here, SAC, is based on the standard implementation in Garage (MIT License) [37]. The policy and value networks are fully connected independent neural networks. The policy is Tanh-Gaussian and the value network has a linear head.

Fig. 2 shows the results on the MuJoCo environments. Overall, we see that HuRL is able to leverage engineered and learned heuristics to significantly improve the learning efficiency. This trend is consistent across all environments that we tested on.

For the sparse-reward experiments, we see that SAC and PBRS struggle to learn, while HuRL is able to converge to the optimal performance much faster. For the dense reward experiments, similarly HuRL-MC converges much faster, though the gain in HalfCheetah-v2 is minor and it might have converged to a worse local maximum in Swimmer-v2. In addition, we see that warm-starting SAC using BC (i.e. SAC w/ BC) can improve the learning efficiency compared with the vanilla SAC, but using BC alone does not result in a good policy. Lastly, we see that using the zero heuristic (HuRL-zero) with extra $\lambda$-scheduling does not further improve the performance of SAC w/ BC. This comparison verifies that the learned Monte-Carlo heuristic provides non-trivial information.

Interestingly, we see that applying PBRS to SAC leads to even worse performance than running SAC with the original reward. There are two reasons why SAC+PBRS is less desirable than SAC+HuRL as we discussed before: 1) PBRS changes the reward/value scales in the induced MDP, and popular RL algorithms like SAC are very sensitive to such changes. In contrast, HuRL induces values on the same scale as we show in Proposition 4.2. 2) In HuRL, we are effectively providing the algorithm some more side-information to let SAC shorten the horizon when the heuristic is good.

The results in Fig. 2 also have another notable aspect. Because the datasets used in the dense reward experiments contain trajectories collected by a range of policies, it is likely that BC suffers from disagreement in action selection among different policies. Nonetheless, training a heuristic using a basic Monte-Carlo regression seems to be less sensitive to these conflicts and still results in a helpful heuristic for HuRL. One explanation can be that heuristics are only functions of states, not of states and actions, and therefore the conflicts are minor. Another plausible explanation is that HuRL only uses the heuristic to guide learning, and does not completely rely on it to make decisions Thus, HuRL can be more robust to the heuristic quality, or, equivalently, to the quality of prior knowledge.

To further verify that the acceleration in Fig. 2 is indeed due to horizon shortening, we conducted an ablation study for HuRL-MC on Hopper-v2, whose results are presented in Fig. 2(f). HuRL-MC’s best $\lambda$-schedule hyperparameters on Hopper-v2, which are reflected in its performance in the aforementioned Fig. 2(c), induced a near-constant schedule at $\lambda=0.95$; to obtain the curves in Fig. 2(f), we ran HuRL-MC with constant-$\lambda$ schedules for several more $\lambda$ values. Fig. 2(f) shows that increasing $\lambda$ above $0.98$ leads to a performance drop. Since using a large $\lambda$ decreases bias and makes the reshaped MDP more similar to the original MDP, we conclude that the increased learning speed on Hopper-v2 is due to HuRL’s horizon shortening (coupled with the guidance provided by its heuristic).