LS-Imagine: Overcoming Short-Sightedness in Open-World Reinforcement Learning

TL;DR

Training agents to operate in open-world environments like Minecraft is notoriously difficult due to the "short-sightedness" of current Model-Based Reinforcement Learning (MBRL). LS-Imagine breaks this bottleneck by allowing agents to "jump" through time in their imagination. By combining a Long Short-Term World Model with affordance-driven intrinsic rewards, the agent can visualize distant goals and achieve SOTA success rates with significantly higher sample efficiency.

The "Short-Sighted" Bottleneck in MBRL

State-of-the-art MBRL frameworks, such as DreamerV3, typically train policies by rolling out short sequences of imagined states. While effective for localized tasks, this approach fails in the "Open World" for three reasons:

1. Vast State Space: The agent spends too much time on irrelevant local interactions.
2. Sparse Long-Term Rewards: Beneficial outcomes (like finding a tree) are often "invisible" to a short-term imagination horizon.
3. High Trial-and-Error: Without guidance, the agent relies on random exploration that rarely hits long-horizon targets.

Methodology: The Architecture of Long-Sight

The core innovation of LS-Imagine is its dual-branch transition model. Instead of just predicting step $t+1$, the model can trigger a Jumpy Transition to a distant future state $t+\Delta$.

1. Affordance Maps via Virtual Exploration

To know where to jump, the agent needs a spatial prior. The authors propose a clever "zoom-in" trick: the model scans the current observation, crops potential target areas, and resizes them—simulating the visual effect of moving toward a goal. A pre-trained MineCLIP model evaluates these "virtual videos" against text instructions (e.g., "cut a tree") to generate an Affordance Map.

Figure 1: The framework enables agents to switch between instant and jumpy state transitions within their imagination.

2. The Long Short-Term World Model

The world model is augmented with:

• Jump Predictor (${\hat{j}}_t$): A flag that triggers long-term branch imagination based on the kurtosis (peakiness) of the affordance map.
• Interval Predictor (${\hat{\Delta }}_t$): Estimates how many real-world steps a "jump" represents.
• Reward Predictor (${\hat{G}}_t$): Estimates the cumulative discounted reward across that jump.

Experiments: SOTA Achievement in MineDojo

LS-Imagine was tested on five challenging Minecraft tasks. The results show a dominant lead over traditional MBRL and foundation models.

Performance Highlights:

• Success Rate: In the "Harvest log" task, LS-Imagine reached an 80.6% success rate, compared to DreamerV3's 53.3%.
• Efficiency: The agent completes tasks in significantly fewer steps, as the "jumpy" imagination allows it to optimize for long-term value directly.

Figure 2: Success rate comparison across various open-world Minecraft tasks.

Deep Insight: Why "Series" Imagination Matters

An interesting finding in the ablation studies was the comparison between Series and Parallel imagination.

• Parallel: The model imagines a normal sequence AND a jump independently.
• Series (LS-Imagine): The jump is integrated into the sequence.

The Series approach performed much better because it allows the "post-jump" value (the vision of reaching the tree) to propagate back to the "pre-jump" actions, providing a direct gradient for distal exploration.

Critical Analysis & Future Outlook

Takeaway: LS-Imagine proves that MBRL agents don't just need to dream; they need to dream across multiple timescales. By identifying high-value "affordance" areas, the agent transforms a needle-in-a-haystack search into a targeted navigation problem.

Limitations: The current "zoom-in" heuristic for affordance mapping is tailored for 3D embodied navigation. Adapting this to complex manipulation (e.g., fine-grained tool assembly) or fixed-viewpoint tasks remains a challenge. Additionally, the computational overhead of the U-Net and jumpy rollouts adds approximately 23GB of VRAM usage, making it hardware-intensive.

Future Work: Integrating this temporal abstraction with Large Language Models (LLMs) for task decomposition could create an even more powerful agent capable of both high-level planning and efficient low-level visual control.