OMEGA: Efficient Occlusion-Aware Navigation for Air-Ground Robot in Dynamic Environments via State Space Model

In plain words

A robot that flies and drives has to avoid obstacles it cannot see yet, such as people walking behind a wall. It builds a map from its sensors, and the parts hidden behind other things are guessed. The paper's first piece, OccMamba, is a network that takes a point cloud from a laser scanner and guesses which parts of the surrounding space are filled and what each filled part is. It handles the two questions separately, one branch for what things are and one for where the solid space is, and then merges the two views from above. The network uses a newer sequence-model layer, called Mamba, which keeps the cost low while still letting distant parts of the scene influence each other.

The second piece is a planner. Many planners first build a distance map of the whole scene and then search it, which takes most of the planning time. This one starts from a rough path, searches for a cheaper route only where the path hits something, and then smooths the result with an optimizer. It also prefers ground driving when that saves energy and switches to flight only when needed.

The authors report that the perception network beats the strongest prior network on a standard driving benchmark, and that the full system succeeds in most trials in simulated and real dynamic scenes. The perception code is public; the planner code is not in the linked repository.

TL;DR

• OMEGA is an occlusion-aware navigation system for air-ground robots. Its perception part, OccMamba, predicts 3D semantic occupancy from LiDAR, and its planner, AGR-Planner, avoids building an ESDF map §1 fig. 2.
• OccMamba has a semantic branch and a geometric branch, each with Mamba blocks, and a BEV U-Net that fuses the two. Training uses cross-entropy plus Lovász losses, with an auxiliary head per stage §3.1 §3.2 fig. 3.
• Author-reported on SemanticKITTI hidden test: IoU 59.9, mIoU 25.0, 22.1 FPS tab. 2. Validation set: IoU 58.6, mIoU 25.2 tab. 3.
• Author-reported navigation: 98% planning success over 200 simulated trials §5.3 and 96% average success in five real dynamic scenes, ten runs each §5.4.
• Only the OccMamba code is in the linked repository. The planner is described in the paper but not released there.

Problem

Air-ground robots (AGRs) navigate with a local map. Occluded areas must be predicted before they can be planned around §1. The authors argue that prior systems do two things poorly in dynamic scenes: their perception networks are slow or inaccurate on moving obstacles, and their planners spend most of the local planning time building an ESDF map, which the authors place at about 70% of local planning time in prior systems §1.

The paper assumes LiDAR point clouds as input for the perception network §3.1 §5.1 and an onboard computer that can run a TensorRT-optimised network at real-time rates §5.1. The real robot used a depth camera, a visual-inertial tracking camera and a Jetson Xavier NX §5.1.

Method

Overview. OMEGA has two modules. OccMamba predicts the occupancy of the local space, and the predicted occupied cells are merged into the local map with a query-based update method borrowed from AGRNav §1 fig. 2. AGR-Planner then plans on that map. The two run asynchronously on the onboard computer fig. 2.

OccMamba: branches. The network has three branches §3 fig. 3.

• Semantic branch §3.1. A voxelisation layer assigns each point to voxels at several resolutions and max-pools the point features into a voxel feature. Three encoder stages follow. The paper does not name the layer type; the code uses sparse 3D convolutions semantic_segmentation.py. After each stage the voxels are projected to a bird's-eye-view (BEV) grid by taking the maximum over voxels that share a BEV index, and the sparse BEV grid is densified. A Sem-Mamba block then runs on the dense BEV features.
• Geometric branch §3.1. A 7×7×7 input convolution is followed by three Geo-Mamba blocks. Each block has a residual 3D block and a Mamba module on the voxels. The dense 3D features are then flattened along the height axis into BEV features and compressed with 2D convolutions.
• BEV fusion branch §3.1. A U-Net with 2D convolutions takes the fused BEV features. At each of the first three encoder resolutions, a feature fusion module takes the output of the previous stage together with the semantic and geometric features at that scale. The decoder upsamples with skip connections and ends in a convolution that predicts the occupancy classes §3.1.

Mamba layer. The paper writes the state-space model as a continuous system, discretises it with a time step Δ, and gives the recurrence and the equivalent global convolution kernel eq. 1 eq. 2 eq. 3 eq. 4. It states that Mamba uses data-dependent SSM parameters §2, but the equations shown are the fixed (time-invariant) form.

# OccMamba forward pass  [§3] [fig. 3]
# Input:  points (N, 4) as x, y, z, intensity; occupancy (B, D, H, W)
# Output: class logits (B, C, D, H, W), with D = 32 height bins

