Uranus: Building the Next-Generation Simulation Infrastructure for Embodied AI

AI summary

Overview

• Research area: Robotics — data-driven (learned) simulation for embodied AI, specifically video world models that act as robot simulators.
• Technical level: Advanced. The paper assumes familiarity with diffusion models, flow matching, diffusion transformers (DiT), autoregressive generation, KV caching, and robot kinematics (forward kinematics, MJCF/URDF, qpos).
• Scope: A technical report describing Uranus, a joint-trajectory-conditioned autoregressive diffusion simulator for multi-view robot interaction, together with its architecture, data curation pipeline, training strategy, and serving infrastructure.

What This Paper Is About

Simulation is needed to generate robot data, train policies, run evaluations, and iterate safely, but real-world interaction is costly, slow, and hard to reproduce, while conventional simulators like Isaac Sim and MuJoCo require heavy manual construction of assets, environments, and physics models, and still suffer a sim-to-real gap. Uranus takes a data-driven alternative: rather than hand-specifying physics, it learns to predict future multi-view observations directly from real robot interaction data, conditioned on a supplied future joint-position trajectory. The goal is a practical, deployable simulator with interactivity, real-time generation, long-horizon stability, and scalability, which the authors argue existing video world models do not yet achieve because they are trained on narrow embodiment/task/scene distributions and inherit offline video-generation designs such as predicting long frame chunks per call.

Key Contributions

1. A joint-trajectory-conditioned autoregressive diffusion simulator. Uranus conditions on future joint configurations (qpos), camera calibrations, an embodiment description (MJCF or URDF), reference observations, and an optional language prompt, then predicts the next synchronized multi-view observation and feeds it back as context. Each autoregressive step produces one latent frame that decodes into four RGB frames per camera, with the four qpos samples aligned one-to-one with those four RGB timestamps. It predicts the visual consequences of a specified kinematic trajectory rather than integrating torques or velocities.

2. Three claimed system capabilities. Streaming, open-ended rollout with no fixed horizon (one latent frame, corresponding to four RGB frames, generated per step from online joint-position trajectories); low-latency generation at 24 FPS after inference optimization; and scalable, extensible robot control through a unified interface for synchronized multi-view generation across diverse robot embodiments and camera configurations.

3. A cross-embodiment, multi-view conditioning scheme. Embodiment-specific kinematics are decoupled from the learned representation: forward kinematics and camera projection render joint states into shared image-space 2D skeleton maps (with the end-effector marker encoding gripper opening as radius and a spherical-harmonics color derived from end-effector-to-camera relative pose). Camera geometry is encoded as dense 6D Plücker ray embeddings computed from calibration, so variable viewpoints need no learned camera identifiers. Four ordered RGB-time condition maps are compressed into one latent-time token slice by learned 3D convolutions with temporal kernel and stride of four.

4. An end-to-end infrastructure stack and an open release. The report describes data curation (Ray and Daft over MP4 plus Lance tables), training infrastructure (hybrid sharded data parallelism, sequence parallelism, distributed tiled VAE encoding), inference (KV cache, streaming VAE decoding, CUDA graph execution, diffusion distillation), and serving. Code, model weights, SDK, and demo data are released, including an SFT model (Uranus-1.3B) and a distillation model (Uranus-1.3B-Distillation).

Main Findings

• Training corpus scale: After curation, the training corpus contains more than 3,300 hours of high-quality real-world robot data — a total of 245,934 episodes and 351,296,877 frames across AgiBot World Beta (AgiBot G1: 141,355 episodes, 257,410,042 frames, 2,383.4 h), AgiBot World 2026 (AgiBot G2: 1,685 episodes, 2,881,796 frames, 26.7 h), DROID (Franka Panda: 44,349 episodes, 13,982,053 frames, 15 FPS, 258.9 h), and RoboChallenge Table30 v1 and v2 (AgileX, ARX5, Franka Panda, UR5, DOS-W1). The largest single source is AgiBot World Beta.

