OccSTeP-WM: Tokenizer-Free 4D Scene Forecasting

• OccSTeP-WM is a tokenizer-free world model that integrates dense voxel embeddings with linear-complexity attention and recurrent state modules for 4D occupancy forecasting.
• It combines reactive forecasting (predicting imminent scene evolution) with proactive, action-conditioned forecasting to handle noisy and incomplete sensor inputs.
• The model employs SE(3) warping, gated fusion, and a lightweight 3D-UNet decoder, achieving notable improvements in occupancy IoU and semantic mIoU over previous methods.

OccSTeP-WM is a tokenizer-free world model designed for spatio-temporal persistence in 4D occupancy forecasting, particularly for autonomous driving scenarios that demand robust, temporally persistent scene understanding under sensor disturbance and future action conditioning. It incrementally fuses dense voxel-based scene states across time using a linear-complexity attention backbone and a recurrent state-space module with ego-motion compensation, enabling both reactive ("what will happen next") and proactive ("what would happen given a specific future action") forecasting. OccSTeP-WM provides robust, online inference even when historical inputs are missing or noisy, and it has shown substantial gains over prior methods in challenging scenarios (Zheng et al., 17 Dec 2025).

1. Core Forecasting Objectives and Formulation

OccSTeP-WM addresses two complementary tasks:

• Reactive forecasting: Given observed sensor histories $X_{1:t}$ and ego-poses $P_{1:t}$, it predicts the imminent scene occupancy grids $\hat O_{t+1:t+T}$ and the "most likely safe" future ego-motion $\hat P_{t+1:t+T}$, formalized as

$$(\tilde X_{t+1:t+T}, \hat P_{t+1:t+T}) = \mathcal{W}(X_{1:t}, P_{1:t})$$

• Proactive forecasting: Conditioned on $X_{1:t}$, $P_{1:t}$, and a user-specified future ego-motion $P_{t+1:t+T}$, it predicts the counterfactual occupancy $\tilde O_{t+1:t+T}$, given by

$$\tilde X_{t+1:t+T} = \mathcal{W}(X_{1:t}, P_{1:t}, P_{t+1:t+T})$$

This architectural duality enables modelling both the passive evolution of scenes and action-conditioned counterfactuals, a requirement for planning and robust autonomy (Zheng et al., 17 Dec 2025).

2. Voxel-Based Scene Representation and Embedding

The core scene representation is a dense semantic occupancy tensor $O_t \in \{0,\ldots,K{-}1\}^{D \times H \times W}$, where each voxel index maps to a semantic class ($0$ denotes free, $1\ldots K{-}1$ for semantic categories). Feature construction proceeds as follows:

• Each class index $c=O_t(d,h,w)$ is mapped by a learnable embedding $E \in \mathbb{R}^{K \times C_e}$, and combined with a fixed 3D Fourier positional code $P\in\mathbb{R}^{D\times H\times W\times C_p}$:

$$X_t (d,h,w,:) = [E_c,\,P_{d,h,w}] \in \mathbb{R}^{C_e+C_p}$$

• The resulting tensor $X_t \in \mathbb{R}^{D\times H\times W\times C}$ is flattened to a sequence of length $L = D\cdot H\cdot W$ by a tiled Morton (Z-order) permutation $\pi$ to preserve spatial locality:

$$\pi: \mathbb{R}^{D\times H\times W\times C} \to \mathbb{R}^{L\times C}$$

This tokenizer-free embedding enables direct dense scene encoding without reliance on discrete semantic tokens, fostering robustness against typical semantic perturbations (Zheng et al., 17 Dec 2025).

3. Linear-Complexity Attention Backbone

Long-range spatial dependencies are captured efficiently by a linear-complexity (“Mamba”) attention backbone, which replaces quadratic self-attention with a state-space model (SSM).

• Standard attention for sequence $X$: $O(L^2\cdot d)$ complexity.
• In Mamba:

$$h_n = \exp(-\operatorname{softplus}(A)\Delta t)\odot h_{n-1} + (1 -\exp(-\operatorname{softplus}(A)\Delta t))\odot B\odot(W_{in} x_n)$$

$y_n = C\odot h_n + D x_n$

Each token update is $O(1)$ in sequence length $L$, hence total cost is $O(L\cdot d)$. This facilitates tractable scene reasoning over high-resolution voxel grids. Two Mamba blocks are used: a pre-fusion encoder ($\text{MB}_{\text{pre}}$) and a post-fusion encoder ($\text{MB}_{\text{post}}$), enabling progressive spatial context refinement (Zheng et al., 17 Dec 2025).

4. Incremental Spatio-Temporal Priors Fusion (ISTPF)

Temporal scene memory is managed by a recurrent state-space module, maintaining a hidden voxel grid $S_t \in \mathbb{R}^{C_h\times D\times H\times W}$. At each timestep:

