How bytedance/irasim works
1. What this repo is
- IRASim is a PyTorch-based system for learning a robot action-to-video simulator: given an initial frame and a robot action trajectory, it generates a realistic future video of the arm executing that action.
- The repo includes the core model, dataset loaders, training/evaluation entrypoints, demo applications, and benchmark scripts for three robot datasets.
2. Top-level execution flow
- main.py is the primary entrypoint for both training and evaluation/generation.
- Behavior is driven by YAML configs under configs/, which select dataset, task, and model variant.
- The config tree separates:
- configs/base/ for shared data and diffusion settings
- configs/train/* for training jobs
- configs/evaluation/* for short/long rollout evaluation across datasets
3. Core model and generative backbone
- models/irasim.py contains the main IRASim model implementation.
- diffusion/ provides the diffusion machinery used to synthesize videos, including masking and timestep respacing utilities.
- The naming in configs (frame_ada, video_ada, lvdm, vdm) indicates multiple diffusion-backed variants are supported, with IRASim likely wrapping one of these generation modes around robot-trajectory conditioning.
4. Data layer and dataset boundaries
- dataset/ contains the data pipeline:
- dataset_2D.py and dataset_3D.py indicate support for both frame-based and video/latent-video inputs.
- dataset_util.py and video_transforms.py handle preprocessing and augmentation.
- The repo distinguishes between:
- raw videos
- frame-level inputs
- latent videos, which are used to speed up training and generation
- This boundary is important: the model can operate either on decoded video or pre-encoded latent representations depending on pre_encode.
5. Robot benchmark data organization
- robotdata/opensource_robotdata/ shows the local benchmark layout for the three datasets: Language-Table, RT-1, and Bridge.
- Each dataset is organized into:
- annotation/ for metadata and trajectory labels
- videos/ for RGB videos
- latent_videos/ for precomputed latent files
- dataset_structure.txt documents the expected on-disk format, which is central to wiring data loaders correctly.
6. Training pipeline
- Training examples in the README run through main.py with a training config, such as configs/train/rt1/frame_ada.yaml.
- The repo supports both:
- single-GPU training
- distributed multi-GPU training via torchrun
- The main architectural choice is whether to pre-encode videos into latents before training, which reduces compute and storage bandwidth.
7. Evaluation pipeline: short and long rollouts
- Evaluation is split into two modes:
- Short trajectory evaluation: generate one video per test example and compute metrics
- Long trajectory evaluation: generate videos autoregressively over multiple rollout chunks
- Short-mode generation is driven by main.py with evaluation configs.
- Long-mode generation uses scripts/generate_command.py to create multi-process commands, then runs shell scripts such as the generated RT-1 job.
8. Autoregressive rollout subsystem
- sample/sample_autoregressive.py and sample/pipeline_trajectory2videogen.py indicate a dedicated sampling pipeline for longer-horizon generation.
- This is a separate inference path from training because it must chain predictions over time rather than generate a single fixed-length clip.
- Architecturally, this is the key bridge from “single trajectory video synthesis” to “multi-step interactive simulator.”
9. Application/demo boundary
- application/languagetable.py is an interactive demo for the Language-Table setting.
- The README describes it as a keyboard-controlled 2D plane interface over an initial image, which suggests this app is a lightweight human-in-the-loop frontend for testing action-conditioned generation.
- application/vive_controller_traj2video.py suggests an additional trajectory-to-video demo using external controller input.
10. Metrics and benchmark tooling
- evaluate/ contains the automated metric stack:
- compute_psnr_ssim.py
- compute_latent_l2.py
- scripts for short and long evaluation
- The repo vendors or reuses external metric implementations via:
- pytorch-fid/ for FID
- stylegan-v/ for FVD
- These are distinct from model code and form the evaluation boundary for benchmark reporting.
11. Baseline and comparison implementations
- baselines/vdm/ contains a baseline video diffusion implementation.
- The presence of configs for multiple methods shows the repo is structured to compare IRASim against alternative diffusion-based video models under the same dataset and evaluation protocols.
12. Asset and checkpoint management
- assets/ holds static figures and example outputs.
- pretrained_models/ contains scheduler configuration, implying dependency on pretrained diffusion scheduling components.
- The dataset/checkpoint download flow is externalized through scripts/download.sh, and the README points to Bytedance-hosted archives plus Hugging Face for the full dataset release.
13. What a diagram should emphasize
- The central architecture is:
- robot dataset files → dataset loaders/transforms → IRASim diffusion model → generated video/latent rollout
- Surrounding that core are:
- training via main.py
- interactive apps for manual rollout
- evaluation scripts and metric computation
- autoregressive long-horizon sampling
- The main architectural boundary is between data preparation / latent encoding and generative video synthesis, with configs selecting dataset-specific and mode-specific behavior.
Main entry points: main.py main.py · evaluate Metrics · train Train configs · evaluation Eval configs
Components 23
Configs
YAML-driven experiment definitions for training and evaluation.
- Train configs · Dataset- and method-specific training jobs.
configs/train - Eval configs · Dataset- and method-specific evaluation jobs.
configs/evaluation - Base configs · Shared data and diffusion settings.
configs/base
Data
Dataset loaders, transforms, and robot benchmark storage.
- 2D data · Frame-based dataset pipeline.
dataset/dataset_2D.py - 3D data · Video or latent-video dataset pipeline.
dataset/dataset_3D.py - Preprocess · Shared dataset utilities and video transforms.
dataset/dataset_util.py - Transforms · Video preprocessing and augmentation helpers.
dataset/video_transforms.py - Robot data · On-disk robot benchmark corpus for Language-Table, RT-1, and Bridge.
robotdata/opensource_robotdata
Core model
IRASim generation stack and diffusion backbone.
- IRASim · Main robot action-to-video simulator.
models/irasim.py - Diffusion · Masked diffusion and timestep respacing machinery.
diffusion - Masked diffusion · Core noise process and masking-aware synthesis.
diffusion/mask_gaussian_diffusion.py - Respacing · Timestep respacing and schedule utilities.
diffusion/mask_respace.py
Inference
Single-shot, long-horizon, and autoregressive generation paths.
- main.py · Primary training and evaluation entrypoint.
main.py - Autoregressive · Long-horizon rollout sampler that chains predictions.
sample/sample_autoregressive.py - Trajectory pipeline · Trajectory-to-video generation pipeline used for rollout.
sample/pipeline_trajectory2videogen.py
Evaluation
Metric computation and benchmark scripts.
- Short eval · Single-clip generation and benchmark execution.
evaluate/evaluate_short_script.py - Long eval · Multi-step rollout generation and benchmark execution.
evaluate/evaluate_long_script.py - Metrics · PSNR, SSIM, and latent L2 computation.
evaluate
Applications
Interactive demo frontends for human-in-the-loop rollout.
- Language-Table · Keyboard-controlled 2D human-in-the-loop demo.
application/languagetable.py - Vive demo · External controller-driven trajectory-to-video demo.
application/vive_controller_traj2video.py
External
Vendored or downstream baseline dependencies used for comparison and metrics.
- VDM baseline · Reference video diffusion implementation for comparison.
baselines/vdm/video_diffusion_pytorch.py - FID · Vendored FID implementation.
pytorch-fid - FVD · Vendored StyleGAN-V metric implementation.
stylegan-v
Connections 27
- Base configs → Train configs: shared settings · Training jobs inherit base data and diffusion configuration.
- Base configs → Eval configs: shared settings · Evaluation jobs inherit base data and diffusion configuration.
- Train configs → main.py: launch · Training configs drive the main entrypoint.
- Eval configs → main.py: launch · Evaluation configs drive the main entrypoint.
- main.py → 2D data: loads · Frame-based training or evaluation uses the 2D loader.
- main.py → 3D data: loads · Video or latent-video jobs use the 3D loader.
- 2D data → Preprocess: uses · 2D pipeline relies on shared preprocessing utilities.
- 3D data → Preprocess: uses · 3D pipeline relies on shared preprocessing utilities.
- Preprocess → Transforms: applies · Shared preprocessing invokes video transforms.
- 2D data → Robot data: reads · Frame datasets consume benchmark files on disk.
- 3D data → Robot data: reads · Video and latent-video datasets consume benchmark files on disk.
- main.py → IRASim: trains/evals · Main entrypoint instantiates the IRASim model.
- IRASim → Diffusion: wraps · IRASim is built around diffusion-backed generation.
- Diffusion → Masked diffusion: contains · Diffusion backbone provides masked Gaussian diffusion.
- Diffusion → Respacing: contains · Diffusion backbone provides timestep respacing utilities.
- main.py → Autoregressive: rollout · Long-horizon generation uses the autoregressive sampler.
- Autoregressive → Trajectory pipeline: chains · Autoregressive sampling builds on the trajectory-to-video pipeline.
- main.py → Short eval: runs · Short rollout evaluation is launched from the main workflow.
- main.py → Long eval: runs · Long rollout evaluation is launched from the main workflow.
- Short eval → Metrics: scores · Short evaluation computes image/video quality metrics.
- Long eval → Metrics: scores · Long evaluation computes image/video quality metrics.
- Metrics → FID: uses · Metric stack depends on FID tooling.
- Metrics → FVD: uses · Metric stack depends on FVD tooling.
- Language-Table → Trajectory pipeline: drives · Interactive demo feeds actions into the generation pipeline.
- Vive demo → Trajectory pipeline: drives · Controller-based demo feeds actions into the generation pipeline.
- Train configs → VDM baseline: compares · Training configs support baseline method comparison.
- Eval configs → VDM baseline: compares · Evaluation configs support baseline method comparison.