Paper: Collaborative World Models: An Online-Offline Transfer RL Approach (ICLR-2024)

Reviewer comment (Experiments):
One main challenge in the proposed approach is how to choose a source RL simulator. Clearly, the key of the online-to-offline transfer lies in the quality of the source agent or the alignment between the source and target tasks. A priori, 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? Needless to say, one cannot choose it arbitrarily. So how do you specify the threshold according to the values in the transfer matrix? Also, what is the impact of the source task quality on the learning performance? Many of these details are missing.

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

Reviewer comment (Reproducibility):
This work is compared with the fine-tune method, where in offline RL, there are many possible ways to fine-tune the model. It is important to specify the details of these methods, e.g., how many steps for fine-tuning for each task?

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

| Method   | Meta-World | DMC  | Meta-World $\rightarrow$ RoboDesk |
| -------- |:----------:|:----:|:---------------------:|
| Finetune |    300k    | 600k |    300k    |
| CoWorld  |    300k    | 600k |   300k   |

Reviewer comment (Evaluation):
What is the variance in Figure 3(b)?

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

Reviewer comment (Evaluation):
In both Figure 5(a), would you explain why in iteration around 150k, the CoWorld performance is worse or similar compared with CoWorld w/o State A/B/C?

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

Reviewer comment (Theory):
What is the computational complexity of the training phase of the auxiliary agent?

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

Reviewer comment (Evaluation):
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?

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

| Task | DrQ-v2 Value | DrQ-v2+CQL Value |
| --- | :---: | :---: |
| Door Close | 959 | 639 |
| Button Press | 836 | 766 |
| Handle Press | 1104 | 921 |
| Window Close | 903 | 892 |
| Button Topdown | 1211 | 1076 |
| Drawer Close | 1001 | 924 |
| Avg. | 1002 | 870 |

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.

Reviewer comment (Experiments):
It is not clear how the source domain should be selected in practice. Although the paper provides a selection metric based on the ratio between the finetune model and the online model, it is not practical because we do not have access to the test-time domain (if you have them, why do you do offline learning?) Here the $R_{Online}$ should not be available during training time. It makes the experiment results less convincing, as the source domain selection is actually **leaking information from test to train**.

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

Reviewer comment (Theory):
From the formulation and the algorithm, it is not clear how the source and the target domains are related, which makes it hard to justify the feasibility of the method.

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

Reviewer comment (Theory):
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. Based on the formulation provided, the source domain and the target domain may have different observation spaces $O^{(S)}$ and $O^{(T)}$. Why does it make sense to align $p(s_t^{(S)}|o_t^{(T)})$ and $p(s_t^{(T)}|o_t^{(T)})$? Does it mean you assume the similarity between $s_t^{(S)}$ and $s_t^{(T)}$? The relation between two domains are not clear to me (e.g. how similar are $O^{(S)}$ and $O^{(T)}$). The same issue exists across multiple equations, where the regularization is to align the t-th step of the source and the target, without explaining why they CAN be aligned.

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

Reviewer comment (Experiments):
Can you evaluate the performance of the algorithm without the designed "source domain selection" (for the reason I specified in weakness)? What if the source domain does not match the requirement? Will it hurt the performance by a lot?

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

Reviewer comment (Writing):
It's not clear if CoWorld is a method or a new problem statement. In the reviewer's opinion, the paper presents two things: (1) a new problem statement where we have a simulator where we can do online RL with the objective of performing well on an offline dataset from a target environment, and (2) a method for doing well in this new problem statement.

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

Reviewer comment (Theory):
**Unrealistic assumptions**. The paper tries to tackle offline RL by assuming the existence of an off-the-shelf simulator from a source domain. On one hand, as discussed in part 4.2, the success of this online-to-offline paradigm highly hinges on the similarity between the source and target domains. On the other hand, offline RL research is motivated by the assumption that online data collection can be too costly or dangerous on the target domain. These two conditions seem contradictory since if we cannot query large amount of online data from a target domain due to safety or economic issues, it's unlikely that we can build a high-quality simulator on a very similar source domain.

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

Reviewer comment (Experiments):
**Marginal or unclear performance improvement of the method**. The paper compares CoWorld with a number of offline RL methods trained **on the target domain only** with only one exception, which is a simple transfer learning strategy which finetunes the source model on the target domain. I find the performance improvement of CoWorld not convincing enough due to the following considerations: There are many potentially better baseline methods that can leverage data from both domains, which the paper doesn't compare. To name a few: finetuning the model with early stopping/soft update/elastic weight consolidation [1], or finetuning the model and source agent simultaneously, or simply co-training a policy on both domains simultaneously.

Author response:
> Q4: The paper compares CoWorld with a number of offline RL methods trained on the target domain only with only one exception. There are many potentially better baseline methods that can leverage data from both domains.

Following the reviewer's suggestion, we additionally compare CoWorld with the following baselines:
- **More Finetuning Baselines:** We pretrained the baseline models, including Offline DV2, DrQ+BC, and LOMPO, using the source domain data, and then finetuned the models on the target dataset. We found that all of the baseline models suffer from the so-called "negative transfer" effect when incorporating a source pretraining stage. --- A naive transfer learning scheme might degenerate the target performance by introducing unexpected bias. Please refer to our **General Response (Q3)** for details of the quantitative results. 
- **Finetune+EWC:** We introduced a baseline model utilizing the continual learning approach, EWC, to enforce regularization for preserving knowledge from the online source domain. However, our observations reveal that, without additional model designs for domain adaptation, maintaining knowledge from the source domain could affect the performance in the target domain. 
- **One World Model on Both Domains:** We implemented a baseline model named "Multi-Task DV2", which involves training DreamerV2 on both offline and online data with a joint world model and separate actor-critic models. As we can see from the results below, CoWorld consistently performs better.
- **One Policy Model on Both Domains:** We also implemented the "OnePolicy" model, which employs separate world models while learning one shared policy across both source and target domains.

