Open-World Reinforcement Learning over Long Short-Term Imagination

We thank the reviewer for the constructive comments.

Q1: Clarification of "jumpy" and "affordance map".

The jumpy transition enables the agent to skip intermediate states from current state and directly simulate a future state that is more relevant to the task. We have revised the introduction section to define the term "jumpy" when it first appears in the paper. Besides, we have introduced the affordance map at the end of Page 2.

Q2: Some other unclear aspects/minor mistakes.

We have made the following changes in the revised paper:

• Line 326: We have revised the sentence.
• Line 351: We have corrected the grammar.
• Line 354: We have highlighted the difference from the DreamerV3 loss.
• Line 361: We have checked the line and confirm that there is no typo; we originally wrote "challenging".
• Line 500: We have simplified the phrasing.
• We have compared and cited "Learning to Move with Affordance Maps".

Q3: Effectiveness of LS-Imagine in sparse environments.

The long-term transitions in our approach leverage the affordance map to identify high-value exploration areas. Even in environments where targets are sparse, our approach continues to benefit from the affordance maps. Unlike object recognition algorithms, which highlight areas only when targets are present, our method remains effective even in the absence of clear target cues. Thanks to MineCLIP's pretraining on extensive expert demonstration videos, our approach can generate affordance maps that provide guidance even when the target is sparse or even completely occluded.

For instance, as illustrated in Figure 13 in the revised paper, throughout the task of locating a village, the affordance map consistently provides effective guidance to the agent, suggesting exploration of the forest to the right or the open area on the left hillside, even when the village is not visible in the current observation. Similarly, in mining tasks where ores are typically underground, the affordance map directs the agent to dig into the mountain area on the right. As we can see, even when the target is occluded, the affordance map enables the agent to continue exploring effectively.

Potential research direction: Nevertheless, we agree that due to the complexity of open-world environments, the affordance map may fail to provide effective guidance in scenarios that the MineCLIP model has not encountered before. To address this issue, we plan to progressively finetune the MineCLIP model with the collected new data and introduce a new prompt to the agent: "Explore the widest possible area to find {target}" when the affordance map fails to identify high-value areas. This prompt, combined with intrinsic rewards generated by MineCLIP, encourages the agent to conduct extensive exploration.

Q4: Limitations of LS-Imagine.

We have included the discussion about the limitations of our approach in Section 6.

Q5: On the imagination horizon of LS-Imagine.

It is a typo. The imagination horizon of our approach is 15 steps. We have corrected this in the revised paper.

Q6: Could the limitations be expanded upon to cover the general application of LS-Imagine beyond Minecraft?

Apart from open-world environments like Minecraft, our method is also applicable to navigation environments such as MiniWorld [1], Habitat [2], and DeepMind Lab [3]. As long as embodied agents are navigating a 3D environment where rewards are obtained by approaching objects associated with specific goals, our method can be effectively utilized.

References:

[1] Chevalier-Boisvert et al. "Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks", NeurIPS Track on Datasets and Benchmarks, 2024.

[2] Savva et al. "Habitat: A platform for embodied ai research", ICCV, 2019.

[3] Beattie et al. "Deepmind lab" arXiv preprint arXiv:1612.03801, 2016.

Q5: The dynamics threshold on lines 248-250 needs clarification, and the "jumpy transition" should be explained further, including how often it's triggered and the average number of short-term steps saved.

(1) On the dynamic threshold:

We design the dynamic threshold $P_\text{thresh}$ for two purposes:

• At the beginning of the training process, we do not have an accurate estimation of $P_{\text{jump}}$. To stabilize the training process in the imagination phase, we only trigger jumpy state transitions when $P_{\text{jump}} > P_{\text{thresh}}$, which is computed as the moving average of all $P_{\text{jump}}$ collected in the experience replay buffer.
• Manually selecting suitable $P_{\text{thresh}}$ values for different tasks is a time-consuming work, because the threshold should account for the varying guidance strength provided by the affordance map across different tasks (as shown in Figure 12). Therefore, we propose to compute $P_{\text{thresh}}$ automatically.

Specifically, from the very beginning of training, we store the $P_\text{jump}$ values for every interaction step in a data buffer. The threshold $P_\text{thresh}$ is then dynamically calculated as the mean of all $P_\text{jump}$ values currently in the buffer plus their standard deviation. This dynamic adjustment ensures that the threshold adapts to task-specific features and remains robust throughout training.