def occmamba(points, occupancy):
    # Semantic branch: max-pooled voxel features  [§3.1]
    voxel_feats = voxelize_max_pool(points)          # (N_vox, C)
    sem_bev = []
    for stage in semantic_stages:
        # Sparse 3D convolutions over the voxels, three stages  [§3.1]
        voxel_feats = stage(voxel_feats)
        # Max-pool voxels that share a BEV index, then densify  [§3.1]
        bev = max_pool_to_bev(voxel_feats)           # (B, C, H, W)
        # Sem-Mamba block on the dense BEV features  [§3.1]
        bev = sem_mamba(bev)
        sem_bev.append(bev)
    # Geometric branch: dense 3D convolutions, Geo-Mamba after each  [§3.1]
    geo_bev = []
    voxels = occupancy
    for stage in geometry_stages:
        voxels = stage(voxels)
        voxels = geo_mamba(voxels)                   # Geo-Mamba  [§3.1]
        # Flatten the height axis into channels to form BEV  [§3.1]
        geo_bev.append(flatten_height(voxels))
    # BEV U-Net fuses the three scales of both branches  [§3.1]
    # The code also concatenates the input occupancy with the dense
    # BEV features; the paper does not say so  [code: networks/dsc.py#L66-L72]
    bev_input = concat(occupancy, dense_bev_of_points)
    return bev_unet(bev_input, sem_bev, geo_bev)  # (B, C, D, H, W)

OccMamba: losses. The semantic branch uses cross-entropy plus Lovász loss at each of its three stages eq. 5. The completion (geometric) branch uses binary cross-entropy plus Lovász loss at each of its three stages eq. 6. The BEV output uses three times the sum of cross-entropy and Lovász loss eq. 7. The total is the sum of the three terms eq. 8. Auxiliary MLP heads sit after each encoder block of the semantic and completion branches during training only, and are removed at inference §3 fig. 3.

# Training loss  [§3.2] [eq. 5] [eq. 6] [eq. 7] [eq. 8]
# Input:  per-stage logits of each branch; BEV logits; labels
# Output: scalar total loss

def total_loss(sem_logits, com_logits, bev_logits, labels):
    # Semantic branch: cross-entropy plus Lovász at each stage  [eq. 5]
    loss_sem = 0
    for logits, target in zip(sem_logits, sem_targets):
        loss_sem += cross_entropy(logits, target)
        loss_sem += lovasz_softmax(logits, target)
    # Completion branch: binary cross-entropy plus Lovász  [eq. 6]
    loss_com = 0
    for logits, target in zip(com_logits, occ_targets):
        loss_com += binary_cross_entropy(logits, target)
        loss_com += lovasz_softmax(logits, target)
    # BEV output: three times (cross-entropy plus Lovász)  [eq. 7]
    loss_bev = 3 * (cross_entropy(bev_logits, labels)
                    + lovasz_softmax(bev_logits, labels))
    # The network is trained end to end on the sum  [eq. 8]
    return loss_bev + loss_sem + loss_com

AGR-Planner: initial trajectory and guidance. The planner builds on EGO-Swarm §4. It first makes an initial trajectory whose inner control points are random and anchored at the start and goal, ignoring obstacles §4.1. For each control-point segment that collides, it runs a kinodynamic A* search over motion primitives, with an extra energy cost for flight primitives, and uses the result as a guidance segment §4.1. The obstacle distance is taken from the control point to the obstacle surface along a safe direction vector eq. 11. A second anchor point is found along the opposite direction, so that several candidate paths can be compared §4.1.

AGR-Planner: trajectory optimisation. The trajectory is a uniform B-spline with degree p_b and N_c control points §4.2 eq. 9. Velocity, acceleration and jerk control points are finite differences of the control points, divided by the knot spacing Δt eq. 10. In ground mode the control points are two-dimensional, and in aerial mode they are three-dimensional §4.2. The cost is a weighted sum of four EGO-Swarm terms, smoothness (s), collision (c), dynamic feasibility (d) and terminal progress (t) eq. 12, plus a curvature term that penalises ground curvature above a threshold C_max eq. 13 eq. 14 eq. 15. The problem is solved with NLopt §4.2. Tracking and mode switching reuse the methods of AGRNav §4.2.

# AGR-Planner: initial guess, kinodynamic A*, then refinement  [§4]
# Input:  local map (OccMamba output); start and goal positions
# Output: smooth B-spline trajectory

def plan(local_map, start, goal):
    # Random inner control points anchored at start and goal  [§4.1]
    initial = random_control_points(start, goal)
    # Replace each colliding segment by a kinodynamic A* guidance  [§4.1]
    guidance = kinodynamic_astar(local_map, initial, energy_cost=True)
    # Refine B-spline control points by minimising J_AGR with NLopt  [§4.2]
    trajectory = minimize_nlopt(j_agr, initial, guidance)
    return trajectory