Here, we summarize the results of the newly added baseline models for the cross-environment experiments (Meta-World to RoboDesk). We present the mean returns and standard deviations that are calculated over 10 episodes with 3 random seeds. It is important to note that all of these models can access the source domain data as CoWorld does, and our approach still achieves the best performance.

| Meta-World $\rightarrow$ RoboDesk | Button Press  $\rightarrow$ Push Button | Window Close $\rightarrow$ Open Silde |
|:--------------------------------- |:--------------------------------------- |:------------------------------------- |
| Multi-Task DV2                    | 342 $\pm$ 29                            | 173 $\pm$   22                        |
| Finetune+EWC                      | 259 $\pm$ 14                            | 161 $\pm$   6                         |
| OnePolicy                         | 364 $\pm$  51                           | 186 $\pm$  27                         |
| CoWorld                           | **428 $\pm$ 42**                            | **202 $\pm$ 19**                          |

We have included these comparisons in the revised paper.

Reviewer comment (Experiments):
The performance of CoWorld seems to be highly sensitive to hyperparameters such as the Target-inclined reward factor $k$. According to Figure 11, by simply changing $k=0.3$ to $k=0.7$ results in about 600 decline to ~3000 in episode return. In this case, CoWorld is inferior to both Finetune and Offline DV2 on DC $\rightarrow$ BP task according to Table 1. The improvement in Figure 9 also looks very marginal compared to the naive Finetune baseline.

Author response:
> Q5: Hyperparameter sensitivity of $k$.

We have refined the hyperparameter analysis results in Figure 11, providing a more detailed exploration of $k$. The results below showcase the robustness of CoWorld in the range of $0.1 \le k \le 0.5$, achieving more than 3200 episode returns on average.

| $k$ | 0 | 0.1 | 0.2 | 0.3 | 0.4 | 0.5 | 0.6 | 0.7 | 0.8 | 0.9 | 1 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| CoWorld's Returns | 2550 $\pm$ 493 | **3212 $\pm$ 801** | **3392 $\pm$ 771** | **3623 $\pm$ 543** | **3489 $\pm$ 751** | **3451 $\pm$ 463** | 3052 $\pm$ 1138 | 3036 $\pm$ 698 | 2608 $\pm$ 998 | 3056 $\pm$ 623 | 3136 $\pm$ 867 |

We'll provide more analyses with finer intervals on the other two hyperparameters.

Reviewer comment (Experiments):
Your method relies on the availability of an online RL simulator. 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?

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

Reviewer comment (Evaluation):
For all experiments, competitive model-free methods such as CQL consistently underperform model-based methods by a large margin. However, we generally expect that model-free methods achieve higher asymptotic performance compared to model-based methods at the cost of sample efficiency. Any explanations?

Author response:
> Q7: Why do the model-free methods such as CQL consistently underperform model-based methods by a large margin?

It is an interesting question, and we will explain this in two aspects:
- First, in the visual continuous control tasks on the Meta-World/DMC/RoboDesk benchmarks, the base model, DreamerV2, typically outperforms the model-free methods such as CURL and DrQ-V2 (i.e., our base model for CQL). Such results are illustrated in Figure 7 from the arXiv version of the DrQ-V2 paper ("Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning" by Yarats et al., 2021). 
- Second, in offline visual RL scenarios, model-free approaches are more likely to be affected by the overfitting problem when learning latent states from limited offline visual observations. This challenge can be partly alleviated by the model-based approaches, as the "imagination" process in the policy learning procedure in DreamerV2 (and our approach as well) can naturally augment the offline data.

Reviewer comment (Theory):
*Dependence on Auxiliary Simulators:* 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. This reliance may limit the method's applicability in situations where such simulators are unavailable or where the source and target domains do not share sufficient similarity to ensure sufficient knowledge transfer.

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

Reviewer comment (Experiments):
*Decoupling Rationale*: While it is easy to understand how getting more data from an online simulator in the source domain is beneficial, the rationale behind the decision to maintain two separate world models (rather than a single, jointly trained model on both offline and online data) is not sufficiently explained or supported by evidence.

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

| Meta-World $\rightarrow$ RoboDesk               | Button Press  $\rightarrow$ Push Button | Window Close $\rightarrow$ Open Silde |
|:------------------- |:-------------------------------------------- |:------------------------------------- |
| Multi-Task DV2 (One world model) | 342 $\pm$ 29                                   | 173 $\pm$   22                        |
| CoWorld             | 428 $\pm$ 42                                 | 202 $\pm$ 19                          |

Reviewer comment (Experiments):
Were all other approaches you compare to (except for the "Finetune" baseline) trained only on the offline data?

Author response:
Yes, they were. However, to further demonstrate the effectiveness of CoWorld, we conducted additional experiments to compare CoWorld with new baselines that also involve the source data for model pretraining. Please refer to the results in our **General Response (Q3)**. From the results, we may conclude that a naive transfer learning method without any domain adaptation designs may introduce unexpected bias from the source domain and degenerate the target performance.