The one-paragraph version
This paper asks what happens if a robot policy that dreams about the future also learns to dream in 3D. The starting point is a family called World Action Models, which are robot policies built on top of large video-generation models. They look at a language instruction and current camera views, and then jointly generate what will happen next and what the robot should do. Until now those models only dreamed in color images, with no explicit sense of distance or shape. The authors introduce a Spatially Aware World Action Model, called SA-WAM, which adds depth alongside color and actions inside the same generative model. The trick is to squeeze real-world depth into a picture-like format that a frozen, pretrained image compressor already understands, with extra detail saved for close-up manipulation where it matters most. With that change alone, SA-WAM sets new best scores on RoboCasa with very little training data, on the robustness test LIBERO-Plus, and on a real UR5 arm, especially when distractors make color alone ambiguous.
The problem they're solving, and why it matters
Robot policies need to predict both motion through time and contact in space. Two lines of work have approached this from opposite sides. On one side, Vision-Language-Action models, meaning models that take images and language and output robot actions, generalize well from internet pretraining but are often built from still images and text, so they lack a strong prior for how things move and deform. On the other side, video-based policies and World Action Models inherit motion and physics priors from large video generators, but almost all of them operate only on flat color images, called RGB observations. That is a mismatch for manipulation: picking a condiment bottle from a counter into a sink, turning a faucet, or hanging a pink mug depends on where surfaces are in 3D, not just what color they are. Separately, other work has shown that explicit 3D helps policies, for example by using point clouds or 3D position encodings, but that 3D awareness had not been brought into the unified world-and-action diffusion design. The paper tries to close that gap without throwing away the valuable video pretraining.
The key idea, in plain words
Think of the pretrained video model as a fluent speaker of “video latent,” a compressed internal language for image sequences. Instead of teaching it a whole new language for 3D, SA-WAM slips depth in as if it were just more video frames. Concretely, for each camera view the model sees the current color frame and the current depth frame side by side in that compressed language, and it is asked to denoise, meaning to gradually clean up noise into, the future color frames, future depth frames, and the next chunk of robot actions all at once. The hard part is that depth is measured in meters and is in principle unbounded — far walls can be very far — while the frozen compressor expects picture inputs in a fixed bounded range, here from minus one to one. The authors solve this with a carefully chosen squashing function, called normalization. Their preferred choice is log-scale: it stretches out small nearby distances where gripper precision matters, like under 0.3 meters from the wrist camera, while still keeping a usable signal for far background. Because this squashing is reversible, predicted squashed depth can be mapped back to meters for evaluation.
How it actually works, step by step
Start with sensing. There are three views: a wrist camera on the arm and left and right external cameras. For each camera at the current time, there is a color image and a metric depth map, meaning per-pixel distance in meters. There is also the robot’s internal body sense, called proprioception, here joint or end-effector pose plus gripper state, and a language task description such as “pick the condiment bottle from the counter and place it in the sink.”
Next, make depth look like an image. For each camera separately, the authors collect all depth values seen in training and pick robust low and high cutoffs, specifically the 2nd and 98th percentiles, to ignore outliers. In words: clip each new depth map to that camera’s typical range, then apply a monotone, meaning order-preserving, squashing function that maps that range to zero to one. Then rescale zero to one to minus one to one and copy the single value into three color-like channels. The result is a three-channel depth picture that can be fed to the same frozen compressor used for color.
There are three candidate squashings. In words before any formula: linear spreads detail evenly across near and far; inverse-depth spends almost all detail very close up and blurs far away; log-scale is in between. Their sensitivities explain this: linear has constant sensitivity to a change in meters, inverse-depth sensitivity falls off very steeply with distance, and log-scale sensitivity falls off gently, proportional to one over distance.
Then, encode and arrange. A frozen spatiotemporal compressor, called a VAE tokenizer, turns each color and depth picture into compact latent frames. In words, a VAE tokenizer is a pretrained encoder-decoder that compresses images or video into a smaller learned representation and can reconstruct them. SA-WAM places those latents, plus tokens for proprioception and for the action chunk, meaning a short sequence of future actions predicted together, into one long sequence for a Diffusion Transformer, called DiT. A Diffusion Transformer is a transformer network trained to reverse a noising process. Conditioning, meaning the information kept clean and given as context, is the instruction, current proprioception, and current color-plus-depth latents. The targets to be denoised are the action chunk and the future color-plus-depth latents at a horizon ahead, for example 32 steps ahead in simulation with the first 16 executed.
Finally, train with one shared loss. The backbone is initialized from the open-source Cosmos-Predict2 video model, with its Wan2.1 tokenizer kept frozen. Training adds Gaussian noise at different noise levels, following a formulation called EDM, and asks the network to recover the clean action plus future latents. Noise levels are sampled mostly from the base Cosmos distribution and sometimes uniformly to cover large noise. One objective supervises all slots — actions, future color, future depth — with no separate heads or 3D-specific encoder.
What they showed — results, honestly, with limitations
The authors first test which squashing preserves shape through the frozen compressor. They do a round-trip: squash, encode, decode, unsquash, and measure Absolute Relative Error, called AbsRel, meaning average absolute depth error divided by true depth, lower is better. Linear is poor up close at 6.68% error below 0.3 meters on the wrist camera but good farther out. Inverse-depth is excellent up close at 0.62% but collapses far away at 4.16% mid-range and 13.19% far. Log-scale is the balanced winner at 0.63% overall, and that ranking carries over to policy success on an eight-task RoboCasa subset: color-only 52.7%, plus linear depth 56.5%, plus inverse-depth 59.5%, plus log depth 61.7%.
On the full 24-task RoboCasa benchmark, SA-WAM reaches 76.6% average success with only 50 demonstrations per task. For context reported in the paper, the matched color-only Cosmos-Policy baseline gets 67.1% on the same budget, so the gain is 9.5 points and is consistent across categories, with especially large gains on pick-and-place, up 16.7 points from 51.8% to 68.5%, and turn-or-toggle tasks, up 11.2 points from 66.3% to 77.5%. Other prior methods score lower despite using 6 to 20 times more data per task, for example 72.6% for a GR00T variant with World2Act and 71.3% for FLARE, but the paper notes category breakdowns are not available for all baselines, so those comparisons are overall-average only.
On standard LIBERO the task is largely saturated and SA-WAM at 98.4% is described as on par with the strongest policies. The sharper test is LIBERO-Plus, which trains only on LIBERO and tests zero-shot under perturbations to layout, viewpoint, initial state, language, lighting, background, and sensor noise. There SA-WAM gets 86.6% weighted average, ahead of the strongest vision-language-action baseline cited at 84.6% and ahead of Cosmos-Policy at 81.4%.
World-model quality also improves: predicted future color matches true simulator renders better on pixel fidelity, called PSNR, higher is better, structural similarity, called SSIM, higher is better, and perceptual distance, called LPIPS, lower is better, especially for the moving wrist camera and when measured only on robot-plus-object masks. The authors then show prediction quality correlates with success. Using wrist-camera depth error inside the target-object mask around the grasp, successful pick-and-place rollouts have lower 3D error, and a simple threshold detects 80% of failures at 13% false alarms, with area under the curve around 0.88. This is presented as a diagnostic, not proof that better prediction causes success.
Real-world tests on a 6-DoF UR5 arm with a parallel gripper and one fixed RGB-D camera use waypoint actions, meaning absolute end-effector pose plus gripper commands rather than dense joint torques, trained on 20 demos for each of 10 tasks. On four composite categories, SA-WAM reaches 90.0% completion in clean scenes versus 75.0% for Cosmos-Policy and 21.3% for fine-tuned pi-zero, and retains 77.5% under distractors and visual changes versus 48.8% and 11.3%. Qualitative examples, like grapes into a yellow plate then banana into a pink plate, suggest depth helps disambiguate visually similar objects.
Limitations are stated plainly: inconsistent futures and mismatched actions still occur, geometric consistency is not explicitly enforced during training or checked at test time, and diffusion-based world-action models remain expensive at inference. Sample sizes, especially 10 trials per real-world category, are small, so real-world gaps should be read as suggestive rather than precise.
Why this might matter to me
For spatial reasoning in vision-language models, the lesson is that adding an explicit geometric channel can beat asking color latents to implicitly infer depth, especially under visual confusion and viewpoint motion. The wrist-camera result matters: that is where parallax, occlusion, and contact happen, and where a world model most needs grounding.
For world models, SA-WAM offers a pattern for upgrading a 2D video prior without retraining its compressor: keep the backbone and tokenizer frozen, add geometry as extra latent frames, and supervise futures and actions jointly. The correlation between object-centric 3D error and failure also points toward failure detectors and test-time verification — for example, reject or replan when predicted object depth diverges from observed depth.
For text-to-CAD and parametric geometry, the log-scale trick is familiar in spirit: allocate representation where precision pays, preserve invertibility back to metric units, and calibrate per sensor range using robust percentiles rather than global minima and maxima. It is a reminder that unit handling and normalization are first-class design choices when bridging generative latents and physical quantities.
For 3D generation and scene representation, the work suggests dense image-aligned depth is a pragmatic middle ground between flat video and full point clouds or radiance fields: it reuses 2D architectures, stays aligned with pixels for masking and evaluation, yet provides enough 3D signal to improve both action accuracy and future prediction. The open direction the authors flag — geometrically consistent training and verification — is exactly where neural fields, depth, and action generation could meet.
Terms worth knowing
- World Action Model: A policy that uses one generative model to jointly predict future observations and future actions, so visual foresight and control share representations.
- RGB-D observation: A color image plus a per-pixel depth map in meters; here from wrist and two external cameras.
- Proprioception: The robot’s internal state, such as end-effector position, orientation as a quaternion, and gripper opening.
- Action chunk: A short horizon of actions predicted together and partially executed open-loop before replanning; 32 steps with 16 executed in simulation, 3 steps with 2 executed on the UR5.
- Diffusion model / EDM: A generator trained by adding Gaussian noise at various levels and learning to denoise; EDM refers to a specific formulation of noise schedule and weighting used here.
- VAE tokenizer: A pretrained encoder that compresses images or video into compact latents and can decode them back; here the Wan2.1 tokenizer is kept frozen.
- Diffusion Transformer (DiT): A transformer backbone used as the denoiser; here initialized from Cosmos-Predict2, a video model built for physical AI scenarios.
- Normalization in this paper: The reversible squashing of clipped metric depth into the minus-one-to-one range the tokenizer expects; log-scale is the retained choice.
- AbsRel: Absolute Relative Error for depth, meaning absolute error divided by true depth, lower is better.
- PSNR / SSIM / LPIPS: Standard image-prediction scores: pixel-level fidelity, structural similarity, and learned perceptual distance; higher is better for the first two, lower is better for LPIPS.