Collaborative World Models: An Online-Offline Transfer RL Approach

Q5: In Figure 5(a), why CoWorld performance is worse or similar compared with CoWorld w/o State A/B/C?

At approximately 150k iterations, the convergence of the world model and policy in CoWorld is not complete, resulting in unstable performance. We observed a significant decrease in reward variance and evaluation returns around this point, suggesting a potential overfitting process within the reward predictor. However, by continuing the training of CoWorld, we successfully mitigated this overfitting issue. Additionally, around 170k iterations in CoWorld without Stage B, a similar decrease was noted, although without a subsequent increase in reward variance. This discrepancy might have contributed to the final result decline in CoWorld without Stage B.

Q6: What is the computational complexity of the training phase of the auxiliary agent?

We have provided the comparisons of the training time and the inference time of the compared models in the above General Response (Q4). Please note that our approach yields comparable inference time to the baseline models such as LOMPO.

Q7: In Figure 5(b), the CQL (Kumar et al.) method seems to overestimate the value function for all the tasks, why is this the case considering that CQL is a conservative way for offline RL?

As there is no official code for CQL with visual inputs, we adopted the DrQ-v2+CQL implementation from the work of V-D4RL (Lu et al., 2023). Below, we compare the value estimated by DrQ-v2+CQL and its base DrQ-v2 model:

We have used the above results to update Figure 5(b). In comparison to the naive DrQ-v2, using the CQL regularization term can partly mitigate the issue of overestimated values. However, there are still overestimates due to the complexity of offline visual RL problems.

Reference: Lu et al. "Challenges and opportunities in offline reinforcement learning from visual observations", TMLR, 2023.

Thank you for your valuable comments. We hope our responses below can help address your concerns on this paper.

Q1: How to choose a source RL simulator？ How do you set the criteria for choosing the source task and how do you ensure the source task is sufficiently informative for the target tasks?

We understand the reviewer's concern regarding the impact of the source task quality on the learning performance. We would like to address this question from two perspectives:

(1) CoWorld's performance with a random source domain:

In the revised paper, we have updated Figure 3(a) to include the Transfer Matrix among the 6 tasks on the Meta-World benchmark. Within all domain transfer cases, there are instances of pairs with less related source and target tasks. Despite these challenging scenarios, our observations indicate that CoWorld outperforms Offline-DV2 in the majority of cases (26 out of 30).

(2) Results with the REVISED "source domain selection" method:

Admittedly, our initial proposal of the "source domain selection" method is not practical in real-world scenarios, as $R_\text{online}$ should not be available during training in the offline RL setup. To address this issue, we have thoroughly revised this method and introduced a new adaptive domain selection approach for scenarios with multiple source domains available.

For detailed technical insights and corresponding results, please refer to our General Response (Q2). In summary, the "adaptive source domain selection" method yields results comparable to models trained with manually designated online simulators. This highlights our approach's ability to automatically identify a useful online simulator from a set of both related and less related source domains, thereby expanding its applicability in practical scenarios.

We have included the aforementioned results in the revised paper.

Q2: Some details of the proposed method are missing or simply punted to another paper.

Thank you for your suggestion. We have incorporated additional details about the proposed method for world model learning and behavior learning in Appendix A. This includes further explanations regarding the insight behind the objective functions and a more in-depth description of the network architecture.

Q3: How many steps for fine-tuning for each task?

Below, we present the finetuning steps of the "Finetune" baseline model compared with the training steps of CoWorld in the target dataset. Please refer to Appendix B for more details of the "Finetune" method.

Q4: What is the variance in Figure 3(b)?

We trained the compared models with three random seeds and presented the mean results along with standard deviations. The experimental results in Figure 3(b) have been updated in the revised paper.

We appreciate your great efforts in reviewing our paper and hope that the following responses can address most of your concerns.

Q1: It is not clear how the source domain should be selected in practice, as the source domain selection is actually leaking information from test to train.

Yes, indeed, our initial proposal of the "source domain selection" method is not practical in real-world scenarios. We greatly appreciate that the reviewer pointed it out. To address this problem, we have thoroughly modified the domain selection method, which can extend CoWorld to scenarios with multiple source domains available. This is achieved by measuring the distance of the latent states between the OFFLINE target dataset and each source domain provided by different world models. We have included more technical details and corresponding results in General Response (Q2).

In summary, the "adaptive source domain selection" method yields results comparable to models trained with manually designated online simulators. This highlights our approach's ability to automatically identify a useful online simulator from a set of both related and less related source domains, thereby expanding its applicability in practical scenarios.

