DreamFly: Causal Memory and Receding-Horizon Diffusion Planning for Aerial Vision-Language Navigation

DreamFly tackles aerial vision-language navigation (VLN) by addressing three key limitations in VLA models: limited historical context, short planning horizons, and unreliable implicit termination. It builds on Dream-VLA and introduces (1) a causally aligned memory module, (2) receding-horizon diffusion planning, and (3) LiteStop for explicit stop prediction.

Causal Historical Memory§

Given a streaming history of observations $o_1, \dots, o_t$, the current visual feature $f_t$ is augmented with a causal memory summary $m_{t-1}$ that only depends on observations before the current decision step. This prevents future information leakage while enabling temporal reasoning. The memory is updated via a lightweight recurrent mechanism: $$ m_t = \mathrm{GRU}(f_t, m_{t-1}). $$ The augmented representation $h_t = f_t \oplus m_{t-1}$ is fed to the policy.

Receding-Horizon Diffusion Planning§

Instead of predicting a single action, the policy denoises a $K$-step action chunk $a_{t:t+K-1}$ conditioned on $h_t$ and the language instruction. During inference only the first action $a_t$ is executed, then the environment returns a new observation and the process repeats. This "plan-$K$, execute-one" scheme provides auxiliary future action targets during training while preserving closed-loop feedback. The diffusion loss is: $$ \mathcal{L}_{\mathrm{diff}} = \mathbb{E}_{\epsilon \sim \mathcal{N}(0, I), k \sim \mathrm{Uniform}} \left[ \| \epsilon - \epsilon_\theta(\bar{a}_{t:t+K-1}^{(k)}, h_t, k) \|^2 \right], $$ where $\bar{a}^{(k)}$ is the noised action chunk at diffusion step $k$.

LiteStop: Explicit Termination§

Termination is decoupled from action generation. The stop probability is estimated directly from the initial all-mask action-token logits produced by the diffusion decoder: $$ p_{\mathrm{stop}} = \sigma(W_{\mathrm{stop}} \cdot \mathrm{logits}_{\mathrm{mask}} + b). $$ This avoids the ambiguity of generating a special "stop" action and improves navigation precision.

Training and Implementation§

The full objective combines the diffusion loss, memory prediction loss (if any), and stop-head cross-entropy. The model is trained on the OpenFly dataset. At inference, the agent replans after each executed action, and stops when $p_{\mathrm{stop}} > 0.5$.

A simplified PyTorch-style forward pass:

class DreamFly(nn.Module):
    def __init__(self, base, K=5, T=1000):
        super().__init__()
        self.base = base  # Dream-VLA backbone
        self.gru = nn.GRU(input_size=base.hidden_size, hidden_size=base.hidden_size)
        self.diffusion_head = DiffusionHead(base.hidden_size, action_dim, K, T)
        self.stop_head = nn.Linear(base.hidden_size, 1)

    def forward(self, obs_feat, lang_emb, prev_memory):
        fused = self.base.crossmodal(obs_feat, lang_emb)
        _, memory = self.gru(fused.unsqueeze(0), prev_memory)
        cond = torch.cat([fused, prev_memory.squeeze(0)], dim=-1)
        action_chunk = self.diffusion_head.sample(cond)
        stop_logit = self.stop_head(cond.detach())
        return action_chunk[:, 0], stop_logit, memory

Results§

On OpenFly, DreamFly achieves 32.04%/29.46% SR and 28.22%/23.54% SPL on test-seen/test-unseen, outperforming all baselines on both metrics and attaining the lowest navigation error. Ablations confirm that each component contributes to the gains.