Open-World Reinforcement Learning over Long Short-Term Imagination
Authors: Jiajian Li
Qi Wang
Yunbo Wang
Xin Jin
Yang Li
Wenjun Zeng
Xiaokang Yang
Venue: ICLR 2025 Oral (2025)
Keywords: World models
reinforcement learning
visual control
Primary area: reinforcement learning

Abstract: Training visual reinforcement learning agents in a high-dimensional open world presents significant challenges. While various model-based methods have improved sample efficiency by learning interactive world models, these agents tend to be “short-sighted”, as they are typically trained on short snippets of imagined experiences. We argue that the primary challenge in open-world decision-making is improving the exploration efficiency across a vast state space, especially for tasks that demand consideration of long-horizon payoffs. In this paper, we present LS-Imagine, which extends the imagination horizon within a limited number of state transition steps, enabling the agent to explore behaviors that potentially lead to promising long-term feedback. The foundation of our approach is to build a $\textit{long short-term world model}$. To achieve this, we simulate goal-conditioned jumpy state transitions and compute corresponding affordance maps by zooming in on specific areas within single images. This facilitates the integration of direct long-term values into behavior learning. Our method demonstrates significant improvements over state-of-the-art techniques in MineDojo.

Published as a conference paper at ICLR 2025
OPEN-WORLDY REINFORCEMENT LEARNING OVER LONG TORTT-TERM IMAGINATION
Wenjun Zeng² Xiaokang Yang' 1 MoE Key Lab of Artificial Intelligence, Al Institute, Shanghai Jiao Tong University, Shanghai, China
Jiajian Li Qi Wang Yunbo Wang't Xin Jin
Yang Li
2 Ningbo Institute of Digital Twin, Eastern Institute of Technology, Ningbo, China
3 School of Computer Science and Technology. East China Normal University, Shanghai, China
ht htttp::/qlwangeee wit..u....tttb... -iaigngnn
ABSTRACT
Training visual reinforcement learning agents in high-dimensional open world presents significant challenges. While various model-based methods have im- proved sample efficiency by learning interactive world models, these agents tend to be "short-sighted", as they are typically trained on short snippets of imagined experiences. We argue that the primary challenge in open-world decision-making is improving the exploration efficiency across a vast state space, especially for tasks that demand consideration of long-horizon payoffis. In this paper, we present LS- LS-Imagine, which extends the imagination horizon within limited number of state transition steps, enabling the agent to explore behaviors that potentially lead to promising long-term feedback. The foundation of our approach is to build a long short-term world model. To achieve this, we simulate goal-conditioned jumpy state transitions and compute corresponding affordance maps by zooming in specific areas within single images. This facilitates the integration of direct long-term values into behavior learning. Our method demonstrates significant improvements over state-of- of-thh-art techniques in MineDojo.
INTRODUCTION
Open-world mecision-maaing in the context of reinforcement learning (RL) involves the following characteristics: (i) The agent operates within interactive environment that features vast state space; (ii) The learned policy presents a high degree of flexibility, allowing interaction with various objects in the environment; (iii) The agent lacks full visibility of the internal states and physical dynamics of the external world, meaning that its perception of the environment (e.g., raw images) carries substantial uncertainty. For example, Minecraft serves as a typical open-world game.
Building upon recent progress in visual control, open-world mecision-making aims to train agents to ap- proach human-level intelligence based solely on high-dimensional visual observations. However, this presents significant challenges. For example, in Minecraft tasks, existing methods like Voyager (Wang al. 2024a) employ specific Minecraft APIs as the high-level controller, which incompatible with standard visual control settings. While approaches such as PPO-with-MineCLII (Fan et al. 2022) and DECKARD (Nottingham al. (202) perform low-level visual control, these model-free RL methods struggle to grasp the underlying mechanics of the environment. This may result in high trial-and-error costs, leading to inefficiencies in both exploration and sample usage. Although DreamerV3 (Hafner ct al. 2023) employs a model-based RL (MBRL) approach to improve sample efficiency, it is often "short-sighted" since the policy is optimized using short-term
15 time steps- generated by the world model. The absence of long-term guidance significantly hampers an effective exploration of the vast solution space of the open world.
To improve the behavior learning efficiency of MBRL, in this paper, introduce a novel method named Long Short-Term Imagination (LS--maaiine. Our key approach involves enabling the world model to efficiently simulate the long-term effects of specific behaviors without the need repeatedly
Equual contribution.
Corresponding author; Yunbo Wang <yunbow@s jjtoooeawjtu.... edw.

---

Published as a conference paper at ICLR 2025
environment
M+*k
As+
R4
O:
M
Mess
Mesa
Short-Term Imagination
Cut tree
Long-Term Imagination
The general framework of LS- LS-Imagine, an MBRL agent that operates solely on raw pixels. The fundamental idea is to extend the imagination horizon within limited number of state transition steps, enabling the agent to explore behaviors that potentially lead to promising long-term feedback.
Figure I:
rolling one-step predictions. As illustrated in Figure once trained, the world model provides both instant and jumpy state transitions along with corresponding (intrinsic) rewards, facilitating policy optimization in a joint space of short- and long-term imaginations. This encourages the agent to explore behaviors that lead to promising long-term outcomes.
The foundation of LS-Imagine to train long short-term world model, which requires integrating task-specific guidance into the representation learning phase based on off-policy experience replay. However, this creates a classic Shnicken-and-egg dilemma: without true data showing the agent has reached the goal, how can we effectively train the model to simulate jumpy transitions from current states to pivotal future states that suggest a high likelihood of achieving that goal? To address this issue, we first continuously zoom in on individual images to simulate the consecutive video frames as agent approaches the goal. We then generate affordance maps by evaluating the relevance the pseudo video to task-specifis goals presented in textual instructions, using the established MineCLIP reward model (Fan et al., 2022). Subsequently, we train specific branches of the world model to capture both instant and jumpy state transitions, using pairs of image observations from adjacent time steps as well as those across longer intervals. Finally, we optimize the agent's policy based on finite sequence of imagined latent states generated by the world model, integrating more direct estimate of long-term values into decision-making.
Let's the example in Figure further elaborate the novel aspects of the behavior learning process: After receiving the instruction "cut a tree", the agent simulates near-future states based on the current real observation. It initially performs several single-step rollouts until identifies point in time for long-distance state jump that allows it to approach the tree. The agent then executes this jump and optimizes its policy network to maximize the long-sight value function.
We evaluate our approach in the challenging open- woeld tasks from MineDojo (Fan et al. 2022 L.S- Imagine demonstrates superior performance compared existing visual RL methods.
The contributions of this work are summarized as follows:
We present novel model-based RL method that captures both instant and jumpy state transitions and leverages them in behavior learning to improve exploration efficiency in the open world.
Our approach presents four concrete contributions: (i) long short-term world model architecture, (ii) method for generating affordance maps through image zoom-in, (iii) novel form of intrinsic rewards based on the affordance map, and (iv) an improved behavior learning method that integrates long-term values and operates on mixed long short-term imagination pathway.
PROBLEM FORMULATION AND NOTATIONS
We solve visual reinforcement learning as a partially observable Markov decision process (POMDP), using MineDojo as the test bench. Specifically, our method manipulates now-dimensional control signals a, while receiving only sequential high-dimensional visual observations Oct and episodic sparse rewards eny without access to the internal APIs of the open- world games. In comparison, as
As shown Figure a jumpy transition allows the agent bypass intermediate states and directly simulate a task-relevant future state 8e+ æ in one step. This process occurs exclusively during world model imagination. Afforrdacce maps highlight regions within observation that potentially relevant to the task QQi al.
2020
Wang al. 2022.
2

---

Published as a conference paper at ICLR 2025
Table 1: Experimental setups of the Minecraft Al agents. IL is short for imitation learning.
Model
Controller
Ohservation
Video Demos
DECKARD (0033
RL + RI
Pixels & Inventory
Auto MC- Reward 2024a)
Pixels & GPS
Voyager 2024a)
GPT-4
Minecraft simulation Error trace
DEPS (0223)
IL.
Pixels & Yaw/pitch angle & GPS & Voxel
STEVE-1 (2023
Generative model
Pixels
VPT 2022
IL RL
Pixels
Dreamer 2023
RL
Pixels
AS-Imagine
RL
Pixels
shown in Table T existing Minecraft agents present notable distinctions in learning paradigms (i.e., controller), observation data, and the use of expert demonstrations.
The world model presented in this paper consists two main components: a short-term transition branch and long-term imagination branch. As result, employs complex notation system. We now introduce the key notations that will be frequently used throughout the paper:
Mr represents the affordance map.
Ct denotes the episode continuation flag.
jr the jumping flag that triggers jumpy state imaginations.
A, represents the number environmental steps between the jumpy transitions.
G₁ is the cumulative reward over A..
We use (0/, o4 M,, G. A, G2) to represent the simulated environment data that are used to train the long-term imagination branch of the world model. The policy is learned on trajectories of mixed long- and short-term imaginations {(S, â,, ft, Gt, Jt. Å., G.)}. where &, represents the latent state, and the variables predicted by the model are indicated using the superscript (*).
METHOD
3.1
OVERVIEW OF LS-IMAGINE
In this section, we present the details of LS- LS--mmgine, which involves the following algorithm steps, including world model learning, behavior learning. and environment interaction:
1. Affordance map computation (Sec. I.2.Ip: We employ sliding bounding box to individual images and execute continuous zoom-ins inside the bounding box, simulating consecutive video frames that correspond long-distance state transitions. We then create affordance maps by assessing the relevance of the fake video clips to task-specific goals expressed in text using the established MineCLIP reward model (Fan al. 2022).
Rapid affordance map generation (Sec. 3.2.2 Given that affordance maps will be frequently used in subsequent Step 5 to evaluate the necessities for jumpy state transitions, we train a U-Net module to approximate the affordance maps annotated in Step 1 for the sake of efficiency.
World model training (Sec. 3 We train the world model capture short- and long-term state transitions, using replay data with high responses from the affordance map. Each trajectory from the buffer includes pairs of samples from both adjacent time steps and long-distance intervals.
Behavior learning (Sec. 3.4): We perform an actor-critic algorithm to optimize the agent's policy based on finite sequence of long short-term imaginations generated by the world model.
5. Data update: We apply the agent to interact with the environment and gather new data. Next, we leverage the generated affordance map to efficiently filter sample pairs suitable for long-term modeling, incorporating both short- and long term sample pairs update the replay buffer.
6. Iterate Steps 3-5.
Below, we discuss each training step in detail. The full algorithm can be found in Appendix C.4
3.2 AFFORDANCE MAP AND INTRINSIC REWARD
We generate affordance maps using visual observations and textual task definitions to improve the sample efficiency of model- -ased reinforcement learning in open-world tasks. The core idea is to di- rect the agent's 's attention task-relevant areas of the visual observation, leading to higher exploration
3

---

Published as a conference paper at ICLR 2025
env
Cut tree
fotore vieen Ioneprostt the
miblee opplonaion
Mutrrdd
MineCLIP
Corr - 0.60
Corr: 0.90
Corr: 0.15
Maz6n, R)
agent
(b) Rapid affordance map generation Figure 2: The two steps for on-the-fly fly affordance map estimation: (a) Simulate exploration via image zoom-in and calculate the task-correlation scores of the virtual explorations using MineCLIP. (b) Learn to generate affordance maps more efficiently using multimodal U-Net.
(a) Affordance map calculation
efficiency. Let (w, h) be the affordance map that represents the potential exploration value at pixel position (w, h) on the image observation oz, given textual instruction I (e.g., "cut a tree"). The affordance map highlights the relevance between regions of the observation and the task description, serving as a spatial prior that effectively directs the agent's exploration toward areas of interest.
3.2.1 AFFORDANCE MAP COMPUTATION VIA VIRTUAL EXPLORATION
To create the affordance map, as shown in Figure 2(a) we simulate and evaluate the agent's exploration without relying on real successful trajectories. Concretely, we first adopt random agent to interact with task-specific environments for data collection. Starting with the agent's observation or time step 4, we use sliding bounding box with dimensions scaled to 15% of the observation's 's width and height to traverse the entire observation from left right and top to bottom. The sliding bounding box moves horizontally and vertically in 9 steps, respectively, covering every potential region in both dimensions. For each position on the sliding bounding box of the observation os, we crop 16 images from 0z, These cropped images narrow the field of view to focus on the region and are resized back to the original image dimensions. These resized images are denoted as x (where 0 < k < 16). The ordered set X₂ [of k 0, 15] represents a sequence of 16 frames simulating the visual transition as the agent moves towards the position specified by the current sliding bounding box. Subsequently, images, simulating the virtual exploration process, and the task description I. In this way, we quantify the affordance value of the sliding bounding box, indicating the potential exploration interest of the area. After calculating the correlation score for each sliding bounding box, we fuse these values to obtain smooth affordance map Mo,. For pixels that are covered by multiple sliding bounding boxes due to overlapping regions, the integrated affordance value is obtained by averaging the values from all the overlapping windows.
employ the MineCLIP mode!" to calculate the correlation between the X, of
3.2.2 MULTIMODAL U-NET FOR RAPID AFFORDANCE MAP GENERATION
The annotation of affordance maps, previously described, involves extensive window traversal and computations for each window position using pre-trained video-text alignment model. This method computationally demanding and time-consuming, making real-time applications challenging. To address this issue, we first use a random agent to interact with the environment for data collection. Next, we annotate the affordance maps for the collected images using the aforementioned method based on virtual exploration. We gather a dataset of tuples (o,, I, I.Mo,.) and use it to train multimodal U-Net based on Swin-Unet (Cao et al., 2022). To handle multimodal inputs, we extract text features from the language instructions and image features from the dow downsampling process of Swin-Unet, and fuse them with multi-head attention. We present architecture details in Figure / in the appendix. In this way, with the pretrained multimodal U-Net, maps at each time step using visual observations and language instructions.
efficiently generate affordance
IIP (ann et al. compute the correlation between text string and a 16-frame video segment.
2022 pretrains video-language representation using Minecraft videos, enabling it
4

---

Published as a conference paper at ICLR 2025
3.2.3 AFFORDANCE- DRIVEN INTRINSIC REWARD
To leverage the task-relevant prior knowledge presented by the affordance map for efficient exploration in the open world, we introduce the following intrinsic reward function:
W MΣΣ h=1
Mw, h) ((w, ht),
WH
where W and H denote the width and height of the visual observation. 9 represents a Gaussian matrix with dimensions matching those of the affordance map. corresponds 2D Gaussian distribution, with its peak located at the center of the affordance map. The values in the matrix are determined by standard deviations (T while the mean is uniformly set across the entire matrix. We present visualizations of G and conduct hyperparameter analyses on (Øa, Ty) in Appendix D.6 The intuition behind this design is to encourage the agent to move toward the target.
Overall, the agent receives composite reward consisting of the episodic sparse reward from the environment, the reward from MineCLIP (Fan et al. 2022), and the intrinsic reward from the affordance map:
p + Ft MineCLLPP + art where & is a hyperparameter. In contrast to the MineCLIP reward, which relies on the agent's 's past performance, our affordance-driven intrinsic reward emphasizes long-term values derived from future virtual exploration. It encourages the agent to adjust the policy to pursue task-related targets when they appear in its view, ensuring these targets centrally positioned in future visual observations to maximize this reward function.
3.3 LONG SHORT-TERM WORLD MODEL
3.3.1 LEARNING JUMPING FLAGS
In ISS-mmagee,, the world model is customized for long-term and short-term state transitions. decides which type of transition to adopt based on the current state and predicts the next state with the selected transition branch. To facilitate the switch between long-term and short-term state transitions, we introduce jumping flag jr. which indicates whether a jumpy transition or long-term state transition, should be adopted at time step t. When distant task-related target appears in the agent's observation, which can be reflected by a higher kurtosis in the affordance map, jumpy transition allows the agent to imagine the future state of approaching the target. To this end, we define relative kurtosis K2 which measures whether there are significantly higher target areas than the surrounding areas in the affordance map, and absolute kurtosis Ka represents the confidence level of target presence in that area. Formally,
W Σ
( (Mo,i(w,)) - -moan(Mau))
K, WH w=1h- hml
std(Mo..)
K₂ max(Mo,r) mean(M..r).
To normalize the relative kurtosis, we apply the sigmoid function to it, and then multiply it by the absolute kurtosis to calculate the jumping probability:
Prump sigmoid(K,) x K
(3)
The jumping probability measures the confidence in the presence of task-relevant targets far from the agent the visual observation. To determine whether to employ long-term state transition, we use a dynamic threshold, which is the mean of the collected jumping probabilities at each time step, plus one standard deviation. For a detailed explanation, please refer to Co If Pjump exceeds this threshold, the jump flag je is True and the agent switches to jumpy state transitions in the imagination phase.
3.3.2 LEARNING JUMPY STATE TRANSITIONS
In LS- Imagine, the state transition model includes both short-term and long-term branches. As shown in Figure (a), the short-term transition model integrates the previous deterministic recurrent state h4-1, stochastic state 2;-1, and action 4;-1 to adopt the single-step transition. In contrast, the long-term branch simulates jumpy state transitions toward the target. It is important to clarify that the index t does not denote the time step during real environmental interactions but instead represents the positional order of states within the imagination sequence. The overall world model of LS-Ima is
5

---

Published as a conference paper at ICLR 2025
predictor
Jump predictor
predictor
Short-Term Branch
Short- Terrt Imagination
8G, Mr
Mr
Actor
Pavl
Mhs XL.
op-1
Av
h--1
Av-1
å/. 6
f-11
F/1 &
Cmii
Oies
Maa
Artor
enc
8/, M2
Fiyt
M2
Long-Term Branch
Long-Term Imagination
(a) World model of LS-Imagine
(b) Behavior learning in latent imagination
Figure 3: The overall architecture of the world model and the behavior learning process.
primarily based on DreamerV3 (Hafner et al. 2023 with novel components specifically designed to capture jumpy state transitions:
Short-term transition model:
h. = fo 21-1,
Long-term transition model:
W a0(1-1,2-
Encoder:
2, g0(2) Or,
Dynamics predictor:
2, ~ p hha)
Reward predictor:
Pr. é. po Ct lhe, 26)
Decoder:
ō M, ~ ps My he,
Jump predictor:
3. po he,
Interval predictor:
2t Å G p G/ 14-1, 11,
At time step we feed the recurrent state f4p the observation or, and the affordance map M, into the encoder to obtain posterior state ž₁. We also use the affordance map as an input of the encoder, which serves as the goal-conditioned prior guidance the agent. Notably, the prediction of prior state & does not involve the current observation or affordance map, relying solely on historical information. We use (h,, 2x) to reconstruct the visual observation ò and the affordance map Mr, and predict the reward fr. episode continuation flag é and jumping flag jr. For long-term state transitions, we use an interval predictor to estimate the expected number of interaction steps A required to transition from the pre-jump state (/t,1, 1t- to the post-jump state (M,, 2(), along with the expected cumulative reward G² that the agent may receive during this time interval. We detail the approach to annotate A and G², using the real interaction data Appendix [.I
We collect short-term tuples D: (ot, at, Mr, Ft, Je. At, Gt) from each interaction with the environment using the current policy. When observing ji = 1, we additionally construct long-term tuples D/+1 = (o[+1) a'41: 1 1 6441:1141, A/+ Gi(+1) based on D.. More details for this process can be found in Appendix C.I During representation learning, we sample short-term tuples (D.)T= and the long-term tuples following jumpy transitions {D/+ from the replay buffer B, where T denotes the set of time steps at which long-term state transitions are required. The loss functions for each component of the short-term and long-term world model branch are as follows:
Layn max (1, KL [sg (gs (Ze he, ot, Mc)) |l Pa (2x he)l
Lene = max (1, KL [gs (Et Me, || sg (po (z. ho)l)
Short-term branch:
(5)
Lake In Pa (or, Me h4, 21)
Sprod - In Po (Te, Ct hey 2x) In ps (je he, zx)

---

Starting from the initial state encoded from the sampled observation and the affordance map, we dynamically select either the long-term transition model or the short-term transition model to predict subsequent states based on the jumping flag jt. For the long short-term imagination sequence {(e, at) ) with an imagination horizon of L, we predict reward sequence ÂI.L and the continuation flag sequence ÊL.Z through the reward predictor. Similar to Eq. { the index t does not represent the time step in the environment, but rather the positional order of the states in the imagination sequence. Published as a conference paper at ICLR 2025
Cayn = max (1, KL (sg (as (2l M1. of M,)) Po (z4 k)D
f'enc = max (1, KL KL 190 (2t 11 M,) ) (po (2) k(DID)
Long-term branch:
f'Lec = In pa (o 1 24)
(6)
fipred = In polri, d 114 244 In pol' 11f =tt)
Cim = --In In pa (A,, G2 2 h1 24)
We can optimize the world model W& by minimizing over replay buffer B:
L = E D(D.T CSyyn BencCene + Specd (Laee + Epred)) +
Brong ExP (dy + Benc + Speed (E'tec + Eirest + )
3.4 BEHAVIOR LEARNING OVER MIXED LONG SHORT-TERY IMAGINATIONS
As shown in Figure 3 (b), LS- -maainee employs an actor-critic algorithm to learn behavior from the latent state sequences predicted by the world model. The goal of the actor is to optimize the policy to maximize the discounted cumulative reward Rg while the role of the critic is to estimate the discounted cumulative rewards using the current policy for each state & (h,, 2:
Actor: āg ~ T (āz 8z), Critic:
(Rt, 8:)
Specifically, starting from state 8, any subsequent state obtained via either short-term transition or a long- term transition is indexed sequentially as + 1.
For jumpy states predicted by long-term imagination, the interval predictor estimates (i) the number of steps Å, from 6 to 8, and (ii) the potential discounted cumulative reward G2 over the time interval of Á, Otherwise, for states obtained via short-term imagination, which correspond to single-step transitions in the environment, we set Å, 1 and G, = Å.. Consequently, within one imagination episode, we obtain sequence of step intervals Â and sequence of predicted rewards G2.L between consecutive imagination states.
We employ modified bootstrapped X-returns that considers both long-tern and short-term imagina- tions to calculate the discounted cumulative rewards for each state:
R
GfGG4 + Ar+i [(1 (8 + AR (5
if TikK L ift=
The critic uses the maximum likelihood loss to predict the distribution of the return estimates R.
C(v) = Σ Inp (R² | (a)
(10)
1=1
Following DreamerV3 (Hafner et al., 2023), we train the actor to maximize the return estimates RAA Notably, since long- long-term term imagination does not involve actions, do not optimize the actor at time steps when jumpy state transitions are adopted. Therefore, unlike DreamerV3, we apply an additional factor of (1 11-18 je) to ignore updates at long-term imagination steps:
Σ sg (1 (1-j.)
I.
R² 2g max S) log (ā 5t) + t (o(a 6c).
((0) =
(11)
7

---

Published as a conference paper at ICLR 2025
VPT
STEVE-1
PTGM
Director
DreamerV3
LS IS-mmagine
10
1.0
0.8
0.8
Suicess Rate
Success Rate
0.44
9.4
0.2
0.0
0.0
Environment steps
Environment steps
Environment steps
(a) Harvest log in plains
(b) Harvest water with bucket
(c) Harvest sand
0.4
0.6
0.3
MineCLIP Score
0.8
Success Rate
0.6-
0.4
0.2
0.0
log
water
sand
wool iioo_roe ore
Environment steps
Environment steps (e) Mine iron ore
(d) Shear sheep
(1) Comparison in MineCLIP reward
Figure 4: Comparison of L.S--SSSSL- Imagine against strong Minecraft agents, including DreamerV3 (Hafner et al. 2023), VPT (Baker et al. 2022), STEVE-I (Lifshitz et al. (023). PTGM (Yuan et al. (2024). and Director (Hafner et al. 2022). We present the numerical results Table 3aleein iin the appendix.
EXPERIMENTS
We explore LS Imagine on the challenging MijeDojo (Fan al. 2022) benehmark on top of the popular Minecraft game, which is comprehensive simulation platform with various open-ender tasks. We use tasks, i.e,, harvest log in plains, harvest water with bucket, harvest sand, shear sheep, and mine iron ore. These tasks demand numerous steps to complete and present significant challenges for agent learning. We adopt a binary reward that indicates whether the task was completed, along with the MineCLIP reward (Fan al. 2022b. Further details of the environmental setups are provided in Appendix A Besides, we introduce the compared models in Appendix
Implementation details. We conduct our experiments on the MineDojo environment, where both visual observation and corresponding affordance maps are resized to 64 x 64 pixels. To generate accurate affordance maps, we collect 2, 2,000 images from the environment using random agent under the current task instruction and generate discrete set of (o,, T, M.,...., which are then used to finetune the multimodal U-Net for 200 epochs. For tasks in the MineDojo benchmark, we train the agent for 1 x 10 environment steps. Each training of LS Imagine takes approximately 23 GB of VRAM and requires around 1.7 days to complete on single RTX 3090 GPU.
4.1
MAIN COMPARISON
We evaluate all the Minecraft agents in terms of success rate shown in Figure and per-episode steps shown in F We find that LS-Imagine significantly outperforms the compared models, particularly in scenarios where sparse targets are distributed in the task. In Figure (f), showcase the MineCLIP values achieved by LS- Im--mne and DreamerV3. Specifically, a sliding window of length 16 used compute the local MineCLIP values for each segment. The mean value is then calculated from all sliding windows. We can see that agents trained using our method achieve higher MineCI.IP values within a single episode compared DreamerV3. This suggests that I--Imne facilitates quicker detection of task-relevant visual targets in open-world world environments.
Additionally, we present qualitative results in Figure 6(a) In the top row, we decode the latent states before and after the jumpy state transitions back to the pixel space. To better understand how affordance maps facilitate the jumpy state transitions and whether they can provide effective goal- conditioned guidance, the bottom rows visualize the affordance maps reconstructed from the latent states. These visualizations demonstrate that the proposed world model can adaptively determine when to utilize long-term imagination based the current visual observation. Furthermore, the
8

---

Published as a conference paper at ICLR 2025
STEVE-1
PTGM
Director
DreamerV3
LS I--amaiine
2000
Strpp pttppper- 4pts00e
900
600
1900
0014
600
200
100
1700
as
24
D2
15000
Mine steps 168
Harvest log plains
Harvest with bucket
Harvest
(d) Sherr shoop
Figure 5: The number of steps per episode for task completion.
magixttion
longtem Imaginttinn
1,0
Predebi
0.4-
0.2
USimagine (xeries)
LSimaghe (parale))
0.0
steps
(a) Visualization of long short-term imaginations
(b) Series vs. parallel connections
Figure 6: (a) Visualization of long short-term imaginations and (b) further discussion possible
architecture designs of Series and Parallel connections of these two imagination pathways.
generated affordance maps align effectively with areas that highly relevant to the final goal, thereby enabling the agent to perform more efficient policy exploration.
4.2 MODEL ANALYSES
Ablation studies. We conduct the ablation studies to validate the effect of the affordance-drrven intrinsic re- ward and long short-term imagination. Figure presents corresponding results in the challenging MineDojo tasks. As shown by the blue curve, removing the long-term imagination of LS-Imagine leads to a performance de- cline, which indicates the necessity introducing long- term imagination and switching between it and short- term imagination adaptively. For the model represented by the green curve, we do not employ affordance-driven intrinsic reward. shows that the affordance-driven in- trinsic reward also plays important role during the early training stage of agents. Additionally, unlike the MineCLIP reward being calculated based on series of states, the affordance-ding driven intrinsic reward relies solely on a single independent state. This approach enables accurate estimation of the reward for the post-jumpy -raaassmtttaanitiiionn state.
Success rate
0.6
0.4
LS-Imagine
0.2
w/o long-term imagination
w/o intrinsic reward
0.0
0.2
0.6 Environment steps
1.0
Figure 7: Ablation study results.
Alternative pathways of mixed imaginations. is worth highlighting that the long short-term imagination implemented sequentially. In Figure [O(a] in the appendix, we provide visualization illustrating how the agent sequentially performs short-term and long-term imaginations within single sequence. Alternatively, illustrated in Figure [O(b] we could structure long- and short-term imagination pathways in parallel. Specifically, we begin by applying short-ters term imagination within single sequence. For each predicted state, we examine the jumping flag: If }r 1, we initiate new imagination sequence starting from the post-jump state, which predicted by the long-term transition model and the dynamics predictor. In other words, whenever long-term state jump occurs, the world model generates 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 conduct sequences. To evaluate the advantages of using sequential long short-term imagination, an experimental comparison between LS- Imagine (series) and LS I.S-mne (parallel). Figure 6(b) shows that the I---mneene (series) outperforms Im-gine (parallel) by large margins. This implies

---

Published as a conference paper at ICLR 2025
that the parallel imagination sequences are independent of one another, meaning that the sequence starting with a post-jumping state does not guide the prior-jumping transitions.
In the appendix, further include (i) experiments the long-horizon "Tech Tree" task, (ii) analyses of the long-term imagination frequency and corresponding state jumping intervals å predicted by the model, and (iii) visualization of affordance maps with occluded target objects.
RELATED WORK
Visual MBRL. Recently, learning control policies from images, L.e,, visual RI. has been used widely,
whereas previous RL algorithhs learn policies from low-dimensional states. Existing approaches
can be grouped by the use of model-free RL methods (Laskin et al. 2020, |Schwarzer et al. 2021
Stooke et al., 2021 Xiao et al.] 2022, Parisi et al. [0022; Yarats al. [2022] Zheng et al. 2023
or model-
RL methods (Hatner et al. 2019; 2020; 20 Seo et al. |0222] Pan et al. 2022
Zhang et al. 2023a Mazzaglia al.) 2023 Micheli al. 2023] Zhang et al. 2023b Ying et al.
|023 Seo et al.]]223)) Alonso et al.] [224]] Hansen et al. [024:] Wang et al.] 2024b). The following
methods specifically enhance the modeling of long-term dynamics in visual MBRL. [ee et (2024b)
proposed the prediction of temporally smoothed rewards to address long-horizon sparse-reward tasks. R2I (Samsami et al., 2024) improves long-term memory and long-horizon credit assignment in MBRL. Unlike existing methods, our work presents a long short-term world model architecture specifically designed for visual control in open-world environments.
Affordance maps for robot learning, Our work is also related to the affordance map for robot
learning (Mo al. [221]] Jiang et al. [202]] Yarats et al. [202]] Mo al. [0022] Geng al., 2022
Xu et al. 2022a; Wang et al.) 2022 Wu et al. 2022, Ha & Song 2022, Xu et al. 2022b Cheng
et al. 2024: Lee et al. 2024a Li et al. 2024b). Where2Fxplore Ning et al. (023) introduces a
cross-category fhw-shot affordance learning framework that leverages the similaria Maritesi in geometries across different categories. DualAfford (Zhao et al., 2023) learns collaborative actionable affordance for dual-gripper manipulation tasks over various 3D shapes. VoxPoser (Huang et al. 2023) unleashes the power of large language models and vision- language models for extracting affordances and constraints of real-world manipulation tasks, which are grounded in 3D perceptual space. VRB (Bahl) eet al. 2023) trains a visual affordance model with videos of human interactions and deploys the model in real-world robotic tasks directly. Qi al. (2020) adopted spatial affordance map that is trained by interacting with the environment for navigation. However, our approach distinguishes itself by employing visual observation to generate affordance maps as guidance mitigate the low exploration efficiency in open-world environments.
Hierarchical methods. Like our approach, Director (Hafner et al., 2022 learns hierarchical behaviors in the latent space, which adopts high-level policy (manager) to produce latent goals to guide low-level policy (worker). Dr. Strategy (Hamed et al.) 2024) proposes strategic dreaming with latent landmarks to learn highway policy that enables the agent to move to landmark in the dream. Gumbsch et al. (2024) presented hierarchy of world models, which perform high-level and lowelevel prediction adaptively, and the high-level predictions depend on the low-level predictions. Our method distinguishes itself by generating affordance maps through image zoom-in to encourage the agent to explicitly execute long-term imagination in the world model.
6 CONCLUSIONS AND LIMITATIONS
In this paper, we presented novel approach to overcoming the challenges of training visual rein- forcement learning agents in digh--imensional open worlds. By extending the imagination horizon and leveraging a long short-term world model, our method facilitates efficient off-policy exploration across expansive state spaces. The incorporation of goal-conditioned jumpy state transitions and affordance maps allows agents to better grasp long-term value, enhancing their eeciiioonmmaking abilitics. Our results demonstrate substantial improvements over existing state-of-the-an techniques in MineDojo, highlighting the potential of our approach for open-world reinforcement learning and inspiring future research in this domain.
A limitation of LS- Imagine is the computational overhead it introduces. Additionally, its effectiveness has only been validated in 3D navigation environments with embodied agents. We aim to enhance the generalization of our approach across a wider range of tasks.
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
In this work, we are committed to upholding ethical research practices. This work does not involve human subjects, personal data, sensitive information. All environments and datasets used are synthetic and publicly available. We recognize the potential for reinforcement learning models to be misused, particularly in decision-making scenarios where unintended outcomes could arise. To mitigate these risks, we emphasize responsible deployment and encourage careful consideration of the broader impact of such systems, restricting the use our work strictly to research purposes.
REPRODUCIBILITY STATEMENT
We prioritize the reproducibility of our work. All results can be reproduced on publicly available RL environments by following the experimental details presented in Sec. and Appendix D.6 We provide the source code at
com wttp://wwwwwwww///////// /LSS- magine
ACKNOWLEDGMENTS
This work was supported by the National Natural Science Foundation of China (Grants 62250062, 62302246), the Smart Grid National Science and Technology Major Project (Grant 2024ZD0801 1200). the Shanghai Municipal Science and Technology Major Project (Grant 202 2021SHZDZX0100 Fundamental Research Funds for the Central Universities, and the CCF-Tencent Rhino- Bird Open Research Fund. Additional support provided by the Natural Science Foundation of Zhejiang Province, China (Grant LQ23F010008) the High Performance Computing Center at Eastern Institute of Technology, Ningbo, and Ningbo Institute of Digital Twin.
the
REFERENCES
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in atari. In NeurlPS, 2024.
Shikhar Bahl, Russell Mendonca, Lili Chen, Unnat Jain, and Deepak Pathak. Affordances from human videos as versatile representation for robotics. In CVPR, pp. 13778-13790, 2023.
Bowen Baker, Ilge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. In NeurlPS, 2022.
Hu Cao, Yueyue Wang, Joy Chen, Dongsheng Jiang, Xiaopeng Zhang, Qi Tian, and Manning Wang. Swin-unet: Unet-like pure transformer for medical image segmentation. In ECCVW, 2022.
Guangran Cheng, Chuheng Zhang, Wenzhe Cai, Li Zhao, Changyin Sun, and Jiang Bian. Empow- cring large language models on robotic manipulation with affordance prompting. arXiv preprint arXiv:2404. / 1027, 2024.
Linxi Fan, Guanzhi Wang, Yunfan Jiang, Ajay Mandlekar, Yuncong Yang. Haoyi Zhu, Andrew Tang, De-An Huang, Yuke Zhu, and Anima Anandkumar, Minedojo: Building open-ended embodied agents with internet-scale knowledge. In NeurIPS, volume 35, pp. 18343-18362, 2022.
Yiran Geng, Boshi An, Haoran Geng, Yuanpei Chen, Yaodong Yang. and Hao Dong. End-to-end affordance learning for robotic manipulation. arXiv preprint arXiv:2209.12944 2022.
Christian Gumbsch, Noor Sajid, Georg Martius, and Martin V Butz. Learning hierarchical world models with adaptive temporal abstractions from discrete latent dynamics. In ICI.R 2024.
Huy Ha and Shuran Song. Flingbot: The unreasonable effectiveness of dynamic manipulation for cloth unfolding. In CoRL, pp. 24-33, 2022.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In ICMI. pp. 2355-2565, 2019.
11

---

Published as a conference paper at ICLR 2025
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In ICLR, 2020.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. ICLR, 2021.
Danijar Hafner, Suang-Huei Lee, Ian Fischer, and Pieter Abbeel. Deep hierarchical planning from pixels. NeurIPS, 2022.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04 2023.
Hany Hamed, Subin Kim, Dongyeong Kim, Jaesik Yoon, and Sungjin Ahn. Dr. strategy: Model-based generalist agents with strategic dreaming. In ICML, 2024.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpe2: Scalable, robust world models for continuous control. In ICLR, 2024.
Wenlong Huang. Chen Wang, Ruohan Zhang, Yunzhu Li, Jiajun Wu, and Li Fei-Fei. Voxposer: Composable 3d value maps for robotic manipulation with language models. In CoRL, 2023.
Zhenyu Jiang, Yifeng Zhu, Maxwell Svetlik, Kuan Fang. and Yuke Zhu. Synergies between affordance and geometry: 6-dof grasp detection via implicit representations, Robotics: science and systems, 2021.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In ICML, pp. 5639-5650, 2020.
Olivia Y Lee, Annie Xie. Kuan Fang, Karl Pertsch, and Chelsea Finn. Affordance-guided reinforce- ment learning via visual prompting. arXiv preprint urXiv:2407.0334, 2024a.
Vint Lee, Pieter Abbeel, and Youngwoon Lee. Dreamsmooth: Improving model-basee bbsee reinforcement learning via reward smoothing. In ICLR, 2024b.
Hao Li, Xue Yang, Zhaokai Wang, Xizhou Zhu, Jie Zhou, Yu Qiao, Xiaogang Wang, Hongsheng Li, Lewei Lu, and Jifeng Dai. Auto me-reward: Automated dense reward design with large language models for minecraft. In CVPR, pp. 16426-16435, 2024a.
Xiaoqi Li, Mingxu Zhang, Yiran Geng. Haoran Geng, Yuxing Long, Yan Shen, Renrui Zhang, Jiaming Liu, and Hao Dong. Manipllm: Embodied multimodal large language model for object-centric robotic manipulation. In CVPR, pp. 18061-18070, 2024b.
Shalev Lifshitz, Keiran Paster, Harris Chan, Jimmy Ba, and Sheila Mcllraith. Steve-1: A generative model for text-to-behavior in minecraft. In NeurlPS, 2023.
Pietro Mazzaglia, Tim Verbelen, Bart Dhoedt, Alexandre Lacoste, and Sai Rajeswar. Choreographer: Learning and adapting skills in imagination. In ICLR, 2023.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. In ICLR, 2023.
Kaichun Mo. Leonidas J Guibas. Mustafa Mukadam, Abhinav Gupta, and Shubham Tulsiani. WhereZact: From pixels to actions for articulated 3d objects. In ICCV, pp. 6813-6823, 2021.
Kaichun Mo, Yuzhe Qin, Fanbo Xiang, Hao Su, and Leonidas Guibas. O2o-afford: Annotation-free large-scale object-object affordance learning. In CoRL, pp. 1666-1677, 2022.
Chuanruo Ning, Ruihai Wu, Haoran Lu, Kaichun Mo, and Hao Dong- Where2explore: Few-shot affordance learning for unseen novel categories of articulated objects. NeurlPS, 2023.
Kolby Nottingham, Prithviraj Ammanabrolu, Alane Suhr, Yejin Choi, Hannanch Hajishirzi, Sameer Singh, and Roy Fox. Do embodied agents dream of pixelated sheep: Embodied decision making using language guided world modelling. In ICMI., pp. 26311-26325 1-26325, 2023.
12

---

Published as a conference paper at ICLR 2025
Minting Pan, Xiangming Zhu, Yunbo Wang, and Xiaokang Yang. Iso-dream: Isolating and leveraging noncontrollable visual dynamics in world models. In NeurIPS, volume 35, pp. 2022.
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsurprising effectiveness of pre-trained vision models for control. In ICML, pp. 17359-17371, 2022.
William Qi, Ravi Teja Mullapudi. Saurabh Gupta, and Deva Ramanan. Learning to move with affordance maps. In ICLR, 2020.
Mohammad Reza Samsami, Artem Zholus, Janarthanan Rajendran, and Sarath Chandar. Mastering memory tasks with world models. In ICLR, 2024.
Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R De- von Hjelm, Philip Bachman, and Aaron C Courville. Pretraining representations for data-efficient reinforcement learning. In NeurIPS, volume 34, pp. 12686-12699, 2021.
Younggyo Seo, Kimin Lee, Stephen L James, and Pieter Abbeel. Reinforcement learning with action-free pre-training from videos. In ICML, pp. 19561-19579, 2022.
Younggyo Seo, Junsu Kim, Stephen James, Kimin Lee, Jinwoo Shin, and Pieter Abbeel. Multi-view masked world models for visual robotic manipulation. In ICML, pp. 30613-30632, 2023.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In ICML, pp. 9870-9879, 2021.
Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language models. TMLR, 2024a.
Qi Wang, Junming Yang, Yunbo Wang, Xin Jin, Wenjun Zeng, and Xiaokang Yang. Making offline rl online: Collaborative world models for offline visual reinforcement learning. In NeurIPS, 2024b.
Yian Wang, Ruihai Wu, Kaichun Mo, Jiaqi Ke, Qingnan Fan, Leonidas Guibas, and Hao Dong. Adaafford: Learning adapt manipulation affordance for 3d articulated objects via few-shot interactions. In ECCV, pp. 90-107, 2022.
Zihao Wang, Shaofei Cai, Guanzhou Chen, Anji Liu, Xiaojian Ma, and Yitao Liang. Describe, explain, plan and select: Interactive planning with large language models enables open-world multi-task agents. In NeurIPS, 2023.
Ruihai Wu, Yan Zhao, Kaichun Mo. Zizheng Guo, Yian Wang, Tianhao Wu, Qingnan Fan, Xuelin Chen, Leonidas Guibas, and Hao Dong. Vat-mart: Learning visual action trajectory proposals for manipulating 3d articulated objects. ICLR, 2022.
Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv preprint (tt.:220330677.. 2022.
Chao Xu, Yixin Chen, He Wang, Song-Chun Zhu, Yixin Zhu, and Siyuan Huang. Partafford:
Part-level affordance discovery from 3d objects. arXiv preprint a3Xiv:2202.335/9 2022a.
Zhenjia Xu, Zhanpeng He, and Shuran Song. Universal manipulation policy network for articulated objects. IEEE robotics and automation letters, 72):2447-2454, 2022b.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In ICLR, 2021.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. In ICLR, 2022.
Chengyang Ying, Zhongkai Hao, Xinning Zhou, Hang Su, Songming Liu, Jialian Li. Dong Yan, and Jun Zhu. Reward informed dreamer for task generalization reinforcement learning. arXiv preprint ar arXiv:2303.0500 2023.
13

---

Published as a conference paper at ICLR 2025
Haoqi Yuan, Zhancun Mu, Feiyang Xie, and Zongging Lu. Pre-training goal-based models for sample-effcient reinforcement learning. In ICLR, 2024.
Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. Storm: Efficient stochastic transformer based world models for reinforcement learning. In NeurIPS, 2023a.
Wendong Zhang, Geng Chen, Xiangming Zhu, Siyu Gao, Yunbo Wang. and Xiaokang Yang. Predictive experience replay for continual visual control and forecasting. arXiv preprint 0665: 2023b.
Yan Zhao, Ruihai Wu, Zhehuan Chen, Yourong Zhang, Qingnan Fan, Kaichun Mo, and Hao Dong. Dualafford: Learning collaborative visual affordance for dual-gripper manipulation. ICLR, 2023.
Ruijie Zheng, Xiyao Wang, Yanchao Sun, Shuang Ma, Jieyu Zhao, Huazhe Xu, Hal Daumé III, and Furong Huang. Taco: Temporal latent action-driven contrastive loss for visual reinforcement learning. NeurIPS, volume 36, 2023.
14

---

Published as a conference paper at ICLR 2025
APPENDIX
A ENVIRONMENT DETAILS
As illustrated in Table language description employed for calculating the MineCLIP reward (Fan al. 2022). Initial tools are the items provided in the inventory at the beginning of each episode. Initial mobs and distance specifies the types of mobs present the start of each episode and their initial distance from the agent. Max steps refers to the maximum allowed steps per episode.
Table 2: Details of the MineDojo tasks.
Task
Language description
Initial tools Initial mobs and distance Max steps
Harvest log in plains
"Cut tree."
1000
Harvest water with bucket
"Obtain water."
bucket
1000
Harvest sand
"Obtain sand."
1000
Shear sheep
"Obtain wool."
shear stone pickaxe
sheep.
1000
Mine iron ore
"Mine iron ore."
2000
B COMPARED METHODS
We compare LS- -S--magine with strong Minecraft agents, including:
DreamerV3 (Hafer et al. 2023): An MBRL approach that learns directly from the step-by-step imaginations of future latent states generated by the world model.
VPT (Baker et al. )022): foundation model designed for Minecraft trained through behavior cloning, on a dataset consisting of 70,000 hours of game playing collected from the Internet.
STEVE-/ (Lifshitz et al., 2023): An instruction-follooing Minecraft agent that translates lan- guage instructions into specific goals. To evaluate its effectiveness, we assess Steve-1's zero-shot performance on our tasks by supplying it with task instructions.
Director (Hafner et al. 2022): An agent learns hierarchical behaviors by leveraging a world model plan within its latent space.
PTGM (Yuan et al., (024)) An RL method that pretrains goal-based policy and adopts temporal abstractions and behavior regularization.
C MODEL DETAILS
C.I ENVIRONMENTAL INTERACTION AND DATA COLLECTION
To train LS- J.--mmgisn's 's world model, collect both short-term and long-term transition data through interactions with the environment. As shown Figure(8] at each time step t, the agent interacts with the environment following the current policy. At each time step, the data buffer collects tuple D., which includes (or, Mr, Tt, Jt, A., G.):
represents the observed image.
represents the agent's is action taken given Oz-
Mr is the affordance map generated by multimodal U-Net given or and task instructions
r, defined Sec. 3.2.3 which is the immediate reward computed a weighted sum of the sparse environmental reward r. after executing 4,-1, the MineCLIP reward r MineCTIPP from pretrained scoring model (Fan et al. 2022). and the intrinsic reward intr defined in Eq. I and based on M,
ca the continuation flag received from the environment, which indicates whether further interac- tion is required after this step.
je is the jumping flag, which is used to train the world model to trigger long-term imagination during model-based behavior learning. We first estimate the jumping probability Prump using Eq. (3) based on M, To stabilize training, we establish a dynamic threshold Pitresh which accounts for the varying guidance strength provided by the affordance map across different tasks, resulting in
15

---

Published as a conference paper at ICLR 2025
env
A
C2
F
Multimodal U-Net
Multimodal U-Net
Muttimodal U-Net
Je
f
M
M
Masa
A
Alua
Gest
Pyip
Patandh
Pramp Piteme Bent False
P Prmpp' Painee False
True
agent
Figure 8: Environmental interaction and data collection.
task-specific distributions of Pjup-- Specifically, from the beginning of training. we store the Pjump values for every interaction step in dedicated buffer. The threshold Pawesh is then dynamically calculated as the mean of all fjump values currently in the buffer plus their standard deviation. This dynamic adjustment ensures that the threshold adapts to the characteristics of the task and remains robust throughout training. If Pjump >
Pihresh- we set j: = 1; otherwise, jr 0.
A represents the expected number of step intervals in the jumpy state transitions during long-term imaginations. Specifically, we set A, = 1 by default, corresponding to short-term transition.
G, represents the expected cumulative reward between the pre- and post-jump states when long- term imagination occurs. Specifically, for short-term transition, we set G. Ft by default.
jt 0, D, defined as the starting point of short-term transition within the pair (D,, D(+1). During world model training, (D,, Dz+1) ) is replayed to train the related modules associated with short-term dynamics. Once we obtain J: during interactions, we define the current step as the starting point of simulated long-term transition (Dr, D/+1). Notably, we use D441 to differentiate from its short-term counterparts.
We define 1 (c++ d'+ M2 ( puted same manner as in short-term tuples but with 0/41 and M as inputs. Similarly, d'41 and j/+1 are also computed in the same way as in short-term tuples. We record them in the data buffer for better training of the reward predictor and the jump predictor.
1441 A, G.+.). G,41,G44
where r'+1 and c/ are com-
The next question is how to annotate A/+1. G/+1, and 0'42 to train the long-term branch.
de+1 i is simulated image rather than a real-captured image. is obtained by cropping the original observation oe based on the high-value regions in the affordance map Mr-
A/+1 is estimation of the number of real interaction steps between the pre-jump state and the post-jump state. Since the post-jump state is not real data obtained from the environment, we first identify real state that closely resembles the post-jump state. We then calculate the number of steps required to transition from the pre-jump state to this identified real post-jump state. Specifically, we use the intrinsic reward as a measurement. Starting from the pre-jump state, during subsequent interactions with the environment, if the agent reaches real state where the intrinsic reward satisfies intr (+A; 2 intr we take this state as the real post-jump state and take 2 as the long-term jumping interval.
G 1 is the cumulative reward within 4/41 interaction steps, i.e., G'+1 = E gi y
16

---

Published as a conference paper at ICLR 2025
Compute affordance map with MineCLIP
Cut a tree
Min E
Text Encoder
W
H
x
48
Patch Partition
Linear Projection
W H Class
Linear Embedding
Patch Expanding
Wx H x O(4x)
W
H
W H x xe
C
Swin Transformer
Swin Transformer
Block
Block
Patch Merging
Patch Expanding
8
x
20
Swin Transformer
TIA
Swin Transformer
8
x
20
Patch Merging
Patch Expanding
W
17
W
17 16
Transformer Black 2
Swin Transformer
18
16
TIA
Block
16
Patch Merging
Patch Expanding
W
8C'
Swin Transformer Black
Swin Transformer Block
x
801
32
32
TIA
32
32
Figure 9: The architecture of multimodal U-Net.
C.2 FRAMEWORK OF MULTIMODAL U-NET
As described in Sec. 3.2.2 we train a multimodal U-Net to rapidly generate affordance maps based on observation images and task instructions. Our enhanced multimodal U-Net architecture, illustrated in Figure is based on Swin-Unet (Cao et al.) 2022), a U-shaped encoder-decoder architecture built on Swin Transformer blocks. The enhanced multimodal U-Net consists of an encoder, decoder, bridge layer, and text processing module. In the Swin-Unet-inspired structure, the basic unit is the Swin Transformer block. For the encoder, the input image is divided into non-overlapping patches of size 4 X 4 to convert the input into a sequence of patch embeddings. Through this method, each patch has feature dimension of 4 4 x 48. The patch embeddings are then projected through linear embedding layer (denoted as C),, and the transformed patch tokens are passed through several Swin Transformer blocks and patch merging layers to produce hierarchical feature representations. The patch merging layers are responsible for downsampling and increasing the dimensionality, while the Swin Transformer blocks handle feature representation learning.
For the task instruction, the text description is processed through the text encoder of MineCLIP (Cao) (et al. 2022) to obtain text embeddings, which are integrated with the image features extracted at each layer of the encoder via the Text-Image Attention (TIA) module. The TIA module employs multi-head attention mechanism to fuse image features (as keys and values) with text features (as queries) in multi-scale attention based fusion. The resulting fused text-image features are passed through the bridge layer and are subsequently combined with the corresponding features during the upsampling process in the decoder.
The decoder comprises Swin Transformer blocks and patch-expanding layers. The extracted context features are combined through the bridge layer with the multi-scale text-image features from the encoder to compensate for the spatial information lost during downsampling and to integrate the text information. Unlike the patch merging layers, the patch expanding layers are specifically designed for upsampling. They reshape the adjacent feature maps by performing a 2 upsampling of the resolution, expanding the feature maps into larger ones. Finally, final patch expanding layer performs 4x upsampling to restore the resolution of the feature map to the input resolution W x H followed by a linear projection layer applied on the upsampled features to produce pixel-level affordance maps.
17

---

Published as a conference paper at ICLR 2025
61
6g
A
G
84
Ais,
Gu
015
81
82
3x
83
30
3+
85
Jis
814
3u
8
315
(a) L.-Imagine (series))
10
100
100
10
3
4
611
200
000
000
10
2
2
3
36
912
40M
81
16
(b) LS- LS-Imagine (parallel)
Figure 10: Comparison with series and parallel variants of mixed imaginations.
C.3 VARIANTS OF LONG SHORT-TERM IMAGINATIONS
We compare two alternative pathways of the long short-term imaginations in Figure 10
C.4 FULL ALGORITHM
We present the training pipeline of LS-Imagine in Algorithm 1
Algorithm The training pipeline of IS-Immgine.
1: Initialize parameters ø, 0, d.
2: Compute affordance map with MineCLIP.
> Affordance map generation
3: Train multimodal U-Net with annotated data. To enable real-time interaction with the affordance maps
4: Train random agent and collect replay buffer B.
5: while not converged do
6: Sample long short-term transitions from B.
Representation learning
7: Update the world model á using Eq. (7).
8: Generate (81, 01, j1) using Ne and Wg.
9:
for time step = 2 2- L do
> Behavior lcarning
jump flag jt-1 then
11:
Generate (s), dg de, fí, A,, Gi) using me and long-tern imagination of W..
Update (&z, ae, Ee, jr, Å., G.) +- (sh, db R. Ar, G)).
13:
else
14:
Generate (ă,, āz, Fr, če, J() using To and short-term imagination of Wo- Update (&e, âe, 8 Je, Å., G.) + (3,, az, J., fT).
15:
16:
end if
17:
end for
18:
Caleulate value estimate using Eq.
19:
Optimize actor using Eq. 1 over {(5t, âr, &, J., & Go)
20:
Optimize critic By using Fq. 10 over [(A,, à, G. 3r, A. G.)
21:
for time step ........ Sample â; no (ā. 8.)
Environment interaction
22:
ri 1.. G <- env ste.s))
24:
Generate affordance map M, with multimodal U-Net for each os. Calculate intrinsic reward r and jump flag je based on the affordance map.
25:
26:
Collect short-term data (or, Me, Jt, Ac, Gt).
27:
jumpy flag j. then Construct long-term data (o(+1, d Miss, P4a d4a Han Alti G(1)).
28:
29:
end if
30:
end for
31:
Append long short-term transitions B.
end while
18

---

Published as a conference paper at ICLR 2025
Table 3:
The success rate and the number of steps per episode for task completion.
Harvest log in plains
Harvest water with bucket
Harvest sand
Shear sheep
Mine iron ore
Model
succ. (%)
suce. step
suce. (5)
succ. step
succ. (%)
succ: step
succ. (%)
succ, step
succ. (5)
suce. step
VPT
6.97
963.32
0.61
987.65
12.99
880.54
1.94
987.49
0.00
STEVE-1
57,00
752.47
6.00
989.07
37,00
770.40
3.00
992.36
0.00
PTGM
41,86
811.19
2.78
977.78
17.71
833.64
21.54
887.03
15.14
1586.03
Director
8.67
968.09
20,90
931.74
36.36
825.35
1.27
995.99
7.82
1906.31
DreamerV3
53.33
711.22
55.72
628.79
59.88
548.76
25.13
841.14
16.79
1789.06
ImS-ine
601.18
633.78
20.28
1748.55
Forved
0.10
400
0.08
0.06
0.04
0.05
100
8.
8.
Environment
366
Enviranment steps
160
Environment steps
346
(a) Jumping frequency
(h) Interval Á
(c) Dynamic threshold Phresh
Figure 11: Analyses of long-term imaginations throughout training.
C.5 CLARIFICATION ON STOCHASTIC LONG-TERM IMAGINATION
One might argue that long-term imagination could skip essential intermediate steps that gradually lead to the objective, potentially resulting in a lack of learning for these crucial actions. To address this issue, we adopt a probabilistic mechanism. Specifically, even when j. True, indicating that a long-term transition is to be executed, we implement probability of 0.7 for executing the jump and 0.3 for not jumping. This allocation ensures a 30% chance that the transition will execute the short- term imagination with gradient feedback attached to the actions. This stochastic decision-making based on a uniform distribution, providing a balanced approach between leveraging long-term imagination and capturing essential short-term behaviors.
C.6 ADDITIONAL LIMITATION
It is worth mentioning that LS- Amagmnee simulates the agent's 's state when approaching a target object in 3D navigation environments with embodied agents by zooming in on the observed image, and sets intrinsic rewards based whether the agent is close to and has positioned the target object at the center of the observation. Therefore, LS-Imagine is not suitable for environments with fixed viewpoints, 2D environments, or those where the reward mechanism is more complex than approaching objects (e.g.. driving).
D ADDITIONAL RESULTS
D.1 NUMERICAL COMPARISONS
Table 3 compares existing approaches the challenging MineDojo environment.
D.2 ANALYSES LONG-TERM IMAGINATIONS
We use the task harvest log in plains as an example to facilitate the understanding of the long short- term imagination process. Figure (a) we first track the frequency long-term imaginations and the corresponding predicted state intervals Â throughout the training process. The curve shows the proportion of imagination sequences involving jumpy state transitions relative to the total number of imagination sequences. Initially, the jumping frequency is low because the world model has not yet learned to identify when jump is necessary based on the state. As the model's predictions improve in the early stages of training, the frequency increases, likely due to the agent's 's underdeveloped policies, which result in more observations far from the goal and necessitate long-term exploration. Over time, as the agent learns policies that bring it closer to the target, the frequency of observations far from the goal decreases, reducing the need for jumps.
19

---

Published as a conference paper at ICLR 2025
Short-term imagination
Long-term imagination
Adeetass Posteske say
Harvest log in plains
Harvest
Harvest sand
(d) Shear sheep
(e) Mine iron ore
Figure 12: Visualization of the complete long short-term imagination sequences.
Additionally, we find that among all sequences with jumpy state transitions, the average number of jumpy transitions sequence, within horizon of 15 steps, 1.02. This indicates that, in most cases of these tasks, a single jumpy transition is sufficient to bring the agent close to the target.
In Figure 11(b) we track the variations of the jumping state intervals, Å, throughout training. At the beginning, A, high, indicating that the policy requires many steps reach the target. As the policy improves, fewer steps are needed to approach the target, leading to a gradual decrease A. Notably, as Å, evolves with the updated policy, also ensures minimal misalignment in Eq. (9) between the future cumulative rewards computed with jumpy imaginations and the behavior policy.
Furthermore, Figure [[[]] we track the variation curves the dynamic threshold Pitresh during training in different tasks, and observe that:
20

---

Published as a conference paper at ICLR 2025
Explore for village
Mine
(a) Explore village
(b) Mine ore
Figure 13: Affordance maps when the target is invisible or occluded.
Dreamerve
LS Imagine
Dreomerva
15 simgghn
300
1000
223
25
43
11
855
Spers ap2h112
700
5.2
403
3B1
00/45
45%
anodee
cosbestere
200
moriing
COAA
(a) Success rate
(b) The number of steps required task completion
Figure 14: Comparison of LS-Imagine and Dreamer V3 on a long-horizon "Tech Tree" task.
For task such as hurvest log in plains, the variance thresh high during the early stages training. Since Pihresh serves as temporal smoothing of Pump- this reflects the significant fluctuations of Pjump at the beginning of training, highlighting the importance adopting dynamic threshold.
Across various tasks, Pitresh consistently converges in the later stages of training, demonstrating its effectiveness in improving the stability of exploratory imaginations.
The converged values Piteesh differ across tasks, indicating that involving an automated computa- tion of Pitrssh enables us avoid tedious hyperparameter tuning.
D.3 VISUALIZATION OF LONG SHORT-TERM IMAGINATIONS
As illustrated in Figure 12 visualize the complete long short-term imagination sequences for the agent across various tasks. This visualization further demonstrates how the affordance map accurately identifies regions of high exploration potential in the image, and how the long short-term imagination approach provides reasonable and applicable guidance for the agent's task execution. These qualitative results reinforce the effectiveness of method in guiding the agent toward its goal with greater precision and efficiency.
D.4 DEPENDENCE ON THE VISIBILITY OF OBJECTS
The long-term transitions of our approach rely the affordance map to identify high-value ex- ploration areas. However, it is crucial to note that our affordance map generation method is not merely object 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 affordance maps that provide guidance even when the target is completely occluded.
21

---

Published as a conference paper at ICLR 2025
Long-term branch lass scale
reward weight
Gaussian distribution std
0.8
0.8
Success rate
0.6
0.5
9.6
Slicess success
0.4
Aug
Avg
0.2
Aog +1.0
0.5
1.0
s-l (64,
5.0
Aom 10.0
10.0
(du 0y--1 1601
0.0
0.4
0.6
0.8
1.0
68
02
0.0
0.2
0.4
066
08
Environment
led
Enviranment steps
105
Environment steps
1n5
Figure 15: Results of hyperparameter sensitivity analyses.
For instance, as illustrated in Figure 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 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 occluded, the affordance map enables the agent continue exploring effectively.
Further research direction. 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 ftarget/" when the affordance map fails identify high-value areas. This prompt, combined with intrinsic rewards generated by MineCLIP, encourages the agent to conduct extensive exploration.
D.5 RESULTS ON LONG-HORIZON TASKS
To demonstrate the potential application of LS--magine more complex tasks, we conduct experi- ments on a "Tech Tree" task in MineDojo, specifically crafting a stone pickaxe from scratch. This task involves seven subgoals: log, planks, crafting table, stick, wooden pickaxe, cobblestone, and stone pickaxe. Since LS- IS--ine primarily designed to focus on environmental interactions and task execution under fixed objectives, rather than task decomposition and planning, we adopt the DECKARD method (Nottingham et al. (203) for task planning. This method provides top-level guidance, with L.S- IS-IImgiine executing the corresponding subtasks. Each subtask was trained for million steps and then tested within 1, 1,000 steps per episode. The results are shown Figure 14 which demonstrate that our LS-Imagine consistently outperforms DreamerV3, achieving higher success rates and requiring fewer steps to complete each subgoal.
D.6 HYPERPARAMETER ANALYSES
We conduct sensitivity analyses on three hyperparameters:
The long-term branch loss scale 3hog: As shown in Figure 15 (Leff), we observe that when Frong for the long-term branch is too small or too large, it impedes the learning of long-term imagination, leading to a decline in performance.
The intrinsic reward weight az From Figure 15 (Middle), we observe that if the hyperparameter a for intrinsic reward is excessively small, it may result in insufficient guidance and inaccurate reward estimation for the post-jumpy-transitte state.
The intrinsic reward Gaussian parameters Oy): As shown in Figure 16 (a aTy) control the
standard deviations of the Gaussian distribution along the horizontal and vertical axes, respectively. Intuitively, setting these hyperparameters too low may cause the model to overlook targets located at the edges of the observed images. Conversely, excessively high (a,,. Fy)) 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. From Figure [5] (Right), we observe that the final performance is robust to the tested parameters, with all configurations outperforming the baseline models presented in previous experiments.
The final hyperparameters of LS- Imaine are shown in Table
22

---

Published as a conference paper at ICLR 2025
[d. o,) - 164,40 40)
(dg o,) - (126.80)
[o,, o,) 1256,160
40
15
40-
80-
10
64
126
192
256
128
152
256
6
120
1.2
256
Figure 16: Visualization of Gaussian matrices with different standard deviations.
Table 4: Hyperparameters of LS-Imagine.
Name
Notation
Value
Affordance map generation
Sliding window size
0.15 0.15
Sliding steps
9 x 9 500
U-Net train epochs
U-Net initial learning rate
5 10
U-Net learning rate decay epochs
50
U-Net learning rate decay rate
0.10
Text feature dimensions
512
Gaussian distribution standard deviations
(a2: oy)
(128, 80)
General
Replay capacity
1 x 106
Batch size
B
16
Batch length
T
32
Train ratio
16
World Model
Intrinsic reward weight
1 4,096
Deterministic latent dimensions
Stochastic latent dimensions
32
Discrete latent classes
32
RSSM number of units
1,024
World model learning rate Long-term branch loss scale Reconstruction loss scale Dynamics loss scale Representation loss scale
1 x 10
Blong
1
Sprod
Bdyn
Prep
0.1
Behavior Learning
Imagination horizon
L
15
Discount
0.997
X
X-target
0.95
Actor learning rate
3 10
Critic learning rate
3 10-5
23


Reviews:
[{"review_id": "CQKc8YBsz7", "reviewer": "Reviewer_hBEP", "summary": "The paper introduces a novel hierarchical model-based reinforcement learning (MBRL) framework named Long Short-Term Imagination (LS-Imagine), designed to address the challenge of short-sightedness in open-world decision-making with high-dimensional visual inputs, such as MineDojo. In LS-Imagine, imagination horizon can be extended with a specialized world model. Specifically, a affordance maps are predicted to guide the jumpy switch. Two world models are trained to capture transitions at different temporal resolutions, short-term state transition and long-term state transition. Agent learning is conducted during imagination, a common strategy utilized in most background planning method.  The authors provided experimental results on Harvest tasks from the MineDojo benchmark, showed superior results compared with baseline methods.", "questions": "1. Could the authors elaborate more on how the long-term transition data is collected? The phrase from the appendix still confusing to me, specifically, how is the reward of the state helps measure the long-term transition.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The jumpy prediction technique within the long-term imagination framework is innovative as it departs from the fixed interval approach prevalent in previous work, offering increased flexibility in jumpy prediction\n2. The paper is well-organized and clearly written.", "weaknesses": "1. The proposed method employs a hierarchical structure, yet the baseline comparisons are made with flat learning methods. Including comparisons with hierarchical MBRL methods like Director[1] could greatly strengthen the paper.\n2. Equation 9 appears to have an inconsistency in the time indexing; should the bootstrapping term $R^\\lambda_{t+1}$ be $R^\\lambda_{t+\\hat{\\Delta}_{t+1}+1}$ ? \n3. The use of  $\\lambda$ -return in evaluating the policy might introduce bias since it should be evaluated with on-policy data,, but the predicted jumpy state, $\\hat{z}_{t+1}$, might not aligned with the learning policy.\n4. The paper focuses on Harvest tasks. Including results from other complex, long-horizon tasks, such as the Tech Tree task group from MineDojo, would better demonstrate the framework’s effectiveness.\n\n[1]: Hafner, Danijar, et al. \"Deep hierarchical planning from pixels.\" Advances in Neural Information Processing Systems 35 (2022): 26091-26104.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Jvq6y1SfOO", "reviewer": "Reviewer_ynNo", "summary": "This paper studies Model-based RL in Open-World environment, specifically Minecraft. The authors propose a long short-term imagination method to improve the imagination process of the dreamer algorithm. The main idea is to use an affordance map to identify the distance object of interest and train a long-term transition model to skip the approaching steps so that the imagination can be more task-relevant. Results from 5 tasks in Minecraft showcase the improvement of the method over Dreamerv3 and other two baselines.", "questions": "- The term \"Gaussian matrix\" in line 215 is confusing. Based on the context later, I think you are referring to Gaussian kernel. If that is correct, what is the $\\sigma$ you use in the paper? Is the method sensitive to different value of $\\sigma$?\n- In Section 3.2.2, you mentioned the affordance map generator is pretrained on the random exploration data. How do you make sure that the task relevant observations are present in the random data to train a meaningful generator? For example, if the task is to collect a diamond, and there certainly won't be any diamond seen by the random agent. Will the method still help to solve the task?\n- I am confused by the model analysis part with the parallel pathway. Things explained in Section 4.3 sound the same with Section 3.4. Could you elaborate what is the difference between the sequential pathway and the parallel pathway?\n- Some relevant citations are missing: [1] studies the hierarchical world model which uses a similar strategy of doing the long-term imagination on the higher-level states; [2] studies a similar problem that short imagination horizon may not enough to cover behaviour of interest. I would suggest the author to discuss these works.\n- There are no error bars for the middle and right plots in Figure 7. Are these results only based on one seed?\n- What is the computation cost of this method?\n- Maybe I missed some parts, but how does the policy learn the behaviour that the long-term imagination skip? For example, in the \"Harvest log in plains\" task, the go toward the tree is always skipped by the long-term imagination if the model learns well. Since the skip transition doesn't have gradient attached to the action, I wonder how does the policy learn to walk toward the tree.\n\n[1] Gumbsch *et al.*, Learning Hierarchical World Models with Adaptive Temporal Abstractions from Discrete Latent Dynamics, ICLR 2024 \n\n[2] Hamed *et al.*, Dr. Strategy: Model-Based Generalist Agents with Strategic Dreaming, ICML 2024", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is mostly well written. \n- The method proposed is novel and the results are promising comparing to the baselines.", "weaknesses": "- Although the high-level idea is straight-forward, the implementation is overcomplicated. \n- The method feels very ad-hoc to the Minecraft tasks studied in this paper. It doesn't come into my mind about any other relevant tasks other than Minecraft where the proposed method can be applied.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "O2YbGf1unT", "reviewer": "Reviewer_3G6A", "summary": "This paper introduces an exploration approach, LS-Imagine, to improve model-based reinforcement learning from pixels on open-world environments (with large state spaces). \n\nThe paper aims to address the short-hoirzon limitation of model-based approaches relying on repeated autoregressive single timestep prediction. This is achieved by first generating affordance maps of observations by zooming into different patches and determining if there are objects associated with the current task in those patches using a vision-language reward model (MineCLIP in this work) distilled into an affordance map predictor. This affordance map is used to generate an intrinsic reward associated with the observation. The kurtosis of the affordance map is then used to determine if the object associated with reward is near or far away (binary, based on a threshold). If the object is deemed to be near, a standard 'short-term' world model using autoregressive single timestep prediction is used to generate trajectories for an actor and critic to learn from. Alternatively, if the object deemed to be far away, a 'long-term' world model using autoregressive multi-timestep interval prediction is used (although ignoring actions). This enables the world model to generate a trajectory over a larger number of environment timesteps in the same number of autoregressive prediction steps, therefore potentially enabling the generated trajectory to get closer to the object of interest. This long-term trajectory is used to train the critic but not the actor. The combination of these models is termed a Long Short-Term World Model, giving the approach its name: LS-Imagine.\n\nLS-Imagine is applied to 5 MineDojo tasks and outperforms the baselines in terms of success rate and per-episode steps after training. Ablations demonstrate the benefit of both affordance-based intrinsic reward and the use of long-term imagination.\n\nIn summary, the paper attempts to tackle the important problem of long horizon world modelling and provides an interesting and novel approach to address this problem in 3D environments. While it is a relatively complex method and somewhat limited to Minecraft, I believe the ideas and insights warrant acceptance.", "questions": "- Why was the word 'jumpy' chosen? Could a more precise word be used instead?\n- The authors say on L42 that MBRL approaches optimize short-term experiences of \"typically 50 time steps\". Where did 50 come from? My understanding is that MBRL methods such as DreamerV3 commonly use an even shorter horizon of $\\sim 16$ timesteps.\n- Could the limitations be expanded upon to cover the general application of LS-Imagine beyond Minecraft?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- **Significance**\n  - Long-horizon world modeling and reinforcement learning in open-world environments are important problems.\n  - The proposed approach is insightful and successfully addresses these problems.\n- **Originality**\n  - The proposed approach involves the combination of multiple novel and inisightful components.\n- **Quality**\n  - Overall the quality of the paper is relatively high, with the method reasonably clearly explained and analyzed.", "weaknesses": "- **Clarity**\n  - Some aspects of the paper are not particularly clear. The main one is the use of the word 'jumpy' throughout the paper. The meaning of this word is assumed, but is not defined in the paper or standard usage as far as I'm aware, and is relatively unscientific, so I feel it is not the right word to use. 'Multi-step' state transitions seems more appropriate. If the authors were attempting to highlight that the number of steps can vary, then 'variable-step' transitions would be better. At the very least, 'jumpy' should be properly defined at the beginning of the paper.\n  - Similarly affordance maps may not be familiar to all readers and the exact meaning of this term can vary. For example, a short clarification early in the paper such as ''...affordance maps, that elucidate which parts of an observation may be associated with reward, ...\" would be helpful.\n  - Some other unclear aspects/minor mistakes include:\n    - L326: \"employs *an* actor-critic algorithm to learn behavior *from* the latent state sequences...\"\n    - L351: Grammar is slightly wrong and confusing, should be: \"Notably, since long-term imagination does *not* involve actions, we do not optimize the the actor when long-term imagnation is adopted.\" \n    - L354: Also worth highlighting the difference with the DreamerV3 loss. \"The loss of the actor is therefore equivalent to DreamerV3, with an additional factor of $(1-\\hat j_t)$ to ignore updates from long-term imagination:\"\n    - L361: \"on the chellenging...\"\n    - L500: Doesn't make sense. Maybe \"Our work is also related to affordance maps for robotics.\" is sufficient?\n    - \"*Learning to Move with Affordance Maps*\" [1] should likely also be compared and cited here.\n\n- **Limitations**\n\n  - This approach has important limitations that are not mentioned. In particular, the approach is limited to embodied agents navigating a 3D environment in which there are objects associated for which reward is obtained by approaching them. Therefore the approach assumes, for example:\n    - Observations are of a 3D environment\n    - Actions are navigation actions of an embodied agent\n    - Rewards are assoiated with identifying or moving towards objects\n    - A reward model is available to identify high reward regions of observations\n  - The experiments are limited to Minecraft for which these assumptions hold. This approach would likely not work as well even in Crafter [2] for example, which provides a 2D 'open-world' analogue of Minecraft, since objects do not become larger as you move towards them.\n  - The approach also relies on both the long-term and short-term models being used, given only the short-term model is able to update the actor. While the thresholding of $P_{jump}$ can partially be used to address this, this is not particularly robust, and still requires some close-up objects in initial exploration for the standard one-step world model to be used, so the approach may not work as well in very sparse environments.\n  - There is still significant value of the approach despite these limitations, and the paper is reasonably scoped, but they should be included in the limitations at the end of the paper, which are currently overly brief and narrow.\n\n  **References:**\n\n  [1] \"*Learning to Move with Affordance Maps*\", Qi et al., 2020\n\n  [2] \"*Benchmarking the Spectrum of Agent Capabilities*\", Hafner, 2021", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OqKNW9Mswm", "reviewer": "Reviewer_22FD", "summary": "This paper proposes a long short-term world model to address the limited imagination horizon of traditional world models.\nThe authors leverage MineCLIP to detect task goals and generate intrinsic rewards, enabling the model to jump directly to these goal states when feasible.\nFor policy optimization over these jump-based transitions, they train modules to predict future value, interval timesteps, and cumulative rewards.\n\nThe paper is well-organized, and its key contributions are:\n(1) It introduces a method for generating target states in MineDojo (or possibly in other 3D-RPG games).\n(2) It demonstrates the feasibility of training a world model with jumping transitions and optimizing policies over such transitions.", "questions": "1. Jumping transitions\n\n    (a) I don't get the reason and rule for this dynamics threshold on lines 248 to 250. Could you clarify this?\n\n    (b) The jumpy transition in the world models is a novel contribution, and I think it's worth more explanation for intuitive understanding. Like how many times would the jumping transition be triggered during the training, and what is the predicted average $\\Delta'_t$ in the imagination (how many short-term steps are saved)?\n\n2. I suggest adding a numerical result (possibly in the appendix) for the results in Figures 4 and 5.\n\n3. For the parallel imagination in section 4.3. Is the post-jumping state connected with the bootstrapping $\\lambda$-return of the pre-jumping states? How is the conclusion on lines 464 to 466 be made?\n\n4. For the model-based RL review on lines 487-489, I suggest adding another two world model papers for completeness: (1) Zhang, Weipu, et al. \"STORM: Efficient stochastic transformer based world models for reinforcement learning.\" Advances in Neural Information Processing Systems 36 (2024). and (2) Alonso, Eloi, et al. \"Diffusion for World Modeling: Visual Details Matter in Atari.\" arXiv preprint arXiv:2405.12399 (2024). (NeurIPS24 Spotlight)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. It introduces a method for generating target states in MineDojo (or possibly in other 3D-RPG games).\n2. It demonstrates the feasibility of training a world model with jumping transitions and optimizing policies over such transitions.\n3. The illustration is clear and the writing organization is good.", "weaknesses": "1. Since MineCLIP is an important tool for this work, I suggest the author include a brief introduction of what MineCLIP can do in section 3.2/3.2.1 or an appropriate position. This would help readers who are not familiar with research on MineDojo to understand this paper.\n\n2. In the abstract, \"We argue that the primary obstacle in open-world decision-making is improving the efficiency of off-policy exploration across an extensive state space.\" This seems not closely connected to the main contribution of this paper. I suggest paraphrasing it to highlight \"across a long horizon\" or something that is more related to the topic.\n\n3. Though the method sounds promising for solving MineDojo tasks, it may not be a general method for all kinds of open-world games. Such as in 2D games or fixed camera control tasks.\n\n    (a)  Before seeing the target for the first time in the training process, there won't be a reasonable goal-directed reward or jumping option, the exploration still requires extensive enumerates.\n\n    (b) The crop and resize operation (or assumption) is only useful for 3D visual navigation tasks.\n\n    (c) When the target is occluded, the world model still needs to perform step-by-step imagination.\n\nIf these are true, I suggest the authors include a sentence or so in the limitation section to clarify it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)