Paper: Advantage-Aware Policy Optimization for Offline Reinforcement Learning (ICLR-2024)

Reviewer comment (Experiments):
The design of adding advantage values into a jointly trained CVAE can lead to extra training instability. As the value functions are gradually learned, the advantage value can keep changing during the CVAE training. This means that the conditioning information in CVAE is not stationary, causing the learning of CVAE to also fluctuate. It can become worse when using the latent representation from a non-stationary CVAE, which can further cause instability in the value and policy learning. This is evident as the authors have admitted that the training step $K$ of CVAE needs to be carefully tuned, and as shown in Fig. 5, the proposed method can fail miserably if using an improper $K$ value.

Author response:
**[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.

Reviewer comment (Experiments):
The references on offline RL algorithms are mostly restricted to those before 2022. As this field is rapidly developing, there are many strong offline RL algorithms have been proposed in 2023 but lack acknowledgment or comparison. The reported scores are only comparable to some of the newer but more lightweight in-sample learning algorithms like XQL[1] and SQL[2], and less performant than some recent methods that use diffusion policies [3, 4]. Given the proposed method is somewhat heavy (needs to learn an additional, potentially unstable CVAE), I don't think the proposed method offers impressive performance.

Author response:
**[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.

Reviewer comment (Novelty):
The value and policy learning procedure is essentially TD3+BC[5], but conducted in the latent space. However, TD3+BC is never mentioned in Section 4.2. Even in the codes provided by the author, the corresponding method names are "TD3BC_critic_loss" and "TD3BC_actor_loss". Same as TD3+BC, the actual implementation needs to use a $\lambda$ hyperparameter to trade off max Q and minimize the BC penalty. However, this is not mentioned in Eq.(8). This extra hyperparameter is intentionally omitted and corresponding hyperparameter tuning is not discussed. If this hyperparameter is carefully tuned for each task, even the original TD3+BC is likely to have strong performance.

Author response:
**[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."

Reviewer comment (Experiments):
I feel the proposed method is essentially doing some kind of representation learning over TD3+BC, but unlike other representation learning methods, it could be more unstable. It is suggested that the proposed method should also be compared with existing offline RL baselines with representation learning.

Author response:
**[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. 
| Env                                 | BPR                   | A2PO                  |
| ----------------------------------- | --------------------- | --------------------- |
| halfcheetah-medium-v2               | **47.25 $\pm$ 0.34**  | 47.09 $\pm$ 0.17      |
| hopper-medium-v2                    | 58.16 $\pm$ 4.32      | **80.29 $\pm$ 3.95**  |
| walker2d-medium-v2                  | 82.74 $\pm$ 1.83      | **84.88 $\pm$ 0.23**  |
| halfcheetah-medium-replay-v2        | 41.20 $\pm$ 2.43      | **44.74 $\pm$ 0.22**  |
| hopper-medium-replay-v2             | 41.86 $\pm$ 7.98      | **101.59 $\pm$ 1.25** |
| walker2d-medium-replay-v2           | 83.30 $\pm$ 25.11     | **82.82 $\pm$ 1.70**  |
| halfcheetah-medium-expert-v2        | 95.16 $\pm$ 0.94      | **95.61 $\pm$ 0.54**  |
| hopper-medium-expert-v2             | **110.18 $\pm$ 2.72** | 105.44 $\pm$ 0.56     |
| walker2d-medium-expert-v2           | 109.25 $\pm$ 0.28     | **112.13 $\pm$ 0.24** |
| halfcheetah-random-expert-v2        | 3.10 $\pm$ 1.07       | **90.32 $\pm$ 1.63**  |
| hopper-random-expert-v2             | 53.96 $\pm$ 16.22     | **105.19 $\pm$ 4.54** |
| walker2d-random-expert-v2           | 27.60 $\pm$ 26.97     | **91.96 $\pm$ 10.98** |
| halfcheetah-random-medium-v2        | 39.86 $\pm$ 1.37      | **45.20 $\pm$ 0.21**  |
| hopper-random-medium-v2             | 3.99 $\pm$ 0.60       | **7.14 $\pm$ 0.35**   |
| walker2d-random-medium-v2           | 31.73 $\pm$ 29.98     | **75.80 $\pm$ 2.12**  |
| halfcheetah-random-medium-expert-v2 | 26.40 $\pm$ 11.76     | **90.58 $\pm$ 1.44**  |
| hopper-random-medium-expert-v2      | 46.67 $\pm$ 7.65      | **107.84 $\pm$ 0.42** |
| walker2d-random-medium-expert-v2    | **105.25 $\pm$ 2.08** | 97.71 $\pm$ 6.74      |
| **Total**    | 1007.66 |  **1466.33**    |

Reviewer comment (Experiments):
In the experiments, adding additional experiments on datasets mixed with random data is not very meaningful for practical settings. As nobody will use bad random data to train their offline RL policy in practice. Moreover, BPPO is a flawed offline RL algorithm and is not worth comparison. In its code implementation, it hides online evaluation during offline training, which is a kind of cheating in the offline RL setting. It is a shame that such a paper even gets accepted in ICLR 2023. On the other hand, some recent strong and rigorous offline RL algorithms should be compared.

Author response:
**[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.

Reviewer comment (Evaluation):
The scores for antmaze-m-d and antmze-l-d look strange, CQL and IQL are reported to have reasonable scores in the IQL paper as well as in previous references [1-2,4], however, their scores are all zero in this paper. Moreover, the proposed A2PO's performance looks very bad, inferior to scores of baselines as reported in other papers. Why is that?

Author response:
**[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. 


| Env               | CQL              | IQL                   | A2PO                  |
| ------------------------- | ---------------- | --------------------- | --------------------- |
| maze2d-umaze-v1           | 17.02 $\pm$ 1.87 | 56.17 $\pm$ 9.86      | **133.27 $\pm$ 9.58** |
| maze2d-medium-v1          | 22.45 $\pm$ 6.70 | 25.67 $\pm$ 16.93     | **83.95 $\pm$ 10.56** |
| maze2d-large-v1           | 2.53 $\pm$ 6.58  | 45.67 $\pm$ 18.91     | **127.61 $\pm$ 5.35** |
| antmaze-umaze-diverse-v1  | 80.00 $\pm$ 8.16 | 86.67 $\pm$ 12.47     | **96.66 $\pm$ 4.71**  |
| antmaze-medium-diverse-v1 | 0.00 $\pm$ 0.00  | 46.67 $\pm$ 18.86     | **50.00 $\pm$ 15.25** |
| antmaze-large-diverse-v1  | 6.67 $\pm$ 4.71  | **43.33 $\pm$ 12.47** | 6.16 $\pm$ 4.90       |
| **Total** | 128.67  | 304.18 | **497.49**      |

Reviewer comment (Theory):
Follow the last point, if the $\xi$ is normalized to [0,1], the largest value of $\xi=1$ only corresponds to single data point in the mixed dataset, how do you learn the action distribution from CVAE with this single data point?

Author response:
**[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.

Reviewer comment (Experiments):
Some offline-RL baselines, such as BEAR/AWAC are missed in experiment, and the recent popular SPOT [1] method is not considered in experiments as well. It is suggested to consider the baselines in offline-RL methods.

Author response:
**[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.

| Env                                 | BEAR             | AWAC                 | SPOT                | A2PO                  |
| ----------------------------------- | ---------------- | -------------------- | ------------------- | --------------------- |
| halfcheetah-medium-v2               | *41.7*           | *43.5*               | ***58.4***          | 47.09 $\pm$ 0.17      |
| hopper-medium-v2                    | *52.1*           | *57.0*               | ***86.0***          | 80.29 $\pm$ 3.95      |
| walker2d-medium-v2                  | *59.1*           | *72.4*               | ***86.4***          | 84.88 $\pm$ 0.23      |
| halfcheetah-medium-replay-v2        | *38.6*           | *40.5*               | ***52.2***          | 44.74 $\pm$ 0.22      |
| hopper-medium-replay-v2             | *19.2*           | *37.2*               | *100.2*             | **101.59 $\pm$ 1.25** |
| walker2d-medium-replay-v2           | *33.7*           | *27.0*               | ***91.6***          | 82.82 $\pm$ 1.70      |
| halfcheetah-medium-expert-v2        | *53.4*           | *42.8*               | *86.9*              | **95.61 $\pm$ 0.54**  |
| hopper-medium-expert-v2             | *40.1*           | *55.8*               | *99.3*              | **105.44 $\pm$ 0.56** |
| walker2d-medium-expert-v2           | *96.3*           | *74.5*               | *112.0*             | **112.13 $\pm$ 0.24** |
| halfcheetah-random-expert-v2        | 4.18 $\pm$ 1.79  | 87.32 $\pm$ 2.91     | 60.42 $\pm$ 3.69    | **90.32 $\pm$ 1.63**  |
| hopper-random-expert-v2             | 5.73 $\pm$ 3.63  | 84.70 $\pm$ 19.83    | 98.60 $\pm$ 5.32    | **105.19 $\pm$ 4.54** |
| walker2d-random-expert-v2           | -0.36 $\pm$ 0.00 | 11.70 $\pm$ 12.04    | 7.20 $\pm$ 9.96     | **91.96 $\pm$ 10.98** |
| halfcheetah-random-medium-v2        | 13.78 $\pm$ 4.87 | **46.54 $\pm$ 0.10** | 46.12 $\pm$ 0.25    | 45.20 $\pm$ 0.21      |
| hopper-random-medium-v2    | 1.40 $\pm$ 0.58  | 19.45 $\pm$ 11.87    | **7.78 $\pm$ 0.87** | 7.14 $\pm$ 0.35       |
| walker2d-random-medium-v2       | -0.46 $\pm$ 0.07 | -0.03 $\pm$ 0.08     | 7.77 $\pm$ 4.10     | **75.80 $\pm$ 2.12**  |
| halfcheetah-random-medium-expert-v2 | 2.25 $\pm$ 0.00  | 89.65 $\pm$ 1.67     | 80.18 $\pm$ 7.47    | **90.58 $\pm$ 1.44**  |
| hopper-random-medium-expert-v2      | 8.62 $\pm$ 2.10  | 30.31 $\pm$ 7.46     | 40.19 $\pm$ 23.07   | **107.84 $\pm$ 0.42** |
| walker2d-random-medium-expert-v2    | -0.41 $\pm$ 0.65 | -0.31 $\pm$ 0.07     | 10.31 $\pm$ 4.63    | **97.71 $\pm$ 6.74**  |
| **Total** | 468.93 | 1205.23 | 1131.57   | **1466.33**  |

Reviewer comment (Experiments):
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.

Author response:
**[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.

Reviewer comment (Experiments):
In the experiment, the authors manually combines the single-quality datasets to setup the mixed datasets. The details of mixed data combination should be provided, and some ablation studies about the amount of each single-quality data are also suggested.

Author response:
**[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.
|  Env           | medium:expert=3:1 | medium:expert=2:1 | medium:expert=1:1     | medium:expert=1:2     | medium:expert=1:3 |
| ----------- | ----------------- | ----------------- | --------------------- | --------------------- | ----------------- |
| halfcheetah | 93.78 $\pm$ 0.70  | 94.28 $\pm$ 0.09  | **95.61 $\pm$ 0.54**  | 95.30 $\pm$ 0.25      | 95.06 $\pm$ 0.22  |
| hopper      | 106.94 $\pm$ 0.67 | 74.72 $\pm$ 25.18 | 105.44 $\pm$ 0.56     | **112.19 $\pm$ 0.09** | 110.17 $\pm$ 1.37 |
| walker2d    | 110.94 $\pm$ 0.36 | 110.77 $\pm$ 0.57 | **112.13 $\pm$ 0.24** | 111.26 $\pm$ 0.49     | 110.67 $\pm$ 0.21 |
| **Total**                               | 311.66           | 279.77     | 313.18  | **318.75**    | 315.90 |

Reviewer comment (Novelty):
The contribution is a bit incremental, by mainly moving the advantage values from weight coefficient of the loss function (as in LAPO equation [7]) to condition variables (as in this paper equation [5]).

Author response:
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.

Reviewer comment (Writing):
A minor typo on Page 8: “The performance comparison of different discrete advantage conditions for test is given in Figure **4**”.

Author response:
Thanks. We have corrected it in the paper.

Reviewer comment (Experiments):
I wonder how A2PO compares to LAPO if the policy optimization steps are made the same.

Author response:
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.
|  Environment                                 | LAPO                  | A2PO w/o BC           | A2PO              |
| ----------------------------------- | --------------------- | --------------------- | ----------------- |
| halfcheetah-medium-v2               | 45.58 $\pm$ 0.06      | **46.81 $\pm$ 0.09**  | 47.09 $\pm$ 0.17  |
| hopper-medium-v2                    | 52.53 $\pm$ 2.61      | **70.08 $\pm$ 4.03**  | 80.29 $\pm$ 3.95  |
| walker2d-medium-v2                  | 80.46 $\pm$ 1.25      | **81.97 $\pm$ 1.12**  | 84.88 $\pm$ 0.23  |
| halfcheetah-medium-replay-v2        | 41.94 $\pm$ 0.47      | **42.01 $\pm$ 0.28**  | 44.74 $\pm$ 0.22  |
| hopper-medium-replay-v2             | 50.14 $\pm$ 11.16     | **96.51 $\pm$ 1.47**  | 101.59 $\pm$ 1.25 |
| walker2d-medium-replay-v2           | 60.55 $\pm$ 10.45     | **71.09 $\pm$ 7.98**  | 82.82 $\pm$ 1.70  |
| halfcheetah-medium-expert-v2        | 94.22 $\pm$ 0 .46     | **94.29 $\pm$ 0.03**  | 95.61 $\pm$ 0.54  |
| hopper-medium-expert-v2             | **111.04 $\pm$ 0.36** | 107.27 $\pm$ 1.93     | 105.44 $\pm$ 0.56 |
| walker2d-medium-expert-v2           | 110.88 $\pm$ 0.15     | **111.61 $\pm$ 0.07** | 112.13 $\pm$ 0.24 |
| halfcheetah-random-expert-v2        | **52.58 $\pm$ 17.30** | 31.37 $\pm$ 6.27      | 90.32 $\pm$ 1.63  |
| hopper-random-expert-v2             | 82.33 $\pm$ 18.95     | **113.20 $\pm$ 1.20** | 105.19 $\pm$ 4.54 |
| walker2d-random-expert-v2           | 0.89 $\pm$ 0.53       | **66.82 $\pm$ 11.01** | 91.96 $\pm$ 10.98 |
| halfcheetah-random-medium-v2        | 18.53 $\pm$ 0.99      | **43.19 $\pm$ 0.54**  | 45.20 $\pm$ 0.21  |
| hopper-random-medium-v2             | **4.17 $\pm$ 3.11**   | 1.59 $\pm$ 0.92       | 7.14 $\pm$ 0.35   |
| walker2d-random-medium-v2           | 23.65 $\pm$ 33.97     | **72.32 $\pm$ 4.43**  | 75.80 $\pm$ 2.12  |
| halfcheetah-random-medium-expert-v2 | **71.09 $\pm$ 0.47**  | 70.77 $\pm$ 4.21      | 90.58 $\pm$ 1.44  |
| hopper-random-medium-expert-v2      | 66.59 $\pm$ 19.29     | **86.54 $\pm$ 7.25**  | 107.84 $\pm$ 0.42 |
| walker2d-random-medium-expert-v2    | 60.41 $\pm$ 43.22     | **110.35 $\pm$ 1.20** | 97.71 $\pm$ 6.74  |
| **Total**  | 1027.58   | **1317.79** | 1466.33  |

Reviewer comment (Novelty):
The idea of conditioning on a proxy for policy return (i.e., the advantage) itself is not novel. It would be helpful for the authors to include a reference to reward-conditioned supervised learning [Brandfonbrener et. al., 2023].

Author response:
**[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.

Reviewer comment (Experiments):
Why do the authors ablate discrete vs continuous advantage variables within the CVAE? For a distribution that is inherently categorical, I could understand why this ablation would be useful, but because the space of advantages is continuous anyways I don't understand what the reader should takeaway from this experiment.

Author response:
**[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.

Reviewer comment (Reproducibility):
The margins for improvement between A2PO and the baselines is quite narrow in some cases. It's possible that the ordering of the results could change with more or less tuning. Can the authors clarify what hyper parameters were tuned within A2PO vs the baselines?

Author response:
**[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.

Reviewer comment (Experiments):
A comparison with diffusion-based methods is lacking. Figure 2 shows that one feature of the proposed method is its ability to model multi-modal actions conditioned on the given state. Recently, diffusion models have become a popular approach for this purpose. How does your method compare to diffusion-based methods? Here are some related papers:  
  - Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning  
  - Planning with Diffusion for Flexible Behavior Synthesis  
  - IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies

Author response:
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. 

| Environment                                 | Diffusion-QL         | IDQL              | A2PO                  |
| ----------------------------------- | -------------------- | ----------------- | --------------------- |
| halfcheetah-medium-v2               | ***51.1***           | *51*              | 47.09 $\pm$ 0.17      |
| hopper-medium-v2                    | ***90.5***           | *65.4*            | 80.29 $\pm$ 3.95      |
| walker2d-medium-v2                  | ***87.0***           | *82.5*            | 84.88 $\pm$ 0.23      |
| halfcheetah-medium-replay-v2        | ***47.8***           | *45.9*            | 44.74 $\pm$ 0.22      |
| hopper-medium-replay-v2             | *95.5*               | *92.1*            | **101.59 $\pm$ 1.25** |
| walker2d-medium-replay-v2           | ***101.3***          | *85.1*            | 82.82 $\pm$ 1.70      |
| halfcheetah-medium-expert-v2        | ***96.8***           | *95.9*            | 95.61 $\pm$ 0.54      |
| hopper-medium-expert-v2             | ***111.1***          | *108.6*           | 105.44 $\pm$ 0.56     |
| walker2d-medium-expert-v2           | *110.1*              | ***112.7***       | 112.13 $\pm$ 0.24     |
| halfcheetah-random-expert-v2        | 86.07 $\pm$ 1.49     | 32.55 $\pm$ 2.94  | **90.32 $\pm$ 1.63**  |
| hopper-random-expert-v2             | 101.96 $\pm$ 7.03    | 19.73 $\pm$ 13.80 | **105.19 $\pm$ 4.54** |
| walker2d-random-expert-v2           | 56.33 $\pm$ 31.44    | 0.18 $\pm$ 0.08   | **91.96 $\pm$ 10.98** |
| halfcheetah-random-medium-v2        | **48.43 $\pm$ 0.32** | 6.26 $\pm$ 1.18   | 45.20 $\pm$ 0.21      |
| hopper-random-medium-v2             | 6.93 $\pm$ 0.75      | 2.78 $\pm$ 2.01   | **7.14 $\pm$ 0.35**   |
| walker2d-random-medium-v2           | 3.27 $\pm$ 2.35      | 3.82 $\pm$ 3.21   | **75.80 $\pm$ 2.12**  |
| halfcheetah-random-medium-expert-v2 | 81.15 $\pm$ 7.21     | 36.24 $\pm$ 13.27 | **90.58 $\pm$ 1.44**  |
| hopper-random-medium-expert-v2      | 70.09 $\pm$ 2.05     | 6.17 $\pm$ 3.31   | **107.84 $\pm$ 0.42** |
| walker2d-random-medium-expert-v2    | 56.56 $\pm$ 23.02    | 18.55 $\pm$ 8.13  | **97.71 $\pm$ 6.74**  |
| **Total**  | 1301.99  | 865.48  | **1466.33**  |

Reviewer comment (Experiments):
As shown in Figure 45, the CVAE over the action space is sensitive to the number of training steps. How can we determine the optimal number of training steps before training an offline RL agent? Is this stage time-consuming?

Author response:
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.

Reviewer comment (Experiments):
Is it possible to use this algorithm in various robot manipulation environments, such as Meta-World and ManiSkill? Additionally, is this stage time-consuming?

Author response:
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.

| Environment         | BC     | BCQ    | CQL       | IQL        | TD3BC            | SPOT              | BPPO              | LAPO                  | A2PO                  |
| ------------------- | ------ | ------ | --------- | ---------- | ---------------- | ----------------- | ----------------- | --------------------- | --------------------- |
| kitchen-complete-v0 | *33.8* | *8.1*  | *43.8*    | ***62.5*** | 0.83 $\pm$ 1.18  | 41.67 $\pm$ 11.40 | 0.00 $\pm$ 0.00   | 50.83 $\pm$ 10.27     | 60.00 $\pm$ 2.04      |
| kitchen-partial-v0  | *33.8* | *18.9* | *49.8*    | *46.3*     | 0.00 $\pm$ 0.00  | 0.00 $\pm$ 0.00   | 16.67 $\pm$ 2.36  | **58.75 $\pm$ 16.25** | 48.33 $\pm$ 4.08      |
| kitchen-mixed-v0    | *47.5* | *10.6* | *51.0*    | *51.0*     | 0.83 $\pm$ 1.18  | 0.00 $\pm$ 0.00   | 0.00 $\pm$ 0.00   | 52.33 $\pm$ 6.27      | **53.33 $\pm$ 2.36**  |
| pen-human-v1        | *34.4* | *12.3* | *37.5*    | ***71.5*** | -3.69 $\pm$ 0.38 | 32.70 $\pm$ 11.40 | 25.77 $\pm$ 19.25 | 68.06 $\pm$ 18.01     | 68.94 $\pm$ 5.90      |
| hammer-human-v1     | *1.5*  | *1.2*  | ***4.4*** | *1.4*      | 0.71 $\pm$ 0.25  | 1.99 $\pm$ 0.07   | 0.50 $\pm$ 0.35   | 1.12 $\pm$ 0.29       | 1.84 $\pm$ 0.35       |
| door-human-v1       | *0.5*  | *0.4*  | ***9.9*** | *4.3*      | -0.33 $\pm$ 0.01 | -0.33 $\pm$ 0.01  | -0.02 $\pm$ 0.03  | 6.07 $\pm$ 4.60       | 8.51 $\pm$ 3.66       |
| relocate-human-v1   | *0.0*  | *0.0*  | *0.2*     | *0.1*      | -0.30 $\pm$ 0.01 | -0.07 $\pm$ 0.13  | -0.08 $\pm$ 0.08  | 0.04 $\pm$ 0.01       | **0.49 $\pm$ 0.58**   |
| pen-cloned-v1       | *56.9* | *28.0* | *39.2*    | *37.3*     | 1.74 $\pm$ 2.90  | 2.54 $\pm$ 7.96   | 21.77 $\pm$ 5.88  | 55.84 $\pm$ 20.29     | **84.80 $\pm$ 21.43** |
| hammer-cloned-v1    | *0.8*  | *0.4*  | ***2.1*** | ***2.1***  | 0.26 $\pm$ 0.02  | 0.46 $\pm$ 0.03   | 0.25 $\pm$ 0.22   | 0.84 $\pm$ 0.35       | 0.43 $\pm$ 0.15       |
| door-cloned-v1      | *-0.1* | *0.0*  | *0.4*     | ***1.6***  | -0.35 $\pm$ 0.02 | -0.35 $\pm$ 0.02  | -0.04 $\pm$ 0.05  | 0.22 $\pm$ 0.40       | 0.31 $\pm$ 0.04       |
| relocate-cloned-v1  | -0.1   | *-0.2* | -0.1      | *-0.2*     | -0.31 $\pm$ 0.01 | -0.30 $\pm$ 0.01  | -0.14 $\pm$ 0.11  | -0.05 $\pm$ 0.07      | **0.00 $\pm$ 0.04**   |
| **Total**  | 209.00   | 79.70 | 238.20  | 277.90   | -0.61 | 78.31  | 64.68  | 294.05  | **326.98**  |