JackLiu0406/da3-xvla-pipeline

DA3-XVLA — Training & Evaluation Pipeline

Self-contained, reproducible pipeline for DA3-XVLA: X-VLA (Florence-2 VLM + DiT flow-matching action expert) augmented with Depth-Anything-3 geometry conditioning and pluggable spatial-conditioning architectures (spatial cross-attention adapters at every layer, per-token gated fusion, T5-language K/V bank, etc.).

Clone this repo onto a fresh cluster node, point it at data + base checkpoints, and you can recreate every ablation we've validated. Tune the major axes (input resolution, K geometry tokens, action-expert structure, spatial-conditioning variant, geometry backbone, language conditioning) by changing a handful of flags.

TL;DR

1. conda env create -f training/environment.yml && conda activate xvla (see §2 for cluster-specific notes incl. Blackwell sm_120 / GB10).
2. Pull data + base ckpt to disk (§3).
3. Pick a launch script from launchers/proven_recipes/, edit the env block at the top, then run it (or wrap in sbatch for SLURM clusters).
4. Each ablation cell changes ONE axis vs the baseline — see §6 for the knob table.

Models trained from this pipeline: JackLiu0406/DA3-XVLA-roboreal-ablations (one subfolder per config; each has its own eval-contract README). Training data: JackLiu0406/roboreal_lerobot (RoboReal 15,999 episodes, lerobot v2.0, countertop+left+right cams).

---

1. Repository layout

training/              DA3-XVLA training codebase (entry: train.py, peft_train.py)
  models/
    configuration_xvla.py         ← ALL architecture defaults live here
    modeling_xvla.py              ← XVLA model: VLM + spatial cross-attn + action expert
    transformer.py                ← DiT action expert blocks
    geometry_conditioning.py      ← DA3 → DualDPT → Perceiver → K geometry tokens
    da3_inline.py                 ← inlined Depth-Anything-3 backbone (frozen)
    vggt_inline.py                ← VGGT-Omega backbone (alternative geometry)
    action_hub.py                 ← flow-matching action head
    processing_xvla.py            ← input preprocessing (images, instructions)
    t5_language.py                ← optional T5-base language K/V bank
  datasets/                       ← lerobot dataset + meta-json loader
  third_party/
    Depth-Anything-3/             ← vendored DA3 source (clone-and-run)
    vggt_omega/                   ← vendored VGGT-Omega source
  evaluation/                     ← per-benchmark eval configs (robotwin-2.0, libero, …)
  train.py                        ← entry point (flow-matching trainer)
  peft_train.py                   ← LoRA / PEFT variant
  environment.yml                 ← conda env spec
  requirements.txt                ← pip extras (transformers, accelerate, peft, …)
  requirements_xvla.txt           ← full pinned env (use only as reference)
launchers/
  proven_recipes/                 ← the 9 launch scripts for ablation cells we've trained:
    baseline_DA3BASE_K160_native_aspect.sh
    baseline_DA3BASE_K160_dptfused_aspect.sh
    baseline_DA3LARGE_K160.sh
    baseline_DA3GIANT_K160_dptfused_xattn_all_aspect.sh         ← scalar gate, K=160, ALL 24 layers
    baseline_DA3GIANT_K160_dptfused_pertokengate_xattn_all_aspect.sh  ← per-token gated fusion (paper Eq.5)
    baseline_DA3GIANT_K160_dptfused_fullseqxattn_all_aspect.sh  ← full-sequence cross-attn target
    baseline_DA3GIANT_K160_dptfused_t5lang_fullseqxattn_all_aspect.sh  ← + T5-base language K/V bank
    baseline_DA3GIANT_K160_spatialxattn_aspect.sh               ← spatial expert variant
    baseline_DA3GIANT_K240_dptfused_aspect.sh                   ← K=240 geometry tokens
    meta_all_ct_3cam_clean.json   ← the data spec used by every recipe above
xvla/                  upstream X-VLA source (reference; the stack builds on it)
robotwin/              RoboTwin/RoboPRO eval adaptation (NOT needed to train; eval-side)
DA3-VLA_SETUP.md       eval-side environment setup (RoboTwin/SAPIEN + curobo)
README.md              this file

---

2. Environment setup

2a. Conda env (training only)

conda env create -f training/environment.yml
conda activate xvla
pip install -e training/third_party/Depth-Anything-3
pip install -r training/requirements.txt

The env pins Python 3.10, PyTorch 2.1 / CUDA 12.1 by default. This works on Ampere (A100), Hopper (H100/H200), and Ada (RTX 4090). Verified daily on H200.

2b. Blackwell / RTX 5080 / GB10 (Grace-Blackwell) clusters

CUDA 12.1 does NOT support Blackwell sm_120. On Blackwell hardware (GB10 Superchip, RTX 5080, B100/B200), use a newer torch:

