A2PO: Towards Effective Offline Reinforcement Learning from an Advantage-aware Perspective

Table R1. Test returns of our A2PO and return-conditioned and data rebalancing baselines. Bold indicates the best performance among the two algorithms. Italic indicates that the score are directly taken from the original paper. For the newly-constructed dataset, we rerun %BC, DT with the official code from [1]. As ReD[2] is a method without publicly available source code, we reimplement it ourselves.

Table R2. Test returns of recent policy constraints Offline RL method PRDC and our A2PO.

Thanks for your support on the idea, writing, and experiment. And we will solve your doubts with text description and experimental support.

[Q1: Although this paper presents a new approach, the comparison with existing techniques may not be in depth enough and lacks an analysis of why the algorithm works.]

Sorry for the confusion. In our original manuscript, we have compared our A2PO with other advanced offline RL methods. These techniques include policy constraint methods such as TD3+BC, value regularization methods like CQL, model-based method MOPO, diffusion-based method Diffusion-QL, and advantage-weighted methods like LAPO. Moreover, we have additionally conducted experiments to evaluate the return conditioned methods DT and %BC [1], as well as the data rebalancing techniques ReD [2] and its variation DeReD, as presented in Table R1. To validate the efficacy of the A2PO components, we have also conducted various ablation studies in the original paper. And we have the following key observations:

1. In Figure 1 of the paper, the toy demo showcases the erroneous constraint issue of the AW method and highlights the precise behavior distribution modeling capability of our A2PO, while the following comparing experiments further confirmed this point.
2. In comparison to conventional offline rl methods like TD3+BC/EQL and return-conditioned methods such as DT/%BC, the superiority of our A2PO becomes increasingly apparent as the substantial gaps between behavior policies become larger from medium, medium-expert to random-expert, random-expert-medium.
3. Although diffusion-based method Diffusion-QL precisely models the action distribution and gets comparable performance in several scenarios, our A2PO outperforms Diffusion-QL in most cases, particularly on mixed-quality datasets while being less time-consuming, as shown in Figure 7 of the paper.
4. The ablation study on the CVAE and the advantage-aware policy constraint (Appendices E, F) demonstrates the effectiveness of advantage-aware policy optimization.

These key observations highlight the effectiveness of A2PO to capture high-quality interactions within the dataset for a reasonable advantage-aware policy constraint and policy optimization. Unlike previous AW methods focusing on redistributing the dataset but inadvertently reduce data diversity, our A2PO precisely disentangles and models the mixed behavior policies with advantage-condition technique on CVAE, then the agent can follow such disentangled action distribution constraints to optimize the advantage-aware policy towards high advantage values, therefore consistently improving the performance.

[Q2: In addition, could the authors further discuss the conditions under which A2PO may fail?]

Thank you for the constructive comments. The results of the original paper indicate that our A2PO does not perform well when the offline dataset distribution is narrow, such as halfcheetah-medium. When the dataset distribution is severely restricted, A2PO struggles to efficiently train both the CVAE and the agent model, which can result in task failure.

[Q3: This paper focuses on the analysis and comparison of advantage-weighted offline RL methods, ignoring recent work on policy constraints in the same relative as A2PO, such as PRDC: Policy Regularization with Dataset Constraint for Offline Reinforcement Learning.]

Thanks for the insight suggestion! We have additionally conducted experiments to compare the recently proposed policy-constrained method, PRDC [3] with our A2PO The results are presented in Table R1, which show that our A2PO consistently achieves superior performance on the majority of the gym, maze, antmaze, and adroit tasks. This comparison clearly underscores the superior performance and effectiveness of advantage-aware policy constraints employed in A2PO for offline policy optimization.

[Q4: Regarding implementation details, how can we ensure a sufficient number of high-quality samples in the dataset to support training when selecting ξ = 1? ]

It is challenging to ensure that the dataset comprises an adequate number of high-quality samples. We address this limitation by utilizing the generalization capabilities of the generative model CVAE. By selecting $\xi=1$, the advantage-aware CVAE can generate and infer the optimal actions. As indicated in Table 6 of Appendix E, by directly inputting $\xi=1$, the CVAE-generated action $a\sim p_\psi(\cdot|z\sim\mathcal N(0,I),c)$ can also achieve high performance in various scenarios. Furthermore, we have also investigated the scenario where the dataset contains only a small quantity of high-quality samples, as shown in Appendix H. The results demonstrate that our A2PO algorithm is capable of achieving expert-level performance, even with a limited number of high-quality samples. These findings underscore the generalization ability of the advantage-aware CVAE and the robustness of our A2PO algorithm in deriving optimal policies across diverse structures of offline datasets.

