Advantage-Aware Policy Optimization for Offline Reinforcement Learning

Table R1. Test returns of our A2PO and new state-of-the-art lightweight baselines. Italics indicates that the results are obtained from the original paper, while the other results are evaluated over 3 random seeds based on the corresponding official code.

Table R2. Test returns of our A2PO and diffusion-based baselines.

We greatly appreciate the reviewer for the insightful comments, which help us improve the quality of the paper significantly. Below we address the main points raised in the review. We will update the additional experimental results in the revision later.

[Weakness 1]: The design of adding advantage values into a jointly trained CVAE can lead to extra training instability.

We agree that introducing an additional network may lead to instability, but we believe it is worthwhile for the enhancement in performance. The experimental results show that our A2PO method exhibits superior performance, particularly in mixed-quality datasets. Moreover, regarding the hyperparameter of CVAE training steps, our A2PO can also demonstrate promising robustness. The number of training steps is a crucial hyperparameter for any deep learning model and is not unique to our proposed CVAE. A lower number of training steps may result in non-convergence of the model, whereas a higher number can lead to overfitting. The ablation study conducted in our paper actually involved selecting training steps at two extreme ends, which is why there was a drastic drop in performance. As shown in Table S4 of our original paper, except for the specific requirements in the antmaze and maze2d environments, a value of $K=2*10^5$ proves to be sufficient for all other scenarios. This highlights the robustness of the CVAE training step across various tasks and datasets.

[Weakness 2]: I don't think the proposed method offers impressive performance.

We have additionally compared our A2PO to the mentioned lightweight baselines, including XQL, SQL, and EQL, and the diffusion-based policy baselines, including IDQL and Diffusion-QL. The corresponding results are shown in Table R1 and Table R2.

The results in Table R1 indicate that our A2PO method can achieve comparable performance to the state-of-the-art lightweight baselines. The results in Table R2 show that our A2PO exhibits performance on par with these more heavyweight diffusion-based methods. Notably, diffusion-based methods often perform poorly when there is a significant difference in data quality within the mixed-quality dataset. This observation underscores the effectiveness of our advantage-aware mechanism, which allows our lightweight CVAE model to capture the multi-modal characteristics from the offline dataset more effectively compared to the heavyweight diffusion models.

While the diffusion-based methods solely rely on the diffusion models to implicitly learn the distribution of offline data, our advantage-aware mechanism explicitly disentangles the action distributions of intertwined behavior policies by modeling the advantage values of all training data as conditional variables. It is noted that these two approaches are orthogonal. This implies that our advantage-aware mechanism can be effectively integrated with the diffusion-based methods by introducing the conditional diffusion models [1], thereby enhancing their capacity to process multi-modal data. This is an interesting direction for future work.

[Weakness 5]: (1) Experiments on datasets mixed with random data is not very meaningful for practical settings. (2) BPPO is a flawed offline RL algorithm and is not worth comparison. (3) Some recent strong and rigorous offline RL algorithms should be compared.

(1) We would like to point out that conducting experiments under various dataset settings is of great significance, as the quality of the dataset is often diverse in many practical scenarios. Random data has a low acquisition cost and also contains some exploration information. The Offline RL technique should be able to utilize datasets of different qualities to improve the agent performance as much as possible instead of only considering high-quality data. Moreover, our A2PO method can leverage arbitrary transition samples to estimate the advantage value for training. This eliminates the need for explicit trajectory return signals, rendering our approach more adaptable and practical for real-world applications where only transition data is accessible and the returns of entire trajectories are not known.

(2) We compare A2PO with advantage-weighted BPPO because there is limited existing research on advantage-based methods. We also agree that it is unreasonable that the official implementation of BPPO uses online evaluation for offline BC model selection. For a fair comparison, in our original reproduction experiment, we have removed online evaluation for offline BC model selection in BPPO. Thus, the reported results in our original manuscript are convincing.

(3) We have added more lightweight baselines (XQL [3], EQL [6], SQL [6]), Diffusion-policy baselines (Diffusion-QL [4], IDQL [7]), and representation learning baseline (BPR [2]) for a much more rigorous comparison as mentioned above. Compared to the additional baselines, our A2PO demonstrates promising results.

[Question 1]: The scores for antmaze-m-d and antmze-l-d look strange.