We have included the aforementioned results in the revised paper.

Q2: Can you evaluate the performance of the algorithm without the designed "source domain selection"? What if the source domain does not match the requirement? Will it hurt the performance by a lot?

(1) CoWorld's performance with a random source domain:

In the revised paper, we have updated Figure 3(a) to include the Transfer Matrix among the 6 tasks on the Meta-World benchmark. Within all domain transfer cases, there are instances of pairs with less related source and target tasks. Despite these challenging scenarios, our observations indicate that CoWorld outperforms Offline-DV2 in the majority of cases (26 out of 30).

(2) Results with the REVISED "source domain selection" method:

As discussed above, we have revised the source domain selection method. The updated method can easily extend CoWorld to scenarios with both related and less related source domains by automatically identifying a useful domain as the auxiliary simulator. We observe that:

• In Table 1 in the revised manuscipt, the multi-source CoWorld achieves comparable results to the models trained with manually designated online simulators.
• In Figure 3(a), the multi-source CoWorld achieves positive improvements over Offline-DV2 in all cases, approaching the best results of models using each individual source task as the auxiliary domain.
• In Figure 3(b), the multi-source CoWorld consistently outperforms the Finetune baseline, even when the single-source CoWorld faces challenges with specific undesirable source domains.

These results demonstrate our approach's ability to operate without strict assumptions about domain similarity.

Q3: Time complexity of the method compared to baselines.

Please refer to the comparisons of the training/inference time complexity in our General Response (Q4).

Q4: From the formulation and the algorithm, it is not clear how the source and the target domains are related.

We propose the following two techniques to establish a relationship between the source and target MDPs, mitigating cross-domain discrepancies:

• Latent Space Alignment: To mitigate the cross-domain discrepancies in visual observations. This is achieved by closing the distance of latent states produced by the world models with the same visual inputs.
• Target-Inclined Source Model Tuning: To mitigate the cross-domain discrepancies in reward functions. This is achieved by incorporating the target reward information into the learning process of the source reward predictor. As we assume the latent states in different domains have been successfully aligned, the objective in this stage is to enable the source reward predictor to estimate the target returns based on the aligned state inputs.

Q5: For latent state alignment, it is confusing to me why you can feed the same target domain observations into the two world models and close the distance of latent states.

In "model transfer" setups, as suggested by Knowledge Flow (Liu et al., 2019) and Transferrable Memory Unit (Yao et al., 2020), it is a common practice to pretrain a teacher model in the source domain, and subsequently transfer the knowledge from the teacher model to the student model in the target domain by providing the teacher model with target domain data. These existing approaches typically assume the presence of distribution shifts between $O^{(S)}$ and $O^{(T)}$.

References:

• Liu et al. "Knowledge flow: Improve upon your teachers", ICLR, 2019.
• Yao et al. "Unsupervised transfer learning for spatiotemporal predictive networks", ICML, 2020.

We greatly appreciate your valuable comments. We hope the responses below can help address your concerns.

Q1: In Figure 4, what are the source and target domains for (a) and (b)?

Figure 4 provides cross-environment results, where the source domains (a: Press Button, b: Close Window) are from the Meta-World environment, while the target domains (a: Push Button, b: Open Slide) are from another environment named Robodesk.

The data distribution shifts between the source domain and the target domain reside in visual observation, physical dynamics, reward definition, and even the action space of the robots. Let's take the experimental setup in Figure 4(a) for example:

Q2: Could you elaborate on how dependent the performance of CoWorld is on the quality of this simulator? How would the method perform if the simulator is not a perfect representation of the real-world task?

We understand the reviewer's concern regarding the dependencies between the source domain and the target domain. We would like to address this question from two perspectives:

First, it is important to emphasize that the cross-environment experiments (Meta-World to RoboDesk) we clarified in the above response are INDEED the scenarios that "the simulator is not a perfect representation of the real-world task". Our experimental results in Figure 4 reveal that our approach can effectively mitigate the cross-domain discrepancies in visual observation, physical dynamics, reward definition, or even the action space of the robots. In methodology, this is achieved by 1) Latent Space Alignment and 2) Target-Inclined Source Model Tuning.

Second, in Figure 3(a) in the revised paper, we have updated the Transfer Matrix among the 6 tasks on the Meta-World benchmark. Among all domain transfer cases, there exist pairs of seemingly unrelated source and target tasks. However, we observe that our approach outperforms Offline-DV2 in the majority of scenarios (26/30 cases).