In Figure 11$c$ in the revised appendix, we track the variation curves of $P_\text{thresh}$ during training in different tasks, and observe that:

• For task such as "harvest log in plains", the variance of $P_\text{thresh}$ is high during the early stages of training. Since $P_\text{thresh}$ serves as a temporal smoothing of $P_\text{jump}$, this reflects the significant fluctuations of $P_\text{jump}$ at the beginning of training, highlighting the importance of adopting a triggering threshold.
• Across various tasks, $P_\text{thresh}$ consistently converges in the later stages of training, demonstrating its effectiveness in improving the stability of exploratory imaginations.
• The converged values of $P_\text{thresh}$ differ across tasks, indicating that involving an automated computation of $P_\text{thresh}$ enables us to avoid tedious hyperparameter tuning.

Please refer to Page 19 in our revision for detailed results.

(2) On the frequency and jumping state intervals of the "jumpy transitions":

In Figure 11$a-b$, we showcase the frequency of jumpy transitions in the imagination process and the predicted $\hat{\Delta}_t$ throughout training. We have included further analyses in Appendix D.2 (Pages 18-19). Please refer to our revision.

Q6: Add a numerical result for the results in Figures 4 and 5.

As suggested by the reviewer, we have added numerical results in Table 3 in the revised appendix.

Q7: Clarify if the post-jumping state is connected with the bootstrapping $\lambda$-return of the pre-jumping states and explain the basis for the conclusion on lines 464-466.

(1) The connection of post-jumping state and pre-jumping states.

As indicated by Equation 9, the long-term rewards from the post-jump imaginations directly affect the bootstrapping $\lambda$-return of the pre-jump states within the same sequence, thereby influencing the value estimation of those pre-jump states.

For further details, please refer to our responses to Reviewer hBEP's Q3.

(2) Why post-jumping state does not guide the prior-jumping transitions in parallel variant of LS-Imagine?

As illustrated in Figure 10(b), an alternative method is to structure long- and short-term imagination pathways in parallel. Specifically, we begin by applying short-term imagination within a single sequence. For each predicted state, we examine the jumping flag: If $\hat{j}_t = 1$, we initiate a new imagination sequence starting from the post-jump state, which is predicted by the long-term transition model and the dynamics predictor. In other words, whenever a long-term state jump occurs, the world model generates a new sequence from the post-jump state, while the intermediate state transitions within the sequence are governed exclusively by short-term dynamics. Importantly, we optimize the actor independently for each sequence, ensuring that there is no gradient or value transfer between sequences.

We have included these discussions in Section 4.3 in the revision.

Q8: Add another two world model papers for completeness.

We have included the suggested papers in Section 5 in the revision.

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

Q1: Add an introduction of MineCLIP in Section 3.2/3.2.1.

We have added a brief introduction of MineCLIP to section 3.2.1 in the revised paper.

Q2: The abstract should be revised to emphasize "across a long horizon" to better align with the paper's main contribution.

The main contribution of our approach is improving the efficiency of exploration across an extensive state space in open-world decision-making, especially for tasks that demand consideration of long-horizon payoffs. We refine the abstract in the revised paper.

Q3: Dependence on the visibility of objects. The long-term transitions of our approach rely on the affordance map to identify high-value exploration areas. However, it is crucial to note that our affordance map generation method is not merely a target recognition algorithm that highlights areas only when the target is present. Thanks to MineCLIP's pretraining on extensive expert demonstration videos, our approach can generate value maps that provide guidance even when the target is completely occluded.

For instance, as illustrated in Figure 13 in the revised paper, throughout the task of locating a village, the affordance map consistently provides effective guidance to the agent, suggesting exploration of the forest to the right or the open area on the left hillside, even when the village is not visible in the current observation. Similarly, in mining tasks where ores are typically underground, the affordance map directs the agent to dig into the mountain area on the right. As we can see, even when the target is occluded, the affordance map enables the agent to continue exploring effectively.

Furthermore, we agree that due to the complexity of open-world environments, the affordance map may fail to provide effective guidance in scenarios that the MineCLIP model has not encountered before. To address this issue, we plan to progressively finetune the MineCLIP model with the collected new data and introduce a new prompt to the agent: "Explore the widest possible area to find {target}" when the affordance map fails to identify high-value areas. This prompt, combined with intrinsic rewards generated by MineCLIP, encourages the agent to conduct extensive exploration.

