diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/README.md b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..046d62283af8bf958c9c32cb8c1995ca07b84fd8
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/README.md
@@ -0,0 +1,12 @@
+# Flappy QwenOFT H1 with RTX3090 mean latency
+
+Training complete: 5000 optimizer updates, seed42, global batch128 (2 Blackwell GPUs x64, accumulation1), BF16/ZeRO2, gradient checkpointing false. Final5000 checkpoint is published; full optimizer recovery is retained on the training host.
+
+The data use fixed75.87417450998383ms from the archived RTX3090 OpenVLA profile, quantized to one raw frame at10Hz (100ms effective). This is mean profiling. Binary causal eight-decision history24D, current RGB image, native2D NOOP/FLAP action, horizon1. 180train/20validation episodes,3240000/360000rows,200 unique seeds.
+
+Final offline validation action loss: 0.03496555 on 6400 sampled validation frames. Closed-loop latency-bench evaluation is pending until all four mean VLA trainings complete. Offline loss is not task-performance acceptance.
+
+Weights: `checkpoints/model.pt` (pure tensor state_dict; weights_only CPU load and finite-tensor checks passed). `config.yaml` is portable inference config; `config.full.yaml` records actual training configuration. Backbone/tokenizer source: Qwen/Qwen3-VL-4B-Instruct@ebb281ec70b05090aa6165b016eac8ec08e71b17. Task/history contract and dataset normalization are included.
+
+Data: https://huggingface.co/datasets/latency-sensitive-bench/Standard-Pipeline/tree/a568aa3792b49189d72c02781956c46f50d81b4e/flappy/profile_latency/QwenOFT/mean_rtx3090_s42_200ep_v1
+Training: https://wandb.ai/dongqianyu99-zhejiang-university/starvla_tasks/runs/mean3090-flappy-vla-blackwell-20260930
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/checkpoints/model.pt b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/checkpoints/model.pt
new file mode 100644
index 0000000000000000000000000000000000000000..a2e8ebbcf158b732cfa5cf51d3fd3c75dffda68a
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/checkpoints/model.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f784cb7d6c19872b4ac419f8a31459b395c12f23809f561ae134080a74c6e80a
+size 9785137569
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/collection_quality_report.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/collection_quality_report.json
new file mode 100644
index 0000000000000000000000000000000000000000..4f5fcc434563f4e01bab4fe41da69c7aea3e3632
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/collection_quality_report.json
@@ -0,0 +1,18 @@
+{
+  "dataset_repo": "latency-sensitive-bench/Standard-Pipeline",
+  "revision": "a568aa3792b49189d72c02781956c46f50d81b4e",
+  "verified_hashes": 29,
+  "episodes": 200,
+  "train_episodes": 180,
+  "val_episodes": 20,
+  "train_rows": 3240000,
+  "val_rows": 360000,
+  "state_dim": 24,
+  "action_dim": 2,
+  "action_horizon": 1,
+  "mean_latency_ms": 75.87417450998383,
+  "effective_latency_raw_frames": 1,
+  "effective_latency_ms": 100,
+  "data_root": "/home/ubuntu/lzj/mean-profiling/flappy/data",
+  "canonical_data_recollected": false
+}
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/config.full.yaml b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/config.full.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..41a6b161406b35124bd3c10014ba5843535a95e6
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/config.full.yaml
@@ -0,0 +1,271 @@
+framework:
+  name: QwenOFT
+  qwenvl:
+    base_vlm: /home/ubuntu/lzj/models/Qwen3-VL-4B-Instruct
+    attn_implementation: flash_attention_2
+    enable_gradient_checkpointing: false
+  action_model:
+    action_model_type: MLP
+    action_dim: 2
+    state_dim: 24
+    action_hidden_dim: 2560
+    future_action_window_size: 0
+    past_action_window_size: 0
+    loss_type: discrete_ce
+    state_encoding: continuous_projector
+    task_objective: null
+    action_horizon: 1
+    action_env_dim: 2
+  kv_memory:
+    enabled: false
+    window: 4
+    rollout_len: 8
+    packed_train: false
+    rebased_sink: true
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    include_state: true
+    data_root_dir: /home/ubuntu/lzj/mean-profiling/flappy/data
+    data_mix: lerobot
+    eval_data_mix: lerobot__val
+    custom_mixtures_path: /home/ubuntu/lzj/mean-profiling/flappy/data/_generated_mixtures/lerobot.json
+    action_type: discrete
+    sequential_step_sampling: false
+    eval_sequential_step_sampling: null
+    num_workers: 2
+    eval_num_workers: 2
+    prefetch_factor: 4
+    persistent_workers: true
+    pin_memory: true
+    shuffle: true
+    action_balance:
+      enabled: false
+      strategy: balanced_epoch
+      action_key: action_id
+      target_flap_fraction: 0.3
+      noop_id: 0
+      flap_id: 1
+    latency_curriculum:
+      enabled: false
+      strategy: exclusive
+      latencies: null
+      phase_steps: null
+      phase_distributions: null
+      new_latency_passes: 1.0
+      replay_passes: 0.25
+      target_total_passes: 2.0
+      final_equalization: true
+      step_budget_mode: auto
+      eval_at_phase_end: false
+      save_at_phase_end: false
+      computed_plan: null
+    per_device_batch_size: 64
+    load_all_data_for_training: true
+    num_obs_frames: 1
+    image_mode: single
+    prompt_mode: raw
+    stitch_grid:
+    - 2
+    - 2
+    obs_image_size: null
+    video_backend: torchvision_av
+dataset:
+  source_hf: ''
+  config_name: null
+  source_subdir: null
+  converted_name: flappy_train
+  single_source_hf: ''
+  mixed_source_hf: ''
+  single_converted_name: flappy_train
+  mixed_converted_name: flappy_mixed_latency_train
+  single_latency_filter: null
+  mixed_latency_filter: null
+  force_download: false
+  setup_force: false
+  skip_verification: false
+  target_latency_unit: observation_steps
+  verify_rows: 200
+  max_episodes: null
+  episodes_per_latency: null
+  latency_filter: null
+  debug_subset:
+    enabled: false
+    max_episodes: 5
+    suffix: debug
+base_model:
+  repo_id: Qwen/Qwen3-VL-4B-Instruct
+initialization:
+  checkpoint_local_dir: null
+  checkpoint_hf_repo_id: null
+  checkpoint_filename: null
+trainer:
+  max_train_steps: 5000
+  num_warmup_steps: 100
+  save_interval: 500
+  eval_interval: 500
+  eval_num_batches: 50
+  per_latency_eval_num_batches: null
+  eval_action_classification: false
+  eval_action_classification_interval: null
+  cc_f1_tolerance: 1
+  learning_rate:
+    base: 3.0e-05
+    qwen_vl_interface: 3.0e-06
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr_rate: 0.03333333333333333
+  freeze_modules: ''
+  freeze_vit: false
+  freeze_tied_embedding: false
+  freeze_llm_layers: []
+  loss_scale:
+    vla: 1.0
+    vlm: 0.0
+  max_grad_norm: 1.0
+  weight_decay: 0.0
+  logging_frequency: 1
+  profile_timing:
+    enabled: false
+    log_interval: 10
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 1
+  distributed_backend: deepspeed
+  is_resume: false
+  pretrained_checkpoint: null
+  resume_step: 0
+  reload_modules: null
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+    fused: true
+  save_format: pt
+  hist8_diagnostics:
+    steps:
+    - 0
+    - 500
+    - 1000
+    - 2500
+    - 5000
+    sample_count: 2048
+    noise_seeds:
+    - 1729
+    sample_seed: 42
+    batch_size: 8
+    bootstrap_repeats: 4000
+workspace_dir: WORKSPACE_DIR
+run_root_dir: /home/ubuntu/lzj/mean-profiling/flappy/vla
+seed: 42
+wandb_entity: dongqianyu99-zhejiang-university
+wandb_project: starvla_tasks
+auth:
+  env_file: null
+  hf_token_env: HF_TOKEN
+  wandb_api_key_env: WANDB_API_KEY
+paths:
+  run_root_dir: results/Checkpoints
+  dataset_local_dir: playground/Datasets/rl_games
+  dataset_cache_dir: null
+  base_model_dir: /home/ubuntu/lzj/models/Qwen3-VL-4B-Instruct
+  accelerate_config: starVLA/config/deepseeds/deepspeed_zero2.yaml
+launch:
+  use_accelerate: true
+  gpus: null
+  num_processes: 1
+  dry_run: false
+conda:
+  enabled: true
+  env_name: null
+rl_games:
+  model_alias: openvla
+  env_eval:
+    image_size: 224
+    frameskip: 1
+    image_transform: raw_rgb
+    prompt_mode: raw
+    ghost_trail:
+      history_frames: 5
+      gamma: 1.3
+      min_alpha: 35
+      scroll_px_per_step: 4.0
+      ground_fraction: 0.22
+    seed: 42
+    fixed_episode_seeds: true
+    latency_seed_stride: 0
+    task_seed_stride: 0
+    task_description: ''
+    eval_parallel_envs: 5
+    action_chunk_execution:
+      enabled: false
+      chunk_size: null
+    enabled: false
+    eval_backend: latency_bench
+    distributed_mode: none
+    vectorized:
+      enabled: false
+      batch_size: 1
+    latency:
+      prompt_map_path: null
+      mode: single
+      values:
+      - 0
+    mid_train:
+      enabled: false
+      interval_steps: 100
+      latencies:
+      - 0
+      num_episodes: 5
+      max_steps_per_episode: 3600
+    post_train:
+      enabled: false
+      latencies:
+      - 0
+      - 1
+      - 2
+      - 3
+      - 4
+      num_episodes: 5
+      max_steps_per_episode: 3600
+  task: flappy
+  initialization_mode: scratch
+  action_carrier: native
+model: openvla
+env: flappy
+init: scratch
+mode: single
+checkpoint:
+  load: none
+  hf_repo_id: null
+  save_best_model: false
+  save_final_model: true
+  save_pt_file: true
+  save_training_state: true
+  save_safetensors_file: false
+  local:
+    keep_last_n: 2
+  sync:
+    enabled: false
+    repo_id: null
+    keep_last_n: 0
+    sync_every_n_checkpoints: 1
+    resume_policy: local_latest
+run_id: flappy-qwenoft-mean-3090-s42
+output_dir: /home/ubuntu/lzj/mean-profiling/flappy/vla/flappy-qwenoft-mean-3090-s42
+config_yaml: /home/ubuntu/lzj/mean-profiling/flappy/vla.yaml
+is_debug: false
+version_id: '0.21'
+wandb_name: flappy-qwenoft-mean-3090-s42
+wandb_group: mean-profiling-3090-blackwell
+wandb_tags:
+- flappy
+- qwenoft
+- hist8-v2
+- history01
+- scratch
+- profile-latency
+training_latency_condition: mean-profile
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/config.yaml b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..a5bd02b498b3702c16436ae4721301020a7abb4b
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/config.yaml
@@ -0,0 +1,45 @@
+framework:
+  name: QwenOFT
+  qwenvl:
+    base_vlm: Qwen/Qwen3-VL-4B-Instruct
+    attn_implementation: flash_attention_2
+    enable_gradient_checkpointing: false
+  action_model:
+    action_model_type: MLP
+    action_dim: 2
+    state_dim: 24
+    action_hidden_dim: 2560
+    future_action_window_size: 0
+    past_action_window_size: 0
+    loss_type: discrete_ce
+    state_encoding: continuous_projector
+    task_objective: null
+    action_horizon: 1
+    action_env_dim: 2
+  kv_memory:
+    enabled: false
+    window: 4
+    rollout_len: 8
+    packed_train: false
+    rebased_sink: true
+version_id: '0.21'
+datasets:
+  vla_data:
+    include_state: true
+    obs_image_size: null
+    image_mode: single
+    stitch_grid:
+    - 2
+    - 2
+    num_obs_frames: 1
+trainer:
+  pretrained_checkpoint: null
+backbone:
+  repo_id: Qwen/Qwen3-VL-4B-Instruct
+  revision: ebb281ec70b05090aa6165b016eac8ec08e71b17
+  allow_patterns:
+  - '*.json'
+  - '*.safetensors'
+  - '*.txt'
+  - '*.jinja'
+  - '*.model'
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/dataset_statistics.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/dataset_statistics.json
new file mode 100644
index 0000000000000000000000000000000000000000..6736f190a37ec0d0731999f6be04246820d76d39
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/dataset_statistics.json
@@ -0,0 +1,194 @@
+{
+  "new_embodiment": {
+    "action": {
+      "mean": [
+        0.9008839726448059,
+        0.09911604970693588
+      ],
+      "std": [
+        0.30056583881378174,
+        0.30056583881378174
+      ],
+      "max": [
+        1.0,
+        1.0
+      ],
+      "min": [
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true
+      ]
+    },
+    "state": {
+      "mean": [
+        0.9995555281639099,
+        0.09905370324850082,
+        0.09904845803976059,
+        0.9996111392974854,
+        0.09906204044818878,
+        0.09905370324850082,
+        0.999666690826416,
+        0.09907067567110062,
+        0.09906204044818878,
+        0.9997222423553467,
+        0.09909043461084366,
+        0.09907067567110062,
+        0.9997777938842773,
+        0.09909876435995102,
+        0.09909043461084366,
+        0.999833345413208,
+        0.09910740703344345,
+        0.09909876435995102,
+        0.9998888969421387,
+        0.09911327064037323,
+        0.09910740703344345,
+        0.9999444484710693,
+        0.09911450743675232,
+        0.09911327064037323
+      ],
+      "std": [
+        0.02107249759137688,
+        0.30049213767051697,
+        0.30048587918281555,
+        0.019712645560502656,
+        0.30050185322761536,
+        0.30049213767051697,
+        0.018251324072479345,
+        0.30051174759864807,
+        0.30050185322761536,
+        0.016662079840898097,
+        0.3005351126194,
+        0.30051174759864807,
+        0.014903813600540161,
+        0.3005451261997223,
+        0.3005351126194,
+        0.012907790020110239,
+        0.3005552887916565,
+        0.3005451261997223,
+        0.01053975615650543,
+        0.30056232213974,
+        0.3005552887916565,
+        0.0074531468562811926,
+        0.3005639910697937,
+        0.30056232213974
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ]
+    },
+    "num_transitions": 3240000,
+    "num_trajectories": 180
+  }
+}
\ No newline at end of file
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/dataset_statistics_eval.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/dataset_statistics_eval.json
new file mode 100644
index 0000000000000000000000000000000000000000..1e5ffe2d5a01eca477cd301def0bbcbecc6bf721
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/dataset_statistics_eval.json
@@ -0,0 +1,194 @@
+{
+  "new_embodiment": {
+    "action": {
+      "mean": [
+        0.9010527729988098,
+        0.09894721955060959
+      ],
+      "std": [
+        0.2985910475254059,
+        0.2985910475254059
+      ],
+      "max": [
+        1.0,
+        1.0
+      ],
+      "min": [
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true
+      ]
+    },
+    "state": {
+      "mean": [
+        0.9995555281639099,
+        0.09889999777078629,
+        0.09889444708824158,
+        0.9996111392974854,
+        0.09890277683734894,
+        0.09889999777078629,
+        0.999666690826416,
+        0.09890833497047424,
+        0.09890277683734894,
+        0.9997222423553467,
+        0.09892222285270691,
+        0.09890833497047424,
+        0.9997777938842773,
+        0.09893333166837692,
+        0.09892222285270691,
+        0.999833345413208,
+        0.09894166886806488,
+        0.09893333166837692,
+        0.9998888969421387,
+        0.09894444793462753,
+        0.09894166886806488,
+        0.9999444484710693,
+        0.09894721955060959,
+        0.09894444793462753
+      ],
+      "std": [
+        0.02107250504195682,
+        0.29852861166000366,
+        0.2985208034515381,
+        0.019712593406438477,
+        0.2985321879386902,
+        0.29852861166000366,
+        0.01825134269893312,
+        0.2985396385192871,
+        0.2985321879386902,
+        0.01666202954947938,
+        0.2985576391220093,
+        0.2985396385192871,
+        0.014903815463185194,
+        0.29857251048088074,
+        0.2985576391220093,
+        0.01290779560804152,
+        0.29858362674713135,
+        0.29857251048088074,
+        0.010539757087829197,
+        0.2985873222351074,
+        0.29858362674713135,
+        0.0074531468562811926,
+        0.2985910475254059,
+        0.2985873222351074
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ]
+    },
+    "num_transitions": 360000,
+    "num_trajectories": 20
+  }
+}
\ No newline at end of file
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/files.sha256.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/files.sha256.json
new file mode 100644
index 0000000000000000000000000000000000000000..319c5ce60ad20d325db1b70348a477403b0ec0b0
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/files.sha256.json
@@ -0,0 +1,102 @@
+{
+  "README.md": {
+    "bytes": 1513,
+    "sha256": "1f4036fe39bcc14657ad3dce517a5ee46ae764e030e38c60c808b7d29859b255"
+  },
+  "checkpoints/model.pt": {
+    "bytes": 9785137569,
+    "sha256": "f784cb7d6c19872b4ac419f8a31459b395c12f23809f561ae134080a74c6e80a"
+  },
+  "collection_quality_report.json": {
+    "bytes": 532,
+    "sha256": "0741f81a3b2bd9cf7b55d7b8478695532968bcf64abb7640535bea425bac6590"
+  },
+  "config.full.yaml": {
+    "bytes": 6375,
+    "sha256": "d6da13357b405aafb9f5b6e07934a2cd59a28ab3a7578fd81b01bcfea40e8eb2"
+  },
+  "config.yaml": {
+    "bytes": 979,
+    "sha256": "6c33e3ebf0f2a5327163329289b4b225c418bc01f65730ff8f7599702e50d61f"
+  },
+  "dataset_statistics.json": {
+    "bytes": 3299,
+    "sha256": "2c7e8a74a07ca7bec408a0e32cd00957d1af68ecc58fb83571cea5258fe178d7"
+  },
+  "dataset_statistics_eval.json": {
+    "bytes": 3302,
+    "sha256": "a57b1ff610f31dbd61bdb28f5c176e5945adfaa03be75f590e7cd550f658c3dc"
+  },
+  "latency_prompt_map.json": {
+    "bytes": 189,
+    "sha256": "e5bdad8c4747fb2e8101f5d7265e89e2c8794c573ee42bac764365a42b6bc3a4"
+  },
+  "launch_audit.json": {
+    "bytes": 2140,
+    "sha256": "6441da61d2da9190ddc6fd4d7a031acd1d455fc102c3ac40ae3e997aada1fb95"
+  },
+  "manifest.json": {
+    "bytes": 2295,
+    "sha256": "0ee6164ded3e87ffb0ea956b700bc1fae82c1a4de2c3f4b432163fa070f40fc4"
+  },
+  "provenance.json": {
+    "bytes": 3432,
+    "sha256": "17416c9020e50a6abac5e82f2001ec7e0962b479cbde84702061af45ec5116d9"
+  },
+  "summary.jsonl": {
+    "bytes": 159,
+    "sha256": "2905acb79c5fffb33fd4689d87af3bd05d74f4184b9d36ff3eae0c64deecd930"
+  },
+  "task_contract.json": {
+    "bytes": 982,
+    "sha256": "958afa1b69927f50ef9f92e50875b0f6ec937e98ce0a7a86732342a05b94f76c"
+  },
+  "training-code/QWen3.py": {
+    "bytes": 17586,
+    "sha256": "5fe76fc98eaa84e215c3c25c11326815499d2a4b0683c6ab95382b881f535370"
+  },
+  "training-code/QwenOFT.py": {
+    "bytes": 62854,
+    "sha256": "acad1b0dcde6e9fed6e2912b6534db0f8e33a74ec7f6630bbf1b0161fc298b0d"
+  },
+  "training-code/datasets.py": {
+    "bytes": 154347,
+    "sha256": "33e22556471128a2bad18a1ef8b84c399847b4d98b2979d785d139de0cdf36df"
+  },
+  "training-code/gym-language-contract.patch": {
+    "bytes": 688,
+    "sha256": "9f5dee2c5b52d5507aff27373a8cf6def4959361d2c5264d5079b416f28a4c49"
+  },
+  "training-code/loader-spawn-cache.patch": {
+    "bytes": 711,
+    "sha256": "24532ce84e5aa9aef381493fe699b4bb4987afd954422191ae852b0f5eb5b2cd"
+  },
+  "training-code/loader-spawn-test.patch": {
+    "bytes": 1500,
+    "sha256": "a36db15b13d68defeecb75ddc9efa1ac202d9aa0a40f66ba5baaee2dcaef21f1"
+  },
+  "training-code/mean-profile-preparation.patch": {
+    "bytes": 7053,
+    "sha256": "de5fe558376c98515a1cfb186ffdc09b618acc207eb8a5f667b0cbb303767c09"
+  },
+  "training-code/train_starvla.py": {
+    "bytes": 125159,
+    "sha256": "f872648f941cec272af724772e9d6fc3721e9ed6eaf0c8683ab285eb9396212e"
+  },
+  "training-code/trainer_tools.py": {
+    "bytes": 29144,
+    "sha256": "9a48bdfb99b405b03fececfc1e051d3ebb179eb423c15141f02b9be62e2e1791"
+  },
+  "training-code/worker_context.py": {
+    "bytes": 2279,
+    "sha256": "8493631104269f0630655ee2faae592a84ebfb3119aca142c8ac2701f5214464"
+  },
+  "training_code_snapshot.json": {
+    "bytes": 626,
+    "sha256": "136c2cbb92daa93a6dab1aed5644e57491da4f95075c7f4f5d12b3ca8aa16adb"
+  },
+  "training_results.json": {
+    "bytes": 512,
+    "sha256": "b08fa8927249eed1bb0a1c82fc7a75e59e05d7c2fbaf4c6c8ed6f14dc214d488"
+  }
+}
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/latency_prompt_map.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/latency_prompt_map.json
new file mode 100644
index 0000000000000000000000000000000000000000..465788c1cf15573a239a3146b6aca2bb8b5deaa4
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/latency_prompt_map.json
@@ -0,0 +1,7 @@
+{
+  "1": {
+    "prompt": "You are playing Flappy Bird. Pass through the pipe gaps and stay alive. Choose the action: NOOP, FLAP.",
+    "latency_raw_frames": 1,
+    "latency_ms": 100.0
+  }
+}
\ No newline at end of file
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/launch_audit.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/launch_audit.json
new file mode 100644
index 0000000000000000000000000000000000000000..7a977514c1f2e540ddf1577263d34a389c78b83f
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/launch_audit.json
@@ -0,0 +1,62 @@
+{
+  "host": "ubuntu@136.115.188.185",
+  "reference_repo_revision": "180da0effdaaaeacef218cf5fbaa4db41af67495",
+  "reference": "/home/ubuntu/lzj/mean-profiling/reference/checkpoint-configs/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-hist8-scratch-s42-v2/config.full.yaml",
+  "config_sha256": "0d8e33d18da69ac470a3aac5973babdb36b3eb8bfaa02c0f0d3ae8aeb912850f",
+  "base_revision": "ebb281ec70b05090aa6165b016eac8ec08e71b17",
+  "physical_gpus": [
+    0,
+    1
+  ],
+  "micro_batch_per_gpu": 64,
+  "gradient_accumulation": 1,
+  "global_batch": 128,
+  "gradient_checkpointing": false,
+  "optimizer_updates": 5000,
+  "seed": 42,
+  "horizon": 1,
+  "preflight_optimizer_updates": 20,
+  "preflight_discarded": true,
+  "formal_start_step": 0,
+  "preflight_gpu_peaks": [
+    {
+      "rank": "0",
+      "logical_gpu": 0,
+      "gpu_peak_allocated_gib": 92.39531517028809,
+      "gpu_peak_reserved_gib": 92.84375,
+      "main_peak_rss_gib": 18.482311248779297,
+      "wall_seconds_with_initialization": 66.88179596900227,
+      "optimizer_updates": 20,
+      "formal_updates": 0,
+      "global_batch": 128,
+      "micro_per_gpu": 64,
+      "gradient_accumulation": 1,
+      "gradient_checkpointing": false,
+      "status": "passed"
+    },
+    {
+      "rank": "1",
+      "logical_gpu": 1,
+      "gpu_peak_allocated_gib": 92.39531517028809,
+      "gpu_peak_reserved_gib": 93.5078125,
+      "main_peak_rss_gib": 18.478984832763672,
+      "wall_seconds_with_initialization": 66.89783567699851,
+      "optimizer_updates": 20,
+      "formal_updates": 0,
+      "global_batch": 128,
+      "micro_per_gpu": 64,
+      "gradient_accumulation": 1,
+      "gradient_checkpointing": false,
+      "status": "passed"
+    }
+  ],
+  "data_workers_per_rank": 2,
+  "eval_workers_per_rank": 2,
+  "spawn_cache_payload_fix_test": "9 relevant tests passed",
+  "memory_high_gib": 400,
+  "memory_max_gib": 450,
+  "post_train_rollout_eval": false,
+  "evaluation": "pending until four VLA trainings complete",
+  "allocator": "expandable_segments:True",
+  "formal_retry_reason": "CUDA allocator fragmentation/near-capacity initial failure, 20 same-path steps verified"
+}
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/manifest.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/manifest.json
new file mode 100644
index 0000000000000000000000000000000000000000..4c89937ccfd7ef0b0c61ca46519b17784507ba0a
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/manifest.json
@@ -0,0 +1,111 @@
+{
+  "dataset_name": "lerobot",
+  "env_name": "flappy",
+  "episodes": 180,
+  "frames": 3240000,
+  "task_prompts": [
+    "You are playing Flappy Bird. Pass through the pipe gaps and stay alive. Choose the action: NOOP, FLAP."
+  ],
+  "format": "starvla_lerobot_v2_decision_v1",
+  "source": "/home/ubuntu/lzj/mean-profiling/flappy/data/raw",
+  "robot_type": "rl_games_flappy",
+  "action_layout": "flappy_discrete",
+  "source_action_layout": "flappy_discrete",
+  "action_labels": [
+    "NOOP",
+    "FLAP"
+  ],
+  "action_dim": 2,
+  "active_action_dim": 2,
+  "state_source": "decision_v1",
+  "state_dim": 24,
+  "state_labels": [
+    "decision_history_0",
+    "decision_history_1",
+    "decision_history_2",
+    "decision_history_3",
+    "decision_history_4",
+    "decision_history_5",
+    "decision_history_6",
+    "decision_history_7",
+    "decision_history_8",
+    "decision_history_9",
+    "decision_history_10",
+    "decision_history_11",
+    "decision_history_12",
+    "decision_history_13",
+    "decision_history_14",
+    "decision_history_15",
+    "decision_history_16",
+    "decision_history_17",
+    "decision_history_18",
+    "decision_history_19",
+    "decision_history_20",
+    "decision_history_21",
+    "decision_history_22",
+    "decision_history_23"
+  ],
+  "state_normalization": {
+    "min": [
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0
+    ],
+    "max": [
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0
+    ]
+  },
+  "action_history_decisions": 8,
+  "rows_unit": "decision_step",
+  "obs_fps": 10.0,
+  "obs_stride_raw_frames": 1,
+  "latency_prompt_map_path": "/home/ubuntu/lzj/mean-profiling/flappy/data/lerobot/latency_prompt_map.json",
+  "custom_mixtures_path": "/home/ubuntu/lzj/mean-profiling/flappy/data/_generated_mixtures/lerobot.json",
+  "validation_dataset_name": "lerobot__val",
+  "validation_episodes": 20,
+  "validation_frames": 360000
+}
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/provenance.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/provenance.json
new file mode 100644
index 0000000000000000000000000000000000000000..096ffd593812c859d934bd03e9828e322adb7979
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/provenance.json
@@ -0,0 +1,93 @@
+{
+  "task": "flappy",
+  "model": "qwenoft",
+  "run_id": "flappy-qwenoft-mean-3090-s42",
+  "training_latency_condition": "mean-profile",
+  "training_hardware": {
+    "gpu": "NVIDIA RTX PRO6000 Blackwell Server Edition",
+    "count": 2,
+    "physical_indices": [
+      0,
+      1
+    ]
+  },
+  "latency_profile_hardware": "1x-RTX3090",
+  "fixed_latency_ms": 75.87417450998383,
+  "effective_latency_raw_frames": 1,
+  "env_fps": 10,
+  "obs_fps": 10,
+  "horizon": 1,
+  "history_encoding": "binary_0_1",
+  "state_dim": 24,
+  "action_dim": 2,
+  "global_batch": 128,
+  "micro_per_gpu": 64,
+  "gradient_accumulation": 1,
+  "gradient_checkpointing": false,
+  "precision": "bf16",
+  "optimizer_backend": "ZeRO2",
+  "teacher": {
+    "repo_id": "latency-sensitive-bench/benchmark-models",
+    "revision": "3ca75d71b55ff42889a4f058d6ad90b2e9b3668f",
+    "path": "/home/ubuntu/lzj/mean-profiling/flappy/restored-teacher/latency-aware/flappy/small-policy/sample-factory-qwenoft-rtx3090-mean-hist8-capacity1-s42-v1",
+    "local_checkpoint": "/home/ubuntu/lzj/mean-profiling/flappy/restored-teacher/latency-aware/flappy/small-policy/sample-factory-qwenoft-rtx3090-mean-hist8-capacity1-s42-v1/best_000048808_25034752_reward_2216.841.pth",
+    "env_steps": 25034752,
+    "hash_verified_files": 8,
+    "weights_only_load": true
+  },
+  "dataset": {
+    "dataset_repo": "latency-sensitive-bench/Standard-Pipeline",
+    "revision": "a568aa3792b49189d72c02781956c46f50d81b4e",
+    "verified_hashes": 29,
+    "episodes": 200,
+    "train_episodes": 180,
+    "val_episodes": 20,
+    "train_rows": 3240000,
+    "val_rows": 360000,
+    "state_dim": 24,
+    "action_dim": 2,
+    "action_horizon": 1,
+    "mean_latency_ms": 75.87417450998383,
+    "effective_latency_raw_frames": 1,
+    "effective_latency_ms": 100,
+    "data_root": "/home/ubuntu/lzj/mean-profiling/flappy/data",
+    "canonical_data_recollected": false
+  },
+  "source_reference_revision": "180da0effdaaaeacef218cf5fbaa4db41af67495",
+  "runtime": {
+    "repo": "c3c6a39365a151e9b7a5e215452fd64e957c2b29",
+    "starvla": "ccca13c5177fe3d3c884b6e2de4965d916016649",
+    "runtime_fixes": [
+      "mean-profile-preparation.patch",
+      "gym-language-contract.patch",
+      "loader-spawn-cache.patch",
+      "loader-spawn-test.patch"
+    ],
+    "pytorch3d": {
+      "revision": "33824be3cbc87a7dd1db0f6a9a9de9ac81b2d0ba",
+      "build": "transforms-only, no native render extension; QwenOFT uses transforms only"
+    },
+    "decord": {
+      "version": "0.6.0",
+      "build": "official source CPU decoder CP310",
+      "wheel_sha256": "e193b356b1e984b4eff08d23b62e482c2c9e5037a6efdc0b1af47079ae2e4c47"
+    }
+  },
+  "wandb_url": "https://wandb.ai/dongqianyu99-zhejiang-university/starvla_tasks/runs/mean3090-flappy-vla-blackwell-20260930",
+  "training_results": {
+    "training_complete": true,
+    "optimizer_updates": 5000,
+    "selection": "final_5000",
+    "train_loss": 0.028774918988347054,
+    "validation_action_loss": 0.034965552389621735,
+    "validation_samples": 6400,
+    "final_backbone_lr": 1e-07,
+    "final_action_lr": 3.3333333333333333e-06,
+    "closed_loop_evaluation": "pending until all four VLA trainings complete",
+    "tensor_keys": 732,
+    "serialized_tensor_elements_including_tied_aliases": 4892433410,
+    "weights_only_load": "passed",
+    "all_tensors_finite": true
+  },
+  "optimizer_recovery_states": "retained on training host; this bundle contains inference weights"
+}
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/summary.jsonl b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/summary.jsonl
new file mode 100644
index 0000000000000000000000000000000000000000..a0e21985e6b7ed5b88fddcc5f74f6ec45063b603
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/summary.jsonl
@@ -0,0 +1,10 @@
+{"steps": 500}
+{"steps": 1000}
+{"steps": 1500}
+{"steps": 2000}
+{"steps": 2500}
+{"steps": 3000}
+{"steps": 3500}
+{"steps": 4000}
+{"steps": 4500}
+{"steps": 5000}
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/task_contract.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/task_contract.json
new file mode 100644
index 0000000000000000000000000000000000000000..98570c6679a6f0d6328807e819406d4bb2c7999d
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/task_contract.json
@@ -0,0 +1,43 @@
+{
+  "task": "flappy",
+  "model": "qwenoft",
+  "action_horizon": 1,
+  "image_frames": 1,
+  "state_source": "transport",
+  "state_dim": 24,
+  "history_decisions": 8,
+  "history_encoding": "binary_0_1",
+  "history_normalization": "identity",
+  "history_reset": "all_zero",
+  "history_order": "oldest_to_newest",
+  "history_entry_fields": [
+    "admission_bit",
+    "admitted_command_non_noop_bits",
+    "last_applied_action_non_noop_bits"
+  ],
+  "history_action_bit_order": [
+    "FLAP"
+  ],
+  "action_dim": 2,
+  "action_layout": "flappy_discrete",
+  "action_labels": [
+    "NOOP",
+    "FLAP"
+  ],
+  "task_prompts": [
+    "You are playing Flappy Bird. Pass through the pipe gaps and stay alive. Choose the action: NOOP, FLAP."
+  ],
+  "env_fps": 10,
+  "obs_fps": 10,
+  "evaluation": {
+    "episodes": 100,
+    "seeds": [
+      1000000,
+      1000099
+    ],
+    "parallel_envs": 32,
+    "capacity": 1,
+    "max_raw_frames": 3600,
+    "protocol_version": "hist8-history01-par32-v3"
+  }
+}
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/QWen3.py b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/QWen3.py
new file mode 100644
index 0000000000000000000000000000000000000000..21484c44de5403fbc0f71ad3c5bfec321d6c738e
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/QWen3.py
@@ -0,0 +1,444 @@
+# Copyright 2025 starVLA community. All rights reserved.
+# Licensed under the MIT License, Version 1.0 (the "License");
+# Implemented by [Jinhui YE / HKUST University] in [2025].
+
+import time
+from collections import OrderedDict
+from functools import wraps
+from inspect import signature
+from typing import Optional
+
+import torch
+from starVLA.model.profiling import stage_timer as _stage
+from starVLA.training.trainer_utils import initialize_overwatch
+from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+logger = initialize_overwatch(__name__)
+
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = 151655
+VIDEO_TOKEN_INDEX = 151656
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_VIDEO_TOKEN = "<video>"
+
+_ACTION_TOKEN_MIN = 151669  # how can we know this range? check how you add fast tokens into VLM
+_ACTION_TOKEN_MAX = (
+    153716  # here only for fast_tokenizer, see starVLA/model/modules/vlm/tools/add_qwen_special_tokens/README.md
+)
+
+
+import torch.nn as nn
+
+
+def _patch_qwen3vl_flash_attention_position_ids() -> None:
+    from transformers.models.qwen3_vl import modeling_qwen3_vl as qwen3_vl
+
+    text_attention_cls = qwen3_vl.Qwen3VLTextAttention
+    if hasattr(text_attention_cls, "_starvla_flash_attention_position_ids_patched"):
+        return
+
+    original_forward = text_attention_cls.forward
+
+    def forward_without_flash_attention_position_ids(self, *args, **kwargs):
+        # Qwen3-VL already applies M-RoPE through position_embeddings before
+        # attention. Passing the repeated temporal M-RoPE ids into HF's FA2
+        # wrapper makes Transformers 4.57 mis-detect packed sequences and build
+        # empty cu_seqlens on KV-memory cache steps.
+        if self.config._attn_implementation == "flash_attention_2" and "position_ids" in kwargs:
+            del kwargs["position_ids"]
+        return original_forward(self, *args, **kwargs)
+
+    text_attention_cls.forward = forward_without_flash_attention_position_ids
+    text_attention_cls._starvla_flash_attention_position_ids_patched = True
+
+
+# The fixed Triton stage counts keep the PyTorch 2.11 CUDA 12.8 kernel
+# within the shared-memory budget on sm_120.
+_FLEX_KERNEL_OPTIONS = {
+    "BACKEND": "TRITON",
+    "fwd_num_stages": 2,
+    "bwd_num_stages": 1,
+}
+
+
+def _patch_qwen3vl_flex_attention_support() -> None:
+    """Allow loading Qwen3-VL with ``attn_implementation='flex_attention'``.
+
+    Qwen3-VL's attention already dispatches through the generic
+    ``ALL_ATTENTION_FUNCTIONS`` interface (which includes the ``flex_attention``
+    integration), but the model classes conservatively declare
+    ``_supports_flex_attn = False``, so Transformers 4.57 refuses to load with
+    flex at __init__. Flipping the flag lets the official dispatch path run; the
+    packed KV-memory training (`QwenOFT._forward_memory_packed`) passes a
+    BlockMask that FlexAttention consumes natively.
+
+    Also wraps the registered ``flex_attention`` interface to inject the fixed
+    Triton kernel options. Wrapping the dispatch entry — rather than threading
+    ``kernel_options`` through the top-level forward kwargs — guarantees both the
+    forward and the gradient-checkpointing backward recompile pick it up.
+    """
+    from transformers.models.qwen3_vl import modeling_qwen3_vl as qwen3_vl
+
+    for cls_name in (
+        "Qwen3VLForConditionalGeneration",
+        "Qwen3VLModel",
+        "Qwen3VLTextModel",
+        "Qwen3VLPreTrainedModel",
+    ):
+        cls = getattr(qwen3_vl, cls_name, None)
+        if cls is not None:
+            cls._supports_flex_attn = True
+
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    base_flex = ALL_ATTENTION_FUNCTIONS["flex_attention"]
+    if getattr(base_flex, "_starvla_kernel_options_patched", False):
+        return
+
+    def flex_attention_with_kernel_options(*args, **kwargs):
+        kwargs["kernel_options"] = _FLEX_KERNEL_OPTIONS
+        return base_flex(*args, **kwargs)
+
+    flex_attention_with_kernel_options._starvla_kernel_options_patched = True
+    ALL_ATTENTION_FUNCTIONS["flex_attention"] = flex_attention_with_kernel_options
+
+
+def _position_tensor_key(tensor):
+    if tensor is None:
+        return None
+    return (tensor.device, tensor.dtype, tuple(tensor.shape),
+            tuple(tensor.detach().cpu().reshape(-1).tolist()))
+
+
+def _cache_position_method(module, name, weight):
+    """Cache layout-only outputs on one native model instance, never image features."""
+    original = getattr(module, name)
+    parameters = signature(original)
+    cache = OrderedDict()
+
+    @wraps(original)
+    def cached(*args, **kwargs):
+        if module.training or torch.is_grad_enabled():
+            cache.clear()
+            return original(*args, **kwargs)
+        arguments = parameters.bind(*args, **kwargs)
+        arguments.apply_defaults()
+        parameter = weight()
+        key = (
+            parameter.device, parameter.dtype, parameter._version,
+            tuple(_position_tensor_key(value) for value in arguments.arguments.values()),
+        )
+        if key not in cache:
+            cache[key] = original(*args, **kwargs)
+            if len(cache) > 8:
+                cache.popitem(last=False)
+        cache.move_to_end(key)
+        return cache[key]
+
+    setattr(module, name, cached)
+    return cache
+
+
+class _QWen3_VL_Interface(nn.Module):
+    """
+    This exists because of the diversity of VLMs, so we encapsulate the changes here.
+    Lightweight wrapper around Qwen3-VL (Qwen3VLForConditionalGeneration).
+
+    Purpose:
+        - Unify interface with other VLM backends (CausalLM-like usage).
+        - Centralize preprocessing (tokenization + multimodal packing).
+        - Provide consistent forward / generate signatures.
+
+    """
+
+    def __init__(self, config: Optional[dict] = None, **kwargs):
+        """
+        Initialize the Qwen3-VL wrapper.
+        Following https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct
+
+        """
+        super().__init__()
+
+        qwenvl_config = config.framework.get("qwenvl", {})
+        model_id = qwenvl_config.get("base_vlm", "Qwen/Qwen3-VL-4B-Instruct")
+        attn_implementation = qwenvl_config.get("attn_implementation", "sdpa")
+        enable_grad_ckpt = qwenvl_config.get("enable_gradient_checkpointing", False)
+        print(
+            f"[QWen3] loading {model_id} with gradient_checkpointing={enable_grad_ckpt}",
+            flush=True,
+        )
+
+        if attn_implementation == "flash_attention_2":
+            import flash_attn_2_cuda  # noqa: F401
+
+            _patch_qwen3vl_flash_attention_position_ids()
+
+        if attn_implementation == "flex_attention":
+            _patch_qwen3vl_flex_attention_support()
+
+        model = Qwen3VLForConditionalGeneration.from_pretrained(
+            model_id,
+            attn_implementation=attn_implementation,
+            dtype=torch.bfloat16,
+            ignore_mismatched_sizes=True, # resize image no longer needed? @TODO check bug
+        )
+        processor = AutoProcessor.from_pretrained(model_id)
+        processor.tokenizer.padding_side = "left"
+
+        if enable_grad_ckpt:
+            model.config.text_config.use_cache = False
+            model.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs={"use_reentrant": False}
+            )
+            model.enable_input_require_grads()
+            print(
+                "[QWen3] gradient_checkpointing ENABLED "
+                f"(use_reentrant=False, active={model.is_gradient_checkpointing}, "
+                f"text_use_cache={model.config.text_config.use_cache})",
+                flush=True,
+            )
+
+        self.model = model
+        self.processor = processor
+        self.config = config
+        # align qwen3 with qwen2.5
+        self.model.config.hidden_size = self.model.config.text_config.hidden_size
+        self._last_build_timing = {}
+        visual = model.model.visual
+        self._position_caches = [
+            _cache_position_method(model.model, "get_rope_index", lambda: visual.rotary_pos_emb.inv_freq),
+            _cache_position_method(visual, "rot_pos_emb", lambda: visual.rotary_pos_emb.inv_freq),
+            _cache_position_method(visual, "fast_pos_embed_interpolate", lambda: visual.pos_embed.weight),
+        ]
+
+        # only for fast base model
+        if "-Action" in model_id:
+            self._ACTION_TOKEN_MIN = _ACTION_TOKEN_MIN
+            self._ACTION_TOKEN_MAX = _ACTION_TOKEN_MAX
+
+    def _clear_position_caches(self):
+        for cache in self._position_caches:
+            cache.clear()
+
+    def train(self, mode=True):
+        self._clear_position_caches()
+        return super().train(mode)
+
+    def _apply(self, fn, recurse=True):
+        self._clear_position_caches()
+        return super()._apply(fn, recurse=recurse)
+
+    def _load_from_state_dict(self, *args, **kwargs):
+        self._clear_position_caches()
+        return super()._load_from_state_dict(*args, **kwargs)
+
+    def _profile_timing_enabled(self) -> bool:
+        return self.config.trainer.profile_timing.enabled
+
+    def _profile_sync(self) -> None:
+        if self._profile_timing_enabled() and torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    def forward(
+        self,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        """
+        Forward pass delegating to underlying Qwen2.5-VL backbone.
+        """
+
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            outputs = self.model(
+                **kwargs,
+            )
+
+        return outputs
+
+    def forward_last_hidden(self, **kwargs) -> torch.Tensor:
+        """
+        Run the Qwen3-VL backbone and return only the final hidden state.
+
+        OFT-style action heads only consume the last hidden layer. Calling the
+        full conditional-generation wrapper with ``output_hidden_states=True``
+        returns every layer's activations, which can largely erase the memory
+        benefit of gradient checkpointing.
+        """
+
+        kwargs = dict(kwargs)
+        kwargs.pop("output_hidden_states", None)
+        kwargs.pop("output_attentions", None)
+        kwargs.setdefault("return_dict", True)
+        kwargs.setdefault("use_cache", False)
+
+        def _wrapper_forward_last_hidden() -> torch.Tensor:
+            lm_head = getattr(self.model, "lm_head", None)
+            captured = {}
+            handle = None
+            if lm_head is not None:
+                def _capture_lm_head_input(_module, inputs):
+                    captured["last_hidden"] = inputs[0]
+
+                handle = lm_head.register_forward_pre_hook(_capture_lm_head_input)
+
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                try:
+                    outputs = self.model(
+                        **kwargs,
+                        output_attentions=False,
+                        output_hidden_states=False,
+                        return_dict=True,
+                    )
+                finally:
+                    if handle is not None:
+                        handle.remove()
+
+            if "last_hidden" in captured:
+                return captured["last_hidden"]
+
+            # Last-resort compatibility path for unusual model wrappers that do
+            # not expose lm_head. This preserves correctness, but returns every
+            # layer and therefore saves less memory.
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                outputs = self.model(
+                    **kwargs,
+                    output_attentions=False,
+                    output_hidden_states=True,
+                    return_dict=True,
+                )
+            return outputs.hidden_states[-1]
+
+        backbone = getattr(self.model, "model", None)
+        if backbone is None:
+            return _wrapper_forward_last_hidden()
+
+        try:
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                outputs = backbone(**kwargs)
+        except TypeError:
+            # Some Transformers Qwen-VL wrappers inject visual embeddings before
+            # calling the text backbone, so the backbone itself may not accept
+            # pixel_values/image_grid_thw. Keep those versions working.
+            return _wrapper_forward_last_hidden()
+
+        if hasattr(outputs, "last_hidden_state"):
+            return outputs.last_hidden_state
+        return outputs[0]
+
+    def generate(
+        self,
+        **kwargs,
+    ):
+        """
+        High-level generation interface (auto-regressive decoding), optionally vision-conditioned.
+
+        Args:
+            **kwargs: fully follow raw model.generate() signature.
+        Returns:
+            GenerateOutput | Model-dependent generation return.
+        """
+        with torch.autocast("cuda", dtype=torch.float16):
+            generation_output = self.model.generate(
+                **kwargs,
+            )
+        return generation_output
+
+    def build_qwenvl_inputs(self, images, instructions, solutions=None, profiler=None, **kwargs):
+        """
+        Build model inputs from raw data (images + instructions + optional solutions).
+        Follow Oficial Qwen3-VL Instruct format: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct
+        """
+
+        # Create messages: one message per sample
+        with _stage(profiler, "starvla_qwen_message_build_ms"):
+            messages = []
+            assert len(images) == len(instructions), "Images and instructions must have the same length"
+            for imgs, instruction in zip(images, instructions):
+                content = [{"type": "image", "image": img} for img in imgs]
+
+                if "CoT_prompt" in self.config.datasets.vla_data:  # If using a grounding prompt to task
+                    CoT_prompt = self.config.datasets.vla_data.get("CoT_prompt", "")
+                    prompt = CoT_prompt.replace("{instruction}", instruction)
+                else:
+                    prompt = instruction
+
+                content.append({"type": "text", "text": prompt})
+                msg = [{"role": "user", "content": content}]
+
+                if solutions is not None:
+                    solution = solutions[len(messages)]
+                    msg.append({"role": "assistant", "content": [{"type": "text", "text": solution}]})
+                messages.append(msg)
+
+        profile_timing = self._profile_timing_enabled()
+        build_timing = {}
+
+        t_processor = time.perf_counter() if profile_timing else None
+        with _stage(profiler, "starvla_qwen_processor_ms"):
+            batch_inputs = self.processor.apply_chat_template(
+                messages, tokenize=True, padding=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
+            )
+        if profile_timing:
+            build_timing["timing/qwen_processor"] = time.perf_counter() - t_processor
+
+        # if solutions, mask out the solution tokens in labels
+        if solutions is not None:  #  here only for fast_tokenizer now.
+            action_token_min = _ACTION_TOKEN_MIN  # how can we know this range? --> we has other way for this, but is slower see qwenhelix branch
+            action_token_max = _ACTION_TOKEN_MAX  # here only for fast_tokenizer, see starVLA/model/modules/vlm/tools/add_qwen_special_tokens/README.md
+            labels = batch_inputs["input_ids"].clone()
+            # For each sequence in the batch, find the first occurrence of an action token.
+            for i in range(labels.size(0)):
+                seq = labels[i]
+                # Create a mask for tokens within the action token range.
+                mask_seq = (seq >= action_token_min) & (seq <= action_token_max)
+                nonzero_indices = torch.nonzero(mask_seq, as_tuple=False)
+                if nonzero_indices.numel() > 0:
+                    first_action_index = nonzero_indices[0].item()
+                    # Mask out all tokens before the first action token.
+                    seq[:first_action_index] = IGNORE_INDEX
+                else:
+                    # If no action token is found, mask the entire sequence.
+                    seq[:] = IGNORE_INDEX
+                    RuntimeWarning(
+                        "action token are on in yout tokenizer, plz see starVLA/model/modules/vlm/tools/add_qwen_special_tokens/README.md."
+                    )
+
+            labels[labels == self.processor.tokenizer.pad_token_id] = -100  ## mask out pad tokens as well
+            batch_inputs["labels"] = labels
+
+        t_h2d = time.perf_counter() if profile_timing else None
+        with _stage(profiler, "starvla_qwen_h2d_ms"):
+            batch_inputs = batch_inputs.to(self.model.device)
+        if profile_timing:
+            self._profile_sync()
+            build_timing["timing/qwen_h2d"] = time.perf_counter() - t_h2d
+        self._last_build_timing = build_timing
+        return batch_inputs
+
+
+if __name__ == "__main__":
+    import argparse
+    import os
+
+    from omegaconf import OmegaConf
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config_yaml",
+        type=str,
+        default="examples/SimplerEnv/train_files/starvla_cotrain_oxe.yaml",
+        help="Path to YAML config",
+    )
+    args, clipargs = parser.parse_known_args()
+
+    if os.getenv("DEBUGPY_ENABLE", "0") == "1":
+        import debugpy
+        debugpy.listen(("0.0.0.0", 10092))
+        print("Rank 0 waiting for debugger attach on port 10092...")
+        debugpy.wait_for_client()
+
+    cfg = OmegaConf.load(args.config_yaml)
+
+    cfg.framework.qwenvl.base_vlm = "./playground/Pretrained_models/Qwen3-VL-4B-Instruct"
+    qwen_vl = _QWen3_VL_Interface(cfg)
+    pass
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/QwenOFT.py b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/QwenOFT.py
new file mode 100644
index 0000000000000000000000000000000000000000..e5e8dd55670bc550dd5f5652c89af2a44e970ce0
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/QwenOFT.py
@@ -0,0 +1,1244 @@
+# Copyright 2025 starVLA community. All rights reserved.
+# Licensed under the MIT License, Version 1.0 (the "License");
+# Implemented by [Jinhui YE / HKUST University] in [2025].
+
+"""
+Qwen-OFT Framework
+
+A lightweight implementation that uses an action special token to parallelly predict continuous actions
+conditioned on multi-view images plus a language instruction (shares parameters with the VLM).
+Inspired by OpenVLA-OFT
+Key Points:
+  - Qwen2.5 vision-language backbone
+  - Injects an action special token into the VLM
+  - Continuous action prediction via L1 regression over the action special token hidden states
+
+
+Note: How to add special tokens to Qwen2.5:
+  download our model checkpoint with special tokens added: https://huggingface.co/StarVLA/Qwen2.5-VL-3B-Instruct-Action
+  or /starVLA/model/modules/vlm/tools/add_qwen_special_tokens/README.md （adpat a little code)
+
+"""
+
+from dataclasses import dataclass, field
+import time
+from typing import Any, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+
+from deployment.model_server.tools.image_tools import to_pil_preserve
+from starVLA.model.tools import FRAMEWORK_REGISTRY
+from starVLA.training.trainer_utils import initialize_overwatch
+
+logger = initialize_overwatch(__name__)
+
+# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels)
+IGNORE_INDEX = -100
+
+from starVLA.model.framework.base_framework import baseframework
+from starVLA.model.framework.share_tools import add_discretized_state_to_instruction, merge_framework_config
+from starVLA.model.modules.action_model.MLP_ActionHeader import get_action_model
+from starVLA.model.modules.vlm import get_vlm_model
+from starVLA.training.trainer_utils.trainer_tools import resize_images
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Default Config for QwenOFT
+#  - Documents every framework-level parameter with type + description
+#  - YAML values override these defaults; extra YAML keys are preserved
+# ──────────────────────────────────────────────────────────────────────
+@dataclass
+class QwenOFTDefaultConfig:
+    """QwenOFT framework default parameters.
+
+    All fields can be overridden by the corresponding key in the YAML
+    ``framework:`` section.  Extra YAML keys not listed here are kept
+    as-is (Config-as-API flexibility).
+    """
+
+    # --- Registry identifier (must match @FRAMEWORK_REGISTRY.register) ---
+    name: str = "QwenOFT"
+
+    # === VLM backbone (Qwen2.5-VL / Qwen3-VL) ===
+    qwenvl: dict = field(
+        default_factory=lambda: {
+            # Path to base VLM checkpoint (local or HF hub id)
+            "base_vlm": "./playground/Pretrained_models/Qwen3-VL-4B-Instruct-Action",
+            # Attention implementation: "flash_attention_2" | "eager" | "sdpa"
+            "attn_implementation": "flash_attention_2",
+        }
+    )
+
+    # === Action head (MLP regression over action special tokens) ===
+    action_model: dict = field(
+        default_factory=lambda: {
+            # Action head architecture type
+            "action_model_type": "MLP",
+            # Dimensionality of each action vector (e.g., 7 for 6-DoF + gripper)
+            "action_dim": 7,
+            # Proprioceptive state dimension for continuous_projector mode
+            "state_dim": 7,
+            # Hidden dim for the action MLP (auto-set from VLM hidden_size at runtime)
+            "action_hidden_dim": 2560,
+            # How many future steps to predict
+            "future_action_window_size": 8,
+            # How many past steps included in action chunk (usually 0)
+            "past_action_window_size": 0,
+            # Loss over predicted actions. Supported: "l1", "discrete_ce",
+            # "multibinary_bce", "asterix_factorized_ce".
+            # "multibinary_ce" is accepted as a CLI alias for "multibinary_bce".
+            "loss_type": "l1",
+            # State conditioning: preserve the existing text-bin path by default.
+            "state_encoding": "discretized_text",
+        }
+    )
+
+
+@FRAMEWORK_REGISTRY.register("QwenOFT")
+class Qwenvl_OFT(baseframework):
+    """
+    Multimodal vision-language-action model (OFT variant).
+
+    Components:
+      - Qwen2.5-VL / Qwen3-VL backbone for fused language/vision token embeddings
+      - Action special token injected into the VLM sequence
+      - MLP regression head over action token hidden states (L1 loss)
+
+    Focus: Predict future continuous actions conditioned on images + instruction.
+    """
+
+    def __init__(
+        self,
+        config: Optional[dict] = None,
+        **kwargs,
+    ) -> None:
+        """
+        Construct all submodules and cache key configuration values.
+
+        Args:
+            config: Hierarchical configuration (OmegaConf/dict) containing framework + trainer sections.
+            **kwargs: Reserved for future overrides (unused).
+        """
+        super().__init__()
+        # Merge framework defaults with YAML config (YAML wins on conflicts)
+        self.config = merge_framework_config(QwenOFTDefaultConfig, config)
+        self.qwen_vl_interface = get_vlm_model(config=self.config)
+        # align action_hidden_dim to VLM hidden_size at runtime
+        self.config.framework.action_model.action_hidden_dim = self.qwen_vl_interface.model.config.hidden_size
+        self.action_model = get_action_model(config=self.config)
+
+        # `action_horizon` is the single source of truth for chunk length.
+        # Legacy aliases (`future_action_window_size`, `past_action_window_size`)
+        # are normalised upstream by `share_tools.apply_config_compat`, so we
+        # only ever read `action_horizon` here.
+        self.action_horizon = int(self.config.framework.action_model.action_horizon)
+        self.chunk_len = self.action_horizon
+        self.action_dim = int(self.config.framework.action_model.action_dim)
+        self.action_env_dim = int(getattr(self.config.framework.action_model, "action_env_dim", self.action_dim))
+        self.action_loss_type = str(getattr(self.config.framework.action_model, "loss_type", "l1")).lower()
+        self.state_encoding = str(self.config.framework.action_model.state_encoding)
+        if self.state_encoding not in {"discretized_text", "continuous_projector"}:
+            raise ValueError(f"Unsupported state_encoding: {self.state_encoding}")
+        if self.state_encoding == "continuous_projector":
+            self.action_model.state_projector = nn.Linear(
+                int(self.config.framework.action_model.state_dim),
+                int(self.config.framework.action_model.action_hidden_dim),
+            )
+        cross_task_cfg = getattr(getattr(self.config, "rl_games", None), "cross_task", None)
+        self.loss_by_task = self._to_plain_dict(getattr(cross_task_cfg, "loss_by_task", None))
+        self.loss_weight_by_task = self._to_plain_dict(getattr(cross_task_cfg, "loss_weight_by_task", None))
+        # self.hidden_dim = config.framework.action_model.action_hidden_dim
+
+        self.action_token = "🔍"  # TODO also can add spacail token to Qwen, but too complex
+        self.action_token_id = self.qwen_vl_interface.processor.tokenizer("🔍", add_special_tokens=False)["input_ids"][0]
+
+        # --- Fixed-size KV memory; default OFF so existing paths are untouched. ---
+        kv_cfg = getattr(self.config.framework, "kv_memory", None)
+        self.kv_memory_enabled = bool(getattr(kv_cfg, "enabled", False)) if kv_cfg is not None else False
+        self.kv_window = int(getattr(kv_cfg, "window", 4)) if kv_cfg is not None else 4
+        self.kv_rollout_len = int(getattr(kv_cfg, "rollout_len", max(self.kv_window + 1, 2))) if kv_cfg is not None else self.kv_window + 1
+        # Packed single-forward training path (StreamingLLM-style sliding-window+sink
+        # mask via FlexAttention) — replaces the per-frame Python loop. Eval/live still
+        # uses the streaming KV memory (`_predict_action_memory`). OFF by default so the
+        # legacy loop stays the reference; enable per experiment via
+        # `framework.kv_memory.packed_train: true`.
+        self.kv_train_packed = bool(getattr(kv_cfg, "packed_train", False)) if kv_cfg is not None else False
+        # Per-step re-based sink for the packed path: replicate the (short) text
+        # prefix once per step and re-base its M-RoPE positions so the sink→frame
+        # distance matches the eval rolling-KV layout exactly (StreamingLLM
+        # "positions within the cache"), instead of one global sink whose distance
+        # grows with the step index. ON by default — the growing-distance sink was
+        # the train/eval forward mismatch behind the packed run's eval collapse.
+        # Set false to fall back to the single global sink for A/B.
+        self.kv_rebased_sink = bool(getattr(kv_cfg, "rebased_sink", True)) if kv_cfg is not None else True
+        if self.state_encoding == "continuous_projector" and self.kv_memory_enabled:
+            raise ValueError("continuous_projector requires framework.kv_memory.enabled=false")
+        # Per-stream (eval slot) memories; rebuilt per sample during training.
+        self._kv_memories: dict = {}
+
+        # L1 loss
+        self.l1_loss = nn.L1Loss()
+
+    def _profile_timing_enabled(self) -> bool:
+        return self.config.trainer.profile_timing.enabled
+
+    def _profile_sync(self) -> None:
+        if self._profile_timing_enabled() and torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    def _profile_start(self) -> float:
+        self._profile_sync()
+        return time.perf_counter()
+
+    def _profile_elapsed(self, start: float) -> float:
+        self._profile_sync()
+        return time.perf_counter() - start
+
+    @staticmethod
+    def _to_plain_dict(value: Any) -> dict[str, Any]:
+        if value is None or value == "":
+            return {}
+        if isinstance(value, dict):
+            return {str(key): item for key, item in value.items()}
+        if hasattr(value, "items"):
+            return {str(key): item for key, item in value.items()}
+        return {}
+
+    def _prepare_action_array(self, action: Any) -> np.ndarray:
+        values = np.asarray(action)
+        if values.ndim == 1:
+            values = values[None, :]
+        if values.shape[-1] == self.action_dim:
+            return values
+        if values.shape[-1] > self.action_dim:
+            return values[..., : self.action_dim]
+        pad_width = [(0, 0)] * values.ndim
+        pad_width[-1] = (0, self.action_dim - values.shape[-1])
+        return np.pad(values, pad_width, mode="constant")
+
+    def _compute_action_loss_for_type(
+        self,
+        pred_actions: torch.Tensor,
+        actions_target: torch.Tensor,
+        *,
+        loss_type: str,
+        action_env_dims: Optional[List[int]] = None,
+    ) -> torch.Tensor:
+        loss_type = str(loss_type).lower()
+        effective_dim = min(self.action_env_dim, pred_actions.shape[-1], actions_target.shape[-1])
+        if effective_dim <= 0:
+            raise ValueError(
+                f"Invalid action_env_dim={self.action_env_dim} for predicted shape={pred_actions.shape} "
+                f"and target shape={actions_target.shape}"
+            )
+
+        if loss_type in {"discrete_ce", "ce", "cross_entropy"}:
+            if action_env_dims is not None:
+                dims = torch.as_tensor(action_env_dims, device=pred_actions.device, dtype=torch.long)
+                if dims.numel() != pred_actions.shape[0]:
+                    raise ValueError(f"Expected {pred_actions.shape[0]} action_env_dims, got {dims.numel()}")
+                losses = []
+                weights = []
+                for dim_value in torch.unique(dims).tolist():
+                    dim = int(min(dim_value, pred_actions.shape[-1], actions_target.shape[-1]))
+                    if dim < 2:
+                        raise ValueError(f"discrete_ce requires at least 2 classes, got action_env_dim={dim}")
+                    mask = dims == int(dim_value)
+                    logits = pred_actions[mask, ..., :dim]
+                    target_class = actions_target[mask, ..., :dim].argmax(dim=-1).long()
+                    losses.append(F.cross_entropy(logits.reshape(-1, dim), target_class.reshape(-1), reduction="sum"))
+                    weights.append(target_class.numel())
+                total_weight = max(1, int(sum(weights)))
+                return sum(losses) / total_weight
+
+            if effective_dim < 2:
+                raise ValueError(
+                    f"action_model.loss_type={loss_type!r} requires at least 2 action classes, "
+                    f"got effective_dim={effective_dim}"
+                )
+            logits = pred_actions[..., :effective_dim]
+            target_class = actions_target[..., :effective_dim].argmax(dim=-1).long()
+            return F.cross_entropy(
+                logits.reshape(-1, effective_dim),
+                target_class.reshape(-1),
+            )
+
+        if loss_type in {"asterix_factorized_ce", "factorized_6_ce", "factorized_ce"}:
+            if effective_dim < 6:
+                raise ValueError(
+                    f"action_model.loss_type={loss_type!r} requires at least 6 action logits, "
+                    f"got effective_dim={effective_dim}"
+                )
+            vertical_logits = pred_actions[..., 0:3]
+            vertical_target = actions_target[..., 0:3].argmax(dim=-1).long()
+            horizontal_logits = pred_actions[..., 3:6]
+            horizontal_target = actions_target[..., 3:6].argmax(dim=-1).long()
+            vertical_loss = F.cross_entropy(
+                vertical_logits.reshape(-1, 3),
+                vertical_target.reshape(-1),
+            )
+            horizontal_loss = F.cross_entropy(
+                horizontal_logits.reshape(-1, 3),
+                horizontal_target.reshape(-1),
+            )
+            return 0.5 * (vertical_loss + horizontal_loss)
+
+        if loss_type in {"multibinary_bce", "multibinary_ce", "bce", "binary_cross_entropy"}:
+            logits = pred_actions[..., :effective_dim]
+            targets = (actions_target[..., :effective_dim] > 0).to(dtype=logits.dtype)
+            return F.binary_cross_entropy_with_logits(logits, targets)
+
+        if loss_type in {"mse", "l2"}:
+            return F.mse_loss(
+                pred_actions[..., :effective_dim],
+                actions_target[..., :effective_dim],
+            )
+
+        if loss_type not in {"l1", "mae"}:
+            raise ValueError(
+                f"Unsupported action_model.loss_type={loss_type!r}; "
+                "expected one of: l1, discrete_ce, multibinary_bce, asterix_factorized_ce"
+            )
+
+        return self.l1_loss(
+            pred_actions[..., :effective_dim],
+            actions_target[..., :effective_dim],
+        )
+
+    def _compute_action_loss(
+        self,
+        pred_actions: torch.Tensor,
+        actions_target: torch.Tensor,
+        action_env_dims: Optional[List[int]] = None,
+        rl_games_tasks: Optional[List[str]] = None,
+        action_loss_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if action_loss_mask is not None:
+            error = (pred_actions[..., :self.action_env_dim] - actions_target[..., :self.action_env_dim]).abs().mean(dim=-1)
+            error = error.masked_fill(~action_loss_mask, 0)
+            return (error.sum(dim=-1) / action_loss_mask.sum(dim=-1)).mean()
+        if self.loss_by_task and rl_games_tasks is not None:
+            if len(rl_games_tasks) != pred_actions.shape[0]:
+                raise ValueError(f"Expected {pred_actions.shape[0]} task labels, got {len(rl_games_tasks)}")
+            losses = []
+            weights = []
+            for task_name in sorted(set(str(task) for task in rl_games_tasks)):
+                mask_values = [str(task) == task_name for task in rl_games_tasks]
+                mask = torch.as_tensor(mask_values, device=pred_actions.device, dtype=torch.bool)
+                if not bool(mask.any()):
+                    continue
+                task_loss_type = str(self.loss_by_task.get(task_name, self.action_loss_type)).lower()
+                task_dims = (
+                    [int(dim) for dim, keep in zip(action_env_dims, mask_values) if keep]
+                    if action_env_dims is not None
+                    else None
+                )
+                task_loss = self._compute_action_loss_for_type(
+                    pred_actions[mask],
+                    actions_target[mask],
+                    loss_type=task_loss_type,
+                    action_env_dims=task_dims,
+                )
+                task_weight = float(self.loss_weight_by_task.get(task_name, 1.0))
+                sample_count = int(mask.sum().item())
+                losses.append(task_loss * task_weight * sample_count)
+                weights.append(task_weight * sample_count)
+            if losses:
+                return sum(losses) / max(1.0, float(sum(weights)))
+
+        return self._compute_action_loss_for_type(
+            pred_actions,
+            actions_target,
+            loss_type=self.action_loss_type,
+            action_env_dims=action_env_dims,
+        )
+
+    def _forward_qwen_last_hidden(self, qwen_inputs: dict) -> torch.Tensor:
+        """
+        Return only the final VLM hidden state when the interface supports it.
+
+        Requesting ``output_hidden_states=True`` from Hugging Face materializes
+        every layer's activations. For OFT we only need the last layer, so the
+        Qwen3 interface exposes a leaner path that preserves the practical memory
+        savings from gradient checkpointing.
+        """
+
+        if hasattr(self.qwen_vl_interface, "forward_last_hidden"):
+            return self.qwen_vl_interface.forward_last_hidden(**qwen_inputs)
+
+        qwenvl_outputs = self.qwen_vl_interface(
+            **qwen_inputs,
+            output_attentions=False,
+            output_hidden_states=True,
+            return_dict=True,
+            use_cache=False,
+        )
+        return qwenvl_outputs.hidden_states[-1]
+
+    def _state_conditioned_instructions(self, instructions, state):
+        if state is not None and self.state_encoding == "discretized_text":
+            return self.add_discretized_state_to_instruction(instructions, state)
+        return instructions
+
+    def _condition_action_queries(self, action_queries, state):
+        if self.state_encoding == "discretized_text":
+            return action_queries
+        state_tensor = torch.as_tensor(
+            np.stack(state), device=action_queries.device, dtype=action_queries.dtype
+        ).squeeze(1)
+        return action_queries + self.action_model.state_projector(state_tensor).unsqueeze(1)
+
+    def forward(
+        self,
+        examples: List[dict] = None,
+        **kwargs,
+    ) -> Tuple:
+        """
+        Training forward: directly regress future actions (no diffusion).
+
+        Flow:
+          1. Build QwenVL inputs (images + instruction tokens)
+          2. Extract hidden states from configured layer range
+          7. Predict action and compute L1 loss
+
+        Args:
+            examples: List[dict], each dict requires:
+                - image: List[PIL.Image] (multi-view)
+                - lang: str instruction
+                - action: np.ndarray or list shaped [T, action_dim]
+            **kwargs: Reserved.
+
+        Returns:
+            dict:
+                action_loss (torch.Tensor): Scalar diffusion noise prediction loss.
+        """
+        if self.kv_memory_enabled:
+            if self.kv_train_packed:
+                return self._forward_memory_packed(examples)
+            return self._forward_memory(examples)
+
+        batch_images = [example["image"] for example in examples]  #  [B，[PLT]]
+        instructions = [example["lang"] for example in examples]  # [B, str]
+        actions = [self._prepare_action_array(example["action"]) for example in examples]  # label [B， len, action_dim]
+        action_env_dims = [int(example.get("action_env_dim", self.action_env_dim)) for example in examples]
+        rl_games_tasks = [str(example.get("rl_games_task", "")) for example in examples]
+        if not any(rl_games_tasks):
+            rl_games_tasks = None
+        state = (
+            [example["state"] for example in examples] if "state" in examples[0] else None
+        )  # List[ndarray (1, state_dim)] or None
+
+        instructions = self._state_conditioned_instructions(instructions, state)
+
+        # step 0: add special action token to instruction
+        action_tokens = (
+            self.action_token * self.chunk_len
+        )  # can't add " " between two tokens, otherwise will be tokenized to multiple tokens
+        prompt_suffix = f" Please predict the next {self.chunk_len} robot actions: <action>{action_tokens}<action>."
+        instructions = [instruction + prompt_suffix for instruction in instructions]
+
+        profile_timing = self._profile_timing_enabled()
+        timing_metrics = {}
+        batch_stats = {}
+
+        # Step 1: QWenVL input format
+        t_inputs = self._profile_start() if profile_timing else None
+        qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions)
+        if profile_timing:
+            timing_metrics["timing/qwen_input_build_total"] = self._profile_elapsed(t_inputs)
+            timing_metrics.update(self.qwen_vl_interface._last_build_timing)
+            input_ids = qwen_inputs["input_ids"]
+            attention_mask = qwen_inputs["attention_mask"]
+            effective_tokens = attention_mask.sum()
+            effective_token_count = effective_tokens.detach().float().item()
+            batch_stats["batch/size"] = float(len(examples))
+            batch_stats["batch/input_len_max"] = float(input_ids.shape[1])
+            batch_stats["batch/input_len_mean"] = float(effective_token_count / input_ids.shape[0])
+            batch_stats["batch/padding_ratio"] = float(1.0 - effective_token_count / float(attention_mask.numel()))
+            batch_stats["batch/effective_tokens"] = float(effective_token_count)
+            batch_stats["batch/image_count"] = float(sum(len(images) for images in batch_images))
+            if "pixel_values" in qwen_inputs:
+                batch_stats["batch/pixel_values_rows"] = float(qwen_inputs["pixel_values"].shape[0])
+
+        t_vlm = self._profile_start() if profile_timing else None
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            last_hidden = self._forward_qwen_last_hidden(qwen_inputs)  # [B, L, H]
+        if profile_timing:
+            timing_metrics["timing/vlm_forward"] = self._profile_elapsed(t_vlm)
+
+        # Step 4: Action Expert Forward and Loss
+        t_action = self._profile_start() if profile_timing else None
+        with torch.autocast("cuda", dtype=torch.float32):
+            # Extract action token embeddings as action prediction queries
+            input_ids = qwen_inputs["input_ids"]
+            action_queries = self._gather_action_token_embeddings(
+                last_hidden, input_ids, action_token_id=self.action_token_id
+            )  # [B, chunk_len, H]
+            action_queries = self._condition_action_queries(action_queries, state)
+            pred_actions = self.action_model.predict_action(action_queries)  # (B, chunk_len, action_dim)
+
+            # Label alignment: take the last chunk_len segment
+            actions = torch.tensor(
+                np.array(actions), device=pred_actions.device, dtype=pred_actions.dtype
+            )  # [B, T_full, action_dim]
+            actions_target = actions[:, -self.action_horizon :, :]  # (B, action_horizon, action_dim)
+
+            action_loss_mask = None
+            if "action_loss_mask" in examples[0]:
+                action_loss_mask = torch.as_tensor(
+                    np.asarray([example["action_loss_mask"] for example in examples]),
+                    device=pred_actions.device,
+                    dtype=torch.bool,
+                )
+
+            action_loss = self._compute_action_loss(
+                pred_actions,
+                actions_target,
+                action_env_dims=action_env_dims,
+                rl_games_tasks=rl_games_tasks,
+                action_loss_mask=action_loss_mask,
+            )
+        if profile_timing:
+            timing_metrics["timing/action_head_loss"] = self._profile_elapsed(t_action)
+
+        return {"action_loss": action_loss, "loss_weight": float(len(examples)), "timing": timing_metrics, "batch_stats": batch_stats}
+
+    @torch.inference_mode()
+    def predict_action(
+        self,
+        examples: List[dict] = None,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """
+
+        Steps:
+          1. Resize images to training resolution (if specified)
+          2. Encode with QwenVL (hidden states retained)
+          6. Return normalized action trajectory
+
+        Returns:
+            dict:
+                normalized_actions (np.ndarray): Shape [B, T, action_dim], diffusion-sampled normalized actions.
+        """
+        if type(examples) is not list:
+            examples = [examples]
+        if self.kv_memory_enabled:
+            return self._predict_action_memory(examples)
+        batch_images = [to_pil_preserve(example["image"]) for example in examples]  #  [B，[PLT]]
+        instructions = [example["lang"] for example in examples]  # [B, str]
+        state = (
+            [example["state"] for example in examples] if "state" in examples[0] else None
+        )  # List[ndarray (1, state_dim)] or None
+
+        instructions = self._state_conditioned_instructions(instructions, state)
+
+        train_obs_image_size = getattr(self.config.datasets.vla_data, "obs_image_size", None)
+        if train_obs_image_size:
+            batch_images = resize_images(batch_images, target_size=train_obs_image_size)
+
+        # step 0: add special action token to instruction
+        action_tokens = (
+            self.action_token * self.chunk_len
+        )  # can't add " " between two tokens, otherwise will be tokenized to multiple tokens
+        prompt_suffix = f" Please predict the next {self.chunk_len} robot actions: <action>{action_tokens}<action>."
+        instructions = [instruction + prompt_suffix for instruction in instructions]
+
+        # Step 1: QWenVL input format
+        qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            last_hidden = self._forward_qwen_last_hidden(qwen_inputs)  # [B, L, H]
+
+        # Step 4: Action Expert Forward and Loss
+        with torch.autocast("cuda", dtype=torch.float32):
+            # Extract action token embeddings as action prediction queries
+            input_ids = qwen_inputs["input_ids"]
+            action_queries = self._gather_action_token_embeddings(
+                last_hidden, input_ids, action_token_id=self.action_token_id
+            )  # [B, chunk_len, H]
+            action_queries = self._condition_action_queries(action_queries, state)
+            pred_actions = self.action_model.predict_action(action_queries)  # (B, chunk_len, action_dim)
+
+        normalized_actions = pred_actions.detach().cpu().numpy()
+        return {"normalized_actions": normalized_actions}
+
+    # ──────────────────────────────────────────────────────────────────
+    #  Fixed-size KV memory. Default OFF. The cache/rotation core is
+    #  unit-proven (tests/test_kv_memory_equivalence.py); the processor-driven
+    #  segment construction below still needs end-to-end validation against the
+    #  real Qwen3-VL-4B checkpoint on GPU.
+    # ──────────────────────────────────────────────────────────────────
+    def reset_memory(self, slot_id=None) -> None:
+        """Clear the KV memory for one eval slot (episode boundary), or all."""
+        if slot_id is None:
+            self._kv_memories.clear()
+        else:
+            self._kv_memories.pop(slot_id, None)
+
+    def _kv_components(self):
+        hf = self.qwen_vl_interface.model  # Qwen3VLForConditionalGeneration
+        backbone = hf.model  # Qwen3VLModel
+        text_model = backbone.language_model  # Qwen3VLTextModel
+        return backbone, text_model.rotary_emb, len(text_model.layers), text_model.config
+
+    def _new_kv_memory(self):
+        from starVLA.model.modules.vlm.kv_memory import FrameKVMemory
+
+        _, rotary_emb, num_layers, text_config = self._kv_components()
+        return FrameKVMemory(rotary_emb, self.kv_window, num_layers, text_config)
+
+    def _build_text_first_inputs(self, frames, instruction):
+        """Tokenize a text-first window: [instruction][frame imgs...][action suffix].
+
+        Returns the processor batch (input_ids, pixel_values, image_grid_thw, ...).
+        Static text comes first so its KV stays positionally stable across steps.
+        """
+        action_tokens = self.action_token * self.chunk_len
+        action_suffix = f" Please predict the next {self.chunk_len} robot actions: <action>{action_tokens}<action>."
+        content = [{"type": "text", "text": instruction}]
+        content += [{"type": "image", "image": img} for img in frames]
+        content += [{"type": "text", "text": action_suffix}]
+        messages = [[{"role": "user", "content": content}]]
+        batch = self.qwen_vl_interface.processor.apply_chat_template(
+            messages, tokenize=True, padding=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
+        )
+        return batch.to(self.qwen_vl_interface.model.device)
+
+    def _kv_segments(self, input_ids_1d):
+        """Partition a text-first window's ids into (prefix_len, [frame_lens], action_len).
+
+        A frame block spans its vision_start .. vision_end (inclusive); the prefix is
+        everything before the first vision_start; the action suffix is everything
+        after the last vision_end. Using the real special-token ids keeps eviction
+        leak-free (each evicted frame removes exactly its own markers).
+        """
+        cfg = self.qwen_vl_interface.model.config
+        vstart = int(getattr(cfg, "vision_start_token_id", 151652))
+        vend = int(getattr(cfg, "vision_end_token_id", 151653))
+        ids = input_ids_1d.tolist()
+        starts = [i for i, t in enumerate(ids) if t == vstart]
+        ends = [i for i, t in enumerate(ids) if t == vend]
+        if not starts or not ends:
+            raise RuntimeError("KV memory: window inputs have no vision_start/end tokens")
+        action_start = ends[-1] + 1
+        bounds = starts + [action_start]
+        frame_lens = [bounds[i + 1] - bounds[i] for i in range(len(starts))]
+        return starts[0], frame_lens, len(ids) - action_start
+
+    def _kv_forward_step_batched(self, memories, frames, instruction):
+        """Run ONE streaming KV-memory step for a group of same-layout streams and
+        return ``(last_hidden_new [B, new_len, H], new_input_ids [B, new_len])``.
+
+        ``memories`` and ``frames`` are aligned lists (one entry per stream / eval
+        slot); all entries share ``instruction`` and history layout (grouped by the
+        caller). Two efficiencies vs a naive per-frame, per-step encode:
+
+        - Single-frame preprocessing: the image processor runs once per stream on its
+          OWN current frame. The window's multi-frame token layout (needed only for
+          M-RoPE positions via ``get_rope_index``) is reconstructed by replicating the
+          frame's token block, so duplicate frames are never re-preprocessed.
+        - Cross-stream batching: the streams' canonical KV blocks are stacked along the
+          batch dim, so the whole group is one forward; results are sliced back to each
+          stream's memory afterwards.
+
+        Past frames come from each stream's memory; only the current frame's pixels are
+        encoded fresh. ``B == 1`` (training rollout) reuses the same path with no
+        stacking, so gradients flow through the in-place memory updates.
+        """
+        from starVLA.model.modules.vlm.kv_memory import FrameKVMemory, memory_step
+
+        backbone = self.qwen_vl_interface.model.model  # Qwen3VLModel
+        B = len(memories)
+        n_past = memories[0].num_past_frames()
+        n_visible = n_past + 1
+
+        # One processor call per stream's current frame (no duplicate-frame work).
+        per_stream = [self._build_text_first_inputs([frame], instruction) for frame in frames]
+        ids1 = per_stream[0]["input_ids"][0]
+        cur_grid = per_stream[0]["image_grid_thw"][:1]
+        prefix_len, frame_lens, action_len = self._kv_segments(ids1)
+        frame_len = frame_lens[0]
+
+        # Reconstruct the n_visible-frame window ids by replicating the frame block,
+        # then let get_rope_index assign the (content-independent) M-RoPE positions.
+        prefix, action = ids1[:prefix_len], ids1[prefix_len + frame_len :]
+        block = ids1[prefix_len : prefix_len + frame_len]
+        window_ids = torch.cat([prefix] + [block] * n_visible + [action]).unsqueeze(0)
+        full_pos1, _ = backbone.get_rope_index(window_ids, image_grid_thw=cur_grid.repeat(n_visible, 1))
+        full_pos = full_pos1.expand(3, B, full_pos1.shape[-1]).contiguous()
+
+        batched = FrameKVMemory.stack(memories) if B > 1 else memories[0]
+        past_len = batched.past_length()
+        is_prefill = not batched.has_text()
+        expected_past = 0 if is_prefill else prefix_len + n_past * frame_len
+        if past_len != expected_past:
+            raise RuntimeError(
+                f"KV memory slice misaligned: memory past_length={past_len} but window "
+                f"expects {expected_past}. Token boundaries drifted across steps."
+            )
+        if is_prefill:
+            new_start, frame_start_in_new, prefill_text_len = 0, prefix_len, prefix_len
+        else:
+            new_start, frame_start_in_new, prefill_text_len = past_len, 0, 0
+
+        new_ids = window_ids[:, new_start:].expand(B, -1).contiguous()
+        pixels = torch.cat([b["pixel_values"] for b in per_stream], dim=0)
+        grid = cur_grid.repeat(B, 1)
+
+        last_hidden = memory_step(
+            backbone,
+            batched,
+            new_input_ids=new_ids,
+            pixel_values=pixels,
+            image_grid_thw=grid,
+            full_position_ids=full_pos,
+            frame_start_in_new=frame_start_in_new,
+            frame_len=frame_len,
+            prefill_text_len=prefill_text_len,
+        )
+
+        if B > 1:
+            for i, memory in enumerate(memories):
+                sliced = batched.slice(i)
+                memory.text, memory.frames = sliced.text, sliced.frames
+        return last_hidden, new_ids
+
+    # ------------------------------------------------------------------ #
+    # Packed single-forward training path (StreamingLLM-style mask).
+    #
+    # Instead of the per-frame Python loop (`_forward_memory`), pack the whole
+    # rollout into ONE sequence
+    #     [text/sink] [frame_1][action_1] ... [frame_R][action_R]
+    # and run a SINGLE Qwen3-VL forward under a FlexAttention BlockMask that
+    # reproduces the streaming visibility:
+    #   * the text prefix is an always-visible attention sink (StreamingLLM);
+    #   * each frame attends only the last `kv_window` frames (sliding window);
+    #   * each action read-out block is isolated (seen only by its own step), so
+    #     the frame KV stream is bit-identical to the streaming loop / eval.
+    # Custom M-RoPE `position_ids` keep frame->frame and action->frame relative
+    # distances exactly matching eval; only the text-sink->frame absolute distance
+    # grows with step (accepted approximation — the sink is distance-robust).
+    # Eval/live (`_predict_action_memory`) is untouched and keeps the rolling KV.
+    # ------------------------------------------------------------------ #
+    def _packed_window_layout(self, frames, instruction):
+        """Build the packed [text][f1][a1]...[fR][aR] window for ONE stream.
+
+        Token structure is content-independent (image placeholders are identical
+        per frame), so the prefix / frame / action blocks are reconstructed by
+        token-replication (matching `_kv_forward_step_batched`'s `window_ids`),
+        which guarantees each action block equals the per-step layout exactly.
+
+        Returns dict with: input_ids [1,L], pixel_values, image_grid_thw [R,3],
+        position_ids [3,1,L] (custom M-RoPE), kind [L] (0=text,1=frame,2=action),
+        step [L] (frame/step index for frame&action tokens, -1 for text),
+        action_idx [R, chunk_len] (packed indices of each step's action tokens).
+        """
+        R = len(frames)
+        backbone = self.qwen_vl_interface.model.model  # Qwen3VLModel
+        device = self.qwen_vl_interface.model.device
+
+        # One processor call per frame: identical token block, distinct pixels.
+        per_frame = [self._build_text_first_inputs([f], instruction) for f in frames]
+        ids1 = per_frame[0]["input_ids"][0]
+        grid1 = per_frame[0]["image_grid_thw"][:1]  # [1,3]
+        prefix_len, frame_lens, action_len = self._kv_segments(ids1)
+        frame_len = frame_lens[0]
+        for pf in per_frame:
+            pl, fl, al = self._kv_segments(pf["input_ids"][0])
+            if (pl, fl[0], al) != (prefix_len, frame_len, action_len):
+                raise RuntimeError(
+                    "Packed KV memory: frames have inconsistent token layout "
+                    f"({(pl, fl[0], al)} vs {(prefix_len, frame_len, action_len)}); "
+                    "all frames in a rollout must share resolution/grid."
+                )
+
+        prefix = ids1[:prefix_len]
+        frame_block = ids1[prefix_len : prefix_len + frame_len]
+        action_block = ids1[prefix_len + frame_len :]  # incl. assistant header
+        block_len = frame_len + action_len
+
+        pixel_values = torch.cat([pf["pixel_values"] for pf in per_frame], dim=0)
+        grid = grid1.repeat(R, 1)  # [R, 3]
+
+        # Frames-only M-RoPE positions: frames laid out consecutively (no action
+        # in between), so frame→frame / action→frame relative distances match eval
+        # ("positions within the cache"). These are the canonical frame positions
+        # reused regardless of where the frame sits in the packed token sequence.
+        frames_only_ids = torch.cat([prefix] + [frame_block] * R).unsqueeze(0)
+        full_pos_f, _ = backbone.get_rope_index(frames_only_ids, image_grid_thw=grid)  # [3,1,Lf]
+        text_pos = full_pos_f[:, :, :prefix_len]
+
+        action_offsets = (action_block == self.action_token_id).nonzero(as_tuple=False).flatten()
+        if action_offsets.numel() < self.chunk_len:
+            raise RuntimeError(
+                f"Packed KV memory: action block has {action_offsets.numel()} action "
+                f"tokens (< chunk_len={self.chunk_len})."
+            )
+        action_offsets = action_offsets[-self.chunk_len :]  # last chunk_len, matches gather
+
+        def _frame_pos(k):
+            s = prefix_len + k * frame_len
+            return full_pos_f[:, :, s : s + frame_len]
+
+        def _action_pos(frame_pos_k):
+            fk_max = int(frame_pos_k.max().item())
+            ap = torch.arange(1, action_len + 1, device=device) + fk_max
+            return ap.view(1, 1, -1).expand(3, 1, -1)
+
+        W = self.kv_window
+        parts, pos_parts, kind_parts, step_parts, action_idx_rows = [], [], [], [], []
+
+        if self.kv_rebased_sink:
+            # Per-step re-based sink, with the redundant pre-eviction sinks deduped.
+            # Steps k < W all have j0=0 → o_k=0 → an identical sink at [0,P), so they
+            # share ONE front sink (step=-2, visible to queries with step < W). This
+            # is exactly eval before any eviction (sink pinned at the front, frames
+            # growing from P). Steps k >= W each carry their own re-based sink
+            # ([sink_k][frame_k][action_k]) at o_k so the sink→frame distance equals
+            # eval's rolling-KV layout (bit-exact). Net: 1 + max(0,R-W) sinks instead
+            # of R, saving (W-1) prefix copies while staying numerically identical.
+            base0 = int(full_pos_f[0, 0, prefix_len].item())  # frame 0's temporal start
+
+            # Shared front sink for the pre-eviction steps (step=-2).
+            parts.append(prefix)
+            pos_parts.append(text_pos)
+            kind_parts.append(torch.zeros(prefix_len, dtype=torch.long, device=device))
+            step_parts.append(torch.full((prefix_len,), -2, dtype=torch.long, device=device))
+            cursor = prefix_len
+
+            for k in range(R):
+                if k >= W:
+                    # Re-based sink_k sits right before step k's oldest visible frame.
+                    j0 = k - W + 1
+                    o_k = int(full_pos_f[0, 0, prefix_len + j0 * frame_len].item()) - base0
+                    parts.append(prefix)
+                    pos_parts.append(text_pos + o_k)  # shift all 3 M-RoPE dims by o_k
+                    kind_parts.append(torch.zeros(prefix_len, dtype=torch.long, device=device))
+                    step_parts.append(torch.full((prefix_len,), k, dtype=torch.long, device=device))
+                    cursor += prefix_len
+                frame_pos_k = _frame_pos(k)
+                parts += [frame_block, action_block]
+                pos_parts += [frame_pos_k, _action_pos(frame_pos_k)]
+                kind_parts.append(
+                    torch.cat([
+                        torch.ones(frame_len, dtype=torch.long, device=device),
+                        torch.full((action_len,), 2, dtype=torch.long, device=device),
+                    ])
+                )
+                step_parts.append(torch.full((block_len,), k, dtype=torch.long, device=device))
+                action_idx_rows.append(cursor + frame_len + action_offsets)
+                cursor += block_len
+        else:
+            # Legacy single global sink: [sink][f0 a0]...[f_{R-1} a_{R-1}]. Sink at
+            # the front (step=-1, always visible); sink→frame distance grows with k.
+            parts.append(prefix)
+            pos_parts.append(text_pos)
+            kind_parts.append(torch.zeros(prefix_len, dtype=torch.long, device=device))
+            step_parts.append(torch.full((prefix_len,), -1, dtype=torch.long, device=device))
+            cursor = prefix_len
+            for k in range(R):
+                frame_pos_k = _frame_pos(k)
+                parts += [frame_block, action_block]
+                pos_parts += [frame_pos_k, _action_pos(frame_pos_k)]
+                kind_parts.append(
+                    torch.cat([
+                        torch.ones(frame_len, dtype=torch.long, device=device),
+                        torch.full((action_len,), 2, dtype=torch.long, device=device),
+                    ])
+                )
+                step_parts.append(torch.full((block_len,), k, dtype=torch.long, device=device))
+                action_idx_rows.append(cursor + frame_len + action_offsets)
+                cursor += block_len
+
+        input_ids = torch.cat(parts).unsqueeze(0)  # [1, L]
+        position_ids = torch.cat(pos_parts, dim=2)  # [3,1,L]
+        kind = torch.cat(kind_parts)
+        step = torch.cat(step_parts)
+        action_idx = torch.stack(action_idx_rows, dim=0)  # [R, chunk_len]
+
+        return {
+            "input_ids": input_ids,
+            "pixel_values": pixel_values,
+            "image_grid_thw": grid,
+            "position_ids": position_ids,
+            "kind": kind,
+            "step": step,
+            "action_idx": action_idx,
+        }
+
+    def _packed_block_mask(self, kind, step, L, batch):
+        """Build a FlexAttention BlockMask for the packed sliding-window+sink layout."""
+        from torch.nn.attention.flex_attention import create_block_mask
+        from starVLA.model.modules.vlm.kv_memory import make_packed_sliding_mask_mod
+
+        mask_mod = make_packed_sliding_mask_mod(kind, step, self.kv_window)
+        # B=H=None → the mask is shared across batch and heads (it depends only on
+        # token kind/step, not content), so one BlockMask covers the whole batch.
+        return create_block_mask(mask_mod, B=None, H=None, Q_LEN=L, KV_LEN=L, device=str(kind.device))
+
+    def _assert_flex_attention(self):
+        """Packed training needs the decoder dispatched through FlexAttention.
+
+        We pass a FlexAttention BlockMask straight to the model, so the decoder
+        (forward AND gradient-checkpointing recompute during backward) must use
+        the flex backend the whole time — a transient toggle around the forward
+        would leave the recompute on sdpa and choke on the BlockMask. The eval
+        streaming path also runs fine under flex, so the whole model is simply
+        loaded with `attn_implementation: flex_attention`.
+        """
+        _, _, _, text_config = self._kv_components()
+        impl = getattr(text_config, "_attn_implementation", None)
+        if impl != "flex_attention":
+            raise RuntimeError(
+                "Packed KV-memory training requires the VLM loaded with "
+                f"attn_implementation='flex_attention' (got {impl!r}). Set "
+                "framework.qwenvl.attn_implementation: flex_attention."
+            )
+
+    def _forward_memory_packed(self, examples: List[dict]) -> dict:
+        """Training forward via a single packed sliding-window+sink masked pass.
+
+        Per-frame supervision identical to `_forward_memory` (every real frame is
+        supervised at the window it naturally has), but the expensive VLM forward
+        runs ONCE per rollout instead of R times. Full BPTT across the rollout
+        (no per-step detach) — pair with gradient checkpointing for memory.
+        """
+        self._assert_flex_attention()
+        profile = self._profile_timing_enabled()
+        timing = {"timing/mem_nograd_fwd": 0.0, "timing/mem_grad_fwd": 0.0, "timing/mem_action_head": 0.0}
+
+        instructions = [
+            self._instruction_with_state(ex["lang"], ex.get("state") if "state" in ex else None)
+            for ex in examples
+        ]
+        groups: dict = {}
+        for idx, instruction in enumerate(instructions):
+            valid = np.asarray(examples[idx]["valid"], dtype=bool)
+            groups.setdefault((instruction, tuple(valid.tolist())), []).append(idx)
+
+        total_weighted_loss = None
+        total_weight = 0.0
+        for (instruction, valid_key), indices in groups.items():
+            valid = np.asarray(valid_key, dtype=bool)
+            real_positions = [k for k in range(len(valid)) if valid[k]]
+            R = len(real_positions)
+            if R == 0:
+                continue
+            action_env_dims = [int(examples[i].get("action_env_dim", self.action_env_dim)) for i in indices]
+            rl_games_tasks = [str(examples[i].get("rl_games_task", "")) for i in indices]
+            if not any(rl_games_tasks):
+                rl_games_tasks = None
+
+            # Build per-sample packed layouts (identical token structure, distinct pixels).
+            layouts = [
+                self._packed_window_layout([examples[i]["image"][k] for k in real_positions], instruction)
+                for i in indices
+            ]
+            base = layouts[0]
+            L = base["input_ids"].shape[1]
+            B = len(indices)
+            input_ids = base["input_ids"].expand(B, -1).contiguous()
+            position_ids = base["position_ids"].expand(3, B, -1).contiguous()
+            pixel_values = torch.cat([lay["pixel_values"] for lay in layouts], dim=0)
+            grid = torch.cat([lay["image_grid_thw"] for lay in layouts], dim=0)  # [B*R, 3]
+            action_idx = base["action_idx"]  # [R, chunk_len]
+            block_mask = self._packed_block_mask(base["kind"], base["step"], L, B)
+
+            t_fwd = self._profile_start() if profile else None
+            last_hidden = self._forward_qwen_last_hidden(
+                {
+                    "input_ids": input_ids,
+                    "position_ids": position_ids,
+                    "attention_mask": block_mask,
+                    "pixel_values": pixel_values,
+                    "image_grid_thw": grid,
+                    "use_cache": False,
+                }
+            )  # [B, L, H]
+            if profile:
+                timing["timing/mem_grad_fwd"] += self._profile_elapsed(t_fwd)
+
+            t_head = self._profile_start() if profile else None
+            with torch.autocast("cuda", dtype=torch.float32):
+                for j, k in enumerate(real_positions):
+                    action_queries = last_hidden[:, action_idx[j], :]  # [B, chunk_len, H]
+                    group_preds = self.action_model.predict_action(action_queries)  # [B, chunk, dim]
+                    targets = [self._prepare_action_array(examples[i]["actions_per_frame"][k]) for i in indices]
+                    targets = torch.tensor(np.array(targets), device=group_preds.device, dtype=group_preds.dtype)
+                    targets = targets[:, -self.action_horizon :, :]
+                    step_loss = self._compute_action_loss(
+                        group_preds, targets, action_env_dims=action_env_dims, rl_games_tasks=rl_games_tasks
+                    )
+                    weight = float(sum(float(np.asarray(examples[i]["density_weight"])[k]) for i in indices))
+                    contribution = step_loss * weight
+                    total_weighted_loss = contribution if total_weighted_loss is None else total_weighted_loss + contribution
+                    total_weight += weight
+            if profile:
+                timing["timing/mem_action_head"] += self._profile_elapsed(t_head)
+
+        action_loss = total_weighted_loss / max(total_weight, 1e-6)
+        return {"action_loss": action_loss, "loss_weight": total_weight, "timing": timing, "batch_stats": {}}
+
+    def _instruction_with_state(self, instruction, state):
+        if state is None:
+            return instruction
+        return self.add_discretized_state_to_instruction([instruction], [state])[0]
+
+    def _forward_memory(self, examples: List[dict]) -> dict:
+        """Training forward with the streaming KV memory: per-frame supervision.
+
+        Each sample is an R-frame rollout `[t-R+1 .. t]`; the dataloader marks clamp
+        padding via ``valid`` and supplies one action chunk per frame in
+        ``actions_per_frame`` ([R, H, dim]). We skip padding frames (never feed memory,
+        never supervise) and supervise EVERY real frame at the window it naturally has
+        (1 -> window, then evict), which matches eval (`_predict_action_memory`) frame
+        for frame, including the unfilled window at an episode's start. History KV is
+        detached between steps (truncated BPTT depth 1), so each step's loss flows only
+        through that step's forward; the per-step weighted losses are summed and the
+        trainer backpropagates once.
+
+        Same-instruction, same-`valid` rollouts share a token layout and window
+        evolution, so they run each step in one batched forward via
+        `_kv_forward_step_batched` (degrades to per-sample otherwise), mirroring
+        `_predict_action_memory`.
+        """
+        profile = self._profile_timing_enabled()
+        timing = {"timing/mem_nograd_fwd": 0.0, "timing/mem_grad_fwd": 0.0, "timing/mem_action_head": 0.0}
+
+        instructions = [
+            self._instruction_with_state(ex["lang"], ex.get("state") if "state" in ex else None)
+            for ex in examples
+        ]
+
+        # Group rollouts that share the prompt AND the same real/padding pattern: those
+        # march through identical window states each step, so they batch cleanly.
+        groups: dict = {}
+        for idx, instruction in enumerate(instructions):
+            valid = np.asarray(examples[idx]["valid"], dtype=bool)
+            groups.setdefault((instruction, tuple(valid.tolist())), []).append(idx)
+
+        total_weighted_loss = None
+        total_weight = 0.0
+        for (instruction, valid_key), indices in groups.items():
+            valid = np.asarray(valid_key, dtype=bool)
+            real_positions = [k for k in range(len(valid)) if valid[k]]
+            memories = [self._new_kv_memory() for _ in indices]
+            rollout = [examples[i]["image"] for i in indices]
+            action_env_dims = [int(examples[i].get("action_env_dim", self.action_env_dim)) for i in indices]
+            rl_games_tasks = [str(examples[i].get("rl_games_task", "")) for i in indices]
+            if not any(rl_games_tasks):
+                rl_games_tasks = None
+
+            for step_pos, k in enumerate(real_positions):
+                if step_pos > 0:
+                    for memory in memories:
+                        memory.detach_()  # truncated BPTT: history is a fixed memory
+                step_frames = [frames[k] for frames in rollout]
+
+                t_fwd = self._profile_start() if profile else None
+                last_hidden, new_ids = self._kv_forward_step_batched(memories, step_frames, instruction)
+                if profile:
+                    timing["timing/mem_grad_fwd"] += self._profile_elapsed(t_fwd)
+
+                t_head = self._profile_start() if profile else None
+                with torch.autocast("cuda", dtype=torch.float32):
+                    action_queries = self._gather_action_token_embeddings(
+                        last_hidden, new_ids, action_token_id=self.action_token_id
+                    )
+                    group_preds = self.action_model.predict_action(action_queries)  # [Bg, chunk, dim]
+                    targets = [
+                        self._prepare_action_array(examples[i]["actions_per_frame"][k]) for i in indices
+                    ]
+                    targets = torch.tensor(
+                        np.array(targets), device=group_preds.device, dtype=group_preds.dtype
+                    )  # [Bg, H, dim]
+                    targets = targets[:, -self.action_horizon :, :]
+                    step_loss = self._compute_action_loss(
+                        group_preds,
+                        targets,
+                        action_env_dims=action_env_dims,
+                        rl_games_tasks=rl_games_tasks,
+                    )
+                if profile:
+                    timing["timing/mem_action_head"] += self._profile_elapsed(t_head)
+
+                # Per-supervised-step density weight: sum of the group's per-frame
+                # weights at frame k. Uniform weights reduce this to a sample-count
+                # weighting, i.e. the exact mean over all supervised (sample, step).
+                weight = float(
+                    sum(float(np.asarray(examples[i]["density_weight"])[k]) for i in indices)
+                )
+                contribution = step_loss * weight
+                total_weighted_loss = contribution if total_weighted_loss is None else total_weighted_loss + contribution
+                total_weight += weight
+
+        action_loss = total_weighted_loss / max(total_weight, 1e-6)
+        return {"action_loss": action_loss, "loss_weight": total_weight, "timing": timing, "batch_stats": {}}
+
+    @torch.inference_mode()
+    def _predict_action_memory(self, examples: List[dict]) -> dict:
+        """Eval-time streaming prediction with per-slot KV memory (req3, milestone 4).
+
+        The newest frame (``example['image'][-1]``) is encoded fresh; past frames are
+        served from the slot's memory. ``reset_memory(slot_id)`` (called from the
+        policy at episode boundaries) clears a slot. Streams that currently share the
+        same instruction and history layout are processed in one batched forward; the
+        rest fall back to their own group (no left padding needed).
+        """
+        train_obs_image_size = getattr(self.config.datasets.vla_data, "obs_image_size", None)
+
+        # Resolve each example to its (slot memory, current frame, instruction).
+        entries = []
+        for idx, example in enumerate(examples):
+            slot_id = example.get("slot_id", idx)
+            frames = [to_pil_preserve(f) for f in example["image"]]
+            if train_obs_image_size:
+                frames = resize_images(frames, target_size=train_obs_image_size)
+            state = example.get("state") if "state" in example else None
+            instruction = self._instruction_with_state(example["lang"], state)
+            memory = self._kv_memories.get(slot_id)
+            if memory is None:
+                memory = self._new_kv_memory()
+                self._kv_memories[slot_id] = memory
+            entries.append({"idx": idx, "frame": frames[-1], "instruction": instruction, "memory": memory})
+
+        # Group same-layout streams (same prompt + same cached layout) into one forward.
+        groups: dict = {}
+        for entry in entries:
+            key = (entry["instruction"], entry["memory"].layout_key())
+            groups.setdefault(key, []).append(entry)
+
+        preds: List[Optional[torch.Tensor]] = [None] * len(examples)
+        for group in groups.values():
+            instruction = group[0]["instruction"]
+            memories = [entry["memory"] for entry in group]
+            frames = [entry["frame"] for entry in group]
+            last_hidden, new_ids = self._kv_forward_step_batched(memories, frames, instruction)
+            with torch.autocast("cuda", dtype=torch.float32):
+                action_queries = self._gather_action_token_embeddings(
+                    last_hidden, new_ids, action_token_id=self.action_token_id
+                )
+                pred_actions = self.action_model.predict_action(action_queries)
+            for i, entry in enumerate(group):
+                preds[entry["idx"]] = pred_actions[i : i + 1]
+
+        normalized = torch.cat(preds, dim=0).detach().cpu().numpy()
+        return {"normalized_actions": normalized}
+
+    def _gather_action_token_embeddings(
+        self,
+        last_hidden: torch.Tensor,  # [B, L, H]
+        input_ids: torch.Tensor,  # [B, L]
+        action_token_id=None,  # Can be int or List[int]
+    ) -> torch.Tensor:
+        """
+        Vectorized batch extraction of action token embeddings:
+          - No per-sample for loop
+          - Select the last chunk_len action placeholder tokens from each sample
+        Args:
+            last_hidden: [B, L, H]
+            input_ids:   [B, L]
+            action_token_id: int or List[int]
+        Returns:
+            action_queries: [B, chunk_len, H]
+        """
+        if action_token_id is None:
+            raise ValueError("action_token_id must not be None")
+
+        device = input_ids.device
+        B, L, H = last_hidden.shape
+
+        # Support multiple ids (e.g., multiple variants)
+        if isinstance(action_token_id, (list, tuple, set)):
+            id_list = torch.tensor(list(action_token_id), device=device, dtype=input_ids.dtype)
+            # torch.isin requires PyTorch >=1.10
+            mask = torch.isin(input_ids, id_list)
+        else:
+            mask = input_ids == action_token_id  # [B, L]
+
+        counts = mask.sum(dim=1)  # [B]
+        if (counts < self.chunk_len).any():
+            insufficient = (counts < self.chunk_len).nonzero(as_tuple=False).flatten().tolist()
+            raise RuntimeError(
+                f"The following samples have insufficient action tokens (< {self.chunk_len}): {insufficient} |"
+                f" counts={counts.tolist()}"
+            )
+
+        # Position indices
+        idx = torch.arange(L, device=device).unsqueeze(0).expand(B, L)  # [B, L]
+        masked_pos = torch.where(mask, idx, torch.full_like(idx, -1))  # Set non-action positions to -1
+
+        # Take the last chunk_len positions (higher indices = later in sequence)
+        # Note: count sufficiency already verified, so -1 won't be incorrectly selected
+        topk_pos = masked_pos.topk(k=self.chunk_len, dim=-1).values  # [B, chunk_len] unsorted
+        # Sort in temporal order
+        selected_pos = topk_pos.sort(dim=-1).values  # [B, chunk_len]
+
+        # Gather
+        expanded_index = selected_pos.unsqueeze(-1).expand(-1, -1, H)  # [B, chunk_len, H]
+        action_queries = last_hidden.gather(dim=1, index=expanded_index)  # [B, chunk_len, H]
+        return action_queries
+
+    # Discretised state → instruction prefix (π₀.5 style); shared with QwenPI_v3.
+    add_discretized_state_to_instruction = staticmethod(add_discretized_state_to_instruction)
+
+
+if __name__ == "__main__":
+    import argparse
+    import os
+
+    from omegaconf import OmegaConf
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config_yaml",
+        type=str,
+        default="examples/LIBERO/train_files/starvla_cotrain_libero.yaml",
+        help="Path to YAML config",
+    )
+    args, clipargs = parser.parse_known_args()
+
+    if os.getenv("DEBUGPY_ENABLE", "0") == "1":
+        import debugpy
+
+        debugpy.listen(("0.0.0.0", 10092))
+        print("Rank 0 waiting for debugger attach on port 10092...")
+        debugpy.wait_for_client()
+
+    cfg = OmegaConf.load(args.config_yaml)
+
+    model = Qwenvl_OFT(cfg)
+    print(model)
+
+    image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
+    sample = {
+        "action": np.random.uniform(-1, 1, size=(16, 7)).astype(np.float16),
+        "image": [image],
+        "lang": "This is a fake instruction for testing.",
+        "state": np.random.uniform(-1, 1, size=(1, 7)).astype(np.float16),  # chunk, state_dim
+    }
+    sample2 = sample.copy()
+    sample2["lang"] = "Another fake instruction for testing."
+
+    batch = [sample, sample2]
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    forward_output = model(batch)
+    action_loss = forward_output["action_loss"]
+    print(f"[train] Action Loss (with state): {action_loss.item()}")
+
+    predict_output = model.predict_action(examples=[batch[0]])
+    normalized_actions = predict_output["normalized_actions"]
+    print(f"[infer] Predicted Action shape: {normalized_actions.shape}")
+
+    # Backward-compat: examples without `state` should still work.
+    sample_no_state = {k: v for k, v in sample.items() if k != "state"}
+    forward_no_state = model([sample_no_state, sample_no_state])
+    print(f"[train] Action Loss (no state): {forward_no_state['action_loss'].item()}")
+    predict_no_state = model.predict_action(examples=[sample_no_state])
+    print(f"[infer] Predicted Action shape (no state): {predict_no_state['normalized_actions'].shape}")
+
+    print("Finished")
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/datasets.py b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/datasets.py
new file mode 100644
index 0000000000000000000000000000000000000000..2346230737b152846b9c28904284fbf39219e201
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/datasets.py
@@ -0,0 +1,3429 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""
+In this file, we define 3 types of datasets:
+1. LeRobotSingleDataset: a single dataset for a given embodiment tag
+2. LeRobotMixtureDataset: a mixture of datasets for a given list of embodiment tags
+3. CachedLeRobotSingleDataset: a single dataset for a given embodiment tag,
+                                with caching for the video frames
+
+See `scripts/load_dataset.py` for examples on how to use these datasets.
+"""
+import os
+import hashlib
+import io
+import json, torch
+import copy
+from starVLA.training.trainer_utils.trainer_tools import stitch_frames
+from collections import defaultdict
+from pathlib import Path
+from typing import Sequence
+import os, random
+import numpy as np
+import pandas as pd
+from pydantic import BaseModel, Field, ValidationError
+from torch.utils.data import Dataset
+from tqdm import tqdm
+from PIL import Image
+import torch.distributed as dist
+
+from starVLA.dataloader.gr00t_lerobot.video import get_all_frames, get_frames_by_timestamps
+
+from starVLA.dataloader.gr00t_lerobot.embodiment_tags import EmbodimentTag
+from starVLA.dataloader.gr00t_lerobot.schema import (
+    DatasetMetadata,
+    DatasetStatisticalValues,
+    LeRobotModalityMetadata,
+    LeRobotStateActionMetadata,
+)
+from starVLA.dataloader.gr00t_lerobot.transform import ComposedModalityTransform
+
+from functools import partial
+from typing import Tuple, List
+import pickle
+import gc
+
+# LeRobot v2.0 dataset file names 
+LE_ROBOT_MODALITY_FILENAME = "meta/modality.json"
+LE_ROBOT_EPISODE_FILENAME = "meta/episodes.jsonl"
+LE_ROBOT_TASKS_FILENAME = "meta/tasks.jsonl"
+LE_ROBOT_INFO_FILENAME = "meta/info.json"
+LE_ROBOT_STATS_FILENAME = "meta/stats_gr00t.json"
+LE_ROBOT_DATA_FILENAME = "data/*/*.parquet"
+LE_ROBOT_STEPS_FILENAME = "meta/steps.pkl"
+LE_ROBOT_STATS_FORMAT_VERSION = 2
+EPSILON = 5e-4
+LATENCY_PROMPT_MARKER = " Current action latency is "
+
+
+def _as_bool(value, default: bool = False) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        return value.strip().lower() not in {"false", "0", "no", "off"}
+    return bool(value)
+
+
+def _apply_prompt_mode(prompt: str, prompt_mode: str | None) -> str:
+    mode = str(prompt_mode or "").strip().lower()
+    if mode in {"", "default", "none", "raw"}:
+        return prompt
+    if mode == "latency_neutral":
+        head, marker, _tail = str(prompt).partition(LATENCY_PROMPT_MARKER)
+        return head.rstrip() if marker else str(prompt)
+    raise ValueError(f"Unsupported datasets.vla_data.prompt_mode={prompt_mode!r}")
+
+
+def _concatenate_action_fields(data: dict, keys: list[str]) -> np.ndarray:
+    action = np.concatenate([data[key] for key in keys], axis=-1)
+    return action[0] if action.ndim == 3 else action
+
+
+#  LeRobot v3.0 dataset file names 
+LE_ROBOT3_TASKS_FILENAME = "meta/tasks.parquet"
+LE_ROBOT3_EPISODE_FILENAME = "meta/episodes/*/*.parquet"
+
+
+def calculate_dataset_statistics(parquet_paths: list[Path]) -> dict:
+    """Calculate the dataset statistics of all columns for a list of parquet files."""
+    # Dataset statistics
+    all_low_dim_data_list = []
+    # Collect all the data
+    # parquet_paths = parquet_paths[:3]
+    for parquet_path in tqdm(
+        sorted(list(parquet_paths)),
+        desc="Collecting all parquet files...",
+    ):
+        # Load the parquet file
+        parquet_data = pd.read_parquet(parquet_path)
+        parquet_data = parquet_data
+        all_low_dim_data_list.append(parquet_data)
+    
+    all_low_dim_data = pd.concat(all_low_dim_data_list, axis=0)
+    # Compute dataset statistics
+    dataset_statistics = {}
+    for le_modality in tqdm(all_low_dim_data.columns, desc="Processing modalities"):
+        print(le_modality)
+        if "task_info" in le_modality:
+            continue
+        print(f"Computing statistics for {le_modality}...")
+        try:
+            np_data = np.vstack(
+                [np.asarray(x, dtype=np.float32) for x in all_low_dim_data[le_modality]]
+            )
+        except Exception as e:
+            print(f"Warning: Failed to process modality {le_modality} due to error: {e}")
+            continue  
+
+        dataset_statistics[le_modality] = {
+            "mean": np.mean(np_data, axis=0).tolist(),
+            "std": np.std(np_data, axis=0).tolist(),
+            "min": np.min(np_data, axis=0).tolist(),
+            "max": np.max(np_data, axis=0).tolist(),
+            "q01": np.quantile(np_data, 0.01, axis=0).tolist(),
+            "q99": np.quantile(np_data, 0.99, axis=0).tolist(),
+        }
+    return dataset_statistics
+
+
+def _normalize_action_mode(mode: str) -> str:
+    """Normalize action mode names to {abs, delta, rel}.""" 
+    # @gaoning plz move this, we want dataloader to be independent of the action mode logic, we can move this to transform or a separate utils tool to handle lerobot dataset
+    mode = str(mode).lower()
+    if mode in {"absolute", "raw"}:
+        mode = "abs"
+    if mode not in {"abs", "delta", "rel"}:
+        mode = "abs"
+    return mode
+
+
+def _normalize_action_mode_apply_keys(
+    action_mode_apply_keys: Sequence[str] | None,
+    fallback_keys: Sequence[str] | None = None,
+) -> list[str]:
+    source_keys = action_mode_apply_keys if action_mode_apply_keys else (fallback_keys or [])
+    normalized = []
+    for key in source_keys:
+        key = str(key)
+        if not key.startswith("action."):
+            key = f"action.{key}"
+        normalized.append(key)
+    return normalized
+
+
+def _normalize_action_mode_state_map(action_mode_state_map: dict[str, str] | None) -> dict[str, str]:
+    normalized = {}
+    for action_key, state_key in (action_mode_state_map or {}).items():
+        action_key = str(action_key)
+        state_key = str(state_key)
+        if not action_key.startswith("action."):
+            action_key = f"action.{action_key}"
+        if not state_key.startswith("state."):
+            state_key = f"state.{state_key}"
+        normalized[action_key] = state_key
+    return normalized
+
+
+def _build_stats_cache_config(
+    action_mode: str,
+) -> dict:
+    return {
+        "mode": action_mode,
+    }
+
+
+def _invalidate_legacy_stats_cache(stats_path: Path, reason: str) -> None:
+    if not stats_path.exists():
+        return
+    print(f"Removing stale dataset statistics cache at {stats_path}: {reason}")
+    stats_path.unlink()
+
+
+def _load_stats_cache(
+    stats_path: Path,
+    expected_config: dict,
+    *,
+    invalidate_legacy: bool,
+) -> dict | None:
+    if not stats_path.exists():
+        return None
+
+    try:
+        with open(stats_path, "r") as f:
+            payload = json.load(f)
+    except Exception as exc:
+        if invalidate_legacy:
+            _invalidate_legacy_stats_cache(stats_path, f"failed to load JSON ({exc})")
+        return None
+
+    if not isinstance(payload, dict):
+        if invalidate_legacy:
+            _invalidate_legacy_stats_cache(stats_path, "unexpected top-level format")
+        return None
+
+    format_version = payload.get("__format_version")
+    cache_config = payload.get("__cache_config")
+    statistics = payload.get("statistics")
+    if format_version != LE_ROBOT_STATS_FORMAT_VERSION or cache_config is None or statistics is None:
+        if invalidate_legacy:
+            _invalidate_legacy_stats_cache(stats_path, "legacy statistics format detected")
+        return None
+
+    if cache_config != expected_config:
+        if invalidate_legacy:
+            _invalidate_legacy_stats_cache(stats_path, "statistics config mismatch, rebuilding cache")
+        return None
+
+    return statistics
+
+
+def _save_stats_cache(stats_path: Path, cache_config: dict, statistics: dict) -> None:
+    payload = {
+        "__format_version": LE_ROBOT_STATS_FORMAT_VERSION,
+        "__cache_config": cache_config,
+        "statistics": statistics,
+    }
+    tmp_path = stats_path.with_suffix(".tmp")
+    stats_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(tmp_path, "w") as f:
+        json.dump(payload, f, indent=4)
+    os.replace(tmp_path, stats_path)
+
+
+def _compute_statistics_for_mode(
+    parquet_paths: list[Path],
+    dataset_name: str,
+    action_mode: str,
+    lerobot_modality_meta: "LeRobotModalityMetadata",
+    action_keys_full: list[str],
+    state_keys_full: list[str],
+    action_indices: list[int] | None,
+    state_indices: list[int] | None,
+    action_mode_apply_keys: list[str] | None,
+    action_mode_state_map: dict[str, str] | None,
+) -> dict:
+    print(f"[RANK 0] Calculating dataset statistics for {dataset_name} (mode={action_mode})")
+
+    base_stats = calculate_dataset_statistics(parquet_paths)
+    
+    if action_mode == "abs":
+        return base_stats
+
+    if action_indices is None or state_indices is None:
+        raise ValueError(
+            "Both action and state modalities are required to compute "
+            f"{action_mode} action mode statistics."
+        )
+
+    if action_mode == "delta":
+        return calculate_delta_action_statistics(
+            parquet_paths=parquet_paths,
+            lerobot_modality_meta=lerobot_modality_meta,
+            action_keys_full=action_keys_full,
+            state_keys_full=state_keys_full,
+            action_indices=action_indices,
+            state_indices=state_indices,
+            action_mode_apply_keys=action_mode_apply_keys,
+            action_mode_state_map=action_mode_state_map,
+            base_stats=base_stats,
+        )
+    if action_mode == "rel":
+        return calculate_rel_action_statistics(
+            parquet_paths=parquet_paths,
+            lerobot_modality_meta=lerobot_modality_meta,
+            action_keys_full=action_keys_full,
+            state_keys_full=state_keys_full,
+            action_indices=action_indices,
+            state_indices=state_indices,
+            action_mode_apply_keys=action_mode_apply_keys,
+            action_mode_state_map=action_mode_state_map,
+            base_stats=base_stats,
+        )
+    raise ValueError(f"Unsupported action mode for statistics: {action_mode}")
+
+
+def _load_or_compute_statistics(
+    stats_path: Path,
+    stats_cache_config: dict,
+    parquet_paths: list[Path],
+    dataset_name: str,
+    action_mode: str,
+    lerobot_modality_meta: "LeRobotModalityMetadata",
+    action_keys_full: list[str],
+    state_keys_full: list[str],
+    action_indices: list[int] | None,
+    state_indices: list[int] | None,
+    action_mode_apply_keys: list[str] | None,
+    action_mode_state_map: dict[str, str] | None,
+) -> dict:
+    le_statistics = _load_stats_cache(
+        stats_path,
+        stats_cache_config,
+        invalidate_legacy=True,
+    )
+    if le_statistics is not None:
+        return le_statistics
+
+    le_statistics = _compute_statistics_for_mode(
+        parquet_paths=parquet_paths,
+        dataset_name=dataset_name,
+        action_mode=action_mode,
+        lerobot_modality_meta=lerobot_modality_meta,
+        action_keys_full=action_keys_full,
+        state_keys_full=state_keys_full,
+        action_indices=action_indices,
+        state_indices=state_indices,
+        action_mode_apply_keys=action_mode_apply_keys,
+        action_mode_state_map=action_mode_state_map,
+    )
+    _save_stats_cache(stats_path, stats_cache_config, le_statistics)
+    return le_statistics
+
+
+def _get_action_col_slices(
+    lerobot_modality_meta: "LeRobotModalityMetadata",
+    action_keys_full: list[str],
+    state_keys_full: list[str],
+    action_mode_apply_keys: list[str] | None = None,
+    action_mode_state_map: dict[str, str] | None = None,
+) -> dict[str, list[tuple[tuple[int, int], str, tuple[int, int], str, str]]]:
+    apply_keys = _normalize_action_mode_apply_keys(action_mode_apply_keys, action_keys_full)
+    action_mode_state_map = _normalize_action_mode_state_map(action_mode_state_map)
+
+    action_meta = lerobot_modality_meta.action
+    state_meta = lerobot_modality_meta.state
+
+    # Build per-column mapping: action column -> list of (action_slice, state_column, state_slice)
+    action_col_slices: dict[str, list[tuple[tuple[int, int], str, tuple[int, int]]]] = {}
+    for action_key in apply_keys:
+        if not action_key.startswith("action."):
+            raise ValueError(f"Invalid action key {action_key}. Expected prefix 'action.'.")
+        state_key = action_mode_state_map.get(action_key, action_key.replace("action.", "state.", 1))
+        if state_key not in state_keys_full:
+            raise ValueError(
+                f"State key {state_key} not found for action key {action_key}. "
+                f"Add it to action_mode_state_map or remove {action_key} from action_mode_apply_keys."
+            )
+
+        action_subkey = action_key.replace("action.", "", 1)
+        state_subkey = state_key.replace("state.", "", 1)
+        if action_subkey not in action_meta or state_subkey not in state_meta:
+            raise ValueError(f"Action/state key missing in metadata: {action_key} -> {state_key}")
+
+        action_cfg = action_meta[action_subkey]
+        state_cfg = state_meta[state_subkey]
+        action_col = action_cfg.original_key or action_subkey
+        state_col = state_cfg.original_key or state_subkey
+        action_slice = (action_cfg.start, action_cfg.end)
+        state_slice = (state_cfg.start, state_cfg.end)
+        action_padding = "first_last" if action_cfg.absolute else "zero"
+        state_padding = "first_last" if state_cfg.absolute else "zero"
+        action_col_slices.setdefault(action_col, []).append(
+            (action_slice, state_col, state_slice, action_padding, state_padding)
+        )
+
+    return action_col_slices
+
+
+def calculate_delta_action_statistics(
+    parquet_paths: list[Path],
+    lerobot_modality_meta: "LeRobotModalityMetadata",
+    action_keys_full: list[str],
+    state_keys_full: list[str],
+    action_indices: list[int],
+    state_indices: list[int],
+    action_mode_apply_keys: list[str] | None = None,
+    action_mode_state_map: dict[str, str] | None = None,
+    base_stats: dict | None = None,
+) -> dict:
+    """
+    Calculate action statistics using delta mode.
+
+    Rule:
+      - For t>0: a_t - a_{t-1}
+      - For t=0: a_0 - s_0
+
+    Mapping rule (only two cases):
+      1) Use explicit action_mode_state_map if provided.
+      2) Otherwise, replace 'action.' with 'state.' directly.
+    """
+    if base_stats is None:
+        base_stats = calculate_dataset_statistics(parquet_paths)
+
+    action_col_slices = _get_action_col_slices(
+        lerobot_modality_meta, action_keys_full, state_keys_full, action_mode_apply_keys, action_mode_state_map
+    )
+    if not action_col_slices:
+        raise ValueError("No action columns found in the dataset.")
+
+    def _get_chunk(array: np.ndarray, step_indices: np.ndarray, padding_strategy: str) -> np.ndarray:
+        max_length = array.shape[0]
+        front_padding = step_indices < 0
+        end_padding = step_indices >= max_length
+        padding_positions = np.logical_or(front_padding, end_padding)
+        output = np.zeros((len(step_indices), array.shape[1]), dtype=array.dtype)
+        if (~padding_positions).any():
+            output[~padding_positions] = array[step_indices[~padding_positions]]
+        if padding_positions.any():
+            if padding_strategy == "first_last":
+                output[front_padding] = array[0]
+                output[end_padding] = array[-1]
+            elif padding_strategy == "zero":
+                output[padding_positions] = 0
+            else:
+                raise ValueError(f"Invalid padding strategy: {padding_strategy}")
+        return output
+
+    accum: dict[str, list[np.ndarray]] = {col: [] for col in action_col_slices.keys()}
+    for parquet_path in tqdm(sorted(list(parquet_paths)), desc="Collecting delta action stats"):
+        data = pd.read_parquet(parquet_path)
+        trajectory_length = len(data)
+        for action_col, slice_list in action_col_slices.items():
+            if action_col not in data.columns:
+                raise ValueError(f"{action_col} not found in parquet columns.")
+            action_matrix = np.stack(data[action_col])
+            action_padding_ref = slice_list[0][3]
+            prepared_slices = []
+            for a_slice, state_col, s_slice, action_padding, state_padding in slice_list:
+                if state_col not in data.columns:
+                    raise ValueError(f"{state_col} not found in parquet columns.")
+                state_matrix = np.stack(data[state_col])
+                state_part_full = state_matrix[:, s_slice[0] : s_slice[1]]
+                prepared_slices.append((a_slice, state_part_full, state_padding))
+            for base_index in range(trajectory_length):
+                action_steps = np.array(action_indices) + base_index
+                action_chunk_full = _get_chunk(action_matrix, action_steps, action_padding_ref)
+
+                for a_slice, state_part_full, state_padding in prepared_slices:
+                    action_part_chunk = action_chunk_full[:, a_slice[0] : a_slice[1]]
+                    state_chunk = _get_chunk(state_part_full, np.array(state_indices) + base_index, state_padding)
+                    if action_part_chunk.shape[1] != state_chunk.shape[1]:
+                        raise ValueError(f"Action/state dim mismatch for {action_col}:{a_slice}")
+
+                    out = action_part_chunk.copy()
+                    if len(out) > 1:
+                        out[1:] = action_part_chunk[1:] - action_part_chunk[:-1]
+                    out[0] = action_part_chunk[0] - state_chunk[0]
+                    action_chunk_full[:, a_slice[0] : a_slice[1]] = out
+
+                accum[action_col].append(action_chunk_full)
+
+    delta_stats = copy.deepcopy(base_stats)
+    for action_col, series_list in accum.items():
+        if not series_list:
+            continue
+        all_values = np.concatenate(series_list, axis=0).astype(np.float32)
+        delta_stats[action_col] = {
+            "mean": np.mean(all_values, axis=0).tolist(),
+            "std": np.std(all_values, axis=0).tolist(),
+            "min": np.min(all_values, axis=0).tolist(),
+            "max": np.max(all_values, axis=0).tolist(),
+            "q01": np.quantile(all_values, 0.01, axis=0).tolist(),
+            "q99": np.quantile(all_values, 0.99, axis=0).tolist(),
+        }
+    return delta_stats
+
+
+def calculate_rel_action_statistics(
+    parquet_paths: list[Path],
+    lerobot_modality_meta: "LeRobotModalityMetadata",
+    action_keys_full: list[str],
+    state_keys_full: list[str],
+    action_indices: list[int],
+    state_indices: list[int],
+    action_mode_apply_keys: list[str] | None = None,
+    action_mode_state_map: dict[str, str] | None = None,
+    base_stats: dict | None = None,
+) -> dict:
+    """
+    Calculate action statistics using rel mode.
+
+    Rule:
+      - For all t: a_t - s_0
+
+    Mapping rule (only two cases):
+      1) Use explicit action_mode_state_map if provided.
+      2) Otherwise, replace 'action.' with 'state.' directly.
+    """
+    if base_stats is None:
+        base_stats = calculate_dataset_statistics(parquet_paths)
+
+    action_col_slices = _get_action_col_slices(
+        lerobot_modality_meta, action_keys_full, state_keys_full, action_mode_apply_keys, action_mode_state_map
+    )
+    if not action_col_slices:
+        raise ValueError("No action columns found in the dataset.")
+
+    def _get_chunk(array: np.ndarray, step_indices: np.ndarray, padding_strategy: str) -> np.ndarray:
+        max_length = array.shape[0]
+        front_padding = step_indices < 0
+        end_padding = step_indices >= max_length
+        padding_positions = np.logical_or(front_padding, end_padding)
+        output = np.zeros((len(step_indices), array.shape[1]), dtype=array.dtype)
+        if (~padding_positions).any():
+            output[~padding_positions] = array[step_indices[~padding_positions]]
+        if padding_positions.any():
+            if padding_strategy == "first_last":
+                output[front_padding] = array[0]
+                output[end_padding] = array[-1]
+            elif padding_strategy == "zero":
+                output[padding_positions] = 0
+            else:
+                raise ValueError(f"Invalid padding strategy: {padding_strategy}")
+        return output
+
+    accum: dict[str, list[np.ndarray]] = {col: [] for col in action_col_slices.keys()}
+    for parquet_path in tqdm(sorted(list(parquet_paths)), desc="Collecting rel action stats"):
+        data = pd.read_parquet(parquet_path)
+        trajectory_length = len(data)
+        for action_col, slice_list in action_col_slices.items():
+            if action_col not in data.columns:
+                raise ValueError(f"{action_col} not found in parquet columns.")
+            action_matrix = np.stack(data[action_col])
+            action_padding_ref = slice_list[0][3]
+            prepared_slices = []
+            for a_slice, state_col, s_slice, action_padding, state_padding in slice_list:
+                if state_col not in data.columns:
+                    raise ValueError(f"{state_col} not found in parquet columns.")
+                state_matrix = np.stack(data[state_col])
+                state_part_full = state_matrix[:, s_slice[0] : s_slice[1]]
+                prepared_slices.append((a_slice, state_part_full, state_padding))
+            for base_index in range(trajectory_length):
+                action_steps = np.array(action_indices) + base_index
+                action_chunk_full = _get_chunk(action_matrix, action_steps, action_padding_ref)
+
+                for a_slice, state_part_full, state_padding in prepared_slices:
+                    action_part_chunk = action_chunk_full[:, a_slice[0] : a_slice[1]]
+                    state_chunk = _get_chunk(state_part_full, np.array(state_indices) + base_index, state_padding)
+                    if action_part_chunk.shape[1] != state_chunk.shape[1]:
+                        raise ValueError(f"Action/state dim mismatch for {action_col}:{a_slice}")
+
+                    out = action_part_chunk - state_chunk[0]
+                    action_chunk_full[:, a_slice[0] : a_slice[1]] = out
+
+                accum[action_col].append(action_chunk_full)
+
+    rel_stats = copy.deepcopy(base_stats)
+    for action_col, series_list in accum.items():
+        if not series_list:
+            continue
+        all_values = np.concatenate(series_list, axis=0).astype(np.float32)
+        rel_stats[action_col] = {
+            "mean": np.mean(all_values, axis=0).tolist(),
+            "std": np.std(all_values, axis=0).tolist(),
+            "min": np.min(all_values, axis=0).tolist(),
+            "max": np.max(all_values, axis=0).tolist(),
+            "q01": np.quantile(all_values, 0.01, axis=0).tolist(),
+            "q99": np.quantile(all_values, 0.99, axis=0).tolist(),
+        }
+    return rel_stats
+
+class ModalityConfig(BaseModel):
+    """Configuration for a modality."""
+
+    delta_indices: list[int]
+    """Delta indices to sample relative to the current index. The returned data will correspond to the original data at a sampled base index + delta indices."""
+    modality_keys: list[str]
+    """The keys to load for the modality in the dataset."""
+
+
+class LeRobotSingleDataset(Dataset):
+    """
+    Base dataset class for LeRobot that supports sharding.
+    """
+    def __init__(
+        self,
+        dataset_path: Path | str,
+        modality_configs: dict[str, ModalityConfig],
+        embodiment_tag: str | EmbodimentTag,
+        video_backend: str = "decord",
+        video_backend_kwargs: dict | None = None,
+        transforms: ComposedModalityTransform | None = None,
+        delete_pause_frame: bool = False,
+        data_cfg = None,
+        **kwargs,
+    ):
+        """
+        Initialize the dataset.
+
+        Args:
+            dataset_path (Path | str): The path to the dataset.
+            modality_configs (dict[str, ModalityConfig]): The configuration for each modality. The keys are the modality names, and the values are the modality configurations.
+                See `ModalityConfig` for more details.
+            video_backend (str): Backend for video reading.
+            video_backend_kwargs (dict): Keyword arguments for the video backend when initializing the video reader.
+            transforms (ComposedModalityTransform): The transforms to apply to the dataset.
+            embodiment_tag (EmbodimentTag): Overload the embodiment tag for the dataset. e.g. define it as "new_embodiment"
+        """
+        # first check if the path directory exists
+        self.data_cfg = data_cfg
+        if not Path(dataset_path).exists():
+            raise FileNotFoundError(f"Dataset path {dataset_path} does not exist")
+        # indict letobot version
+        self._lerobot_version =  self.data_cfg.get("lerobot_version", "v2.0") #self._indict_lerobot_version(**kwargs)
+
+        self._action_mode = None
+        self._action_mode_state_map = {}
+        self._action_mode_apply_keys = None
+
+        self.delete_pause_frame = delete_pause_frame
+
+        self.modality_configs = modality_configs
+        self.video_backend = video_backend
+        self.video_backend_kwargs = video_backend_kwargs if video_backend_kwargs is not None else {}
+        self.transforms = (
+            transforms if transforms is not None else ComposedModalityTransform(transforms=[])
+        )
+
+        self._dataset_path = Path(dataset_path)
+        self._dataset_name = self._dataset_path.name
+        if isinstance(embodiment_tag, EmbodimentTag):
+            self.tag = embodiment_tag.value
+        else:
+            self.tag = embodiment_tag
+
+        self._init_action_mode()
+        self._metadata = self._get_metadata(EmbodimentTag(self.tag))
+
+        # LeRobot-specific config
+        self._lerobot_modality_meta = self._get_lerobot_modality_meta()
+        self._lerobot_info_meta = self._get_lerobot_info_meta()
+        self._data_path_pattern = self._get_data_path_pattern()
+        self._video_path_pattern = self._get_video_path_pattern()
+        self._chunk_size = self._get_chunk_size()
+        self._tasks = self._get_tasks()
+        # self._episodes = self._get_episode_info() # TODO why we need this func
+        self.curr_traj_data = None
+        self.curr_traj_id = None
+        # ponytail: retain visited immutable tables; use a bounded cache if datasets outgrow RAM.
+        self._parquet_cache: dict[Path, pd.DataFrame] = {}
+
+        self._trajectory_ids, self._trajectory_lengths = self._get_trajectories()
+        self._modality_keys = self._get_modality_keys()
+        self._delta_indices = self._get_delta_indices()
+        self._all_steps = self._get_all_steps()
+        self.set_transforms_metadata(self.metadata)
+        self.set_epoch(0)
+
+        print(f"Initialized dataset {self.dataset_name} with {embodiment_tag}")
+
+
+        # Check if the dataset is valid
+        self._check_integrity()
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["_parquet_cache"] = {}
+        state["curr_traj_data"] = None
+        state["curr_traj_id"] = None
+        return state
+
+    @property
+    def dataset_path(self) -> Path:
+        """The path to the dataset that contains the METADATA_FILENAME file."""
+        return self._dataset_path
+
+    @property
+    def metadata(self) -> DatasetMetadata:
+        """The metadata for the dataset, loaded from metadata.json in the dataset directory"""
+        return self._metadata
+
+    @property
+    def trajectory_ids(self) -> np.ndarray:
+        """The trajectory IDs in the dataset, stored as a 1D numpy array of strings."""
+        return self._trajectory_ids
+
+    @property
+    def trajectory_lengths(self) -> np.ndarray:
+        """The trajectory lengths in the dataset, stored as a 1D numpy array of integers.
+        The order of the lengths is the same as the order of the trajectory IDs.
+        """
+        return self._trajectory_lengths
+
+    @property
+    def all_steps(self) -> list[tuple[int, int]]:
+        """The trajectory IDs and base indices for all steps in the dataset.
+        Example:
+            self.trajectory_ids: [0, 1, 2]
+            self.trajectory_lengths: [3, 2, 4]
+            return: [
+                ("traj_0", 0), ("traj_0", 1), ("traj_0", 2),
+                ("traj_1", 0), ("traj_1", 1),
+                ("traj_2", 0), ("traj_2", 1), ("traj_2", 2), ("traj_2", 3)
+            ]
+        """
+        return self._all_steps
+
+    @property
+    def modality_keys(self) -> dict:
+        """The modality keys for the dataset. The keys are the modality names, and the values are the keys for each modality.
+
+        Example: {
+            "video": ["video.image_side_0", "video.image_side_1"],
+            "state": ["state.eef_position", "state.eef_rotation"],
+            "action": ["action.eef_position", "action.eef_rotation"],
+            "language": ["language.human.task"],
+            "timestamp": ["timestamp"],
+            "reward": ["reward"],
+        }
+        """
+        return self._modality_keys
+
+    @property
+    def delta_indices(self) -> dict[str, np.ndarray]:
+        """The delta indices for the dataset. The keys are the modality.key, and the values are the delta indices for each modality.key."""
+        return self._delta_indices
+
+    @property
+    def dataset_name(self) -> str:
+        """The name of the dataset."""
+        return self._dataset_name
+
+    @property
+    def lerobot_modality_meta(self) -> LeRobotModalityMetadata:
+        """The metadata for the LeRobot dataset."""
+        return self._lerobot_modality_meta
+
+    @property
+    def lerobot_info_meta(self) -> dict:
+        """The metadata for the LeRobot dataset."""
+        return self._lerobot_info_meta
+
+    @property
+    def data_path_pattern(self) -> str:
+        """The path pattern for the LeRobot dataset."""
+        return self._data_path_pattern
+
+    @property
+    def video_path_pattern(self) -> str:
+        """The path pattern for the LeRobot dataset."""
+        return self._video_path_pattern
+
+    @property
+    def chunk_size(self) -> int:
+        """The chunk size for the LeRobot dataset."""
+        return self._chunk_size
+
+    @property
+    def tasks(self) -> pd.DataFrame:
+        """The tasks for the dataset."""
+        return self._tasks
+
+    def _get_metadata(self, embodiment_tag: EmbodimentTag) -> DatasetMetadata:
+        """Get the metadata for the dataset.
+
+        Returns:
+            dict: The metadata for the dataset.
+        """
+
+        # 1. Modality metadata
+        modality_meta_path = self.dataset_path / LE_ROBOT_MODALITY_FILENAME
+        assert (
+            modality_meta_path.exists()
+        ), f"Please provide a {LE_ROBOT_MODALITY_FILENAME} file in {self.dataset_path}"
+        # 1.1. State and action modalities
+        simplified_modality_meta: dict[str, dict] = {}
+        with open(modality_meta_path, "r") as f:
+            le_modality_meta = LeRobotModalityMetadata.model_validate(json.load(f))
+        for modality in ["state", "action"]:
+            simplified_modality_meta[modality] = {}
+            le_state_action_meta: dict[str, LeRobotStateActionMetadata] = getattr(
+                le_modality_meta, modality
+            )
+            for subkey in le_state_action_meta:
+                state_action_dtype = np.dtype(le_state_action_meta[subkey].dtype)
+                if np.issubdtype(state_action_dtype, np.floating):
+                    continuous = True
+                else:
+                    continuous = False
+                simplified_modality_meta[modality][subkey] = {
+                    "absolute": le_state_action_meta[subkey].absolute,
+                    "rotation_type": le_state_action_meta[subkey].rotation_type,
+                    "shape": [
+                        le_state_action_meta[subkey].end - le_state_action_meta[subkey].start
+                    ],
+                    "continuous": continuous,
+                }
+
+        # 1.2. Video modalities
+        le_info_path = self.dataset_path / LE_ROBOT_INFO_FILENAME
+        assert (
+            le_info_path.exists()
+        ), f"Please provide a {LE_ROBOT_INFO_FILENAME} file in {self.dataset_path}"
+        with open(le_info_path, "r") as f:
+            le_info = json.load(f)
+        simplified_modality_meta["video"] = {}
+        for new_key in le_modality_meta.video:
+            original_key = le_modality_meta.video[new_key].original_key
+            if original_key is None:
+                original_key = new_key
+            le_video_meta = le_info["features"][original_key]
+            height = le_video_meta["shape"][le_video_meta["names"].index("height")]
+            width = le_video_meta["shape"][le_video_meta["names"].index("width")]
+            # NOTE(FH): different lerobot dataset versions have different keys for the number of channels and fps
+            try:
+                channels = le_video_meta["shape"][le_video_meta["names"].index("channel")]
+                fps = le_video_meta["video_info"]["video.fps"]
+            except (ValueError, KeyError):
+                try:
+                    channels = le_video_meta["info"]["video.channels"]
+                    fps = le_video_meta["info"]["video.fps"]
+                except (ValueError, KeyError):
+                    # Fallback for image-only datasets (e.g. VLA-Arena) that lack video_info
+                    channels = 3
+                    fps = le_info.get("fps", 30)
+            simplified_modality_meta["video"][new_key] = {
+                "resolution": [width, height],
+                "channels": channels,
+                "fps": fps,
+            }
+
+
+        # 2. Dataset statistics
+        def is_main():
+            return (not dist.is_initialized()) or dist.get_rank() == 0
+        
+        action_mode = _normalize_action_mode(self.data_cfg.get("action_mode", "abs") if self.data_cfg else "abs")
+
+        stats_path = self.dataset_path / LE_ROBOT_STATS_FILENAME
+        action_cfg = self.modality_configs.get("action")
+        state_cfg = self.modality_configs.get("state")
+        action_keys_full = list(action_cfg.modality_keys) if action_cfg else []
+        state_keys_full = list(state_cfg.modality_keys) if state_cfg else []
+        action_indices = list(action_cfg.delta_indices) if action_cfg else None
+        state_indices = list(state_cfg.delta_indices) if state_cfg else None
+
+        apply_keys = _normalize_action_mode_apply_keys(
+            self.data_cfg.get("action_mode_apply_keys", None) if self.data_cfg else None,
+            action_keys_full,
+        )
+        normalized_state_map = _normalize_action_mode_state_map(
+            self.data_cfg.get("action_mode_state_map", {}) if self.data_cfg else {}
+        )
+        stats_cache_config = _build_stats_cache_config(
+            action_mode=action_mode,
+        )
+        parquet_files = list(self.dataset_path.glob(LE_ROBOT_DATA_FILENAME))
+        parquet_files_filtered = [
+            pf for pf in parquet_files if "episode_033675.parquet" not in pf.name
+        ]
+
+        if is_main():
+            le_statistics = _load_or_compute_statistics(
+                stats_path,
+                stats_cache_config=stats_cache_config,
+                parquet_paths=parquet_files_filtered,
+                dataset_name=self.dataset_name,
+                action_mode=action_mode,
+                lerobot_modality_meta=le_modality_meta,
+                action_keys_full=action_keys_full,
+                state_keys_full=state_keys_full,
+                action_indices=action_indices,
+                state_indices=state_indices,
+                action_mode_apply_keys=apply_keys,
+                action_mode_state_map=normalized_state_map,
+            )
+        else:
+            le_statistics = None
+
+        if dist.is_initialized():
+            dist.barrier()
+
+        if le_statistics is None:
+            le_statistics = _load_stats_cache(
+                stats_path,
+                stats_cache_config,
+                invalidate_legacy=False,
+            )
+            if le_statistics is None:
+                raise RuntimeError(
+                    f"Dataset statistics cache is missing or invalid after sync: {stats_path}"
+                )
+
+        manifest_path = self.dataset_path / "manifest.json"
+        if manifest_path.exists():
+            manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+            if manifest.get("state_source") == "decision_v1":
+                le_statistics["observation.state"]["min"] = manifest["state_normalization"]["min"]
+                le_statistics["observation.state"]["max"] = manifest["state_normalization"]["max"]
+
+        for stat in le_statistics.values():
+            DatasetStatisticalValues.model_validate(stat)
+
+
+        dataset_statistics = {}
+        for our_modality in ["state", "action"]:
+            dataset_statistics[our_modality] = {}
+            for subkey in simplified_modality_meta[our_modality]:
+                dataset_statistics[our_modality][subkey] = {}
+                state_action_meta = le_modality_meta.get_key_meta(f"{our_modality}.{subkey}")
+                assert isinstance(state_action_meta, LeRobotStateActionMetadata)
+                le_modality = state_action_meta.original_key
+                for stat_name in le_statistics[le_modality]:
+                    indices = np.arange(
+                        state_action_meta.start,
+                        state_action_meta.end,
+                    )
+                    stat = np.array(le_statistics[le_modality][stat_name])
+                    dataset_statistics[our_modality][subkey][stat_name] = stat[indices].tolist()
+
+        # 3. Full dataset metadata
+        metadata = DatasetMetadata(
+            statistics=dataset_statistics,  # type: ignore
+            modalities=simplified_modality_meta,  # type: ignore
+            embodiment_tag=embodiment_tag,
+        )
+
+        return metadata
+
+    def _get_trajectories(self) -> tuple[np.ndarray, np.ndarray]:
+        """Get the trajectories in the dataset."""
+        # Get trajectory lengths, IDs, and whitelist from dataset metadata
+        # v2.0
+        if self._lerobot_version == "v2.0":
+            file_path = self.dataset_path / LE_ROBOT_EPISODE_FILENAME
+            with open(file_path, "r") as f:
+                episode_metadata = [json.loads(line) for line in f]
+            trajectory_ids = []
+            trajectory_lengths = []
+            for episode in episode_metadata:
+                trajectory_ids.append(episode["episode_index"])
+                trajectory_lengths.append(episode["length"])
+            return np.array(trajectory_ids), np.array(trajectory_lengths)
+        # v3.0
+        elif self._lerobot_version == "v3.0":
+            file_paths = sorted(list((self.dataset_path).glob(LE_ROBOT3_EPISODE_FILENAME)))
+            trajectory_ids = []
+            trajectory_lengths = []
+            # data_chunck_index = []
+            # data_file_index = []
+            # vido_from_index = []
+            self.trajectory_ids_to_metadata = {}
+            for file_path in file_paths:
+                episodes_data = pd.read_parquet(file_path)
+                timestamp_cols = [
+                    c
+                    for c in episodes_data.columns
+                    if str(c).startswith("videos/") and str(c).endswith("/from_timestamp")
+                ]
+                for index, episode in episodes_data.iterrows():
+                    trajectory_ids.append(episode["episode_index"])
+                    trajectory_lengths.append(episode["length"])
+
+                    from_timestamps = {}
+                    for col in timestamp_cols:
+                        value = episode[col]
+                        if pd.isna(value):
+                            continue
+                        # videos/{video_key}/from_timestamp -> {video_key}
+                        video_key = str(col)[len("videos/") : -len("/from_timestamp")]
+                        from_timestamps[video_key] = float(value)
+
+                    # TODO auto map key 
+                    # Collect video file indices for each video key
+                    #已修改的lerobotv3.0的视频索引（提取视频和文件的索引）
+                    video_file_indices = {}
+                    for col in timestamp_cols:
+                        video_key = str(col)[len("videos/") : -len("/from_timestamp")]
+                        chunk_col = f"videos/{video_key}/chunk_index"
+                        file_col = f"videos/{video_key}/file_index"
+                        if chunk_col in episode and file_col in episode:
+                            video_file_indices[video_key] = {
+                                "chunk_index": int(episode[chunk_col]),
+                                "file_index": int(episode[file_col]),
+                            }
+                    print(video_file_indices)
+                    episode_meta = {
+                        "data/chunk_index": episode["data/chunk_index"],
+                        "data/file_index": episode["data/file_index"],
+                        "data/file_from_index": index,
+                        "videos/from_timestamps": from_timestamps,
+                        "videos/file_indices": video_file_indices,
+                    }
+                    # episode_meta = {
+                    #     "data/chunk_index": episode["data/chunk_index"],
+                    #     "data/file_index": episode["data/file_index"],
+                    #     "data/file_from_index": index,
+                    #     "videos/from_timestamps": from_timestamps,
+                    # }
+                    self.trajectory_ids_to_metadata[trajectory_ids[-1]] = episode_meta
+
+            # Should be able to directly read the saved index info here
+            return np.array(trajectory_ids), np.array(trajectory_lengths)
+
+    def _get_all_steps(self) -> list[tuple[int, int]]:
+        """Get the trajectory IDs and base indices for all steps in the dataset.
+
+        Returns:
+            list[tuple[str, int]]: A list of (trajectory_id, base_index) tuples.
+        """
+        def is_main():
+            return (not dist.is_initialized()) or dist.get_rank() == 0
+    
+        config_key = self._get_steps_config_key()
+        steps_filename = "steps_data_index.pkl"
+        steps_path = self.dataset_path / "meta" / steps_filename
+    
+        # ---------- try to read from cache  ----------
+        if steps_path.exists():
+            try:
+                with open(steps_path, "rb") as f:
+                    cached_data = pickle.load(f)
+                return cached_data["steps"]
+            except Exception as e:
+                # include EOFError / PickleError / KeyError
+                print(
+                    f"[RANK {os.environ.get('RANK', 'NA')}] "
+                    f"Failed to load cached steps ({e}), will rebuild."
+                )
+    
+        # ---------- only build by rank0  ----------
+        if is_main():
+            all_steps = self._get_all_steps_single_process()
+    
+            cache_data = {
+                "config_key": config_key,
+                "steps": all_steps,
+                "num_trajectories": len(self.trajectory_ids),
+                "total_steps": len(all_steps),
+                "computed_timestamp": pd.Timestamp.now().isoformat(),
+                "delete_pause_frame": self.delete_pause_frame,
+            }
+    
+            steps_path.parent.mkdir(parents=True, exist_ok=True)
+            tmp_path = steps_path.with_suffix(".tmp")
+    
+            with open(tmp_path, "wb") as f:
+                pickle.dump(cache_data, f, protocol=pickle.HIGHEST_PROTOCOL)
+            os.replace(tmp_path, steps_path)
+    
+            print(f"[RANK 0] Cached steps saved to {steps_path}")
+    
+        # ---------- sync after rank0  ----------
+        if dist.is_initialized():
+            dist.barrier()
+    
+        # ---------- read by all rank ----------
+        with open(steps_path, "rb") as f:
+            cached_data = pickle.load(f)
+    
+        return cached_data["steps"]
+
+    def _get_steps_config_key(self) -> str:
+        """Generate a configuration key for steps caching."""
+        config_dict = {
+            "delete_pause_frame": self.delete_pause_frame,
+            "dataset_name": self.dataset_name,
+        }
+        # Create a hash of the configuration
+        config_str = str(sorted(config_dict.items()))
+        return hashlib.md5(config_str.encode()).hexdigest()[:12]  #
+
+
+    def _get_all_steps_single_process(self) -> list[tuple[int, int]]:
+        """Original single-process implementation as fallback."""
+        all_steps: list[tuple[int, int]] = []
+        skipped_trajectories = 0
+        processed_trajectories = 0
+        
+        # Check if language modality is configured
+        has_language_modality = 'language' in self.modality_keys and len(self.modality_keys['language']) > 0
+        # TODO why trajectory_length here, why not use data length?
+        for trajectory_id, trajectory_length in tqdm(zip(self.trajectory_ids, self.trajectory_lengths), total=len(self.trajectory_ids), desc="Getting All Step"):
+            try:
+                if self._lerobot_version == "v2.0":
+                    data = self.get_trajectory_data(trajectory_id)
+                elif self._lerobot_version == "v3.0":
+                    data = self.get_trajectory_data_lerobot_v3(trajectory_id)
+                
+                trajectory_skipped = False
+            
+                # Check if trajectory has valid language instruction (if language modality is configured)
+                if has_language_modality:
+                    self.curr_traj_data = data  # Set current trajectory data for get_language to work
+
+                    language_instruction = self.get_language(trajectory_id, self.modality_keys['language'][0], 0)
+                    if not language_instruction or language_instruction[0] == "":
+                        print(f"Skipping trajectory {trajectory_id} due to empty language instruction")
+                        skipped_trajectories += 1
+                        trajectory_skipped = True
+                        continue
+
+            except Exception as e:
+                print(f"Skipping trajectory {trajectory_id} due to read error: {e}")
+                skipped_trajectories += 1
+                trajectory_skipped = True
+                continue
+        
+            if not trajectory_skipped:
+                processed_trajectories += 1
+        
+            for base_index in range(trajectory_length):
+                all_steps.append((trajectory_id, base_index))
+                
+        # Print summary statistics
+        print(f"Single-process summary: Processed {processed_trajectories} trajectories, skipped {skipped_trajectories} empty trajectories")
+        print(f"Total steps: {len(all_steps)} from {len(self.trajectory_ids)} trajectories")
+                   
+        return all_steps
+
+    def _get_position_and_gripper_values(self, data: pd.DataFrame) -> tuple[list, list]:
+        """Get position and gripper values based on available columns in the dataset."""
+        # Get action keys from modality_keys
+        action_keys = self.modality_keys.get('action', [])
+        
+        # Extract position data
+        delta_position_values = None
+        position_candidates = ['delta_eef_position']
+        coordinate_candidates = ['x', 'y', 'z']
+        
+        # First try combined position fields
+        for pos_key in position_candidates:
+            full_key = f"action.{pos_key}"
+            if full_key in action_keys:
+                try:
+                    # Get the lerobot key for this modality
+                    le_action_cfg = self.lerobot_modality_meta.action
+                    subkey = pos_key
+                    if subkey in le_action_cfg:
+                        le_key = le_action_cfg[subkey].original_key or subkey
+                        if le_key in data.columns:
+                            data_array = np.stack(data[le_key])
+                            le_indices = np.arange(le_action_cfg[subkey].start, le_action_cfg[subkey].end)
+                            filtered_data = data_array[:, le_indices]
+                            delta_position_values = filtered_data.tolist()
+                            break
+                except Exception:
+                    continue
+        
+        # If combined fields not found, try individual x,y,z coordinates
+        if delta_position_values is None:
+            x_data, y_data, z_data = None, None, None
+            for coord in coordinate_candidates:
+                full_key = f"action.{coord}"
+                if full_key in action_keys:
+                    try:
+                        le_action_cfg = self.lerobot_modality_meta.action
+                        if coord in le_action_cfg:
+                            le_key = le_action_cfg[coord].original_key or coord
+                            if le_key in data.columns:
+                                data_array = np.stack(data[le_key])
+                                le_indices = np.arange(le_action_cfg[coord].start, le_action_cfg[coord].end)
+                                coord_data = data_array[:, le_indices].flatten()
+                                if coord == 'x':
+                                    x_data = coord_data
+                                elif coord == 'y':
+                                    y_data = coord_data
+                                elif coord == 'z':
+                                    z_data = coord_data
+                    except Exception:
+                        continue
+            
+            if x_data is not None and y_data is not None and z_data is not None:
+                delta_position_values = np.column_stack((x_data, y_data, z_data)).tolist()
+        
+        if delta_position_values is None:
+            # Fallback to the old hardcoded approach if metadata approach fails
+            if 'action.delta_eef_position' in data.columns:
+                delta_position_values = data['action.delta_eef_position'].to_numpy().tolist()
+            elif all(col in data.columns for col in ['action.x', 'action.y', 'action.z']):
+                x_vals = data['action.x'].to_numpy()
+                y_vals = data['action.y'].to_numpy() 
+                z_vals = data['action.z'].to_numpy()
+                delta_position_values = np.column_stack((x_vals, y_vals, z_vals)).tolist()
+            else:
+                raise ValueError(f"No suitable position columns found. Available columns: {data.columns.tolist()}")
+        
+        # Extract gripper data
+        gripper_values = None
+        gripper_candidates = ['gripper_close', 'gripper']
+        
+        for grip_key in gripper_candidates:
+            full_key = f"action.{grip_key}"
+            if full_key in action_keys:
+                try:
+                    le_action_cfg = self.lerobot_modality_meta.action
+                    if grip_key in le_action_cfg:
+                        le_key = le_action_cfg[grip_key].original_key or grip_key
+                        if le_key in data.columns:
+                            data_array = np.stack(data[le_key])
+                            le_indices = np.arange(le_action_cfg[grip_key].start, le_action_cfg[grip_key].end)
+                            gripper_data = data_array[:, le_indices].flatten()
+                            gripper_values = gripper_data.tolist()
+                            break
+                except Exception:
+                    continue
+        
+        if gripper_values is None:
+            # Fallback to the old hardcoded approach if metadata approach fails
+            if 'action.gripper_close' in data.columns:
+                gripper_values = data['action.gripper_close'].to_numpy().tolist()
+            elif 'action.gripper' in data.columns:
+                gripper_values = data['action.gripper'].to_numpy().tolist()
+            else:
+                raise ValueError(f"No suitable gripper columns found. Available columns: {data.columns.tolist()}")
+        
+        return delta_position_values, gripper_values
+
+    def _get_modality_keys(self) -> dict:
+        """Get the modality keys for the dataset.
+        The keys are the modality names, and the values are the keys for each modality.
+        See property `modality_keys` for the expected format.
+        """
+        modality_keys = defaultdict(list)
+        for modality, config in self.modality_configs.items():
+            modality_keys[modality] = config.modality_keys
+        return modality_keys
+
+    def _get_delta_indices(self) -> dict[str, np.ndarray]:
+        """Restructure the delta indices to use modality.key as keys instead of just the modalities."""
+        delta_indices: dict[str, np.ndarray] = {}
+        for config in self.modality_configs.values():
+            for key in config.modality_keys:
+                delta_indices[key] = np.array(config.delta_indices)
+        return delta_indices
+
+    def _init_action_mode(self) -> None:
+        if self.data_cfg is None:
+            self._action_mode = "abs"
+            return
+
+        action_mode = self.data_cfg.get("action_mode", "abs")
+        if action_mode is None:
+            action_mode = "abs"
+        action_mode = _normalize_action_mode(action_mode)
+        if action_mode not in {"abs", "delta", "rel"}:
+            raise ValueError(f"Invalid action_mode: {action_mode}. Expected one of: abs, delta, rel.")
+        self._action_mode = action_mode
+
+        apply_keys = _normalize_action_mode_apply_keys(self.data_cfg.get("action_mode_apply_keys", None))
+        if apply_keys:
+            self._action_mode_apply_keys = apply_keys
+
+        self._action_mode_state_map = _normalize_action_mode_state_map(
+            self.data_cfg.get("action_mode_state_map", {}) or {}
+        )
+
+    def _infer_state_key_for_action(self, action_key: str) -> str | None:
+        if action_key in self._action_mode_state_map:
+            return self._action_mode_state_map[action_key]
+
+        if not action_key.startswith("action."):
+            return None
+        base = action_key.replace("action.", "", 1)
+        if f"state.{base}" in self.modality_keys.get("state", []):
+            return f"state.{base}"
+        return None
+
+    def _apply_action_mode(self, data: dict) -> dict:
+        if self._action_mode in (None, "abs"):
+            return data
+
+        action_keys = self._action_mode_apply_keys or self.modality_keys.get("action", [])
+        for action_key in action_keys:
+            if action_key not in data:
+                print(f"[WARNING] Action key {action_key} not found in data")
+                continue
+            state_key = self._infer_state_key_for_action(action_key)
+
+            # for safety, check if the state key is valid
+            if state_key is None or state_key not in data:
+                continue
+
+            action_values = np.asarray(data[action_key])
+            state_values = np.asarray(data[state_key])
+            if action_values.ndim != 2 or state_values.ndim != 2:
+                raise ValueError(
+                    f"Expected 2D arrays for action/state, got {action_key}: {action_values.shape}, {state_key}: {state_values.shape}"
+                )
+            if action_values.shape[1] != state_values.shape[1]:
+                raise ValueError(
+                    f"Action/state dim mismatch for {action_key} vs {state_key}: {action_values.shape} vs {state_values.shape}"
+                )
+
+            state0 = state_values[0]
+            if self._action_mode == "delta":
+                out = action_values.copy()
+                if len(out) > 1:
+                    out[1:] = action_values[1:] - action_values[:-1]
+                out[0] = action_values[0] - state0
+            elif self._action_mode == "rel":
+                out = action_values - state0
+            else:
+                out = action_values
+
+            data[action_key] = out
+
+        return data
+
+    def _get_lerobot_modality_meta(self) -> LeRobotModalityMetadata:
+        """Get the metadata for the LeRobot dataset."""
+        modality_meta_path = self.dataset_path / LE_ROBOT_MODALITY_FILENAME
+        assert (
+            modality_meta_path.exists()
+        ), f"Please provide a {LE_ROBOT_MODALITY_FILENAME} file in {self.dataset_path}"
+        with open(modality_meta_path, "r") as f:
+            modality_meta = LeRobotModalityMetadata.model_validate(json.load(f))
+        return modality_meta
+
+    def _get_lerobot_info_meta(self) -> dict:
+        """Get the metadata for the LeRobot dataset."""
+        info_meta_path = self.dataset_path / LE_ROBOT_INFO_FILENAME
+        with open(info_meta_path, "r") as f:
+            info_meta = json.load(f)
+        return info_meta
+
+    def _get_data_path_pattern(self) -> str:
+        """Get the data path pattern for the LeRobot dataset."""
+        return self.lerobot_info_meta["data_path"]
+
+    def _get_video_path_pattern(self) -> str:
+        """Get the video path pattern for the LeRobot dataset."""
+        return self.lerobot_info_meta["video_path"]
+
+    def _get_chunk_size(self) -> int:
+        """Get the chunk size for the LeRobot dataset."""
+        return self.lerobot_info_meta["chunks_size"]
+
+    def _get_tasks(self) -> pd.DataFrame:
+        """Get the tasks for the dataset."""
+        if self._lerobot_version == "v2.0":
+            tasks_path = self.dataset_path / LE_ROBOT_TASKS_FILENAME
+            with open(tasks_path, "r") as f:
+                tasks = [json.loads(line) for line in f]
+            df = pd.DataFrame(tasks)
+            return df.set_index("task_index")
+        
+        elif self._lerobot_version == "v3.0":
+            tasks_path = self.dataset_path / LE_ROBOT3_TASKS_FILENAME
+            df = pd.read_parquet(tasks_path)
+            df = df.reset_index()  # convert index to a column, typically named 'index'
+            df = df.rename(columns={'index': 'task'})  # rename 'index' column to 'task'
+            df = df[['task_index', 'task']]  # reorder columns
+            return df
+    def _check_integrity(self):
+        """Use the config to check if the keys are valid and detect silent data corruption."""
+        ERROR_MSG_HEADER = f"Error occurred in initializing dataset {self.dataset_name}:\n"
+
+        for modality_config in self.modality_configs.values():
+            for key in modality_config.modality_keys:
+                if key == "lapa_action" or key == "dream_actions":
+                    continue  # no need for any metadata for lapa actions because it comes normalized
+                # Check if the key is valid
+                try:
+                    self.lerobot_modality_meta.get_key_meta(key)
+                except Exception as e:
+                    raise ValueError(
+                        ERROR_MSG_HEADER + f"Unable to find key {key} in modality metadata:\n{e}"
+                    )
+
+    def set_transforms_metadata(self, metadata: DatasetMetadata):
+        """Set the metadata for the transforms. This is useful for transforms that need to know the metadata, such as the normalization values."""
+        self.transforms.set_metadata(metadata)
+
+    def set_epoch(self, epoch: int):
+        """Set the epoch for the dataset.
+
+        Args:
+            epoch (int): The epoch to set.
+        """
+        self.epoch = epoch
+
+    def __len__(self) -> int:
+        """Get the total number of data points in the dataset.
+
+        Returns:
+            int: the total number of data points in the dataset.
+        """
+        return len(self.all_steps)
+
+    def __str__(self) -> str:
+        """Get the description of the dataset."""
+        return f"{self.dataset_name} ({len(self)} steps)"
+
+
+    def __getitem__(self, index: int) -> dict:
+        """Get the data for a single step in a trajectory.
+
+        Args:
+            index (int): The index of the step to get.
+
+        Returns:
+            dict: The data for the step.
+        """
+        trajectory_id, base_index = self.all_steps[index]
+        raw_data = self.get_step_data(trajectory_id, base_index)
+        data = self.transforms(raw_data)
+        sample = self._pack_sample(data, trajectory_id=trajectory_id, base_index=base_index)
+        self._attach_rl_games_metadata(sample, base_index)
+        return sample
+
+    def _pack_sample(self, data: dict, trajectory_id: int | None = None, base_index: int | None = None) -> dict:
+        """Pack transformed modality data into training sample format."""
+        num_obs_frames = int(self.data_cfg.get("num_obs_frames", 1) or 1) if self.data_cfg is not None else 1
+        image_mode = str(self.data_cfg.get("image_mode", "single")) if self.data_cfg is not None else "single"
+        if self.data_cfg is not None and self.data_cfg.get("pack_image_sequence", False) not in ["False", False]:
+            num_obs_frames = int(self.data_cfg.get("image_sequence_length", num_obs_frames) or num_obs_frames)
+            if image_mode == "single":
+                image_mode = "multiframe"
+        stitch_grid = tuple(self.data_cfg.get("stitch_grid", [2, 2])) if self.data_cfg is not None else (2, 2)
+        kv_memory = bool(self.data_cfg.get("kv_memory", False)) if self.data_cfg is not None else False
+
+        step_images = []
+        view_frames = []
+        for video_key in self.modality_keys["video"]:
+            frames = data[video_key]  # (T, H, W, C), chronological (oldest .. newest)
+            if image_mode in {"single", "stitch_views"}:
+                selected = [frames[-1]]
+            else:
+                n = min(num_obs_frames, len(frames))
+                selected = [frames[i] for i in range(len(frames) - n, len(frames))]
+            pil_frames = [Image.fromarray(f) for f in selected]
+            if image_mode == "stitch":
+                step_images.append(stitch_frames(pil_frames, grid=stitch_grid, size=(224, 224)))
+            elif image_mode == "stitch_views":
+                view_frames.append(pil_frames[-1])
+            else:  # "single" or "multiframe": each frame tokenized independently
+                step_images.extend(p.resize((224, 224)) for p in pil_frames)
+
+        if image_mode == "stitch_views":
+            step_images.append(
+                stitch_frames(view_frames, grid=(1, len(view_frames)), size=(224, 224))
+            )
+
+        language = data[self.modality_keys["language"][0]][0]
+        language = _apply_prompt_mode(language, self.data_cfg.get("prompt_mode") if self.data_cfg is not None else None)
+        action = _concatenate_action_fields(data, self.modality_keys["action"]).astype(np.float16)
+
+        actions_per_frame = None
+        valid = None
+        density_weight = None
+        if kv_memory and base_index is not None and num_obs_frames > 1:
+            # Scheme B per-frame supervision. The action delta was widened to
+            # range(-(R-1), H) (see make_LeRobotSingleDataset), so `action` holds
+            # [R-1+H, dim] and frame k's chunk is the H-length slice at offset k.
+            R = num_obs_frames
+            horizon = action.shape[0] - (R - 1)
+            actions_per_frame = np.stack(
+                [action[k : k + horizon] for k in range(R)], axis=0
+            ).astype(np.float16)  # [R, H, dim]
+            # Keep `action` as the base-index chunk (last H rows) for any consumer
+            # that still reads the single-target field.
+            action = action[-horizon:]
+
+            # valid mask: a frame is real unless its global step was clamped (front
+            # padding at episode start). Mirrors get_video's clamp at datasets.py:1753.
+            video_key = self.modality_keys["video"][0]
+            video_delta = np.asarray(self.delta_indices[video_key])
+            step_idx = video_delta + int(base_index)
+            traj_len = int(self.trajectory_lengths[self.get_trajectory_index(trajectory_id)])
+            valid = ((step_idx >= 0) & (step_idx <= traj_len - 1)).astype(bool)  # [R]
+
+            # density_weight: per-frame loss weight (event balancing). Defaults to
+            # uniform; populated from per-row event metadata when available.
+            density_weight = self._kv_density_weight(step_idx, traj_len).astype(np.float32)  # [R]
+
+        sample = {
+            "action": action,
+            "image": step_images,
+            "lang": language,
+            "robot_tag": self.tag
+        }
+        if actions_per_frame is not None:
+            sample["actions_per_frame"] = actions_per_frame
+            sample["valid"] = valid
+            sample["density_weight"] = density_weight
+        if hasattr(self, "rl_games_task"):
+            sample["rl_games_task"] = self.rl_games_task
+        if hasattr(self, "rl_games_action_env_dim"):
+            sample["action_env_dim"] = int(self.rl_games_action_env_dim)
+
+        if self.data_cfg is not None and self.data_cfg.get("include_state", False) not in ["False", False]:
+            state = []
+            for state_key in self.modality_keys["state"]:
+                state.append(data[state_key])
+            state = np.concatenate(state, axis=1).astype(np.float32)
+            sample["state"] = state
+
+        auxiliary_fields = self.data_cfg.get("auxiliary_fields", {}) if self.data_cfg is not None else {}
+        if auxiliary_fields:
+            row = self.curr_traj_data.iloc[base_index]
+            for output_key, column in auxiliary_fields.items():
+                # Arrow nested lists arrive as object arrays of per-substep arrays.
+                sample[output_key] = np.stack(row[column])
+
+        return sample
+
+    def _kv_density_weight(self, step_idx: np.ndarray, traj_len: int) -> np.ndarray:
+        """Per-frame loss weight for KV-memory scheme-B supervision.
+
+        Replicates clean_v1's event balancing as a per-supervised-step weight instead
+        of dropping rows, so the rollout history stays contiguous. Reads a per-row
+        event-weight column when the exported dataset carries one (configurable via
+        ``datasets.vla_data.density_weight_key``); otherwise weights are uniform.
+        """
+        weights = np.ones(len(step_idx), dtype=np.float32)
+        key = self.data_cfg.get("density_weight_key", None) if self.data_cfg is not None else None
+        if not key or self.curr_traj_data is None or key not in self.curr_traj_data.columns:
+            return weights
+        column = self.curr_traj_data[key].to_numpy()
+        clamped = np.clip(step_idx, 0, traj_len - 1)
+        return column[clamped].astype(np.float32)
+
+    def _attach_rl_games_metadata(self, sample: dict, base_index: int) -> None:
+        """Attach non-modality RL-games metadata columns when converted datasets provide them."""
+        if self.curr_traj_data is None:
+            return
+        if "latency" in self.curr_traj_data.columns and 0 <= int(base_index) < len(self.curr_traj_data):
+            sample["latency"] = int(self.curr_traj_data.iloc[int(base_index)]["latency"])
+        elif "latency_id" in self.curr_traj_data.columns and 0 <= int(base_index) < len(self.curr_traj_data):
+            sample["latency"] = int(self.curr_traj_data.iloc[int(base_index)]["latency_id"])
+
+    def get_step_data(self, trajectory_id: int, base_index: int) -> dict:
+        """Get the RAW data for a single step in a trajectory. No transforms are applied.
+
+        Args:
+            trajectory_id (int): The name of the trajectory.
+            base_index (int): The base step index in the trajectory.
+
+        Returns:
+            dict: The RAW data for the step.
+
+        Example return:
+            {
+                "video": {
+                    "video.image_side_0": [B, T, H, W, C],
+                    "video.image_side_1": [B, T, H, W, C],
+                },
+                "state": {
+                    "state.eef_position": [B, T, state_dim],
+                    "state.eef_rotation": [B, T, state_dim],
+                },
+                "action": {
+                    "action.eef_position": [B, T, action_dim],
+                    "action.eef_rotation": [B, T, action_dim],
+                },
+            }
+        """
+        data = {}
+        # Get the data for all modalities # just for action base data
+        self.curr_traj_data = self.get_trajectory_data(trajectory_id)
+        # TODO @JinhuiYE The logic below is poorly implemented. Data reading should be directly based on curr_traj_data.
+        for modality in self.modality_keys:
+            # Get the data corresponding to each key in the modality
+            for key in self.modality_keys[modality]:
+                data[key] = self.get_data_by_modality(trajectory_id, modality, key, base_index)
+        data = self._apply_action_mode(data)
+        return data
+
+    def _read_trajectory_parquet(self, parquet_path: Path) -> pd.DataFrame:
+        if parquet_path not in self._parquet_cache:
+            self._parquet_cache[parquet_path] = pd.read_parquet(parquet_path)
+        return self._parquet_cache[parquet_path]
+
+    def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame:
+        """Get the data for a trajectory."""
+        if self._lerobot_version == "v2.0":
+        
+            if self.curr_traj_id == trajectory_id and self.curr_traj_data is not None:
+                return self.curr_traj_data
+            else:
+                chunk_index = self.get_episode_chunk(trajectory_id)
+                parquet_path = self.dataset_path / self.data_path_pattern.format(
+                    episode_chunk=chunk_index, episode_index=trajectory_id
+                )
+                self.curr_traj_data = self._read_trajectory_parquet(parquet_path)
+                self.curr_traj_id = trajectory_id
+                return self.curr_traj_data
+        elif self._lerobot_version == "v3.0":
+            return self.get_trajectory_data_lerobot_v3(trajectory_id)
+    
+    def get_trajectory_data_lerobot_v3(self, trajectory_id: int) -> pd.DataFrame:
+        """Get the data for a trajectory from lerobot v3."""
+        if self.curr_traj_id == trajectory_id and self.curr_traj_data is not None:
+            return self.curr_traj_data
+        else: #TODO check detail later
+            episode_meta = self.trajectory_ids_to_metadata[trajectory_id]
+            chunk_index = episode_meta["data/chunk_index"]
+            file_index = self.get_episode_file_index(trajectory_id)
+            # file_from_index = self.get_episode_file_from_index(trajectory_id)
+            
+            
+            parquet_path = self.dataset_path / self.data_path_pattern.format(
+                chunk_index=chunk_index, file_index=file_index
+            )
+            file_data = self._read_trajectory_parquet(parquet_path)
+            
+            # filter by trajectory_id
+            episode_data = file_data.loc[file_data["episode_index"] == trajectory_id].copy()
+            self.curr_traj_id = trajectory_id
+            self.curr_traj_data = episode_data
+            return self.curr_traj_data
+
+    def get_trajectory_columns(self, trajectory_id: int, columns: Sequence[str]) -> pd.DataFrame:
+        """Read selected columns for one trajectory without loading image/action payloads."""
+        requested_columns = list(dict.fromkeys(str(column) for column in columns))
+        if self._lerobot_version == "v2.0":
+            chunk_index = self.get_episode_chunk(trajectory_id)
+            parquet_path = self.dataset_path / self.data_path_pattern.format(
+                episode_chunk=chunk_index, episode_index=trajectory_id
+            )
+            assert parquet_path.exists(), f"Parquet file not found at {parquet_path}"
+            return pd.read_parquet(parquet_path, columns=requested_columns)
+
+        if self._lerobot_version == "v3.0":
+            episode_meta = self.trajectory_ids_to_metadata[trajectory_id]
+            chunk_index = episode_meta["data/chunk_index"]
+            file_index = self.get_episode_file_index(trajectory_id)
+            parquet_path = self.dataset_path / self.data_path_pattern.format(
+                chunk_index=chunk_index, file_index=file_index
+            )
+            assert parquet_path.exists(), f"Parquet file not found at {parquet_path}"
+            columns_with_episode = list(dict.fromkeys([*requested_columns, "episode_index"]))
+            file_data = pd.read_parquet(parquet_path, columns=columns_with_episode)
+            return file_data.loc[file_data["episode_index"] == trajectory_id, requested_columns].copy()
+
+        raise ValueError(f"Unsupported LeRobot version: {self._lerobot_version}")
+
+    def get_action_ids_for_all_steps(self, action_key: str = "action_id") -> list[int]:
+        """Return action IDs aligned one-to-one with ``self.all_steps``."""
+        cache_attr = f"_cached_step_action_ids_{action_key}"
+        cached = getattr(self, cache_attr, None)
+        if cached is not None:
+            return cached
+
+        action_ids_by_trajectory: dict[int, np.ndarray] = {}
+        for trajectory_id, _ in self.all_steps:
+            if trajectory_id in action_ids_by_trajectory:
+                continue
+            try:
+                action_data = self.get_trajectory_columns(trajectory_id, [action_key])
+            except Exception as exc:
+                raise ValueError(
+                    f"Unable to read action balance key `{action_key}` from dataset "
+                    f"`{self.dataset_name}` trajectory {trajectory_id}."
+                ) from exc
+            if action_key not in action_data.columns:
+                raise ValueError(
+                    f"Action balance key `{action_key}` not found in dataset `{self.dataset_name}`."
+                )
+            action_ids_by_trajectory[trajectory_id] = action_data[action_key].to_numpy()
+
+        action_ids: list[int] = []
+        for trajectory_id, base_index in self.all_steps:
+            trajectory_actions = action_ids_by_trajectory[trajectory_id]
+            if base_index >= len(trajectory_actions):
+                raise IndexError(
+                    f"Step index {base_index} is out of range for action key `{action_key}` "
+                    f"in dataset `{self.dataset_name}` trajectory {trajectory_id}."
+                )
+            action_ids.append(int(trajectory_actions[base_index]))
+
+        setattr(self, cache_attr, action_ids)
+        return action_ids
+
+    def get_latencies_for_all_steps(self) -> list[int]:
+        """Return latency IDs aligned one-to-one with ``self.all_steps``."""
+        cached = getattr(self, "_cached_step_latencies", None)
+        if cached is not None:
+            return cached
+
+        latencies_by_trajectory: dict[int, np.ndarray] = {}
+        latency_key_by_trajectory: dict[int, str] = {}
+        for trajectory_id, _ in self.all_steps:
+            trajectory_id = int(trajectory_id)
+            if trajectory_id in latencies_by_trajectory:
+                continue
+            for latency_key in ("latency", "latency_id"):
+                try:
+                    latency_data = self.get_trajectory_columns(trajectory_id, [latency_key])
+                except Exception:
+                    continue
+                if latency_key in latency_data.columns:
+                    latencies_by_trajectory[trajectory_id] = latency_data[latency_key].to_numpy()
+                    latency_key_by_trajectory[trajectory_id] = latency_key
+                    break
+            if trajectory_id not in latencies_by_trajectory:
+                raise ValueError(
+                    f"Dataset `{self.dataset_name}` trajectory {trajectory_id} has no `latency` "
+                    "or `latency_id` column required for latency curriculum sampling."
+                )
+
+        latencies: list[int] = []
+        for trajectory_id, base_index in self.all_steps:
+            trajectory_id = int(trajectory_id)
+            trajectory_latencies = latencies_by_trajectory[trajectory_id]
+            if base_index >= len(trajectory_latencies):
+                latency_key = latency_key_by_trajectory[trajectory_id]
+                raise IndexError(
+                    f"Step index {base_index} is out of range for `{latency_key}` in dataset "
+                    f"`{self.dataset_name}` trajectory {trajectory_id}."
+                )
+            latencies.append(int(trajectory_latencies[base_index]))
+
+        setattr(self, "_cached_step_latencies", latencies)
+        return latencies
+
+
+    def get_trajectory_index(self, trajectory_id: int) -> int:
+        """Get the index of the trajectory in the dataset by the trajectory ID.
+        This is useful when you need to get the trajectory length or sampling weight corresponding to the trajectory ID.
+
+        Args:
+            trajectory_id (str): The ID of the trajectory.
+
+        Returns:
+            int: The index of the trajectory in the dataset.
+        """
+        trajectory_indices = np.where(self.trajectory_ids == trajectory_id)[0]
+        if len(trajectory_indices) != 1:
+            raise ValueError(
+                f"Error finding trajectory index for {trajectory_id}, found {trajectory_indices=}"
+            )
+        return trajectory_indices[0]
+
+    def get_episode_chunk(self, ep_index: int) -> int:
+        """Get the chunk index for an episode index."""
+        return ep_index // self.chunk_size
+    def get_episode_file_index(self, ep_index: int) -> int:
+        """Get the file index for an episode index."""
+        episode_meta = self.trajectory_ids_to_metadata[ep_index]
+        return episode_meta["data/file_index"]
+    
+    def get_episode_file_from_index(self, ep_index: int) -> int:
+        """Get the file from index for an episode index."""
+        episode_meta = self.trajectory_ids_to_metadata[ep_index]
+        return episode_meta["data/file_from_index"]
+
+
+    def retrieve_data_and_pad(
+        self,
+        array: np.ndarray,
+        step_indices: np.ndarray,
+        max_length: int,
+        padding_strategy: str = "first_last",
+    ) -> np.ndarray:
+        """Retrieve the data from the dataset and pad it if necessary.
+        Args:
+            array (np.ndarray): The array to retrieve the data from.
+            step_indices (np.ndarray): The step indices to retrieve the data for.
+            max_length (int): The maximum length of the data.
+            padding_strategy (str): The padding strategy, either "first" or "last".
+        """
+        # Get the padding indices
+        front_padding_indices = step_indices < 0
+        end_padding_indices = step_indices >= max_length
+        padding_positions = np.logical_or(front_padding_indices, end_padding_indices)
+        # Retrieve the data with the non-padding indices
+        # If there exists some padding, Given T step_indices, the shape of the retrieved data will be (T', ...) where T' < T
+        raw_data = array[step_indices[~padding_positions]]
+        assert isinstance(raw_data, np.ndarray), f"{type(raw_data)=}"
+        # This is the shape of the output, (T, ...)
+        if raw_data.ndim == 1:
+            expected_shape = (len(step_indices),)
+        else:
+            expected_shape = (len(step_indices), *array.shape[1:])
+
+        # Pad the data
+        output = np.zeros(expected_shape)
+        # Assign the non-padded data
+        output[~padding_positions] = raw_data
+        # If there exists some padding, pad the data
+        if padding_positions.any():
+            if padding_strategy == "first_last":
+                # Use first / last step data to pad
+                front_padding_data = array[0]
+                end_padding_data = array[-1]
+                output[front_padding_indices] = front_padding_data
+                output[end_padding_indices] = end_padding_data
+            elif padding_strategy == "zero":
+                # Use zero padding
+                output[padding_positions] = 0
+            else:
+                raise ValueError(f"Invalid padding strategy: {padding_strategy}")
+        return output
+
+    def get_video_path(self, trajectory_id: int, key: str) -> Path:
+        chunk_index = self.get_episode_chunk(trajectory_id)
+        original_key = self.lerobot_modality_meta.video[key].original_key
+        if original_key is None:
+            original_key = key
+        if self._lerobot_version == "v2.0":
+            video_filename = self.video_path_pattern.format(
+                episode_chunk=chunk_index, episode_index=trajectory_id, video_key=original_key
+            )
+        elif self._lerobot_version == "v3.0":
+            episode_meta = self.trajectory_ids_to_metadata[trajectory_id]
+
+            video_file_indices = episode_meta.get("videos/file_indices", {})
+            # print(f"{video_file_indices=}")
+            #已修改的lerobotv3.0的视频索引
+            if original_key in video_file_indices:
+                video_chunk_index = video_file_indices[original_key]["chunk_index"]
+                video_file_index = video_file_indices[original_key]["file_index"]
+            else:
+                video_chunk_index = episode_meta["data/chunk_index"]
+                video_file_index = episode_meta["data/file_index"]
+            video_filename = self.video_path_pattern.format(
+                video_key=original_key,
+                chunk_index=episode_meta["data/chunk_index"],
+                file_index=episode_meta["data/file_index"],
+            )
+        return self.dataset_path / video_filename
+
+    def get_video(
+        self,
+        trajectory_id: int,
+        key: str,
+        base_index: int,
+    ) -> np.ndarray:
+        """Get the video frames for a trajectory by a base index.
+
+        Args:
+            dataset (BaseSingleDataset): The dataset to retrieve the data from.
+            trajectory_id (str): The ID of the trajectory.
+            key (str): The key of the video.
+            base_index (int): The base index of the trajectory.
+
+        Returns:
+            np.ndarray: The video frames for the trajectory and frame indices. Shape: (T, H, W, C)
+        """
+        # Get the step indices
+        step_indices = self.delta_indices[key] + base_index
+        # print(f"{step_indices=}")
+        # Get the trajectory index
+        trajectory_index = self.get_trajectory_index(trajectory_id)
+        # Ensure the indices are within the valid range
+        # This is equivalent to padding the video with extra frames at the beginning and end
+        step_indices = np.maximum(step_indices, 0)
+        step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1)
+        assert key.startswith("video."), f"Video key must start with 'video.', got {key}"
+        # Get the sub-key
+        key = key.replace("video.", "")
+
+        # Image-only LeRobot datasets (e.g. VLA-Arena) may store frames directly
+        # in parquet columns and have total_videos == 0 (no mp4 files). In this
+        # case, load frames from the original image column and pad by step indices.
+        original_key = self.lerobot_modality_meta.video[key].original_key
+        if original_key is None:
+            original_key = key
+
+        def _decode_image_entry(entry):
+            if isinstance(entry, np.ndarray):
+                return entry
+            if isinstance(entry, Image.Image):
+                return np.array(entry)
+            if isinstance(entry, (bytes, bytearray, memoryview)):
+                return np.array(Image.open(io.BytesIO(bytes(entry))).convert("RGB"))
+            if isinstance(entry, dict):
+                img_bytes = entry.get("bytes", None)
+                img_path = entry.get("path", None)
+
+                if img_bytes is not None:
+                    return np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
+
+                if img_path is not None:
+                    path_obj = Path(img_path)
+                    if not path_obj.is_absolute():
+                        path_obj = self.dataset_path / path_obj
+                    return np.array(Image.open(path_obj).convert("RGB"))
+
+            raise TypeError(f"Unsupported image entry type: {type(entry)}")
+
+        context_images_column = None
+        if self.data_cfg is not None:
+            configured_context_images_column = self.data_cfg.get("context_images_column", None)
+            if configured_context_images_column not in (None, ""):
+                context_images_column = str(configured_context_images_column)
+        if context_images_column is not None and self.curr_traj_data is not None:
+            missing_columns = [
+                column
+                for column in (context_images_column, original_key)
+                if column not in self.curr_traj_data.columns
+            ]
+            if missing_columns:
+                raise ValueError(
+                    f"Configured context image columns are missing from dataset {self.dataset_name} "
+                    f"trajectory {trajectory_id}: {missing_columns}"
+                )
+            image_sequence_length = int(
+                self.data_cfg.get("image_sequence_length", 1) if self.data_cfg is not None else 1
+            )
+            expected_context_images = image_sequence_length - 1
+            row = self.curr_traj_data.iloc[int(base_index)]
+            context_values = list(row[context_images_column])
+            if len(context_values) != expected_context_images:
+                raise ValueError(
+                    f"Expected {expected_context_images} context image frames, got {len(context_values)}"
+                )
+            return np.stack([_decode_image_entry(item) for item in [*context_values, row[original_key]]])
+
+        if self.lerobot_info_meta["features"][original_key]["dtype"] == "image":
+            image_entries = self.curr_traj_data[original_key].tolist()
+
+            frames = []
+            for idx in step_indices:
+                safe_idx = int(min(max(idx, 0), len(image_entries) - 1))
+                frames.append(_decode_image_entry(image_entries[safe_idx]))
+
+            return np.stack(frames)
+
+        video_path = self.get_video_path(trajectory_id, key)
+        # Get the action/state timestamps for each frame in the video
+        assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}"
+        assert "timestamp" in self.curr_traj_data.columns, f"No timestamp found in {trajectory_id=}"
+        timestamp: np.ndarray = self.curr_traj_data["timestamp"].to_numpy()
+        # Get the corresponding video timestamps from the step indices
+        video_timestamp = timestamp[step_indices]
+        if self._lerobot_version == "v3.0":
+            episode_meta = self.trajectory_ids_to_metadata.get(trajectory_id, {})
+            from_timestamps = episode_meta.get("videos/from_timestamps", {})
+            original_video_key = self.lerobot_modality_meta.video[key].original_key
+            if original_video_key is None:
+                original_video_key = key
+            from_timestamp = float(from_timestamps.get(original_video_key, 0.0))
+            video_timestamp = video_timestamp + from_timestamp
+
+        return get_frames_by_timestamps(
+            video_path.as_posix(),
+            video_timestamp,
+            video_backend=self.video_backend, # TODO
+            video_backend_kwargs=self.video_backend_kwargs,
+        )
+
+    def get_state_or_action(
+        self,
+        trajectory_id: int,
+        modality: str,
+        key: str,
+        base_index: int,
+    ) -> np.ndarray:
+        """Get the state or action data for a trajectory by a base index.
+        If the step indices are out of range, pad with the data:
+            if the data is stored in absolute format, pad with the first or last step data;
+            otherwise, pad with zero.
+
+        Args:
+            dataset (BaseSingleDataset): The dataset to retrieve the data from.
+            trajectory_id (int): The ID of the trajectory.
+            modality (str): The modality of the data.
+            key (str): The key of the data.
+            base_index (int): The base index of the trajectory.
+
+        Returns:
+            np.ndarray: The data for the trajectory and step indices.
+        """
+        # Get the step indices
+        step_indices = self.delta_indices[key] + base_index
+        # Get the trajectory index
+        trajectory_index = self.get_trajectory_index(trajectory_id)
+        # Get the maximum length of the trajectory
+        max_length = self.trajectory_lengths[trajectory_index]
+        assert key.startswith(modality + "."), f"{key} must start with {modality + '.'}, got {key}"
+        # Get the sub-key, e.g. state.joint_angles -> joint_angles
+        key = key.replace(modality + ".", "")
+        # Get the lerobot key
+        le_state_or_action_cfg = getattr(self.lerobot_modality_meta, modality)
+        le_key = le_state_or_action_cfg[key].original_key
+        if le_key is None:
+            le_key = key
+        # Read only the requested steps. Some game episodes contain 18,000
+        # embedded RGB rows; stacking the whole trajectory for each sampled
+        # state/action makes dataloader time proportional to episode length.
+        assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}"
+        assert le_key in self.curr_traj_data.columns, f"No {le_key} found in {trajectory_id=}"
+        padded_indices = np.clip(step_indices, 0, max_length - 1)
+        data_array = np.stack(
+            [np.stack(self.curr_traj_data[le_key].iloc[int(index)]) for index in padded_indices]
+        )
+        le_indices = np.arange(
+            le_state_or_action_cfg[key].start,
+            le_state_or_action_cfg[key].end,
+        )
+        data_array = data_array[..., le_indices]
+        # Get the state or action configuration
+        state_or_action_cfg = getattr(self.metadata.modalities, modality)[key]
+        output = np.zeros(data_array.shape)
+        output[:] = data_array
+        if not state_or_action_cfg.absolute:
+            output[(step_indices < 0) | (step_indices >= max_length)] = 0
+        return output
+
+    def get_language(
+        self,
+        trajectory_id: int,
+        key: str,
+        base_index: int,
+    ) -> list[str]:
+        """Get the language annotation data for a trajectory by step indices.
+
+        Args:
+            dataset (BaseSingleDataset): The dataset to retrieve the data from.
+            trajectory_id (int): The ID of the trajectory.
+            key (str): The key of the annotation.
+            base_index (int): The base index of the trajectory.
+
+        Returns:
+            list[str]: The annotation data for the trajectory and step indices. If no matching data is found, return empty strings.
+        """
+        assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}"
+        # Get the step indices
+        step_indices = self.delta_indices[key] + base_index
+        # Get the trajectory index
+        trajectory_index = self.get_trajectory_index(trajectory_id)
+        # Get the maximum length of the trajectory
+        max_length = self.trajectory_lengths[trajectory_index]
+        # Get the end times corresponding to the closest indices
+        step_indices = np.maximum(step_indices, 0)
+        step_indices = np.minimum(step_indices, max_length - 1)
+        # Get the annotations
+        task_indices: list[int] = []
+        assert key.startswith(
+            "annotation."
+        ), f"Language key must start with 'annotation.', got {key}"
+        subkey = key.replace("annotation.", "")
+        annotation_meta = self.lerobot_modality_meta.annotation
+        assert annotation_meta is not None, f"Annotation metadata is None for {subkey}"
+        assert (
+            subkey in annotation_meta
+        ), f"Annotation key {subkey} not found in metadata, available annotation keys: {annotation_meta.keys()}"
+        subkey_meta = annotation_meta[subkey]
+        original_key = subkey_meta.original_key
+        if original_key is None:
+            original_key = key
+        for i in range(len(step_indices)): # 
+            # task_indices.append(self.curr_traj_data[original_key][step_indices[i]].item())
+            value = self.curr_traj_data[original_key].iloc[step_indices[i]] # TODO check v2.0 
+            task_indices.append(value if isinstance(value, (int, float)) else value.item())
+
+        return self.tasks.loc[task_indices]["task"].tolist()
+
+    def get_data_by_modality(
+        self,
+        trajectory_id: int,
+        modality: str,
+        key: str,
+        base_index: int,
+    ):
+        """Get the data corresponding to the modality for a trajectory by a base index.
+        This method will call the corresponding helper method based on the modality.
+        See the helper methods for more details.
+        NOTE: For the language modality, the data is padded with empty strings if no matching data is found.
+
+        Args:
+            dataset (BaseSingleDataset): The dataset to retrieve the data from.
+            trajectory_id (int): The ID of the trajectory.
+            modality (str): The modality of the data.
+            key (str): The key of the data.
+            base_index (int): The base index of the trajectory.
+        """
+        if modality == "video":
+            return self.get_video(trajectory_id, key, base_index)
+        elif modality == "state" or modality == "action":
+            return self.get_state_or_action(trajectory_id, modality, key, base_index)
+        elif modality == "language":
+            return self.get_language(trajectory_id, key, base_index)
+        else:
+            raise ValueError(f"Invalid modality: {modality}")
+
+    def _save_dataset_statistics_(self, save_path: Path | str, format: str = "json") -> None:
+        """
+        Save dataset statistics to specified path in the required format.
+        Only includes statistics for keys that are actually used in the dataset.
+        Gripper-related keys will be placed at the end.
+        
+        Args:
+            save_path (Path | str): Path to save the statistics file
+            format (str): Save format, currently only supports "json"
+        """
+        save_path = Path(save_path)
+        save_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        # Build the data structure to save
+        statistics_data = {}
+        
+        # Get used modality keys
+        used_action_keys, used_state_keys = get_used_modality_keys(self.modality_keys)
+        
+        # Organize statistics by tag
+        tag = self.tag
+        tag_stats = {}
+        
+        # Process action statistics (only for used keys)
+        if hasattr(self.metadata.statistics, 'action') and self.metadata.statistics.action:
+            action_stats = self.metadata.statistics.action
+            
+            # Filter to only include used action keys and reorder: non-gripper first, gripper last
+            non_gripper_keys = []
+            gripper_keys = []
+            
+            for key in action_stats.keys():
+                if key in used_action_keys:
+                    if "gripper" in key.lower():
+                        gripper_keys.append(key)
+                    else:
+                        non_gripper_keys.append(key)
+            
+            # Reorder: non-gripper first, gripper last
+            reordered_keys = non_gripper_keys + gripper_keys
+            
+            filtered_action_stats = {}
+            for key in reordered_keys:
+                filtered_action_stats[key] = action_stats[key]
+            
+            if filtered_action_stats:
+                # Combine statistics from filtered action sub-keys
+                combined_action_stats = combine_modality_stats(filtered_action_stats)
+                
+                # Add mask field based on whether it's gripper or not
+                mask = generate_action_mask_for_used_keys(
+                    self.metadata.modalities.action, filtered_action_stats.keys()
+                )
+                combined_action_stats["mask"] = mask
+                
+                tag_stats["action"] = combined_action_stats
+        
+        # Process state statistics (only for used keys)
+        if hasattr(self.metadata.statistics, 'state') and self.metadata.statistics.state:
+            state_stats = self.metadata.statistics.state
+            
+            # Filter to only include used state keys, optionally reorder gripper to end
+            non_gripper_keys = []
+            gripper_keys = []
+            
+            for key in state_stats.keys():
+                if key in used_state_keys:
+                    if "gripper" in key.lower():
+                        gripper_keys.append(key)
+                    else:
+                        non_gripper_keys.append(key)
+            
+            # Reorder: non-gripper first, gripper last
+            reordered_keys = non_gripper_keys + gripper_keys
+            
+            filtered_state_stats = {}
+            for key in reordered_keys:
+                filtered_state_stats[key] = state_stats[key]
+            
+            if filtered_state_stats:
+                combined_state_stats = combine_modality_stats(filtered_state_stats)
+                tag_stats["state"] = combined_state_stats
+        
+        # Add dataset counts
+        tag_stats["num_transitions"] = len(self)
+        tag_stats["num_trajectories"] = len(self.trajectory_ids)
+        
+        statistics_data[tag] = tag_stats
+        
+        # Save as JSON file
+        if format.lower() == "json":
+            if not str(save_path).endswith('.json'):
+                save_path = save_path.with_suffix('.json')
+            with open(save_path, 'w', encoding='utf-8') as f:
+                json.dump(statistics_data, f, indent=2, ensure_ascii=False)
+        else:
+            raise ValueError(f"Unsupported format: {format}. Currently only 'json' is supported.")
+        
+        print(f"Single dataset statistics saved to: {save_path}")
+        print(f"Used action keys (reordered): {list(used_action_keys)}")
+        print(f"Used state keys (reordered): {list(used_state_keys)}")
+
+
+class CachedLeRobotSingleDataset(LeRobotSingleDataset):
+    def __init__(self, img_resize: tuple[int, int] | None = None, *args, **kwargs):
+        """
+        This class caches the video frames for each trajectory and key.
+        It is recommended to use this class if the video frames need to be accessed multiple times.
+
+        Args:
+            resize_img (tuple[int, int], optional): The size to resize the video frames to reduce memory usage.
+        """
+        # Convert img_resize to tuple if it is not already
+        if img_resize is not None and not isinstance(img_resize, tuple):
+            img_resize = tuple(img_resize)
+            assert len(img_resize) == 2, f"Expected tuple of length 2, got {img_resize}"
+        self.img_resize = img_resize
+
+        # Initialize img_resize attribute first to ensure it exists
+        super().__init__(*args, **kwargs)
+        cached_frames: dict[str, np.ndarray] = {}
+
+        for key in self.modality_keys["video"]:
+            all_frames = []
+            original_key = key
+            key = key.replace("video.", "")
+            for trajectory_id, trajectory_length in tqdm(
+                zip(self.trajectory_ids, self.trajectory_lengths),
+                total=len(self.trajectory_ids),
+                desc=f"Caching {key} frames",
+            ):
+                video_path = self.get_video_path(trajectory_id, key)
+                frames = get_all_frames(
+                    video_path.as_posix(),
+                    video_backend=self.video_backend,
+                    video_backend_kwargs=self.video_backend_kwargs,
+                    resize_size=img_resize,
+                )
+                assert frames.ndim == 4, f"Expected 4D array, got {frames.shape} array"
+                assert frames.shape[3] == 3, f"Expected 3 channels, got {frames.shape[3]} channels"
+                
+                # Apply image cropping if enabled and the video key is base_view
+                # Note: crop_obs_camera functionality has been removed
+                
+                # assert (
+                #     frames.shape[0] == trajectory_length
+                # ), f"Expected {trajectory_length} frames, got {frames.shape[0]} frames"
+                all_frames.append(frames)
+            cached_frames[key] = np.concatenate(all_frames, axis=0)
+            print(f"{key}: {cached_frames[key].shape}")
+        self.cached_frames = cached_frames
+        self.start_indices = np.cumsum(self.trajectory_lengths) - self.trajectory_lengths
+
+    def get_video(self, trajectory_id: int, key: str, base_index: int) -> np.ndarray:
+        step_indices = self.delta_indices[key] + base_index
+        # Get the trajectory index
+        trajectory_index = self.get_trajectory_index(trajectory_id)
+        # Ensure the indices are within the valid range
+        # This is equivalent to padding the video with extra frames at the beginning and end
+        step_indices = np.maximum(step_indices, 0)
+        step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1)
+        assert key.startswith("video."), f"Video key must start with 'video.', got {key}"
+        # Get the sub-key
+        key = key.replace("video.", "")
+        # Calculate the absolute indices
+        absolute_indices = self.start_indices[trajectory_index] + step_indices
+        return self.cached_frames[key][absolute_indices]
+
+    def get_step_data(self, trajectory_id: int, base_index: int) -> dict:
+        """Get the RAW data for a single step. No transforms are applied.
+
+        Args:
+            trajectory_id (str): The ID of the trajectory.
+            base_index (int): The base index of the step.
+
+        Returns:
+            dict: The data for the step.
+        """
+        data = {}
+        self.curr_traj_data = self.get_trajectory_data(trajectory_id)
+        # Get the data for all modalities
+        for modality in self.modality_keys:
+            # Get the data corresponding to each key in the modality
+            for key in self.modality_keys[modality]:
+                data[key] = self.get_data_by_modality(trajectory_id, modality, key, base_index)
+        return data
+
+    def set_transforms_metadata(self, metadata: DatasetMetadata):
+        """Set the metadata for the transforms. This is useful for transforms that need to know the metadata, such as the normalization values."""
+        if self.img_resize is not None:
+            all_video_keys = [key for key in self.modality_keys["video"]]
+            for key in metadata.modalities.video:
+                if key in all_video_keys:
+                    metadata.modalities.video[key].resolution = self.img_resize
+        super().set_transforms_metadata(metadata)
+
+
+def safe_hash(input_tuple):
+    # keep 128 bits of the hash
+    tuple_string = repr(input_tuple).encode("utf-8")
+    sha256 = hashlib.sha256()
+    sha256.update(tuple_string)
+
+    seed = int(sha256.hexdigest(), 16)
+
+    return seed & 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF
+
+
+class MixtureSpecElement(BaseModel):
+    dataset_path: list[Path] | Path = Field(..., description="The path to the dataset.")
+    dataset_weight: float = Field(..., description="The weight of the dataset in the mixture.")
+    distribute_weights: bool = Field(
+        default=False,
+        description="Whether to distribute the weights of the dataset across all the paths. If True, the weights will be evenly distributed across all the paths.",
+    )
+
+
+# Helper functions for dataset statistics
+
+def combine_modality_stats(modality_stats: dict) -> dict:
+    """
+    Combine statistics from all sub-keys under a modality.
+    
+    Args:
+        modality_stats (dict): Statistics for a modality, containing multiple sub-keys.
+                               Each sub-key contains DatasetStatisticalValues object.
+        
+    Returns:
+        dict: Combined statistics
+    """
+    combined_stats = {
+        "mean": [],
+        "std": [],
+        "max": [],
+        "min": [],
+        "q01": [],
+        "q99": []
+    }
+    
+    # Combine statistics in sub-key order
+    for subkey in modality_stats.keys():
+        subkey_stats = modality_stats[subkey]  # This is a DatasetStatisticalValues object
+        
+        # Convert DatasetStatisticalValues to dict-like access
+        for stat_name in ["mean", "std", "max", "min", "q01", "q99"]:
+            stat_value = getattr(subkey_stats, stat_name)
+            if isinstance(stat_value, (list, tuple)):
+                combined_stats[stat_name].extend(stat_value)
+            else:
+                # Handle NDArray case - convert to list
+                if hasattr(stat_value, 'tolist'):
+                    combined_stats[stat_name].extend(stat_value.tolist())
+                else:
+                    combined_stats[stat_name].append(float(stat_value))
+    
+    return combined_stats
+
+def generate_action_mask_for_used_keys(action_modalities: dict, used_action_keys_ordered) -> list[bool]:
+    """
+    Generate mask based on action modalities, but only for used keys.
+    Gripper-related are False, others are True.
+    
+    Args:
+        action_modalities (dict): Configuration information for action modalities.
+        used_action_keys_ordered: Iterable of actually used action keys in the correct order.
+        
+    Returns:
+        list[bool]: List of mask values
+    """
+    mask = []
+    
+    # Generate mask in the same order as the statistics were combined
+    for subkey in used_action_keys_ordered:
+        if subkey in action_modalities:
+            subkey_config = action_modalities[subkey]
+            
+            # Get dimension count from shape
+            if hasattr(subkey_config, 'shape') and len(subkey_config.shape) > 0:
+                dim_count = subkey_config.shape[0]
+            else:
+                dim_count = 1
+            
+            # Check if it's gripper-related
+            is_gripper = "gripper" in subkey.lower()
+            
+            # Generate mask value for each dimension
+            for _ in range(dim_count):
+                mask.append(not is_gripper)  # gripper is False, others are True
+    
+    return mask
+
+def get_used_modality_keys(modality_keys: dict) -> tuple[list, list]:
+    """Extract used action and state keys from modality configuration."""
+    used_action_keys = []
+    used_state_keys = []
+    
+    # Extract action keys (remove "action." prefix)
+    for action_key in modality_keys.get("action", []):
+        if action_key.startswith("action."):
+            clean_key = action_key.replace("action.", "")
+            used_action_keys.append(clean_key)
+    
+    # Extract state keys (remove "state." prefix)  
+    for state_key in modality_keys.get("state", []):
+        if state_key.startswith("state."):
+            clean_key = state_key.replace("state.", "")
+            used_state_keys.append(clean_key)
+    
+    return used_action_keys, used_state_keys
+
+
+def _strip_video_capture_metadata(obj):
+    """Drop source-capture video metadata for cross-task compatibility checks.
+
+    RL-games mixtures can combine tasks with different native FPS/resolutions;
+    the training transform resizes observations before model input, so these
+    fields are descriptive rather than a distinct model contract.
+    """
+    if isinstance(obj, dict):
+        return {
+            key: _strip_video_capture_metadata(value)
+            for key, value in obj.items()
+            if key not in {"fps", "resolution"}
+        }
+    if isinstance(obj, list):
+        return [_strip_video_capture_metadata(value) for value in obj]
+    return obj
+
+
+class LeRobotMixtureDataset(Dataset):
+    """
+    A mixture of multiple datasets. This class samples a single dataset based on the dataset weights and then calls the `__getitem__` method of the sampled dataset.
+    It is recommended to modify the single dataset class instead of this class.
+    """
+
+    def __init__(
+        self,
+        data_mixture: Sequence[tuple[LeRobotSingleDataset, float]],
+        mode: str,
+        balance_dataset_weights: bool = True,
+        balance_trajectory_weights: bool = True,
+        seed: int = 42,
+        metadata_config: dict = {
+            "percentile_mixing_method": "min_max",
+        },
+        **kwargs,
+    ):
+        """
+        Initialize the mixture dataset.
+
+        Args:
+            data_mixture (list[tuple[LeRobotSingleDataset, float]]): Datasets and their corresponding weights.
+            mode (str): If "train", __getitem__ will return different samples every epoch; if "val" or "test", __getitem__ will return the same sample every epoch.
+            balance_dataset_weights (bool): If True, the weight of dataset will be multiplied by the total trajectory length of each dataset.
+            balance_trajectory_weights (bool): If True, sample trajectories within a dataset weighted by their length; otherwise, use equal weighting.
+            seed (int): Random seed for sampling.
+        """
+        datasets: list[LeRobotSingleDataset] = []
+        dataset_sampling_weights: list[float] = []
+        for dataset, weight in data_mixture:
+            # Check if dataset is valid and has data
+            if len(dataset) == 0:
+                print(f"Warning: Skipping empty dataset {dataset.dataset_name}")
+                continue
+            datasets.append(dataset)
+            dataset_sampling_weights.append(weight)
+        
+        if len(datasets) == 0:
+            raise ValueError("No valid datasets found in the mixture. All datasets are empty.")
+        
+        self.datasets = datasets
+        self.balance_dataset_weights = balance_dataset_weights
+        self.balance_trajectory_weights = balance_trajectory_weights
+        self.seed = seed
+        self.mode = mode
+        self.data_cfg = kwargs["data_cfg"] if "data_cfg" in kwargs else None
+
+        # Set properties for sampling
+
+        # 1. Dataset lengths
+        self._dataset_lengths = np.array([len(dataset) for dataset in self.datasets])
+        print(f"Dataset lengths: {self._dataset_lengths}")
+        self._getitem_count = 0
+        # 2. Dataset sampling weights
+        self._dataset_sampling_weights = np.array(dataset_sampling_weights)
+        
+        if self.balance_dataset_weights:
+            self._dataset_sampling_weights *= self._dataset_lengths
+        
+        # Check for zero or negative weights before normalization
+        if np.any(self._dataset_sampling_weights <= 0):
+            print(f"Warning: Found zero or negative sampling weights: {self._dataset_sampling_weights}")
+            # Set minimum weight to prevent division issues
+            self._dataset_sampling_weights = np.maximum(self._dataset_sampling_weights, 1e-8)
+        
+        # Normalize weights
+        weights_sum = self._dataset_sampling_weights.sum()
+        if weights_sum == 0 or np.isnan(weights_sum):
+            print(f"Error: Invalid weights sum: {weights_sum}")
+            # Fallback to equal weights
+            self._dataset_sampling_weights = np.ones(len(self.datasets)) / len(self.datasets)
+            print(f"Fallback to equal weights")
+        else:
+            self._dataset_sampling_weights /= weights_sum
+
+        # 3. Trajectory sampling weights
+        self._trajectory_sampling_weights: list[np.ndarray] = []
+        for i, dataset in enumerate(self.datasets):
+            trajectory_sampling_weights = np.ones(len(dataset.trajectory_lengths))
+            if self.balance_trajectory_weights:
+                trajectory_sampling_weights *= dataset.trajectory_lengths
+            
+            # Check for zero or negative weights before normalization
+            if np.any(trajectory_sampling_weights <= 0):
+                print(f"Warning: Dataset {i} has zero or negative trajectory weights")
+                trajectory_sampling_weights = np.maximum(trajectory_sampling_weights, 1e-8)
+            
+            # Normalize weights
+            weights_sum = trajectory_sampling_weights.sum()
+            if weights_sum == 0 or np.isnan(weights_sum):
+                print(f"Error: Dataset {i} has invalid trajectory weights sum: {weights_sum}")
+                # Fallback to equal weights
+                trajectory_sampling_weights = np.ones(len(dataset.trajectory_lengths)) / len(dataset.trajectory_lengths)
+            else:
+                trajectory_sampling_weights /= weights_sum
+            
+            self._trajectory_sampling_weights.append(trajectory_sampling_weights)
+
+        # 4. Primary dataset indices
+        self._primary_dataset_indices = np.array(dataset_sampling_weights) == 1.0
+        if not np.any(self._primary_dataset_indices):
+            print(f"Warning: No dataset with weight 1.0 found. Original weights: {dataset_sampling_weights}")
+            # Fallback: use the dataset(s) with maximum weight as primary
+            max_weight = max(dataset_sampling_weights)
+            self._primary_dataset_indices = np.array(dataset_sampling_weights) == max_weight
+            print(f"Using datasets with maximum weight {max_weight} as primary: {self._primary_dataset_indices}")
+            
+        if not np.any(self._primary_dataset_indices):
+            # This should never happen, but just in case
+            print("Error: Still no primary dataset found. Using first dataset as primary.")
+            self._primary_dataset_indices = np.zeros(len(self.datasets), dtype=bool)
+            self._primary_dataset_indices[0] = True
+
+        self._sequential_step_sampling = True
+        self._shuffle_steps = True
+        self._action_balance_cfg = {}
+        self._action_balance_enabled = False
+        self._active_latency_filter: set[int] | None = None
+        self._active_latency_quotas: dict[int, float] | None = None
+        self._active_latency_quota_starts: dict[int, float] = {}
+        if self.data_cfg is not None:
+            seq_cfg = self.data_cfg.get("sequential_step_sampling", True)
+            self._sequential_step_sampling = _as_bool(seq_cfg, default=True)
+            self._shuffle_steps = _as_bool(self.data_cfg.get("shuffle", True), default=True)
+            self._action_balance_cfg = self.data_cfg.get("action_balance", {}) or {}
+            self._action_balance_enabled = _as_bool(self._action_balance_cfg.get("enabled", False), default=False)
+
+        if self._action_balance_enabled and not self._sequential_step_sampling:
+            raise ValueError("action_balance.enabled=true requires sequential_step_sampling=true.")
+        if self._action_balance_enabled:
+            strategy = str(self._action_balance_cfg.get("strategy", "balanced_epoch"))
+            if strategy != "balanced_epoch":
+                raise ValueError(f"Unsupported action_balance.strategy={strategy!r}; expected 'balanced_epoch'.")
+
+        self._step_order: np.ndarray = np.array([], dtype=np.int64)
+        self._flat_steps: list[tuple[int, int]] = []
+        self._flat_action_ids: np.ndarray = np.array([], dtype=np.int64)
+        self._noop_step_order: np.ndarray = np.array([], dtype=np.int64)
+        self._flap_step_order: np.ndarray = np.array([], dtype=np.int64)
+        self._other_step_order: np.ndarray = np.array([], dtype=np.int64)
+
+        # Persistent spawn workers share the epoch but retain their own tables and sample order.
+        self._shared_epoch = torch.zeros((), dtype=torch.int64, device="cpu").share_memory_()
+        # Set the epoch after initializing sequential sampling state.
+        self.set_epoch(0)
+
+        self.update_metadata(metadata_config)
+
+    @property
+    def dataset_lengths(self) -> np.ndarray:
+        """The lengths of each dataset."""
+        return self._dataset_lengths
+
+    @property
+    def dataset_sampling_weights(self) -> np.ndarray:
+        """The sampling weights for each dataset."""
+        return self._dataset_sampling_weights
+
+    @property
+    def trajectory_sampling_weights(self) -> list[np.ndarray]:
+        """The sampling weights for each trajectory in each dataset."""
+        return self._trajectory_sampling_weights
+
+    @property
+    def primary_dataset_indices(self) -> np.ndarray:
+        """The indices of the primary datasets."""
+        return self._primary_dataset_indices
+
+    def _replace_data_mixture(
+        self,
+        data_mixture: Sequence[tuple[LeRobotSingleDataset, float]],
+    ) -> None:
+        """Replace the active sources while retaining this dataset object's worker state."""
+        self.datasets = [dataset for dataset, _weight in data_mixture]
+        raw_sampling_weights = np.asarray(
+            [weight for _dataset, weight in data_mixture], dtype=np.float64
+        )
+        self._dataset_lengths = np.asarray(
+            [len(dataset) for dataset in self.datasets], dtype=np.int64
+        )
+        self._dataset_sampling_weights = raw_sampling_weights.copy()
+        if self.balance_dataset_weights:
+            self._dataset_sampling_weights *= self._dataset_lengths
+        self._dataset_sampling_weights /= self._dataset_sampling_weights.sum()
+
+        self._trajectory_sampling_weights = []
+        for dataset in self.datasets:
+            weights = np.ones(len(dataset.trajectory_lengths), dtype=np.float64)
+            if self.balance_trajectory_weights:
+                weights *= dataset.trajectory_lengths
+            weights /= weights.sum()
+            self._trajectory_sampling_weights.append(weights)
+
+        self._primary_dataset_indices = raw_sampling_weights == 1.0
+        self._rebuild_step_order(self.epoch)
+
+    def __str__(self) -> str:
+        dataset_descriptions = []
+        for dataset, weight in zip(self.datasets, self.dataset_sampling_weights):
+            dataset_description = {
+                "Dataset": str(dataset),
+                "Sampling weight": float(weight),
+            }
+            dataset_descriptions.append(dataset_description)
+        return json.dumps({"Mixture dataset": dataset_descriptions}, indent=2)
+
+    def set_epoch(self, epoch: int):
+        """Set the epoch for the dataset.
+
+        Args:
+            epoch (int): The epoch to set.
+        """
+        self.epoch = epoch
+        self._rebuild_step_order(epoch)
+        self._shared_epoch.fill_(epoch)
+
+    def set_active_latency_filter(self, latencies: Sequence[int] | None) -> None:
+        """Restrict sequential step sampling to the provided latency IDs."""
+        if not self._sequential_step_sampling:
+            raise ValueError(
+                "latency curriculum requires datasets.vla_data.sequential_step_sampling=true "
+                "so the active latency subset can be changed deterministically."
+            )
+        self._active_latency_filter = None if latencies is None else {int(value) for value in latencies}
+        self._active_latency_quotas = None
+        self._active_latency_quota_starts = {}
+        self._rebuild_step_order(getattr(self, "epoch", 0))
+
+    def set_latency_quota_filter(
+        self,
+        quotas: dict[int, float] | None,
+        *,
+        quota_starts: dict[int, float] | None = None,
+    ) -> None:
+        """Restrict sequential step sampling with per-latency pass quotas."""
+        if not self._sequential_step_sampling:
+            raise ValueError(
+                "latency quota curriculum requires datasets.vla_data.sequential_step_sampling=true "
+                "so the phase subset can be changed deterministically."
+            )
+        if quotas is None:
+            self._active_latency_quotas = None
+            self._active_latency_filter = None
+            self._active_latency_quota_starts = {}
+        else:
+            cleaned = {int(latency): float(quota) for latency, quota in quotas.items() if float(quota) > 0.0}
+            if not cleaned:
+                raise ValueError("latency quota curriculum requires at least one positive latency quota.")
+            self._active_latency_quotas = cleaned
+            starts = quota_starts or {}
+            self._active_latency_quota_starts = {
+                latency: float(starts.get(latency, 0.0))
+                for latency in cleaned
+            }
+            self._active_latency_filter = set(cleaned)
+        self._rebuild_step_order(getattr(self, "epoch", 0))
+
+    def get_latency_step_counts(self) -> dict[int, int]:
+        """Return the number of sequential training steps available per latency."""
+        counts: dict[int, int] = {}
+        for dataset in self.datasets:
+            for latency in dataset.get_latencies_for_all_steps():
+                latency = int(latency)
+                counts[latency] = counts.get(latency, 0) + 1
+        return counts
+
+    def _quota_step_indices(
+        self,
+        indices: list[int],
+        quota: float,
+        *,
+        latency: int,
+        pass_start: float = 0.0,
+    ) -> list[int]:
+        """Select deterministic rows from the pass interval [pass_start, pass_start + quota)."""
+        if quota <= 0.0 or not indices:
+            return []
+        selected: list[int] = []
+        cursor = float(pass_start)
+        remaining = float(quota)
+        num_indices = len(indices)
+        eps = 1e-9
+        while remaining > eps:
+            pass_idx = int(np.floor(cursor))
+            frac_start = cursor - pass_idx
+            if 1.0 - frac_start <= eps:
+                cursor = float(pass_idx + 1)
+                continue
+            segment = min(remaining, 1.0 - frac_start)
+            frac_end = frac_start + segment
+            pass_indices = list(indices)
+            seed = safe_hash((self.seed, "latency_quota_pass", latency, pass_idx)) & 0xFFFFFFFF
+            rng = np.random.default_rng(seed)
+            rng.shuffle(pass_indices)
+            start_idx = int(np.ceil(float(num_indices) * frac_start - eps))
+            end_idx = int(np.ceil(float(num_indices) * frac_end - eps))
+            selected.extend(pass_indices[start_idx:min(num_indices, end_idx)])
+            cursor += segment
+            remaining -= segment
+        return selected
+
+    def _rebuild_step_order(self, epoch: int) -> None:
+        """Build deterministic step order for sequential sampling."""
+        if not getattr(self, "_sequential_step_sampling", False):
+            return
+
+        flat_steps: list[tuple[int, int]] = []
+        if self._active_latency_quotas is not None:
+            candidate_steps_by_latency: dict[int, list[tuple[int, int]]] = {
+                int(latency): [] for latency in self._active_latency_quotas
+            }
+            for dataset_index, dataset in enumerate(self.datasets):
+                step_latencies = dataset.get_latencies_for_all_steps()
+                for step_index, latency in enumerate(step_latencies):
+                    latency = int(latency)
+                    if latency in candidate_steps_by_latency:
+                        candidate_steps_by_latency[latency].append((dataset_index, step_index))
+            for latency, quota in sorted(self._active_latency_quotas.items()):
+                candidates = candidate_steps_by_latency.get(int(latency), [])
+                flat_steps.extend(
+                    self._quota_step_indices(
+                        candidates,
+                        float(quota),
+                        latency=int(latency),
+                        pass_start=float(self._active_latency_quota_starts.get(int(latency), 0.0)),
+                    )
+                )
+        else:
+            for dataset_index, dataset in enumerate(self.datasets):
+                if self._active_latency_filter is None:
+                    flat_steps.extend(
+                        (dataset_index, step_index)
+                        for step_index in range(len(dataset.all_steps))
+                    )
+                else:
+                    step_latencies = dataset.get_latencies_for_all_steps()
+                    flat_steps.extend(
+                        (dataset_index, step_index)
+                        for step_index, latency in enumerate(step_latencies)
+                        if int(latency) in self._active_latency_filter
+                    )
+
+        self._flat_steps = flat_steps
+        self._flat_action_ids = np.array([], dtype=np.int64)
+        self._noop_step_order = np.array([], dtype=np.int64)
+        self._flap_step_order = np.array([], dtype=np.int64)
+        self._other_step_order = np.array([], dtype=np.int64)
+        if (self._active_latency_filter is not None or self._active_latency_quotas is not None) and len(flat_steps) == 0:
+            available = sorted(
+                {
+                    int(latency)
+                    for dataset in self.datasets
+                    for latency in dataset.get_latencies_for_all_steps()
+                }
+            )
+            raise ValueError(
+                f"Latency curriculum selected latencies={sorted(self._active_latency_filter)} "
+                f"but no matching training steps were found. Available latencies={available}."
+            )
+        if self._action_balance_enabled:
+            self._ensure_action_balance_indices()
+            self._step_order = self._build_action_balanced_step_order(epoch)
+        else:
+            self._step_order = np.arange(len(flat_steps), dtype=np.int64)
+        if self.mode == "train" and self._shuffle_steps and len(self._step_order) > 0:
+            seed = safe_hash((epoch, self.seed, "sequential_step_sampling", self._action_balance_enabled)) & 0xFFFFFFFF
+            rng = np.random.default_rng(seed)
+            rng.shuffle(self._step_order)
+
+    def _ensure_action_balance_indices(self) -> None:
+        """Build cached action-class indices aligned to ``self._flat_steps``."""
+        if len(self._flat_action_ids) == len(self._flat_steps):
+            return
+
+        action_key = str(self._action_balance_cfg.get("action_key", "action_id"))
+        action_ids_by_dataset = [
+            np.asarray(dataset.get_action_ids_for_all_steps(action_key), dtype=np.int64)
+            for dataset in self.datasets
+        ]
+
+        flat_action_ids = []
+        for dataset_index, step_index in self._flat_steps:
+            flat_action_ids.append(int(action_ids_by_dataset[dataset_index][step_index]))
+
+        self._flat_action_ids = np.asarray(flat_action_ids, dtype=np.int64)
+        noop_id = int(self._action_balance_cfg.get("noop_id", 0))
+        flap_id = int(self._action_balance_cfg.get("flap_id", 1))
+        all_indices = np.arange(len(self._flat_action_ids), dtype=np.int64)
+        self._noop_step_order = all_indices[self._flat_action_ids == noop_id]
+        self._flap_step_order = all_indices[self._flat_action_ids == flap_id]
+        self._other_step_order = all_indices[
+            (self._flat_action_ids != noop_id) & (self._flat_action_ids != flap_id)
+        ]
+
+        if len(self._noop_step_order) == 0 or len(self._flap_step_order) == 0:
+            raise ValueError(
+                "action_balance requires at least one noop and one flap sample, "
+                f"but got noop={len(self._noop_step_order)}, flap={len(self._flap_step_order)}."
+            )
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            target = float(self._action_balance_cfg.get("target_flap_fraction", 0.30))
+            desired_noop_count = int(np.ceil(len(self._flap_step_order) * (1.0 - target) / target))
+            selected_noop_count = min(len(self._noop_step_order), max(1, desired_noop_count))
+            selected_epoch_size = len(self._flap_step_order) + selected_noop_count + len(self._other_step_order)
+            print(
+                "Action-balanced epoch sampling enabled: "
+                f"noop={len(self._noop_step_order)}, flap={len(self._flap_step_order)}, "
+                f"other={len(self._other_step_order)}, target_flap_fraction={target:.3f}, "
+                f"selected_noop_per_epoch={selected_noop_count}, epoch_size={selected_epoch_size}"
+            )
+
+    def _build_action_balanced_step_order(self, epoch: int) -> np.ndarray:
+        """Build one balanced training epoch: all flaps plus a rotating noop slice."""
+        target_flap_fraction = float(self._action_balance_cfg.get("target_flap_fraction", 0.30))
+        if not 0.0 < target_flap_fraction < 1.0:
+            raise ValueError(
+                "action_balance.target_flap_fraction must be in (0, 1), "
+                f"got {target_flap_fraction}."
+            )
+
+        flap_indices = self._flap_step_order
+        noop_indices = self._noop_step_order
+        desired_noop_count = int(np.ceil(len(flap_indices) * (1.0 - target_flap_fraction) / target_flap_fraction))
+        noop_count = min(len(noop_indices), max(1, desired_noop_count))
+
+        noop_start = (int(epoch) * noop_count) % len(noop_indices)
+        if noop_start + noop_count <= len(noop_indices):
+            selected_noops = noop_indices[noop_start : noop_start + noop_count]
+        else:
+            wrap_count = (noop_start + noop_count) - len(noop_indices)
+            selected_noops = np.concatenate((noop_indices[noop_start:], noop_indices[:wrap_count]))
+
+        return np.concatenate((flap_indices, selected_noops, self._other_step_order)).astype(np.int64, copy=False)
+
+    def sample_step(self, index: int) -> tuple[LeRobotSingleDataset, int, int]:
+        """Sample a single step from the dataset."""
+        # return self.sampled_steps[index]
+
+        epoch = self._shared_epoch.item()
+        if self.epoch != epoch:
+            self.epoch = epoch
+            self._rebuild_step_order(epoch)
+
+        if self._sequential_step_sampling:
+            if len(self._step_order) == 0:
+                raise ValueError("Cannot sample from an empty sequential step order.")
+            ordered_index = int(self._step_order[index % len(self._step_order)])
+            dataset_index, step_index = self._flat_steps[ordered_index]
+            dataset = self.datasets[dataset_index]
+            trajectory_id, base_index = dataset.all_steps[step_index]
+            return dataset, int(trajectory_id), int(base_index)
+
+        # Set seed
+        seed = index if self.mode != "train" else safe_hash((self.epoch, index, self.seed))
+        rng = np.random.default_rng(seed)
+
+        # Sample dataset
+        dataset_index = rng.choice(len(self.datasets), p=self.dataset_sampling_weights)
+        dataset = self.datasets[dataset_index]
+
+        # Sample trajectory
+        trajectory_index = rng.choice(
+            len(dataset.trajectory_ids), p=self.trajectory_sampling_weights[dataset_index]
+        )
+        trajectory_id = dataset.trajectory_ids[trajectory_index]
+
+        # Sample step
+        base_index = rng.choice(dataset.trajectory_lengths[trajectory_index])
+        return dataset, trajectory_id, base_index
+
+    
+
+    def __getitem__(self, index: int) -> dict:
+        """Get the data for a single trajectory and start index.
+
+        Args:
+            index (int): The index of the trajectory to get.
+
+        Returns:
+            dict: The data for the trajectory and start index.
+        """
+        self._getitem_count += 1
+        if self._getitem_count % 1000 == 0:
+            gc.collect()
+
+        max_retries = 10
+        last_exception = None
+
+        for attempt in range(max_retries):
+            try:
+                sample_tries = 0
+                max_sample_tries = 200
+                while True:
+                    sample_tries += 1
+                    if sample_tries > max_sample_tries:
+                        raise RuntimeError(
+                            f"Unable to sample a valid item after {max_sample_tries} attempts. "
+                            f"dataset={self.datasets[0].dataset_name if len(self.datasets)>0 else 'unknown'}"
+                        )
+
+                    dataset, trajectory_id, step = self.sample_step(index)
+                    # If dataset has no physical videos (e.g., image frames in parquet
+                    # for VLA-Arena), do not gate sampling on mp4 existence.
+                    total_videos = int(dataset.lerobot_info_meta.get("total_videos", 0))
+                    if total_videos == 0:
+                        break
+
+                    key = dataset.modality_keys["video"][0].replace("video.", "")
+                    video_path = dataset.get_video_path(trajectory_id, key)
+                    if os.path.exists(video_path):
+                        break
+                    index = random.randint(0, len(self) - 1)
+                    
+                raw_data = dataset.get_step_data(trajectory_id, step)
+                data = dataset.transforms(raw_data)
+                sample = dataset._pack_sample(data, trajectory_id=trajectory_id, base_index=step)
+                dataset._attach_rl_games_metadata(sample, step)
+                
+                return sample
+                
+            except Exception as e:
+                last_exception = e
+                if attempt < max_retries - 1:
+                    # Log the error but continue trying
+                    print(f"Attempt {attempt + 1}/{max_retries} failed for index {index}: {e}")
+                    print(f"Retrying with new sample...")
+                    # For retry, we can use a slightly different index to get a new sample
+                    # This helps avoid getting stuck on the same problematic sample
+                    index = random.randint(0, len(self) - 1)
+                else:
+                    # All retries exhausted
+                    print(f"All {max_retries} attempts failed for index {index}")
+                    print(f"Last error: {last_exception}")
+                    # Return a dummy sample or re-raise the exception
+                    raise last_exception
+
+    def __len__(self) -> int:
+        """Get the length of a single epoch in the mixture.
+
+        Returns:
+            int: The length of a single epoch in the mixture.
+        """
+        if self._sequential_step_sampling:
+            return len(self._step_order)
+
+        # Check for potential issues
+        if len(self.datasets) == 0:
+            return 0
+            
+        # Check if any dataset lengths are 0 or NaN
+        if np.any(self.dataset_lengths == 0) or np.any(np.isnan(self.dataset_lengths)):
+            print(f"Warning: Found zero or NaN dataset lengths: {self.dataset_lengths}")
+            # Filter out zero/NaN length datasets
+            valid_indices = (self.dataset_lengths > 0) & (~np.isnan(self.dataset_lengths))
+            if not np.any(valid_indices):
+                print("Error: All datasets have zero or NaN length")
+                return 0
+        else:
+            valid_indices = np.ones(len(self.datasets), dtype=bool)
+        
+        # Check if any sampling weights are 0 or NaN
+        if np.any(self.dataset_sampling_weights == 0) or np.any(np.isnan(self.dataset_sampling_weights)):
+            print(f"Warning: Found zero or NaN sampling weights: {self.dataset_sampling_weights}")
+            # Use only valid weights
+            valid_weights = (self.dataset_sampling_weights > 0) & (~np.isnan(self.dataset_sampling_weights))
+            valid_indices = valid_indices & valid_weights
+            if not np.any(valid_indices):
+                print("Error: All sampling weights are zero or NaN")
+                return 0
+        
+        # Check primary dataset indices
+        primary_and_valid = self.primary_dataset_indices & valid_indices
+        if not np.any(primary_and_valid):
+            print(f"Warning: No valid primary datasets found. Primary indices: {self.primary_dataset_indices}, Valid indices: {valid_indices}")
+            # Fallback: use the largest valid dataset
+            if np.any(valid_indices):
+                max_length = self.dataset_lengths[valid_indices].max()
+                print(f"Fallback: Using maximum dataset length: {max_length}")
+                return int(max_length)
+            else:
+                return 0
+        
+        # Calculate the ratio and get max
+        ratios = (self.dataset_lengths / self.dataset_sampling_weights)[primary_and_valid]
+        
+        # Check for NaN or inf in ratios
+        if np.any(np.isnan(ratios)) or np.any(np.isinf(ratios)):
+            print(f"Warning: Found NaN or inf in ratios: {ratios}")
+            print(f"Dataset lengths: {self.dataset_lengths[primary_and_valid]}")
+            print(f"Sampling weights: {self.dataset_sampling_weights[primary_and_valid]}")
+            # Filter out invalid ratios
+            valid_ratios = ratios[~np.isnan(ratios) & ~np.isinf(ratios)]
+            if len(valid_ratios) == 0:
+                print("Error: All ratios are NaN or inf")
+                return 0
+            max_ratio = valid_ratios.max()
+        else:
+            max_ratio = ratios.max()
+        
+        result = int(max_ratio)
+        if result == 0:
+            print(f"Warning: Dataset mixture length is 0")
+        return result
+
+    @staticmethod
+    def compute_overall_statistics(
+        per_task_stats: list[dict[str, dict[str, list[float] | np.ndarray]]],
+        dataset_sampling_weights: list[float] | np.ndarray,
+        percentile_mixing_method: str = "weighted_average",
+    ) -> dict[str, dict[str, list[float]]]:
+        """
+        Computes overall statistics from per-task statistics using dataset sample weights.
+
+        Args:
+            per_task_stats: List of per-task statistics.
+            Example format of one element in the per-task statistics list:
+                {
+                    "state.gripper": {
+                        "min": [...],
+                        "max": [...],
+                        "mean": [...],
+                        "std": [...],
+                        "q01": [...],
+                        "q99": [...],
+                    },
+                    ...
+                }
+            dataset_sampling_weights: List of sample weights for each task.
+            percentile_mixing_method: The method to mix the percentiles, either "weighted_average" or "weighted_std".
+
+        Returns:
+            A dict of overall statistics per modality.
+        """
+        # Normalize the sample weights to sum to 1
+        dataset_sampling_weights = np.array(dataset_sampling_weights)
+        normalized_weights = dataset_sampling_weights / dataset_sampling_weights.sum()
+
+        # Initialize overall statistics dict
+        overall_stats: dict[str, dict[str, list[float]]] = {}
+
+        # Get the list of modality keys
+        modality_keys = per_task_stats[0].keys()
+
+        for modality in modality_keys:
+            # Number of dimensions (assuming consistent across tasks)
+            num_dims = len(per_task_stats[0][modality]["mean"])
+
+            # Initialize accumulators for means and variances
+            weighted_means = np.zeros(num_dims)
+            weighted_squares = np.zeros(num_dims)
+
+            # Collect min, max, q01, q99 from all tasks
+            min_list = []
+            max_list = []
+            q01_list = []
+            q99_list = []
+
+            for task_idx, task_stats in enumerate(per_task_stats):
+                w_i = normalized_weights[task_idx]
+                stats = task_stats[modality]
+                means = np.array(stats["mean"])
+                stds = np.array(stats["std"])
+
+                # Update weighted sums for mean and variance
+                weighted_means += w_i * means
+                weighted_squares += w_i * (stds**2 + means**2)
+
+                # Collect min, max, q01, q99
+                min_list.append(stats["min"])
+                max_list.append(stats["max"])
+                q01_list.append(stats["q01"])
+                q99_list.append(stats["q99"])
+
+            # Compute overall mean
+            overall_mean = weighted_means.tolist()
+
+            # Compute overall variance and std deviation
+            overall_variance = weighted_squares - weighted_means**2
+            overall_std = np.sqrt(overall_variance).tolist()
+
+            # Compute overall min and max per dimension
+            overall_min = np.min(np.array(min_list), axis=0).tolist()
+            overall_max = np.max(np.array(max_list), axis=0).tolist()
+
+            # Compute overall q01 and q99 per dimension
+            # Use weighted average of per-task quantiles
+            q01_array = np.array(q01_list)
+            q99_array = np.array(q99_list)
+            if percentile_mixing_method == "weighted_average":
+                weighted_q01 = np.average(q01_array, axis=0, weights=normalized_weights).tolist()
+                weighted_q99 = np.average(q99_array, axis=0, weights=normalized_weights).tolist()
+                # std_q01 = np.std(q01_array, axis=0).tolist()
+                # std_q99 = np.std(q99_array, axis=0).tolist()
+                # print(modality)
+                # print(f"{std_q01=}, {std_q99=}")
+                # print(f"{weighted_q01=}, {weighted_q99=}")
+            elif percentile_mixing_method == "min_max":
+                weighted_q01 = np.min(q01_array, axis=0).tolist()
+                weighted_q99 = np.max(q99_array, axis=0).tolist()
+            else:
+                raise ValueError(f"Invalid percentile mixing method: {percentile_mixing_method}")
+
+            # Store the overall statistics for the modality
+            overall_stats[modality] = {
+                "min": overall_min,
+                "max": overall_max,
+                "mean": overall_mean,
+                "std": overall_std,
+                "q01": weighted_q01,
+                "q99": weighted_q99,
+            }
+
+        return overall_stats
+
+    @staticmethod
+    def merge_metadata(
+        metadatas: list[DatasetMetadata],
+        dataset_sampling_weights: list[float],
+        percentile_mixing_method: str,
+    ) -> DatasetMetadata:
+        """Merge multiple metadata into one."""
+        # Convert to dicts
+        metadata_dicts = [metadata.model_dump(mode="json") for metadata in metadatas]
+        # Create a new metadata dict
+        merged_metadata = {}
+
+        # Check all metadata have the same embodiment tag
+        assert all(
+            metadata.embodiment_tag == metadatas[0].embodiment_tag for metadata in metadatas
+        ), "All metadata must have the same embodiment tag"
+        merged_metadata["embodiment_tag"] = metadatas[0].embodiment_tag
+
+        # Merge the dataset statistics
+        dataset_statistics = {}
+        dataset_statistics["state"] = LeRobotMixtureDataset.compute_overall_statistics(
+            per_task_stats=[m["statistics"]["state"] for m in metadata_dicts],
+            dataset_sampling_weights=dataset_sampling_weights,
+            percentile_mixing_method=percentile_mixing_method,
+        )
+        dataset_statistics["action"] = LeRobotMixtureDataset.compute_overall_statistics(
+            per_task_stats=[m["statistics"]["action"] for m in metadata_dicts],
+            dataset_sampling_weights=dataset_sampling_weights,
+            percentile_mixing_method=percentile_mixing_method,
+        )
+        merged_metadata["statistics"] = dataset_statistics
+
+        # Merge the modality configs
+        modality_configs = defaultdict(list)
+        for metadata in metadata_dicts:
+            for modality, configs in metadata["modalities"].items():
+                modality_configs[modality].append(configs)
+        merged_metadata["modalities"] = {}
+        for modality, configs in modality_configs.items():
+            comparable = {
+                json.dumps(_strip_video_capture_metadata(config), sort_keys=True)
+                for config in configs
+            }
+            assert (
+                len(comparable) == 1
+            ), f"Multiple modality configs for modality {modality}: {configs}"
+            merged_metadata["modalities"][modality] = configs[0]
+
+        return DatasetMetadata.model_validate(merged_metadata)
+
+    def update_metadata(self, metadata_config: dict, cached_statistics_path: Path | str | None = None) -> None:
+        """
+        Merge multiple metadatas into one and set the transforms with the merged metadata.
+
+        Args:
+            metadata_config (dict): Configuration for the metadata.
+                "percentile_mixing_method": The method to mix the percentiles, either "weighted_average" or "min_max".
+                    weighted_average: Use the weighted average of the percentiles using the weight used in sampling the datasets.
+                    min_max: Use the min of the 1st percentile and max of the 99th percentile.
+        """
+        # If cached path is provided, try to load and apply
+        if cached_statistics_path is not None:
+            try:
+                cached_stats = self.load_merged_statistics(cached_statistics_path)
+                self.apply_cached_statistics(cached_stats)
+                return
+            except (FileNotFoundError, KeyError, ValidationError) as e:
+                print(f"Failed to load cached statistics: {e}")
+                print("Falling back to computing statistics from scratch...")
+
+        self.tag = EmbodimentTag.NEW_EMBODIMENT.value
+        self.merged_metadata: dict[str, DatasetMetadata] = {}
+        # Group metadata by tag
+        all_metadatas: dict[str, list[DatasetMetadata]] = {}
+        for dataset in self.datasets:
+            if dataset.tag not in all_metadatas:
+                all_metadatas[dataset.tag] = []
+            all_metadatas[dataset.tag].append(dataset.metadata)
+        for tag, metadatas in all_metadatas.items():
+            self.merged_metadata[tag] = self.merge_metadata(
+                metadatas=metadatas,
+                dataset_sampling_weights=self.dataset_sampling_weights.tolist(),
+                percentile_mixing_method=metadata_config["percentile_mixing_method"],
+            )
+        for dataset in self.datasets:
+            dataset.set_transforms_metadata(self.merged_metadata[dataset.tag])
+
+    def save_dataset_statistics(self, save_path: Path | str, format: str = "json") -> None:
+        """
+        Save merged dataset statistics to specified path in the required format.
+        Only includes statistics for keys that are actually used in the datasets.
+        Gripper-related keys will be placed at the end.
+        
+        Args:
+            save_path (Path | str): Path to save the statistics file
+            format (str): Save format, currently only supports "json"
+        """
+        save_path = Path(save_path)
+        save_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        # Build the data structure to save
+        statistics_data = {}
+        
+        # Collect actually used keys from all datasets
+        all_used_action_keys = []
+        all_used_state_keys = []
+        
+        for dataset in self.datasets:
+            used_action_keys, used_state_keys = get_used_modality_keys(dataset.modality_keys)
+            for used_action_key in used_action_keys:
+                if used_action_key not in all_used_action_keys:
+                    all_used_action_keys.append(used_action_key)
+            for used_state_key in used_state_keys:
+                if used_state_key not in all_used_state_keys:
+                    all_used_state_keys.append(used_state_key)
+        
+        # Organize statistics by tag
+        for tag, merged_metadata in self.merged_metadata.items():
+            tag_stats = {}
+            
+            # Process action statistics
+            if hasattr(merged_metadata.statistics, 'action') and merged_metadata.statistics.action:
+                action_stats = merged_metadata.statistics.action
+                
+                # Filter and reorder keys - iterate in all_used_action_keys order
+                non_gripper_keys = []
+                gripper_keys = []
+                
+                for key in all_used_action_keys:
+                    if key in action_stats:
+                        non_gripper_keys.append(key)
+                
+                reordered_keys = non_gripper_keys + gripper_keys
+                
+                filtered_action_stats = {}
+                for key in reordered_keys:
+                    filtered_action_stats[key] = action_stats[key]
+                
+                if filtered_action_stats:
+                    combined_action_stats = combine_modality_stats(filtered_action_stats)
+                    
+                    mask = generate_action_mask_for_used_keys(
+                        merged_metadata.modalities.action, filtered_action_stats.keys()
+                    )
+                    combined_action_stats["mask"] = mask
+                    
+                    tag_stats["action"] = combined_action_stats
+            
+            # Process state statistics
+            if hasattr(merged_metadata.statistics, 'state') and merged_metadata.statistics.state:
+                state_stats = merged_metadata.statistics.state
+                
+                # Filter and reorder keys - iterate in all_used_state_keys order
+                # Filter and reorder keys - iterate in all_used_state_keys order
+                non_gripper_keys = []
+                gripper_keys = []
+                
+                for key in all_used_state_keys:
+                    if key in state_stats:
+                        non_gripper_keys.append(key)
+                
+                reordered_keys = non_gripper_keys + gripper_keys
+                
+                filtered_state_stats = {}
+                for key in reordered_keys:
+                    filtered_state_stats[key] = state_stats[key]
+                
+                if filtered_state_stats:
+                    combined_state_stats = combine_modality_stats(filtered_state_stats)
+                    tag_stats["state"] = combined_state_stats
+            
+            # Add dataset counts
+            tag_stats.update(self._get_dataset_counts(tag))
+            
+            statistics_data[tag] = tag_stats
+        
+        # Save file
+        if format.lower() == "json":
+            if not str(save_path).endswith('.json'):
+                save_path = save_path.with_suffix('.json')
+            with open(save_path, 'w', encoding='utf-8') as f:
+                json.dump(statistics_data, f, indent=2, ensure_ascii=False)
+        else:
+            raise ValueError(f"Unsupported format: {format}. Currently only 'json' is supported.")
+        
+        print(f"Merged dataset statistics saved to: {save_path}")
+        print(f"Used action keys (reordered): {list(all_used_action_keys)}")
+        print(f"Used state keys (reordered): {list(all_used_state_keys)}")
+
+
+    def _combine_modality_stats(self, modality_stats: dict) -> dict:
+        """Backward compatibility wrapper."""
+        return combine_modality_stats(modality_stats)
+
+    def _generate_action_mask_for_used_keys(self, action_modalities: dict, used_action_keys_ordered) -> list[bool]:
+        """Backward compatibility wrapper."""
+        return generate_action_mask_for_used_keys(action_modalities, used_action_keys_ordered)
+
+    def _get_dataset_counts(self, tag: str) -> dict:
+        """
+        Get dataset count information for specified tag.
+        
+        Args:
+            tag (str): embodiment tag
+            
+        Returns:
+            dict: Dictionary containing num_transitions and num_trajectories
+        """
+        num_transitions = 0
+        num_trajectories = 0
+        
+        # Count dataset information belonging to this tag
+        for dataset in self.datasets:
+            if dataset.tag == tag:
+                num_transitions += len(dataset)
+                num_trajectories += len(dataset.trajectory_ids)
+        
+        return {
+            "num_transitions": num_transitions,
+            "num_trajectories": num_trajectories
+        }
+
+    @classmethod
+    def load_merged_statistics(cls, load_path: Path | str) -> dict:
+        """
+        Load merged dataset statistics from file.
+        
+        Args:
+            load_path (Path | str): Path to the statistics file
+            
+        Returns:
+            dict: Dictionary containing merged statistics
+        """
+        load_path = Path(load_path)
+        if not load_path.exists():
+            raise FileNotFoundError(f"Statistics file not found: {load_path}")
+        
+        if load_path.suffix.lower() == '.json':
+            with open(load_path, 'r', encoding='utf-8') as f:
+                return json.load(f)
+        elif load_path.suffix.lower() == '.pkl':
+            import pickle
+            with open(load_path, 'rb') as f:
+                return pickle.load(f)
+        else:
+            raise ValueError(f"Unsupported file format: {load_path.suffix}")
+
+    def apply_cached_statistics(self, cached_statistics: dict) -> None:
+        """
+        Apply cached statistics to avoid recomputation.
+        
+        Args:
+            cached_statistics (dict): Statistics loaded from file
+        """
+        # Validate that cached statistics match current datasets
+        if "metadata" in cached_statistics:
+            cached_dataset_names = set(cached_statistics["metadata"]["dataset_names"])
+            current_dataset_names = set(dataset.dataset_name for dataset in self.datasets)
+            
+            if cached_dataset_names != current_dataset_names:
+                print("Warning: Cached statistics dataset names don't match current datasets.")
+                print(f"Cached: {cached_dataset_names}")
+                print(f"Current: {current_dataset_names}")
+                return
+        
+        # Apply cached statistics
+        self.merged_metadata = {}
+        for tag, stats_data in cached_statistics.items():
+            if tag == "metadata":  # Skip metadata field
+                continue
+                
+            # Convert back to DatasetMetadata format
+            metadata_dict = {
+                "embodiment_tag": tag,
+                "statistics": {
+                    "action": {},
+                    "state": {}
+                },
+                "modalities": {}
+            }
+            
+            # Convert action statistics back
+            if "action" in stats_data:
+                action_data = stats_data["action"]
+                # This is simplified - you may need to split back to sub-keys
+                metadata_dict["statistics"]["action"] = action_data
+            
+            # Convert state statistics back
+            if "state" in stats_data:
+                state_data = stats_data["state"]
+                metadata_dict["statistics"]["state"] = state_data
+            
+            self.merged_metadata[tag] = DatasetMetadata.model_validate(metadata_dict)
+        
+        # Update transforms metadata for each dataset
+        for dataset in self.datasets:
+            if dataset.tag in self.merged_metadata:
+                dataset.set_transforms_metadata(self.merged_metadata[dataset.tag])
+        
+        print(f"Applied cached statistics for {len(self.merged_metadata)} embodiment tags.")
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/gym-language-contract.patch b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/gym-language-contract.patch
new file mode 100644
index 0000000000000000000000000000000000000000..f9f3acfafddcff18c8a6120610a410e67e346605
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/gym-language-contract.patch
@@ -0,0 +1,17 @@
+diff --git a/examples/rl_games/train_files/data_registry/data_config.py b/examples/rl_games/train_files/data_registry/data_config.py
+index 2ef8b39..6d76aeb 100644
+--- a/examples/rl_games/train_files/data_registry/data_config.py
++++ b/examples/rl_games/train_files/data_registry/data_config.py
+@@ -53,10 +53,10 @@ class AtlantisDataConfig(FlappyDataConfig):
+ 
+ 
+ class GymnasiumDataConfig(FlappyDataConfig):
+-    pass
++    language_keys = ["annotation.human.task_description"]
+ 
+ 
+-class GymnasiumNativeDataConfig(FlappyDataConfig):
++class GymnasiumNativeDataConfig(GymnasiumDataConfig):
+     """Configure native Gymnasium state and action vectors."""
+ 
+     state_keys = ["state.native"]
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/loader-spawn-cache.patch b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/loader-spawn-cache.patch
new file mode 100644
index 0000000000000000000000000000000000000000..9de0849e5e86a244a19a4e159446a4e63152e3b0
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/loader-spawn-cache.patch
@@ -0,0 +1,18 @@
+diff --git a/starVLA/dataloader/gr00t_lerobot/datasets.py b/starVLA/dataloader/gr00t_lerobot/datasets.py
+index 17ce5f1..2346230 100644
+--- a/starVLA/dataloader/gr00t_lerobot/datasets.py
++++ b/starVLA/dataloader/gr00t_lerobot/datasets.py
+@@ -667,6 +667,13 @@ class LeRobotSingleDataset(Dataset):
+         # Check if the dataset is valid
+         self._check_integrity()
+ 
++    def __getstate__(self):
++        state = self.__dict__.copy()
++        state["_parquet_cache"] = {}
++        state["curr_traj_data"] = None
++        state["curr_traj_id"] = None
++        return state
++
+     @property
+     def dataset_path(self) -> Path:
+         """The path to the dataset that contains the METADATA_FILENAME file."""
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/loader-spawn-test.patch b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/loader-spawn-test.patch
new file mode 100644
index 0000000000000000000000000000000000000000..10de677ea1f7ca7bed563f47049ad703de930a1a
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/loader-spawn-test.patch
@@ -0,0 +1,31 @@
+diff --git a/tests/data/test_starvla_tasks.py b/tests/data/test_starvla_tasks.py
+index 1403fa13..62468aa6 100644
+--- a/tests/data/test_starvla_tasks.py
++++ b/tests/data/test_starvla_tasks.py
+@@ -395,3 +395,26 @@ def test_balance_masked_l1_uses_equal_sample_weights_and_no_padding_gradient(mon
+     all_valid = torch.ones((2, 8), dtype=torch.bool)
+     torch.testing.assert_close(model._compute_action_loss(pred, target, action_loss_mask=all_valid),
+                                (pred - target).abs().mean())
++
++
++def test_lerobot_spawn_reload_preserves_data_without_copying_parent_cache(tmp_path, monkeypatch):
++    import pickle
++    pytest.importorskip("accelerate")
++    monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "third_party/starVLA"))
++    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotSingleDataset
++
++    path = tmp_path / "episode.parquet"
++    table = pd.DataFrame({"image": [b"rgb" * (1024 * 1024)], "action": [1]})
++    table.to_parquet(path)
++    parent = object.__new__(LeRobotSingleDataset)
++    parent._parquet_cache = {path: table}
++    parent.curr_traj_data = table
++    parent.curr_traj_id = 0
++    parent._all_steps = [(0, 0)]
++
++    payload = pickle.dumps(parent)
++    assert len(payload) < 64 * 1024
++    worker = pickle.loads(payload)
++    assert worker._all_steps == parent._all_steps
++    pd.testing.assert_frame_equal(worker._read_trajectory_parquet(path), table)
++    pd.testing.assert_frame_equal(parent.curr_traj_data, table)
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/mean-profile-preparation.patch b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/mean-profile-preparation.patch
new file mode 100644
index 0000000000000000000000000000000000000000..fecf07a56192eac5998a58823884494aa8d34fb8
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/mean-profile-preparation.patch
@@ -0,0 +1,147 @@
+diff --git a/latency_bench/data/sf_teacher_rollout.py b/latency_bench/data/sf_teacher_rollout.py
+index ef45d95b..af0a9fa9 100644
+--- a/latency_bench/data/sf_teacher_rollout.py
++++ b/latency_bench/data/sf_teacher_rollout.py
+@@ -964,6 +964,11 @@ def _replay_accepted_specs(
+         row_buffer_size=int(row_buffer_size),
+         export_image_size=export_image_size,
+     )
++    stream_current_images = (
++        image_transform_config.normalized_transform() == "raw_rgb"
++        and int(context_window) <= 1
++        and filter_preset == "none"
++    )
+     image_resolution: list[int] | None = None
+ 
+     def assign_specs_to_slots(slots: list[int]) -> dict[int, int]:
+@@ -1051,7 +1056,10 @@ def _replay_accepted_specs(
+                         next_raw_frame=next_decision_step * obs_stride_raw_frames,
+                         max_steps_per_episode=int(max_steps_per_episode),
+                     )
+-                    active_episode_frames[slot].append(np.array(image_frame, copy=True))
++                    if stream_current_images:
++                        image_writer.write(image_frame, image_path)
++                    else:
++                        active_episode_frames[slot].append(np.array(image_frame, copy=True))
+                     row = teacher_dataset_row(
+                         env_name=env_name,
+                         spec=spec,
+@@ -1104,7 +1112,9 @@ def _replay_accepted_specs(
+                             )
+                         else:
+                             shard_rows = [dict(row) for row in final_rows]
+-                        if _is_demon_ghost_trail:
++                        if stream_current_images:
++                            image_writer.flush()
++                        elif _is_demon_ghost_trail:
+                             # Frames are pre-built ghost trail images; map directly.
+                             ep_steps = [int(row["decision_step"]) for row in final_rows]
+                             frames_by_decision_step = {
+@@ -1118,12 +1128,13 @@ def _replay_accepted_specs(
+                                 frames=active_episode_frames[slot],
+                                 config=image_transform_config,
+                             )
+-                        shard_rows = _write_filtered_episode_images(
+-                            image_writer,
+-                            shard_rows,
+-                            frames_by_decision_step,
+-                            context_window=int(context_window),
+-                        )
++                        if not stream_current_images:
++                            shard_rows = _write_filtered_episode_images(
++                                image_writer,
++                                shard_rows,
++                                frames_by_decision_step,
++                                context_window=int(context_window),
++                            )
+                         write_episode_row_shard(output_dir, shard_rows)
+                         progress.update(1)
+                         replayed += 1
+diff --git a/tests/data/test_sf_teacher_dataset_export.py b/tests/data/test_sf_teacher_dataset_export.py
+index e6b25233..7d5537ac 100644
+--- a/tests/data/test_sf_teacher_dataset_export.py
++++ b/tests/data/test_sf_teacher_dataset_export.py
+@@ -1273,7 +1273,8 @@ def test_collect_teacher_dataset_saves_demon_attack_latest_raw_rgb_image(monkeyp
+ 
+ 
+ @pytest.mark.parametrize("export_history", [False, True])
+-def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch, tmp_path, export_history):
++@pytest.mark.parametrize("filter_preset", ["clean_v1", "none"])
++def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch, tmp_path, export_history, filter_preset):
+     import torch
+     from gymnasium.spaces import Discrete
+     from PIL import Image
+@@ -1307,6 +1308,9 @@ def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch,
+     class FakeEnv:
+         action_space = Discrete(2)
+ 
++        def __init__(self, exporting):
++            self.exporting = exporting
++
+         def reset(self, seed=None):
+             self.steps = 0
+             obs = {"obs": np.zeros((1, 12, 2, 2), dtype=np.uint8)}
+@@ -1323,6 +1327,9 @@ def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch,
+ 
+         def step(self, action):
+             self.steps += 1
++            if self.exporting and self.steps == 2 and filter_preset == "none":
++                image_path = output_dir / "images/train/episode_000000/decision_step_000000.png"
++                assert image_path.is_file(), "First frame must be on disk before the episode ends"
+             obs = {"obs": np.ones((1, 12, 2, 2), dtype=np.uint8)}
+             next_stack = np.stack([np.full((6, 5, 3), 50 + idx, dtype=np.uint8) for idx in range(4)], axis=0)
+             return (
+@@ -1354,7 +1361,7 @@ def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch,
+         exporter,
+         "build_env_and_policy",
+         lambda module, cfg, config, num_envs: {
+-            "envs": [FakeEnv()],
++            "envs": [FakeEnv(cfg.export_env_raw_rgb_frames)],
+             "env_info": object(),
+             "actor_critic": FakeActorCritic(),
+             "device": torch.device("cpu"),
+@@ -1382,6 +1389,7 @@ def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch,
+         export_image_size=None,
+         image_writer_workers=0,
+         row_buffer_size=1,
++        filter_preset=filter_preset,
+     )
+     train_rows = _read_shard_rows(output_dir, "train")
+     val_rows = _read_shard_rows(output_dir, "val")
+diff --git a/third_party/MIKASA-Robo b/third_party/MIKASA-Robo
+--- a/third_party/MIKASA-Robo
++++ b/third_party/MIKASA-Robo
+@@ -1 +1 @@
+-Subproject commit 017490bbb787c6efa8c025f9558e1a691965c0e0
++Subproject commit 017490bbb787c6efa8c025f9558e1a691965c0e0-dirty
+diff --git a/third_party/starVLA b/third_party/starVLA
+--- a/third_party/starVLA
++++ b/third_party/starVLA
+@@ -1 +1 @@
+-Subproject commit ccca13c5177fe3d3c884b6e2de4965d916016649
++Subproject commit ccca13c5177fe3d3c884b6e2de4965d916016649-dirty
+diff --git a/training/gymnasium_sf/integration.py b/training/gymnasium_sf/integration.py
+index 61c65ae9..874f655f 100644
+--- a/training/gymnasium_sf/integration.py
++++ b/training/gymnasium_sf/integration.py
+@@ -64,7 +64,6 @@ from training.common.sf_raw_reward_stats_compat import (
+ )
+ from training.common.utils import env_instance_seed, seed_everything, str2bool
+ from training.common.wandb_compat import patch_sample_factory_wandb
+-from sample_factory.algo.mc_ppo.env import EpisodeEnv
+ 
+ 
+ SF_ENV_NAME = "latency_gymnasium"
+@@ -385,7 +384,11 @@ def make_gymnasium_sf_env(
+         env_id=int(env_config.env_id),
+         issued_action_labels=action_labels,
+     )
+-    return EpisodeEnv(env, cfg.ppo) if cfg.algo == "PPO" else env
++    if cfg.algo == "PPO":
++        from sample_factory.algo.mc_ppo.env import EpisodeEnv
++
++        return EpisodeEnv(env, cfg.ppo)
++    return env
+ 
+ 
+ def register_gymnasium_envs() -> None:
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/train_starvla.py b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/train_starvla.py
new file mode 100644
index 0000000000000000000000000000000000000000..e5ee9a5128b3b576244b35c527e947f06e530c34
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/train_starvla.py
@@ -0,0 +1,2692 @@
+# Copyright 2025 starVLA community. All rights reserved.
+# Licensed under the MIT License, Version 1.0 (the "License");
+# Implemented by [Jinhui YE / HKUST University] in [2025].
+
+"""
+StarVLA’s trainer is built directly on native PyTorch + Accelerate + DeepSpeed, keeping the loop explicit and easy to hack.
+Conventions:
+1. Store runtime state in dicts where possible (simplifies data info, procesing info, config, etc).
+2. Use multiple dataloaders to adapt heterogeneous data types / task mixtures.
+3. Put each training strategy in its own `trainer_*.py` file (avoid large if‑else chains).
+"""
+
+# Standard Library
+import argparse
+import json
+import logging
+import math
+import os
+import re
+import shutil
+import signal
+import sys
+import time
+from pathlib import Path
+from typing import Any, Dict, Tuple
+
+# Third-Party Libraries
+import numpy as np
+import torch
+import torch._dynamo
+import torch.distributed as dist
+import wandb
+from accelerate.utils import DistributedType, set_seed
+from omegaconf import OmegaConf
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+from transformers import AutoProcessor, get_scheduler
+
+# Local Modules
+from starVLA.dataloader import build_dataloader
+from starVLA.dataloader.worker_context import build_cpu_only_dataloader_kwargs
+from starVLA.model.framework.base_framework import build_framework
+from starVLA.model.framework.share_tools import apply_config_compat
+from starVLA.training.rl_games import CheckpointSyncManager, apply_action_spec, apply_model_alias, build_rl_games_eval_runner, sync_kv_memory_obs_window, validate_rl_games_config
+from starVLA.training.rl_games.auth import login_training_services
+from starVLA.training.rl_games.eval_core import EvalResult
+from starVLA.training.rl_games import action_cc_f1
+from starVLA.training.train_step_events import (
+    calculate_epoch_progress,
+    should_run_optional_step_interval_event,
+    should_run_step_interval_event,
+)
+from starVLA.training.trainer_utils.config_tracker import AccessTrackedConfig, wrap_config
+from starVLA.training.trainer_utils.trainer_tools import TrainerUtils, build_accelerator, build_param_lr_groups, setup_optimizer_and_scheduler, normalize_dotlist_args
+
+# Sane Defaults
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# Initialize logger
+logger = logging.getLogger(__name__)
+
+_DYNAMO_RECOMPILE_LIMIT = 64
+_DYNAMO_ACCUMULATED_RECOMPILE_LIMIT = 1024
+_GRACEFUL_STOP_REQUESTED = False
+_GRACEFUL_STOP_SIGNAL = None
+
+
+def _configure_torch_dynamo_recompile_limits() -> None:
+    config = torch._dynamo.config
+    config.recompile_limit = max(config.recompile_limit, _DYNAMO_RECOMPILE_LIMIT)
+    config.accumulated_recompile_limit = max(
+        config.accumulated_recompile_limit,
+        _DYNAMO_ACCUMULATED_RECOMPILE_LIMIT,
+    )
+
+
+def _as_bool(value, default: bool = False) -> bool:
+    if value in (None, ""):
+        return default
+    if isinstance(value, bool):
+        return value
+    return str(value).strip().lower() in {"1", "true", "yes", "y", "on"}
+
+
+def _request_graceful_stop(signum, frame) -> None:
+    global _GRACEFUL_STOP_REQUESTED, _GRACEFUL_STOP_SIGNAL
+    _GRACEFUL_STOP_REQUESTED = True
+    _GRACEFUL_STOP_SIGNAL = signum
+
+
+def _install_graceful_stop_handlers() -> dict[int, Any]:
+    global _GRACEFUL_STOP_REQUESTED, _GRACEFUL_STOP_SIGNAL
+    _GRACEFUL_STOP_REQUESTED = False
+    _GRACEFUL_STOP_SIGNAL = None
+    previous_handlers = {
+        signal.SIGINT: signal.getsignal(signal.SIGINT),
+        signal.SIGTERM: signal.getsignal(signal.SIGTERM),
+    }
+    signal.signal(signal.SIGINT, _request_graceful_stop)
+    signal.signal(signal.SIGTERM, _request_graceful_stop)
+    return previous_handlers
+
+
+def _restore_signal_handlers(previous_handlers: dict[int, Any]) -> None:
+    for signum, handler in previous_handlers.items():
+        signal.signal(signum, handler)
+
+
+def _distributed_barrier() -> None:
+    if not dist.is_initialized():
+        return
+    backend = str(dist.get_backend()).lower()
+    if backend == "nccl" and torch.cuda.is_available():
+        device_idx = int(torch.cuda.current_device())
+        dist.barrier(device_ids=[device_idx])
+        return
+    dist.barrier()
+
+
+def _destroy_distributed_process_group(*, use_barrier: bool) -> None:
+    if not dist.is_initialized():
+        return
+    if use_barrier:
+        try:
+            _distributed_barrier()
+        except Exception as exc:
+            logger.warning("Distributed barrier failed during shutdown: %s", exc)
+    try:
+        dist.destroy_process_group()
+        logger.info("Destroyed distributed process group")
+    except Exception as exc:
+        logger.warning("Failed to destroy distributed process group: %s", exc)
+
+
+def _sample_latency(sample: dict) -> int | None:
+    if not isinstance(sample, dict):
+        return None
+    value = sample.get("latency", sample.get("latency_id", None))
+    if value is None:
+        return None
+    if isinstance(value, torch.Tensor):
+        value = value.detach().cpu().reshape(-1)[0].item()
+    if isinstance(value, np.ndarray):
+        value = value.reshape(-1)[0].item()
+    return int(value)
+
+
+def _sample_task(sample: dict) -> str | None:
+    if not isinstance(sample, dict):
+        return None
+    value = sample.get("rl_games_task", sample.get("task", None))
+    if value is None:
+        return None
+    if isinstance(value, torch.Tensor):
+        value = value.detach().cpu().reshape(-1)[0].item()
+    if isinstance(value, np.ndarray):
+        value = value.reshape(-1)[0].item()
+    if isinstance(value, bytes):
+        value = value.decode("utf-8")
+    return str(value)
+
+
+def _group_examples_by_latency(examples: list[dict]) -> dict[int | None, list[dict]]:
+    grouped: dict[int | None, list[dict]] = {}
+    for example in examples:
+        grouped.setdefault(_sample_latency(example), []).append(example)
+    return grouped
+
+
+def _group_examples_by_task_latency(examples: list[dict]) -> dict[tuple[str | None, int | None], list[dict]]:
+    grouped: dict[tuple[str | None, int | None], list[dict]] = {}
+    for example in examples:
+        grouped.setdefault((_sample_task(example), _sample_latency(example)), []).append(example)
+    return grouped
+
+
+def _optional_positive_int(value) -> int | None:
+    if value in (None, ""):
+        return None
+    value = int(value)
+    return value if value > 0 else None
+
+
+def _optional_int_list(value) -> list[int] | None:
+    if value in (None, ""):
+        return None
+    if isinstance(value, str):
+        stripped = value.strip()
+        if stripped.startswith("[") and stripped.endswith("]"):
+            stripped = stripped[1:-1]
+        if not stripped:
+            return []
+        return [int(item.strip()) for item in stripped.split(",") if item.strip()]
+    return [int(item) for item in value]
+
+
+def _quota_interval_row_count(base_count: int, pass_start: float, quota: float) -> int:
+    """Rows selected by the pass interval [pass_start, pass_start + quota)."""
+    if int(base_count) <= 0 or float(quota) <= 0.0:
+        return 0
+    cursor = float(pass_start)
+    remaining = float(quota)
+    selected = 0
+    eps = 1e-9
+    while remaining > eps:
+        pass_idx = math.floor(cursor)
+        frac_start = cursor - pass_idx
+        if 1.0 - frac_start <= eps:
+            cursor = float(pass_idx + 1)
+            continue
+        segment = min(remaining, 1.0 - frac_start)
+        frac_end = frac_start + segment
+        start_idx = int(math.ceil(float(base_count) * frac_start - eps))
+        end_idx = int(math.ceil(float(base_count) * frac_end - eps))
+        selected += max(0, min(int(base_count), end_idx) - min(int(base_count), start_idx))
+        cursor += segment
+        remaining -= segment
+    return int(selected)
+
+
+def _latency_curriculum_cfg_from_config(cfg):
+    vla_data = getattr(getattr(cfg, "datasets", None), "vla_data", None)
+    if vla_data is None:
+        return None
+    curriculum_cfg = getattr(vla_data, "latency_curriculum", None)
+    if curriculum_cfg is None or not _as_bool(getattr(curriculum_cfg, "enabled", False), default=False):
+        return None
+    return curriculum_cfg
+
+
+def _dataset_latency_step_counts(dataset) -> dict[int, int]:
+    if dataset is not None and callable(getattr(dataset, "get_latency_step_counts", None)):
+        return {int(k): int(v) for k, v in dataset.get_latency_step_counts().items()}
+
+    counts: dict[int, int] = {}
+    single_datasets = list(getattr(dataset, "datasets", None) or [dataset])
+    for ds in single_datasets:
+        if ds is None or not callable(getattr(ds, "get_latencies_for_all_steps", None)):
+            continue
+        for latency in ds.get_latencies_for_all_steps():
+            latency = int(latency)
+            counts[latency] = counts.get(latency, 0) + 1
+    return counts
+
+
+def _build_quota_cumulative_plan(cfg, dataset, effective_batch_size: int) -> list[dict[str, Any]]:
+    curriculum_cfg = _latency_curriculum_cfg_from_config(cfg)
+    if curriculum_cfg is None:
+        return []
+    strategy = str(getattr(curriculum_cfg, "strategy", "") or "").strip().lower()
+    if strategy != "quota_cumulative":
+        return []
+
+    latencies = _optional_int_list(getattr(curriculum_cfg, "latencies", None))
+    if not latencies:
+        latencies = _dataset_latency_values(dataset)
+    latencies = sorted(dict.fromkeys(int(value) for value in latencies or []))
+    if not latencies:
+        raise ValueError("quota_cumulative requires configured or discoverable curriculum latencies.")
+
+    counts = _dataset_latency_step_counts(dataset)
+    missing = [latency for latency in latencies if int(counts.get(latency, 0)) <= 0]
+    if missing:
+        raise ValueError(f"quota_cumulative found no training rows for latencies={missing}; counts={counts}.")
+
+    new_latency_passes = float(getattr(curriculum_cfg, "new_latency_passes", 1.0))
+    replay_passes = float(getattr(curriculum_cfg, "replay_passes", 0.25))
+    target_total_passes = float(getattr(curriculum_cfg, "target_total_passes", 2.0))
+    final_equalization = _as_bool(getattr(curriculum_cfg, "final_equalization", True), default=True)
+    if new_latency_passes <= 0.0:
+        raise ValueError("quota_cumulative.new_latency_passes must be positive.")
+    if replay_passes < 0.0:
+        raise ValueError("quota_cumulative.replay_passes must be non-negative.")
+    if target_total_passes <= 0.0:
+        raise ValueError("quota_cumulative.target_total_passes must be positive.")
+    if effective_batch_size <= 0:
+        raise ValueError(f"quota_cumulative effective batch size must be positive, got {effective_batch_size}.")
+
+    cumulative_passes = {latency: 0.0 for latency in latencies}
+    phases: list[dict[str, Any]] = []
+    cursor = 0
+
+    def append_phase(name: str, quotas: dict[int, float]) -> None:
+        nonlocal cursor
+        cleaned = {int(latency): float(quota) for latency, quota in quotas.items() if float(quota) > 0.0}
+        if not cleaned:
+            return
+        quota_starts = {
+            latency: float(cumulative_passes.get(latency, 0.0))
+            for latency in cleaned
+        }
+        rows_by_latency = {
+            latency: _quota_interval_row_count(int(counts[latency]), quota_starts[latency], quota)
+            for latency, quota in cleaned.items()
+        }
+        total_rows = int(sum(rows_by_latency.values()))
+        if total_rows <= 0:
+            return
+        steps = int(math.ceil(total_rows / effective_batch_size))
+        start_step = cursor
+        cursor += steps
+        phases.append(
+            {
+                "phase_idx": len(phases),
+                "name": name,
+                "quotas": cleaned,
+                "quota_starts": quota_starts,
+                "rows_by_latency": rows_by_latency,
+                "dataset_size": total_rows,
+                "steps": steps,
+                "start_step": start_step,
+                "end_step": cursor,
+            }
+        )
+        for latency, quota in cleaned.items():
+            cumulative_passes[latency] = cumulative_passes.get(latency, 0.0) + float(quota)
+
+    for idx, latency in enumerate(latencies):
+        quotas = {prev_latency: replay_passes for prev_latency in latencies[:idx]}
+        quotas[latency] = new_latency_passes
+        append_phase(f"introduce_latency_{latency}", quotas)
+
+    if final_equalization:
+        top_up = {
+            latency: max(0.0, target_total_passes - cumulative_passes.get(latency, 0.0))
+            for latency in latencies
+        }
+        append_phase("equalize_total_coverage", top_up)
+
+    if not phases:
+        raise ValueError("quota_cumulative produced an empty phase plan.")
+    return phases
+
+
+def _configure_quota_cumulative_training_steps(cfg, dataloader, accelerator) -> None:
+    curriculum_cfg = _latency_curriculum_cfg_from_config(cfg)
+    if curriculum_cfg is None:
+        return
+    strategy = str(getattr(curriculum_cfg, "strategy", "") or "").strip().lower()
+    if strategy != "quota_cumulative":
+        return
+
+    grad_accum_steps = accelerator.gradient_accumulation_steps
+    per_device_bs = int(getattr(cfg.datasets.vla_data, "per_device_batch_size"))
+    effective_batch_size = per_device_bs * int(accelerator.num_processes) * grad_accum_steps
+    dataset = getattr(dataloader, "dataset", None)
+    plan = _build_quota_cumulative_plan(cfg, dataset, effective_batch_size)
+    total_steps = int(plan[-1]["end_step"])
+
+    step_budget_mode = str(getattr(curriculum_cfg, "step_budget_mode", "auto") or "auto").strip().lower()
+    if step_budget_mode != "auto":
+        raise ValueError("quota_cumulative currently supports step_budget_mode=auto only.")
+
+    cfg.trainer.max_train_steps = total_steps
+    curriculum_cfg.computed_plan = plan
+    if not dist.is_initialized() or dist.get_rank() == 0:
+        logger.info(
+            "quota_cumulative derived max_train_steps=%s from %s phases and effective_batch_size=%s",
+            total_steps,
+            len(plan),
+            effective_batch_size,
+        )
+        for phase in plan:
+            logger.info(
+                "quota_cumulative phase %s %s: steps=%s rows=%s quotas=%s rows_by_latency=%s end_step=%s",
+                phase["phase_idx"],
+                phase["name"],
+                phase["steps"],
+                phase["dataset_size"],
+                phase["quotas"],
+                phase["rows_by_latency"],
+                phase["end_step"],
+            )
+
+
+def _parse_phase_distributions(raw: Any) -> list[dict[int, float]] | None:
+    if raw is None:
+        return None
+    phases = list(raw)
+    if not phases:
+        return None
+    return [{int(k): float(v) for k, v in dict(phase).items()} for phase in phases]
+
+
+def _dataset_latency_values(dataset) -> list[int]:
+    values: set[int] = set()
+    single_datasets = list(getattr(dataset, "datasets", None) or [dataset])
+    for ds in single_datasets:
+        all_steps = getattr(ds, "all_steps", None) or []
+        seen_trajectories: set[int] = set()
+        for trajectory_id, _ in all_steps:
+            trajectory_id = int(trajectory_id)
+            if trajectory_id in seen_trajectories:
+                continue
+            seen_trajectories.add(trajectory_id)
+            for latency_key in ("latency", "latency_id"):
+                try:
+                    latency_data = ds.get_trajectory_columns(trajectory_id, [latency_key])
+                except Exception:
+                    continue
+                if latency_key not in latency_data.columns or len(latency_data) == 0:
+                    continue
+                values.update(int(value) for value in latency_data[latency_key].dropna().unique().tolist())
+                break
+    return sorted(values)
+
+
+def _dataset_task_latency_values(dataset) -> dict[str, list[int]]:
+    values: dict[str, set[int]] = {}
+    single_datasets = list(getattr(dataset, "datasets", None) or [dataset])
+    for ds in single_datasets:
+        task = getattr(ds, "rl_games_task", None)
+        if task in (None, ""):
+            continue
+        task = str(task)
+        for latency in _dataset_latency_values(ds):
+            values.setdefault(task, set()).add(int(latency))
+    return {task: sorted(latencies) for task, latencies in values.items()}
+
+
+def _episode_latency(ds, trajectory_id: int) -> int | None:
+    for latency_key in ("latency", "latency_id"):
+        try:
+            latency_data = ds.get_trajectory_columns(int(trajectory_id), [latency_key])
+        except Exception:
+            continue
+        if latency_key not in latency_data.columns or len(latency_data) == 0:
+            continue
+        values = latency_data[latency_key].dropna().unique().tolist()
+        if values:
+            return int(values[0])
+    return None
+
+
+def _identity_collate(batch):
+    return batch
+
+
+class _ActionCCF1FrameDataset(Dataset):
+    def __init__(self, single_datasets: list, frame_records: list[tuple[int, int, str, int, str, int | None]]):
+        self.single_datasets = single_datasets
+        self.frame_records = frame_records
+
+    def __len__(self) -> int:
+        return len(self.frame_records)
+
+    def __getitem__(self, index: int):
+        dataset_index, flat_idx, episode_key, base_idx, ep_task, _episode_latency_value = self.frame_records[index]
+        sample = self.single_datasets[dataset_index][flat_idx]
+        sample_task = _sample_task(sample) or ep_task
+        return sample_task, episode_key, base_idx, _sample_latency(sample), sample
+
+
+def load_fast_tokenizer():
+    return AutoProcessor.from_pretrained("physical-intelligence/fast", trust_remote_code=True)
+
+
+def setup_directories(cfg) -> Path:
+    """Create output directory and checkpoint directory."""
+    cfg.output_dir = os.path.join(cfg.run_root_dir, cfg.run_id)
+    output_dir = Path(cfg.output_dir)
+
+    if not dist.is_initialized() or dist.get_rank() == 0:
+        os.makedirs(output_dir, exist_ok=True)
+        os.makedirs(output_dir / "checkpoints", exist_ok=True)
+
+    return output_dir
+
+
+def prepare_data(cfg, accelerator, output_dir) -> tuple[DataLoader, DataLoader | None]:
+    """Prepare VLA training data."""
+    logger.info(f"Creating VLA Dataset with Mixture `{cfg.datasets.vla_data.data_mix}`")
+    vla_train_dataloader = build_dataloader(cfg=cfg, dataset_py=cfg.datasets.vla_data.dataset_py)
+    vla_eval_dataloader = None
+    eval_data_mix = getattr(cfg.datasets.vla_data, "eval_data_mix", None)
+    if eval_data_mix:
+        logger.info(f"Creating VLA Eval Dataset with Mixture `{eval_data_mix}`")
+        vla_eval_dataloader = build_dataloader(
+            cfg=cfg,
+            dataset_py=cfg.datasets.vla_data.dataset_py,
+            data_mix=str(eval_data_mix),
+            mode="eval",
+            save_statistics_filename="dataset_statistics_eval.json",
+        )
+
+    accelerator.dataloader_config.dispatch_batches = False
+    _distributed_barrier()
+    return vla_train_dataloader, vla_eval_dataloader
+
+
+def build_initial_framework(cfg):
+    """Initialize the backbone and fresh action modules with the configured seed."""
+    set_seed(cfg.seed)
+    return build_framework(cfg)
+
+
+def setup_optimizer_and_scheduler(model, cfg) -> Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler]:
+    """Set optimizer and scheduler."""
+    param_groups = build_param_lr_groups(model=model, cfg=cfg)
+    fused = bool(getattr(cfg.trainer.optimizer, "fused", False))
+    optimizer = torch.optim.AdamW(
+        param_groups,
+        lr=cfg.trainer.learning_rate.base,
+        betas=tuple(cfg.trainer.optimizer.betas),
+        weight_decay=cfg.trainer.optimizer.weight_decay,
+        eps=cfg.trainer.optimizer.eps,
+        fused=fused,
+    )
+
+    if dist.is_initialized() and dist.get_rank() == 0:
+        for group in optimizer.param_groups:
+            logger.info(f"LR Group {group['name']}: lr={group['lr']}, num_params={len(group['params'])}")
+
+    # Strip keys unknown to transformers' get_scheduler before passing kwargs.
+    sched_kwargs = {k: v for k, v in cfg.trainer.scheduler_specific_kwargs.items()}
+    lr_scheduler = get_scheduler(
+        name=cfg.trainer.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=cfg.trainer.num_warmup_steps,
+        num_training_steps=cfg.trainer.max_train_steps,
+        scheduler_specific_kwargs=sched_kwargs,
+    )
+
+    return optimizer, lr_scheduler
+
+
+def _pin_cuda_device_from_local_rank() -> None:
+    if not torch.cuda.is_available():
+        return
+    local_rank = os.environ.get("LOCAL_RANK")
+    if local_rank in (None, ""):
+        return
+    try:
+        device_idx = int(local_rank)
+    except ValueError:
+        logger.warning("Ignoring invalid LOCAL_RANK=%r for CUDA device pinning", local_rank)
+        return
+
+    visible_devices = torch.cuda.device_count()
+    if device_idx < 0 or device_idx >= visible_devices:
+        logger.warning(
+            "LOCAL_RANK=%s cannot be mapped to a CUDA device; visible CUDA device count is %s",
+            local_rank,
+            visible_devices,
+        )
+        return
+    torch.cuda.set_device(device_idx)
+
+
+def _preload_model_checkpoint_before_accelerator(cfg, model):
+    """Load model-only checkpoints before ZeRO-3 can replace params with shards."""
+    trainer_cfg = getattr(cfg, "trainer", None)
+    if trainer_cfg is None:
+        return model
+
+    pretrained_checkpoint = getattr(trainer_cfg, "pretrained_checkpoint", None)
+    is_resume = bool(getattr(trainer_cfg, "is_resume", False))
+    if not pretrained_checkpoint or is_resume:
+        return model
+
+    reload_modules = getattr(trainer_cfg, "reload_modules", None)
+    model = TrainerUtils.load_pretrained_backbones(
+        model,
+        pretrained_checkpoint,
+        reload_modules=reload_modules,
+    )
+    trainer_cfg.pretrained_checkpoint = None
+    logger.info(
+        "Preloaded model checkpoint before Accelerator/DeepSpeed initialization: %s",
+        pretrained_checkpoint,
+    )
+    return model
+
+
+class VLATrainer(TrainerUtils):
+    def __init__(self, cfg, model, vla_train_dataloader, vla_eval_dataloader, optimizer, lr_scheduler, accelerator):
+        self.config = cfg
+        self.model = model
+        self.vla_train_dataloader = vla_train_dataloader
+        self.vla_eval_dataloader = vla_eval_dataloader
+        self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler
+        self.accelerator = accelerator
+
+        self.completed_steps = 0
+        self.total_batch_size = self._calculate_total_batch_size()
+        self._checkpoint_sync_manager = CheckpointSyncManager(cfg=self.config)
+        self._rl_games_eval_runner = None
+        self._pending_full_state_resume = None
+        self._save_best_model_enabled = _as_bool(
+            getattr(getattr(self.config, "checkpoint", {}), "save_best_model", None),
+            default=True,
+        )
+        self._save_final_model_enabled = _as_bool(
+            getattr(getattr(self.config, "checkpoint", {}), "save_final_model", None),
+            default=True,
+        )
+        self._save_pt_file_enabled = _as_bool(
+            getattr(getattr(self.config, "checkpoint", {}), "save_pt_file", None),
+            default=False,
+        )
+        self._save_training_state_enabled = _as_bool(
+            getattr(getattr(self.config, "checkpoint", {}), "save_training_state", None),
+            default=True,
+        )
+        self._save_safetensors_file_enabled = _as_bool(
+            getattr(getattr(self.config, "checkpoint", {}), "save_safetensors_file", None),
+            default=False,
+        )
+        self._best_score = float("-inf")
+        self._best_step = 0
+        self._best_state_path = None
+        self._best_metadata_path = None
+        self._latency_curriculum_phase = None
+        self._action_cc_f1_frame_loader = None
+        self._action_cc_f1_frame_loader_key = None
+        self._train_loss_sum = 0.0
+        self._train_loss_weight = 0.0
+        self.graceful_stop_requested = False
+        self._graceful_stop_logged = False
+        self._interrupt_checkpoint_saved = False
+        self._dagger_control_dir = (
+            Path(self.config.trainer.dagger_control_dir)
+            if "dagger_control_dir" in self.config.trainer
+            else None
+        )
+        stop_file = os.environ.get("STARVLA_STOP_FILE")
+        self._graceful_stop_file = Path(stop_file) if stop_file else None
+        if hasattr(self.config, "rl_games") and hasattr(self.config.rl_games, "env_eval"):
+            enabled = bool(getattr(self.config.rl_games.env_eval, "enabled", False))
+            if enabled:
+                self._rl_games_eval_runner = build_rl_games_eval_runner(
+                    cfg=self.config,
+                    output_dir=self.config.output_dir,
+                )
+
+    def _save_periodic_checkpoints_enabled(self) -> bool:
+        return not self._save_best_model_enabled
+
+    def prepare_training(self):
+        _configure_torch_dynamo_recompile_limits()
+        self.accelerator.register_for_checkpointing(self.lr_scheduler)
+
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        seed = self.config.seed + rank if hasattr(self.config, "seed") else rank + 3047
+        set_seed(seed)
+
+        # Save config snapshots upfront so that even if a later setup step
+        # (ckpt load / DeepSpeed init / dataloader build) crashes, the
+        # produced run dir is still introspectable / from_pretrained-able.
+        self._save_initial_configs()
+
+        self._init_checkpointing()
+
+        freeze_modules = (
+            self.config.trainer.freeze_modules
+            if (self.config and hasattr(self.config.trainer, "freeze_modules"))
+            else None
+        )
+        self.model = self.freeze_backbones(self.model, freeze_modules=freeze_modules)
+        self.model = self.freeze_vit_and_llm_layers(self.model, self.config)
+        self.print_trainable_parameters(self.model)
+
+        if self.vla_eval_dataloader is not None:
+            self.model, self.optimizer, self.vla_train_dataloader, self.vla_eval_dataloader = self.setup_distributed_training(
+                self.accelerator,
+                self.model,
+                self.optimizer,
+                self.vla_train_dataloader,
+                self.vla_eval_dataloader,
+            )
+        else:
+            self.model, self.optimizer, self.vla_train_dataloader = self.setup_distributed_training(
+                self.accelerator,
+                self.model,
+                self.optimizer,
+                self.vla_train_dataloader,
+            )
+
+        if self._pending_full_state_resume:
+            self._load_checkpoint(self._pending_full_state_resume)
+        else:
+            self._adjust_lr_scheduler_for_resume()
+
+        self._init_wandb()
+        self.accelerator.wait_for_everyone()
+
+    def _calculate_total_batch_size(self):
+        """Calculate global batch size."""
+        return (
+            self.config.datasets.vla_data.per_device_batch_size
+            * self.accelerator.num_processes
+            * self.accelerator.gradient_accumulation_steps
+        )
+
+    def _init_wandb(self):
+        """Initialize Weights & Biases."""
+        if self.accelerator.is_main_process:
+            wandb_name = self.config.wandb_name if "wandb_name" in self.config else self.config.run_id
+            wandb.init(
+                name=wandb_name,
+                dir=os.path.join(self.config.output_dir, "wandb"),
+                project=self.config.wandb_project,
+                entity=self.config.wandb_entity,
+                group=self.config.wandb_group,
+                tags=self.config.wandb_tags,
+            )
+            wandb.config.update({
+                "micro_batch": self.config.datasets.vla_data.per_device_batch_size,
+                "gradient_accumulation_steps": self.accelerator.gradient_accumulation_steps,
+                "global_batch": self.total_batch_size,
+                "training_latency_condition": self.config.training_latency_condition,
+            }, allow_val_change=True)
+            wandb.define_metric("global_step")
+            wandb.define_metric("*", step_metric="global_step")
+
+    def _save_initial_configs(self):
+        """Save full config and training script at the very start of training."""
+        if not self.accelerator.is_main_process:
+            return
+
+        output_dir = Path(self.config.output_dir)
+
+        # 1. Save config.full.yaml — the complete merged config (all parameters)
+        if isinstance(self.config, AccessTrackedConfig):
+            full_cfg = self.config.unwrap()
+        else:
+            full_cfg = self.config
+        full_yaml_path = output_dir / "config.full.yaml"
+        OmegaConf.save(full_cfg, full_yaml_path, resolve=True)
+        logger.info(f"📝 Full config saved at {full_yaml_path}")
+
+        # 2. Save config.yaml — accessed-only snapshot (will be updated at checkpoints)
+        if isinstance(self.config, AccessTrackedConfig):
+            self.config.save_accessed_config(output_dir / "config.yaml", use_original_values=False)
+            logger.info(f"📊 Accessed config snapshot saved at {output_dir / 'config.yaml'}")
+
+    def _init_checkpointing(self):
+        """Initialize checkpoint directory and handle checkpoint loading."""
+        self.checkpoint_dir = os.path.join(self.config.output_dir, "checkpoints")
+        os.makedirs(self.checkpoint_dir, exist_ok=True)
+        self._best_state_path = os.path.join(self.checkpoint_dir, "best_state")
+        self._best_metadata_path = os.path.join(self.checkpoint_dir, "best_model_metadata.json")
+        self._load_best_checkpoint_metadata()
+
+        pretrained_checkpoint = getattr(self.config.trainer, "pretrained_checkpoint", None)
+        is_resume = getattr(self.config.trainer, "is_resume", False)
+        self.resume_from_checkpoint = pretrained_checkpoint
+
+        if is_resume:
+            if pretrained_checkpoint:
+                resume_from_checkpoint = str(pretrained_checkpoint)
+                self.completed_steps = int(
+                    getattr(self.config.trainer, "resume_step", None) or self._step_from_checkpoint_path(resume_from_checkpoint)
+                )
+            else:
+                resume_from_checkpoint, self.completed_steps = self._get_latest_checkpoint(self.checkpoint_dir)
+            if resume_from_checkpoint:
+                self.resume_from_checkpoint = resume_from_checkpoint
+                if os.path.isdir(self.resume_from_checkpoint):
+                    self._pending_full_state_resume = self.resume_from_checkpoint
+                    logger.info(
+                        f"Will resume full training state from checkpoint: "
+                        f"{self.resume_from_checkpoint}, steps: {self.completed_steps}"
+                    )
+                else:
+                    self.model = self.load_pretrained_backbones(self.model, self.resume_from_checkpoint, reload_modules=None)
+                    logger.info(
+                        f"Resuming model weights from checkpoint: "
+                        f"{self.resume_from_checkpoint}, steps: {self.completed_steps}"
+                    )
+                return
+
+            logger.warning(f"No valid checkpoint found in {self.checkpoint_dir}. Starting training from scratch.")
+            self.completed_steps = 0
+
+        if pretrained_checkpoint:
+            reload_modules = getattr(self.config.trainer, "reload_modules", None)
+            self.model = self.load_pretrained_backbones(self.model, pretrained_checkpoint, reload_modules=reload_modules)
+            self.completed_steps = 0
+            self.resume_from_checkpoint = pretrained_checkpoint
+            logger.info(f"Loaded pretrained checkpoint: {pretrained_checkpoint}, steps: {self.completed_steps}")
+        else:
+            logger.info("No pretrained checkpoint provided. Starting training from scratch.")
+            self.completed_steps = 0
+
+    @staticmethod
+    def _step_from_checkpoint_path(checkpoint_path: str) -> int:
+        match = re.search(r"steps_(\d+)", os.path.basename(str(checkpoint_path)))
+        return int(match.group(1)) if match else 0
+
+    def _adjust_lr_scheduler_for_resume(self):
+        """Adjust LR scheduler state after resuming from non-zero steps."""
+        if self.completed_steps > 0:
+            logger.info(f"Adjusting LR scheduler for resume from step {self.completed_steps}")
+            for _ in range(self.completed_steps):
+                self.lr_scheduler.step()
+            logger.info(
+                f"LR scheduler adjusted to step {self.completed_steps}, current LR: {self.lr_scheduler.get_last_lr()}"
+            )
+
+    def _load_checkpoint(self, checkpoint_path):
+        """Load checkpoint."""
+        self.accelerator.load_state(checkpoint_path)
+        self.accelerator.print(f"Resumed from checkpoint: {checkpoint_path}")
+
+    def _load_best_checkpoint_metadata(self):
+        if not self._best_metadata_path or not os.path.isfile(self._best_metadata_path):
+            return
+        try:
+            with open(self._best_metadata_path, "r", encoding="utf-8") as handle:
+                metadata = json.load(handle)
+            self._best_score = float(metadata.get("best_score", self._best_score))
+            self._best_step = int(metadata.get("best_step", self._best_step))
+            logger.info(
+                "Loaded best checkpoint metadata: step=%s score=%s",
+                self._best_step,
+                self._best_score,
+            )
+        except Exception as exc:
+            logger.warning("Could not load best checkpoint metadata from %s: %s", self._best_metadata_path, exc)
+
+    @staticmethod
+    def _rl_games_best_score(eval_result) -> float:
+        if "macro_mean_reward" in eval_result.aggregate:
+            return float(eval_result.aggregate["macro_mean_reward"])
+        bucket_scores = [
+            float(metrics["mean_reward"])
+            for metrics in eval_result.per_latency.values()
+            if metrics is not None and "mean_reward" in metrics
+        ]
+        if bucket_scores:
+            return float(np.mean(bucket_scores))
+        return float(eval_result.aggregate.get("mean_reward", 0.0))
+
+    @staticmethod
+    def _eval_result_payload(eval_result) -> Dict[str, Any]:
+        return {
+            "per_latency": eval_result.per_latency,
+            "aggregate": eval_result.aggregate,
+        }
+
+    def _distributed_rl_games_eval_enabled(self) -> bool:
+        if self._rl_games_eval_runner is None:
+            return False
+        env_eval = getattr(self.config.rl_games, "env_eval", None)
+        mode = str(getattr(env_eval, "distributed_mode", "none") or "none").strip().lower()
+        return mode in {"rank_sharded", "distributed", "sharded"} and int(self.accelerator.num_processes) > 1
+
+    def _run_rl_games_eval_with_model_mode(
+        self,
+        *,
+        stage: str,
+        shard_rank: int = 0,
+        shard_count: int = 1,
+        save: bool = True,
+    ):
+        was_training = self.model.training
+        self.model.eval()
+        unwrapped = self.accelerator.unwrap_model(self.model)
+        try:
+            run_eval = torch._dynamo.disable(self._rl_games_eval_runner.run)
+            return run_eval(
+                model=unwrapped,
+                step=self.completed_steps,
+                stage=stage,
+                shard_rank=int(shard_rank),
+                shard_count=int(shard_count),
+                save=save,
+            )
+        finally:
+            try:
+                reset = getattr(unwrapped, "reset_memory", None)
+                if callable(reset):
+                    reset()
+            finally:
+                if was_training:
+                    self.model.train()
+
+    def _run_distributed_rl_games_eval(self, stage: str):
+        if self._rl_games_eval_runner is None:
+            return None
+
+        local_result = self._run_rl_games_eval_with_model_mode(
+            stage=stage,
+            shard_rank=int(self.accelerator.process_index),
+            shard_count=int(self.accelerator.num_processes),
+            save=False,
+        )
+        local_payload = self._eval_result_payload(local_result)
+
+        if dist.is_available() and dist.is_initialized():
+            gathered_payloads = [None for _ in range(dist.get_world_size())]
+            dist.all_gather_object(gathered_payloads, local_payload)
+        else:
+            gathered_payloads = [local_payload]
+
+        partial_results = [
+            EvalResult(per_latency=payload["per_latency"], aggregate=payload["aggregate"])
+            for payload in gathered_payloads
+            if payload is not None
+        ]
+        merged_result = self._rl_games_eval_runner.merge_results(
+            partial_results,
+            step=self.completed_steps,
+            stage=stage,
+        )
+        if self.accelerator.is_main_process:
+            self._rl_games_eval_runner.save(result=merged_result, step=self.completed_steps, stage=stage)
+        self.accelerator.wait_for_everyone()
+        return merged_result
+
+    def _run_mid_train_rl_games_eval(self, step_metrics: dict[str, float]) -> dict[str, float]:
+        if self._rl_games_eval_runner is None or not self._rl_games_eval_runner.is_enabled(stage="mid_train"):
+            return step_metrics
+        if self._refresh_graceful_stop_requested():
+            if self.accelerator.is_main_process:
+                logger.info("Skipping RL-games mid-train eval at step %s due to graceful stop", self.completed_steps)
+            return step_metrics
+        t_profile_eval = self._profile_start()
+        if self._distributed_rl_games_eval_enabled():
+            eval_result = self._run_distributed_rl_games_eval(stage="mid_train")
+        else:
+            eval_result = self._run_rl_games_eval_with_model_mode(stage="mid_train")
+        step_metrics = self._append_rl_games_eval_metrics(
+            step_metrics=step_metrics,
+            eval_result=eval_result,
+            stage="mid_train",
+        )
+        if self._profile_timing_should_log():
+            step_metrics["timing/rl_games_mid_train_eval_total"] = self._profile_elapsed(t_profile_eval)
+        eval_score = self._rl_games_best_score(eval_result)
+        is_new_best = self._save_best_checkpoint(
+            eval_result=eval_result,
+            score=eval_score,
+            stage="mid_train",
+        )
+        step_metrics["checkpoint/current_eval_score"] = float(eval_score)
+        step_metrics["checkpoint/best_score"] = float(self._best_score)
+        step_metrics["checkpoint/best_step"] = float(self._best_step)
+        step_metrics["checkpoint/is_new_best"] = float(is_new_best)
+        self._checkpoint_sync_manager.sync_eval_result(
+            eval_path=eval_result.path,
+            stage="mid_train",
+            step=self.completed_steps,
+        )
+        return step_metrics
+
+    def _best_config_paths(self) -> list[str]:
+        output_dir = Path(self.config.output_dir)
+        return [
+            str(output_dir / "config.full.yaml"),
+            str(output_dir / "config.yaml"),
+        ]
+
+    def _save_best_checkpoint(self, eval_result, score: float, stage: str) -> bool:
+        if not self._save_best_model_enabled:
+            return False
+        if score <= self._best_score:
+            return False
+        if not self._best_state_path or not self._best_metadata_path:
+            raise RuntimeError("Best checkpoint paths were not initialized")
+
+        if self.accelerator.is_main_process and os.path.exists(self._best_state_path):
+            shutil.rmtree(self._best_state_path)
+        self.accelerator.wait_for_everyone()
+        self.accelerator.save_state(self._best_state_path, safe_serialization=True)
+        self.accelerator.wait_for_everyone()
+
+        self._best_score = float(score)
+        self._best_step = int(self.completed_steps)
+
+        if self.accelerator.is_main_process:
+            if isinstance(self.config, AccessTrackedConfig):
+                output_dir = Path(self.config.output_dir)
+                self.config.save_accessed_config(output_dir / "config.yaml", use_original_values=False)
+            metadata = {
+                "best_step": self._best_step,
+                "best_score": self._best_score,
+                "score_name": "mid_train_macro_mean_reward",
+                "stage": stage,
+                "eval_result_path": getattr(eval_result, "path", None),
+                "aggregate": eval_result.aggregate,
+                "per_latency_scores": {
+                    key: float(metrics.get("mean_reward", 0.0))
+                    for key, metrics in eval_result.per_latency.items()
+                },
+            }
+            with open(self._best_metadata_path, "w", encoding="utf-8") as handle:
+                json.dump(metadata, handle, indent=2)
+            self.accelerator.print(
+                f"✅ New best checkpoint saved at {self._best_state_path} "
+                f"(step={self._best_step}, score={self._best_score:.6f})"
+            )
+            self._checkpoint_sync_manager.sync_best_checkpoint(
+                state_path=self._best_state_path,
+                metadata_path=self._best_metadata_path,
+                config_paths=self._best_config_paths(),
+            )
+        self.accelerator.wait_for_everyone()
+        return True
+
+    def _save_checkpoint(self):
+        """Save current training state."""
+        checkpoint_path = os.path.join(self.checkpoint_dir, f"steps_{self.completed_steps}")
+        state_checkpoint_path = checkpoint_path + "_state"
+
+        saved_state_path = None
+        if self._save_training_state_enabled:
+            if self.accelerator.is_main_process and os.path.exists(state_checkpoint_path):
+                shutil.rmtree(state_checkpoint_path)
+            self.accelerator.wait_for_everyone()
+            self.accelerator.save_state(state_checkpoint_path, safe_serialization=True)
+            self.accelerator.wait_for_everyone()
+            saved_state_path = state_checkpoint_path
+            if self.accelerator.is_main_process:
+                self.accelerator.print(f"✅ Full training state saved at {state_checkpoint_path}")
+
+        if self.accelerator.is_main_process:
+            model_checkpoint_path = None
+            if self._save_pt_file_enabled or self._save_safetensors_file_enabled:
+                state_dict = self.accelerator.get_state_dict(self.model)
+
+            if self._save_pt_file_enabled:
+                model_checkpoint_path = checkpoint_path + "_pytorch_model.pt"
+                torch.save(state_dict, model_checkpoint_path)
+
+            if self._save_safetensors_file_enabled:
+                from safetensors.torch import save_file
+
+                safetensors_state_dict = {
+                    name: tensor.detach().to(torch.bfloat16) if torch.is_floating_point(tensor) else tensor.detach()
+                    for name, tensor in state_dict.items()
+                }
+                model_checkpoint_path = checkpoint_path + "_model.safetensors"
+                save_file(safetensors_state_dict, model_checkpoint_path)
+
+            summary_data = {"steps": self.completed_steps}
+            with open(os.path.join(self.config.output_dir, "summary.jsonl"), "a") as f:
+                f.write(json.dumps(summary_data) + "\n")
+            if saved_state_path is not None:
+                self.accelerator.print(f"✅ Checkpoint state saved at {saved_state_path}")
+            if model_checkpoint_path is not None:
+                self.accelerator.print(f"✅ Model checkpoint saved at {model_checkpoint_path}")
+            self._checkpoint_sync_manager.register_local_checkpoint(
+                step=self.completed_steps,
+                state_path=saved_state_path,
+                model_path=model_checkpoint_path,
+            )
+
+            if isinstance(self.config, AccessTrackedConfig):
+                logger.info("📊 Saving accessed configuration...")
+                output_dir = Path(self.config.output_dir)
+                self.config.save_accessed_config(output_dir / "config.yaml", use_original_values=False)
+                logger.info("✅ Configuration files saved")
+
+        self.accelerator.wait_for_everyone()
+
+    def _log_wandb(self, metrics):
+        wandb.log({**metrics, "global_step": self.completed_steps})
+
+    def _log_metrics(self, metrics):
+        """Record training metrics."""
+        if self.completed_steps % self.config.trainer.logging_frequency == 0 and self.accelerator.is_main_process:
+            last_lrs = self.lr_scheduler.get_last_lr()
+            for i, group in enumerate(self.optimizer.param_groups):
+                group_name = group["name"]
+                metrics[f"learning_rate/{group_name}"] = last_lrs[i]
+            dataset_size = len(self.vla_train_dataloader.dataset)
+            metrics["epoch"] = round(
+                calculate_epoch_progress(
+                    completed_steps=self.completed_steps,
+                    total_batch_size=self.total_batch_size,
+                    dataset_size=dataset_size,
+                ),
+                2,
+            )
+            self._log_wandb(metrics)
+            logger.info(f"Step {self.completed_steps}, Loss: {metrics})")
+
+    @staticmethod
+    def _total_grad_norm(parameters) -> float:
+        grads = [p.grad.detach() for p in parameters if p.grad is not None]
+        if not grads:
+            return 0.0
+        device = grads[0].device
+        norms = torch.stack([torch.linalg.vector_norm(g.to(device), 2) for g in grads])
+        return float(torch.linalg.vector_norm(norms, 2).item())
+
+    def _record_train_loss(self, action_loss, loss_weight) -> None:
+        weight = float(loss_weight)
+        self._train_loss_sum += float(action_loss.detach().float().item()) * weight
+        self._train_loss_weight += weight
+
+    def _finalize_train_loss(self, device) -> float:
+        loss_stats = torch.tensor(
+            [self._train_loss_sum, self._train_loss_weight],
+            device=device,
+            dtype=torch.float32,
+        )
+        if dist.is_initialized():
+            dist.all_reduce(loss_stats, op=dist.ReduceOp.SUM)
+        self._train_loss_sum = 0.0
+        self._train_loss_weight = 0.0
+        return float((loss_stats[0] / loss_stats[1]).item())
+
+    @staticmethod
+    def _append_rl_games_eval_metrics(step_metrics: Dict[str, float], eval_result, stage: str) -> Dict[str, float]:
+        aggregate = eval_result.aggregate
+        prefix = f"rl_games_eval/{stage}"
+        step_metrics[f"{prefix}/total_episodes"] = float(aggregate.get("total_episodes", 0))
+        step_metrics[f"{prefix}/mean_reward"] = float(aggregate.get("mean_reward", 0.0))
+        step_metrics[f"{prefix}/mean_length"] = float(aggregate.get("mean_length", 0.0))
+        step_metrics[f"{prefix}/std_reward"] = float(aggregate.get("std_reward", 0.0))
+        step_metrics[f"{prefix}/std_length"] = float(aggregate.get("std_length", 0.0))
+        step_metrics[f"{prefix}/macro_mean_reward"] = float(
+            aggregate.get("macro_mean_reward", aggregate.get("mean_reward", 0.0))
+        )
+        step_metrics[f"{prefix}/macro_mean_length"] = float(
+            aggregate.get("macro_mean_length", aggregate.get("mean_length", 0.0))
+        )
+        step_metrics[f"{prefix}/task_count"] = float(aggregate.get("task_count", 0))
+
+        for key, metrics in eval_result.per_latency.items():
+            key_slug = key.replace("/", "__")
+            step_metrics[f"{prefix}/{key_slug}/mean_reward"] = float(metrics.get("mean_reward", 0.0))
+            step_metrics[f"{prefix}/{key_slug}/mean_length"] = float(metrics.get("mean_length", 0.0))
+            step_metrics[f"{prefix}/{key_slug}/std_reward"] = float(metrics.get("std_reward", 0.0))
+            step_metrics[f"{prefix}/{key_slug}/std_length"] = float(metrics.get("std_length", 0.0))
+            step_metrics[f"{prefix}/{key_slug}/num_episodes"] = float(metrics.get("num_episodes", 0))
+        return step_metrics
+
+    def _create_data_iterators(self):
+        """Create data iterators."""
+        self.vla_iter = iter(self.vla_train_dataloader)
+
+    def _wait_for_dagger_round(self) -> bool:
+        """Serve the resident model, then activate the controller's next mixture."""
+        continue_path = self._dagger_control_dir / f"step_{self.completed_steps}.continue.json"
+        if continue_path.is_file():
+            command = json.loads(continue_path.read_text(encoding="utf-8"))
+        else:
+            # latency_bench is an optional integration for upstream StarVLA training.
+            from latency_bench.integrations.starvla_resident_dagger import serve_resident_inference
+            from latency_bench.policy.starvla import build_live_starvla_policy
+
+            resident_config = json.loads(
+                (self._dagger_control_dir / "resident_policy.json").read_text(encoding="utf-8")
+            )
+            model_config = OmegaConf.to_container(
+                self.config.unwrap() if isinstance(self.config, AccessTrackedConfig) else self.config,
+                resolve=True,
+            )
+            was_training = self.model.training
+            self.model.eval()
+            unwrapped = self.accelerator.unwrap_model(self.model)
+            policy = build_live_starvla_policy(
+                framework=unwrapped,
+                model_cfg=model_config,
+                config=resident_config,
+                action_resolver=None,
+            )
+            try:
+                with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+                    command = serve_resident_inference(
+                        policy=policy,
+                        control_dir=self._dagger_control_dir,
+                        socket_dir=Path(resident_config["socket_dir"]),
+                        step=self.completed_steps,
+                        rank=int(self.accelerator.process_index),
+                        world_size=int(self.accelerator.num_processes),
+                        stop_requested=self._refresh_graceful_stop_requested,
+                    )
+            finally:
+                reset = getattr(unwrapped, "reset_memory", None)
+                if callable(reset):
+                    reset()
+                if was_training:
+                    self.model.train()
+
+        if command is None:
+            self._mark_graceful_stop_requested("dagger_round_wait")
+            return False
+
+        self.vla_train_dataloader.dataset.refresh_custom_mixture()
+        self._create_data_iterators()
+        _distributed_barrier()
+        if self.accelerator.is_main_process:
+            self._log_wandb(command["metrics"])
+            logger.info("DAgger round %s resumed with refreshed data", self.completed_steps)
+        return True
+
+    def _mark_graceful_stop_requested(self, reason: str) -> None:
+        self.graceful_stop_requested = True
+        if self._graceful_stop_file is not None and not self._graceful_stop_file.exists():
+            try:
+                self._graceful_stop_file.parent.mkdir(parents=True, exist_ok=True)
+                self._graceful_stop_file.write_text(f"reason={reason}\ntime={time.time():.6f}\n", encoding="utf-8")
+            except Exception as exc:
+                logger.warning("Failed to write graceful stop file %s: %s", self._graceful_stop_file, exc)
+        if self.accelerator.is_main_process and not self._graceful_stop_logged:
+            logger.info("Graceful stop requested (%s); stopping at the next safe training boundary.", reason)
+            self._graceful_stop_logged = True
+
+    def _refresh_graceful_stop_requested(self) -> bool:
+        if _GRACEFUL_STOP_REQUESTED:
+            signal_name = (
+                signal.Signals(_GRACEFUL_STOP_SIGNAL).name
+                if _GRACEFUL_STOP_SIGNAL is not None
+                else "signal"
+            )
+            self._mark_graceful_stop_requested(signal_name)
+        if self._graceful_stop_file is not None and self._graceful_stop_file.exists():
+            self._mark_graceful_stop_requested(f"stop_file:{self._graceful_stop_file}")
+        return self.graceful_stop_requested
+
+    def _save_interrupt_checkpoint(self) -> None:
+        if self._interrupt_checkpoint_saved or self.completed_steps <= 0:
+            return
+        self._interrupt_checkpoint_saved = True
+        if self.accelerator.is_main_process:
+            logger.info("Saving interrupt checkpoint at step %s", self.completed_steps)
+        self._save_checkpoint()
+
+    def _shutdown_dataloader_workers(self) -> None:
+        seen: set[int] = set()
+
+        def shutdown(obj) -> None:
+            if obj is None:
+                return
+            obj_id = id(obj)
+            if obj_id in seen:
+                return
+            seen.add(obj_id)
+            shutdown_fn = getattr(obj, "_shutdown_workers", None)
+            if callable(shutdown_fn):
+                try:
+                    shutdown_fn()
+                except Exception as exc:
+                    logger.warning("Failed to shutdown dataloader workers for %s: %s", type(obj).__name__, exc)
+            for attr in ("_iterator", "dataloader", "_dataloader", "base_dataloader", "_loader"):
+                shutdown(getattr(obj, attr, None))
+
+        shutdown(getattr(self, "vla_iter", None))
+        shutdown(self.vla_train_dataloader)
+        shutdown(self.vla_eval_dataloader)
+        shutdown(self._action_cc_f1_frame_loader)
+
+    def cleanup_runtime(self) -> None:
+        self._shutdown_dataloader_workers()
+        unwrapped = None
+        try:
+            unwrapped = self.accelerator.unwrap_model(self.model)
+        except Exception:
+            unwrapped = self.model
+        reset = getattr(unwrapped, "reset_memory", None)
+        if callable(reset):
+            try:
+                reset()
+            except Exception as exc:
+                logger.warning("Failed to reset model memory during cleanup: %s", exc)
+
+    def _profile_timing_enabled(self) -> bool:
+        return self.config.trainer.profile_timing.enabled
+
+    def _profile_timing_log_interval(self) -> int:
+        return self.config.trainer.profile_timing.log_interval
+
+    def _profile_timing_should_log(self, step: int | None = None) -> bool:
+        if not self._profile_timing_enabled():
+            return False
+        step = self.completed_steps if step is None else step
+        return step % self._profile_timing_log_interval() == 0
+
+    def _profile_sync(self) -> None:
+        if self._profile_timing_enabled() and torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    def _profile_start(self) -> float:
+        self._profile_sync()
+        return time.perf_counter()
+
+    def _profile_elapsed(self, start: float) -> float:
+        self._profile_sync()
+        return time.perf_counter() - start
+
+    def _latency_curriculum_cfg(self):
+        vla_data = getattr(getattr(self.config, "datasets", None), "vla_data", None)
+        if vla_data is None:
+            return None
+        cfg = getattr(vla_data, "latency_curriculum", None)
+        if cfg is None or not _as_bool(getattr(cfg, "enabled", False), default=False):
+            return None
+        return cfg
+
+    def _distribution_curriculum_phase_steps(self, cfg, num_phases: int) -> list[int]:
+        explicit = _optional_int_list(getattr(cfg, "phase_steps", None))
+        if explicit:
+            if len(explicit) != num_phases:
+                raise ValueError(
+                    "datasets.vla_data.latency_curriculum.phase_steps must have one value per phase "
+                    f"(got {len(explicit)} steps for {num_phases} phases)."
+                )
+            if any(step <= 0 for step in explicit):
+                raise ValueError("latency_curriculum.phase_steps values must be positive.")
+            return explicit
+        total_steps = int(getattr(self.config.trainer, "max_train_steps"))
+        base = total_steps // num_phases
+        remainder = total_steps % num_phases
+        return [base + (1 if idx < remainder else 0) for idx in range(num_phases)]
+
+    def _latency_curriculum_phase_steps(self, cfg, latencies: list[int]) -> list[int]:
+        explicit = _optional_int_list(getattr(cfg, "phase_steps", None))
+        if explicit:
+            if len(explicit) != len(latencies):
+                raise ValueError(
+                    "datasets.vla_data.latency_curriculum.phase_steps must have one value per latency "
+                    f"(got {len(explicit)} steps for {len(latencies)} latencies)."
+                )
+            if any(step <= 0 for step in explicit):
+                raise ValueError("latency_curriculum.phase_steps values must be positive.")
+            return explicit
+
+        total_steps = int(getattr(self.config.trainer, "max_train_steps"))
+        if total_steps < len(latencies):
+            raise ValueError(
+                "trainer.max_train_steps must be at least the number of latency curriculum phases "
+                f"when phase_steps is not set (got max_train_steps={total_steps}, latencies={latencies})."
+            )
+        base = total_steps // len(latencies)
+        remainder = total_steps % len(latencies)
+        return [base + (1 if idx < remainder else 0) for idx in range(len(latencies))]
+
+    def _quota_curriculum_plan(self) -> list[dict[str, Any]]:
+        cfg = self._latency_curriculum_cfg()
+        if cfg is None:
+            return []
+        plan = getattr(cfg, "computed_plan", None)
+        if plan in (None, ""):
+            return []
+        if isinstance(plan, list):
+            return plan
+        return list(plan)
+
+    def _latency_curriculum_phase_for_step(
+        self,
+        step: int,
+    ) -> tuple[int, list[int], dict[int, float] | None, dict[int, float] | None, str] | None:
+        cfg = self._latency_curriculum_cfg()
+        if cfg is None:
+            return None
+        strategy = str(getattr(cfg, "strategy", "exclusive") or "exclusive").strip().lower()
+        if strategy == "quota_cumulative":
+            plan = self._quota_curriculum_plan()
+            if not plan:
+                raise ValueError("quota_cumulative requires a computed_plan; did setup derive max_train_steps?")
+            phase = plan[-1]
+            for candidate in plan:
+                if int(step) < int(candidate["end_step"]):
+                    phase = candidate
+                    break
+            quotas = {int(k): float(v) for k, v in dict(phase["quotas"]).items()}
+            quota_starts = {int(k): float(v) for k, v in dict(phase.get("quota_starts", {})).items()}
+            active = sorted(quotas)
+            return int(phase["phase_idx"]), active, quotas, quota_starts, str(phase.get("name", f"phase_{phase['phase_idx']}"))
+
+        if strategy == "curriculum_cumulative_distribution":
+            phase_distributions = _parse_phase_distributions(getattr(cfg, "phase_distributions", None))
+            if not phase_distributions:
+                raise ValueError("curriculum_cumulative_distribution requires phase_distributions to be configured.")
+            num_phases = len(phase_distributions)
+            phase_steps = self._distribution_curriculum_phase_steps(cfg, num_phases)
+            cursor = 0
+            phase_idx = num_phases - 1
+            for idx, num_steps in enumerate(phase_steps):
+                cursor += int(num_steps)
+                if int(step) < cursor:
+                    phase_idx = idx
+                    break
+            distribution = phase_distributions[phase_idx]
+            active = sorted(distribution)
+            return phase_idx, active, distribution, {}, f"phase_{phase_idx}"
+
+        latencies = _optional_int_list(getattr(cfg, "latencies", None))
+        if not latencies:
+            latencies = _dataset_latency_values(getattr(self.vla_train_dataloader, "dataset", None))
+        if not latencies:
+            raise ValueError("latency_curriculum.enabled=true but no curriculum latencies were configured or found.")
+        latencies = sorted(dict.fromkeys(int(value) for value in latencies))
+        if strategy not in {"exclusive", "cumulative"}:
+            raise ValueError(
+                f"Unsupported latency_curriculum.strategy={strategy!r}; "
+                "expected exclusive, cumulative, quota_cumulative, or curriculum_cumulative_distribution."
+            )
+
+        phase_steps = self._latency_curriculum_phase_steps(cfg, latencies)
+        cursor = 0
+        phase_idx = len(latencies) - 1
+        for idx, num_steps in enumerate(phase_steps):
+            cursor += int(num_steps)
+            if int(step) < cursor:
+                phase_idx = idx
+                break
+
+        if strategy == "exclusive":
+            active = [latencies[phase_idx]]
+        else:
+            active = latencies[: phase_idx + 1]
+        return phase_idx, active, None, None, f"phase_{phase_idx}"
+
+    def _latency_curriculum_phase_end_at_step(self, step: int) -> bool:
+        cfg = self._latency_curriculum_cfg()
+        if cfg is None or int(step) <= 0:
+            return False
+        strategy = str(getattr(cfg, "strategy", "exclusive") or "exclusive").strip().lower()
+        if strategy == "quota_cumulative":
+            plan = self._quota_curriculum_plan()
+            final_step = int(getattr(self.config.trainer, "max_train_steps"))
+            return any(int(phase["end_step"]) == int(step) and int(step) < final_step for phase in plan)
+
+        if strategy == "curriculum_cumulative_distribution":
+            phase_distributions = _parse_phase_distributions(getattr(cfg, "phase_distributions", None))
+            if not phase_distributions:
+                return False
+            phase_steps = self._distribution_curriculum_phase_steps(cfg, len(phase_distributions))
+            cursor = 0
+            final_step = int(getattr(self.config.trainer, "max_train_steps"))
+            for num_steps in phase_steps[:-1]:
+                cursor += int(num_steps)
+                if cursor == int(step) and int(step) < final_step:
+                    return True
+            return False
+
+        latencies = _optional_int_list(getattr(cfg, "latencies", None))
+        if not latencies:
+            latencies = _dataset_latency_values(getattr(self.vla_train_dataloader, "dataset", None))
+        if not latencies:
+            return False
+        phase_steps = self._latency_curriculum_phase_steps(cfg, sorted(dict.fromkeys(int(value) for value in latencies)))
+        cursor = 0
+        final_step = int(getattr(self.config.trainer, "max_train_steps"))
+        for num_steps in phase_steps[:-1]:
+            cursor += int(num_steps)
+            if cursor == int(step) and int(step) < final_step:
+                return True
+        return False
+
+    def _apply_latency_curriculum(self, *, force: bool = False) -> dict[str, float]:
+        phase = self._latency_curriculum_phase_for_step(self.completed_steps)
+        if phase is None:
+            return {}
+        phase_idx, active_latencies, active_quotas, quota_starts, phase_name = phase
+        active_key = (
+            phase_idx,
+            tuple(active_latencies),
+            tuple(sorted((active_quotas or {}).items())),
+            tuple(sorted((quota_starts or {}).items())),
+        )
+        if not force and active_key == self._latency_curriculum_phase:
+            return {}
+
+        dataset = getattr(self.vla_train_dataloader, "dataset", None)
+        if dataset is None:
+            raise ValueError("latency curriculum requires a train dataset.")
+        if active_quotas is not None:
+            if not callable(getattr(dataset, "set_latency_quota_filter", None)):
+                raise ValueError("latency curriculum with quotas requires a train dataset with set_latency_quota_filter(...).")
+            dataset.set_latency_quota_filter(active_quotas, quota_starts=quota_starts)
+        else:
+            if not callable(getattr(dataset, "set_active_latency_filter", None)):
+                raise ValueError("latency curriculum requires a train dataset with set_active_latency_filter(...).")
+            dataset.set_active_latency_filter(active_latencies)
+        self._latency_curriculum_phase = active_key
+        if hasattr(self.vla_train_dataloader, "sampler") and callable(getattr(self.vla_train_dataloader.sampler, "set_epoch", None)):
+            self.vla_train_dataloader.sampler.set_epoch(phase_idx)
+        self._create_data_iterators()
+        if self.accelerator.is_main_process:
+            logger.info(
+                "Latency curriculum phase %s (%s) at step %s: active_latencies=%s, active_quotas=%s, dataset_size=%s",
+                phase_idx,
+                phase_name,
+                self.completed_steps,
+                active_latencies,
+                active_quotas,
+                len(dataset),
+            )
+        metrics = {
+            "latency_curriculum/phase": float(phase_idx),
+            "latency_curriculum/active_count": float(len(active_latencies)),
+            "latency_curriculum/max_active_latency": float(max(active_latencies)),
+            "latency_curriculum/dataset_size": float(len(dataset)),
+        }
+        if active_quotas is not None:
+            for latency, quota in active_quotas.items():
+                metrics[f"latency_curriculum/quota/latency_{latency}"] = float(quota)
+                if quota_starts and latency in quota_starts:
+                    metrics[f"latency_curriculum/quota_start/latency_{latency}"] = float(quota_starts[latency])
+            for phase_info in self._quota_curriculum_plan():
+                if int(phase_info["phase_idx"]) == int(phase_idx):
+                    for latency, rows in dict(phase_info.get("rows_by_latency", {})).items():
+                        metrics[f"latency_curriculum/rows/latency_{latency}"] = float(rows)
+                    metrics["latency_curriculum/phase_optimizer_steps"] = float(phase_info.get("steps", 0))
+                    break
+        return metrics
+
+    def _get_next_batch(self):
+        """Get next batch (automatically handle data loop)."""
+        try:
+            batch_vla = next(self.vla_iter)
+        except StopIteration:
+            if not hasattr(self, "vla_epoch_count"):
+                self.vla_epoch_count = 0
+            self.vla_iter, self.vla_epoch_count = TrainerUtils._reset_dataloader(
+                self.vla_train_dataloader, self.vla_epoch_count
+            )
+            batch_vla = next(self.vla_iter)
+
+        return batch_vla
+
+    def train(self):
+        """Execute training loop."""
+        self.model.train()
+        self._log_training_config()
+        self._create_data_iterators()
+        self._apply_latency_curriculum(force=True)
+        # DAgger pauses at round boundaries without shortening the LR schedule.
+        stop_step = (
+            self.config.trainer.stop_after_steps
+            if "stop_after_steps" in self.config.trainer
+            else self.config.trainer.max_train_steps
+        )
+        progress_bar = tqdm(
+            total=stop_step,
+            initial=self.completed_steps,
+            disable=not self.accelerator.is_local_main_process,
+        )
+
+        try:
+            save_interval = int(self.config.trainer.save_interval)
+            resume_continue_path = (
+                self._dagger_control_dir / f"step_{self.completed_steps}.continue.json"
+                if self._dagger_control_dir is not None
+                else None
+            )
+            if (
+                resume_continue_path is not None
+                and self.completed_steps > 0
+                and self.completed_steps < stop_step
+                and self.completed_steps % save_interval == 0
+                and not resume_continue_path.is_file()
+                and not self._wait_for_dagger_round()
+            ):
+                return
+            while self.completed_steps < stop_step:
+                if self._refresh_graceful_stop_requested():
+                    self._save_interrupt_checkpoint()
+                    break
+
+                profile_enabled = self._profile_timing_enabled()
+                t_profile_data = self._profile_start() if profile_enabled else None
+                t_start_data = time.perf_counter()
+                batch_vla = self._get_next_batch()
+                t_end_data = time.perf_counter()
+                profile_data_seconds = self._profile_elapsed(t_profile_data) if profile_enabled else None
+
+                t_profile_model = self._profile_start() if profile_enabled else None
+                t_start_model = time.perf_counter()
+                step_metrics = self._train_step(batch_vla)
+                t_end_model = time.perf_counter()
+                profile_model_seconds = self._profile_elapsed(t_profile_model) if profile_enabled else None
+
+                optimizer_stepped = bool(step_metrics.pop("_optimizer_step"))
+                stop_requested = False
+                ran_mid_train_rl_eval = False
+                if optimizer_stepped:
+                    progress_bar.update(1)
+                    self.completed_steps += 1
+                    stop_requested = self._refresh_graceful_stop_requested()
+                    phase_ended = self._latency_curriculum_phase_end_at_step(self.completed_steps)
+                    curriculum_cfg = self._latency_curriculum_cfg()
+                    if not stop_requested:
+                        if (
+                            phase_ended
+                            and curriculum_cfg is not None
+                            and _as_bool(getattr(curriculum_cfg, "eval_at_phase_end", False), default=False)
+                        ):
+                            if self.accelerator.is_main_process:
+                                logger.info("Starting phase-end RL-games mid-train eval at step %s", self.completed_steps)
+                            step_metrics = self._run_mid_train_rl_games_eval(step_metrics)
+                            ran_mid_train_rl_eval = True
+                            if self.accelerator.is_main_process:
+                                logger.info("Finished phase-end RL-games mid-train eval at step %s", self.completed_steps)
+                        if (
+                            phase_ended
+                            and curriculum_cfg is not None
+                            and _as_bool(getattr(curriculum_cfg, "save_at_phase_end", False), default=False)
+                            and self._save_periodic_checkpoints_enabled()
+                        ):
+                            self._save_checkpoint()
+                        curriculum_metrics = self._apply_latency_curriculum()
+                        if curriculum_metrics:
+                            step_metrics.update(curriculum_metrics)
+                    if self._profile_timing_should_log():
+                        step_metrics["timing/dataloader_next"] = profile_data_seconds
+                        step_metrics["timing/train_step_total"] = profile_model_seconds
+                        if "batch/size" in step_metrics:
+                            step_metrics["throughput/samples_per_sec"] = step_metrics["batch/size"] / profile_model_seconds
+                        if "batch/effective_tokens" in step_metrics:
+                            step_metrics["throughput/effective_tokens_per_sec"] = (
+                                step_metrics["batch/effective_tokens"] / profile_model_seconds
+                            )
+
+                if self.accelerator.is_local_main_process:
+                    progress_bar.set_postfix(
+                        {
+                            "data_times": f"{t_end_data - t_start_data:.3f}",
+                            "model_times": f"{t_end_model - t_start_model:.3f}",
+                        }
+                    )
+
+                gradients_synced = optimizer_stepped
+                if gradients_synced and not stop_requested:
+                    if should_run_step_interval_event(
+                        completed_steps=self.completed_steps,
+                        interval=self.config.trainer.eval_interval,
+                        gradients_synced=gradients_synced,
+                    ):
+                        if self.accelerator.is_main_process:
+                            logger.info("Starting held-out action loss eval at step %s", self.completed_steps)
+                        t_profile_eval = self._profile_start()
+                        step_metrics = self.eval_action_loss(step_metrics)
+                        if self._profile_timing_should_log():
+                            step_metrics["timing/eval_action_loss_total"] = self._profile_elapsed(t_profile_eval)
+                        if self.accelerator.is_main_process:
+                            logger.info("Finished held-out action loss eval at step %s", self.completed_steps)
+                    action_classification_interval = getattr(
+                        self.config.trainer,
+                        "eval_action_classification_interval",
+                        None,
+                    )
+                    if action_classification_interval is None:
+                        action_classification_interval = self.config.trainer.eval_interval
+                    eval_action_classification_enabled = _as_bool(
+                        getattr(self.config.trainer, "eval_action_classification", True),
+                        default=True,
+                    )
+                    if eval_action_classification_enabled and should_run_step_interval_event(
+                        completed_steps=self.completed_steps,
+                        interval=action_classification_interval,
+                        gradients_synced=gradients_synced,
+                    ):
+                        if self.accelerator.is_main_process:
+                            logger.info("Starting held-out action classification eval at step %s", self.completed_steps)
+                        t_profile_eval = self._profile_start()
+                        step_metrics = self.eval_action_cc_f1(step_metrics)
+                        if self._profile_timing_should_log():
+                            step_metrics["timing/eval_action_classification_total"] = self._profile_elapsed(t_profile_eval)
+                        if self.accelerator.is_main_process:
+                            logger.info("Finished held-out action classification eval at step %s", self.completed_steps)
+                    if self._rl_games_eval_runner is not None:
+                        eval_every = self._rl_games_eval_runner.interval_steps()
+                        if eval_every > 0 and should_run_step_interval_event(
+                            completed_steps=self.completed_steps,
+                            interval=eval_every,
+                            gradients_synced=gradients_synced,
+                        ) and not ran_mid_train_rl_eval:
+                            step_metrics = self._run_mid_train_rl_games_eval(step_metrics)
+                elif gradients_synced and stop_requested and self.accelerator.is_main_process:
+                    logger.info("Skipping eval at step %s due to graceful stop", self.completed_steps)
+
+                step_metrics["timing/data"] = t_end_data - t_start_data
+                step_metrics["timing/model"] = t_end_model - t_start_model
+                if stop_requested:
+                    step_metrics["train/graceful_stop_requested"] = 1.0
+                if should_run_step_interval_event(
+                    completed_steps=self.completed_steps,
+                    interval=self.config.trainer.logging_frequency,
+                    gradients_synced=gradients_synced,
+                ):
+                    t_profile_log = self._profile_start() if self._profile_timing_should_log() else None
+                    self._log_metrics(step_metrics)
+                    if self._profile_timing_should_log() and self.accelerator.is_main_process:
+                        self._log_wandb({"timing/log_metrics_total": self._profile_elapsed(t_profile_log)})
+
+                if stop_requested:
+                    self._save_interrupt_checkpoint()
+                    break
+
+                if self._save_periodic_checkpoints_enabled() and should_run_optional_step_interval_event(
+                    completed_steps=self.completed_steps,
+                    interval=self.config.trainer.save_interval,
+                    gradients_synced=gradients_synced,
+                ):
+                    t_profile_checkpoint = self._profile_start() if self._profile_timing_should_log() else None
+                    self._save_checkpoint()
+                    if self._profile_timing_should_log() and self.accelerator.is_main_process:
+                        self._log_wandb({"timing/checkpoint_total": self._profile_elapsed(t_profile_checkpoint)})
+                    if (
+                        self._dagger_control_dir is not None
+                        and self.completed_steps < self.config.trainer.max_train_steps
+                        and not self._wait_for_dagger_round()
+                    ):
+                        break
+
+                if self.completed_steps >= stop_step:
+                    break
+        finally:
+            progress_bar.close()
+            if self.graceful_stop_requested:
+                self.cleanup_runtime()
+
+        if self.graceful_stop_requested:
+            if self.accelerator.is_main_process:
+                logger.info("Training stopped gracefully at step %s; skipping final/post-train eval.", self.completed_steps)
+            return
+
+        self._finalize_training()
+
+    def eval_action_loss(self, step_metrics: dict = None) -> dict:
+        """Compute held-out action loss over the validation LeRobot dataset."""
+        step_metrics = step_metrics or {}
+        if self._refresh_graceful_stop_requested():
+            if self.accelerator.is_main_process:
+                logger.info("Skipping held-out action loss eval at step %s due to graceful stop", self.completed_steps)
+            return step_metrics
+        if self.vla_eval_dataloader is None:
+            return step_metrics
+
+        was_training = self.model.training
+        self.model.eval()
+        eval_start = time.perf_counter()
+        total_loss_sum = torch.zeros((), device=self.accelerator.device, dtype=torch.float32)
+        total_loss_count = torch.zeros((), device=self.accelerator.device, dtype=torch.float32)
+        latency_loss_sums: dict[int, torch.Tensor] = {}
+        latency_loss_counts: dict[int, torch.Tensor] = {}
+        task_loss_sums: dict[str, torch.Tensor] = {}
+        task_loss_counts: dict[str, torch.Tensor] = {}
+        task_latency_loss_sums: dict[tuple[str, int], torch.Tensor] = {}
+        task_latency_loss_counts: dict[tuple[str, int], torch.Tensor] = {}
+        num_batches = int(getattr(self.config.trainer, "eval_num_batches", 20))
+        per_latency_batches = _optional_positive_int(getattr(self.config.trainer, "per_latency_eval_num_batches", None))
+        batch_size = int(self.config.datasets.vla_data.per_device_batch_size)
+        per_latency_sample_budget = per_latency_batches * batch_size if per_latency_batches is not None else None
+        seen_latency_counts: dict[int, int] = {}
+        seen_task_latency_counts: dict[tuple[str, int], int] = {}
+        eval_task = str(getattr(getattr(self.config, "rl_games", None), "task", "") or "").lower()
+        cross_task_mode = eval_task == "cross_task"
+        expected_latencies = (
+            _dataset_latency_values(getattr(self.vla_eval_dataloader, "dataset", None))
+            if per_latency_sample_budget is not None
+            else []
+        )
+        expected_task_latencies = (
+            [
+                (task_name, latency)
+                for task_name, latencies in _dataset_task_latency_values(
+                    getattr(self.vla_eval_dataloader, "dataset", None)
+                ).items()
+                for latency in latencies
+            ]
+            if cross_task_mode and per_latency_sample_budget is not None
+            else []
+        )
+        use_per_latency_budget = per_latency_sample_budget is not None and bool(expected_latencies)
+        with torch.no_grad():
+            for batch_idx, batch_vla in enumerate(self.vla_eval_dataloader):
+                if not use_per_latency_budget and batch_idx >= num_batches:
+                    break
+                batch_examples = list(batch_vla)
+                grouped = _group_examples_by_task_latency(batch_examples)
+                for (task_name, latency), examples in grouped.items():
+                    if latency is not None and use_per_latency_budget:
+                        if cross_task_mode and task_name is not None:
+                            task_latency_key = (str(task_name), int(latency))
+                            remaining = per_latency_sample_budget - seen_task_latency_counts.get(task_latency_key, 0)
+                        else:
+                            task_latency_key = None
+                            remaining = per_latency_sample_budget - seen_latency_counts.get(int(latency), 0)
+                        if remaining <= 0:
+                            continue
+                        examples = examples[:remaining]
+                        if task_latency_key is not None:
+                            seen_task_latency_counts[task_latency_key] = (
+                                seen_task_latency_counts.get(task_latency_key, 0) + len(examples)
+                            )
+                        else:
+                            seen_latency_counts[int(latency)] = seen_latency_counts.get(int(latency), 0) + len(examples)
+                    if not examples:
+                        continue
+                    with torch.autocast("cuda", dtype=torch.bfloat16):
+                        output_dict = self.model.forward(examples)
+                        action_loss = output_dict["action_loss"].detach().float()
+                    weight = torch.tensor(float(len(examples)), device=self.accelerator.device, dtype=torch.float32)
+                    total_loss_sum += action_loss * weight
+                    total_loss_count += weight
+                    if latency is not None:
+                        latency_loss_sums.setdefault(
+                            int(latency),
+                            torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                        )
+                        latency_loss_counts.setdefault(
+                            int(latency),
+                            torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                        )
+                        latency_loss_sums[int(latency)] += action_loss * weight
+                        latency_loss_counts[int(latency)] += weight
+                    if task_name is not None:
+                        task_name = str(task_name)
+                        task_loss_sums.setdefault(
+                            task_name,
+                            torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                        )
+                        task_loss_counts.setdefault(
+                            task_name,
+                            torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                        )
+                        task_loss_sums[task_name] += action_loss * weight
+                        task_loss_counts[task_name] += weight
+                        if latency is not None:
+                            task_latency_key = (task_name, int(latency))
+                            task_latency_loss_sums.setdefault(
+                                task_latency_key,
+                                torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                            )
+                            task_latency_loss_counts.setdefault(
+                                task_latency_key,
+                                torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                            )
+                            task_latency_loss_sums[task_latency_key] += action_loss * weight
+                            task_latency_loss_counts[task_latency_key] += weight
+                if (
+                    use_per_latency_budget
+                    and (
+                        (
+                            cross_task_mode
+                            and expected_task_latencies
+                            and all(
+                                seen_task_latency_counts.get(task_latency, 0) >= per_latency_sample_budget
+                                for task_latency in expected_task_latencies
+                            )
+                        )
+                        or (
+                            not cross_task_mode
+                            and expected_latencies
+                            and all(
+                                seen_latency_counts.get(latency, 0) >= per_latency_sample_budget
+                                for latency in expected_latencies
+                            )
+                        )
+                    )
+                ):
+                    break
+
+        if was_training:
+            self.model.train()
+
+        if dist.is_initialized():
+            dist.all_reduce(total_loss_sum, op=dist.ReduceOp.SUM)
+            dist.all_reduce(total_loss_count, op=dist.ReduceOp.SUM)
+            latency_keys = sorted(latency_loss_sums)
+            gathered_keys = [None for _ in range(dist.get_world_size())]
+            dist.all_gather_object(gathered_keys, latency_keys)
+            all_latency_keys = sorted({int(key) for keys in gathered_keys for key in keys})
+            for latency in all_latency_keys:
+                loss_sum = latency_loss_sums.get(
+                    latency,
+                    torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                )
+                loss_count = latency_loss_counts.get(
+                    latency,
+                    torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                )
+                dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
+                dist.all_reduce(loss_count, op=dist.ReduceOp.SUM)
+                latency_loss_sums[latency] = loss_sum
+                latency_loss_counts[latency] = loss_count
+            task_keys = sorted(task_loss_sums)
+            gathered_task_keys = [None for _ in range(dist.get_world_size())]
+            dist.all_gather_object(gathered_task_keys, task_keys)
+            all_task_keys = sorted({str(key) for keys in gathered_task_keys for key in keys})
+            for task_name in all_task_keys:
+                loss_sum = task_loss_sums.get(
+                    task_name,
+                    torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                )
+                loss_count = task_loss_counts.get(
+                    task_name,
+                    torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                )
+                dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
+                dist.all_reduce(loss_count, op=dist.ReduceOp.SUM)
+                task_loss_sums[task_name] = loss_sum
+                task_loss_counts[task_name] = loss_count
+            task_latency_keys = sorted(task_latency_loss_sums)
+            gathered_task_latency_keys = [None for _ in range(dist.get_world_size())]
+            dist.all_gather_object(gathered_task_latency_keys, task_latency_keys)
+            all_task_latency_keys = sorted(
+                {(str(task_name), int(latency)) for keys in gathered_task_latency_keys for task_name, latency in keys}
+            )
+            for task_latency_key in all_task_latency_keys:
+                loss_sum = task_latency_loss_sums.get(
+                    task_latency_key,
+                    torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                )
+                loss_count = task_latency_loss_counts.get(
+                    task_latency_key,
+                    torch.zeros((), device=self.accelerator.device, dtype=torch.float32),
+                )
+                dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
+                dist.all_reduce(loss_count, op=dist.ReduceOp.SUM)
+                task_latency_loss_sums[task_latency_key] = loss_sum
+                task_latency_loss_counts[task_latency_key] = loss_count
+
+        if self.accelerator.is_main_process and total_loss_count.item() > 0:
+            step_metrics["eval/action_loss/seconds"] = time.perf_counter() - eval_start
+            step_metrics["eval/action_loss/samples"] = total_loss_count.item()
+            step_metrics["eval/loss"] = (total_loss_sum / total_loss_count.clamp_min(1.0)).item()
+            for latency in sorted(latency_loss_sums):
+                count = latency_loss_counts[latency]
+                if count.item() > 0:
+                    step_metrics[f"eval/latency_{latency}/loss"] = (
+                        latency_loss_sums[latency] / count.clamp_min(1.0)
+                    ).item()
+            for task_name in sorted(task_loss_sums):
+                count = task_loss_counts[task_name]
+                if count.item() > 0:
+                    step_metrics[f"eval/{task_name}/loss"] = (
+                        task_loss_sums[task_name] / count.clamp_min(1.0)
+                    ).item()
+            for task_name, latency in sorted(task_latency_loss_sums):
+                count = task_latency_loss_counts[(task_name, latency)]
+                if count.item() > 0:
+                    step_metrics[f"eval/{task_name}/latency_{latency}/loss"] = (
+                        task_latency_loss_sums[(task_name, latency)] / count.clamp_min(1.0)
+                    ).item()
+
+        _distributed_barrier()
+        return step_metrics
+
+    def eval_action_model(self, step_metrics: dict = None) -> dict:
+        """Backward-compatible alias for the held-out action loss evaluator."""
+        return self.eval_action_loss(step_metrics)
+
+    def eval_action_cc_f1(self, step_metrics: dict = None) -> dict:
+        """Component-based Control-Critical F1 over the held-out teacher set.
+
+        Unlike ``eval_action_loss`` (which samples frames randomly via the
+        shuffled eval dataloader), CC-F1 needs *contiguous, fully-covered*
+        episodes so the ``±K`` temporal-tolerance matching is meaningful. So this
+        runs a dedicated sequential pass: it walks each underlying LeRobot
+        dataset's ``all_steps`` in order (bypassing the random mixture sampler),
+        shards whole episodes across ranks, runs ``predict_action`` (the same
+        decode path as rollout), decodes per-frame action *components* and matches
+        teacher vs model events per episode. Counts are all-reduced and reduced to
+        CC-F1 on the main process. See ``action_cc_f1`` for the metric definition.
+        """
+        step_metrics = step_metrics or {}
+        if self._refresh_graceful_stop_requested():
+            if self.accelerator.is_main_process:
+                logger.info(
+                    "Skipping held-out action classification eval at step %s due to graceful stop",
+                    self.completed_steps,
+                )
+            return step_metrics
+        if self.vla_eval_dataloader is None:
+            return step_metrics
+        task = str(getattr(getattr(self.config, "rl_games", None), "task", "") or "").lower()
+        cross_task_mode = task == "cross_task"
+        if not cross_task_mode and task not in action_cc_f1.SUPPORTED_TASKS:
+            return step_metrics
+        if not bool(getattr(self.config.trainer, "eval_action_classification", True)):
+            return step_metrics
+
+        # Resolve spec (deadly_corridor depends on its action layout) and tolerance.
+        # Read the layout exactly like action_spec._deadly_action_dim (same default)
+        # so the decode matches the action dim the model was actually trained with.
+        deadly_layout = action_cc_f1.DEADLY_MULTIBINARY_7
+        deadly_cfg = getattr(getattr(getattr(self.config, "rl_games", None), "env_eval", None), "deadly", None)
+        if deadly_cfg is not None:
+            deadly_layout = str(getattr(deadly_cfg, "action_layout", action_cc_f1.DEADLY_MULTIBINARY_7))
+
+        def _spec_for_task(task_name: str):
+            return action_cc_f1.get_spec(task_name, deadly_layout)
+
+        def _tolerance_for_task(task_name: str, spec) -> int:
+            k = spec.default_k
+            if task_name != "flappy":  # flappy stays per-frame (K=0) to match the shipped flap-F1
+                override = getattr(self.config.trainer, "cc_f1_tolerance", None)
+                if override is not None:
+                    k = int(override)
+            return k
+
+        if not cross_task_mode:
+            spec = _spec_for_task(task)
+            k = _tolerance_for_task(task, spec)
+        else:
+            spec = None
+            k = None
+
+        dataset = getattr(self.vla_eval_dataloader, "dataset", None)
+        if dataset is None:
+            return step_metrics
+        single_datasets = list(getattr(dataset, "datasets", None) or [dataset])
+
+        # Enumerate episodes as contiguous same-trajectory runs of all_steps.
+        # episodes: list of (task, dataset_index, episode_key, latency, [(flat_idx, base_index), ...]).
+        episodes = []
+        for dataset_index, ds in enumerate(single_datasets):
+            ds_task = str(getattr(ds, "rl_games_task", task) or task)
+            if ds_task not in action_cc_f1.SUPPORTED_TASKS:
+                continue
+            all_steps = getattr(ds, "all_steps", None)
+            if not all_steps:
+                continue
+            tag = getattr(ds, "tag", "ds")
+            cur_traj = None
+            cur = []
+            for flat_idx, step in enumerate(all_steps):
+                traj_id, base_idx = step
+                if cur and traj_id != cur_traj:
+                    episodes.append((ds_task, dataset_index, f"{ds_task}:{tag}:{cur_traj}", _episode_latency(ds, int(cur_traj)), cur))
+                    cur = []
+                cur_traj = traj_id
+                cur.append((int(flat_idx), int(base_idx)))
+            if cur:
+                episodes.append((ds_task, dataset_index, f"{ds_task}:{tag}:{cur_traj}", _episode_latency(ds, int(cur_traj)), cur))
+        if not episodes:
+            return step_metrics
+
+        # Shard whole episodes across ranks, bounded by a per-rank frame budget.
+        num_procs = int(self.accelerator.num_processes)
+        rank = int(self.accelerator.process_index)
+        bs = int(self.config.datasets.vla_data.per_device_batch_size)
+        per_latency_batches = _optional_positive_int(getattr(self.config.trainer, "per_latency_eval_num_batches", None))
+        shared_frame_budget = max(1, int(getattr(self.config.trainer, "eval_num_batches", 20)) * bs)
+        per_latency_frame_budget = per_latency_batches * bs if per_latency_batches is not None else None
+        assigned = []
+        frame_count = 0
+        latency_frame_counts: dict[int, int] = {}
+        task_latency_frame_counts: dict[tuple[str, int], int] = {}
+        expected_latencies = sorted({int(ep[3]) for ep in episodes if ep[3] is not None})
+        expected_task_latencies = (
+            sorted({(str(ep[0]), int(ep[3])) for ep in episodes if ep[3] is not None})
+            if cross_task_mode
+            else []
+        )
+        use_per_latency_budget = per_latency_frame_budget is not None and bool(expected_latencies)
+        for ep_idx, ep in enumerate(episodes):
+            if ep_idx % num_procs != rank:
+                continue
+            ep_task = str(ep[0])
+            latency = ep[3]
+            if latency is not None and use_per_latency_budget:
+                latency = int(latency)
+                if cross_task_mode:
+                    task_latency_key = (ep_task, latency)
+                    if task_latency_frame_counts.get(task_latency_key, 0) >= per_latency_frame_budget:
+                        continue
+                else:
+                    if latency_frame_counts.get(latency, 0) >= per_latency_frame_budget:
+                        continue
+            assigned.append(ep)
+            frame_count += len(ep[4])
+            if latency is not None and use_per_latency_budget:
+                if cross_task_mode:
+                    task_latency_key = (ep_task, int(latency))
+                    task_latency_frame_counts[task_latency_key] = (
+                        task_latency_frame_counts.get(task_latency_key, 0) + len(ep[4])
+                    )
+                    if all(
+                        task_latency_frame_counts.get(value, 0) >= per_latency_frame_budget
+                        for value in expected_task_latencies
+                    ):
+                        break
+                else:
+                    latency_frame_counts[int(latency)] = latency_frame_counts.get(int(latency), 0) + len(ep[4])
+                    if all(
+                        latency_frame_counts.get(value, 0) >= per_latency_frame_budget
+                        for value in expected_latencies
+                    ):
+                        break
+            elif not use_per_latency_budget and frame_count >= shared_frame_budget:
+                break
+
+        was_training = self.model.training
+        self.model.eval()
+        unwrapped = self.accelerator.unwrap_model(self.model)
+        counts = action_cc_f1.new_counts(spec) if spec is not None else None
+
+        per_ep = {}  # episode_key -> list of (base_index, teacher_components, model_components)
+
+        def flush_chunk(chunk):
+            normalized = unwrapped.predict_action(examples=[c[4] for c in chunk])["normalized_actions"]
+            for i, (sample_task, episode_key, base_idx, latency, sample) in enumerate(chunk):
+                sample_task = str(sample_task)
+                sample_spec = _spec_for_task(sample_task)
+                teacher_vec = np.asarray(sample["action"])[0, :]
+                model_vec = normalized[i, 0, :]
+                entry = per_ep.setdefault(episode_key, {"task": sample_task, "latency": latency, "items": []})
+                entry["items"].append((base_idx, sample_spec.comp_fn(teacher_vec), sample_spec.comp_fn(model_vec)))
+
+        frame_records = []
+        for ep_task, dataset_index, episode_key, episode_latency_value, frames in assigned:
+            for flat_idx, base_idx in frames:
+                frame_records.append((
+                    int(dataset_index),
+                    int(flat_idx),
+                    str(episode_key),
+                    int(base_idx),
+                    str(ep_task),
+                    episode_latency_value,
+                ))
+
+        eval_num_workers = getattr(self.config.datasets.vla_data, "eval_num_workers", None)
+        if eval_num_workers is None:
+            eval_num_workers = self.config.datasets.vla_data.num_workers
+        eval_num_workers = int(eval_num_workers)
+        dataloader_kwargs = build_cpu_only_dataloader_kwargs(
+            eval_num_workers,
+            pin_memory=self.config.datasets.vla_data.pin_memory,
+            persistent_workers=self.config.datasets.vla_data.persistent_workers,
+            prefetch_factor=(
+                self.config.datasets.vla_data.prefetch_factor
+                if "prefetch_factor" in self.config.datasets.vla_data
+                else None
+            ),
+        )
+        frame_loader_key = (tuple(frame_records), bs, eval_num_workers)
+        if self._action_cc_f1_frame_loader_key != frame_loader_key:
+            if self.accelerator.is_main_process:
+                logger.info(
+                    "Created action CC-F1 frame dataloader with num_workers=%s, persistent_workers=%s, cpu_only_workers=%s",
+                    eval_num_workers,
+                    dataloader_kwargs.get("persistent_workers", False),
+                    "multiprocessing_context" in dataloader_kwargs,
+                )
+            self._action_cc_f1_frame_loader = DataLoader(
+                _ActionCCF1FrameDataset(single_datasets, frame_records),
+                batch_size=bs,
+                shuffle=False,
+                collate_fn=_identity_collate,
+                num_workers=eval_num_workers,
+                **dataloader_kwargs,
+            )
+            self._action_cc_f1_frame_loader_key = frame_loader_key
+        frame_loader = self._action_cc_f1_frame_loader
+
+        data_wait_seconds = 0.0
+        predict_seconds = 0.0
+        processed_frames = 0
+        with torch.no_grad():
+            t_data_start = time.perf_counter()
+            for chunk in frame_loader:
+                t_data_end = time.perf_counter()
+                data_wait_seconds += t_data_end - t_data_start
+                t_predict_start = time.perf_counter()
+                flush_chunk(chunk)
+                predict_seconds += time.perf_counter() - t_predict_start
+                processed_frames += len(chunk)
+                t_data_start = time.perf_counter()
+
+        if was_training:
+            self.model.train()
+
+        latency_counts: dict[int, np.ndarray] = {}
+        task_counts: dict[str, np.ndarray] = {}
+        task_latency_counts: dict[tuple[str, int], np.ndarray] = {}
+        for entry in per_ep.values():
+            entry_task = str(entry.get("task", task))
+            entry_spec = _spec_for_task(entry_task)
+            entry_k = _tolerance_for_task(entry_task, entry_spec)
+            items = entry["items"]
+            items.sort(key=lambda r: r[0])
+            target_counts = task_counts.setdefault(entry_task, action_cc_f1.new_counts(entry_spec)) if cross_task_mode else counts
+            action_cc_f1.accumulate_episode(
+                entry_spec,
+                frames=[r[0] for r in items],
+                teacher_comps=[r[1] for r in items],
+                model_comps=[r[2] for r in items],
+                k=entry_k,
+                counts=target_counts,
+            )
+            latency = entry.get("latency")
+            if latency is not None:
+                latency = int(latency)
+                if cross_task_mode:
+                    task_latency_key = (entry_task, latency)
+                    task_latency_counts.setdefault(task_latency_key, action_cc_f1.new_counts(entry_spec))
+                    target_latency_counts = task_latency_counts[task_latency_key]
+                else:
+                    latency_counts.setdefault(latency, action_cc_f1.new_counts(entry_spec))
+                    target_latency_counts = latency_counts[latency]
+                action_cc_f1.accumulate_episode(
+                    entry_spec,
+                    frames=[r[0] for r in items],
+                    teacher_comps=[r[1] for r in items],
+                    model_comps=[r[2] for r in items],
+                    k=entry_k,
+                    counts=target_latency_counts,
+                )
+
+        # Counts are small integers (well within float32's exact range); float32
+        # keeps NCCL all_reduce broadly compatible across torch builds.
+        if not cross_task_mode:
+            counts_t = torch.tensor(counts, device=self.accelerator.device, dtype=torch.float32)
+            if dist.is_initialized():
+                dist.all_reduce(counts_t, op=dist.ReduceOp.SUM)
+            if self.accelerator.is_main_process:
+                step_metrics.update(action_cc_f1.reduce_metrics(spec, counts_t.detach().cpu().numpy()))
+
+            if dist.is_initialized():
+                latency_keys = sorted(latency_counts)
+                gathered_keys = [None for _ in range(dist.get_world_size())]
+                dist.all_gather_object(gathered_keys, latency_keys)
+                all_latency_keys = sorted({int(key) for keys in gathered_keys for key in keys})
+            else:
+                all_latency_keys = sorted(latency_counts)
+
+            for latency in all_latency_keys:
+                latency_counts_t = torch.tensor(
+                    latency_counts.get(latency, action_cc_f1.new_counts(spec)),
+                    device=self.accelerator.device,
+                    dtype=torch.float32,
+                )
+                if dist.is_initialized():
+                    dist.all_reduce(latency_counts_t, op=dist.ReduceOp.SUM)
+                if self.accelerator.is_main_process:
+                    latency_metrics = action_cc_f1.reduce_metrics(spec, latency_counts_t.detach().cpu().numpy())
+                    for key, value in latency_metrics.items():
+                        if key.startswith("eval/"):
+                            key = f"eval/latency_{int(latency)}/{key[len('eval/'):]}"
+                        else:
+                            key = f"eval/latency_{int(latency)}/{key}"
+                        step_metrics[key] = value
+        else:
+            if dist.is_initialized():
+                task_keys = sorted(task_counts)
+                gathered_task_keys = [None for _ in range(dist.get_world_size())]
+                dist.all_gather_object(gathered_task_keys, task_keys)
+                all_task_keys = sorted({str(key) for keys in gathered_task_keys for key in keys})
+            else:
+                all_task_keys = sorted(task_counts)
+
+            for task_name in all_task_keys:
+                task_spec = _spec_for_task(task_name)
+                counts_t = torch.tensor(
+                    task_counts.get(task_name, action_cc_f1.new_counts(task_spec)),
+                    device=self.accelerator.device,
+                    dtype=torch.float32,
+                )
+                if dist.is_initialized():
+                    dist.all_reduce(counts_t, op=dist.ReduceOp.SUM)
+                if self.accelerator.is_main_process:
+                    metrics = action_cc_f1.reduce_metrics(task_spec, counts_t.detach().cpu().numpy())
+                    for key, value in metrics.items():
+                        if key.startswith("eval/"):
+                            key = f"eval/{task_name}/{key[len('eval/'):]}"
+                        else:
+                            key = f"eval/{task_name}/{key}"
+                        step_metrics[key] = value
+
+            if dist.is_initialized():
+                task_latency_keys = sorted(task_latency_counts)
+                gathered_task_latency_keys = [None for _ in range(dist.get_world_size())]
+                dist.all_gather_object(gathered_task_latency_keys, task_latency_keys)
+                all_task_latency_keys = sorted(
+                    {(str(task_name), int(latency)) for keys in gathered_task_latency_keys for task_name, latency in keys}
+                )
+            else:
+                all_task_latency_keys = sorted(task_latency_counts)
+
+            for task_name, latency in all_task_latency_keys:
+                task_spec = _spec_for_task(task_name)
+                counts_t = torch.tensor(
+                    task_latency_counts.get((task_name, latency), action_cc_f1.new_counts(task_spec)),
+                    device=self.accelerator.device,
+                    dtype=torch.float32,
+                )
+                if dist.is_initialized():
+                    dist.all_reduce(counts_t, op=dist.ReduceOp.SUM)
+                if self.accelerator.is_main_process:
+                    latency_metrics = action_cc_f1.reduce_metrics(task_spec, counts_t.detach().cpu().numpy())
+                    for key, value in latency_metrics.items():
+                        if key.startswith("eval/"):
+                            key = f"eval/{task_name}/latency_{int(latency)}/{key[len('eval/'):]}"
+                        else:
+                            key = f"eval/{task_name}/latency_{int(latency)}/{key}"
+                        step_metrics[key] = value
+
+        data_wait_t = torch.tensor(data_wait_seconds, device=self.accelerator.device, dtype=torch.float32)
+        predict_t = torch.tensor(predict_seconds, device=self.accelerator.device, dtype=torch.float32)
+        frames_t = torch.tensor(float(processed_frames), device=self.accelerator.device, dtype=torch.float32)
+        if dist.is_initialized():
+            dist.all_reduce(data_wait_t, op=dist.ReduceOp.MAX)
+            dist.all_reduce(predict_t, op=dist.ReduceOp.MAX)
+            dist.all_reduce(frames_t, op=dist.ReduceOp.SUM)
+        if self.accelerator.is_main_process:
+            step_metrics["eval/action_classification/data_wait_seconds"] = data_wait_t.item()
+            step_metrics["eval/action_classification/predict_seconds"] = predict_t.item()
+            step_metrics["eval/action_classification/frames"] = frames_t.item()
+
+        _distributed_barrier()
+        return step_metrics
+
+    def _log_training_config(self):
+        """Record training config."""
+        if self.accelerator.is_main_process:
+            logger.info("***** Training Configuration *****")
+            logger.info(f"  Total optimization steps = {self.config.trainer.max_train_steps}")
+            logger.info(f"  Per device batch size = {self.config.datasets.vla_data.per_device_batch_size}")
+            logger.info(f"  Gradient accumulation steps = {self.accelerator.gradient_accumulation_steps}")
+            logger.info(f"  Total batch size = {self.total_batch_size}")
+
+    def _train_step(self, batch_vla, batch_vlm=None):
+        """Execute single training step."""
+        profile_next_step = self.completed_steps + 1
+        profile_log = self._profile_timing_should_log(profile_next_step)
+        profile_metrics = {}
+        if self.accelerator.distributed_type == DistributedType.DEEPSPEED:
+            t_forward = self._profile_start() if profile_log else None
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                output_dict = self.model.forward(batch_vla)
+                action_loss = output_dict["action_loss"]
+                loss_weight = output_dict["loss_weight"]
+                total_loss = action_loss
+            if profile_log:
+                profile_metrics["timing/forward"] = self._profile_elapsed(t_forward)
+                profile_metrics.update(output_dict["timing"])
+                profile_metrics.update(output_dict["batch_stats"])
+
+            t_backward = self._profile_start() if profile_log else None
+            self.model.backward(total_loss)
+            self._record_train_loss(action_loss, loss_weight)
+            if profile_log:
+                profile_metrics["timing/backward"] = self._profile_elapsed(t_backward)
+
+            optimizer_stepped = bool(self.model.is_gradient_accumulation_boundary())
+            t_optimizer = self._profile_start() if profile_log else None
+            self.model.step()
+            if profile_log:
+                profile_metrics["timing/optimizer_step"] = self._profile_elapsed(t_optimizer)
+
+            t_scheduler = self._profile_start() if profile_log else None
+            if optimizer_stepped:
+                self.lr_scheduler.step()
+            if profile_log:
+                profile_metrics["timing/lr_scheduler"] = self._profile_elapsed(t_scheduler)
+
+            if not optimizer_stepped:
+                return {"_optimizer_step": False}
+            grad_norm = self.model.get_global_grad_norm()
+            metrics = {
+                "train/loss": self._finalize_train_loss(action_loss.device),
+                "_optimizer_step": True,
+            }
+            metrics.update(profile_metrics)
+            if isinstance(grad_norm, torch.Tensor):
+                grad_norm = grad_norm.detach().float().item()
+            metrics["train/grad_norm_pre_clip"] = float(grad_norm)
+            return metrics
+
+        with self.accelerator.accumulate(self.model):
+            t_forward = self._profile_start() if profile_log else None
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                output_dict = self.model.forward(batch_vla)
+                action_loss = output_dict["action_loss"]
+                loss_weight = output_dict["loss_weight"]
+                total_loss = action_loss
+            if profile_log:
+                profile_metrics["timing/forward"] = self._profile_elapsed(t_forward)
+                profile_metrics.update(output_dict["timing"])
+                profile_metrics.update(output_dict["batch_stats"])
+
+            t_backward = self._profile_start() if profile_log else None
+            self.accelerator.backward(total_loss)
+            self._record_train_loss(action_loss, loss_weight)
+            if profile_log:
+                profile_metrics["timing/backward"] = self._profile_elapsed(t_backward)
+
+            grad_norm = None
+            gradients_synced = self.accelerator.sync_gradients
+            t_grad = self._profile_start() if profile_log else None
+            if gradients_synced and self.config.trainer.gradient_clipping is not None:
+                grad_norm = self.accelerator.clip_grad_norm_(self.model.parameters(), self.config.trainer.gradient_clipping)
+            elif gradients_synced:
+                grad_norm = self._total_grad_norm(self.model.parameters())
+            if profile_log:
+                profile_metrics["timing/grad_clip_or_norm"] = self._profile_elapsed(t_grad)
+
+            t_optimizer = self._profile_start() if profile_log else None
+            self.optimizer.step()
+            if profile_log:
+                profile_metrics["timing/optimizer_step"] = self._profile_elapsed(t_optimizer)
+            # Only step the LR scheduler when gradients are actually synced
+            # (i.e., not mid-accumulation). Without this guard the scheduler
+            # runs gradient_accumulation_steps times faster than intended,
+            # causing warmup to end too early and cosine decay to bottom out
+            # at min_lr well before max_train_steps is reached.
+            t_scheduler = self._profile_start() if profile_log else None
+            if gradients_synced:
+                self.lr_scheduler.step()
+            if profile_log:
+                profile_metrics["timing/lr_scheduler"] = self._profile_elapsed(t_scheduler)
+            t_zero_grad = self._profile_start() if profile_log else None
+            self.optimizer.zero_grad()
+            if profile_log:
+                profile_metrics["timing/zero_grad"] = self._profile_elapsed(t_zero_grad)
+
+        if not gradients_synced:
+            return {"_optimizer_step": False}
+        metrics = {
+            "train/loss": self._finalize_train_loss(action_loss.device),
+            "_optimizer_step": True,
+        }
+        metrics.update(profile_metrics)
+        if grad_norm is not None:
+            if isinstance(grad_norm, torch.Tensor):
+                grad_norm = grad_norm.detach().float().item()
+            metrics["train/grad_norm_pre_clip"] = float(grad_norm)
+        return metrics
+
+    def _finalize_training(self):
+        """Training end processing."""
+        save_interval = int(getattr(self.config.trainer, "save_interval", 0) or 0)
+        final_step_already_saved = (
+            self._save_periodic_checkpoints_enabled()
+            and self.completed_steps > 0
+            and save_interval > 0
+            and self.completed_steps % save_interval == 0
+        )
+        if (
+            self._save_final_model_enabled
+            and self.completed_steps > 0
+            and not final_step_already_saved
+        ):
+            self._save_checkpoint()
+
+        if self.accelerator.is_main_process:
+            if self._save_best_model_enabled:
+                logger.info(
+                    "Training complete. Best checkpoint step=%s score=%s path=%s",
+                    self._best_step,
+                    self._best_score,
+                    self._best_state_path,
+                )
+            else:
+                logger.info(f"Training complete. Final training state is saved under {self.checkpoint_dir}")
+
+        if (
+            self._save_best_model_enabled
+            and self._best_state_path
+            and os.path.isdir(self._best_state_path)
+        ):
+            self.accelerator.print(
+                f"Loading best checkpoint for post-train eval: "
+                f"{self._best_state_path} (step={self._best_step}, score={self._best_score:.6f})"
+            )
+            self._load_checkpoint(self._best_state_path)
+        elif self._save_best_model_enabled:
+            logger.warning(
+                "checkpoint.save_best_model=true but no best_state checkpoint was found; "
+                "post-train eval will use the current in-memory model."
+            )
+        self.accelerator.wait_for_everyone()
+
+        if self._rl_games_eval_runner is not None and self._rl_games_eval_runner.is_enabled(stage="post_train"):
+            if self._distributed_rl_games_eval_enabled():
+                eval_result = self._run_distributed_rl_games_eval(stage="post_train")
+            elif self.accelerator.is_main_process:
+                eval_result = self._run_rl_games_eval_with_model_mode(stage="post_train")
+            else:
+                eval_result = None
+
+            if self.accelerator.is_main_process and eval_result is not None:
+                final_metrics = self._append_rl_games_eval_metrics(
+                    step_metrics={},
+                    eval_result=eval_result,
+                    stage="post_train",
+                )
+                if self._save_best_model_enabled:
+                    final_metrics["checkpoint/best_score"] = float(self._best_score)
+                    final_metrics["checkpoint/best_step"] = float(self._best_step)
+                self._checkpoint_sync_manager.sync_eval_result(
+                    eval_path=eval_result.path,
+                    stage="post_train",
+                    step=self.completed_steps,
+                )
+                self._log_wandb(final_metrics)
+
+        if self.accelerator.is_main_process:
+            wandb.finish()
+
+        self.accelerator.wait_for_everyone()
+
+
+def _run_eval_batch_benchmark(trainer) -> None:
+    """Benchmark rl_games (latency_bench) eval wall-clock + VRAM across batch sizes.
+
+    Gated by env var LB_EVAL_BENCH=1. The plan is a comma list of
+    ``parallel:episodes`` pairs in LB_EVAL_BENCH_PLAN (default
+    ``1:4,5:20,10:20,20:20``). Model is loaded ONCE; every config runs against
+    the same live model so the comparison is apples-to-apples. The serial 1:4
+    row is meant to be linearly extrapolated to 20 eps (4x serial is too slow).
+    """
+    import os
+    import time
+
+    plan_str = os.environ.get("LB_EVAL_BENCH_PLAN", "1:4,5:20,10:20,20:20")
+    plan = []
+    for item in plan_str.split(","):
+        p, e = item.strip().split(":")
+        plan.append((int(p), int(e)))
+
+    try:
+        import pynvml
+
+        pynvml.nvmlInit()
+        _nvml_handle = pynvml.nvmlDeviceGetHandleByIndex(torch.cuda.current_device())
+    except Exception:
+        pynvml = None
+        _nvml_handle = None
+
+    def _nvml_used_mb():
+        if _nvml_handle is None:
+            return None
+        return pynvml.nvmlDeviceGetMemoryInfo(_nvml_handle).used / 1024 / 1024
+
+    runner = trainer._rl_games_eval_runner
+    model = trainer.accelerator.unwrap_model(trainer.model)
+
+    def _set(parallel, episodes):
+        # The trainer cfg is an access-tracked wrapper that OmegaConf.update can't
+        # touch, so override the runner's per-stage getters directly instead.
+        runner._eval_parallel_envs = lambda *a, **k: int(parallel)
+        runner._num_episodes = lambda *a, **k: int(episodes)
+
+    import threading
+
+    def _timed_run(parallel, episodes, label):
+        _set(parallel, episodes)
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+        nvml_peak = {"v": 0.0}
+        stop = threading.Event()
+
+        def _sample():
+            while not stop.is_set():
+                u = _nvml_used_mb()
+                if u is not None and u > nvml_peak["v"]:
+                    nvml_peak["v"] = u
+                time.sleep(0.1)
+
+        t = threading.Thread(target=_sample, daemon=True)
+        t.start()
+        t0 = time.perf_counter()
+        result = runner.run(model=model, step=0, stage="mid_train", save=False)
+        torch.cuda.synchronize()
+        dt = time.perf_counter() - t0
+        stop.set()
+        t.join(timeout=1.0)
+        torch_peak = torch.cuda.max_memory_reserved() / 1024 / 1024
+        mean_r = result.aggregate.get("mean_reward")
+        logger.info(
+            "[LB_EVAL_BENCH] %s | parallel=%d eps=%d | wall=%.1fs | torch_reserved_peak=%.0fMB | "
+            "nvml_used_peak=%sMB | mean_reward=%s",
+            label, parallel, episodes, dt, torch_peak,
+            f"{nvml_peak['v']:.0f}" if _nvml_handle is not None else "n/a", mean_r,
+        )
+        return {"parallel": parallel, "episodes": episodes, "wall_s": dt,
+                "torch_peak_mb": torch_peak, "nvml_peak_mb": nvml_peak["v"], "mean_reward": mean_r}
+
+    logger.info("[LB_EVAL_BENCH] warmup (parallel=1 eps=2, untimed) to trigger CUDA init/autotune ...")
+    _set(1, 2)
+    runner.run(model=model, step=0, stage="mid_train", save=False)
+    torch.cuda.synchronize()
+
+    rows = []
+    for parallel, episodes in plan:
+        rows.append(_timed_run(parallel, episodes, label="bench"))
+
+    # Summary table with serial-extrapolation + speedup vs the 1:* serial row.
+    serial = next((r for r in rows if r["parallel"] == 1), None)
+    serial_per_ep = (serial["wall_s"] / serial["episodes"]) if serial else None
+    logger.info("[LB_EVAL_BENCH] ==== SUMMARY ====")
+    logger.info("[LB_EVAL_BENCH] parallel | eps | wall_s | s/ep | extrap_20ep_s | speedup_vs_serial20 | torch_peak_MB | nvml_peak_MB | mean_reward")
+    serial20 = serial_per_ep * 20 if serial_per_ep else None
+    for r in rows:
+        s_per_ep = r["wall_s"] / r["episodes"]
+        extrap20 = s_per_ep * 20
+        speedup = (serial20 / extrap20) if serial20 else float("nan")
+        logger.info(
+            "[LB_EVAL_BENCH] %8d | %3d | %6.1f | %5.2f | %12.1f | %18.2fx | %12.0f | %11.0f | %s",
+            r["parallel"], r["episodes"], r["wall_s"], s_per_ep, extrap20, speedup,
+            r["torch_peak_mb"], r["nvml_peak_mb"], r["mean_reward"],
+        )
+    if serial20 is not None:
+        logger.info("[LB_EVAL_BENCH] serial 20-ep estimate (1:4 x5) = %.1fs", serial20)
+
+
+def main(cfg) -> None:
+    previous_signal_handlers = _install_graceful_stop_handlers()
+    trainer = None
+    normal_completion = False
+    _pin_cuda_device_from_local_rank()
+    try:
+        logger.info("VLA Training :: Warming Up")
+        if hasattr(cfg, "rl_games"):
+            login_training_services(cfg, workspace_dir=getattr(cfg, "workspace_dir", None))
+            validate_rl_games_config(cfg)
+        # Keep the obs window tied to the KV-memory rollout length before the config is
+        # wrapped/consumed by the dataloader and framework.
+        sync_kv_memory_obs_window(cfg)
+        apply_model_alias(cfg)
+        apply_action_spec(cfg)
+
+        cfg = wrap_config(cfg)
+        logger.info("✅ Configuration wrapped for access tracking")
+
+        output_dir = setup_directories(cfg=cfg)
+        vla = build_initial_framework(cfg)
+        vla = _preload_model_checkpoint_before_accelerator(cfg=cfg, model=vla)
+
+        distributed_backend = str(getattr(cfg.trainer, "distributed_backend", "deepspeed")).lower()
+        accelerator = build_accelerator(cfg, use_deepspeed=distributed_backend == "deepspeed")
+
+        vla_train_dataloader, vla_eval_dataloader = prepare_data(cfg=cfg, accelerator=accelerator, output_dir=output_dir)
+        _configure_quota_cumulative_training_steps(cfg=cfg, dataloader=vla_train_dataloader, accelerator=accelerator)
+        optimizer, lr_scheduler = setup_optimizer_and_scheduler(model=vla, cfg=cfg)
+
+        trainer = VLATrainer(
+            cfg=cfg,
+            model=vla,
+            vla_train_dataloader=vla_train_dataloader,
+            vla_eval_dataloader=vla_eval_dataloader,
+            optimizer=optimizer,
+            lr_scheduler=lr_scheduler,
+            accelerator=accelerator,
+        )
+
+        trainer.prepare_training()
+
+        import os as _os
+        if _os.environ.get("LB_EVAL_BENCH"):
+            _run_eval_batch_benchmark(trainer)
+            logger.info("[LB_EVAL_BENCH] done; skipping training.")
+            normal_completion = True
+            return
+
+        trainer.train()
+        normal_completion = not bool(getattr(trainer, "graceful_stop_requested", False))
+
+        if normal_completion:
+            logger.info("... and that's all, folks!")
+    finally:
+        if trainer is not None:
+            trainer.cleanup_runtime()
+        if sys.exc_info()[0] is None:
+            _destroy_distributed_process_group(use_barrier=normal_completion)
+        _restore_signal_handlers(previous_signal_handlers)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config_yaml",
+        type=str,
+        default="examples/SimplerEnv/train_files/starvla_cotrain_oxe.yaml",
+        help="Path to YAML config",
+    )
+    args, clipargs = parser.parse_known_args()
+
+    cfg = OmegaConf.load(args.config_yaml)
+    dotlist = normalize_dotlist_args(clipargs)
+    cli_cfg = OmegaConf.from_dotlist(dotlist)
+    cfg = OmegaConf.merge(cfg, cli_cfg)
+
+    # Normalise legacy YAML keys into the current `version_id == "0.21"` schema.
+    # This is idempotent and does not modify framework class signatures.
+    # See bar/config_收紧.md for the rationale.
+    cfg = apply_config_compat(cfg)
+
+    # Store source config path for later copying to output dir
+    cfg.config_yaml = args.config_yaml
+
+    if cfg.is_debug and dist.is_initialized() and dist.get_rank() == 0:
+        import debugpy
+
+        debugpy.listen(("0.0.0.0", 10092))
+        print("🔍 Rank 0 waiting for debugger attach on port 10092...")
+        debugpy.wait_for_client()
+
+    main(cfg)
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/trainer_tools.py b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/trainer_tools.py
new file mode 100644
index 0000000000000000000000000000000000000000..2b87675816acbc873c5abb8040812abd19831489
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/trainer_tools.py
@@ -0,0 +1,736 @@
+"""
+metrics.py
+
+Utility classes defining a Metrics container and multiple Trackers to enable model/stage-specific logging to various
+endpoints (e.g., JSONL local logs, Weights & Biases).
+"""
+
+from typing import Tuple
+import re
+import json
+import numpy as np
+import torch
+import torch.distributed as dist
+from transformers import get_scheduler
+
+from accelerate import Accelerator, DeepSpeedPlugin
+from accelerate.logging import get_logger
+from starVLA.training.trainer_utils.config_tracker import reload_module_paths
+
+logger = get_logger(__name__)
+
+
+def build_accelerator(cfg, *, use_deepspeed: bool) -> Accelerator:
+    accelerator_kwargs = {
+        "gradient_accumulation_steps": cfg.trainer.gradient_accumulation_steps,
+    }
+    if use_deepspeed:
+        accelerator_kwargs["deepspeed_plugin"] = DeepSpeedPlugin()
+
+    accelerator = Accelerator(**accelerator_kwargs)
+    accelerator.print(accelerator.state)
+    return accelerator
+
+
+_VOCAB_SIZED_WEIGHT_SUFFIXES = (
+    "embed_tokens.weight",
+    "lm_head.weight",
+)
+
+
+def _adapt_vocab_sized_weights(module, state_dict):
+    """Copy overlapping vocab rows when tokenizer-sized Qwen weights changed."""
+    target_state_dict = module.state_dict()
+    adapted = dict(state_dict)
+    adapted_keys = []
+    for key, source_tensor in state_dict.items():
+        target_tensor = target_state_dict.get(key)
+        if target_tensor is None or not hasattr(source_tensor, "shape") or not hasattr(target_tensor, "shape"):
+            continue
+        if tuple(source_tensor.shape) == tuple(target_tensor.shape):
+            continue
+        if not key.endswith(_VOCAB_SIZED_WEIGHT_SUFFIXES):
+            continue
+        if source_tensor.ndim != 2 or target_tensor.ndim != 2:
+            continue
+        if source_tensor.shape[1] != target_tensor.shape[1]:
+            continue
+
+        patched_tensor = target_tensor.detach().clone()
+        rows = min(source_tensor.shape[0], target_tensor.shape[0])
+        patched_tensor[:rows].copy_(source_tensor[:rows].to(dtype=patched_tensor.dtype, device=patched_tensor.device))
+        adapted[key] = patched_tensor
+        adapted_keys.append((key, tuple(source_tensor.shape), tuple(target_tensor.shape), rows))
+    return adapted, adapted_keys
+
+
+# === Define Tracker Interface ===
+#
+
+# utils/cli_parser.py
+
+
+def normalize_dotlist_args(args):
+    """
+    Convert ['--x.y', 'val'] and ['--flag'] → ['x.y=val', 'flag=true']
+    """
+    normalized = []
+    skip = False
+    for i in range(len(args)):
+        if skip:
+            skip = False
+            continue
+
+        arg = args[i]
+        if arg.startswith("--"):
+            key = arg.lstrip("-")
+            if "=" in key:
+                normalized.append(key)
+            elif i + 1 < len(args) and not args[i + 1].startswith("--"):
+                normalized.append(f"{key}={args[i + 1]}")
+                skip = True
+            else:
+                normalized.append(f"{key}=true")
+        else:
+            pass  # skip orphaned values
+    return normalized
+
+
+def setup_optimizer_and_scheduler(model, cfg) -> Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler]:
+    """Build AdamW optimizer + LR scheduler from cfg.trainer.
+
+    Supports:
+    - Per-module learning rates via cfg.trainer.learning_rate
+    - cosine_with_min_lr (or any transformers scheduler) via cfg.trainer.lr_scheduler_type
+    """
+    param_groups = build_param_lr_groups(model=model, cfg=cfg)
+    fused_available = torch.cuda.is_available()
+    optimizer = torch.optim.AdamW(
+        param_groups,
+        lr=cfg.trainer.learning_rate.base,
+        betas=tuple(cfg.trainer.optimizer.betas),
+        weight_decay=cfg.trainer.optimizer.weight_decay,
+        eps=cfg.trainer.optimizer.eps,
+        fused=fused_available,
+    )
+
+    if dist.is_initialized() and dist.get_rank() == 0:
+        for group in optimizer.param_groups:
+            logger.info(f"LR Group {group['name']}: lr={group['lr']}, num_params={len(group['params'])}")
+
+    # Strip keys unknown to transformers' get_scheduler before passing kwargs.
+    sched_kwargs = cfg.trainer.scheduler_specific_kwargs
+    lr_scheduler = get_scheduler(
+        name=cfg.trainer.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=cfg.trainer.num_warmup_steps,
+        num_training_steps=cfg.trainer.max_train_steps,
+        scheduler_specific_kwargs=sched_kwargs,
+    )
+
+    return optimizer, lr_scheduler
+
+
+def _qwen_vl_interface(model):
+    return model.qwen_vl_interface
+
+
+def _llm_language_model(model):
+    return _qwen_vl_interface(model).model.model.language_model
+
+
+def _vit_module(model):
+    # Qwen-VL keeps the visual tower and its vision-to-language connector/merger
+    # under this module.
+    return _qwen_vl_interface(model).model.model.visual
+
+
+def _llm_decoder_layers(model):
+    return _llm_language_model(model).layers
+
+
+def _tied_embedding_modules(model):
+    language_model = _llm_language_model(model)
+    hf_model = _qwen_vl_interface(model).model
+    return language_model.embed_tokens, hf_model.lm_head
+
+
+def resolve_freeze_llm_layer_indices(model, cfg):
+    """Translate ``trainer.freeze_llm_layers=[start,end]`` to an inclusive layer list."""
+    freeze_llm_layers = cfg.trainer.freeze_llm_layers
+    if not freeze_llm_layers:
+        return [], len(_llm_decoder_layers(model))
+    start, end = [int(value) for value in freeze_llm_layers]
+    return list(range(start, end + 1)), len(_llm_decoder_layers(model))
+
+
+def _module_param_ids(module):
+    return {id(param) for param in module.parameters()}
+
+
+def _has_vit_or_llm_freeze_request(cfg) -> bool:
+    return bool(cfg.trainer.freeze_vit or cfg.trainer.freeze_tied_embedding or cfg.trainer.freeze_llm_layers)
+
+
+def vit_llm_frozen_param_ids(model, cfg):
+    """Parameter ids excluded from optimizer by trainer freeze controls."""
+    if not _has_vit_or_llm_freeze_request(cfg):
+        return set()
+
+    frozen_params = set()
+    freeze_vit = cfg.trainer.freeze_vit
+    if freeze_vit:
+        frozen_params.update(_module_param_ids(_vit_module(model)))
+
+    if cfg.trainer.freeze_tied_embedding:
+        for module in _tied_embedding_modules(model):
+            frozen_params.update(_module_param_ids(module))
+
+    frozen_layer_indices, _ = resolve_freeze_llm_layer_indices(model, cfg)
+    layers = list(_llm_decoder_layers(model))
+    for layer_index in frozen_layer_indices:
+        frozen_params.update(_module_param_ids(layers[layer_index]))
+
+    return frozen_params
+
+
+def build_param_lr_groups(model, cfg):
+    """
+    build multiple param groups based on cfg.trainer.learning_rate.
+    support specifying different learning rates for different modules, the rest use base.
+
+    Args:
+        vla: nn.Module model object
+        cfg: config object, requires cfg.trainer.learning_rate dictionary
+
+    Returns:
+        List[Dict]: param_groups that can be used to build optimizer with torch.optim
+    """
+
+    lr_cfg = cfg.trainer.learning_rate
+    base_lr = lr_cfg.get("base", 1e-4)  # default base learning rate
+
+    freeze_modules = cfg.trainer.get("freeze_modules", "")
+    if not isinstance(freeze_modules, str):
+        freeze_modules = ""
+    freeze_patterns = [p.strip() for p in freeze_modules.split(",") if p.strip()]
+
+    used_params = set()
+    frozen_params = set()
+    param_groups = []
+
+    for freeze_path in freeze_patterns:
+        module = model
+        try:
+            for attr in freeze_path.split("."):
+                module = getattr(module, attr)
+            frozen_params.update(id(p) for p in module.parameters())
+        except AttributeError:
+            print(f"⚠️ freeze module path does not exist: {freeze_path}")
+            continue
+
+    # Exclude trainer freeze controls so frozen parameters never enter an
+    # optimizer param group.
+    frozen_params.update(vit_llm_frozen_param_ids(model, cfg))
+
+    for module_name, lr in lr_cfg.items():
+        if module_name == "base":
+            continue
+        # try to find the module under vla by module_name (support nested paths)
+        module = model
+        try:
+            for attr in module_name.split("."):
+                module = getattr(module, attr)
+            # filter out frozen parameters
+            params = [p for p in module.parameters() if id(p) not in frozen_params]
+            if params:  # only add param group if there are trainable parameters
+                param_groups.append({"params": params, "lr": lr, "name": module_name})
+                used_params.update(id(p) for p in params)
+        except AttributeError:
+            ReferenceError(f"⚠️ module path `{module_name}` not found in vla")
+
+    # assign base learning rate to the remaining unused parameters (exclude frozen ones)
+    other_params = [p for p in model.parameters() if id(p) not in used_params and id(p) not in frozen_params]
+    if other_params:
+        param_groups.append({"params": other_params, "lr": base_lr, "name": "base"})
+
+    return param_groups
+
+
+import torch.distributed as dist
+
+
+def only_main_process(func):
+    """
+    decorator: only run in main process (rank=0)
+    """
+
+    def wrapper(*args, **kwargs):
+        if dist.is_initialized() and dist.get_rank() != 0:
+            return None  # non-main process does not execute
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+from torchvision.ops import box_iou
+from PIL import Image
+
+
+def resize_images(images, target_size=(224, 224)):
+    """
+    recursively resize all images in the nested list.
+
+    :param images: nested list of images or single image.
+    :param target_size: target size (width, height) after resizing.
+    :return: resized images list, keeping the original nested structure.
+    """
+    if isinstance(images, Image.Image):  # if it is a single PIL image
+        return images.resize(target_size)
+    elif isinstance(images, list):  # if it is a list, recursively process each element
+        return [resize_images(img, target_size) for img in images]
+    else:
+        raise ValueError("Unsupported image type or structure.")
+
+
+def stitch_frames(frames, grid=(2, 2), size=(224, 224)):
+    """Tile temporal frames into one image, keeping the token count of a single image.
+
+    Used by the ``stitch`` multi-image mode: instead of feeding N frames as N
+    tokenized images, they are laid out on a fixed grid and the whole canvas is
+    returned at ``size``. Both training (``_pack_sample``) and eval
+    (``policy/starvla.py``) call this so the layouts match.
+
+    Args:
+        frames: list of ``PIL.Image`` or ``np.ndarray`` (H, W, C), oldest first.
+        grid: (rows, cols). Cells beyond ``len(frames)`` repeat the last frame so
+            the layout is fixed regardless of how many frames are available.
+        size: (width, height) of the returned canvas.
+    """
+    rows, cols = grid
+    canvas_w, canvas_h = size
+    cell_w, cell_h = canvas_w // cols, canvas_h // rows
+    pil = [f if isinstance(f, Image.Image) else Image.fromarray(f) for f in frames]
+    capacity = rows * cols
+    pil = (pil + [pil[-1]] * (capacity - len(pil))) if len(pil) < capacity else pil[:capacity]
+    canvas = Image.new("RGB", (canvas_w, canvas_h))
+    for idx, frame in enumerate(pil):
+        r, c = divmod(idx, cols)
+        canvas.paste(frame.convert("RGB").resize((cell_w, cell_h)), (c * cell_w, r * cell_h))
+    return canvas
+
+
+import torch.distributed as dist
+
+
+class TrainerUtils:
+    @staticmethod
+    def freeze_backbones(model, freeze_modules=""):
+        """
+        directly freeze the specified submodules based on the relative module path list (patterns), no longer recursively find all submodule names:
+          - patterns: read from config.trainer.freeze_modules, separated by commas to get the "relative path" list
+            for example "qwen_vl_interface, action_model.net",
+            it means to freeze model.qwen_vl_interface and model.action_model.net.
+
+        Args:
+            model: nn.Module model object
+            freeze_modules: relative module path list (patterns)
+
+        Returns:
+            model: nn.Module model object
+        return:
+          - model:
+        """
+        frozen = []
+        print("#"*30)
+        print(freeze_modules)
+        if freeze_modules and type(freeze_modules) == str:
+            # split and remove whitespace
+            patterns = [p.strip() for p in freeze_modules.split(",") if p.strip()] if freeze_modules else []
+
+            for path in patterns:
+                # split the "relative path" by dots, for example "action_model.net" → ["action_model", "net"]
+                attrs = path.split(".")
+                module = model
+                try:
+                    for attr in attrs:
+                        module = getattr(module, attr)
+                    # if the module is successfully get, freeze it and its all submodule parameters
+                    for param in module.parameters():
+                        param.requires_grad = False
+                    frozen.append(path)
+                except AttributeError:
+                    # if the attribute does not exist, skip and print warning
+                    print(f"⚠️ module path does not exist, cannot freeze: {path}")
+                    continue
+
+        # accelerator.wait_for_everyone()  # synchronize when distributed training
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(f"🔒 Frozen modules with re pattern: {frozen}")
+        return model
+
+    @staticmethod
+    def freeze_vit_and_llm_layers(model, cfg):
+        """Apply trainer.freeze_vit/freeze_tied_embedding/freeze_llm_layers.
+
+        trainer.freeze_vit freezes Qwen-VL's visual module, including the
+        vision-to-language connector/merger it contains.
+        """
+        if not _has_vit_or_llm_freeze_request(cfg):
+            return model
+
+        freeze_vit = cfg.trainer.freeze_vit
+        if freeze_vit:
+            for param in _vit_module(model).parameters():
+                param.requires_grad = False
+
+        freeze_tied_embedding = cfg.trainer.freeze_tied_embedding
+        if freeze_tied_embedding:
+            for module in _tied_embedding_modules(model):
+                for param in module.parameters():
+                    param.requires_grad = False
+
+        frozen_layer_indices, total = resolve_freeze_llm_layer_indices(model, cfg)
+        layers = list(_llm_decoder_layers(model))
+        for layer_index in frozen_layer_indices:
+            for param in layers[layer_index].parameters():
+                param.requires_grad = False
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(
+                "🔒 freeze_vit="
+                f"{freeze_vit}, freeze_tied_embedding={freeze_tied_embedding}, "
+                f"freeze_llm_layers={list(frozen_layer_indices)}"
+                f" / total={total}"
+            )
+        return model
+
+    @staticmethod
+    def print_trainable_parameters(model):
+        """
+        print the total number of parameters and trainable parameters of the model
+        :param model: PyTorch model instance
+        """
+        if dist.is_initialized() and dist.get_rank() != 0:
+            return
+        print("📊 model parameter statistics:")
+        num_params = sum(p.numel() for p in model.parameters())
+        num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        print(
+            f"# Parameters (in millions): {num_params / 10**6:.3f} Total, {num_trainable_params / 10**6:.3f} Trainable"
+        )
+        return num_params, num_trainable_params
+
+    @staticmethod
+    def load_pretrained_backbones(model, checkpoint_path=None, reload_modules=None):
+        """
+        load checkpoint:
+        - if reload_modules is set, load by path part
+        - otherwise → load the entire model parameters (overwrite model)
+
+        return:
+            replace, loaded_modules: list of module paths that successfully loaded parameters; if global load, then ["<full_model>"]
+        """
+        if not checkpoint_path:
+            return []
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(f"📦 loading checkpoint: {checkpoint_path}")
+        try:
+            if _is_safetensors_path(checkpoint_path):
+                from safetensors.torch import load_file
+
+                checkpoint = load_file(checkpoint_path)
+            else:
+                checkpoint = torch.load(checkpoint_path, map_location="cpu")
+        except Exception as e:
+            raise RuntimeError(f"❌ loading checkpoint failed: {e}")
+
+        loaded_modules = []
+
+        module_paths = reload_module_paths(reload_modules)
+        if module_paths:  # partial load
+            for path in module_paths:
+                reload_path = path.split(".")
+                module = model
+                try:
+                    for module_name in reload_path:  # find the module to modify level by level
+                        module = getattr(module, module_name)
+                    prefix = path + "."
+                    sub_state_dict = {k[len(prefix) :]: v for k, v in checkpoint.items() if k.startswith(prefix)}
+                    if sub_state_dict:
+                        sub_state_dict, adapted_keys = _adapt_vocab_sized_weights(module, sub_state_dict)
+                        for key, source_shape, target_shape, rows in adapted_keys:
+                            if not dist.is_initialized() or dist.get_rank() == 0:
+                                print(
+                                    f"⚠️ adapted vocab-sized weight '{path}.{key}' from checkpoint shape "
+                                    f"{source_shape} to current shape {target_shape}; copied {rows} rows"
+                                )
+                        module.load_state_dict(sub_state_dict, strict=True)
+                        if not dist.is_initialized() or dist.get_rank() == 0:
+                            print(f"✅ parameters loaded to module '{path}'")
+                        loaded_modules.append(path)
+                    else:
+                        print(f"⚠️ parameters not found in checkpoint '{path}'")
+                except AttributeError:
+                    print(f"❌ cannot find module path: {path}")
+        else:  # full load
+            try:
+                checkpoint, adapted_keys = _adapt_vocab_sized_weights(model, checkpoint)
+                for key, source_shape, target_shape, rows in adapted_keys:
+                    if not dist.is_initialized() or dist.get_rank() == 0:
+                        print(
+                            f"⚠️ adapted vocab-sized weight '{key}' from checkpoint shape "
+                            f"{source_shape} to current shape {target_shape}; copied {rows} rows"
+                        )
+                model.load_state_dict(checkpoint, strict=False)
+                if not dist.is_initialized() or dist.get_rank() == 0:
+                    print("✅ loaded <full_model> model parameters")
+                loaded_modules = ["<full_model>"]
+            except Exception as e:
+                raise RuntimeError(f"❌ loading full model failed: {e}")
+        return model
+
+    @staticmethod
+    def print_freeze_status(model):
+        """
+        print the freezing status of each parameter in the model
+        :param model: PyTorch model instance
+        """
+        for name, param in model.named_parameters():
+            status = "Frozen" if not param.requires_grad else "Trainable"
+            print(f"{name:60s}  |  {status}")
+
+    @staticmethod
+    def setup_distributed_training(accelerator, *components):
+        """
+        use Accelerator to prepare distributed training components
+        :param accelerator: Accelerate instance
+        :param components: any number of components (such as model, optimizer, dataloader, etc.)
+        :return: prepared distributed components (in the same order as input)
+        """
+
+        # use accelerator.prepare method to wrap components
+        prepared_components = accelerator.prepare(*components)
+        return prepared_components
+
+    @staticmethod
+    def euclidean_distance(predicted: np.ndarray, ground_truth: np.ndarray) -> float:
+        return np.linalg.norm(predicted - ground_truth)
+
+    @staticmethod
+    def _reset_dataloader(dataloader, epoch_counter):
+        """safe reset dataloader iterator"""
+        # 1. update epoch counter
+        epoch_counter += 1
+
+        dataset = getattr(dataloader, "dataset", None)
+        if dataset is not None and callable(getattr(dataset, "set_epoch", None)):
+            dataset.set_epoch(epoch_counter)
+
+        # 2. set new epoch (distributed core)
+        if hasattr(dataloader, "sampler") and callable(getattr(dataloader.sampler, "set_epoch", None)):
+            dataloader.sampler.set_epoch(epoch_counter)
+
+        # 3. create new iterator
+        return iter(dataloader), epoch_counter
+
+    @staticmethod
+    def compute_grad_angle_with_stats(grads_a: list[torch.Tensor], grads_v: list[torch.Tensor]) -> Tuple[float, float]:
+        """
+        compute the cosine angle between two groups of gradient vectors (degrees), and calculate the average angle and variance.
+        grads_a, grads_v: gradient Tensor list corresponding to the same parameter list interface_params
+        return:
+            mean_angle_deg: average angle (degrees)
+            angle_variance: angle variance
+        """
+        angle_degs = []
+
+        # compute the cosine angle between each gradient block grads_a[0].shape = 1280, 3, 14, 14
+        # grads_1 = grads_a[0][0]  # [3, 14, 14]
+        # grads_2 = grads_v[0][0]
+        # grads_a = grads_1.view(-1, 3)  # reshape to [196, 3]
+        # grads_v = grads_2.view(-1, 3)
+
+        # lang linear
+        # reshape to 14*14, 3
+        # layer
+        grads_action = grads_a[0]  # [2048, 11008]
+        grads_action = grads_action[
+            :32, :7
+        ]  # only take the first 7 elements, avoid cosim failure in high-dimensional space
+        grads_vl = grads_v[0]  # [2048, 11008]
+        grads_vl = grads_vl[
+            :32, :7
+        ]  # only take the first 32 elements, 7 dimensions, avoid cosim failure in high-dimensional space
+        for g_a, g_v in zip(grads_action, grads_vl):
+            dot = torch.sum(g_a * g_v)
+            norm_a_sq = torch.sum(g_a * g_a)
+            norm_v_sq = torch.sum(g_v * g_v)
+
+            # avoid division by zero
+            norm_a = torch.sqrt(norm_a_sq + 1e-16)
+            norm_v = torch.sqrt(norm_v_sq + 1e-16)
+
+            cos_sim = (dot / (norm_a * norm_v)).clamp(-1.0, 1.0)
+            angle_rad = torch.acos(cos_sim)
+            angle_deg = angle_rad * (180.0 / torch.pi)
+
+            angle_degs.append(angle_deg.item())
+
+        # compute the average angle and variance
+        angle_degs_tensor = torch.tensor(angle_degs)
+        mean_angle_deg = torch.mean(angle_degs_tensor).item()
+        angle_variance = torch.sqrt(torch.var(angle_degs_tensor)).item()
+        # accelerator.wait_for_everyone()
+        return mean_angle_deg, angle_variance
+
+    @staticmethod
+    def pcgrad_project(grads_a: list[torch.Tensor], grads_v: list[torch.Tensor]) -> list[torch.Tensor]:
+        """
+        apply PCGrad projection to the second group of gradients grads_v, suppress negative transfer between grads_a and grads_v
+        if the dot product of two groups of gradients < 0, then:
+            grads_v <- grads_v - (dot / ||grads_a||^2) * grads_a
+        return the new grads_v list
+        """
+        # first compute dot and ||grads_a||^2
+        dot, norm_a_sq = 0.0, 0.0
+        for g_a, g_v in zip(grads_a, grads_v):
+            dot += torch.sum(g_a * g_v)
+            norm_a_sq += torch.sum(g_a * g_a)
+
+        if dot < 0:
+            coeff = dot / (norm_a_sq + 1e-6)
+            # projection
+            grads_v = [g_v - coeff * g_a for g_a, g_v in zip(grads_a, grads_v)]
+
+        return grads_v
+
+    @staticmethod
+    def eval_qwenpi(qwenpi, dataloader, num_batches=20):
+        """
+        evaluate QwenQFormerDiT model, compute IoU and action distance.
+
+        Args:
+            qwenpi: QwenQFormerDiT model instance.
+            dataloader: data loader.
+            num_batches: number of batches to evaluate.
+
+        Returns:
+            dict: contains IoU and action distance evaluation results.
+        """
+        iou_scores = []
+        action_distances = []
+        count = 0
+
+        dataset_iter = iter(dataloader)
+        while count < num_batches:
+            try:
+                batch_samples = next(dataset_iter)
+                count += 1
+            except StopIteration:
+                break
+
+            # extract data
+            images = [example["image"] for example in batch_samples]
+            instructions = [example["lang"] for example in batch_samples]
+            actions = [example["action"] for example in batch_samples]
+            solutions = [example["solution"] for example in batch_samples]
+
+            # model prediction
+            predicted_solutions, normalized_actions = qwenpi.predict_action_withCoT(
+                images=images, instructions=instructions, use_ddim=False, num_ddim_steps=20
+            )
+
+            # extract and convert predicted results
+            parsed_solutions = []
+            for solution in predicted_solutions:
+                parsed_solution = TrainerUtils.extract_json_from_string(solution)
+                parsed_solutions.append(parsed_solution)
+
+            # compute IoU
+            for pred_dict, gt_dict in zip(parsed_solutions, solutions):
+                pred_pick_bbox = torch.tensor(pred_dict["pick"]["bbox_2d"], dtype=torch.float32).unsqueeze(0)
+                gt_pick_bbox = torch.tensor(gt_dict["pick"]["bbox_2d"], dtype=torch.float32).unsqueeze(0)
+                pred_place_bbox = torch.tensor(pred_dict["place"]["bbox_2d"], dtype=torch.float32).unsqueeze(0)
+                gt_place_bbox = torch.tensor(gt_dict["place"]["bbox_2d"], dtype=torch.float32).unsqueeze(0)
+
+                pick_iou = box_iou(pred_pick_bbox, gt_pick_bbox).item()
+                place_iou = box_iou(pred_place_bbox, gt_place_bbox).item()
+
+                iou_scores.append({"pick_iou": pick_iou, "place_iou": place_iou})
+
+            # compute action distance
+            actions = np.array(actions)  # convert to numpy array
+            num_pots = np.prod(actions.shape)  # B*len*dim
+            action_distance = TrainerUtils.euclidean_distance(normalized_actions, actions)
+            average_action_distance = action_distance / num_pots
+            action_distances.append(average_action_distance)
+
+        # summarize results
+        avg_action_distance = np.mean(action_distances)
+        return {"iou_scores": iou_scores, "average_action_distance": avg_action_distance}
+
+    @staticmethod
+    def extract_json_from_string(input_string):
+        """
+        extract valid JSON part from string and convert to dictionary.
+
+        Args:
+            input_string (str): string containing extra characters.
+
+        Returns:
+            dict: dictionary extracted and parsed.
+        """
+        json_match = re.search(r"{.*}", input_string, re.DOTALL)
+        if json_match:
+            json_str = json_match.group(0)
+            try:
+                return json.loads(json_str)
+            except json.JSONDecodeError as e:
+                print(f"JSON decode failed: {e}")
+                return None
+        else:
+            print("No valid JSON part found")
+            return None
+
+    def _get_latest_checkpoint(self, checkpoint_dir):
+        """Find the latest checkpoint in the directory based on step number."""
+        if not os.path.exists(checkpoint_dir):
+            self.accelerator.print(f"No checkpoint directory found at {checkpoint_dir}")
+            return None, 0
+
+        checkpoints = []
+        for name in os.listdir(checkpoint_dir):
+            path = os.path.join(checkpoint_dir, name)
+            state_match = re.match(r"steps_(\d+)_state$", name)
+            if state_match and os.path.isdir(path):
+                checkpoints.append((name, int(state_match.group(1)), 1))
+                continue
+            file_match = re.match(r"steps_(\d+)_(?:pytorch_model\.pt|model\.safetensors)$", name)
+            if file_match and os.path.isfile(path):
+                checkpoints.append((name, int(file_match.group(1)), 0))
+
+        if not checkpoints:
+            self.accelerator.print(f"No checkpoints found in {checkpoint_dir}")
+            return None, 0
+
+        # Sort by step number, preferring full training-state directories over
+        # model-only files when both exist for the same step.
+        checkpoints.sort(key=lambda x: (x[1], x[2]))
+        latest_checkpoint, completed_steps, _ = checkpoints[-1]
+
+        latest_checkpoint_path = os.path.join(checkpoint_dir, latest_checkpoint)
+        self.accelerator.print(f"Latest checkpoint found: {latest_checkpoint_path}")
+        return latest_checkpoint_path, completed_steps
+
+import os
+
+
+def is_main_process():
+    rank = int(os.environ.get("RANK", 0))  # if RANK is not set, default to 0
+    return rank == 0
+
+
+def _is_safetensors_path(path):
+    """Check if a path refers to a safetensors file."""
+    return str(path).endswith(".safetensors")
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/worker_context.py b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/worker_context.py
new file mode 100644
index 0000000000000000000000000000000000000000..1e0e4ddfa240c5c695fac627ff0a4f075b2c2267
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/worker_context.py
@@ -0,0 +1,77 @@
+import multiprocessing.context
+import os
+from typing import Any
+
+
+_DISTRIBUTED_ENV_KEYS = (
+    "LOCAL_RANK",
+    "RANK",
+    "WORLD_SIZE",
+    "MASTER_ADDR",
+    "MASTER_PORT",
+    "GROUP_RANK",
+    "LOCAL_WORLD_SIZE",
+    "ROLE_RANK",
+    "ROLE_WORLD_SIZE",
+    "TORCHELASTIC_RUN_ID",
+    "TORCHELASTIC_RESTART_COUNT",
+    "TORCHELASTIC_MAX_RESTARTS",
+)
+
+
+class _CpuOnlyWorkerSpawnProcess(multiprocessing.context.SpawnProcess):
+    def start(self):
+        had_cuda_visible_devices = "CUDA_VISIBLE_DEVICES" in os.environ
+        if had_cuda_visible_devices:
+            cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        distributed_env = {key: os.environ[key] for key in _DISTRIBUTED_ENV_KEYS if key in os.environ}
+
+        os.environ["CUDA_VISIBLE_DEVICES"] = ""
+        for key in distributed_env:
+            del os.environ[key]
+        try:
+            return super().start()
+        finally:
+            if had_cuda_visible_devices:
+                os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
+            else:
+                del os.environ["CUDA_VISIBLE_DEVICES"]
+            for key, value in distributed_env.items():
+                os.environ[key] = value
+
+
+class CpuOnlyWorkerSpawnContext(multiprocessing.context.SpawnContext):
+    Process = _CpuOnlyWorkerSpawnProcess
+
+
+CPU_ONLY_WORKER_CONTEXT = CpuOnlyWorkerSpawnContext()
+
+
+def _as_bool(value: Any, default: bool = False) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        return value.strip().lower() not in {"false", "0", "no", "off"}
+    return bool(value)
+
+
+def build_cpu_only_dataloader_kwargs(
+    num_workers: int,
+    *,
+    pin_memory: Any = False,
+    persistent_workers: Any = False,
+    prefetch_factor: Any = None,
+) -> dict[str, Any]:
+    dataloader_kwargs: dict[str, Any] = {
+        "pin_memory": _as_bool(pin_memory),
+    }
+    if int(num_workers) <= 0:
+        return dataloader_kwargs
+
+    dataloader_kwargs["multiprocessing_context"] = CPU_ONLY_WORKER_CONTEXT
+    dataloader_kwargs["persistent_workers"] = _as_bool(persistent_workers)
+    if prefetch_factor is not None:
+        dataloader_kwargs["prefetch_factor"] = int(prefetch_factor)
+    return dataloader_kwargs
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training_code_snapshot.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training_code_snapshot.json
new file mode 100644
index 0000000000000000000000000000000000000000..9f3eae00c31cf25e9f57b6117847b0a5d5c70475
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training_code_snapshot.json
@@ -0,0 +1,19 @@
+{
+  "repo": "c3c6a39365a151e9b7a5e215452fd64e957c2b29",
+  "starvla": "ccca13c5177fe3d3c884b6e2de4965d916016649",
+  "runtime_fixes": [
+    "mean-profile-preparation.patch",
+    "gym-language-contract.patch",
+    "loader-spawn-cache.patch",
+    "loader-spawn-test.patch"
+  ],
+  "pytorch3d": {
+    "revision": "33824be3cbc87a7dd1db0f6a9a9de9ac81b2d0ba",
+    "build": "transforms-only, no native render extension; QwenOFT uses transforms only"
+  },
+  "decord": {
+    "version": "0.6.0",
+    "build": "official source CPU decoder CP310",
+    "wheel_sha256": "e193b356b1e984b4eff08d23b62e482c2c9e5037a6efdc0b1af47079ae2e4c47"
+  }
+}
diff --git a/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training_results.json b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training_results.json
new file mode 100644
index 0000000000000000000000000000000000000000..6d2a9d9f3ff94ed1cf2b780ece5fd38d3bbee3bf
--- /dev/null
+++ b/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training_results.json
@@ -0,0 +1,15 @@
+{
+  "training_complete": true,
+  "optimizer_updates": 5000,
+  "selection": "final_5000",
+  "train_loss": 0.028774918988347054,
+  "validation_action_loss": 0.034965552389621735,
+  "validation_samples": 6400,
+  "final_backbone_lr": 1e-07,
+  "final_action_lr": 3.3333333333333333e-06,
+  "closed_loop_evaluation": "pending until all four VLA trainings complete",
+  "tensor_keys": 732,
+  "serialized_tensor_elements_including_tied_aliases": 4892433410,
+  "weights_only_load": "passed",
+  "all_tensors_finite": true
+}