Sorry for the confusion. The discrepancy in scores arises from that we perform training in the maze2d and antmaze tasks without proper hyperparameters as in the official papers. Specifically, for CQL we follow the default setting of Lagrange threshold $\tau=1.0$ for locomotion tasks instead of $\tau=5.0$ for maze2d/antmaze navigation tasks. For IQL we follow the default setting of expectile $\tau=0.9$ and temperature $\beta=10.0$ for locomotion tasks instead of $\tau=0.7$ and $\beta=3.0$ for maze2d and antmaze navigation tasks. We have corrected these settings and performed new experiments. The results in Table R4 show that compared to CQL and IQL with the official settings, our A2PO can still achieve superior performance. We have corrected these results in our revision. Moreover, we use "v1" version for experiments in the antmaze and maze2d tasks as in LAPO [3], while many of the other baselines use "v0" version (XQL [4], Diffusion-QL [5], D4RL [6]) or "v2" version (EQL [7], SQL [7], IDQL [8]) for experiments. Note that performance is not comparable across task versions.

Table R4. Test returns of CQL and IQL following the hyperparameters in their official papers.

References

[1] Ho, J., & Salimans, T. Classifier-free diffusion guidance. arXiv 2022.

[2] Zhang, H. et al. Behavior Prior Representation Learning for Offline Reinforcement Learning. ICLR 2023.

[3] Chen, X. et al. LAPO: Latent-Variable Advantage-Weighted Policy Optimization for Offline Reinforcement Learning. NeurIPS 2022.

[4] Garg, D., Hejna, J., Geist, M., & Ermon, S. Extreme Q-Learning: MaxEnt RL without Entropy. ICLR 2023.

[5] Wang, Z. et al. Diffusion policies as an expressive policy class for offline reinforcement learning. ICLR 2023.

[6] Fu, J. et al. D4RL: Datasets for Deep Data-Driven Reinforcement Learning. arXiv 2020.

[7] Xu, H., et al. Offline RL with no OOD actions: In-sample learning via implicit value regularization. ICLR 2023.

[8] Hansen-Estruch, P. et al. Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv 2023.

[Weakness 3 & Question 2]: Please report the max-Q and BC trade-off hyperparameter values in the TD3+BC style policy improvement step (Eq.(8)). Has it been tuned for different tasks?

Thanks for the suggestions. As stated in Appendix A.2 (Implementation Details) of our original manuscript, we have illustrated that our implementation follows the TD3BC approach. We follow the same trade-off hyperparameter value as in the official TD3BC implementation and keep this hyperparameter on the whole experiment without tuning.

As suggested by the reviewer, we have added more details about TD3BC in the revision:

(1) In Section 4.2 (Agent Policy Optimazation): "For the policy improvement step, the TD3BC-style loss is defined as:..., following TD3BC, we add a normalization coefficient $\lambda={\alpha}/{(\frac{1}{N}\sum\_{(s\_i, a\_i)} |Q(s\_i, a\_i)|)}$ to the first term to keep the scale balance between Q value objective and regularization, where $\alpha$ is a hyperparameter to control the scale of the normalized Q value."

(2) In Appendix B.2 (Implementation Details): "For the hyperparameter value $\alpha$ in Q normalization, we follow the same setting $\alpha=2.5$ as in the official TD3BC implementation and keep this hyperparameter on the whole experiment."

[Weakness 4]: It is suggested that the proposed method should also be compared with existing offline RL baselines with representation learning.

Thanks for the suggestion. We have additionally compared our A2PO with offline RL baseline with representation learning, BPR [2]. The results show that our A2PO can demonstrate promising results.

Table R3. Test returns of our proposed A2PO and BPR.

Table R1. Test returns of our proposed A2PO and lightweight baselines.

Table R2. Test returns of our proposed A2PO and diffusion-based baselines.

[Comment 4]: it will be helpful to know A2PO's performance on antmaze-v0 or v2 tasks for better comparative evaluation.

We have evaluated A2PO on the antmaze-v0 tasks, as shown in Table R3. The results demonstrate that our A2PO achieves superior performance to other baselines in 'umaze' and 'medium' tasks.

Table R3. Test returns of our proposed A2PO and other state-of-the-art baselines.

Thanks for the response!

