PaperlayerView exact revision on arXiv

Source-linked AI summary

SearchWorld: Spatial Value-Grounded Imagination for UAV Object Search via World Models

Yatai Ji, Zhengqiu Zhu, Yong Zhao, Yue Hu, Fanglong Yao, Chen Gao, Pengfei Zhu, Quanjun Yin

arXiv:2610.09335v1cs.AIcs.LG

TL;DR

Aerial object search must identify a target across large urban regions despite partial observability, while existing map-based methods remain reactive and latent world models make spatially grounded planning challenging. SearchWorld links explicit spatial memory with value-guided imagination, reaching 23.8% SR and 35.5% OSR overall on UAV-ON.

  • Problem

    Map-based search agents reason over past observations, while world models keep spatial knowledge implicit, leaving spatial memory and prospective reasoning unlinked.

    Source: §1 · p. 2

  • Method

    SearchWorld combines a recurrent state-space world model, explicit BEV exploration and obstacle memory, and a decoded task-aware spatial value layer that guides imagined action planning.

    Source: §1 · p. 2

  • Results

    23.8% SR and 35.5% OSR overall on UAV-ON, improving over the strongest published agent, with 19.9% success on unseen scenes.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Latent imagination costs a single latent transition per imagined step, and SearchWorld runs an order of magnitude faster than methods invoking a language model per step.

    Source: §5 · p. 10

Abstract

from arXiv · show

Autonomous unmanned aerial vehicle (UAV) object search involves a closed loop of perception, decision-making, and action under partial observability. Urban environments pose several challenges: large search areas and narrow egocentric views limit coverage, dense 3D geometry constrains safe motion, and open-world instructions require identifying a specific target among distractors. Many existing methods mitigate partial observability through explicit maps or memory representations, yet remain largely reactive, reasoning over past observations without explicitly predicting future states. World models enable prospective reasoning through imagined rollouts. However, image-generating world models can incur high inference latency, while spatially grounded planning remains challenging for latent world models. We propose SearchWorld, a recurrent state-space world model that connects explicit spatial memory with value-guided imagination. The model maintains BEV exploration and obstacle memory and decodes a task-aware spatial value layer to guide search. A cognition-action network uses this learned spatial value prior to improve the policy through imagined rollouts, without training a separate scalar critic. Training progresses from world-model learning to expert imitation and imagination-based exploration refinement. On UAV-ON, SearchWorld improves the success rate to 23.8% (19.5% for the strongest published agent) and raises oracle success to 35.5%, while remaining robust on unseen scenes (19.9% success rate). By grounding imagination in explicit spatial representations, SearchWorld enables UAV agents to plan prospectively rather than react.

1 Introduction

Aerial object search requires agents to explore regions larger than their camera view while remembering observations and choosing useful future regions. SearchWorld links explicit spatial memory with value-guided imagination in a three-stage training curriculum.

1 Introduction

  • Aerial search spans regions much larger than the camera view, with similar objects and no externally provided target location or route.Broad exploration can waste the step budget, while following one target hypothesis can cause repeated visits and occlusion traps.Source: §1 · p. 1
  • Map-based agents use explicit spatial memory but reason reactively over past observations, while world models can imagine future states but keep spatial knowledge implicit.The introduction identifies linking spatial memory to prospective reasoning as the shared challenge.Source: §1 · p. 2
  • 23.8% SR and 35.5% OSR overall on UAV-ON, with 19.9% success on unseen scenes; SearchWorld improves over the strongest published agent.The reported values are overall results except the unseen-scenes success rate.Source: §1 · p. 2
  • SearchWorld combines a recurrent state-space world model with BEV exploration and obstacle memory and a task-aware spatial value layer decoded from the latent state.The three-stage curriculum learns state transitions and spatial representations, imitates expert search, then refines spatial exploration through value-guided imagination.Source: §1 · p. 2
  • The cognition–action network uses the decoded spatial value to guide imagined planning and improve the action policy.The value layer is supervised by exploration, obstacle, and training-time target signals; at inference it is not given the target coordinate, and imitated stopping behavior is retained.Source: §1 · p. 2

2 Related Work

Prior navigation world models either generate future observations as pixels or imagine in latent space; aerial search methods use explicit spatial memories but remain reactive. SearchWorld’s overview connects a task-oriented RSSM, spatial memory, and value-guided action refinement across a three-stage curriculum.