Q4: Limitation of LS-Imagine.

We have included the limitation of LS-Imagine in Section 6 in the revised paper.

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

Q1: The method feels very ad-hoc to the Minecraft tasks.

LS-Imagine can be extended to the navigation environments, as long as embodied agents are navigating a 3D environment where rewards are obtained by approaching objects associated with specific goals, such as MiniWorld [1], Habitat [2], and DeepMind Lab [3].

References:

[1] Chevalier-Boisvert et al. "Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks", NeurIPS Track on Datasets and Benchmarks, 2024.

[2] Savva et al. "Habitat: A platform for embodied ai research", ICCV, 2019.

[3] Beattie et al. "Deepmind lab" arXiv preprint arXiv:1612.03801, 2016.

Q2: The term "Gaussian matrix" in line 215 is unclear and likely refers to a Gaussian kernel; please specify the $\sigma$ used in the paper and discuss the method's sensitivity to different $\sigma$ values.

(1) Definition of "Gaussian matrix": The term "Gaussian matrix" mentioned in line 215 is indeed intended to refer to a Gaussian kernel. In our experiments, we use ($\sigma_x = 128,\sigma_y = 80$) as the standard deviations for the Gaussian kernel.

The visualization of Gaussian matrices with different standard deviations is presented in Figure 16. As shown in the visualizations, intuitively, setting these hyperparameters too low may cause the model to overlook targets located at the edges of the observed images. Conversely, excessively high $(\sigma_x, \sigma_y)$ may reduce the reward discrepancy for targets at different positions within the observation, thereby diminishing the agent's incentive to focus on the target precisely.

(2) Sensitivity of hypermeter $\sigma$: We present the sensitivity analysis of $\sigma$ in the Figure 15 (right) in the revised paper. We observe that the final performance is robust to the tested parameters, with all configurations outperforming the baseline models presented in previous experiments.

Q3: How are task-relevant observations ensured in the random exploration data used to pretrain the affordance map generator, particularly for tasks where random agents might not encounter task-related objects?

In our current work, we employ MineCLIP to calculate the correlation between the images of simulating exploration and task description, which provides strong prior knowledge about the environment and helps guide the learning process by associating visual features with meaningful task-related concepts.

We plan to further address this issue from two aspects:

• First, we can incorporate human expert demonstrations, which provide critical examples that random agents may otherwise fail to encounter, particularly when interacting with objects related to specific tasks.
• Second, we can iteratively update the UNet model, allowing it to progressively refine its understanding of affordances and become increasingly effective at recognizing task-relevant observations, even in more challenging scenarios.

Q4: The difference between the sequential pathway and the parallel pathway.

In Figure 10(a) in the revised paper, we provide a visualization illustrating how the agent sequentially performs short-term and long-term imaginations within a single imagination trajectory.

As illustrated in Figure 10(b), an alternative method is to structure long- and short-term imagination pathways in parallel. Specifically, we begin by applying short-term imagination within a single sequence. For each predicted state, we examine the jumping flag: If $\hat{j}_t = 1$, we initiate a new imagination sequence starting from the post-jump state, which is predicted by the long-term transition model and the dynamics predictor. In other words, whenever a long-term state jump occurs, the world model generates a new sequence from the post-jump state, while the intermediate state transitions within the sequence are governed exclusively by short-term dynamics. Importantly, we optimize the actor independently for each sequence, ensuring that there is no gradient or value transfer between sequences.

We have included these discussions in Section 4.3 in the revision.

Q5: Some relevant citations are missing.

We have added relevant citaions in Section 5 and provided further discussions to compare these models.

To limit the growth of parallel sequences, we impose a constraint: among the predicted states that $\hat{j}_t=`True`$ in a single imagination sequence, at most one state is selected for generating a parallel sequence via long-term transition. Since the imagination horizon of a short-term sequence is relatively short, states in the same short-term sequence typically transit to a similar post-jumping state after a long-term transition.

Furthermore, as presented in Figure 11(a) of the revised manuscript, the frequency of long-term transitions is relatively low$\textemdash$approximately less than 10%. This ensures that the number of branching of new sequences is limited.

Another potential cause for batch size growth could stem from generating additional parallel sequences on top of previously generated ones. However, we find that among all sequences with jumpy state transitions, the average number of jumpy transitions per sequence, within a horizon of 15 steps, is 1.02. This indicates that in most cases, a single jumpy transition is enough to enable the agent to approach the target, and subsequent long-term transitions are not needed.