• Datasets deliberately excluded: RoboCOIN, RoboMIND, and RoboMIND 2.0 releases available to the authors lack camera parameters, and some trajectories lack sufficiently complete embodiment-specific joint angles, so they are not included in the reported corpus. Portions of DROID and RoboChallenge contained camera extrinsics yielding visibly inaccurate robot projections (including some from manual calibration-board annotation), and AgiBot World 2026 contains episodes whose camera streams are not mutually synchronized or do not align with the state sequence; the pipeline performs calibration refinement and episode-level filtering rather than taking these sources at face value.

• Attention cost reduction: Because full self-attention over flattened camera, latent-frame, and patch dimensions costs O((N_c·F·H_p·W_p)²), Uranus factorizes attention. A spatial block (cross-view fusion within a latent frame) costs O(F·(N_c·H_p·W_p)²), a factor-F reduction; a temporal block (motion modeling per camera stream) costs O(N_c·(F·H_p·W_p)²), a factor-N_c reduction. The two modes alternate every two DiT blocks.

• Latency: Uranus achieves 24 FPS after inference optimization. Diffusion step distillation reduces the inner diffusion process from many denoising iterations to four steps while preserving the causal interface.

• Rollout length flexibility: The causal VAE represents the leading RGB frame as one latent frame and adds one latent frame per four subsequent RGB frames (F_lat = 1 + floor((F_rgb − 1)/4) for complete (4k+1)-frame clips), which is what allows generation without a fixed horizon under streaming inference.

• Evaluation claims and explicit limitations: The authors report quantitative and qualitative evaluation on in-distribution and out-of-distribution data, including long-horizon interaction on the WorldOlympiad benchmark, a controlled comparison with GE-Sim 2.0, consistency between simulated and real executions, closed-loop interaction with robot policies, and generalization across unseen scenes, trajectories, tasks, camera motions, and robot embodiments. They state the results demonstrate strong action controllability, multi-view and temporal consistency, and slower degradation over long rollouts, while identifying remaining limitations in precise contact dynamics and object-state transitions under challenging distribution shifts. Specific numeric scores, benchmark tables, and metric values are not reported in the available paper content.

Methodology in Plain English

Uranus is built as an LLM-style autoregressive world model for video. At each step, it takes a bundle of context — clean reference images plus previously generated frames, a robot description, camera calibrations, and the next four joint configurations — and predicts the next chunk of video, one latent frame at a time, which then becomes part of the context for the next step.

Two nested processes make this work. The outer loop is causal and autoregressive: it consumes four future joint positions at a time, aligned to four output RGB frames, and keeps going indefinitely. The inner process is a diffusion model that starts from noise and denoises it into the next clean latent, using flow matching, where a diffusion transformer (DiT) predicts a velocity that moves the sample toward the target. The video itself is generated in a compact VAE latent space to keep things fast.

To avoid the cost of attending over every camera, frame, and pixel patch at once, the transformer alternates two attention patterns: spatial attention, in which the synchronized views of a single latent frame exchange information, and temporal attention, in which frames along one camera stream exchange information. This matches the structure of streaming robot simulation and cuts attention cost.

Generality across robots and cameras comes from how the conditions are built, not from changing the architecture. Joint states are pushed through forward kinematics and projected into the image plane as skeleton diagrams, so different robot bodies — UR5, ALOHA, AGIBOT, ARX5, DROID, and others — become the same kind of input. Camera poses and intrinsics become Plücker ray maps, so novel viewpoints work without any learned camera identity. Text prompts enter through a frozen T5 encoder and cross-attention.

Training is done directly as a causal generator with teacher forcing rather than training a bidirectional video model and distilling it later. A teacher-forcing causal mask lets each target see reference anchors and past latent frames but not future ones; latent-frame-relative rotary position encoding keeps training clips consistent with long streaming rollouts; and reference frames sit at the head of the sequence as persistent attention sinks so older history can be evicted from the KV cache without losing visual anchors.