Q3: Unrealistic assumptions --- the success of this online-to-offline paradigm highly hinges on the similarity between the source and target domains... It's unlikely that we can build a high-quality simulator on a very similar source domain.

As discussed earlier, CoWorld does not rely on strict assumptions about domain similarity. The results presented in Figure 3 and Figure 4 illustrate CoWorld's effectiveness in mitigating cross-domain discrepancies and utilizing a less related online simulator to improve offline RL.

Furthermore, CoWorld can be easily extended to online-to-offline scenarios with multiple source domains. As detailed in our General Response (Q2), CoWorld achieves comparable results to the models trained with manually designated online simulators. This demonstrates our approach's ability to automatically identify a useful online simulator from a set of both related and less related source domains, thereby expanding its applicability in practical scenarios.

Thank you for the insightful comments.

Q1: It's not clear is CoWorld is a method or a new problem statement. I think it would be more straightforward and fair to separate these contributions.

Following the reviewer's suggestion, we have refined our introduction to elaborate on the paper's contributions in two key aspects:

• First, we propose a novel online-to-offline transfer RL problem, which aims to improve offline visual RL by leveraging an online simulator as the auxiliary domain.
• Second, we present CoWorld, a new model-based RL approach tailored for the online-to-offline setup. CoWorld effectively transfers domain-sharing knowledge by addressing cross-domain discrepancies. This is achieved through:
  • Latent Space Alignment: To mitigate the domain shifts in visual observations;
  • Target-Inclined Source Model Tuning: To mitigate the domain shifts in rewards;
  • Min-Max Value Regularization: To enable mildly-conservative value estimation.

Q2: Only one baseline is fine-tuned from the source online simulator to the offline data.

First, we have compared CoWorld with more baseline models that are pretained with source domain data and then finetuned in the target dataset. Please refer to our General Response (Q3) for the results. We found that all of the baseline models suffer from the so-called "negative transfer" effect when incorporating a source pretraining stage. This indicates that a naive transfer learning scheme might degenerate the target performance by introducing unexpected bias.

Furthermore, we implemented another baseline model by combining the Finetuning method with elastic weight consolidation (EWC) (Kirkpatrick et al., 2017). EWC allows the model for preserving source domain knowledge. However, we found that, without additional model designs for target domain adaptation, maintaining knowledge from the source domain could affect the performance in the target domain. Detailed quantitative comparisons across different Meta-World tasks can be found in the General Response (Q3) and Table 1 in the revised paper.

For the suggested method, IQL, we did not find a well-performed off-the-shelf implementation that allows for visual inputs. We tried the official code for the work from Cho et al. (2022), which includes an implementation of image-based IQL, but found that it did not converge well. Besides, we should note that IQL was originally designed for the offline-to-online problem, which is inherently different from our problem setup.

References:

• Kirkpatrick et al. "Overcoming catastrophic forgetting in neural networks", Proceedings of the national academy of sciences, 2017.
• Cho et al. "S2P: State-conditioned image synthesis for data augmentation in offline reinforcement learning", NeurIPS, 2022.

Q3: Are the DrQ+BC, CQL, CURL, and LOMPO baselines all trained fully offline?

Yes, they were. However, as discussed above, we have provided new results by pretraining these models in the source domain. Please refer to the quantitative comparisons in our General Response (Q3).

Thank you for your effort in reviewing our paper. Below, you'll find our responses addressing your specific concerns.

Q1: The CoWorld framework's performance is critically dependent on the availability of an online simulator in a domain that is sufficiently similar to the target domain.

Indeed, it is beneficial to leverage an online simulator as the source domain that is very similar to the target domain. However, it is essential to emphasize that (1) our approach does not assume a strict requirement for domain similarity, and (2) it is not that difficult to find a useful auxiliary simulator in practical applications.

(1) Can CoWorld benefit from a less similar source domain?

Our experimental results in Figure 4 demonstrate CoWorld's ability to mitigate cross-environment discrepancies and benefit from a weakly related source domain. As presented in our General Response (Q1), notable differences between the source (Meta-World) and target environments (RoboDesk) exist in visual observation ($\mathcal O$), physical dynamics ($\mathcal T$), reward definition ($\mathcal R$), and even the action space of robots ($\mathcal A$). CoWorld outperforms the second-best approaches by approximately 14%-18% in episode returns. In methodology, this is achieved by 1) Latent Space Alignment and 2) Target-Inclined Source Model Tuning.