[Comment 1]: CVAE increases the burden for hyperparameter tuning. As a fact, in our A2PO implementation and experiments, the only hyperparameter of the advantage-aware CVAE is its training step $K$ with a hyperparameter search only over several values: $1\times10^5$, $2\times10^5$, $4\times10^5$, and $1\times10^6$. As shown in Table S4 of our paper, except for the specific requirements in the antmaze and maze2d environments, a value of $K=2*10^5$ proves to be sufficient for all other scenarios. This highlights the robustness of the CVAE training step across various tasks and datasets, without necessitating excessive hyperparameter tuning. Moreover, the number of training steps is a crucial hyperparameter for any deep learning model and is not unique to our CVAE.

[Comment 2]: the method does not offer impressive performance given its relatively heavy architecture and extra hyperparameters.

We appreciate your insights and are inclined to believe that an additional CVAE, in our view, is not a heavy architecture. Many previous lightweight methods like BCQ and BEAR, as well as recent state-of-the-art methods SPOT and LAPO, all utilize the CVAE to approximate the behavior policy. Our A2PO performance is superior to these baselines and comparable to the heavy diffusion-based methods. It is important to emphasize that the primary efficacy of A2PO stems from this advantage-aware policy, rather than the specific conditional model (be it CVAE or Diffusion). Specifically, we utilize $\pi(a|s,\xi)$ with the advantage value $\xi$ to represent the agent policy, while Diffusion-QL utilizes $\pi(a|s)$ to represent the agent policy. Diffusion-QL introduces a conditional diffusion with only the state as the condition to implicitly capture the multimodal action distribution, overlooking the intrinsic quality of different actions. However, our A2PO is conditioned on both the state and the advantage value. This advantage-aware policy explicitly leverages the advantage value to disentangle the action distributions of intertwined behavior policies, ensuring that the quality of different actions is positively correlated with the advantage value. Thus, our advantage-aware mechanism can also be effectively integrated with the diffusion-based methods, thereby enhancing their capacity to identify multi-modal data with different quality.

[Comment 3]: adding scores of these random-mixed datasets to the total score does not necessarily demonstrate the superiority of the proposed method in practice.

The motivation behind utilizing the combined datasets is to measure the ability of different offline baselines to learn from various types of multi-modal datasets, and more importantly, to reveal whether the agent is able to discriminate the high-quality data for effective learning and not be disturbed by the low-quality data. The final goal is only to highlight the effectiveness of our A2PO when there are substantial gaps between behavior policies. The reason for utilizing the random-incorporated datasets is that the datasets provided by D4RL have only a few types of structures, and we have to combine these finite datasets for more complicated and diverse datasets. Moreover, we give a summarization of the total scores in the locomotion tasks without the random-mixed datasets in Table R1 and Table R2. The results show that our A2PO achieves superior overall performance to the lightweight baselines. The performance of our A2PO is slightly worse than the diffusion-based method Diffusion-QL, but the benefits of diffusion mainly come from its heavyweight foundation model. As stated in comment 2 above, our approach is orthogonal to theirs; our core contribution is not the model itself, but the introduction of an advantage-aware policy.

Dear Reviewer,

We deeply appreciate the time and effort you have invested in reviewing our manuscript. Your expertise and insights are invaluable to us.

As the discussion period is nearing its conclusion, with less than two hours remaining, we are keen to know if our responses have adequately addressed your concerns.

Thank you once again for your attention and consideration. We eagerly await your response.

Sincerely,

Authors of A2PO

[Question 3]: The details of mixed data combination should be provided, and some ablation studies about the amount of each single-quality data are also suggested.

