Policy Decorator: Model-Agnostic Online Refinement for Large Policy Model

We deeply appreciate your review and positive feedback! We are pleased that you find our method simple but effective, and our paper well-written. We value your efforts in helping us enhance the quality of our work. Below, we address your concerns.

---

lacks a couple of related works that need to be discussed

Thank you for recommending these related works! We have cited them in the related works section of the updated PDF. We also discuss them below:

D3IL [1] is a recent benchmark designed to quantitatively evaluate a model’s ability to learn multi-modal behavior. It offers a wide range of multi-modal tasks and demonstrations.

IMC [2] and BESO [3] are both great imitation learning methods that handle multi-modal demonstrations. IMC [2] employs a mixture of experts policy, where each mixture component selects its own subset of the training data for learning. BESO [3] utilizes a generative, score-based diffusion model to learn general-purpose goal-specified policies from large uncurated datasets without rewards.

References:

[1] Jia, Xiaogang, et al. "Towards Diverse Behaviors: A Benchmark for Imitation Learning with Human Demonstrations." The Twelfth International Conference on Learning Representations.

[2] Blessing, Denis, et al. "Information maximizing curriculum: A curriculum-based approach for learning versatile skills." Advances in Neural Information Processing Systems 36 (2024).

[3] Reuss, Moritz, et al. "Goal-conditioned imitation learning using score-based diffusion policies." arXiv preprint arXiv:2304.02532 (2023).

---

However, to my understanding, the residual policy is a Gaussian policy such that a mode collapse will happen once the fine-tuning starts. Does this mean that the overall fine-tuned policy is indeed maximizing performance, but therefore ignores other modes in the behavior?

Adding the residual policy to the multi-modal base policy does not necessarily lead to mode collapse. Instead, we observed that the multi-modal property is typically preserved in our setups due to the residual policy's small scale. We illustrate this point through both an illustrative example and a real case study from our experiments.

Illustrative Example:

As demonstrated in Figure 101, when a bimodal distribution (blue) is combined with a Gaussian distribution (orange), the sum distribution (green) still preserves its bimodal nature. This process effectively shifts the multi-modal distribution and adjusts the standard deviation of its modes. The multi-modal property is maintained as long as the Gaussian distribution's variance remains relatively small compared to the separation between modes.

Implementation Notes:

• The probability density function (PDF) of the bimodal distribution (blue): $f_{Bimodal}(x) = w_1 \cdot \mathcal{N}(x; \mu_1, \sigma_1^2) + w_2 \cdot \mathcal{N}(x; \mu_2, \sigma_2^2)$, where $\mathcal{N}$ represents Gaussian distribution.
• The PDF of the Gaussian distribution (orange): $f_{Gaussian}(x) = \mathcal{N}(x; \mu_3, \sigma_3^2)$
• The PDF of the sum of the two distributions (green) can be computed analytically: $f_{sum}(x) = w_1 \cdot \mathcal{N}(x; \mu_4, \sigma_4^2) + w_2 \cdot \mathcal{N}(x; \mu_5, \sigma_5^2)$, where $\mu_4=\mu_1+\mu_3, \sigma_4=\sqrt{\sigma_1^2+\sigma_3^2}, \mu_5=\mu_2+\mu_3, \sigma_5=\sqrt{\sigma_4^2+\sigma_3^2}$
• The parameters used in the plot: $w_1=0.5, w_2=0.5, \mu_1=0.5, \mu_2=0.5, \mu_3=3, \sigma_1=1, \sigma_2=1, \sigma_3=1$

Real Case Study from Our Experiments:

To demonstrate the preservation of multi-modality in practice, we visualize action distributions from a specific state in the ManiSkill StackCube task, using Behavior Transformer as the base policy. We sampled 1000 actions from both base and residual policies, then applied PCA dimensionality reduction for visualization purposes. We use histograms to visualize these action samples. The results are shown in Figure 102.

We can see that the base policy exhibits a clear bimodal distribution. When combined with the residual policy, the sum distribution maintains its bimodal nature while exhibiting slight shifts in position and variance. Note that the residual policy here is actually a squashed Gaussian distribution (as per SAC [4]) rather than a pure Gaussian, due to SAC's action bounds requirement. This practical example aligns well with our illustrative example, confirming that multi-modal property is preserved in our actual experiments.

We have also included these results in Appendix L of the updated PDF.

References:

[4] Haarnoja, Tuomas, et al. "Soft actor-critic algorithms and applications." arXiv preprint arXiv:1812.05905 (2018).

---

Hope the above clarifications and responses address your concerns. Thank you again for your time and effort!