[ Q5: How can we guarantee the accuracy of the introduced advantage information? ]

Thanks for the insightful comments. The advantage value is determined by both the critic and the advantage estimation method. It is important to clarify that we do not propose a novel advantage computation method; rather, we adopt the widely accepted baseline in offline reinforcement learning, TD3+BC [4], as our framework. We utilize the advantage definition $A(s, a) = Q(s, a) - V(s)$ for advantage estimation and subsequent normalization, ensuring that the agent maintains an effective estimation.

[1] Decision Transformer: Reinforcement Learning via Sequence Modeling. NeurIPS 2021

[2] Boosting Offline Reinforcement Learning via Data Rebalancing. arXiv 2022

[3] Policy Regularization with Dataset Constraint for Offline Reinforcement Learning. ICML 2023

[4] A Minimalist Approach to Offline Reinforcement Learning. NeurIPS 2021

Table R2. Test returns of our A2PO and other return conditioned and data rebalancing baselines. Italic indicates that the scores are directly taken from the original paper. For the newly constructed dataset, we rerun %BC and DT with the official code from [2]. As ReD[3] is a method without publicly available source code, we reimplement it ourselves.

Thanks for your positive comments on the experiment's completeness and the writing's quality. We will address your concerns in this section.

[Q1.1: The Agent policy optimization is unclear: Why is a behavior regularization term (2nd term in (8)) needed when you have a generative model trained on the dataset?]

Sorry for the confusion. First of all, as an offline RL method with policy constraints, A2PO requires the agent policy to closely align with the behavior policy. In order to achieve this, we incorporate a behavior regularization term that constrains the action distribution of the advantage-aware agent based on the computed advantage value.

Meanwhile, although latent action space inherently imposes a constraint on the action output [1], the regularization term ensures that the actions chosen by the A2PO policy align with the advantage condition $\xi$ determined by the critic, thereby providing a more precise policy constraint. To confirm our perspective, we conducted an ablation study by eliminating the regularization term in Table 7 and Appendix F of the original paper. The results of this study demonstrate that the regularization term in A2PO enhances its performance in most cases, providing strong support for the necessity and effectiveness of the regularization term.

[Q1.2: The Agent policy optimization is unclear: Why not condition the 2nd in eg 8 term on c?]

Thanks for your insightful comment. By directly constraining the agent's output on $c^*$, the agent will become fully constrained on the 'optimal' action generated by the CVAE. However, this approach may introduce additional computing errors due to the training of the CVAE. Instead, we propose utilizing the critic-computed $c$ and constraining the agent on the precisely sampled action. This not only avoids errors but also allows the agent to be aware of the multiple action distributions of behavior policies within the datasets.

To further verify our statement, we conducted an ablation study on the 'random-medium', 'random-expert' and 'random-medium' dataset by changing the regularization term from the original advantage-aware form $\mathbb E_{\substack{(s,a)\sim \mathcal D, \tilde{z}\sim \pi_\omega(\cdot|c), \\ a_\xi\sim p_\psi(\cdot|\tilde{z}, c)}}\big[(a-a_\xi)^2\big]$ to $\mathbb E_{\substack{(s,a)\sim \mathcal D, \tilde{z}^*\sim \pi_\omega(\cdot|c^*), \\ a_\xi^*\sim p_\psi(\cdot|\tilde{z}^*, c^*),\\ a_\text{cvae}^*\sim p_\psi(\cdot|z, c^*), z\sim \mathcal N(0,I)}}\big[(a_\xi^*-a_\text{cvae}^*)^2\big]$. Therefore, the output of the A2PO optimal policy $a_\xi^*$ is constrained only by the CVAE-inferred best action $a_\text{cvae}^*$.. The results are presented in Table R1， which show that our original form outperforms the variation in most cases, providing support for our regularization term.

Table R1. Test returns of A2PO and A2PO constrined on $a_\text{cvae}^*\sim p_\psi(\cdot|z=\mathcal N(0,I), c^*)$. Bold indicates the best performance among the two algorithms.