Furthermore, in Figure 3(a) in the revised paper, we have updated the Transfer Matrix among the 6 tasks on the Meta-World benchmark. There exist pairs of seemingly unrelated source and target tasks. Despite the challenging domain transfer setups, we observe that our approach outperforms Offline-DV2 in the majority of scenarios (26/30 cases).

(2) How to choose the auxiliary online simulator given the offline target domain?

Our results demonstrate the ability of CoWorld to mitigate the cross-domain discrepancies and leverage a less related online simulator to improve offline RL. This enhances the convenience of selecting an auxiliary simulator based on the type of robot (major) and task similarity (minor). For example, the Meta-World/RoboDesk/RLBench simulators can all be directly applied to the classic offline RL dataset, D4RL FrankaKitchen. Please refer to our General Response (Q1) for further details.

Q2: The rationale behind the decision to maintain two separate world models is not sufficiently explained or supported by evidence.

(1) Rationale of separate world models:

In light of non-identical observation spaces, diverse physics, varying reward scales, and even disparate action spaces between source and target MDPs (e.g., Meta-World and RoboDesk environments using different robotic arms with distinct action dimensions), we train separate world models for different domains. Our approach involves aligning the learned world models in latent state space and reward prediction. This strategy captures both domain-specific and shared information in the online-to-offline setup.

(2) Empirical evidence for using separate world models vs. One world model:

To investigate the effectiveness of employing separate world models, we conduct an experimental comparison between CoWorld and "Multi-Task DV2". The latter involves training DreamerV2 on both offline and online data with a joint world model and separate actor-critic models. Here, we present the mean returns and standard deviations for the cross-environment experiments (Meta-World to RoboDesk) that are calculated over 10 episodes with 3 random seeds. CoWorld consistently performs better.

Q2: (1) Can CoWorld benefit from less related source tasks? (2) Can CoWorld deal with multi-source scenarios?

(1) Clarification of CoWorld's performance with various source domains:

We have updated the Transfer Matrix among 6 different tasks in Figure 3(a), which presents the ratios of returns achieved by CoWorld compared to those achieved by the Offline DV2 without source domain pretraining. Values larger than 1 indicate positive domain transfers. We can see that our approach improves the offline RL performance in the majority of scenarios (26/30 cases).

(2) Technical details and experimental analyses of Multi-Source CoWorld:

CoWorld can easily accommodate multiple source domains by adaptively selecting suitable tasks as the auxiliary simulator. This extension can be easily achieved by measuring the latent space distance.

Method: The key idea of adaptive domain selection is to allocate a set of one-hot weights $\omega_t^{i=1:N}$ to candidate source domains by calculating their KL divergence in the latent state space to the target domain, where $i \in [1,N]$ is the index of each source domain. The adaptive domain selection procedure includes the following steps:

• World Models Pretraining: We pretrain a world model for each source domain individually.
• Domain Distance Measurement: At each training step in the target domain, we measure the KL divergence between the latent states of the target domain, produced by $e^{(T)}\_{\phi}(o_t^{(T)})$, and corresponding states in each source domain, produced by $e^{(S)}_{\phi^{\prime}_i}(o_t^{(T)})$. Here, $e^{(T)}\_{\phi}$ is the encoder of the target world model and $e^{(S)}\_{\phi^{\prime}_i}$ is the encoder of the world model for source domain $i$.
• Auxiliary Domain Identification: We dynamically identify the closest source domain with the smallest KL divergence. We set $\omega\_t^{i=1:N}$ as a one-hot vector, where $\omega\_t^i=1$ indicates the selected auxiliary domain.
• Rest of Training: With the one-hot weights, we continue the rest of the proposed online-to-offline training approach. For example, in the representation learning stage, we adaptively align the target state space to the selected online simulator by rewriting the domain alignment loss term in Eq. (3) as

Results: Below, we show the multi-source results on the cross-environment benchmark. We take three source domain tasks from Meta-World for adaptive source domain selection, while the target domain is from the RoboDesk environment.

Additionally, we evaluate the multi-source CoWorld across multiple tasks within the Meta-World benchmark. Here, we use "Best-Source" to present the best results obtained on the target dataset among the 6 single-source CoWorld models using each individual source task as the auxiliary domain.

Findings: We observe that CoWorld can flexibly adapt to the transfer learning scenarios with multiple source domains, achieving comparable results to the model that exclusively uses our manually designated auxiliary simulator as the source domain (Best-Source). This study significantly improves the applicability of CoWorld in broader scenarios.

We have included the above results in Table 1, Figure 3, and Figure 4 in the revised paper.