Hyperparameters

Values marked "paper" are stated in the text. Values marked "code only" are in the repository and not in the paper; they are the defaults the code runs with.

Results

Author-reported. Setup: SemanticKITTI, LiDAR-only OccMamba, and navigation in simulation and five real dynamic scenes §5.

Table 5 is the only source for the perception-planner ablation. The abstract's "96% average planning success rate" refers to the real-world runs, not the simulation.

Paper vs code

Official code: the abstract links the project page, and the project page links https://github.com/jmwang0117/Occ-Mamba §1 (project page). The owner's account also hosts the group's AGRNav repository, which the README links README.md. The repository has no LICENSE file, so the code is not licensed for reuse. We did not need a GitHub search, because the project page gives the link. The AGR-Planner is not in the repository. The README says only that OccMamba's code was released README.md.

Matches: the three-branch structure, the 80 epochs, the batch size of 6, the learning rate of 0.001, the grid of 256 × 256 × 32 at 0.2 m, the Geo-Mamba layer after each of the three geometric stages, the BEV loss weight of 3 and the total-loss sum all agree with the code dsc.py completion.py template.yaml template.yaml.

Differences:

• Weight decay: the paper states none. The code uses AdamW with weight decay 0.05 hard-coded optimizer.py; the config file's value is NA.
• Learning-rate schedule: the paper states none. The code multiplies the rate by 0.98 each epoch optimizer.py train.py.
• Class weights: the paper gives plain cross-entropy eq. 5 eq. 7. The code weights the semantic and BEV cross-entropy by 1 / log(f + 0.001), with f from the frequency table in the config dsc.py semantic_segmentation.py. The completion branch uses no weights completion.py.
• Completion loss: the paper says binary cross-entropy eq. 6. The code's completion head produces two channels and applies softmax cross-entropy completion.py completion.py.
• Input to the BEV U-Net: the code concatenates the input occupancy grid with the dense BEV features dsc.py. The paper describes only the fused BEV features.
• BEV U-Net depth: the paper says the decoder upsamples three times §3.1. The code has four upsampling stages bev_net.py bev_net.py.
• Encoder blocks: the paper says the encoder uses residual blocks §3.1. The code's encoder blocks are plain double convolutions with max-pooling bev_net.py bev_net.py.
• BEV fusion module: the paper cites a module from [mei2023ssc] §3.1. The code uses a class named BEVFusionv1, which gates each branch with sigmoid channel attention computed from pooled features and sums the results bev_net.py. We did not check this against the cited source.
• Mamba layer: the code uses the mamba_ssm Mamba module with d_state 16, d_conv 4 and expand 2, preceded by an instance norm and followed by a 1×1 3D convolution completion.py semantic_segmentation.py. The Sem-Mamba layer runs on each BEV map with depth 1 semantic_segmentation.py.
• Voxel input features: the paper says only "max pooling" §3.1. The code uses 10 features per point: the four input values, plus two 3-vectors, one offset from the voxel point mean and one from the voxel centre. The features are max-pooled per voxel preprocess.py preprocess.py.
• Augmentation: the paper says random x-y flips §5.1. The code draws one of four options per sample (none, x, y, both) and applies it to voxels, labels and points semantic_kitti.py semantic_kitti.py semantic_kitti.py.
• Checkpoint: the README says the checkpoints used for testing were the best on the validation set during training README.md. The test script's example uses epoch 76 run_test.sh.
• Unused code: utils/ssc_loss.py defines geo_scal_loss, sem_scal_loss and CE_ssc_loss, but no call site in the training path uses them.
• Config: run_train.sh points to cfgs/2024.6.16.yaml, which is not in the repository run_train.sh. Only cfgs/template.yaml is present.
• Planner: the released code has no planner, so AGR-Planner cannot be checked against code.

Pitfalls

