# bytedance/irasim architecture

> Architecture diagram of the GitHub repository bytedance/irasim, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-17). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/bytedance/irasim
- Repository: https://github.com/bytedance/irasim

## Overview

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`](https://github.com/bytedance/irasim/blob/main/main.py) (main.py), [`evaluate`](https://github.com/bytedance/irasim/tree/main/evaluate) (Metrics), [`configs/train`](https://github.com/bytedance/irasim/tree/main/configs/train) (Train configs), [`configs/evaluation`](https://github.com/bytedance/irasim/tree/main/configs/evaluation) (Eval configs)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_configs["Configs"]
  node_train_cfgs["Train configs<br/>yaml configs"]
  node_eval_cfgs["Eval configs<br/>yaml configs"]
  node_base_cfgs["Base configs<br/>yaml configs"]
end

subgraph group_data["Data"]
  node_dataset_2d["2D data<br/>loader<br/>[dataset_2D.py]"]
  node_dataset_3d["3D data<br/>loader<br/>[dataset_3D.py]"]
  node_dataset_utils["Preprocess<br/>data utils<br/>[dataset_util.py]"]
  node_video_transforms["Transforms<br/>augmentation"]
  node_robot_data[("Robot data<br/>benchmark data")]
end

subgraph group_core["Core model"]
  node_irasim_model{{"IRASim<br/>model<br/>[irasim.py]"}}
  node_diffusion_core["Diffusion<br/>backbone"]
  node_mask_diffusion["Masked diffusion<br/>diffusion impl"]
  node_mask_respace["Respacing<br/>diffusion util<br/>[mask_respace.py]"]
end

subgraph group_inference["Inference"]
  node_main_py["main.py<br/>entrypoint<br/>[main.py]"]
  node_sample_auto["Autoregressive<br/>sampler"]
  node_traj_pipeline["Trajectory pipeline<br/>inference pipeline"]
end

subgraph group_eval["Evaluation"]
  node_short_eval["Short eval<br/>eval runner"]
  node_long_eval["Long eval<br/>eval runner"]
  node_metrics["Metrics<br/>metric stack"]
end

subgraph group_apps["Applications"]
  node_apps_demo["Language-Table<br/>demo app<br/>[languagetable.py]"]
  node_vive_demo["Vive demo<br/>demo app"]
end

subgraph group_external["External"]
  node_vdm_baseline["VDM baseline"]
  node_fid_lib["FID<br/>external metric"]
  node_fvd_lib["FVD<br/>external metric"]
end

node_base_cfgs -->|"shared settings"| node_train_cfgs
node_base_cfgs -->|"shared settings"| node_eval_cfgs
node_train_cfgs -->|"launch"| node_main_py
node_eval_cfgs -->|"launch"| node_main_py
node_main_py -->|"loads"| node_dataset_2d
node_main_py -->|"loads"| node_dataset_3d
node_dataset_2d -->|"uses"| node_dataset_utils
node_dataset_3d -->|"uses"| node_dataset_utils
node_dataset_utils -->|"applies"| node_video_transforms
node_dataset_2d -->|"reads"| node_robot_data
node_dataset_3d -->|"reads"| node_robot_data
node_main_py -->|"trains/evals"| node_irasim_model
node_irasim_model -->|"wraps"| node_diffusion_core
node_diffusion_core -->|"contains"| node_mask_diffusion
node_diffusion_core -->|"contains"| node_mask_respace
node_main_py -->|"rollout"| node_sample_auto
node_sample_auto -->|"chains"| node_traj_pipeline
node_main_py -->|"runs"| node_short_eval
node_main_py -->|"runs"| node_long_eval
node_short_eval -->|"scores"| node_metrics
node_long_eval -->|"scores"| node_metrics
node_metrics -->|"uses"| node_fid_lib
node_metrics -->|"uses"| node_fvd_lib
node_apps_demo -->|"drives"| node_traj_pipeline
node_vive_demo -->|"drives"| node_traj_pipeline
node_train_cfgs -.->|"compares"| node_vdm_baseline
node_eval_cfgs -.->|"compares"| node_vdm_baseline

click node_main_py "https://github.com/bytedance/irasim/blob/main/main.py"
click node_train_cfgs "https://github.com/bytedance/irasim/tree/main/configs/train"
click node_eval_cfgs "https://github.com/bytedance/irasim/tree/main/configs/evaluation"
click node_base_cfgs "https://github.com/bytedance/irasim/tree/main/configs/base"
click node_dataset_2d "https://github.com/bytedance/irasim/blob/main/dataset/dataset_2D.py"
click node_dataset_3d "https://github.com/bytedance/irasim/blob/main/dataset/dataset_3D.py"
click node_dataset_utils "https://github.com/bytedance/irasim/blob/main/dataset/dataset_util.py"
click node_video_transforms "https://github.com/bytedance/irasim/blob/main/dataset/video_transforms.py"
click node_robot_data "https://github.com/bytedance/irasim/tree/main/robotdata/opensource_robotdata"
click node_irasim_model "https://github.com/bytedance/irasim/blob/main/models/irasim.py"
click node_diffusion_core "https://github.com/bytedance/irasim/tree/main/diffusion"
click node_mask_diffusion "https://github.com/bytedance/irasim/blob/main/diffusion/mask_gaussian_diffusion.py"
click node_mask_respace "https://github.com/bytedance/irasim/blob/main/diffusion/mask_respace.py"
click node_sample_auto "https://github.com/bytedance/irasim/blob/main/sample/sample_autoregressive.py"
click node_traj_pipeline "https://github.com/bytedance/irasim/blob/main/sample/pipeline_trajectory2videogen.py"
click node_short_eval "https://github.com/bytedance/irasim/blob/main/evaluate/evaluate_short_script.py"
click node_long_eval "https://github.com/bytedance/irasim/blob/main/evaluate/evaluate_long_script.py"
click node_metrics "https://github.com/bytedance/irasim/tree/main/evaluate"
click node_apps_demo "https://github.com/bytedance/irasim/blob/main/application/languagetable.py"
click node_vive_demo "https://github.com/bytedance/irasim/blob/main/application/vive_controller_traj2video.py"
click node_vdm_baseline "https://github.com/bytedance/irasim/blob/main/baselines/vdm/video_diffusion_pytorch.py"
click node_fid_lib "https://github.com/bytedance/irasim/tree/main/pytorch-fid"
click node_fvd_lib "https://github.com/bytedance/irasim/tree/main/stylegan-v"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_train_cfgs,node_eval_cfgs,node_base_cfgs,node_vdm_baseline,node_fid_lib,node_fvd_lib toneBlue
class node_dataset_2d,node_dataset_3d,node_dataset_utils,node_video_transforms,node_robot_data toneAmber
class node_irasim_model,node_diffusion_core,node_mask_diffusion,node_mask_respace toneMint
class node_main_py,node_sample_auto,node_traj_pipeline toneRose
class node_short_eval,node_long_eval,node_metrics toneIndigo
class node_apps_demo,node_vive_demo toneTeal
```

