For embodied agents to perceive and act intelligently in the physical world, they must understand its 3D structure and predict how it will evolve over time. However, visual observations provide only partial 2D projections of the world, each capturing a limited set of viewpoints, visible surfaces, and changes over time. Robust 4D scene understanding therefore requires inferring what is not directly observed: 3D geometry, camera viewpoint, scene appearance from unseen views, and scene evolution.
In this work, we introduce UGG, a unified generative geometry model that treats visual generation and geometric prediction as different conditional queries of the same underlying 4D scene representation. Appearance (RGB video), 3D structure (depth), and viewpoint (camera geometry) provide complementary observations of a dynamic scene. Depending on what is observed, a joint model can infer missing scene components.
Our generative formulation also enables inference-time search over multiple plausible scene completions. Instead of committing to a single geometric prediction, we sample multiple candidates and select the one with the strongest cross-modal consistency, measured through photometric reprojection. This inference-time search capability is also consistent with the cognitive theory of perception as inference: human scene understanding requires inferring latent scene causes from incomplete observations, and may rely on internal simulation to evaluate possible scene explanations.
Overall, our contributions are threefold: (i) We formulate 4D scene understanding as joint generative modeling over appearance and geometry, treating visual generation and geometric prediction as missing-modality denoising within a single conditional model. (ii) We introduce a unified generative-predictive model that represents RGB, depth, and camera pose as spatially aligned modalities, enabling arbitrary conditioning and learning from heterogeneous datasets with incomplete supervision. (iii) We leverage the generative formulation for inference-time search over multiple 4D scene hypotheses, using geometric consistency to improve geometric prediction.
Method
We propose a unified generative-predictive framework for 4D scene understanding. Given partial observations of a scene, our goal is to infer a complete scene representation containing appearance, depth, and camera geometry. We formulate both visual generation and geometric prediction as conditional completion in a joint generative model over spatially aligned RGB, depth, and camera-ray modalities. We specify each task by a noise pattern over multimodal tokens, with observed tokens as zero-noise context and missing tokens as noisy targets. The framework supports heterogeneous training data with incomplete modality supervision, arbitrary partial conditioning, and test-time search over multiple 4D scene hypotheses.
Unified 4D Scene Representation
The resulting scene representation, a joint representation of appearance and local 3D geometry, contains the appearance and local 3D geometry for each frame, a frame-indexed scene representation.
Generative-Predictive Transformer
Given the spatially aligned representation, the joint RGB, depth, and camera-ray representation, we design a multimodal diffusion transformer that jointly processes appearance, depth, and camera-ray geometry.
Each token is processed with modality-specific patch embedding layers and a spatiotemporal RoPE positional encoding over latent spatial coordinates and the frame index. If the VAE encoder performs temporal downsampling, the number of latents and frames could differ; we omit this from our notation for simplicity and continue to refer to the number of frames or latents as the sequence length.
At the output stage, RGB and depth tokens are decoded back to pixel space using the VAE decoder. Ray-map tokens are directly mapped back to dense ray maps, which can be converted back to camera parameters, including rotation, translation, and intrinsics.
Conditional Completion with Noise as Masking
At zero diffusion time, the token is fully clean, while at the terminal diffusion time, it is pure noise. By varying the noise levels across both modality and temporal position, we realize different conditional generation and prediction tasks without modifying the model architecture.
Inference-Time Search
The generative formulation also enables inference-time search, in contrast with a deterministic predictor that commits to one scene. Given the same observed context, we sample multiple candidate completions from different random initializations. Each sample yields a complete 4D scene hypothesis.
To select among these hypotheses without ground-truth depth or camera pose, we evaluate how well each hypothesis explains geometric evidence extracted from the observed RGB frames. Full equations and hyperparameters are given in the appendix.
We extract pixel correspondences between observed RGB frames using RAFT. To remove issues with occlusions and moving objects, we keep correspondences only if forward mapping followed by backward mapping returns a pixel to its original point. For dense observations, we select frame pairs with informative pixel motion, adapting the temporal separation according to the observed optical flow magnitude. For sparse observations, both generation and verification use only the first and last observed RGB frames. These correspondences provide geometric constraints against which candidate depth and camera predictions can be evaluated.
Camera Pose Estimation
We evaluate camera pose estimation from RGB observations in dense and sparse settings, conditioning on either the full video or only the first and last frames. Camera-ray tokens are denoised and converted back into camera parameters, testing whether camera geometry can be predicted as a missing modality within the same conditional-completion framework.
Depth Prediction
We evaluate depth prediction from RGB observations in the dense and sparse settings described in the previous section. In both cases, depth tokens are initialized from noise and denoised, testing whether the unified conditional-completion model learns a useful geometric posterior over scene depth.
Baselines. We compare with Geo4D, ChronoDepth, Gen3R, Aether, and Marigold. These baselines include specialized depth and 4D geometry models, as well as a unified model, while ours is trained for RGB-depth-camera completion.
Test-Time Search on Geometric Prediction
We evaluate whether test-time search (TTS) improves geometric prediction by sampling multiple RGB-depth-camera completions at inference time. For each input, we generate 16 candidates from different noise initializations and select depth and pose independently using their respective verification scores, as described in the inference-time search method.
UGG (TTS Oracle) selects one of the same 16 candidates per input by minimizing ground-truth scale-and-shift-aligned Abs Rel for depth, ATE for dense pose, or tAE for sparse pose. All metrics are reported for the selected candidate. Overall, these results show that inference-time search is a promising benefit of the generative formulation, while leaving room for stronger search criteria and verification.
Conclusion
We present a unified generative geometry model for 4D scene understanding, formulating visual generation and geometric prediction as conditional completion over RGB, depth, and camera rays. Through noise patterns over modality tokens, a single model supports camera-controlled generation and geometric prediction under dense and sparse observations, while learning from heterogeneous datasets with varying modality availability. Across diverse datasets, UGG achieves strong video generation, depth prediction, and pose estimation performance, and further improves geometric consistency through test-time search. These results suggest that jointly modeling appearance and geometry with generative models is a promising direction toward more general 4D scene understanding.
Our current model is not uniformly stronger than specialized methods: task-specific NVS models still perform better on pixel-aligned metrics, and specialized geometry models remain stronger on several indoor depth metrics. Future work could address these gaps with better dataset balancing, stronger geometric supervision, and improved depth-camera consistency during training. Our test-time search also relies on RGB reprojection error, which is useful but not uniformly reliable; learned or multimodal consistency scores may provide stronger search criteria. Finally, training at larger scale, higher resolution, and with more diverse data could further improve both visual detail and geometric accuracy.
Model Architecture
We initialize from Wan2.2-TI2V-5B and extend its diffusion transformer into a multimodal Mixture-of-Transformers backbone. The pretrained video VAE and text encoder are kept frozen, and we finetune only the diffusion transformer. The model jointly processes four latent streams corresponding to RGB video, ray directions, ray moments, and depth, respectively. The two ray latents are concatenated along the channel dimension and embedded as a single ray token stream, so the transformer processes three token streams: RGB, ray, and depth.
Training Objective and Diffusion Forcing
We train with continuous flow matching. We sample raw diffusion times for the context and target groups, as described below, and apply the shifted schedule.
Given a clean latent the uncorrupted representation and Gaussian noise randomly sampled noise, the noisy input is a mixture of the clean latent and Gaussian noise, and the model predicts the flow target the target direction of change from noise toward the clean latent.
The loss is the mean-squared error between the predicted and target flow, computed separately for RGB, ray direction, ray moment, and depth. For samples without ray-map or depth supervision, the corresponding losses are masked out.
Inference
At inference time, the same clean/noisy pattern specifies the task. Observed RGB, depth, or ray-map tokens are clamped to their clean values at every denoising step and assigned zero diffusion level. Unknown tokens are initialized from Gaussian noise and progressively denoised. We use UniPC sampling with a specified number of denoising steps and the same shifted time schedule used during training. Under this setting, the denoising loop for one 50-frame clip takes approximately 20 seconds on a single NVIDIA H100 or H200 GPU.
Camera Pose Evaluation
In the dense setting, models condition on the full RGB sequence. Following Geo4D, we align predicted trajectories to ground truth using Sim(3) Umeyama alignment. We report absolute trajectory error (ATE), relative pose error in translation (RPE t), and in rotation (RPE r). ATE measures global trajectory agreement after alignment, while RPE measures local frame-to-frame motion consistency.
Acknowledgments and Disclosure of Funding
We thank the Kempner Institute at Harvard University for providing computational resources that supported this work. Ayush Tewari acknowledges the NVIDIA Academic Grant Program.
Frequently Asked Questions
What does UGG model? UGG jointly models RGB appearance, depth, and camera-ray geometry as spatially aligned modalities for 4D scene understanding.
How does UGG handle missing scene information? It uses conditional completion, treating observed tokens as clean context and missing tokens as noisy targets to be denoised.
How does test-time search choose a scene prediction? It samples 16 candidate completions and selects depth and pose predictions using verification scores based on geometric evidence from observed RGB frames.
Which metrics are used for camera pose evaluation? The paper reports absolute trajectory error and relative pose error in translation and rotation.