Making Offline RL Online: Collaborative World Models for Offline Visual Reinforcement Learning

Q1: The availability of an online environment.

There are many real-world scenarios (such as robot control, autonomous driving, and healthcare) where existing simulators can be employed as source domains. See General Response (Q1) for details.

Further, our method is specifically designed to address cross-domain inconsistencies in dynamics, visual appearances, and action space (Table 1), thereby reducing the necessity for a very similar online environment.

Q2: How dependent is CoWorld on the quality of the online simulator?

In Fig 3 (left) of the original paper, we evaluate our model's performance across each pair of source and target domains. We observe positive transfer effects (indicated by the red blocks) in 18/20 single-domain transfer scenarios. Indeed, we also observe negative transfers (blue blocks) in 2/20 scenarios.

In Lines 159-162 and Appendix C, we present the multi-source CoWorld that can automatically discover a beneficial task from multiple source candidates. In Fig 3 (left), the multi-source CoWorld consistently achieves positive transfer (highlighted by the dashed blue box). The right side of Fig 3 demonstrates the effectiveness of the multi-source CoWorld. It yields results comparable to those of the top-performing single-source CoWorld.

Q3: Why maintain separate source and target world models?

First, using separate world models allows the source and target domains to have significantly different observations. See General Response (Q2, Point(1)) for additional empirical results where the source domain has a low-dim observation space.

Additionally, we compare CoWorld (w/ separate models) and Multi-Task DV2 (w/ one model) in Table 7 in the Rebuttal PDF. The latter involves training on offline and online data with a joint world model and separate actor-critic models. CoWorld consistently performs better in terms of the cross-environment results (Meta-World -> RoboDesk).

Q4: Training on multiple source domains.

As suggested, we design two baseline models trained with multiple source domains and compare them with our proposed multi-source CoWorld:

• DV2-FT (multi-source pre-train), which involves pre-training on DC, WC, and DC* as sources for 40K steps on each task and then finetuning on BP for 300K steps.
• CoWorld (multi-source co-train), which involves co-training with DC, WC, and DC*, as well as the target (BP), achieving a return of 3012.

As shown in Table 8 in the Rebuttal PDF, training on multiple sources can lead to negative transfer. In contrast, CoWorld can identify appropriate auxiliary tasks, enabling it to better manage the potential negative impacts of less relevant source tasks

Q5: Computational complexity.

We present the results on Meta-World (BP -> HP) in Fig 1 in the Rebuttal PDF. CoWorld achieves convergence (90% of the highest returns) in approximately 14 hours; while it costs DV2 Finetune (DV2-FT) about 13 hours. These results indicate that CoWorld requires a comparable training wall-clock time to DV2-FT, while consistently maintaining better performance in terms of returns after model convergence.

We further illustrate the efficiency of CoWorld from two aspects:

• Convergence time: From the results above, CoWorld and DV2-FT achieve comparable convergence speeds in wall-clock time. Thereafter, our approach consistently outperforms DV2-FT.
• Improving efficiency: In Table 6 in the Rebuttal PDF, we add new experiments demonstrating that efficiency can be further optimized by flexibly reducing the frequency of source training steps during the co-training phase. The overall performance still surpasses that of DV2-FT (3702 $\pm$ 451) and DV2-FT-EWC (528 $\pm$ 334).

CoWorld, like the baseline models, can be trained using a single 3090 GPU.

Q6: More insight into selecting an appropriate source domain.

Recognizing the challenges of manually selecting an appropriate source domain, we propose a multi-source CoWorld approach with an adaptive domain selection method.

First, we typically have some information about the target task, based on which we can initially construct a set of source domain candidates (e.g., if the target domain is a robotic arm visual control task, we can randomly select several tasks from Meta-World or other environments as potential source domains), and we allow some of these source domains to differ significantly from the target task.

Next, we propose a method for adaptive source domain selection. 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 technical details in Appendix C and corresponding results in Fig 3.

Q7: The performance gains over simpler baselines (fine-tuning).

Notably, finetuning requires a very high similarity between the source and target domains, resulting in poor performance in the following three scenarios:

1. Cross-environment experiments (Fig 4)
2. Adding-noise experiments (Table 6)
3. Scenarios with sparse rewards in the source domain (General Response Q2, Point(2)).

As shown in Table 9 in the Rebuttal PDF, CoWorld demonstrates significant improvement over the fine-tuning method, particularly in cases where there are more pronounced domain discrepancies.

Furthermore, another disadvantage of the finetuning method is that it cannot be applied in scenarios where source and target domains have observation spaces of different dimensions. See General Response (Q2, Point(1)).

Q8: What if the source and target domains have significantly different observation spaces?