[Q2: Missing some return conditioned baselines：1. Comparison with return conditioned approaches such as %BC and Decision Transformers are needed to understand the advantage of the proposed approach; 2.Comparisons with data rebalancing techniques can help understand methods with mixed datasets]

Thanks for your constructive suggestions! We have additionally compared our A2PO with the return condition baselines 10%BC and Decision Transformer [2], and Offline RL with data rebalancing technique ReD and its variation DeReD[3]. The results are presented in Table R2. We directly take the score from the original paper for 'medium', 'medium-replay' and 'medium-expert' datasets. For the new 'random-medium', 'random-expert' and 'random-medium-expert' datasets, due to the time limit, we only evaluate IQL with ReD. The results show that our A2PO outperforms both the return condition baselines and data rebalancing techniques in most cases, which further demonstrates the superiority of our advantage-aware methods to tackle the constrain conflict issue.

[1] PLAS: Latent Action Space for Offline Reinforcement Learning. CoRL 2021

[2] Decision Transformer: Reinforcement Learning via Sequence Modeling. NeurIPS 2021

[3] Boosting Offline Reinforcement Learning via Data Rebalancing. arXiv 2022

Table R1. Test returns of A2PO and other baselines in gym tasks. Bold indicates the best performance among these algorithms.

Table R2. Test returns of A2PO and other baselines in antmaze v0 tasks.

Thank you for affirming the idea's novelty and the experiment's comprehensiveness. In this section, we aim to address your concerns.

[Q1: In line 169, what prior distribution of $p(z)$ did you choose?]

Sorry for the confusion. As stated in line 167 of the original paper, the prior distribution of $p(z)$ is chosen to be $\mathcal N(0,I)$.

[Q2: When training the Q and V functions (Eq(7)), the authors claimed that training simultaneously can stabilize training. Do the authors have any intuition, theory, or literature to back up such a design?]

Thank you for your insightful question. We apologize for the confusion caused by the misdescription in the manuscript. In fact, we did not train the Q and V functions simultaneously; instead, we employed an alternating training approach for policy evaluation following the LAPO method [1]. Specifically, the Q-function is still optimized with the Bellman equation: $L_Q(\theta)=\mathbb E_{(s,a,r,s')\sim \mathcal D}\sum_i [r+\gamma V_{\hat \phi}(s')-Q_{\theta_i}(s,a)]^2$, while the V-function is optimized with the the expectation on Q-value: L_V(\phi)=\mathbb E_{s\sim \mathcal D}[V_\phi(s)-\mathbb E_{ \tilde z^*\sim \pi_\omega (\cdot|c^*), a^*_\xi\sim p_\psi(\cdot|\tilde z^*, c^*) }Q_\hat\theta(s,a_\xi^*)]^2. Corresponding to the implementation presented in the supplementary material, in line 230-238 oflatent_a.py, the V-function update codes are:

adv_sample = torch.ones((next_s.size()[0], 1)).to(self.device)        #get \xi^*
latent_a = self.actor_target(torch.cat((s, adv_sample), dim=1))	      #get \tilde z^*
hat_a = self.vae.decode(torch.cat((s, adv_sample), dim=1), latent_a)  #get a_\xi^*
temp_q1, temp_q2 = self.critic_target(s,hat_a)                        #get Q value
target_v = torch.min(temp_q1, temp_q1) * self.doubleq_min + torch.max(temp_q1, temp_q1) * (1 - self.doubleq_min) 
current_v = self.critic.v(s)
loss = ... + F.mse_loss(current_v, target_v)

[Q3: In Table 1, I found some baseline results inconsistent with the original paper.]

Sorry for the confusion. We have systematically reviewed the score tables for the 'medium', 'medium-replay', and 'medium-expert' datasets, rectifying any recording errors identified. The mismatch scores reported in the original manuscript were obtained by re-running the corresponding source code, which was accidently mixed with the officially reported scores. We give a new comparison table between these baselines and A2PO, as illustrated in Table R1, and we will update the paper accordingly.

[Q4: It is unfair to compare A2PO's v2 results with the v0 results of other baselines.]

Thanks for your constructive comments. We have additionally tested A2PO using the v0 version for the AntMaze tasks. For BC, BCQ, TD3+BC, CQL, AWAC, IQL, and Diffusion-QL, we have directly referenced the scores from the Diffusion-QL paper [1]. For CQL+AW, we have utilized the scores from the CQL+AW paper [2]. For EQL and LAPO—neither of which provide scores for the v0 version—we have re-evaluated their performance, presenting both the scores and standard deviations. The results are summarized in Table R2, indicating that our A2PO continues to achieve performance comparable to that of advanced baselines such as Diffusion-QL.

[Q5: I found that the training curve of hopper random-medium-expert in Figure 4 is unstable. Can the authors provide more training curves for other environments. ]

Thanks for the insightful comment. We have provided various A2PO training curves in the supplementary pdf, in which the A2PO is trained with 1M timesteps while performing online evaluation at each 2K timesteps. The results show that the A2PO maintains overall stability during training on most of the tasks and datasets.

[Q6: My main concern is about the empirical experiments. It seems the algorithm doesn't have competitive performance compared to diffusion-based policy algorithms.]

We appreciate your valuable comment; however, we respectfully disagree with the assertion that A2PO does not exhibit competitive performance compared to diffusion policy for the following reasons:

1. Even taking into account the comparative data from the Diffusion-QL paper, our A2PO method still outperforms diffusion-QL in 11 out of 18 gym tasks, 4 out of 6 maze tasks, 2 out of 3 kitchen tasks, and 1 out of 2 adroit tasks. This demonstrates that our algorithm exhibits a comparative advantage in most task scenarios over the diffusion policy.
2. The primary focus of our research is not on benchmarking against diffusion-based methods but on addressing the conflict constraint issue, which current AW methods fail to solve effectively. Our extensive experiments indicate that A2PO significantly outperforms existing AW methods, such as LAPO.
3. Our approach is orthogonal to diffusion-based methods. Specifically, Diffusion-QL utilizes diffusion as the policy to implicitly capture the multimodal action distribution, while it overlooks the intrinsic quality of different actions. In contrast, our A2PO explicitly leverages the advantage value to disentangle the action distributions of interrelated behavior policies, ensuring that the quality of different actions is positively correlated with the advantage value. Therefore, our advantage-aware mechanism can be effectively integrated with diffusion-based methods, enhancing their capacity to identify multimodal data of varying quality.

In conclusion, our core contribution lies not in the model itself but in the introduction of an advantage-aware policy.

[1] Latent-Variable Advantage-Weighted Policy Optimization for Offline RL. NeurIPS 2022

[2] Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning. ICLR 2023

[3] Harnessing Mixed Offline Reinforcement Learning Datasets via Trajectory Weighting. ICLR 2023

We sincerely appreciate your positive support of our work. In the future, we will make further improvements to the A2PO to contribute to the offline rl community.

Thanks for your supportive review of our method, experiments, and writing skills. We will address your doubts here.

[Q1. since we can directly select the good trajectories according to the rewards, is there any ablation study to show that CVAE can do better than this simple method?]

Thanks for your constructive suggestions! We have additionally conducted ablation study on our CVAE in the A2PO method. To illustrate the effectiveness of our CVAE approach compared to merely selecting high-quality trajectories for training, we eliminate the CVAE module and implement the trajectory selection method proposed in [1,2], which involves selecting only the top 10% of trajectories based on their returns for agent training.

The results shown in Table R1 clearly indicate that A2PO without CVAE [top-10%] consistently lags behind the original A2PO in most cases, especially on the mixed-qualities datasets like 'walker2d-random-medium' and 'walker2d-random-expert'. This is mainly due to the challenge of determining the hyperparameter $k$ for selecting the top-$k$ trajectories for training across different tasks and datasets. A smaller $k$ filters out more data, limiting the agent's awareness of the action distribution and task information. On the other hand, a larger $k$ introduces constraint conflicts that hinder agent policy improvement. Instead, our A2PO utilizes CVAE conditioned on advantage value to obtain disentangled behavior constraints. By optimizing the advantage-aware policy based on these disentangled action distribution constraints, we are able to achieve higher returns by directing the agent toward high advantage values.

Table R1. Test returns of A2PO and A2PO w/o CVAE [selecting top-10% trajectory for training]. Bold indicates the best performance among the two algorithms.

[1] Decision Transformer: Reinforcement Learning via Sequence Modeling. NeurIPS 2021.

[2] Harnessing Mixed Offline Reinforcement Learning Datasets via Trajectory Weighting. ICLR 2023.