## Components

### Configs

YAML-driven experiment definitions for training and evaluation.

- [Train configs](https://github.com/bytedance/irasim/tree/main/configs/train): yaml configs · Dataset- and method-specific training jobs. · `configs/train`
- [Eval configs](https://github.com/bytedance/irasim/tree/main/configs/evaluation): yaml configs · Dataset- and method-specific evaluation jobs. · `configs/evaluation`
- [Base configs](https://github.com/bytedance/irasim/tree/main/configs/base): yaml configs · Shared data and diffusion settings. · `configs/base`

### Data

Dataset loaders, transforms, and robot benchmark storage.

- [2D data](https://github.com/bytedance/irasim/blob/main/dataset/dataset_2D.py): loader · Frame-based dataset pipeline. · `dataset/dataset_2D.py`
- [3D data](https://github.com/bytedance/irasim/blob/main/dataset/dataset_3D.py): loader · Video or latent-video dataset pipeline. · `dataset/dataset_3D.py`
- [Preprocess](https://github.com/bytedance/irasim/blob/main/dataset/dataset_util.py): data utils · Shared dataset utilities and video transforms. · `dataset/dataset_util.py`
- [Transforms](https://github.com/bytedance/irasim/blob/main/dataset/video_transforms.py): augmentation · Video preprocessing and augmentation helpers. · `dataset/video_transforms.py`
- [Robot data](https://github.com/bytedance/irasim/tree/main/robotdata/opensource_robotdata): benchmark 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](https://github.com/bytedance/irasim/blob/main/models/irasim.py): model · Main robot action-to-video simulator. · `models/irasim.py`
- [Diffusion](https://github.com/bytedance/irasim/tree/main/diffusion): backbone · Masked diffusion and timestep respacing machinery. · `diffusion`
- [Masked diffusion](https://github.com/bytedance/irasim/blob/main/diffusion/mask_gaussian_diffusion.py): diffusion impl · Core noise process and masking-aware synthesis. · `diffusion/mask_gaussian_diffusion.py`
- [Respacing](https://github.com/bytedance/irasim/blob/main/diffusion/mask_respace.py): diffusion util · Timestep respacing and schedule utilities. · `diffusion/mask_respace.py`

### Inference

Single-shot, long-horizon, and autoregressive generation paths.

- [main.py](https://github.com/bytedance/irasim/blob/main/main.py): entrypoint · Primary training and evaluation entrypoint. · `main.py`
- [Autoregressive](https://github.com/bytedance/irasim/blob/main/sample/sample_autoregressive.py): sampler · Long-horizon rollout sampler that chains predictions. · `sample/sample_autoregressive.py`
- [Trajectory pipeline](https://github.com/bytedance/irasim/blob/main/sample/pipeline_trajectory2videogen.py): inference pipeline · Trajectory-to-video generation pipeline used for rollout. · `sample/pipeline_trajectory2videogen.py`

### Evaluation

Metric computation and benchmark scripts.

- [Short eval](https://github.com/bytedance/irasim/blob/main/evaluate/evaluate_short_script.py): eval runner · Single-clip generation and benchmark execution. · `evaluate/evaluate_short_script.py`
- [Long eval](https://github.com/bytedance/irasim/blob/main/evaluate/evaluate_long_script.py): eval runner · Multi-step rollout generation and benchmark execution. · `evaluate/evaluate_long_script.py`
- [Metrics](https://github.com/bytedance/irasim/tree/main/evaluate): metric stack · PSNR, SSIM, and latent L2 computation. · `evaluate`

### Applications

Interactive demo frontends for human-in-the-loop rollout.

- [Language-Table](https://github.com/bytedance/irasim/blob/main/application/languagetable.py): demo app · Keyboard-controlled 2D human-in-the-loop demo. · `application/languagetable.py`
- [Vive demo](https://github.com/bytedance/irasim/blob/main/application/vive_controller_traj2video.py): demo app · 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](https://github.com/bytedance/irasim/blob/main/baselines/vdm/video_diffusion_pytorch.py): baseline · Reference video diffusion implementation for comparison. · `baselines/vdm/video_diffusion_pytorch.py`
- [FID](https://github.com/bytedance/irasim/tree/main/pytorch-fid): external metric · Vendored FID implementation. · `pytorch-fid`
- [FVD](https://github.com/bytedance/irasim/tree/main/stylegan-v): external metric · Vendored StyleGAN-V metric implementation. · `stylegan-v`

## Connections

- 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.