We have conducted new experiments using a source task with low-dimensional states (Meta-World) and a target domain characterized by high-dimensional images (RoboDesk). Please refer to General Response (Q2, Point(1)) for the results.

Q9: Clarity in presentation.

We hope Table 10 in the Rebuttal PDF can make the paper easier to understand.

Thank you for your constructive comments. Our responses are provided below. If you have any further questions, please feel free to let us know.

Q1: Not sure if taking 10 episodes over 3 seeds is sufficient.

We extended the results from 3 random seeds to 5 random seeds and tested the performance variations with 3, 5, and 10 random seeds, as well as 10, 20, 50, and 100 episodes. The experiments demonstrate that CoWorld consistently outperforms the baseline model DV2 Finetune (DV2-FT) in terms of performance stability across various tasks and different numbers of random seeds. Please refer to General Response (Q3) for more details.

Q2: The source of success in transfer, particularly in reward alignment, is unclear and seems to impose strict limitations on task relevance, which is not sufficiently discussed.

(1) On the reasoning behind 'Online-to-Offline Reward Alignment'

First, let us review the training details in the reward alignment process:

• Step 1: We sample $\\{o^{(T)}, a^{(T)}, r^{(T)}\\}$ from the target offline dataset;
• Step 2: In Eq. (4), we use the source reward predictor to predict $R_{\phi^\prime}(o^{(T)}, a^{(T)})$ based on the target inputs;
• Step 3 (Important!): Also in Eq. (4), we mix the previous output and the true target reward to generate $\tilde{r}^{(S)}=(1-k) \cdot R_{\phi^\prime}(o^{(T)}, a^{(T)}) + k \cdot r^{(T)}$;
• Step 4 (Important!): In Eq. (5), we use $\tilde{r}^{(S)}$ to supervise the training of $R_{\phi^\prime}(o^{(T)}, a^{(T)})$ based on the target inputs;
• Step 5: Also in Eq. (5), we remain the training term of $R_{\phi^\prime}(o^{(S)}, a^{(S)})$ based on the source inputs to approach the raw source reward data $r^{(S)}$.

Clearly, due to the aforementioned Steps 3-4, the reward model $R_{\phi^\prime}(\cdot)$ learns to approximate the reward combined with the true target reward $r^{(T)}$ conditioned on the target inputs $o^{(T)}$ and $a^{(T)}$, where $\{o^{(T)}, a^{(T)}, r^{(T)}\}$ are from the same target tuple. In this way, the reward alignment process integrates target information into the learned reward function of the source domain. Consequently, it enables the subsequent learning process of the source critic to be informed with rich target domain knowledge.

(2) Can CoWorld benefit from a source domain with a significantly different reward signal?

Yes, CoWorld can benefit from a source domain with a significantly different reward signal. In our experiments, the source rewards are sparse and the target rewards are dense. The experiment results demonstrate that although excessively sparse rewards can hinder the training process, CoWorld still outperforms DV2-FT with a more balanced reward structure. For detailed results and analysis, please refer to General Response (Q2, Point(2)).

Q3: What are the conditions for mitigating value overestimation?

Good question! Regarding the conditions for mitigating value overestimation, Figure 2 illustrates our approach effectively. However, it does indeed assume that the source critic and the ideal target critic perform similarly across a broad range of states. This assumption can be softened by our proposed approaches, namely, 1) state alignment and 2) reward alignment. Through these approaches, we aim to mitigate domain discrepancies between distinct source and target MDPs.

According to Eq. (9) in Appendix A.4, the source critic $v_{\xi^\prime}(\cdot)$ is trained to estimate the expected future rewards \mathbb E[\sum\_{\tau \ge t\}({\gamma}\_{\tau-t} \cdot \hat{r}\_\tau)] given the source latent at timestamp $t$. There are two critical points here:

• For the training supervisions, during the imagination phase, $v_{\xi^\prime}(\cdot)$ is optimized ONLY based on the n-step predicted rewards (Algorithm 1, Lines 22-24), rather than the ground-truth source rewards. As previously demonstrated, the predicted rewards contain a mix of information from both the source and target domains.
• At the input end, $v_{\xi^\prime}(\cdot)$ is indeed trained on the imagined source latents. However, due to the state alignment process, the source latents share close distributions with the latents conditioned on target inputs.

Thus, it is not entirely accurate to suggest that $v_{\xi^\prime}(\cdot)$ is trained to 'maximize the expected sum of source rewards in a range of source latents.' In fact, it learns to maximize the expected sum of target-informed rewards across latents that are closely distributed to the target latents. Consequently, the source critic can be effectively applied to the target domain to guide the training of the target agent, using the target latents as inputs.

Q4: How to ensure that source critic is beneficial to the target task, and is it possible to have a negative transfer when facing conflicting tasks?