2 Related Work

  • Pixel-generating navigation models simulate future viewpoints or trajectories, but autoregressive video generation at each decision creates latency concerns for large-region search.The cited examples include PathDreamer, DreamWalker, DreamNav, NWM, and ANWM.Source: §2 · p. 2
  • Latent-space world models use imagined latent transitions, but the reviewed approaches keep environmental memory implicit and spatially grounded search planning remains challenging.Examples include MILE, X-Mobility, NavMorph, and WMNav.Source: §2 · p. 2
  • Figure 1 shows an RSSM encoding RGB view, instruction, and spatial memory, then decoding observations, spatial memory, value, and UAV action across a three-stage curriculum.The cognition–action network trains the action policy in Stage 2 and refines movement actions in Stage 3 by pooling decoded value over imagined action footprints.Source: Figure 1 · p. 3
  • Recent aerial search agents use explicit maps or adaptive memory, but their execution-time readout relies on rules or model queries and does not model action consequences.The passage names APEX, PRPSearcher, and OctMem-Agent as examples.Source: §2 · p. 3

3 Method

SearchWorld combines an explicit bird’s-eye-view memory with a recurrent world model that imagines search dynamics and decodes spatial values. Its cognition–action network uses imagined values to refine movement actions, within a three-stage training curriculum.

3.1 Problem Formulation

  • The POMDP formulation conditions the policy on the interaction history because the full UAV pose and environment occupancy are not directly observable.The benchmark supplies eight fixed atomic actions, including stop; success is used for evaluation rather than training.Source: §3.1 · p. 4

3.2 Explicit Spatial-Cognition Memory

  • SearchWorld’s allocentric BEV memory combines sensor-derived exploration and obstacle layers with a task-aware value layer decoded from the world model’s latent state.The first two layers record observed coverage and flight constraints; the value layer predicts where to search next and is supervised during training.Source: §3.2 · p. 4

3.3 Task-Oriented RSSM World Model

  • The task-oriented RSSM rolls search forward in latent space using prior-only transitions, while decoder heads predict observations, spatial memory, values, and actions.At decision time, the prior advances imagined states without querying the simulator; pixels are not decoded at decision time.Source: §3.3 · p. 5

3.4 Cognition–Action Network

  • The cognition–action network replaces a learned scalar critic with a decoded spatial value prior and uses imagined action utilities to refine the task policy.The target location is used only in training supervision, and value guidance updates movement actions while leaving stopping behavior under the imitated policy.Source: §3.4 · pp. 5–6

3.5 Three-Stage Training Curriculum

  • Training proceeds through random-action world-model learning, expert behavior cloning, and value-guided imagination planning that refines movement actions without a critic loss.In Stage 3, the RSSM and value decoder are frozen, and only the action head receives gradients.Source: §3.5 · pp. 6–7

4 Experiments

On UAV-ON, SearchWorld achieves the highest reported overall success and oracle success among the compared agents, with strong results on unseen scenes. Its experiments also assess runtime, spatial-memory components, training stages, and a representative search episode.

4.1 Experimental Setup

  • The evaluation uses UAV-ON’s 14 scenes and about 11,000 tasks, with 4 scenes held out; success requires stopping within 20 m of the target.Reported metrics include SR, OSR, and SPL, stratified by object size and scene familiarity; comparisons include benchmark, memory-augmented, and world-model baselines.Source: §4.1 · p. 7

4.2.1 Comparison with SOTA baselines

  • 23.80% SR and 35.52% OSR overall; SearchWorld exceeds OctMem-Agent’s 19.50% SR and AOA-F’s 7.30% SR.The authors attribute the SR advantage to finding the target region more often, not stopping more liberally.Source: §4.2.1 · pp. 7–8
  • 25.14%/19.92% SR on seen/unseen scenes; SearchWorld leads both splits in SR and OSR, with 33.12% unseen OSR.Its seen–unseen SR drop is 5.2 points, compared with 2.6 points for OctMem-Agent.Source: §4.2.1 · p. 8, Table 2 · p. 8

