DIT4DIT

A mimic-video–style video–action model: LoRA-finetune a Cosmos video backbone, then train a lightweight DiT with flow matching for action prediction on LIBERO manipulation tasks (LeRobot-format data).

Early validation of the Stage 1 LoRA finetuning (Left: Ground Truth | Right: DiT Auto-encoder reconstruction from noise):

Early Stage 2 evaluation on libero_object — the model starts to successfully complete manipulation tasks:

Task: pick up the bbq sauce and place it in the basket  |  pick up the tomato sauce and place it in the basket

• Data: Built-in support for all 4 LIBERO suites via --suite (libero_spatial, libero_object, libero_goal, libero_10). Uses 100% of available demonstrations for training to maximize simulated evaluation metrics.
• Cameras: agentview + wrist (observation.images.image / observation.images.wrist_image), 256×256, concatenated side-by-side into a single stream
• State / actions: 8-D proprio, 7-D actions (x, y, z, roll, pitch, yaw, gripper), action chunk size 16
• Video: 17 pixel frames → 5 latent frames (2 conditional + 3 predicted), 10 FPS
• Backbone: nvidia/Cosmos-Predict2-2B-Video2World + LoRA
• Action head: ~67M-parameter DiT + flow matching
• Training: two stages — Stage 1 video prediction (LoRA) → Stage 2 action prediction

pip install -r requirements.txt
# For LIBERO sim eval (WebSocket server + client):
pip install websockets

Python 3.10+, CUDA, and bf16 are recommended (training defaults to bf16).

Hyperparameters and the suite registry live in configs/config.py (DataConfig, ModelConfig, Stage1Config, Stage2Config, LIBERO_SUITES). Edit that file to change global settings, camera keys, W&B project names, or dataset proportions.

Use --preset to apply a pre-tuned micro_batch_size / dtype / gradient_checkpointing for your hardware. gradient_accumulation is computed automatically to keep the effective batch size at 200.

V100 note: V100 has no bf16 hardware support. The v100 preset sets dtype=fp16, but stable fp16 training also requires a GradScaler in the trainers (not yet wired up by default).

Individual flags override the preset:

# A100 80 GB × 2, but keep gradient checkpointing on (less VRAM)
torchrun --nproc_per_node=2 scripts/train_stage1.py \
    --suite libero_object --preset a100_80g \
    --no_gradient_checkpointing  # remove this flag to re-enable GC

# Override just micro_batch on top of a preset
torchrun --nproc_per_node=4 scripts/train_stage1.py \
    --suite libero_object --preset a100_40g --micro_batch_size 32

Override priority (highest wins): --micro_batch_size / --dtype / --no_gradient_checkpointing > --preset > config.py defaults.

Choose a suite to train on: libero_spatial, libero_object, libero_goal, or libero_10. The --suite argument automatically manages Hugging Face repo IDs, episode counts, multi-task T5 embeddings, and checkpoint directories.

# 1. Precompute multi-task T5 text embeddings for the suite
python scripts/precompute_embeddings.py --suite libero_object --cosmos_model_id your_cosmos_model_path

# (Optional) Precompute VAE latents for speed:
# python scripts/precompute_embeddings.py --suite libero_object --latents

# 2. Stage 1: video backbone LoRA (auto-saves to checkpoints/libero_object/stage1/)
torchrun --nproc_per_node=5 scripts/train_stage1.py \
    --suite libero_object \
    --preset 4090 \
    --cosmos_model_id your_cosmos_model_path \
    --wandb_project "dit4dit-stage1" \
    --resume checkpoints/libero_object/stage1/step_5000

# 3. Stage 2: action decoder (auto-loads Stage 1 & saves to checkpoints/libero_object/stage2/)
#
# On a fresh machine, run single-process first to generate action_stats.pt,
# then Ctrl+C after "Saved action stats to ..." appears, and launch torchrun.
# This avoids other ranks timing out at the DDP barrier while rank 0 computes stats.
#
#   python scripts/train_stage2.py --suite libero_object --preset 4090 \
#       --cosmos_model_id your_cosmos_model_path \
#       --stage1_checkpoint checkpoints/libero_object/stage1/final
#   # Ctrl+C after: "Saved action stats to precomputed/libero_object/action_stats.pt"
#
torchrun --nproc_per_node=5 scripts/train_stage2.py \
    --suite libero_object \
    --preset 4090 \
    --cosmos_model_id your_cosmos_model_path \
    --stage1_checkpoint checkpoints/libero_object/stage1/final \
    --wandb_project "dit4dit-stage2" \
    --resume checkpoints/libero_object/stage2/step_1000

(Optional) To evaluate offline Action MSE (requires manually setting val_episodes > 0 in config.py for the suite):

python scripts/evaluate.py --suite libero_object --cosmos_model_id your_cosmos_model_path

Evaluation scripts set HF_HUB_OFFLINE=1 so Hugging Face models are resolved from the local cache when possible; the first run still needs network access to download Cosmos / T5 weights.

Evaluation is performed zero-shot on novel initial state distributions in simulation. Install LIBERO in a separate environment, then run the WebSocket server and sim client (see also LIBERO_evaluation/setup.sh).

Terminal 1 — model server (e.g., your training env):

conda activate mimic
python scripts/eval_server.py \
        --suite libero_object \
        --stage1_checkpoint checkpoints/libero_object/stage1/step_4000 \
        --stage2_checkpoint checkpoints/libero_object/stage2/step_1000 \
        --cosmos_model_id checkpoints/models--nvidia--Cosmos-Predict2-2B-Video2World/snapshots/f50c09f5d8ab133a90cac3f4886a6471e9ba3f18

Terminal 2 — LIBERO client (e.g., your LIBERO simulator env):

python LIBERO_evaluation/libero_client.py \
  --server_url ws://localhost:8765 \
  --suites libero_object \
  --num_episodes 50 \
  --save_video

python scripts/benchmark_latency.py --dry_run
python scripts/benchmark_latency.py --device cuda --warmup 3 --repeats 10

DIT4DIT/
├── configs/
│   └── config.py                 # Data / model / two-stage training config / Suite Registry / GPU Presets
├── mimic_video/
│   ├── data/
│   │   ├── dataset.py            # MimicVideoDataset (LeRobot + Multi-task handling)
│   │   └── transforms.py         # Two-camera concat, normalization, etc.
│   ├── models/
│   │   ├── video_backbone.py     # Cosmos + LoRA + hidden-state hooks
│   │   ├── action_decoder.py     # ActionDecoderDiT (flow matching)
│   │   └── flow_matching.py      # Scheduler and ODE solver
│   ├── training/
│   │   ├── stage1_trainer.py
│   │   └── stage2_trainer.py
│   └── inference/
│       └── policy.py             # MimicVideoPolicy (inference wrapper)
├── LIBERO_evaluation/
│   ├── libero_client.py          # Sim client talking to eval_server
│   └── setup.sh                  # Two-env run notes (edit paths for your machine)
└── scripts/
    ├── precompute_embeddings.py
    ├── train_stage1.py
    ├── train_stage2.py
    ├── evaluate.py
    ├── eval_server.py            # WebSocket server for LIBERO
    └── benchmark_latency.py

Paper: mimic-video: Video-Action Models for Generalizable Robot Control Beyond VLAs. This repo is an independent reimplementation / extension; behavior and defaults are defined by the code and configs/config.py.