SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\hat {\pi} (\boldsymbol {a} | \boldsymbol {s}) = (1 - \lambda_ {s}) \frac {\exp \left(\frac {1}{\hat {\alpha}} \left(Q _ {R} (\boldsymbol {s} , \boldsymbol {a}) + \omega^ {\prime} \log \pi (\boldsymbol {a} | \boldsymbol {s})\right)\right)}{\sum_ {\boldsymbol {a} ^ {\prime} \in \mathcal {A}} \exp \left(\frac {1}{\hat {\alpha}} \left(Q _ {R} (\boldsymbol {s} , \boldsymbol {a}) + \omega^ {\prime} \log \pi (\boldsymbol {a} | \boldsymbol {s})\right)\right)} + \lambda_ {s} \frac {1}{| \mathcal {A} |}. \tag {13}
$$
During backward propagation, the residual Q-value is updated as follows:
$$
Q _ {R} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) = \left\{ \begin{array}{l l} r _ {R} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \gamma R & t = T - 1, \\ r _ {R} \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) + \gamma V _ {R} \left(\boldsymbol {s} _ {t + 1}\right) & t <   T - 1, \end{array} \right. \tag {14}
$$
with  $V_{R}(s_{t + 1}) = \hat{\alpha}\log \sum_{\pmb{a}^{\prime}\in \mathcal{A}}\exp \left(\frac{1}{\hat{\alpha}}\left(Q_{R}(\pmb{s}_{t + 1},\pmb{a}^{\prime}) + \omega^{\prime}\log \pi (\pmb{a}^{\prime}|\pmb{s}_{t + 1})\right)\right)$ . In particular, we use Monte Carlo simulations to estimate the return-to-go  $R$ . The roll-out policy is set as the prior policy  $\pi$ , and  $R$  is computed as the cumulative return of the add-on reward  $r_R$  over the roll-out trajectories. While the estimated  $R$  is biased, it serves as a heuristic to guide the tree search to select nodes with future trajectories collecting high add-on rewards while staying close to the prior policy. The effect of this biased estimation will gradually be mitigated as the search tree expands to the terminal time step.
Our residual Q-learning framework can be considered a principled approach to combine the objectives of IL and RL. Prior works have proposed different approaches for this purpose. In the IL literature, reward augmentation [9] has been introduced as a general framework to provide additional incentives from prior knowledge into the imitation learning process. Under the Generative Adversarial Imitation Learning (GAIL) framework, a surrogate reward can be conveniently incorporated to augment the learned discriminator for policy optimization [35, 3]. The augmented imitative reward is essentially the same as the total reward in our framework. The difference is that ours aims to further customize a trained imitative policy for downstream tasks, instead of regularizing IL with prior knowledge.
There are also a variety of works that directly design the objective function as a weighted mixture of IL and RL objectives. In some methods, like DQfD [21], a supervised classification loss is combined
with a TD loss to train the Q-function. It regularizes the learned Q-function to assign higher values to demonstrated actions. Other methods, such as DAPG [44] and BC-SAC [36], use a weighted mixture of behavior cloning loss and RL objectives in their loss function. SHAIL [25] formulates the imitation problem as a constrained RL problem with a constraint on the lower bound of policy likelihood on demonstrations. These approaches have shown that IL can help accelerate RL learning, especially under sparse rewards [21, 44], and vice versa [25]. As discussed in Sec. 1, similar objectives are also adopted in some RL fine-tuning algorithms [44, 40, 61, 51]. The pre-trained policy is not only used to initialize RL but also leveraged consistently over the fine-tuning stage as regularization. The introduction of RL has also been shown to effectively robustify behavior cloning [36] for challenging and safety-critical autonomous driving tasks. Conversely, IL regularization is used in offline RL approaches, such as TD3+BC [13] and CQL [31], to prevent overestimating the values on out-of-distribution actions and improve the robustness of offline RL policies. Different from these works, where the additional objective is treated as regularization, our framework formulates the learning task as a new MDP and provides a principled way to interpret and design the combined objective.
TS2C [56] is similar to ours from the perspective of having an imperfect expert—the prior policy in RQL can be considered as an imperfect expert for the overall task—and fusing knowledge extracted from this imperfect expert and RL exploration. The difference is that TS2C aims to train a policy towards a known reward; Thus, a value function can be trained to assert if the expert is reliable. In our case, since we rely on the prior policy to encode the basic task reward implicitly, we cannot use such a value function to govern the fusion between imitation and RL. One work in the literature that aligns with our motivation behind MCTS online customization is Deep Imitative Model [45], which aims to direct an imitative policy to arbitrary goals during online usage. It is formulated as a planning problem solved online to find the trajectory that maximizes its posterior likelihood conditioned on the goal. In contrast, we formulate the additional requirements as rewards instead of goals, which are more flexible to specify. Also, RQL is a unified framework for both offline and online customization.
In curriculum RL, the concept of learning the residuals of Q-functions has been adopted in boosted curriculum RL (BCRL) [28, 24]. Under the BCRL framework, a task curriculum is constructed by decomposing a complex task into a sequence of subtasks with increasing difficulty. For each task of the curriculum, its Q-function is modeled and learned as the sum of residuals, with each residual corresponding to a previous task of the curriculum, which was shown to outperform fitting a single function approximator [50, 28]. Similar to RL fine-tuning, BCRL focuses on expediting RL learning on a target task with the Q-functions learned on the previous tasks. The residual is learned to construct the Q-function for the target task with a known reward. Conversely, RQL aims at learning the residual to construct the Q-function for a new task whose reward function combines the unknown underlying reward of the expert and the known add-on reward.
We evaluate the proposed algorithms on four environments selected from different domains: CartPole and Continuous Mountain Car environments from the OpenAI gym classic control suite [5], and Highway and Parking from the highway-env environments [32]. In our experiments, we implemented our algorithms upon Stable-Baselines3 [43] and its imitation library [16]. In Sec. 5.1, we provide the configurations of our experiments, including the settings of policy customization tasks in different environments, baselines, and evaluation metrics. In Sec. 5.2, we present and analyze the experimental results of RL offline policy customization. In Sec. 5.3, we demonstrate an example of zero-shot online customization leveraging the residual maximum-entropy MCTS algorithm. In Sec. 5.4, we investigate two representative RL fine-tuning methods under the context of policy customization and compare them with the proposed algorithms. Please refer to the appendices for detailed experiment configurations, implementation details, ablation studies, and additional experiments in MuJoCo [49].
**Environments.** The detailed configurations of the environments can be found in Appendix D. We mainly describe the policy customization task settings here. In each environment, we design the basic and add-on rewards to illustrate a practical application scenario, where the basic reward corresponds to some basic task requirements and the add-on reward reflects customized specifications such as behavior preference or additional task objectives.
- CartPole balancing: In the cartpole environment, the basic task is to balance the pole as long as possible. During policy customization, we demand an additional task that requires the cart to stay at the center of the rack. To validate if the additional objective is satisfied, we monitor the average absolute position error of the cart which is defined as  $\bar{e}_x = \sum_{t=1}^{T} |x_t| / T$ , where  $x_t$  is the coordinate of the cart with respect to the center of the rack.
- Continuous mountain car: In the mountain car environment, the basic task is to climb the mountain with the least energy consumption. During policy customization, we enforce an additional preference to avoid negative actions whenever possible. This corresponds to use cases where certain actions are less favorable. To monitor the change in action distribution, we count the number of negative actions executed during each episode, denoted as  $n_{\mathrm{neg}}$ .
- Highway navigation: In the highway environment, the basic task is to navigate the vehicle safely and efficiently through the traffic. The vehicle can switch to arbitrary lanes to overtake the other vehicles. During policy customization, we enforce an additional preference to stay on the rightmost lane whenever possible. This reproduces driving habits that human drivers may have in real life. For example, drivers may prefer to drive on the rightmost lane to avoid missing the exits or when they want to drive the car slowly. We compute the average normalized lane index over an episode to validate if the desired behavior emerges, which is defined as  $\bar{I}_{\mathrm{lane}} = \sum_{t=1}^{T} I_{\mathrm{lane},t} / 2T$ , where  $I_{\mathrm{lane},t} \in \{0,1,2\}$  is the lane index with the leftmost lane as the lane 0 and the rightmost lane as the lane 2.
- Parking: In the parking environment, the basic task is to park the car in the target parking slot. During policy customization, we add an additional requirement to avoid touching the boundaries of the parking slots during parking, which is considered good practice to avoid scratching other cars and to leave sufficient space for neighboring parking slots. We compute the non-violation rate to monitor if the requirement is met, which is defined as the percentage of episodes where the constraint is not violated, denoted as  $\gamma_{\mathrm{no - violation}}$ .
Baselines. In our experiments, we mainly compare the performance of four categories of policies:
- RL prior policy: We train an RL policy optimizing the basic reward. It serves as the prior policy for policy customization. For environments with discrete action space (i.e., CartPole and Highway), we train the RL prior policy using soft Q-learning. For environments with continuous action space (i.e., Continuous Mountain Car and Parking), soft actor-critic is used. At test time, we also evaluate its performance on the overall task without customization, which serves as a baseline to show the effectiveness of the proposed algorithms.
- IL prior policy: We train another IL prior policy that imitates the RL prior with GAIL [22]. As discussed in Sec. 2, the proposed RQL algorithms are primarily motivated to customize imitative policy without knowledge of its inherent reward. The IL prior policy serves as a baseline for comparison with the residual-Q policy customized from the IL prior policy. In GAIL, we use soft Q-learning as the policy learning algorithm when the action space is discrete, and use soft actor-critic when the action space is continuous.
- Residual-Q customized policies (Ours): In each environment, we train two residual-Q customized policies leveraging the RL and IL prior policies respectively. They are mainly compared against the corresponding prior policies to validate the effectiveness of policy customization. We apply residual soft Q-learning for environments with discrete action space and apply the residual soft actor-critic for environments with continuous action space. Besides, we evaluate the zero-shot customization ability of the proposed residual maximum-entropy MCTS algorithm in the Highway environment.
- RL with total reward: We also train an RL policy with the total reward (i.e.,  $r + \omega r_R$ ). Note that we do not intend to treat it as a baseline for fair comparison under the policy customization setting. We aim to compare the residual-Q policies with it to validate that the customized policies indeed solve the overall task  $\hat{\mathcal{M}}$ .
Metric. We are mainly interested in validating the performance of the proposed policy customization methods from two perspectives: 1) whether the policy is customized toward the add-on reward; 2) whether the customized policy still possesses the characteristics of the prior policy, in other words, maintains the same level of performance on the basic task. Hence, we use the average basic reward and the success rate of completing the basic task as metrics to evaluate the policies' performance on
the basic task. Meanwhile, we use the average add-on reward and task-specific metrics introduced above as the metrics for evaluating the policies' performance with respect to the customized objective.
5.2 RL Offline Customization
The experimental results of RL offline customization are summarized in Table 1. In all the environments and with either RL or IL prior policies, the customized policies significantly outperform the prior policies on the add-on task objectives. Meanwhile, the customized policies maintain the same level of performance as the prior policies on the basic tasks. It validates that the proposed residual-Q algorithms can tailor the policies for customized requirements while maintaining their original characteristics. One interesting observation is that the policy customization improves the performance of the IL policy on the basic task in the Parking environment. One factor is that GAIL does not perform well in the first place. The Parking environment is indeed more challenging than the other environments given the continuous action space and, more importantly, the nature of the task—navigation in tight space under the non-holonomic constraint. Another aspect we think can explain this phenomenon is that the add-on reward implicitly guides the policy to improve its basic task performance. As shown in Appendix G, it is difficult for the imitative policy to stop the vehicle accurately at the target parking slot. Since GAIL merely learns to match the trajectory distributions under the expert and learned policies without the knowledge of the actual task objective, it is challenging for the imitative policy to precisely localize such a narrow parking space. In contrast, the add-on boundary violation constraint confines the vehicle trajectory within the parking space once the vehicle approaches the target parking slot. Hence, the policy is able to reach the goal point faster and receives higher returns even on the basic parking task.
In Table 1, we also compare the customized policies against the RL policies trained with the total reward function, named RL Full Policy in the table. In environments with discrete action spaces, the residual soft-Q learning policies behave very similarly to the RL full policies. Meanwhile, the residual soft actor-critic agents have different reward patterns compared to the RL full policies. We think the behavior difference is mainly due to the approximation errors in value and policy networks. The proposed residual-Q framework relies on the prior policy to acquire knowledge of the basic task. Thus, a suboptimal prior will affect the performance of the customized policy, especially on the basic task. It also explains the performance difference between the policies customized from RL and IL priors. Inspired by the observation mentioned above on IL policy customization, we think a practical solution to mitigate this issue is partially incorporating some basic task objectives into the add-on reward if available. It is reasonable for many practical use cases even with an imitative prior. While it is difficult to precisely describe the underlying reward function of human experts, we can comfortably assume that they follow some obvious commonsense objectives (e.g., collision avoidance). In this way, our method can be considered a more principled framework for the general synergy between IL and RL beyond policy customization where the reward functions are completely decomposed.
5.3 MCTS Online Customization
We choose the Highway environment to demonstrate the zero-shot online customization ability of residual maximum-entropy MCTS. We implement residual maximum-entropy MCTS based on the MCTS planner provided by the highway-env environment. The planner leverages the ground-truth dynamics model of the environment for roll-out simulation. The results are summarized in Table 1. MCTS online customization is able to achieve similar performance as the RL offline counterpart.
5.4 Ablation Study: Comparison with RL Fine-tuning
In our RQL framework, we formulate policy customization as an MDP whose reward function balances the utilities of the imitative and add-on tasks. It provides the theoretical foundation to determine the synergy between the IL and RL objectives in a principled manner. As discussed in Sec. 1 and Sec. 4, different approaches have been proposed in the RL fine-tuning literature to combine IL and RL objectives. We investigate two representative methods under the context of policy customization. We compare them against RQL and summarize the experimental results in this subsection. Please refer to Appendix F for the complete ablation study.
Greedy Reward Decomposition. One common approach is regularizing the divergence between the trained RL and prior policies during policy updates. It has been widely used in offline RL algorithms
to address the distributional shift issue [30, 54], and Advantage Weighted Actor-Critic (AWAC) [40] demonstrates the benefits of such policy regularization in stabilizing and accelerating RL fine-tuning. As shown in AWAC [40], when using KL-divergence to measure the distance between policies, the regularized optimal policy becomes a maximum-entropy policy with the advantage weighted by the prior policy. If we directly adapt it to solve the policy customization problem, the customized policy is essentially defined as the solution to the following optimization problem at each policy update step:
$$
\tilde {\pi} _ {t} = \arg \max  _ {\tilde {\pi} \in \Pi} \mathbb {E} _ {\boldsymbol {a} \sim \tilde {\pi} (\cdot | \boldsymbol {s})} \left[ \tilde {Q} _ {R, t} (\boldsymbol {s}, \boldsymbol {a}) - \hat {\alpha} \log \tilde {\pi} (\boldsymbol {a} | \boldsymbol {s}) \right] - \lambda \mathcal {D} _ {\mathrm {K L}} (\tilde {\pi} (\cdot | \boldsymbol {s}) \| \pi (\cdot | \boldsymbol {s})) \tag {15}
$$
where  $\tilde{Q}_{R,t}$  is the soft Q-function of the MDP with only the add-on reward  $r_R(s, a)$ , which is defined iteratively with the following update rule:
$$
\tilde {Q} _ {R, t + 1} (\boldsymbol {s}, \boldsymbol {a}) = r _ {R} (\boldsymbol {s}, \boldsymbol {a}) + \gamma \mathbb {E} _ {\boldsymbol {s} ^ {\prime} \sim p (\boldsymbol {s} ^ {\prime} | \boldsymbol {a})} \left[ \mathbb {E} _ {\boldsymbol {a} ^ {\prime} \sim \tilde {\pi} _ {t}} \left[ \tilde {Q} _ {R, t} \left(\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime}\right) - \hat {\alpha} \log \tilde {\pi} _ {t} \left(\boldsymbol {a} ^ {\prime} \mid \boldsymbol {s} ^ {\prime}\right) \right] \right]. \tag {16}
$$
The closed-form solution to the optimization problem is given as:
$$
\tilde {\pi} _ {t} (\boldsymbol {a} | \boldsymbol {s}) \propto \exp \left(\frac {1}{\hat {\alpha} + \lambda} \left(\tilde {Q} _ {R, t} (\boldsymbol {s}, \boldsymbol {a}) + \lambda \log \pi (\boldsymbol {a} | \boldsymbol {s})\right)\right). \tag {17}
$$
We can reduce our RQL policy to the same form by simply replacing the residual Q-function in Eqn. (6) with the soft Q-function  $\tilde{Q}_{R,t}$ . The resulting policy solves the optimization problem in Eqn. (15) but with an entropy weight of  $\hat{\alpha} - \lambda$  in the objective function. If we consider it for the policy customization problem, it is essentially equivalent to heuristically estimating the optimal Q-function as the sum of the optimal Q-functions corresponding to the basic and add-on rewards. Hence, the synthesized policy is not the optimal solution to the target MDP but a greedy approximation [52, 26]. We compare RQL with this greedily customized policy on the Continuous Mountain Car and Parking environments. The results can be found in Appendix F.1. We found that this greedy customized policy performed much worse than RQL in the most challenging Parking environment. Compared to the prior policy, greedy customization compromises the policy's success rate on the basic parking task to customize them towards satisfying the add-on objective. However, the greedy policy still has a significantly lower non-violation rate compared to the RQL policy. It validates that our framework is a more principled method to jointly optimize the IL and RL objectives for policy customization than divergence-based regularization.
Augmenting Reward with Policy KL-Divergence. Alternatively, some works [54, 61] have explored adding the policy divergence to the reward function. It has been shown to outperform policy regularization in offline RL setting [54]. To adopt this line of approaches for policy customization, we formulate the reward function of the target MDP as the sum of the add-on reward and a policy divergence regularization term. In particular, we follow [61] to augment the add-on reward with a penalty whose expectation regularizes the KL-divergence between the customized and prior policies. Formally, the reward function is defined as:
$$
\check {r} _ {t} (\boldsymbol {s}, \boldsymbol {a}) = r _ {R} (\boldsymbol {s}, \boldsymbol {a}) - \beta \log \frac {\check {\pi} _ {t} (\boldsymbol {a} | \boldsymbol {s})}{\pi (\boldsymbol {a} | \boldsymbol {s})}, \tag {18}
$$
where  $\check{\pi}_t$  denotes the customized policy after the  $t^{\text{th}}$  iteration of policy update, and  $\beta$  is a hyperparameter to balance the add-on reward and the policy regularization term. We can then apply standard RL algorithms to customize the policy toward this reward function. However, we found that it made the learning task much more difficult in our policy customization problem—as shown in Appendix F.2, soft actor-critic, for example, failed to find a customized policy that at least performed well on the basic parking task. We believe that one crucial factor contributing to the failure of divergence-based regularization is that the basic and add-on rewards are mostly orthogonal in our setting. The add-on reward only encodes additional task requirements; thus, we heavily rely on the prior policy to embed the desirable behavior on the basic task. Merely regularizing the KL-divergence is therefore insufficient—being close to the prior policy in terms of probabilistic distance does not necessarily imply behavior that accomplishes the basic task objective. Besides, incorporating the policy likelihood into the reward function makes it more challenging to update the Q-function in practice. For example, in Appendix F.3, we show that our RQL algorithm is theoretically equivalent to directly solving an MDP whose reward function is a weighted sum of the add-on reward and the log-likelihood of the prior policy. However, similar to the case of divergence-augmented reward, soft actor-critic failed to find a solution in Parking environment.