4.2.2 Computational efficiency

  • SearchWorld runs an order of magnitude faster per step than methods invoking a language model per step, and faster than the closest competitor.Execution uses frozen-encoder passes, depth back-projection, one RSSM transition, and a value head, without language-model or VLA inference.Source: §4.2.2 · p. 8, Table 3 · p. 8

4.3 Ablation Study

  • Direct expert imitation reaches 14.33% SR, while the Dreamer-style imagined actor–critic reaches 19.07%; SearchWorld’s remaining 4.7-point gap isolates spatial value grounding.The X-Mobility-style latent-only imitation variant reaches 12.91% SR, below direct imitation.Source: §4.3 · pp. 8–9
  • Removing all BEV memory costs 7.8 SR and 10.5 OSR points; removing exploration costs 5.7 SR points, while removing obstacles costs 4.0 OSR points.The authors report re-searching covered ground without exploration memory, detours around unmarked structure without obstacle memory, and lower SPL under every layer removal.Source: §4.3 · p. 9
  • Skipping Stage 1 drops SR by 10.4 points, skipping Stage 2 by 12.8, and skipping Stage 3 leaves 19.71% SR.The stages supply transition coverage, a task prior, and value-guided refinement beyond demonstrations, respectively.Source: §4.3 · p. 9

4.4 Case Study

  • In the traced successful episode, spatial value shifts from the nearest unexplored frontier to the remaining unexplored sector, then rises around the target before the agent stops within τd.The authors also describe a failure in which value concentrates behind tall structures and the agent repeatedly approaches the occluded region.Source: §4.4 · p. 9

5 Conclusion

SearchWorld’s conclusions emphasize map-grounded value guidance, which refines expert search-and-stop behavior rather than replacing it. Its latent imagination is also reported to support faster inference than methods invoking a language model at each step.

5 Conclusion

  • The same world model with a scalar critic recovers half of SearchWorld’s improvement over direct imitation, supporting the conclusion that gains come from deciding on the map rather than the latent state.Imitation alone plateaus below the strongest published agent.Source: §5 · pp. 9–10
  • Expert imitation supplies search-and-stop behavior, while value-guided imagination routes movement more efficiently without replacing that task policy.One imagined step costs a single latent transition, and SearchWorld runs an order of magnitude faster than methods invoking a language model per step.Source: §5 · p. 10

A Implementation Details

The implementation encodes visual, language, and BEV memory streams into a recurrent world model with RGB, memory, value, and action heads. Training proceeds through three stages, with the final stage freezing the world model and updating the action head.

A Implementation Details

  • The observation embedding fuses frozen DINOv2 visual features, frozen SigLIP image and text features, and an encoded BEV exploration-and-obstacle memory.DINOv2 and SigLIP each provide 768-dimensional features; the BEV branch maps two input layers to 1024 dimensions before linear fusion.Source: A Implementation Details · p. 12
  • The RSSM uses a 1024-unit GRU and a 512-dimensional diagonal-Gaussian stochastic state, with RGB, BEV, spatial-value, and categorical action heads.The BEV decoder outputs 256 × 256 exploration and obstacle layers; the value decoder is supervised by the target of equation 5.Source: A Implementation Details · p. 12
  • Training separates world-model pretraining, expert imitation, and value-guided movement refinement, freezing the RSSM and value decoder during Stage 3 while retaining the imitated stop policy.Stages 1 and 2 use random-action rollouts and expert demonstrations, respectively; Stage 3 trains only the action head using prior-only imagined rollouts and a KL objective.Source: A Implementation Details · p. 12
  • A complete three-stage training run takes about 20 hours on 8 NVIDIA H20 GPUs, while evaluation takes about 34.5 seconds per episode on one RTX 5090 workstation.Table 5 lists the spatial-memory grid, geometric construction thresholds, and training regularizer weights; the obstacle layer uses depth-image pinhole back-projection.Source: A Implementation Details · pp. 12–13, Table 5 · p. 12

B Action Discretization Protocol

The benchmark’s fixed discrete actions provide a compact interface and make action footprints computable from the current pose. The theoretical details bound deployment imagination to recent observations and explain the constrained value-guided policy update and its spatial utility.

B Action Discretization Protocol

  • The benchmark supplies eight fixed actions—five translations, two rotations, and stop—whose fixed motion deltas make candidate footprints computable in closed form.Translations are 5 m horizontally and 2 m vertically; rotations are 15°; stop terminates the episode and triggers success evaluation.Source: B Action Discretization Protocol · p. 13