• The planner cannot be reproduced from this card. The paper gives no values for the cost weights λ, the curvature threshold C_max, the barrier threshold ζ, the B-spline degree p_b, the number of control points N_c, Δt, the energy cost of flight primitives, the A* heuristic or the NLopt algorithm §4.2. The definitions of J_s, J_c, J_d and J_t come from EGO-Swarm and are not restated §4.2 eq. 12. Implement them from EGO-Swarm, and treat the rest as unknown.
• Eq. 11 is ambiguous. It writes D_ij = (Q_i − p_ij) × v_ij, with ×. Read as a cross product, this is a vector, not a distance. A dot product or a norm is probably meant. The paper does not say.
• Curvature term. The angle change uses the arctangent of Δy/Δx without a quadrant-aware form eq. 14 §4.2. The paper does not define ΔQ_ti, and C_max is not stated.
• B-spline derivatives eq. 10. The velocity, acceleration and jerk formulas are plain differences divided by Δt. The paper gives no degree factor and does not say whether one is intended.
• Index conventions. Control points are written Q_1 … Q_{N_c} in one place and Q_{t0} … Q_{tM} in another, with M = N + p_b. The sum in eq. 13 runs to M − 1. Eq. 3 uses h(t) on the left and h_t on the right. In eq. 4 the length m is not defined. Read the indices from the intended algorithm, not the typeset symbols.
• Mamba is described two ways. The text says Mamba uses data-dependent parameters §2, and §3.1 presents the fixed-parameter SSM equations eq. 1 eq. 4. The code uses the mamba_ssm selective module completion.py. Implement the selective form, and check the scan order: the code flattens the 3D voxels in row-major order before the scan completion.py. The paper does not state a scan order.
• Completion loss form. The paper's completion loss is binary cross-entropy eq. 6, and the code uses two-channel softmax cross-entropy. Use the code's form if you want to match the released numbers.
• Class weights change the semantic and BEV loss. Using plain cross-entropy, as the paper's equations suggest, may not reproduce the code's behaviour. The weights are computed from the frequency tables in the config template.yaml.
• Numbers that do not match the paper's own tables.
  • §5.2 says OccMamba beats Co-Occ in IoU by 5.8%. Table 2 gives 59.9 − 56.6 = 3.3 points.
  • §5.2 says the model weighs 23.8 MB. Table 3 gives 23.8 M parameters.
  • §5.2 says removing Sem-Mamba drops mIoU by 4.37%. Table 4 gives 25.2 → 24.1, which is 1.1 points, or 4.37% relative.
  • §5.3 says OccMamba improves planning success by "4% and 1%", "respectively". Table 5 gives 98 − 95 = 3 points over the SCONet row and 98 − 94 = 4 over the no-perception row.
  • §5.3 says replacing SCONet in AGRNav with OccMamba adds 2%. Table 5 has no row for AGRNav's own planner, so this cannot be checked.
  • §5.3 cites AGRNav as fan2019autonomous. That key is for HDF. AGRNav is wang2024agrnav.
  • §1 says energy is saved by about 18%. §5.4 gives 22.1% for scene B and 17.28% for scene C.
• Units and symbols. The text says "0.2m^3" for each voxel §5.1. The config makes this a 0.2 m edge template.yaml.
• Test vs validation numbers. The hidden-test numbers tab. 2 come from the benchmark server. Only validation numbers tab. 3 can be reproduced locally. The validation and test IoU differ, 58.6 against 59.9, and so do mIoU, 25.2 against 25.0.
• Wrong attachment in the source bundle. sup.tex in the source bundle is the supplement of a different paper (Co-Occ). It is not this paper's appendix, and it was not used.
• Version. The source read is the arXiv v2 LaTeX (updated 2024-12-05), the camera-ready version. Figure and table numbers were counted from the LaTeX order. Equation numbers were counted the same way. Check them against the PDF if the float placement differs.
• Figure-only results. The per-method comparisons in simulation, including the planning-time comparison, are in Figure 6 and not in a table §5.3 fig. 6. The text gives 98% success and a 89.33% reduction in planning time, citing Figure 6c §5.3.
• Hardware for FPS. The paper reports 22.1 FPS but does not say which GPU or precision produced it tab. 2.
• Real-world platform name is blank. §5.1 says "a for real-time onboard computation", with the platform name missing. It later names a Jetson Xavier NX.
• Energy measurement method in the real-world tests is not described.
• Cost the abstract hides. Table 3 shows 505.1 GFLOPs for OccMamba on the validation set, close to MonoScene's 501.8 GFLOPs. The memory figure is per training batch §5.2. The paper's linear-complexity claim refers to the Mamba layers §3.1; it gives no complexity for the whole network.

Lineage

• Builds on: AGRNav [wang2024agrnav] for the occlusion-aware map update, the B-spline matrix form, tracking and control, and the real-world setup §1 §4.2 §5.1.
• Builds on: EGO-Swarm [zhou2021ego] for the planner's cost terms and the obstacle-distance approximation §1 §4.2.
• Builds on: the state-space model S4 [gu2021efficiently] and Mamba [gu2023mamba, dao2024transformers] §2.
• Compared against: MonoScene, OccFormer, VoxFormer, TPVFormer, LMSCNet, SSC-RS, SCONet, M-CONet, Co-Occ on SemanticKITTI tab. 2 tab. 3; HDF, TABV, M-TABV, AGRNav and HE-Nav in navigation tab. 1.