In Figure 3 (left) of the original manuscript, we evaluate our model's performance across each pair of source and target domains. We observe positive transfer effects (indicated by the red blocks) in most cases (18/20 single-domain transfer scenarios). And indeed, we also observe negative transfers (indicated by the blue blocks) in 2/20 scenarios (e.g., DC$^*$ $\rightarrow$ BP).

To ensure an effective source critic, as described in Lines 159-162, we first manually select several candidates for the source domain that are potentially relevant to the target task. We then propose a method for multi-source training, with technical details provided in Appendix C. By introducing the multi-source CoWorld, we can automatically discover a beneficial task from multiple source candidates. As shown in Figure 3 (left), the multi-source CoWorld consistently achieves positive transfer (indicated by the results inside the dashed blue box).

On the right side of Figure 3, we further demonstrate the effectiveness of the multi-source domain selection method. It achieves comparable results to that of the top-performing single-source CoWorld.

Dear Reviewer Bpkq,

Thank you once again for your time in reviewing our paper. We have made every effort to improve the overall clarity in the rebuttal. We would appreciate it if you could check our feedback and see if any concerns remain.

Best regards,

Authors

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

Q1: Practical situations where the proposed method is applicable.

Good question! CoWorld can be applied in various real-world scenarios such as robot control, autonomous driving, healthcare, advertising bidding, and power systems. For instance, it can leverage established simulators like Meta-World and CARLA for robotics and autonomous driving, or use medical simulators to evaluate clinical treatments in healthcare. Similarly, it can enhance advertising systems through simulated bidding environments and improve power system operations using electrical system simulators. Please refer to our General Response (Q1) for further details.

Q2: How does CoWorld operate when the two world model architectures have greater differences? Can it still effectively transfer information between these two resources?

To investigate the effectiveness of CoWorld under significantly different world model architectures, extended experiments were conducted where the source domain observations were low-dimensional states and the target domain observations were high-dimensional images. CoWorld achieves competitive performance against strong baselines. Please refer to General Response (Q2, Point(1)) for more details.

Q3: How to prove that CoWorld is really significantly better than other methods? Can you provide the results of more experimental runs or statistical analysis?

We extended the results from 3 random seeds to 5 random seeds and tested the performance variations with 3, 5, and 10 random seeds, as well as 10, 20, 50, and 100 episodes. The experiments demonstrate that CoWorld consistently outperforms the baseline model DV2 Finetune (DV2-FT) in terms of performance stability across various tasks and different numbers of random seeds. Please refer to General Response (Q3) for more detail.

Q4: Line 14 in Algorithm 1: The notation $T$ represents both the target environment and the maximum number of time steps.

Sorry for the misleading notations. We will use $N$ to represent the maximum number of time steps instead of $T$ in the revised paper. For example: $\{(o_{t}^{(T)}, a_{t}^{(T)}, r_{t}^{(T)})\}_{t=1:N} \sim \mathcal{B}^{(T)}$.

Q5: Typos in Table 4 (Page 13) and Line 481 (Page 14).

We will correct these typos in the revised paper.

Q6: Line 516 (Page 17): ...as illustrated in Table 1 $\rightarrow$ should be Table 6?

In Line 516, we wanted to emphasize that the cross-environment setup from Table 1 is even more challenging than the adding-noise setup shown in Table 6, due to more significant domain gaps. Corresponding results for the cross-environment setup are presented in Section 4.3 of the main text.

Thank you for your comments. We have conducted new experiments to discuss the training cost of our model. Please refer to the results presented in the Rebuttal PDF. We hope our responses will address your specific concerns regarding training efficiency.

Q1: How does the performance look when the x-axis is training wall-clock time instead of the number of training iterations?

Per the reviewer's request, we present the results on the Meta-World benchmark (Button Press $\rightarrow$ Handle Press) in Figure 1 in the Rebuttal PDF. As shown, CoWorld achieves convergence (90% of the highest returns) in approximately 14 hours; while it costs DV2 Finetune (DV2-FT) about 13 hours. These results indicate that CoWorld requires a comparable training wall-clock time to DV2-FT, while consistently maintaining better performance in terms of returns after model convergence.

We will include these results in the revised paper.

Q2: The cost of training time is higher in comparison to baselines.

We further illustrate the efficiency of CoWorld from two aspects:

• Convergence time: From the results above, CoWorld and DV2-FT achieve comparable convergence speeds in wall-clock time. Thereafter, our approach consistently outperforms DV2-FT.
• Improving efficiency: In Table 6 in the Rebuttal PDF, we add new experiments demonstrating that efficiency can be further optimized by flexibly reducing the frequency of source training steps during the co-training phase. The overall performance still surpasses that of DV2-FT (3702 $\pm$ 451) and DV2-FT-EWC (528 $\pm$ 334).