C Theoretical Details

  • The value-guided target reallocates probability among movement actions under a KL-regularized utility objective while holding stop at its expert-policy probability.As temperature approaches zero, the target favors the utility-maximizing movement; as it approaches infinity, it recovers the demonstrated prior.Source: C Theoretical Details · p. 13
  • Each action’s footprint sweeps its sensing fan along the benchmark-fixed motion delta, and pooling decoded value over that footprint gives its expected value mass.The footprint is mapped into the allocentric grid using the current pose; out-of-grid cells are dropped, and the geometry is computable rather than learned.Source: C Theoretical Details · p. 13
  • Stage 1 trains the RSSM with a β-weighted variational objective that reconstructs RGB, memory, and value while regularizing posterior states toward the prior.The prior also serves as the imagination transition.Source: C Theoretical Details · p. 14
  • At deployment, each real observation re-grounds the recurrent state, limiting prior-only latent extrapolation used by a decision to an H-step window.Longer compounding transition error occurs in training rollouts but is not amplified at deployment because decisions do not consume prior-only chains longer than H.Source: C Theoretical Details · p. 14

D Evaluation Metrics

Evaluation defines episode-level success, oracle success, collision, path length, and target distances. It reports SR, OSR, and SPL in line with recent literature, with DTS included for completeness under the original UAV-ON protocol.

D Evaluation Metrics

  • Success is a stop within τd of the target; oracle success records whether the trajectory passed within τd at any step, regardless of stopping location.The evaluation set contains N episodes, with success and oracle success recorded as binary episode indicators.Source: D Evaluation Metrics · p. 14
  • The OSR − SR gap measures episodes in which the agent reaches the target vicinity but fails to recognize and stop there.This is termed the “exploration without termination” failure mode.Source: D Evaluation Metrics · p. 14
  • The tables report SR, OSR, and SPL following recent literature, while DTS is retained from the original UAV-ON protocol for completeness.DTS is computed over successful and unsuccessful episodes alike.Source: D Evaluation Metrics · p. 14
  • SPL credits successful episodes and penalizes path lengths exceeding the geodesic optimum; a score of 1 means every episode succeeds along a shortest path.DTS is reported in meters across all episodes, and lower is better; it reflects the final pose’s distance to the target at termination.Source: D Evaluation Metrics · p. 14
  • APEX, benchmark baselines, and OctMem-Agent results come from their cited papers, while WMNav and PRPSearcher are reproduced on UAV-ON.Video-generation world models are excluded for per-step latency, and methods without released weights cannot be evaluated under the same protocol; APEX has no reported per-object-size breakdown.Source: D Evaluation Metrics · p. 15

E Training Data Collection

The training data collection uses random valid rollouts to pretrain the world model and successful expert demonstrations for task imitation. Both datasets use the UAV-ON AirSim interface.

Dataset interface

  • Both datasets use the UAV-ON AirSim interface, with four RGB-D views, 6-DoF pose, IMU readings, and discretized actions.Spatial memory uses front-view depth; other views are recorded but not consumed by the current architecture.Source: E Training Data Collection · p. 15

Random-interaction dataset (Stage 1)

  • Stage 1 pretrains the world model on random rollouts of up to 50 steps, discarding collisions and recollecting from the same initial state.Actions are drawn uniformly from seven movement primitives by default.Source: E Training Data Collection · p. 15

Expert-demonstration dataset (Stage 2)

  • Stage 2 demonstrations come from a five-phase expert combining oracle-bearing navigation with reactive depth-based avoidance.The phases are stop, align, height align, approach, and avoid; the stopping threshold is 5 m, tighter than the 20 m evaluation threshold.Source: E Training Data Collection · p. 15
  • Only successful-stop episodes are retained; failed attempts are retried with 15% collision-safe exploration noise, while tasks exhausting retries are skipped.Spawn-mismatched and underground tasks are also excluded automatically.Source: E Training Data Collection · p. 15

F Qualitative Visualizations and Failure Analysis

The visual analysis illustrates both successful search behavior and failure patterns. A representative failure links missed target detection to misdirected exploration, while the taxonomy identifies three recurring modes and proposed fixes.

Qualitative visualizations

  • Figure 2 is cited as a case study illustrating successful search behavior discussed in the main text.The supplied passage does not specify the behavior shown in that case study.Source: F Qualitative Visualizations and… · p. 16

Representative failure case

  • In the representative failure, the UAV explores but remains too far from the target, and the step budget expires in the complex environment.The value layer directs it toward unexplored sectors that do not contain the target, consistent with the ablation finding that spatial memory alone cannot identify the right sector.Source: F Qualitative Visualizations and… · p. 16

Failure taxonomy

  • The failure taxonomy identifies occlusion traps, semantic confusion between target instances, and oscillation between similarly valued frontiers.These cases respectively involve repeated approaches without circumnavigating, under-separated instance embeddings at distance, and exploration-dominated value conditioning.Source: F Qualitative Visualizations and… · p. 16
  • Proposed future fixes are depth-aware value masking, instance-aware goal embeddings, and hysteresis in frontier selection.The fixes correspond respectively to occlusion traps, semantic confusion, and frontier dithering.Source: F Qualitative Visualizations and… · p. 16

G ICLR 2027 Paper Checklist

The checklist scopes the central claim to UAV-ON simulation and points to the primary evidence in Section 4.2, with further mechanism-level evidence in Sections 4.3 and 4.4. It also records evaluation, compute, code, and disclosure details.

Claims

  • The central claim is that explicit BEV representations connected through world-model imagination outperform implicit-memory world models and reactive map-based agents on aerial object search, scoped to UAV-ON simulation.Sections 4.2, 4.3, and 4.4 provide primary and mechanism-level evidence; sim-to-real transfer is identified as a limitation.Source: G ICLR 2027 Paper Checklist · p. 16

Checklist disclosures

  • The checklist marks theory as not applicable, locates limitations after the references, reports AI-assisted editing and code generation checked by two authors, and refers broader impacts to the Ethics statement.The authors state that the research ideas are their own.Source: G ICLR 2027 Paper Checklist · p. 17

Experiment results and evaluation

  • Tables 1, 2, and 4 report final trained configurations evaluated on the full UAV-ON test set; the stated metrics are SR, OSR, and SPL.Results are stratified by object size and scene familiarity.Source: G ICLR 2027 Paper Checklist · p. 17

Infrastructure and code

  • The experimental stack uses PyTorch and Unreal Engine with AirSim via UAV-ON, and the paper provides an anonymous code repository.Training uses eight NVIDIA H20 GPUs; evaluation uses a single NVIDIA RTX 5090 workstation.Source: G ICLR 2027 Paper Checklist · p. 17

Experiment compute

  • A complete three-stage training run takes about 20 hours on eight NVIDIA H20 GPUs, and evaluation takes about 34.5 seconds per episode on one RTX 5090 workstation.Ablation variants repeat the pipeline with noted components removed.Source: G ICLR 2027 Paper Checklist · p. 17

Originality

  • The checklist describes the method as novel and states that, to the authors’ knowledge, no prior work used world models for aerial object search.This is the paper’s stated originality claim.Source: G ICLR 2027 Paper Checklist · p. 17

Limitations

The evaluation is limited to UAV-ON simulation and omits real-world sensing and action errors. Several simulator-provided assumptions, including exact pose and ground-truth depth, remain untested outside that setting.

Evaluation scope

  • All experiments use photorealistic Unreal Engine scenes, with no real-world factors modeled or cross-domain evaluation included.Robustness under sensor noise, pose perturbations, action-execution errors, sim-to-real transfer, and cross-simulator transfer remains unmeasured.Source: Limitations · p. 17

Simulator assumptions

  • SearchWorld assumes exact 6-DoF pose, ground-truth depth, fixed camera intrinsics, and idealized execution of discretized actions.The deterministic memory layers inherit sensing error directly, and actuation dynamics, latency, and control noise are untested.Source: Limitations · p. 17

Representation and policy scope

  • The two-dimensional BEV memory uses a height threshold and abstracts away fine-grained vertical structure such as overhead wires.The value layer receives target-position supervision during training but not the target coordinate at inference; stopping is governed by the imitated task policy.Source: Limitations · p. 17

Related papers

Loading 2610.09335v1…