The paper does not compare to standard, non-large policy model baselines, such as GAIL with an MLP or CNN policy. … it still remains an important baseline to observe whether large policy models are even necessary on these tasks.

Thanks for your suggestions! While our initial submission did not include the exact GAIL + MLP/CNN baseline, we had included ROT (an improved version of GAIL) as a baseline and conducted experiments with MLP and CNN (detailed in Appendix D.1). Following your suggestion, we have now implemented and tested the GAIL + MLP baseline. The results in Figure 104 show that this baseline achieves 0% success rate on StackCube and about 20% success rate on TurnFaucet after 3M environment interactions. These results are expected given that the demonstrations were collected task and motion planning (for StackCube) and model predictive control (for TurnFaucet) - resulting in naturally multi-modal distributions. Additionally, in offline imitation learning scenarios, MLP also performs significantly worse than large policy models, as shown in the table below:

These results together demonstrate that simple MLPs are insufficient for capturing multi-modal distributions and highlight the need for large policy models to effectively utilize multi-modal demonstrations.

We have also included these results in Appendix K of the updated PDF.

---

Is it possible to fix the entropy coefficient during critic pretraining, but then unfreeze it during fine-tuning?

Following the reviewer's suggestion, we fixed the entropy coefficient during the warm-start phase and enabled auto-tuning during subsequent fine-tuning. Results are shown in Figure 105. From six independent runs, three of them still blow up upon unfreezing while the other three remained stable. This result indicates that this unfreezing strategy does not effectively address the training stability issue associated with warm-starting.

We have also included these results in Appendix M of the updated PDF.

---

the paper should identify which tasks may have access to an easily-specified dense reward.

Although we use sparse rewards in our experiments, all tasks in this paper do have existing human-engineered dense rewards:

• ManiSkill StackCube:87 lines of code, 14 tunable hyperparameters
• ManiSkill PegInsertionSide: 82 lines of code, 18 tunable hyperparameters
• ManiSkill TurnFaucet: 41 lines of code, 6 tunable hyperparameters
• ManiSkill PushChiar: 69 lines of code, 18 tunable hyperparameters
• Adroit Door: 18 lines of code, 9 tunable hyperparameters
• Adroit Hammer: 18 lines of code, 10 tunable hyperparameters
• Adroit Pen: 11 lines of code, 8 tunable hyperparameters
• Adroit Relocate: 17 lines of code, 9 tunable hyperparameters

As shown in the above codes, these human-engineered dense rewards are not as “easily-specified” as people may expect. They typically require dozens of lines of Python code and numerous tunable parameters. Designing these rewards manually involves extensive iteration over potential reward terms and tuning hyperparameters through trial and error. This process is laborious but critical for the success of human-engineered rewards.

We have included these dense reward examples in Appendix I to illustrate the challenges in obtaining dense rewards and highlight why sparse-reward settings are preferable.

We are so grateful for such a detailed review and positive evaluation! It is encouraging to receive recognition for our thorough empirical results and clear writing! We sincerely appreciate your efforts to help us improve the quality of our work. Below, we address the concerns you raised.

Note: The outcomes of additional experiments and the related discussions have been added to Appendix I, J, K, L, M of the updated PDF file.

---

In the paper, it remains unjustified whether adding the Gaussian residual policy to the base policy still retains the multimodal capabilities of the large policy model.

In this paper, applying a small residual action to correct a multi-modal base policy typically maintains its multi-modal property. We illustrate this point through both an illustrative example and a real case study from our experiments.

Illustrative Example:

As demonstrated in Figure 101, when a bimodal distribution (blue) is combined with a Gaussian distribution (orange), the sum distribution (green) still preserves its bimodal nature. This process effectively shifts the multi-modal distribution and adjusts the standard deviation of its modes. The multi-modal property is maintained as long as the Gaussian distribution's variance remains relatively small compared to the separation between modes.

Implementation Notes:

• The probability density function (PDF) of the bimodal distribution (blue): $f_{Bimodal}(x) = w_1 \cdot \mathcal{N}(x; \mu_1, \sigma_1^2) + w_2 \cdot \mathcal{N}(x; \mu_2, \sigma_2^2)$, where $\mathcal{N}$ represents Gaussian distribution.
• The PDF of the Gaussian distribution (orange): $f_{Gaussian}(x) = \mathcal{N}(x; \mu_3, \sigma_3^2)$
• The PDF of the sum of the two distributions (green) can be computed analytically: $f_{sum}(x) = w_1 \cdot \mathcal{N}(x; \mu_4, \sigma_4^2) + w_2 \cdot \mathcal{N}(x; \mu_5, \sigma_5^2)$, where $\mu_4=\mu_1+\mu_3, \sigma_4=\sqrt{\sigma_1^2+\sigma_3^2}, \mu_5=\mu_2+\mu_3, \sigma_5=\sqrt{\sigma_4^2+\sigma_3^2}$
• The parameters used in the plot: $w_1=0.5, w_2=0.5, \mu_1=0.5, \mu_2=0.5, \mu_3=3, \sigma_1=1, \sigma_2=1, \sigma_3=1$