conda create -n xvla python=3.11 -y
conda activate xvla

# Torch 2.7+ with CUDA 12.8 (supports sm_120). For GB10 specifically (ARM CPU),
# install from the ARM-64 channel:
pip install --index-url https://download.pytorch.org/whl/cu128 torch torchvision torchaudio

pip install -e training/third_party/Depth-Anything-3
pip install -r training/requirements.txt
# NB: transformers<=4.51.3 is pinned in requirements.txt for X-VLA compatibility.

GB10 notes:

• Single GPU per board — most of our launch scripts assume 4 GPUs. Set CUDA_VISIBLE_DEVICES=0 and --num_processes 1 if you have one GPU per node; or use multi-node accelerate launch --num_machines N for SLURM/MPI clusters.
• Unified memory (Grace CPU shares mem with Blackwell GPU) — VRAM cap is effectively the unified-memory ceiling (~96-128 GiB on the GB10 dev kit, more on production GB10 nodes). Most of our recipes were sized for 140 GiB H200, so reduce batch from 32 → 16-24 if you hit OOM, or enable XVLA_VLM_GRADIENT_CHECKPOINTING=1.
• flash-attn: not buildable on Blackwell+torch2.7+ARM at the moment. The model falls back to SDPA / eager attention via attn_implementation: "sdpa". ~15-20% perf hit but works correctly.

2c. Eval env (separate, optional)

The RoboTwin/SAPIEN evaluation stack has conflicting deps with training (curobo

• SAPIEN renderer). Use DA3-VLA_SETUP.md for the eval env install — it's a separate conda env from xvla.

---

3. Data + base checkpoints

3a. Training data — RoboReal lerobot v2.0 (15,999 episodes)

# Download to a local path (5.5 GB)
mkdir -p /work/jack/DATASETS
cd /work/jack/DATASETS
huggingface-cli download \
    --repo-type dataset \
    JackLiu0406/roboreal_lerobot \
    --local-dir roboreal_merged_lerobot

# Verify structure
ls roboreal_merged_lerobot/{data,videos,meta}
# Expect 16 chunks × 1000 episodes (15,999 total; one dropped from the source corpus).

Schema:

• robot_type: roboreal (Agilex Aloha bimanual)
• fps: 25
• codebase_version: lerobot v2.0
• action_dim / state_dim: 14 (left_waist, left_shoulder, left_elbow, left_forearm_roll, left_wrist_angle, left_wrist_rotate, left_gripper, right_waist, …, right_gripper)
• video keys: observation.images.countertop, observation.images.left, observation.images.right (240×320 RGB)

3b. Meta-json (training data spec)

launchers/proven_recipes/meta_all_ct_3cam_clean.json enumerates which task/scene combinations to include and the (possibly weighted) language paraphrasings. Most recipes reference META=/work/jack/da3xvla_runs/meta_all_ct_3cam_clean.json; update this path in the launch script after staging the dataset locally.

3c. Base checkpoints

Pull the pretrained DA3-XVLA-Pt base checkpoint (the foundation model X-VLA was pretrained on, with our DA3 modifications):

# Per-recipe base ckpt — each launch script's BASE_CKPT="$RUN/ckpt_init" expects this
huggingface-cli download \
    JackLiu0406/xvla-pt-da3-base \
    --local-dir /work/jack/da3xvla_runs/baseline_DA3GIANT_K160_dptfused_xattn_all_aspect/ckpt_init

Each ablation cell's launch.sh has a BASE_CKPT pointing at a ckpt_init/ directory that must contain a valid HF-style XVLA checkpoint (config.json + pytorch_model.bin or safetensors). The exact config.json values are sanity- checked by the launch script before training starts — see §6.

3d. Depth-Anything-3 frozen backbone weights

DA3 weights are pulled automatically by transformers AutoModel on first use:

da3_model_name: "depth-anything/DA3-Giant-1.1"  # ← the default for DA3-Giant variants
# Also valid: "depth-anything/DA3-Base", "depth-anything/DA3-Large-1.1"

⚠️ Use DA3-Giant-1.1, NOT DA3-Giant. Upstream deprecated v1.0 after a bugfix. Same for DA3-Large-1.1. Build new run configs with the -1.1 suffix.

Pre-cache them (otherwise first step has multi-minute download stall):

HF_HOME=/work/jack/da3xvla_workspace/hf_cache \
python -c "from transformers import AutoModel; AutoModel.from_pretrained('depth-anything/DA3-Giant-1.1')"

---

4. Architecture (data flow)

