SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chenran Li $^{1*}$ , Chen Tang $^{1*}$ , Haruki Nishimura $^{2}$ , Jean Mercat $^{2}$ , Masayoshi Tomizuka $^{1}$ , Wei Zhan $^{1}$
<sup>1</sup> University of California Berkeley, <sup>2</sup> Toyota Research Institute, USA {chenran_li, chen_tang, tomizuka, wzhan}@berkeley.edu {haruki.nishimura, jean.mercat}@tri.global
Imitation Learning (IL) is a widely used framework for learning imitative behavior from demonstrations. It is especially appealing for solving complex real-world tasks where handcrafting reward function is difficult, or when the goal is to mimic human expert behavior. However, the learned imitative policy can only follow the behavior in the demonstration. When applying the imitative policy, we may need to customize the policy behavior to meet different requirements coming from diverse downstream tasks. Meanwhile, we still want the customized policy to maintain its imitative nature. To this end, we formulate a new problem setting called policy customization. It defines the learning task as training a policy that inherits the characteristics of the prior policy while satisfying some additional requirements imposed by a target downstream task. We propose a novel and principled approach to interpret and determine the trade-off between the two task objectives. Specifically, we formulate the customization problem as a Markov Decision Process (MDP) with a reward function that combines 1) the inherent reward of the demonstration; and 2) the add-on reward specified by the downstream task. We propose a novel framework, Residual Q-learning (RQL), which can solve the formulated MDP by leveraging the prior policy without knowing the inherent reward or value function of the prior policy. We derive a family of residual Q-learning algorithms that can realize offline and online policy customization, and show that the proposed algorithms can effectively accomplish policy customization tasks in various environments. Demo videos and code are available on our website: https://sites.google.com/view/residualq-learning.
Imitation learning (IL) is a widely used framework for learning imitative behavior from demonstrations. It is especially appealing for solving complicated real-world tasks where handcrafting reward function is difficult, or when the goal is to mimic expert behavior. For instance, IL has been widely studied for synthesizing human-like agents from large-scale human demonstrations on challenging real-world tasks, such as autonomous driving [4, 7, 2] and robot manipulation [59, 37, 39]. However, when applying the imitative policy in open real-world scenarios, we may encounter different circumstances that impose diverse requirements beyond merely following the behavior in the demonstration, such as reaching particular goals [45], enforcing safety [36], or complying with certain behavior preferences [61]. In this case, we need to customize the imitative policy to meet those task-specific objectives. Meanwhile, in many applications, we may still want the customized policy to possess some inherent characteristics of the expert demonstration (e.g., human-like).
To this end, we formulate a new problem setting we refer to as policy customization. In the policy customization problem, we are given a pre-trained policy, and the task is to train a new policy that 1) inherits the properties of the prior policy; 2) satisfies some additional requirements imposed by a given downstream task. At first glance, this newly formulated problem may seem to be seamlessly solved by existing reinforcement learning (RL) fine-tuning algorithms [44, 40, 61, 51], where RL is applied to fine-tune a pre-trained policy obtained from behavior cloning (BC) or offline RL [31, 29, 34]. In particular, some RL fine-tuning methods combine different forms of IL objectives with RL to accelerate policy training, e.g., maximizing the policy likelihood on demonstration [44], regularizing the KL-divergence between the fine-tuned and prior policies [40, 61]. Different from the policy customization problem we consider, the IL objectives are introduced heuristically as regularization in RL fine-tuning. Solving policy customization instead requires a principled way to design the learning objective to jointly optimize the policy's performance on the imitative and downstream tasks, which has not been well-studied in the RL fine-tuning literature.
In this work, we propose such a principled approach to interpret and design the policy customization objective with a solid theoretical basis. Formally, we formulate policy customization as a Markov decision process (MDP) whose reward function is a linear combination of 1) the inherent reward of the demonstration; and 2) an add-on reward specified by the downstream task. The reward weight directly determines the trade-off between the imitative and downstream tasks for policy optimization. However, we cannot solve this MDP with existing RL algorithms, since we can only access the imitative policy without knowing its underlying reward.
To tackle this challenge, we propose a novel Residual Q-Learning (RQL) framework that can solve the formulated MDP without knowing the inherent reward or value function of the prior policy. Specifically, we define a residual Q-function that, together with the log-likelihood of the prior policy, can conveniently construct the maximum-entropy policy solving the target MDP. We show that the residual Q-function can be iteratively learned with an update rule derived from soft Bellman equation [17]. Hence, the proposed RQL method can be applied on top of any value-based maximum-entropy RL algorithms. In particular, we derive residual soft Q-learning and residual soft actor-critic—two model-free policy customization algorithms through RL fine-tuning—which are adapted from soft Q-learning [17] and soft actor-critic [18] respectively. In addition, we introduce a model-based policy customization algorithm adapted from maximum-entropy Monte Carlo Tree Search (MCTS) [55], which enables zero-shot online policy customization when a dynamics model is available. In our experiments, we show that the proposed RQL algorithms can effectively customize the policies toward the additional task objective while maintaining their performance on the basic tasks for which the prior policies are trained.
2 Policy Customization with Residual Q-Learning
2.1 Problem Formulation: Policy Customization
We consider a MDP defined by the tuple  $\mathcal{M} = (\mathcal{S},\mathcal{A},r,p)$ , where  $\mathcal{S}$  is the state space,  $\mathcal{A}$  is the action space, and  $r:S\times \mathcal{A}\to \mathbb{R}$  is the reward function. Without loss of generality, we introduce the formulation under continuous state and action spaces. The state transition probability function  $p:S\times \mathcal{A}\times \mathcal{S}\rightarrow [0,\infty)$  represents the probability density of the next state  $s_{t + 1}\in S$  given the current state  $s_t\in S$  and action  $a_{t}\in \mathcal{A}$ . While the following formulation applies whether the reward function  $r$  is given or not, we are primarily interested in the case when  $r$  remains unknown and, instead, expert demonstrations are available to indicate desirable behaviors. We assume that we have access to a prior policy  $\pi :S\times \mathcal{A}\to [0,\infty)$  pre-trained with imitation learning [15, 22] methods. In particular, we follow the common practice and model  $\pi$  as a maximum-entropy policy to account for sub-optimal demonstration [60], which is especially important for human demonstration. Specifically, the policy  $\pi$  maximizes the entropy-augmented cumulative reward that is inherent in the demonstrations. The resulting optimal policy follows a Boltzmann distribution [17] as follows:
$$
\pi (\boldsymbol {a} | \boldsymbol {s}) = \frac {1}{Z _ {s}} \exp \left(\frac {1}{\alpha} Q ^ {*} (\boldsymbol {s}, \boldsymbol {a})\right), \tag {1}
$$
where  $Q^{*}(s, \pmb{a})$  is the soft Q-function as defined in [17] and it satisfies the soft Bellman equation:
$$
Q ^ {*} (\boldsymbol {s}, \boldsymbol {a}) = r (\boldsymbol {s}, \boldsymbol {a}) + \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime} \sim p (\cdot | \boldsymbol {s}, \boldsymbol {a})} \left[ \alpha \log \int_ {\mathcal {A}} \exp \left(\frac {1}{\alpha} Q ^ {*} (\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime})\right) d \boldsymbol {a} ^ {\prime} \right], \tag {2}
$$
and  $Z_{s}$  is the normalization factor defined as  $\int_{\mathcal{A}}\exp \left(\frac{1}{\alpha} Q^{*}(s,\pmb {a})\right)d\pmb{a}$ . The temperature coefficient  $\alpha$  determines the relative importance of entropy and reward in the maximum-entropy policy [17]. It is straightforward to derive from Eqn. (1) that:
$$
Q ^ {*} (\boldsymbol {s}, \boldsymbol {a}) = \alpha \log \pi (\boldsymbol {a} | \boldsymbol {s}) + \alpha \log Z _ {s}. \tag {3}
$$
Instead of directly deploying the policy  $\pi$ , we aim to customize it to satisfy some additional requirements for a given downstream task, specified by an add-on reward function  $r_R: S \times \mathcal{A} \to \mathbb{R}$ . Meanwhile, we want the customized policy to inherit the characteristics of the original policy. Formally, we aim to leverage the pre-trained prior policy  $\pi$  to train a new policy for the MDP defined by  $\hat{\mathcal{M}} = (\mathcal{S}, \mathcal{A}, \omega r + r_R, p)$ . The new reward function is defined as a weighted sum of the basic reward  $r$  and the add-on reward  $r_R$  so that the customized policy meets both requirements. Under this setting, the main challenge lies in how to find the optimal policy for the new task without knowing the prior reward  $r$  but only the imitative policy  $\pi$ .
2.2 Residual Q-Learning
Since the reward  $r$  is unknown, we cannot directly train the new policy to maximize the total reward through RL. One plausible solution is to learn the inherent reward via inverse RL [12] and then combine the inferred reward with the add-on reward for RL training. However, practical inverse RL algorithms require jointly optimizing the policy and reward function, which could be challenging to yield good performance in certain environments when compared to imitation learning [16]. Instead, we aim to explore a solution without the need to explicitly infer the underlying reward to derive a more flexible policy customization framework. It turns out that we can leverage the imitative policy  $\pi$  to construct the new policy without knowing  $\pi$ 's underlying reward or value function. The proposed method can be applied on top of any value-based maximum-entropy RL algorithms [17, 18, 55], which rely on the soft Bellman update operator defined as:
$$
\hat {Q} _ {t + 1} (\boldsymbol {s}, \boldsymbol {a}) = r _ {R} (\boldsymbol {s}, \boldsymbol {a}) + \omega r (\boldsymbol {s}, \boldsymbol {a}) + \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime} \sim p (\cdot | \boldsymbol {s}, \boldsymbol {a})} \left[ \hat {\alpha} \log \int_ {\mathcal {A}} \exp \left(\frac {1}{\hat {\alpha}} \hat {Q} _ {t} (\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime})\right) d \boldsymbol {a} ^ {\prime} \right], \tag {4}
$$
where  $\hat{Q}_t$  is the estimated soft Q-function for  $\hat{\mathcal{M}}$  at the  $t^{\mathrm{th}}$  iteration. Instead of directly applying the soft Bellman update operator to iterate  $\hat{Q}_t$  during policy learning, we define a residual Q-function as  $Q_{R,t} \coloneqq \hat{Q}_t - \omega Q^*$  and attempt to iterate this residual Q-function during policy learning. We then derive an update rule for  $Q_{R,t}$  that does not depend on the reward or value function of  $\pi$ :
$$
\begin{array}{l} Q _ {R, t + 1} (\boldsymbol {s}, \boldsymbol {a}) \\ = r _ {R} (\boldsymbol {s}, \boldsymbol {a}) + \omega r (\boldsymbol {s}, \boldsymbol {a}) + \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime}} \left[ \hat {\alpha} \log \int_ {\mathcal {A}} \exp \left(\frac {1}{\hat {\alpha}} \hat {Q} _ {t} \left(\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime}\right)\right) d \boldsymbol {a} ^ {\prime} \right] - \omega Q ^ {*} (\boldsymbol {s}, \boldsymbol {a}), (5a) \\ = r _ {R} (\boldsymbol {s}, \boldsymbol {a}) + \omega Q ^ {*} (\boldsymbol {s}, \boldsymbol {a}) - \omega \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime}} \left[ \alpha \log \int_ {\mathcal {A}} \exp \left(\frac {1}{\alpha} Q ^ {*} (\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime})\right) d \boldsymbol {a} ^ {\prime} \right] \\ + \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime}} \left[ \hat {\alpha} \log \int_ {\mathcal {A}} \exp \left(\frac {1}{\hat {\alpha}} \hat {Q} _ {t} \left(\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime}\right)\right) d \boldsymbol {a} ^ {\prime} \right] - \omega Q ^ {*} (\boldsymbol {s}, \boldsymbol {a}), (5b) \\ = r _ {R} (\boldsymbol {s}, \boldsymbol {a}) - \omega \alpha \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime}} \log Z _ {s ^ {\prime}} \\ + \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime}} \left[ \hat {\alpha} \log \int_ {\mathcal {A}} \exp \left(\frac {1}{\hat {\alpha}} \left(Q _ {R, t} \left(\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime}\right) + \omega \alpha \log \pi \left(\boldsymbol {a} ^ {\prime} \mid \boldsymbol {s} ^ {\prime}\right) + \omega \alpha \log Z _ {s ^ {\prime}}\right)\right) d \boldsymbol {a} ^ {\prime} \right], (5c) \\ = r _ {R} (\boldsymbol {s}, \boldsymbol {a}) - \omega \alpha \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime}} \log Z _ {s ^ {\prime}} + \omega \alpha \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime}} \log Z _ {s ^ {\prime}} \\ + \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime}} \left[ \hat {\alpha} \log \int_ {\mathcal {A}} \exp \left(\frac {1}{\hat {\alpha}} \left(Q _ {R, t} \left(\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime}\right) + \omega \alpha \log \pi \left(\boldsymbol {a} ^ {\prime} \mid \boldsymbol {s} ^ {\prime}\right)\right)\right) d \boldsymbol {a} ^ {\prime} \right], (5d) \\ = r _ {R} (\boldsymbol {s}, \boldsymbol {a}) + \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime}} \left[ \hat {\alpha} \log \int_ {\mathcal {A}} \exp \left(\frac {1}{\hat {\alpha}} \left(Q _ {R, t} \left(\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime}\right) + \omega \alpha \log \pi \left(\boldsymbol {a} ^ {\prime} \mid \boldsymbol {s} ^ {\prime}\right)\right)\right) d \boldsymbol {a} ^ {\prime} \right], (5e) \\ \end{array}
$$
where Eqn. (5b) results from expressing the reward  $r$  with the soft Q-function  $Q^{*}$  based on Eqn. (2). Eqn. (5c) results from the definition of the residual Q-function and Eqn. (3). Please refer to Appendix A for a step-by-step breakdown of the derivation process. In each iteration, we can then define the policy corresponding to the current estimated  $\hat{Q}_t$  without computing  $\hat{Q}_t$ :
$$
\hat {\pi} _ {t} (\boldsymbol {a} | \boldsymbol {s}) \propto \exp \left(\frac {1}{\hat {\alpha}} \left(Q _ {R, t} (\boldsymbol {s}, \boldsymbol {a}) + \omega \alpha \log \pi (\boldsymbol {a} | \boldsymbol {s})\right)\right). \tag {6}
$$
One remaining issue is that we do not have access to the underlying temperature coefficient  $\alpha$  of the prior policy  $\pi$ . Note that  $\alpha$  is multiplied upon the reward weight  $\omega$  in Eqn. (5) and Eqn. (6). In practice, the reward weight is a hyperparameter to tune. Hence, we can define  $\omega' = \omega \alpha$  and directly treat  $\omega'$  as the hyperparameter without the need to know the true  $\alpha$ .
3 Practical Algorithms
In this section, we present three practical algorithms for policy customization under different settings based on the update rule of residual Q-function in Eqn. (5) and the maximum-entropy policy defined in Eqn. (6). In these algorithms, the goal is to find a policy that solves the MDP  $\hat{\mathcal{M}}$  given a prior policy  $\pi$  and the add-on reward function  $r_R$ . We first introduce two model-free algorithms to learn the residual Q-function and customize the policy through additional RL steps, which are adapted from soft Q-learning [17] and soft actor-critic [18] respectively. Then we introduce a model-based policy customization algorithm adapted from maximum-entropy MCTS [55], which enables zero-shot online policy customization when a dynamics model is available.
Residual Soft Q-Learning. We adapt the soft Q-learning [17] algorithm to a residual soft Q-learning algorithm for offline policy customization (i.e., policy customization via additional RL training steps) when the action space is discrete. The residual algorithm is essentially the same as the soft Q-learning. The only difference is that we use a function approximator parameterized by  $\theta$  to model the residual Q-function instead of the soft Q-function and denote it as  $Q_{R,\theta}(\boldsymbol{s},\boldsymbol{a})$ . At each iteration, the parameters are updated by taking gradient steps to minimize the temporal-difference (TD) error:
$$
J _ {Q _ {R}} (\theta) = \mathbb {E} _ {\left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}, r _ {R, t}, \boldsymbol {s} _ {t + 1}\right) \sim \mathcal {D}} \left[ \frac {1}{2} \left(Q _ {R, \tilde {\theta}} ^ {\text {t a r g e t}} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) - Q _ {R, \theta} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right)\right) ^ {2} \right], \tag {7}
$$
where  $\mathcal{D}$  is the replay buffer and the target residual Q-value  $Q_{R,\bar{\theta}}^{\mathrm{target}}(s_t,a_t)$  is computed as:
$$
Q _ {R, \bar {\theta}} ^ {\mathrm {t a r g e t}} (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) = r _ {R, t} + \gamma \hat {\alpha} \log \sum_ {\boldsymbol {a} ^ {\prime} \in \mathcal {A}} \exp \left(\frac {1}{\hat {\alpha}} \left(Q _ {R, \bar {\theta}} (\boldsymbol {s} _ {t + 1}, \boldsymbol {a} ^ {\prime}) + \omega^ {\prime} \log \pi (\boldsymbol {a} ^ {\prime} | \boldsymbol {s} _ {t + 1})\right)\right), \tag {8}
$$
with  $\bar{\theta}$  being the target parameters. The above objective function can be derived straightforwardly from the original objective function of soft Q-learning. We provide the detailed derivation in Appendix B for completeness. With a discrete action space, we can compute the policy distribution exactly given the residual Q-function and the prior policy as in Eqn. (6).
Residual Soft Actor-Critic. For tasks with continuous action spaces, we adapt the soft actor-critic [18] algorithm to a residual soft actor-critic algorithm. We aim to train a parameterized residual Q-function  $\hat{Q}_{R,\theta}(s,a)$  and a policy network  $\hat{\pi}_{\phi}(a|s)$  by alternating between policy evaluation and improvement. In the policy evaluation step, we update the residual Q-function for evaluating the current policy. The loss function is essentially the one in Eqn. (7) but with the target residual Q-value computed with respect to the current policy  $\hat{\pi}_{\phi}$ :
$$
Q _ {R, \bar {\theta}} ^ {\text {t a r g e t}} (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) = r _ {R, t} + \gamma \mathbb {E} _ {\boldsymbol {a} ^ {\prime} \sim \hat {\pi} _ {\phi}} \left[ Q _ {R, \bar {\theta}} \left(\boldsymbol {s} _ {t + 1}, \boldsymbol {a} ^ {\prime}\right) + \omega^ {\prime} \log \pi \left(\boldsymbol {a} ^ {\prime} \mid \boldsymbol {s} _ {t + 1}\right) - \hat {\alpha} \log \hat {\pi} _ {\phi} \left(\boldsymbol {a} ^ {\prime} \mid \boldsymbol {s} _ {t + 1}\right) \right] (9)
$$
The formula for the residual Q-value is derived from the modified Bellman update operator of soft actor-critic [18] following a similar procedure as in Eqn. (5). The detailed derivation is attached to Appendix C for completeness. In the policy improvement step, we improve the policy  $\hat{\pi}_{\phi}$  leveraging the current estimated residual Q-value  $Q_{R,\theta}(s,a)$  and the prior policy  $\pi$  by minimizing the expected KL-divergence between  $\hat{\pi}_{\phi}$  and the desired maximum-entropy policy given the current  $Q_{R,\theta}(s,a)$ , which is equivalent to minimizing the loss function below:
$$
J _ {\hat {\pi}} (\phi) = \mathbb {E} _ {\boldsymbol {s} _ {t} \sim \mathcal {D}} \left[ \mathbb {E} _ {\boldsymbol {a} \sim \hat {\pi} _ {\phi}} \left[ \hat {\alpha} \log \hat {\pi} _ {\phi} (\boldsymbol {a} | \boldsymbol {s} _ {t}) - Q _ {R, \theta} (\boldsymbol {s} _ {t}, \boldsymbol {a}) - \omega^ {\prime} \log \pi (\boldsymbol {a} | \boldsymbol {s} _ {t}) \right] \right]. \tag {10}
$$
In our experiments, we implement the algorithm based on the practical soft actor-critic algorithm presented in [19], which incorporates auto entropy adjustment and double residual Q-functions [14].
Residual Maximum-Entropy MCTS. The last algorithm we introduce is a model-based policy customization algorithm adapted from the maximum-entropy MCTS [55] algorithm, which we refer to as residual maximum-entropy MCTS. Our primary interest is in applying this algorithm in the online setting with an offline learned dynamics model. MCTS essentially solves an online planning
problem leveraging the prior policy and add-on reward function, which enables zero-shot online policy customization. Such a paradigm is particularly useful for problem domains like autonomous driving, where large-scale human driving datasets [11, 58, 53] are available for learning good imitative policies [1, 46] as well as accurate traffic predictive models for online planning [10, 33]. Note that MCTS-based methods inherently assume a discrete action space. In practice, we may extend it to tasks with continuous action space by discretizing the action space [47, 57].
To solve the online planning problem, maximum-entropy MCTS constructs a look-ahead tree  $\mathcal{T}$  incrementally via Monte Carlo simulation. In the tree  $\mathcal{T}$ , each node  $n(s)\in \mathcal{T}$  corresponds to a state  $s$ . For each action  $a\in \mathcal{A}$ , the node  $n(s)$  stores an estimated soft Q-value  $\hat{Q} (s,a)$  and a visit count  $N(s,a)$ . At each iteration, simulations are conducted starting from the root node of the search tree, which consists of two stages: 1) querying a tree policy to select actions and traversing the tree until a leaf node is reached; 2) querying an evaluation function (e.g., Monte Carlo simulations leveraging a roll-out policy) at the leaf node to estimate the return-to-go. After simulation, the estimated returns are propagated backward to update all the nodes along the simulated paths. Afterward, the tree is grown by expanding the leaf nodes reached during the simulations. Different from a typical MCTS algorithm [8], maximum-entropy MCTS exploits maximum-entropy policy optimization. Hence, the tree policy used in maximum-entropy MCTS is defined as:
$$
\hat {\pi} (\boldsymbol {a} | \boldsymbol {s}) = (1 - \lambda_ {s}) \frac {\exp \left(\frac {1}{\hat {\alpha}} \hat {Q} (\boldsymbol {s} , \boldsymbol {a})\right)}{\sum_ {\boldsymbol {a} ^ {\prime} \in \mathcal {A}} \exp \left(\frac {1}{\hat {\alpha}} \hat {Q} (\boldsymbol {s} , \boldsymbol {a})\right)} + \lambda_ {s} \frac {1}{| \mathcal {A} |}, \tag {11}
$$
where  $\lambda_s = \epsilon |\mathcal{A}| / \log (\sum_aN(s,a) + 1)$  with  $\epsilon$  as a hyperparameter. Given a simulated trajectory from the root node to a leaf node, i.e.,  $\{\pmb{s}_0,\pmb{a}_0,\dots ,\pmb{s}_T,\pmb{a}_T\}$ , the Q-values of all the visited nodes are updated with a softmax backup:
$$
\hat {Q} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) = \left\{ \begin{array}{l l} r _ {R} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \omega r \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \gamma \hat {R} & t = T - 1, \\ r _ {R} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \omega r \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \gamma \hat {\alpha} \log \sum_ {\boldsymbol {a} ^ {\prime} \in \mathcal {A}} \exp \left(\frac {1}{\hat {\alpha}} \hat {Q} \left(\boldsymbol {s} _ {t + 1}, \boldsymbol {a} ^ {\prime}\right)\right) & t <   T - 1, \end{array} \right. \tag {12}
$$
where  $\hat{R}$  is the return-to-go estimated by querying an evaluation function on the leaf node  $n(s_{T})$ . To adapt the algorithm to residual maximum-entropy MCTS, we instead have each node store an estimated residual Q-value  $Q_{R}(s,a)$ . The tree policy is revised as: