Paper: Effective Offline Environment Reconstruction when the Dataset is Collected from Diversified Behavior Policies (ICLR-2024)

Reviewer comment (Novelty):
The Related Work does not seem comprehensive enough to me. It should include more recent works on OPE and MBORL. I did not find citation of OPE later than 2020. For MBORL, COMBO, ROMI and many more later works are missing. I suggest adding a more comprehensive Related Work.

Author response:
### **Q1: The Related Work does not seem comprehensive enough to me. It should include more recent works later than 2020.**

Thank you for the suggestion. We have included additional related work in the main text, covering topics such as COMBO, ROMI, and others later than 2020.

1. **Inclusion of COMBO, ROMI, etc. in Related Work**: We apologize for the oversight in our previous submission where the updated related work section was mistakenly omitted. We have rectified this error and submitted a new version of the manuscript. The updated related work now includes a comprehensive discussion on models such as COMBO, ROMI, and others, reflecting a broader spectrum of popular MBORL approaches.

Reviewer comment (Experiments):
As an experiment paper, I do not think the experiments conducted are extensive enough, more tasks should be tested.

Author response:
### **Q3: I do not think the experiments conducted are extensive enough, more tasks should be tested.**

Thanks for your valuable suggestion. We have supplemented five datasets for OPE task, including ant-medium-replay, ant-medium-expert, halfcheetah-medium-expert, hopper-medium-expert, and walker2d-medium-expert. Similarly, our method also outperforms other methods significantly on these tasks.  Please see Table 3 to Table 6 for raw results in our revised paper and full discussion in the general response.

2. **Comprehensive Evaluation Concerns**: We acknowledge your concern regarding the comprehensiveness of our evaluation. We want to clarify that the 'random' dataset does not align with the assumptions of our approach, as it is generated from a single random policy rather than the multi-source policies our method is designed for. After adding the results in medium-expert datasets (thanks to the reviewer BstF's suggestion), we believe that our experiments **have covered all tasks in the D4RL dataset that involve multi-source policy collections**. We believe these additions make our experimental evaluation more robust and comprehensive.

We are committed to conducting further investigations to refine our method and extend its applicability. We hope that the revisions and clarifications provided address your concerns and demonstrate the rigor of our approach. Should you have any additional questions or require further clarification, we are readily available to engage in further discussions.

Reviewer comment (Presentation):
Figure 2 and 3 seem to have spacing issues. The legends are partially screened by the captions.

Author response:
### **Q7: Figure 2 and 3 seem to have spacing issues.**

Thanks for pointing out this formatting issue. We have fixed them in our revised paper.

Reviewer comment (Presentation):
Can the authors explain Figure 2(c) in more detail, how are the quantities calculated?

Author response:
### **Q4: Can the authors explain Figure 2(c) in more detail, how are the quantities calculated?**

In our experiment, we constructed a set of 10 distinct policies. These policies were then used to generate datasets on which our PCM was trained, resulting in 10 unique vector policy embeddings corresponding to each policy. We evaluated the performance of each policy using models conditioned on each of the 10 embeddings. The value gap was calculated for each pairing, providing a matrix of outcomes.

Specifically, in Figure 2(c), each cell in row i, column j, quantifies the value gap when policy j is evaluated using the PCM adapted with embedding i. The heatmap's intensity indicates the magnitude of the value gap, with darker colors representing larger gaps. This visual representation allows us to observe how well the model, when conditioned with a specific policy embedding, can adapt and evaluate other policies.

Reviewer comment (Presentation):
In Figure 3, what does the data for percentage 0% and 100% look like. Also, the differences of value gaps does not seem substantial to me, is it the case?

Author response:
### **Q5: In Figure 3, what does the data for percentage 0% and 100% look like. Also, the differences of value gaps does not seem substantial to me, is it the case?**

We would like to directly compare the value gaps of PAM and PCM and also investigate the influence of different levels of dataset diversity on them. To do so, we construct datasets with varying levels of diversity (0\%, 20\%, 50\%, 80\%, 100\%), where the percentage indicates that the dataset is created from the replay buffer of SAC until the policy reaches the specific level of performance.

Reviewer comment (Novelty):
The ideas of assuming mixed-type of data and leveraging evaluation policies to improve state-action visitation are not very interesting. Similar assumptions and problems were investigated by recent OPE work [1], which also conducted experiments on D4RL. The major claim, “model learning from offline datasets collected by numerous behavior policies in fact implies a problem of data fitting from a mixture of multi-source data distributions, which is ignored in current model learning paradigms (in Introduction)” looks overstrong, which should be carefully scoped, given there are already works investigating such scenarios (e.g., [1-2]).

Author response:
Thank you for your insightful comments and the referenced articles. Upon careful consideration, we maintain the validity of our claim and provide the following clarifications:

Regarding the work in [1], while it shares similarities with our approach in encoding historical data, the purposes differ significantly. [1] focuses on “learning a compact and disentangled latent representation space from offline trajectories, which can better capture the dynamics underlying environments”. In contrast, our work is centered on obtaining representations for policy representation. This divergence in objectives leads to distinct training methodologies: [1] employs a variational inference framework to construct an ELBO for a robust environmental representation space, whereas we use the representation as an extra input for model prediction, aiming solely at reconstructing the policy to encode its representation. As we clarified to ****reviewer Ag2q****, our contribution does not lie in the specific technical implementation of using RNNs to encode trajectory features. To our knowledge, this particular application of policy-conditioned methods to enhance policy learning in the context of diverse offline RL datasets is a novel contribution. It surpasses existing methods by focusing on aligning policies with the prediction emphasis of dynamic models, deviating from the primarily trajectory-encoding approach in prior research.

Regarding [2], the discussion revolves around the diversity of dataset sources in a model-free offline policy learning setting. In this context, the issue of policy being multi-source is intuitively obvious, as the direct goal is to imitate and optimize the offline data's policy. However, in a model-based scenario, the hidden multi-source data distribution caused by multiple source policies and the potential to leverage this feature for enhanced model generalization are not as readily apparent.

In summary, our assertion that model learning from offline datasets collected by numerous behavior policies implies a problem of data fitting from a mixture of multi-source data distributions remains unaddressed in current literature. Therefore, our claim is not overstated but rather highlights a novel aspect that has not been formally proposed until now.

Reviewer comment (Experiments):
For off-policy selection, some baselines (e.g., IS) are not included. They are still able to be compared, by using the policy with the maximum OPE values. For off-policy evaluation, it might also be helpful to evaluate if recent works with similar ideas can be considered, as mentioned in Point 1.

Author response:
Thanks for the valuable suggestion. In previous version, we selected FQE as the baseline for comparison since it the best baseline in OPE tasks. Following the suggestion, we have additionally evaluated IS and DICE for off-policy selection, which is shown in Page 23 Table 8-9 in the appendix, and Table 2 in the main body. We observe that IS indeed has more competitive performance (59.4% performance gain) than FQE (20.6% performance gain) but still falls behind PCM (71.7% performance gain).

Reviewer comment (Presentation):
Figure 3: Part of labels on x axis is covered by captions.

Author response:
Thanks for pointing out this formatting issue. We have fixed them in our revised paper.

Reviewer comment (Experiments):
The universality of the dataset collected from diversified behavior policies has not been comprehensively elaborated which limits the meaning of this problem.

Author response:
Thank you for your valuable feedback regarding the universality of the dataset derived from diversified behavior policies in our research. We acknowledge the importance of this aspect and have addressed it more comprehensively in our revised manuscript.

In practical scenarios, the multi-source nature of offline datasets is a common phenomenon, particularly in fields such as robotic manipulation, autonomous driving, and sequential recommendation. These domains often involve data collected from a diverse array of policies. These include parameterized policies, rule-based policies, and human behavior policies, each contributing unique and realistic dimensions to the dataset.

In our revised paper, we have expanded the introduction to include a detailed discussion of this aspect, illustrating the relevance and applicability of our approach to real-world scenarios. To support our argument, we have referenced several key studies:

1. "Roboturk: A crowdsourcing platform for robotic skill learning through imitation" and "What matters in learning from offline human demonstrations for robot manipulation" [1,2] highlight the significance of varied data sources in robotic manipulation.
2. "BDD100K: A diverse driving video database with scalable annotation tooling" and "Scalability in perception for autonomous driving: Waymo open dataset" [3,4] underscore the diversity in autonomous driving datasets.
3. "Open bandit dataset and pipeline: Towards realistic and reproducible off-policy evaluation" and "Kuairec: A fully-observed dataset and insights for evaluating recommendation systems" [5,6] elucidate the variety in datasets used for sequential recommendation.

These references substantiate the universality and practical importance of multi-source datasets in offline RL research, aligning closely with the objectives and scope of our study.

We believe that this enhanced discussion in our paper now adequately addresses the concern raised and underscores the significance and applicability of our research in diverse real-world contexts.

[1] Roboturk: A crowdsourcing platform for robotic skill learning through imitation.

[2] What matters in learning from offline human demonstrations for robot manipulation.

[3] BDD100K: A diverse driving video database with scalable annotation tooling.

[4] Scalability in perception for autonomous driving: Waymo open dataset.

[5] Open bandit dataset and pipeline: Towards realistic and reproducible off-policy evaluation.

[6] Kuairec: A fully-observed dataset and insights for evaluating recommendation systems.

Reviewer comment (Writing):
The relationship between the proposed solution and the challenges in the setting of dataset collected from diversified behavior policies is not clarified well. It is still unknown or at least confusing the difference between the setting of the dataset collected from homogeneous behavior policies and the dataset collected from diversified behavior policies. And what are the additional challenges brought by this new setting and how the proposed approach help solve these challenges?

Author response:
Thank you for the valuable comments. In our revised manuscript, we have expanded our discussion in this regard, particularly in the methodology section, to address these concerns more comprehensively.

The key distinction in our research lies in the comparison between datasets collected from homogeneous behavior policies and those from diversified behavior policies. Traditional approaches often rely on the simplification of using data from a single-behavior policy, which does not fully capture the complexity and variability found in real-world applications. In contrast, our study focuses on the multi-source nature of data, which is more representative of realistic scenarios, as explained in our previous response.

**The unique aspect of our approach is not about addressing new challenges presented by diversified datasets but rather about leveraging the additional information inherent in such settings.** In a multi-source dataset, different trajectories may originate from varying policies, but the transition data within each trajectory are consistently generated by the same policy. This distinct feature provides valuable insights for dynamics model learning, which is a central aspect of our proposed method.

By acknowledging and utilizing this additional information, our approach goes beyond the limitations of  algorithms designed for single-source datasets. It recognizes and exploits the nuances in multi-source data that are often overlooked in more simplified formulations. We believe this refined perspective, now thoroughly discussed in our paper, underscores the novelty and relevance of our research in addressing the complexities of real-world data in AI applications.

Reviewer comment (Experiments):
Though the authors have paid much attention to how to adapt the learned environment model for various policies in different downstream tasks, how to overcome the challenges of the diverse upstream dataset is less covered in the paper.

Author response:
Thank you for your valuable observation regarding our paper's focus on adapting learned environmental models to various policies in different downstream tasks.

We would like to clarify that our research fundamentally revolves around leveraging the diversity of datasets to learn a more accurate and faithful environmental model. The diverse nature of the dataset is not just a condition for our method but rather the cornerstone upon which our approach is built. Hence, our primary aim is to demonstrate how to effectively utilize the inherent diversity in datasets to enhance environmental model learning.

Given this focus, the challenges posed by the diverse nature of the dataset are integral to our research. However, the explicit discussion on overcoming these challenges may appear less prominent as our method inherently embraces and capitalizes on this diversity, rather than viewing it as a hurdle to be overcome.

We appreciate your input and have made efforts in our revised manuscript to articulate this aspect more clearly, ensuring that our approach's alignment with the premise of diverse datasets is well understood.

Reviewer comment (Experiments):
The authors are expected to provide more evidence and supporting materials on why choosing OPE, OPS, and MPC as the downstream tasks to validate the effectiveness of the proposed method.

Author response:
These tasks correspond to the discussion in related works. OPE is the task most relevant to the “value gap”. Since OPE needs to evaluate the value of a target policy and therefore we can expect that a more accurate model that has a smaller value gap should do better in OPE. OPS is a task derived from offline RL. For the offline RL setting, it tends to be not allowed to select a desirable policy from many candidates via evaluating them in the real environment, thus OPS is a key process for offline RL, which is also be attain great importance in other work [7]. MPC is a task to consider whether our model can be applied to decision/control rather than only policy evaluation. We don’t consider policy learning in our paper since policy learning with a learned model in the offline setting always needs to incorporate some additional designs, like ensemble, uncertainty quantification, and penalty  [8, 9], which introduce more variables and make it difficult to evaluate whether the improvement of policy learning is due to the more accurate model.

[7] NeoRL: A near real-world benchmark for offline reinforcement learning.

[8] MOPO: model-based offline policy optimization.

[9] Model-Bellman inconsistency for model-based offline reinforcement learning.

Reviewer comment (Novelty):
**Limited technical contributions**. The central idea of the paper is training policy-conditioned dynamics models instead of a single model on a diverse offline RL dataset to achieve better modeling error. However, the idea of using MDP-conditioned embeddings for training a universal model in deep RL is not new, see for example [9][10][11].

Author response:
While we appreciate the references to existing literature on MDP-conditioned methods in deep RL, we believe that the technical contributions of this paper is beyond the method implementation itself. The paper presents a novel perspective on leveraging MDP-conditioned dynamics models, particularly in the context of offline RL datasets. 

We would like to first highlight the distinct aspects of our approach. 

- The referenced works, such as Universal Successor Features Approximators [11], FOCAL [9], and Policy-Dynamics Value Functions [10], indeed employ MDP-conditioned methods. However, these approaches primarily focus on adapting to diverse environmental conditions and reward structures through modifications prior to policy application.
- In contrast, our work diverges significantly in both intent and implementation. We utilize policy-conditioned methods not as a means to adapt to environmental variability, but rather to enable different policies to learn more effectively from a dynamic model tailored to the dataset's diversity. This approach specifically targets the reduction of value gaps, a critical aspect in offline RL that has not been the primary focus of the aforementioned studies.

To the best of our knowledge, this particular application of policy-conditioned methods to enhance policy learning in the context of diverse offline RL datasets represents a novel contribution. It extends beyond the scope of existing methods by prioritizing the alignment between policies and the prediction focus of dynamic models, which is a departure from the primarily environment-focused adaptations in prior research.

We believe this clarification underscores the unique technical contributions of our work and its significance in advancing the field of deep reinforcement learning.

Reviewer comment (Theory):
**Gap between theory and experiments**: The main theorem of the paper, Proposition 4.2, is considered a relatively trivial result. On the RHS of (8), the training error minus the adaptation gain is precisely the model error on the LHS. It only shows that the generalization error can be bounded given a Lipshitz assumption, without explicitly modeling the effect of policy embedding. The main argument seems to be that the adaptation gain of PCM is positive but zero for PAM. However, one can simply fine-tune PAM given the data of a new policy $\pi$ to achieve a positive adaptation gain, which is also a stronger baseline that should be compared in experiments in 5.3.1 (see point 3). Therefore, it is not convincing enough why policy-conditioned PCM is superior given the theory, which is supposed to be the main contribution of the paper.

Author response:
The analysis is in fact an ideal case where the policy information could be perfectly injected into the model, under which case the generalization to other policies naturally relies on the Lipschitz continuity of the learned model. 

The policy embedding is a practical method to realize such the information injection mechanism, however, it is hard in general to analyze the effect of the concrete embedding, and we resort to the adaptation gain argument, which is an overall description of the model adaptation effect. In the previous submission, we discussed three possible mode of the adaptation in Appendix A.3, where the Case 2 Imaginary retraining is an ideal model of fine-tuning PAM for the target policy. However, it is just “imaginary” since it requires the real target domain data, which is inaccessible in our setting as well as many realistic situations, and we regard it as a ceiling of model accuracy for any practical PCM algorithms. If we cannot obtain the target policy data, the PAM can never be tuning specially for the target policy, and the adaptation gain is exactly zero. We believe the case 3 Intermediate adaptation reflects the adaptation effect of practical PCM methods, which lies between the two extreme cases of perfect adaptation ( case 2. imaginary retraining) and zero adaptation gain (case 1. direct match) and is able to partially reduce the extrapolation error. The empirical results in Sec 5.2 provides evidence to support this discussion.

Reviewer comment (Experiments):
**Missing important model-based baselines**. As mentioned in 5.3.1, naive PAM significantly outperforms other competitive baselines of off-policy evaluation because the authors use better architectures than MLP. While it is appreciated that the authors are honest and transparent about this point, since the key contribution PCM is policy-conditioned design of dynamics model, it is important to see whether PCM outperforms other competitive model-based OPE baselines when using the same architecture. A simple fine-tuned PAM on an unseen task should also be compared as a stronger baseline (see point 2).

Author response:
Thank you for your valuable feedback regarding the comparison of our PCM with other model-based and model-free OPE methods. 

Firstly, it is important to note that among the baselines considered in our study, PAM is the only model-based OPE method. The remaining baselines are model-free OPE methods. This distinction is crucial as it ensures a fair and relevant comparison within the context of model-based OPE approaches.

Regarding the suggestion to compare our approach with a fine-tuned PAM on unseen tasks, we argue that a fine-tuned PAM should not be considered an additional baseline. This is due to its reliance on real interaction trajectories of the target policies, which diverges from the primary objective of OPE methods - to evaluate policies without necessitating further real-world interactions.

Furthermore, to bolster our comparison and validate the efficacy of PCM, we have also implemented PCM with a standard Multilayer Perceptron (MLP) architecture. This implementation demonstrates that the PCM mechanism consistently aids in reducing value gaps across different architectures. The results of this implementation can be found in the general response.

Reviewer comment (Novelty):
The review of related work of model-based offline RL is inadequate. In this submission, the related work only refers to those published before 2021 and there are more recently published works such as PMDB[1], VMG[2], RAMBO[3], COUNT-MORL[4], ROSMO[5], CABI[6], AMPL[7], and CBOP[8]. The author’s categorization of MBORL into two groups—MPC and PL—is disagreed with. For example, AMPL[7] focuses on the model learning process but cannot be categorized into policy learning; CABI[6] focuses on model data augmentation; RAMBO[3] focuses on the mini-max optimization about the offline dataset. Furthermore, MPC methods cover only a relatively small portion of model-based offline RL methods. Additionally, the setting of the paper seems more related to the so-called compositional/functional RL [12], which should be discussed and some of the baselines should be compared.

Author response:
Thank you for your comments regarding the review of related work in our paper. Our approach to categorizing related work is based on how the models are utilized within the offlineRL framework.

Specifically, we distinguish between methods that use models for model predictive control (MPC), where policies are generated through planning, and those used for policy learning (PL). This classification is designed to evaluate the effectiveness of models in facilitating downstream tasks, which is central to our study.

Considering the recent works you mentioned, we acknowledge that they indeed present novel approaches within the model-based offline RL domain. For instance, AMPL [7] and RAMBO [3] integrate dynamic-model training with policy learning, and CABI [6] focuses on generating new datasets for offline RL policy learning in an adversarial manner. According to our classification principle, these methods primarily support policy learning as they leverage dynamics models for this purpose.

We have chosen this classification scheme because it aligns closely with the objectives of our research and informs the design of our experiments. While we recognize other valid methods of classification, we believe our approach is more pertinent to the specific issues and questions our paper addresses. To this end, we have updated our manuscript to include additional related works in the main text, thereby providing a more comprehensive overview.

Regarding your point about compositional/functional RL, we respectfully disagree with the applicability of this concept to our work. Compositional RL typically assumes that tasks can be broken down into subtasks, which is not an assumption we make in our approach. Our method focuses on learning representations about tasks without decomposing them into subfunctions, distinguishing it from compositional methods.

We hope this response adequately addresses your concerns and provides clarity on our classification rationale and its relevance to our research objectives.

Reviewer comment (Writing):
Minor issue — Lack of clarity on notation: What is $W$ in Proposition 4.2?

Author response:
We apologize for the omitted explanation for it. Here, $W_1$ is the Wasserstein-1 distance.

Reviewer comment (Presentation):
Minor issue — Inconsistency in expected outcome illustration: In Fig 2.c, the expected outcome should look like a diagonal matrix since evaluating on the same policy where the embedding is trained (diagonal entry) should give the best performance.

Author response:
The ideal result is that it appears as a diagonal matrix. However, due to some randomness in experiments, it’s hard to get an exact, “perfect” diagonal matrix. The result shown in Fig 2.c is close to a diagonal matrix, in which we can see that the model error near the diagonal is small while the error away from the diagonal is large.