Getting this to run at scale required infrastructure work on both ends. The data pipeline converts heterogeneous sources (raw directories, HDF5, LeRobot v2.1 and v3.0, RLDS) into MP4 video plus Lance tables, validates integrity and temporal alignment, recovers missing intrinsics with MoGe-3 and optimizes extrinsics with bundle adjustment, filters episodes by comparing projected robot arms against SAM 3 segmentation masks using IoU, and randomly samples episodes for VLM review. Data curation is executed as a distributed pipeline on Ray (scheduling, isolated per-stage resource requests, retries) and Daft (columnar, multimodal DataFrame operations over the Lance corpus). On the training side, hybrid sharded data parallelism, sequence parallelism, and distributed tiled VAE encoding make long multi-view sequences fit. On the inference side, a KV cache, streaming VAE decoding, CUDA graphs, and diffusion distillation turn generation into a stateful streaming service.

Why This Matters

Video world models are a plausible escape hatch from the cost structure of both real-world robot interaction and hand-built simulators. This paper matters because it argues that the field's bottleneck is not only model quality but the whole stack — data quality and scale, training feasibility, streaming inference latency, and persistent serving — and it reports that stack in detail alongside the model. It also makes an unusual choice for the video world model literature by pretraining causally with teacher forcing rather than converting a bidirectional model afterward, and by treating geometry (skeletons and calibrated rays) rather than learned per-embodiment or per-camera embeddings as the interface for cross-embodiment and cross-viewpoint generalization. The open release of code, an SDK, demo data, and 1.3B SFT and distillation weights lowers the barrier for others to test and extend the approach.

Real-world applications the paper's framing points toward:

• Robot data generation to augment or replace costly real-world data collection for policy training.
• Policy evaluation and controlled benchmarking in a repeatable simulated environment before committing hardware.
• Safe iteration on controllers and behaviors, including closed-loop interaction with robot policies, without hardware and safety constraints.
• Cross-embodiment deployment support, using the unified multi-view generation interface across single-arm, dual-arm, and mobile manipulators.

Industry relevance: teams training manipulation policies, building robot fleets, or constructing digital twins can treat this as a candidate layer for scalable data, evaluation, and pre-deployment testing — provided the acknowledged gaps in contact dynamics and object-state transitions are acceptable for their task. The paper's own caveat that "action-conditioned" here means conditioning on a future joint-position trajectory rather than torque or velocity integration is important for anyone expecting a classical physics simulator.

Future Directions

• Improve contact dynamics and object-state transitions. The authors explicitly identify precise contact dynamics and object-state transitions under challenging distribution shifts as remaining limitations, which is the most direct gap to close.
• Broaden training distribution. Part of the motivation for the work was that existing world models are trained on a limited set of embodiments, tasks, and scenes, leaving it unclear whether performance reflects transferable world modeling; expanding coverage would test that question directly.
• Recover or re-release datasets that currently fail the pipeline. RoboCOIN, RoboMIND, and RoboMIND 2.0 are excluded for missing camera parameters and incomplete joint angles, so improved calibration recovery or better source metadata would add data.
• Push latency and horizon further. The current numbers — 24 FPS with four-step distilled denoising and one latent frame (four RGB frames) per step — set a baseline for further streaming optimization, and the evaluation already examines how degradation scales over long rollouts.

Target Audience

Robotics and embodied-AI researchers working on learned simulators, video world models, and policy training; machine learning engineers responsible for large-scale distributed training or low-latency serving of generative video models; and practitioners who need scalable data generation, evaluation, or pre-deployment testing infrastructure for manipulation and mobile robots. Readers without background in diffusion models or transformer attention will find the architecture and training sections demanding, though the introduction and data sections are broadly readable.

Authors’ abstract

Scalable simulation is essential for robot data generation, policy training, evaluation, and safe iteration, yet real-world interaction is costly and conventional simulators require labor-intensive construction. We present Uranus, a data-driven robot simulator built around a joint-trajectory-conditioned autoregressive diffusion model. Uranus offers three key capabilities: (1) streaming, open-ended rollout, which receives future joint-position trajectories online and autoregressively generates one latent frame per step, corresponding to four RGB frames, without a fixed horizon; (2) low-latency generation, achieving 24 FPS after inference optimization; and (3) scalable, extensible robot control, providing a unified interface for synchronized multi-view generation across diverse robot embodiments and camera configurations. We conduct comprehensive quantitative and qualitative evaluations on both in-distribution and out-of-distribution data, providing an objective assessment of Uranus and clearly identifying its current limitations. We release the code and model weights to empower the community with practical tools and insights.