• SE(3)-warping: Prior state $S_{t-1}$ is aligned to the current ego frame by trilinear resampling under the estimated transform $T_{t \to t+1}\in$ SE(3):

$\tilde S_t = \mathcal{Q}(S_t, T_{t\to t+1})$

• State update with gating and exponential forgetting: Updates use learned per-channel decay/mix weights:

$$\alpha = \exp(-\operatorname{softplus}(A)\odot\operatorname{softplus}(\Delta t)),\quad \beta=(1-\alpha)\odot B$$

$S_{t+1} = \alpha\odot\tilde S_t + \beta\odot X_t^h$

$$Y_c = C\odot S_{t+1},\quad Y_t = W_{out}(Y_c)\odot G_t + X_t^{skip}\odot(1-G_t)$$

The only persistent memory is $S_t$, resulting in $O(1)$ per-frame state requirements. This architecture supports robust, incremental fusion even under missing or corrupted frames, a property central to the OccSTeP benchmark (Zheng et al., 17 Dec 2025).

5. Spatio-Temporal Fusion and Handling Corruptions

OccSTeP-WM's design incrementally fuses new information while preserving and warping prior context. The process is as follows:

1. SE(3) warping of the prior hidden state.
2. Gated fusion of the current Mamba-projected features.
3. Post-fusion refinement using a second Mamba block.
4. Decoder: A lightweight 3D-UNet upsamples and sharpens voxel-wise predictions.

Robustness mechanisms:

• Discontinuous frames: When sensor frames are dropped, compound transforms are composed, preserving updates across variable time intervals.
• Fragmentary sensor input: Missing LiDAR or RGB views yield sparser voxelizations, but upstream fusion compensates.
• Reductive (semantic label swaps): The gating mechanism enables the model to discount unreliable new labels and rely on persistent memory.

These mechanisms yield resilience against typical perception corruptions encountered in autonomous driving (Zheng et al., 17 Dec 2025).

6. Forecasting Pipelines and Learning Objectives

• Reactive forecasting operates as an autoregressive loop, predicting both grid occupancy and future ego-motion updates.
• Proactive forecasting applies the same architecture, conditioning forward prediction on exogenously specified ego-motions.

The per-frame loss objective is:

$$L = \lambda_{\text{sem}}\cdot\mathrm{CE}(Z^{\text{sem}}_t, Y_t) + \lambda_{\text{pos}}\cdot\mathrm{SmoothL1}([\hat x,\hat y],[x,y]) + \lambda_{\text{rot}}\cdot \|\operatorname{wrap}(\widehat{\Delta\psi}) - \operatorname{wrap}(\Delta\psi)\|_1$$

where $\mathrm{CE}$ denotes voxel cross-entropy, and typical weights are $$\lambda_{\text{sem}}=1.0,\ \lambda_{\text{pos}}=0.1,\ \lambda_{\text{rot}}=0.1$$.

Forecasting proceeds either via model-generated or externally-provided ego-motion sequences, with metrics computed on per-voxel semantic and geometric accuracy (Zheng et al., 17 Dec 2025).

7. Evaluation, Results, and Performance Summary

Evaluation uses the Occ3D dataset and OccSTeP benchmarks, computing:

• Occupancy IoU: $$\operatorname{IoU} = |V_{\text{pred}}\cap V_{\text{gt}}| / |V_{\text{pred}}\cup V_{\text{gt}}|$$
• Semantic mIoU: $$\operatorname{mIoU} = \frac{1}{K}\sum_{c=1}^K \frac{TP_c}{TP_c + FP_c + FN_c}$$

Reported results:

• Proactive pipeline: $\operatorname{mIoU} = 23.70\%$ ($+6.56$ pp), $\operatorname{IoU} = 35.89\%$ ($+9.26$ pp) over previous baselines.
• Robustness under benchmark-specific corruptions is improved, with up to $+12.86$ pp IoU gain on the 'Reverse' scenario (Zheng et al., 17 Dec 2025).

Summary Table: Core Components and Functions

Component	Role	Complexity
Voxel grid	Scene state, semantic encoding	$O(DHW)$ mem
Tokenizer-free embed	Dense feature mapping + 3D position	$O(DHW)$ compute
Mamba backbone	Long-range spatial context	$O(DHW)$
ISTPF module	Spatio-temporal memory via state gating	$O(1)$/frame mem
3D-UNet decoder	Semantic and geometric refinement	$O(DHW)$

OccSTeP-WM delivers an incremental, SE(3)-equivariant, and memory-efficient world model, advancing the state-of-the-art in 4D occupancy forecasting across scenarios with noisy or incomplete historical data, while supporting both reactive and action-conditioned future inference (Zheng et al., 17 Dec 2025).