Overall, these findings ensure that the number of newly generated parallel sequences remains controlled, and batch size variations are not significant. Therefore, we believe the comparison between the parallel and series pathways is still fair.

Thank you for reviewing our paper. Below, we provide detailed responses to your comments.

Q1: Comparison with Director.

We compare the performance of Director and our approach. As shown below, LS-Imagine outperform Director in terms of success rate (%) on all tasks.

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

Q2: Inconsistency in time indexing in Equation 9.

Thank you for your comment. We would like to clarify that the indexing in Equation 9 is indeed correct. Notably, Equation 9 is used in the model-based behavior learning process, where the index $t$ represents positional order of the imagination states predicted by the model, rather than the real time step during environmental interactions. For example, starting from state $\hat{s}\_t$, any subsequent state obtained via either a short-term transition or a long-term transition is indexed sequentially as $\hat{s}\_{t+1}$.

Therefore, in Equation 9, the bootstrapping term is indexed as $R\_{t+1}^\lambda$, which is consistent with the sequential indexing convention we adopt for the imagination process; while we apply $\hat{\Delta}\_{t+1}$ to the discount factor to extend the $\lambda$-return formulation, enabling it to handle temporal spans beyond single imagination steps.

For better illustration, we provide the sketch of a typical imagination sequence in Figure 10(a) in the revision, where the transition from $\hat{s}\_2$ to $\hat{s}\_3$ represents a short-term transition, while that from $\hat{s}\_3$ to $\hat{s}\_4^\prime$ is a long-term transition, and $\hat{\Delta}\_4$ estimates the time step interval between the pre-jump state ($\hat{s}\_3$) and the post-jump state ($\hat{s}\_4^\prime$) when the agent interacts with the real environment.

We hope this clears up any confusion, and we appreciate your careful attention to this detail.

Q3: Evaluating the policy with $\lambda$-return may introduce bias since it relies on on-policy data, and the predicted jumpy states might not align with the learning policy.

Let's take the imagination trajectory in Figure 10(a) in the revision as an example, where:

$s\_1 \rightarrow \hat{a}\_1 \sim \pi\_\theta(s\_1) \rightarrow \hat{s}\_2 \rightarrow \hat{a}\_2\sim \pi\_\theta(\hat{s}\_2) \rightarrow \hat{s}\_3 \rightarrow$ Jump $\rightarrow \hat{s}\_4^\prime \rightarrow \hat{a}\_4\sim \pi\_\theta(\hat{s}\_4^\prime) \rightarrow \hat{s}\_5 \rightarrow \cdots$

Although the long-term transition operation ($\hat{s}\_3 \rightarrow \hat{s}\_4^\prime$) is action-free, the discounted cumulative rewards in Equation 9 remain effective for the following reasons:

• The rewards derived from the pre-jump actions $(\hat{a}\_1, \hat{a}\_2)$ and the post-jump actions $(\hat{a}\_4, \cdots)$ are all consistent with the current policy.
• While ($\hat{s}\_3 \rightarrow \hat{s}\_4^\prime$) is action-free, the post-jump states $\hat{s}^{\prime}\_4$ is indirectly influenced by $\pi\_\theta$ through $\hat{s}\_3$. Therefore, Equation 9 can evaluate the potential long-term advantage of the pre-jump state $\hat{s}\_3$ by performing the $\lambda$-return iterations across the jumping operation.
• In Equation 9, we mitigate the influence of a biased post-jump reward by applying $\hat{\Delta}_{4}$ to the discount factor when updating $R\_3^\lambda$ with $R\_4^\lambda$. Here, $\hat{\Delta}\_4$ estimates the number of intermediate steps in between when we use $\pi\_\theta$ to interact with the real environment.

Technically, the $\lambda$-return in Equation 9 tries to strike a balance between the long-horizon evaluation of the current policy and mitigating reward bias. In open-world tasks, we argue that encouraging long-horizon exploration of the agent is more crucial, as it enables the agent to develop efficient policies for navigating diverse environments.

Dear Reviewer hBEP,

Thank you once again for your time in reviewing our paper. If any concerns remain, please don't hesitate to inform us.

Best regards,

Authors

Thank you! I really appreciate the detailed explanation and additional experiment results from the author. My concerns have been addressed, and I would like to increase my score to 8.