3 RGB views ─┬─► Florence-2 VLM  ───────────────────────► vlm tokens ─┐
             │                                                         │
             └─► DA3-Giant-1.1 (frozen) ─► DualDPT dense feat (128ch)   ├─► spatial cross-attn adapters
                 ─► Perceiver resampler (per-view) ─► K geometry tokens ┘   at EVERY action expert layer
                                                                            (scalar OR per-token gate)
                 (+ optional T5-base language K/V bank)                        │
                                                                       Action expert (DiT, flow-matching)
                                                                            [action | vlm | aux | soft]
                                                                                  │
                                                                       30-step × 14-dim joint qpos chunk

Key components:

• Geometry backbone: Depth-Anything-3 Giant-1.1, frozen, used in dpt_fused mode (DualDPT head's output_conv2 pre-final hook gives a 128-channel dense feature at spatially-recovered resolution). For non-Giant baselines, also supports last, multi (custom multi-layer fuse), vggt_omega.
• Resampler: Per-view Perceiver compresses dense features → K geometry tokens. Default K=160 distributed as [96 countertop, 32 L-wrist, 32 R-wrist] for the 3-camera setup. We have ablations at K=240 and K=320.
• Spatial cross-attention adapter: Inserted at every action-expert layer (24 layers for the Giant variant). Each adapter cross-attends from action tokens (Q) to the geometry tokens (K/V). Outputs gated and residually added back to the action stream.
  • Scalar gate (spatial_gate_type: "scalar"): one learnable β per layer.
  • Per-token gate (spatial_gate_type: "per_token"): small MLP per layer emits per-token sigmoid gate from the action hidden state — implements paper Eq.5 GatedFusion.
• Cross-attention target: Either action tokens only (spatial_cross_attention_target: "action_only"), or full action+vlm+aux sequence ("full_sequence").
• T5-language K/V bank (optional): t5_language_enabled: true adds T5-base encoder over raw instruction strings → 768-dim per-token features → projected to 1024 → concatenated with geometry tokens as combined K/V for the spatial adapters. Adds ~110M frozen + 1.8M trainable params.
• Action head: Flow-matching with rectified-flow velocity loss. 30-step chunks of 14-D joint qpos (absolute joints, NOT EE deltas).

---

5. Training launch (the actual commands)

A typical run:

# 1. Set up the run dir + symlink ckpt_init
RUN=/work/jack/da3xvla_runs/baseline_DA3GIANT_K160_dptfused_pertokengate_xattn_all_aspect
mkdir -p $RUN
# (one-time: download the base ckpt into $RUN/ckpt_init/ as in §3c)

# 2. Stage the meta json
cp launchers/proven_recipes/meta_all_ct_3cam_clean.json /work/jack/da3xvla_runs/

# 3. Edit the launch script's SRC, RUN, META paths at the top, then run:
bash launchers/proven_recipes/baseline_DA3GIANT_K160_dptfused_pertokengate_xattn_all_aspect.sh

Every launch script is self-documenting at the top with the architecture intent and the recipe deltas vs the baseline. Read the comment block before running.

Common envs set by the launch scripts:

CUDA_VISIBLE_DEVICES=4,5,6,7          # 4 GPUs for the meta data scope
XVLA_DDP_STATIC_GRAPH=1
XVLA_POSED_DA3=1                       # DA3 in posed mode (depth + pose)
XVLA_RGB_INPUT=1                       # true-RGB input (don't normalize before DA3)
XVLA_DA3_NATIVE_INPUT=1                # DA3 receives the full native-aspect input
XVLA_DA3_INPUT_H=252 XVLA_DA3_INPUT_W=336   # aspect-preserved 252×336 = exact 3:4 multiple
XVLA_VLM_GRADIENT_CHECKPOINTING=0     # disabled for speed (set =1 if VRAM-bound)
XVLA_NUM_WORKERS=16 XVLA_PREFETCH_FACTOR=4
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
HF_HOME=/work/jack/da3xvla_workspace/hf_cache

Common train.py flags:

--models "$BASE_CKPT"                 # path to ckpt_init/ (HF-style)
--train_metas_path "$META"            # the meta json from §3b
--output_dir "$OUT"                   # this run's output dir
--batch_size 32                       # per-GPU; multiplied by num_processes for global
--learning_rate 1e-4 --learning_coef 0.1
--weight_decay 0.0 --betas 0.9 0.95 --max_grad_norm 1.0
--iters 100000 --freeze_steps 1000 --warmup_steps 2000
--use_cosine_decay --min_lr_ratio 0.5    # NOT 0.1 (memorised rule: 0.1 stalls late convergence)
--save_interval 10000 --log_interval 20
--seed 0 --apply_finetune_policy

---

6. Knobs to tune (the single-knob ablation table)

The proven recipes in launchers/proven_recipes/ each change ONE axis vs the DA3-XVLA baseline:

Axis	Default	Variants in recipes	Where in ckpt_init/config.json
da3_model_name	depth-anything/DA3-Base	DA3-Large-1.1, DA3-Giant-1.1	da3_model_name
da3_feature_layer	last	multi (1×1 conv fuse), dpt_fused (DualDPT head)	da3_feature_layer
da3_input_dim	384	128 (dpt_fused), 1536 (multi for Giant)	da3_input_dim
Input resolution	native_aspect 224	aspect 252×336 (3:4 native)	da3_input_height, da3_input_width
num_geometry_tokens (K)	160	240, 320	num_geometry_tokens
Spatial cross-attn	OFF	ON via use_spatial_cross_attention: true	flag
Layers covered	middle_late (12 layers)	"all" (24 layers)	spatial_cross_attention_layers
spatial_gate_type	scalar (one β per layer)	per_token (paper Eq.5 GatedFusion MLP)	spatial_gate_type
spatial_cross_attention_target	action_only	full_sequence	flag
T5 language K/V	OFF	t5_language_enabled: true + t5_model_name: t5-base	flag

Baselines and their losses (RoboReal meta, 100k steps, bs=32×4=128):

Recipe (folder name)	meta loss @ 100k	RoboTwin SR (avg)
baseline_DA3BASE_K160_native_aspect	0.081	64%
baseline_DA3BASE_K160_dptfused_aspect	0.075	67%
baseline_DA3LARGE_K160	0.078	(TBD)
baseline_DA3GIANT_K240_dptfused_aspect	0.094	(TBD)
baseline_DA3GIANT_K160_dptfused_xattn_all_aspect	0.099	(TBD)
baseline_DA3GIANT_K160_dptfused_pertokengate_xattn_all_aspect	0.096	(TBD)
baseline_DA3GIANT_K160_dptfused_fullseqxattn_all_aspect	0.099	(TBD)
baseline_DA3GIANT_K160_dptfused_t5lang_fullseqxattn_all_aspect	(TBD)	(TBD)

⚠️ Loss is NOT a reliable success-rate predictor. VGGT (loss 0.099) BEAT DA3-Base aspect (loss 0.081) by +6 SR on RoboTwin in our prior eval. Always run the closed-loop eval; never A/B compare ablations by loss alone.

---

7. Eval (RoboTwin / RoboPRO)

The eval-side glue lives in robotwin/customized_robotwin/policy/dxvla/:

# Env setup (separate conda env from training)
# See DA3-VLA_SETUP.md for SAPIEN + curobo install.

# Launch eval client → server
bash robotwin/launchers/sbatch_eval_robotwin_EH_4gpu.sh

Key eval env vars:

• DXVLA_PRIMARY_CAM=countertop → model receives the simulator's countertop_camera as its primary view (mapped to the model's cam_high slot).
• EVAL_FAST_CTRL=1 → critical — without this, TOPP planner blows up to 280k waypoints per step (3.6s/step → eval takes days). Already baked into harness/run_h200.sh.
• DXVLA_VIDEO_CAM=countertop_camera → record video from the model's input view.