Real Case Study from Our Experiments:

To demonstrate the preservation of multi-modality in practice, we visualize action distributions from a specific state in the ManiSkill StackCube task, using Behavior Transformer as the base policy. We sampled 1000 actions from both base and residual policies, then applied PCA dimensionality reduction for visualization purposes. We use histograms to visualize these action samples. The results are shown in Figure 102.

We can see that the base policy exhibits a clear bimodal distribution. When combined with the residual policy, the sum distribution maintains its bimodal nature while exhibiting slight shifts in position and variance. Note that the residual policy here is actually a squashed Gaussian distribution (as per SAC [1]) rather than a pure Gaussian, due to SAC's action bounds requirement. This practical example aligns well with our illustrative example, confirming that multi-modal property is preserved in our actual experiments.

We have also included these results in Appendix L of the updated PDF.

References:

[1] Haarnoja, Tuomas, et al. "Soft actor-critic algorithms and applications." arXiv preprint arXiv:1812.05905 (2018).

---

… how much more costly gradient updates on large policy models are (required by the baselines) than inference (required by both baselines and Policy Decorator).

We conducted benchmarks on the Behavior Transformer's backward pass (gradient update) and forward pass (inference) running times. Results are shown in Figure 103

The results demonstrate that BeT's forward pass is significantly faster than its backward pass, with this gap becoming more pronounced as model size increases. This confirms that the backward pass constitutes the major training time bottleneck.

Implementation Details:

• Batch Size: 1024
• GPU: NVIDIA GeForce RTX 2080 Ti
• Results averaged over 100 independent runs

Additionally, we present the actual training wall-clock time comparison below, demonstrating our Policy Decorator is indeed more time-efficient compared to naive fine-tuning:

We have also included these results in Appendix J of the updated PDF.

Thank you so much for offering thoughtful and constructive feedback! We are delighted to know that you find our experimental evaluation thorough, the performance strong, and the presentation clear. Below, we address the concerns you raised.

---

I would argue that it is essentially the same as the old residual policy learning.

We respectfully contest the viewpoint that “our method is essentially the same as the old residual policy learning.” In contrast, we identify the uncontrolled exploration issue of vanilla residual RL (see example here), and propose a set of controlled exploration strategies to address it. This distinction is discussed throughout the paper, including Sec. 1 (lines 91-101), Sec. 2 (lines 146-148), Sec. 4.2, Sec. 5.3 (lines 427-431), Sec. 5.4.1, Sec. 5.4.2, and Appendix G.1.

More specifically, vanilla residual policy learning is hindered by uncontrolled random exploration in the early stages of training, which results in a lack of success signals to guide learning. To address this limitation, we introduce two controlled exploration strategies: bounded residual action and progressive exploration schedule.

Our empirical results demonstrate that these controlled exploration strategies are crucial for performance and learning efficiency. As illustrated in Figures 6, 7, and 8, our approach significantly outperforms vanilla residual RL, particularly when the base policy's performance is not strong enough. The importance of our controlled exploration strategies is further validated through ablation studies presented in Figures 9, 10, and 11.

---

I don't think it is appropriate for the paper to rebrand this approach as "Policy Decorator".

As explained above, our framework is distinct from residual RL, and we did acknowledge that our framework is built upon residual RL.

The reason why we name our framework “Policy Decorator” is to highlight its model-agnostic nature. While this property may be less significant when using hand-crafted controllers as base policies (e.g., [1][2] in the reviewer's reference), it becomes particularly crucial for improving large policy models. This model-agnostic property is important here because: 1) improving large policy models is an increasingly relevant problem that demands attention; 2) the complicated architectural design of modern large policy models can make the fine-tuning process non-trivial. Therefore, we chose this name to highlight to robot learning researchers that residual-based learning can offer greater versatility and flexibility compared to naive fine-tuning when working with large policy models.

---

Hope the above clarifications and responses address your concerns. If not, feel free to let us know, and we are more than happy to answer additional questions. If you feel that our rebuttal has addressed your concerns, we would be grateful if you would consider revising your score in response. Thanks so much for your time!