Sorry for the confusion. As stated in Section 5.1 (Tasks and Datasets) of our original paper, we construct mixed datasets using single-quality datasets from the D4RL benchmark. The single-quality datasets are generated with random, medium, and expert behavior policies. Since the D4RL benchmark only includes the two mixed-quality datasets ("medium-replay" and "medium-expert"), we follow the official instructions provided by D4RL (https://github.com/Farama-Foundation/D4RL/wiki/Dataset-Reproducibility-Guide) to manually construct the other mixed-quality datasets ("random-medium", "random-expert", and "random-medium-expert") by directly combining the corresponding single-quality datasets. Each single-quality dataset ("random", "medium", "expert") has 1e6 transitions. The mixed-quality "random-medium" and "random-expert" datasets consist of 2e6 transitions, while the "random-medium-expert" dataset contains 3e6 transitions.

As suggested by the reviewer, we have additionally conducted an ablation study to investigate the impact of varying amounts of single-quality samples in mixed-quality datasets, as shown in Table R2. The results demonstrate the robustness of our A2PO model in handling mixed-quality datasets containing different proportions of single-quality samples.

Table R2. Test returns of our proposed A2PO on the locomotion tasks with the medium-expert dataset containing different proportions of single-quality samples.

References

[1] Wu, J. et al. Supported Policy Optimization for Offline Reinforcement Learning. NeurIPS 2022.

[2] Yang, S. et al. A Behavior Regularized Implicit Policy for Offline Reinforcement Learning. arXiv 2022.

[3] Fu, J. et al. D4RL: Datasets for Deep Data-Driven Reinforcement Learning. arXiv 2020.

[Question 1]: It is suggested to consider the baselines in offline-RL methods.

As suggested by the reviewer, we have added the mentioned baselines for comparison, including BEAR, AWAC, and SPOT, as shown in Table R1. The results show that our A2PO exhibits superior performances compared to the baselines, especially in the mixed-quality datasets.

Table R1. Test returns of our proposed A2PO and the suggested baselines on the locomotion tasks. The italicized results of BEAR are obtained from the D4RL paper [3], while the italicized results of AWAC and SPOT are obtained from the SPOT paper [1]. The other results are evaluated over 3 random seeds based on the corresponding official code.

[Question 2]: The author suggested CVAE to disentangle policies from mixed data. However, the CVAE may not be a good choice for distinguish data from different sources from [1] and [2]. So I wonder why A2PO could achieve significant results in empirical experiments.

Thanks for the insightful comments. Our A2PO can utilize CVAE to disentangle behavior policies from mixed data for achieving superior performance, primarily due to the distinct conditional variables of our CVAE in comparison to those described in the CVAE of reference [1, 2].

(1) Firstly, it should be noted that the state-of-the-art SPOT [1] does not illustrate the limitations of CVAE in distinguishing data. Instead, SPOT also utilizes CVAE to estimate behavior policy outputs, which are then used to provide the constraint penalty term. Reference [2] employs a behavior cloning experiment to evaluate the performance of the CVAE, which reveals that the CVAE overestimates unwanted low data-density regions. However, the CVAE of reference [1, 2] conditiones only on the given state, which tries to distinguish data from different sources without discrimination of the data quality.

(2) In contrast, our advantage-aware CVAE disentangles different behavior policies into separate action distributions by modeling the advantage values of collected state-action pairs as conditioned variable. By conditioning on specific states and advantage values, our CVAE generates corresponding actions, ensuring that the quality of these actions is positively correlated with the advantage value. Thus, A2PO can further optimize the agent toward high advantage values to obtain an effective decision-making policy.

We greatly appreciate the reviewer for the constructive comments, which help us improve the quality of the paper significantly. Below we address the main points raised in the review. We will update the additional experimental results in the revision later.

[Weakness 1 & Weakness 2]: more details about the largest advantage condition $\xi^\*$

Sorry for the confusion. Our advantage-aware CVAE disentangles and infers the action distributions of intertwined behavior policies by modeling the advantage values of all training data as conditional variables.

Concretely, Unlike previous methods [1] utilizing CVAE to predict action solely based on the state $s$, we consider both state $s$ and advantage value $\xi$ for the CVAE condition. For a state-action pair $(s, a)$, the advantage condition $\xi$ can be computed with agent critic as $\xi =\tanh(\min\_{i=1,2} Q\_{\theta\_{i}}(s,a)- V\_\phi(s))$. The state-advantage condition $c$ is formulated as $c=s\parallel\xi$, while the CVAE model is trained using the state-advantage condition $c$ and the corresponding action $a$ with ELBO loss function. Since the range $\xi$ is normalized within $(-1,1)$ with $\tanh$, with the fixed largest advantage condition $\xi^*=1$ the CVAE is able to infer high-quality action distributions associated with the behavior policies from the CVAE decoder. As for the advantage-aware policy $\pi\_\omega(\cdot|c)$, it generates a latent representation $\tilde{z}$ based on the state-advantage condition $c$. $\tilde{z}$ is then fed into the CVAE decoder to get a recognizable action. The optimal action $a^*\_\xi$ is obtained with the designated $\xi^*=1$ condition and the corresponding latent representation $\tilde z^*$ from the policy network.

In short, our CVAE directly learns from all training data with different advantage values, benefiting from the generalization capability to infer the high-advantage action distribution. Considering Equation (7), the action $a^*\_\xi$ with designated $\xi^*=1$ is not selected directly from the dataset samples. Instead, the CVAE decoder infers the high-advantage action distribution only based on the state $s$ and fixed advantage value $\xi^*$. Then the agent can select the optimal action $a^*$ from this distribution.

[Weakness 3]: It is suggested to add some important details in Algorithm 1.

Thanks for the suggestion. We have added details in Algorithm 1, illustrating how $c$ and $c^*$ are computed for both behavior policy disentangling and agent policy optimization.

Dear Reviewer,

We deeply appreciate the time and effort you have invested in reviewing our manuscript. Your expertise and insights are invaluable to us.

As the discussion period is nearing its conclusion, with less than two hours remaining, we are keen to know if our responses have adequately addressed your concerns.

Thank you once again for your attention and consideration. We eagerly await your response.

Sincerely,

Authors of A2PO

[Weakness 2 & Question 2]: the differences in policy optimization between A2PO and LAPO

Thanks for the constructive comments. We would like to point out that A2PO and LAPO stand for advantage-aware and advantage-weighted methods, respectively. These methods do not necessarily require identical structures or loss functions in policy optimization.

Technically, LAPO only adopts the deterministic policy gradient loss in policy optimization, while our A2PO adopts the TD3BC-style loss based on advantage conditions, including the deterministic policy gradient loss and the behavior cloning loss. The first term encourages the optimal policy condition on the largest advantage condition $\xi^*$ to select actions that yield the highest expected returns represented by the Q-value. The second behavior cloning term explicitly imposes constraints on the advantage-aware policy, ensuring the policy selects in-sample actions that adhere to the advantage condition $\xi$ determined by the critic.

We have additionally conducted experiments of A2PO with and without the BC term in the policy optimization step, as shown in Table R1. The results indicate that even without the BC regularization term, A2PO consistently outperforms LAPO in the majority of tasks. Moreover, the BC term in A2PO can enhance its performance in most cases. This comparison highlights the superior performance achieved by A2PO, showcasing its effective disentanglement of action distributions from different behavior policies in order to enforce a reasonable advantage-aware policy constraint and obtain an optimal agent policy.

Table R1. Test returns of LAPO and our proposed A2PO without BC term in policy optimization step in the locomotion tasks. The results of A2PO w/o BC are evaluated over 3 random seeds while the LAPO and A2PO results in our original paper are evaluated over 5 random seeds. Bold indicates that the better performance among LAPO and A2PO w/o BC.

[Question 1]: a minor typo

Thanks. We have corrected it in the paper.

References

[1] Chen, X. et al. LAPO: Latent-Variable Advantage-Weighted Policy Optimization for Offline Reinforcement Learning. NeurIPS 2022.

We sincerely thank the reviewer for the constructive comments, which help us improve the quality of the paper significantly. Below we address the main points raised in the review. We will update the additional experimental results in the revision later.

[Weakness 1]: the contribution about the advantage values

Sorry for the confusion. Our main contribution is the first dedicated attempt toward advantage-aware policy optimization to alleviate the constraint conflict issue under the mixed-quality offline dataset. LAPO [1] and our A2PO solve the constraint conflict issue from two quite different aspects.

(1) Methodology. LAPO is an advantage-weighted method, which employs weighted sampling to prioritize training transitions with high advantage values from the offline dataset. In contrast, our proposed A2PO is an advantage-aware method, which directly conditions the agent policy on the advantage values of all training data without any prior preference.

(2) Effectiveness. The advantage-weighted mechanism of LAPO implicitly reduces the diverse behavior policies associated with the offline dataset into a narrow one from the viewpoint of the dataset redistribution. LAPO inherently biases the learning process towards transitions with higher advantage values, potentially overlooking critical information from lower-advantage transitions. As a result, this redistribution operation can be overfitting on high-advantage data during training, thus impeding the advantage estimation for the effective state-action space. These errors in advantage estimation can further lead to unreliable policy optimization. In contrast, our proposed A2PO disentangles different behavior policies into separate action distributions by modeling the advantage values of collected state-action pairs as conditioned variables, which maintains an unbiased perspective of the entire dataset. Thus, A2PO can perform a more effective advantage estimation than LAPO, as illustrated in Figure 1 and Figure S4 of our original paper. Furthermore, A2PO exhibits superior performances compared to LAPO, as illustrated in Table 1 of our original paper.

In summary, "moving the advantage values from the weight coefficient of the loss function to condition variables" represents a significant conceptual leap, not merely a positional change. While LAPO and A2PO share the common goal of addressing the constraint conflict in mixed-quality datasets, the methodology and effectiveness of our A2PO are distinct and innovative, offering a novel perspective on advantage-aware policy optimization.

Dear Reviewer,

We deeply appreciate the time and effort you have invested in reviewing our manuscript. Your expertise and insights are invaluable to us.

As the discussion period is nearing its conclusion, with less than two hours remaining, we are keen to know if our responses have adequately addressed your concerns.

Thank you once again for your attention and consideration. We eagerly await your response.

Sincerely,

Authors of A2PO

We are glad that the reviewer appreciates our work as a useful contribution to the community. Below we address the main points raised in the review. We will update the additional experimental results in the revision later.

[Weakness 1]: more experiments in the robot manipulation environments

Thanks for the constructive suggestions. We have additionally conducted experiments on the robot manipulation tasks of Adroit and FrankaKitchen, as shown in Table R1. Besides the baselines presented in our original paper, the recent state-of-the-art baseline SPOT [1] is also included for a more comprehensive comparison. The results demonstrate that our A2PO method can achieve comparable performance to the state-of-the-art LAPO [2] and SPOT [1], as well as the other baselines. This underscores the versatility of our A2PO method, highlighting its applicability not only in navigation and locomotion tasks but also in a broader spectrum of robotic manipulation tasks.

Table R1. Test returns of our proposed A2PO and baselines in the robot manipulation tasks of Adroit and FrankaKitchen. Italics indicates that the results are obtained from the D4RL paper [3], while the other results are evaluated over 3 random seeds based on the corresponding official code.

[Weakness 2]: more discussion about reward-conditioned supervised learning

Thanks for the insightful suggestions. As suggested by the reviewer, we have discussed the paper "When does return-conditioned supervised learning work for offline reinforcement learning?" [4] in the related works. The return-conditioned method relies on sampling entire trajectories to obtain return signals for supervised learning. In contrast, our A2PO method can leverage arbitrary transition samples to estimate the normalized advantage value for CVAE training. This eliminates the need for explicit trajectory return signals, rendering our approach more adaptable and practical for real-world applications where only transition data is accessible and the returns of entire trajectories are not known.

[Question 1]: Why do the authors ablate discrete vs continuous advantage variables within the CVAE?

Sorry for the confusion. The ablation analysis in Section 5.3 of the original paper demonstrates the effectiveness of the continuous advantage condition with the normalization trick, which enables the CVAE to disentangle mixed-quality datasets. Moreover, the discrete advantage condition is proposed to stabilize the normalized advantage value, as the advantage value of a given state-action pair may fluctuate during different training phases. Nonetheless, the ablation study indicates that such discretization can only enhance agent performance in several simple tasks. This can be attributed to the fact that discretization potentially impedes the accurate estimation of the advantage values, especially in complex tasks.

[Question 2]: Can the authors clarify what hyper parameters were tuned within A2PO vs the baselines?

Sorry for the confusion. We only tune two hyperparameters in our experiments: the use of discrete $\xi$ and the CVAE training step $K$, as indicated in Table S4 in the original paper. For discrete $\xi$, we observed that the simple tasks, specifically antmaze and maze2d, can benefit from its utilization. In contrast, for the other tasks, the direct use of continuous $\xi$ yielded state-of-the-art results. For the CVAE training step $K$, we conducted a hyperparameter search over $1\times10^5$, $2\times10^5$, $4\times10^5$, and $1\times10^6$. Except for the specific requirements in the antmaze and maze2d environments, a value of $K=2*10^5$ proves to be sufficient for all other scenarios. In short, these two hyperparameters showcase robustness across different tasks and datasets, requiring only minimal adjustments. The remaining hyperparameter settings of A2PO follow the official TD3BC [5] implementation for actor-critic and the official LAPO [2] implementation for CVAE.

References

[1] Wu, J. et al. Supported Policy Optimization for Offline Reinforcement Learning. NeurIPS 2022.

[2] Chen, X. et al. LAPO: Latent-Variable Advantage-Weighted Policy Optimization for Offline Reinforcement Learning. NeurIPS 2022.

[3] Fu, J. et al. D4RL: Datasets for Deep Data-Driven Reinforcement Learning. arXiv 2020.

[4] Brandfonbrener, D. et al. When does return-conditioned supervised learning work for offline reinforcement learning? NeurIPS 2022.

[5] Fujimoto, S., & Gu, S. S. A minimalist approach to offline reinforcement learning. NeurIPS 2021.

We are glad that the reviewer appreciates our attempt, and sincerely thank the reviewer for the constructive comments. Below we address the main points raised in the review. We will update the additional experimental results in the revision later.

[Weakness 1]: compare to the diffusion-based methods

Thanks for the suggestion. We have additionally conducted experiments of diffusion-based offline RL methods, including Difussion-QL [1] and IDQL [2]. The results in Table R1 show that our A2PO exhibits performance on par with these more heavyweight diffusion-based methods. Notably, diffusion-based methods often perform poorly when there is a significant difference in data quality within the mixed-quality dataset. This observation underscores the effectiveness of our advantage-aware mechanism, which allows our lightweight CVAE model to capture the multi-modal characteristics from the offline dataset more effectively compared to the heavyweight diffusion models.

While the diffusion-based methods solely rely on the diffusion models to implicitly learn the distribution of offline data, our advantage-aware mechanism explicitly disentangles the action distributions of intertwined behavior policies by modeling the advantage values of all training data as conditional variables. It is noted that these two approaches are orthogonal. This implies that our advantage-aware mechanism can be effectively integrated with the diffusion-based methods by introducing the conditional diffusion models [3], thereby enhancing their capacity to process multi-modal data. This is an interesting direction for future work.

Table R1. Test returns of our proposed A2PO and diffusion-based baselines. Italics indicates that the results are obtained from the original paper, while the other results are evaluated over 3 random seeds based on the corresponding official code.

[Weakness 2]: determine the number of CVAE training steps

Sorry for the confusion. The number of training steps is a crucial hyperparameter for deep learning models, including our CVAE. A lower number of training steps may result in non-convergence of the model, whereas a higher number can lead to overfitting. In our implementation, we conducted a hyperparameter search on the CVAE training step, denoted as $K$, over $1\times10^5$, $2\times10^5$, $4\times10^5$, and $1\times10^6$. As shown in Table S4 of our original paper, except for the specific requirements in the antmaze and maze2d environments, a value of $K=2*10^5$ proves to be sufficient for all other scenarios. This highlights the robustness of the CVAE training step across various tasks and datasets, without necessitating excessive time expenditure.

[Question 1]: more experiments on the robot manipulation tasks

Thanks for the constructive comments. Due to the time limitation, we have additionally conducted experiments on the robot manipulation tasks using Adroit and FrankaKitchen instead of Meta-World and ManiSkill, following the previous offline RL works [4,5]. The results in Table R2 indicate that our A2PO method can achieve comparable performance to the state-of-the-art LAPO [5] and SPOT [6], as well as the other baselines.

Table R2. Test returns of our proposed A2PO and baselines in the robot manipulation tasks of Adroit and FrankaKitchen.

References

[1] Wang, Z. et al. Diffusion policies as an expressive policy class for offline reinforcement learning. ICLR 2023.

[2] Hansen-Estruch, P. et al. Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv 2023.

[3] Ho, J., & Salimans, T. Classifier-free diffusion guidance. arXiv 2022.

[4] Fu, J. et al. D4RL: Datasets for Deep Data-Driven Reinforcement Learning. arXiv 2020.

[5] Chen, X. et al. LAPO: Latent-Variable Advantage-Weighted Policy Optimization for Offline Reinforcement Learning. NeurIPS 2022.

[6] Wu, J. et al. Supported Policy Optimization for Offline Reinforcement Learning. NeurIPS 2022.

Dear Reviewer,

We sincerely appreciate your positive feedback, the time and effort you have put into reviewing our modifications and responses! It is immensely gratifying to learn that our revisions have successfully addressed your concerns. Your suggestion to further explore this direction using diffusion models is invaluable, and we are enthusiastic about incorporating this perspective into our future research.

Thank you once again for your constructive feedback and support.

Best,

Authors