---

8. Troubleshooting

Symptom	Likely cause	Fix
KeyError: cam_high in dataloader	modality.json key mismatch	Add aliased cam_high → observation.images.countertop in meta/modality.json
flash_attn not found, falling back to sdpa	FA2 not installed (Blackwell + torch2.7+)	Accept the 15-20% perf hit, or build FA2 from source matching your torch+cuda.
First-step delay 5+ min	DA3 weights downloading	Pre-cache via §3d before launch.
OOM at step 1 (Blackwell)	VRAM headroom insufficient at bs=32	Drop to bs=16 or set XVLA_VLM_GRADIENT_CHECKPOINTING=1.
Eval is 3.6s/step	TOPP blow-up	Set EVAL_FAST_CTRL=1 in the eval launcher.
Eval planner picks wrong camera	DXVLA_PRIMARY_CAM unset	Set DXVLA_PRIMARY_CAM=countertop (NOT cam_high).
"DA3-Giant" deprecated warning at ckpt-init	Old v1.0 model name	Update da3_model_name to depth-anything/DA3-Giant-1.1 in ckpt_init/config.json.
Loss plateaus at 0.1+ late in training	Cosine min_lr_ratio=0.1 is too aggressive	Set --min_lr_ratio 0.5 (we've validated this empirically).

---

9. License + provenance

Upstream X-VLA: MIT (see xvla/LICENSE). Depth-Anything-3: Apache 2.0 (see training/third_party/Depth-Anything-3/LICENSE). DA3-XVLA additions (geometry conditioning, spatial expert, per-token gate, T5 language K/V): Apache 2.0.

Cite X-VLA: arxiv.org/pdf/2510.10274 Cite DA3: arxiv.org/abs/2509.xxxxx