Publish completed intercept QwenOFT H1 mean-profile 5000-update run

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/README.md

@@ -0,0 +1,8 @@
+# intercept QwenOFT H1 mean-profile VLA
+
+Training complete:5000 optimizer updates, seed0, globalbatch128, microbatch64 perGPU, accumulation1, gradientcheckpointingTrue, BF16/ZeRO2. Hardware:2 RTX PRO6000 Blackwell GPUs. Fixed mean99.05021289731565ms uses originalRTX3090 profile; all actual science/runtime settings and user-authorized throughput changes are recorded in config.full.yaml/launch_audit.json. Horizon1, native state/action dimensions7/7.
+
+Puretensor weights checkCPUweights_only/allfinite passed. Portable inferenceconfig, normalization/statistics, taskcontract, sourcecode/profile/data provenance included. Full optimizer recovery remains local. Finalofflinevalidation loss0.03070492297410965; this is not closed-loop task-performance acceptance. Latency-bench evaluation pending until all four models finish.
+
+Dataset:https://huggingface.co/datasets/latency-sensitive-bench/Standard-Pipeline/tree/bba7d56c0ad5bb2b6d6ed07adb70b4f1afd5d25a/mikasa_intercept_grab_fast/profile_latency/QwenOFT/mean_rtx3090_h1_s123_275ep_v1
+Training:https://wandb.ai/dongqianyu99-zhejiang-university/starvla_tasks/runs/mean3090-intercept-vla-blackwell-20260930

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/checkpoints/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7bea6746bf7973ff7d4e4e8a457a3bfbc0ce981c80fe5f7afadca651ec5ca85e
+size 9785101729

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/collection_quality_report.json

@@ -0,0 +1,306 @@
+{
+  "episodes": 275,
+  "unique_seeds": 275,
+  "rows": 12211,
+  "length_min": 37,
+  "length_max": 55,
+  "all_success_once": true,
+  "state_dim": 7,
+  "action_dim": 7,
+  "rgb_shape": [
+    128,
+    128,
+    6
+  ],
+  "finite_numeric": true,
+  "latency_raw_frames": 2,
+  "applied_command_matches_source_issued_action": true,
+  "action_horizon": 1,
+  "dagger": false,
+  "raw_bytes": 1204453838,
+  "episode_seeds": [
+    0,
+    1,
+    10,
+    180,
+    181,
+    182,
+    183,
+    184,
+    185,
+    186,
+    187,
+    188,
+    189,
+    11,
+    190,
+    191,
+    192,
+    193,
+    194,
+    195,
+    196,
+    197,
+    198,
+    199,
+    12,
+    20,
+    21,
+    22,
+    23,
+    24,
+    25,
+    26,
+    27,
+    28,
+    29,
+    13,
+    200,
+    201,
+    202,
+    203,
+    204,
+    205,
+    206,
+    207,
+    208,
+    209,
+    14,
+    210,
+    211,
+    212,
+    213,
+    214,
+    215,
+    216,
+    217,
+    218,
+    219,
+    15,
+    220,
+    221,
+    222,
+    223,
+    224,
+    225,
+    226,
+    227,
+    228,
+    229,
+    16,
+    230,
+    231,
+    232,
+    233,
+    234,
+    235,
+    236,
+    237,
+    238,
+    239,
+    17,
+    240,
+    241,
+    242,
+    243,
+    244,
+    245,
+    246,
+    247,
+    248,
+    249,
+    18,
+    250,
+    251,
+    252,
+    253,
+    254,
+    255,
+    256,
+    257,
+    258,
+    259,
+    19,
+    260,
+    261,
+    262,
+    263,
+    264,
+    265,
+    266,
+    267,
+    268,
+    269,
+    2,
+    100,
+    270,
+    271,
+    272,
+    273,
+    274,
+    30,
+    31,
+    32,
+    33,
+    34,
+    101,
+    35,
+    36,
+    37,
+    38,
+    39,
+    40,
+    41,
+    42,
+    43,
+    44,
+    102,
+    45,
+    46,
+    47,
+    48,
+    49,
+    50,
+    51,
+    52,
+    53,
+    54,
+    103,
+    55,
+    56,
+    57,
+    58,
+    59,
+    60,
+    61,
+    62,
+    63,
+    64,
+    104,
+    65,
+    66,
+    67,
+    68,
+    69,
+    70,
+    71,
+    72,
+    73,
+    74,
+    105,
+    75,
+    76,
+    77,
+    78,
+    79,
+    80,
+    81,
+    82,
+    83,
+    84,
+    106,
+    85,
+    86,
+    87,
+    88,
+    89,
+    90,
+    91,
+    92,
+    93,
+    94,
+    107,
+    95,
+    96,
+    97,
+    98,
+    99,
+    108,
+    109,
+    3,
+    110,
+    111,
+    112,
+    113,
+    114,
+    115,
+    116,
+    117,
+    118,
+    119,
+    4,
+    120,
+    121,
+    122,
+    123,
+    124,
+    125,
+    126,
+    127,
+    128,
+    129,
+    5,
+    130,
+    131,
+    132,
+    133,
+    134,
+    135,
+    136,
+    137,
+    138,
+    139,
+    6,
+    140,
+    141,
+    142,
+    143,
+    144,
+    145,
+    146,
+    147,
+    148,
+    149,
+    7,
+    150,
+    151,
+    152,
+    153,
+    154,
+    155,
+    156,
+    157,
+    158,
+    159,
+    8,
+    160,
+    161,
+    162,
+    163,
+    164,
+    165,
+    166,
+    167,
+    168,
+    169,
+    9,
+    170,
+    171,
+    172,
+    173,
+    174,
+    175,
+    176,
+    177,
+    178,
+    179
+  ],
+  "train_episodes": 250,
+  "val_episodes": 25,
+  "train_rows": 11119,
+  "val_rows": 1092,
+  "teacher_checkpoint": "/home/ubuntu/lzj/mean-profiling/intercept/teacher-publication/checkpoint.pt",
+  "teacher_sha256": "e47526f03fad3ed251dd78f3cacfba8aa8b60bd1578f40bd7a1393e17298e973",
+  "fixed_latency_ms": 99.05021289731565,
+  "conversion_complete": true
+}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/config.full.yaml

@@ -0,0 +1,102 @@
+run_id: intercept-qwenoft-mean-3090-h1-s0
+run_root_dir: /home/ubuntu/lzj/mean-profiling/intercept/vla
+seed: 0
+wandb_name: intercept-qwenoft-mean-3090-h1-s0
+wandb_project: starvla_tasks
+wandb_entity: dongqianyu99-zhejiang-university
+wandb_group: mean-profiling-3090-blackwell
+wandb_tags:
+- intercept
+- qwenoft
+- mean-profile
+- h1
+training_latency_condition: mean-profile
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenOFT
+  qwenvl:
+    base_vlm: /home/ubuntu/lzj/models/Qwen3-VL-4B-Instruct
+    attn_implementation: flash_attention_2
+    enable_gradient_checkpointing: true
+  action_model:
+    action_model_type: MLP
+    action_dim: 7
+    state_dim: 7
+    action_hidden_dim: 2560
+    future_action_window_size: 0
+    past_action_window_size: 0
+    loss_type: l1
+    state_encoding: continuous_projector
+    task_objective: null
+    action_env_dim: 7
+    action_horizon: 1
+datasets:
+  vla_data:
+    dataset_py: latency_bench_tasks
+    lerobot_version: v3.0
+    include_state: true
+    task_contract_path: /home/ubuntu/lzj/mean-profiling/intercept/data/lerobot/task_contract.json
+    data_root_dir: /
+    data_mix: train
+    eval_data_mix: validation
+    custom_mixtures_path: /home/ubuntu/lzj/mean-profiling/intercept/data/mixture.json
+    sequential_step_sampling: true
+    per_device_batch_size: 64
+    pin_memory: true
+    persistent_workers: true
+    num_workers: 2
+    prefetch_factor: 4
+    obs_image_size:
+    - 224
+    - 224
+    video_backend: pyav
+    image_mode: single
+    stitch_grid:
+    - 1
+    - 2
+    num_obs_frames: 1
+trainer:
+  max_train_steps: 5000
+  num_warmup_steps: 100
+  save_interval: 500
+  eval_interval: 5000
+  eval_num_batches: 1
+  logging_frequency: 5
+  learning_rate:
+    base: 3.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 1.0e-06
+  freeze_modules: ''
+  freeze_vit: false
+  freeze_tied_embedding: false
+  freeze_llm_layers: []
+  loss_scale:
+    vla: 1.0
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 1
+  pretrained_checkpoint: null
+  is_resume: false
+  distributed_backend: deepspeed
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+    fused: true
+  profile_timing:
+    enabled: false
+    log_interval: 10
+checkpoint:
+  save_best_model: false
+  save_final_model: true
+  save_pt_file: true
+  save_training_state: true
+  local:
+    keep_last_n: 1
+config_yaml: /home/ubuntu/lzj/mean-profiling/intercept/vla.yaml
+output_dir: /home/ubuntu/lzj/mean-profiling/intercept/vla/intercept-qwenoft-mean-3090-h1-s0

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/config.yaml

@@ -0,0 +1,41 @@
+framework:
+  name: QwenOFT
+  qwenvl:
+    base_vlm: Qwen/Qwen3-VL-4B-Instruct
+    attn_implementation: flash_attention_2
+    enable_gradient_checkpointing: true
+  action_model:
+    action_model_type: MLP
+    action_dim: 7
+    state_dim: 7
+    action_hidden_dim: 2560
+    future_action_window_size: 0
+    past_action_window_size: 0
+    loss_type: l1
+    state_encoding: continuous_projector
+    task_objective: null
+    action_env_dim: 7
+    action_horizon: 1
+version_id: '0.21'
+datasets:
+  vla_data:
+    include_state: true
+    obs_image_size:
+    - 224
+    - 224
+    image_mode: single
+    stitch_grid:
+    - 1
+    - 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'

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/dataset_statistics.json

@@ -0,0 +1,127 @@
+{
+  "new_embodiment": {
+    "action": {
+      "mean": [
+        -0.048697348684072495,
+        0.24886493384838104,
+        -0.16479986906051636,
+        -0.032609131187200546,
+        -0.501691997051239,
+        0.44703611731529236,
+        0.6793776154518127
+      ],
+      "std": [
+        0.1079770028591156,
+        0.47203660011291504,
+        0.14475314319133759,
+        0.2849433720111847,
+        0.318120539188385,
+        0.35324087738990784,
+        0.6349603533744812
+      ],
+      "max": [
+        0.30179348587989807,
+        0.7795099020004272,
+        0.2688712179660797,
+        0.8059309720993042,
+        0.5333064198493958,
+        1.0,
+        1.0
+      ],
+      "min": [
+        -0.6572320461273193,
+        -1.0,
+        -0.8218780755996704,
+        -0.7686842083930969,
+        -0.9920897483825684,
+        -0.30364304780960083,
+        -1.0
+      ],
+      "q01": [
+        -0.4552440994977951,
+        -1.0,
+        -0.7419710528850555,
+        -0.475678391456604,
+        -0.9305388247966766,
+        -0.20442348539829253,
+        -0.8981475746631622
+      ],
+      "q99": [
+        0.1125871029496192,
+        0.7402709639072418,
+        0.22590837419033025,
+        0.7005072689056396,
+        0.16810789704322815,
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        true
+      ]
+    },
+    "state": {
+      "mean": [
+        -0.14554812014102936,
+        0.07433967292308807,
+        0.1436367779970169,
+        2.700871706008911,
+        0.45681115984916687,
+        -0.48229172825813293,
+        0.06972828507423401
+      ],
+      "std": [
+        0.07678605616092682,
+        0.21988309919834137,
+        0.09388109296560287,
+        1.109084963798523,
+        0.3012843728065491,
+        0.3006882667541504,
+        0.0164483655244112
+      ],
+      "max": [
+        0.014208495616912842,
+        0.5909395217895508,
+        0.3558541536331177,
+        3.1415467262268066,
+        0.9420600533485413,
+        0.09770228713750839,
+        0.07999999821186066
+      ],
+      "min": [
+        -0.37446925044059753,
+        -0.29720592498779297,
+        0.016940969973802567,
+        -3.140481948852539,
+        -0.10707131028175354,
+        -1.2001688480377197,
+        0.03238735347986221
+      ],
+      "q01": [
+        -0.3139808464050293,
+        -0.25229972660541533,
+        0.0188405978679657,
+        -3.134360270500183,
+        -0.04521620273590088,
+        -1.0301372480392457,
+        0.03585549935698509
+      ],
+      "q99": [
+        -0.013631248474121185,
+        0.5275326526165007,
+        0.33348673582077026,
+        3.1367013454437256,
+        0.8922865676879882,
+        0.056985778063535636,
+        0.07999999821186066
+      ]
+    },
+    "num_transitions": 11119,
+    "num_trajectories": 250
+  }
+}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/dataset_statistics_eval.json

@@ -0,0 +1,127 @@
+{
+  "new_embodiment": {
+    "action": {
+      "mean": [
+        -0.048697348684072495,
+        0.24886493384838104,
+        -0.16479986906051636,
+        -0.032609131187200546,
+        -0.501691997051239,
+        0.44703611731529236,
+        0.6793776154518127
+      ],
+      "std": [
+        0.1079770028591156,
+        0.47203660011291504,
+        0.14475314319133759,
+        0.2849433720111847,
+        0.318120539188385,
+        0.35324087738990784,
+        0.6349603533744812
+      ],
+      "max": [
+        0.30179348587989807,
+        0.7795099020004272,
+        0.2688712179660797,
+        0.8059309720993042,
+        0.5333064198493958,
+        1.0,
+        1.0
+      ],
+      "min": [
+        -0.6572320461273193,
+        -1.0,
+        -0.8218780755996704,
+        -0.7686842083930969,
+        -0.9920897483825684,
+        -0.30364304780960083,
+        -1.0
+      ],
+      "q01": [
+        -0.4552440994977951,
+        -1.0,
+        -0.7419710528850555,
+        -0.475678391456604,
+        -0.9305388247966766,
+        -0.20442348539829253,
+        -0.8981475746631622
+      ],
+      "q99": [
+        0.1125871029496192,
+        0.7402709639072418,
+        0.22590837419033025,
+        0.7005072689056396,
+        0.16810789704322815,
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        true
+      ]
+    },
+    "state": {
+      "mean": [
+        -0.14554812014102936,
+        0.07433967292308807,
+        0.1436367779970169,
+        2.700871706008911,
+        0.45681115984916687,
+        -0.48229172825813293,
+        0.06972828507423401
+      ],
+      "std": [
+        0.07678605616092682,
+        0.21988309919834137,
+        0.09388109296560287,
+        1.109084963798523,
+        0.3012843728065491,
+        0.3006882667541504,
+        0.0164483655244112
+      ],
+      "max": [
+        0.014208495616912842,
+        0.5909395217895508,
+        0.3558541536331177,
+        3.1415467262268066,
+        0.9420600533485413,
+        0.09770228713750839,
+        0.07999999821186066
+      ],
+      "min": [
+        -0.37446925044059753,
+        -0.29720592498779297,
+        0.016940969973802567,
+        -3.140481948852539,
+        -0.10707131028175354,
+        -1.2001688480377197,
+        0.03238735347986221
+      ],
+      "q01": [
+        -0.3139808464050293,
+        -0.25229972660541533,
+        0.0188405978679657,
+        -3.134360270500183,
+        -0.04521620273590088,
+        -1.0301372480392457,
+        0.03585549935698509
+      ],
+      "q99": [
+        -0.013631248474121185,
+        0.5275326526165007,
+        0.33348673582077026,
+        3.1367013454437256,
+        0.8922865676879882,
+        0.056985778063535636,
+        0.07999999821186066
+      ]
+    },
+    "num_transitions": 1092,
+    "num_trajectories": 25
+  }
+}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/files.sha256.json

@@ -0,0 +1,106 @@
+{
+  "README.md": {
+    "bytes": 1155,
+    "sha256": "c9ccbe7aa494be2270decacd8f7f3f56afbd0e5c57c58529b6e9c08d72fe5402"
+  },
+  "checkpoints/model.pt": {
+    "bytes": 9785101729,
+    "sha256": "7bea6746bf7973ff7d4e4e8a457a3bfbc0ce981c80fe5f7afadca651ec5ca85e"
+  },
+  "collection_quality_report.json": {
+    "bytes": 3134,
+    "sha256": "ce3effc9d3267edf410642d0ecd3b5b7790c072e1ff5d0f79eeab4ae883fe1a1"
+  },
+  "config.full.yaml": {
+    "bytes": 2574,
+    "sha256": "e9e1231302aa3ae47efd1673f356c0004daf815f47eb9c02ba62714d89b53f0b"
+  },
+  "config.yaml": {
+    "bytes": 871,
+    "sha256": "d0612a41944af616e4750bb1217353a90a562ac978571d9ebd663ec7caae61ed"
+  },
+  "dataset_statistics.json": {
+    "bytes": 2843,
+    "sha256": "7a5fde37c6a0b220321d0acdb73b8b48b7ffe9f6295b7a0bb0e0abcf4849b653"
+  },
+  "dataset_statistics_eval.json": {
+    "bytes": 2841,
+    "sha256": "fcbc27a8dc2d8f6f8b5069300bae4fbf71ddacea91c8b23551fa652b88d0801e"
+  },
+  "launch_audit.json": {
+    "bytes": 1062,
+    "sha256": "afcbecdd81e19e8bf4d3a6fa946e7b891f126cb60d0ea8e24f5badeab961f07c"
+  },
+  "manifest.json": {
+    "bytes": 221,
+    "sha256": "97b5bea337c766d9db49c8e3a01ad93de4721dbdf965acf31e7eba4948f24421"
+  },
+  "profile.source.json": {
+    "bytes": 2190,
+    "sha256": "31a9b15d6b04182439934f0b377cb3d09be16ce21f3cf95c1049918f8a5d4984"
+  },
+  "provenance.json": {
+    "bytes": 8324,
+    "sha256": "67f1a7c42dbd772f40995fc73454cf954dda78c2408741bc9274f5259766f7ce"
+  },
+  "startup_verified.json": {
+    "bytes": 1210,
+    "sha256": "6bd0871b87954dbf4dfba1778e84ba91361741c4e7d106e97299965d2c0ba1c0"
+  },
+  "summary.jsonl": {
+    "bytes": 159,
+    "sha256": "2905acb79c5fffb33fd4689d87af3bd05d74f4184b9d36ff3eae0c64deecd930"
+  },
+  "task_contract.json": {
+    "bytes": 3063,
+    "sha256": "b40d58b27dfa9431f86f83ea5d4f36bcd652b3e3517c044638282df0d1e86c6b"
+  },
+  "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/doom-mean-reset.patch": {
+    "bytes": 9569,
+    "sha256": "4113989e16c547bd1f02f5f10f285c4d511be3a3c18d10a0aa3316bd8ff3175d"
+  },
+  "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_results.json": {
+    "bytes": 610,
+    "sha256": "7165309b2a84a8fa8250cff1dc2ca1717837f114b5aac269ee9a7e3188182e17"
+  }
+}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/launch_audit.json

@@ -0,0 +1,29 @@
+{
+  "physical_gpus": [
+    2,
+    3
+  ],
+  "seed": 0,
+  "horizon": 1,
+  "dagger_updates": 0,
+  "state_dim": 7,
+  "action_dim": 7,
+  "micro_per_gpu": 64,
+  "gradient_accumulation": 1,
+  "global_batch": 128,
+  "gradient_checkpointing": true,
+  "formal_budget": 5000,
+  "post_train_rollout_eval": false,
+  "init": "same pinnedQwen3-VL-4B backbone; scratchactionhead",
+  "user_authorization": "AfterDeadlycompletes startInterceptGPU23; no newtaskGPU01afterAnt",
+  "runtime_capacity_choice": "GCtrue64 methodpassedAnt/Deadly; Interceptmuststillcomplete20boundedthenverifyformalstartup",
+  "dataset_revision": "bba7d56c0ad5bb2b6d6ed07adb70b4f1afd5d25a",
+  "startup_verified": true,
+  "recovery_keep_last_n": 1,
+  "storage_capacity": {
+    "free_before_bytes": 266315395072,
+    "estimated_two_intercept_saves_plus_one_ant_transient_bytes": 218520000000,
+    "policy": "preserveonecompleteoptimizerrecovery;newsavebeforeoldprune;originaldata/artifactsnotdeleted"
+  },
+  "formal_startup_evidence": "/home/ubuntu/lzj/mean-profiling/intercept/vla.startup-verified.json"
+}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/manifest.json

@@ -0,0 +1,10 @@
+{
+  "format": "latency_bench_starvla_h1_v1",
+  "task": "mikasa_intercept_grab_fast",
+  "robot_type": "latency_mikasa_h1",
+  "state_dim": 7,
+  "action_dim": 7,
+  "action_horizon": 1,
+  "vla_fps": 20.0,
+  "dagger": false
+}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/profile.source.json

@@ -0,0 +1,66 @@
+{
+  "burst_model_path": "latency_burst_model.json",
+  "distribution_path": "latency_distribution.json",
+  "env_fps": 20,
+  "frame_ms": 50.0,
+  "gpu_class": "1x-rtx3090",
+  "instance_id": "instance_3a0d42681a03715c",
+  "latency_kind": "observation_to_action_latency",
+  "latency_method": "temporal",
+  "model_id": "qwenoft",
+  "n_admitted_observations": 7505,
+  "n_capacity_drops": 7495,
+  "n_observation_attempts": 15000,
+  "per_slot_summary": {
+    "0": {
+      "admitted_count": 7505,
+      "mean_observation_to_action_latency_ms": 99.05021289731565,
+      "mean_worker_service_time_ms": 98.55270653801072,
+      "p95_observation_to_action_latency_ms": 100.32732703685761,
+      "p95_worker_service_time_ms": 99.84762226343155,
+      "p99_worker_service_time_ms": 100.37557450771331
+    }
+  },
+  "provenance": {
+    "base_config": "configs/examples/mikasa/official_starvla_profile.yaml",
+    "checkpoint_kind": "best",
+    "model_artifact": {
+      "checkpoint": "checkpoints/model.pt",
+      "model_config": "config.yaml",
+      "path_in_repo": "zero-latency/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/mikasa_intercept_grab_fast_qwenoft_formal_latest_r8gpu8_20260907T082154Z",
+      "repo_id": "latency-sensitive-bench/benchmark-models",
+      "source": "hf"
+    },
+    "session_ids": [
+      0,
+      1,
+      2,
+      3,
+      4
+    ]
+  },
+  "sample_model_type": "hidden_regime",
+  "source_run_id": "20260909T044501695676Z",
+  "summary": {
+    "frame_ms": 50.0,
+    "max_ms": 101.19136601686478,
+    "mean_effective_frames": 1.981004257946313,
+    "mean_ms": 99.05021289731565,
+    "min_ms": 57.612886011600494,
+    "n_samples": 7505,
+    "p50_frames": 1.9978350806236267,
+    "p50_ms": 99.89175403118134,
+    "p90_frames": 2.0040290887355803,
+    "p90_ms": 100.20145443677902,
+    "p95_frames": 2.006546540737152,
+    "p95_ms": 100.32732703685761,
+    "p99_frames": 2.0128874400138854,
+    "p99_ms": 100.64437200069428,
+    "prob_latency_gt_1_frame": 1.0,
+    "prob_latency_gt_2_frames": 0.2139906728847435,
+    "prob_latency_gt_3_frames": 0.0,
+    "std_ms": 4.583189247461202
+  },
+  "visualization_path": "latency_profile.png",
+  "workload_id": "mikasa_intercept_grab_fast"
+}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/provenance.json

@@ -0,0 +1,415 @@
+{
+  "task": "intercept",
+  "model": "QwenOFT",
+  "run_id": "intercept-qwenoft-mean-3090-h1-s0",
+  "training_status": "complete",
+  "training_hardware": {
+    "gpu": "RTX PRO6000 Blackwell Server Edition",
+    "count": 2,
+    "physical_indices": [
+      2,
+      3
+    ]
+  },
+  "profile": {
+    "mean_ms": 99.05021289731565,
+    "profile": "/home/ubuntu/lzj/profiles/intercept-published/profiles/qwenoft/1x-rtx3090/mikasa_intercept_grab_fast/instance_3a0d42681a03715c/profile.json",
+    "sha256": "31a9b15d6b04182439934f0b377cb3d09be16ce21f3cf95c1049918f8a5d4984"
+  },
+  "profile_hardware": "RTX3090",
+  "condition": "fixed_mean_profile",
+  "horizon": 1,
+  "seed": 0,
+  "state_dim": 7,
+  "action_dim": 7,
+  "micro_per_gpu": 64,
+  "gradient_accumulation": 1,
+  "global_batch": 128,
+  "gradient_checkpointing": true,
+  "precision": "BF16",
+  "backend": "ZeRO2",
+  "teacher": {
+    "repo_id": "latency-sensitive-bench/benchmark-models",
+    "revision": "ba7e4f3b6047b82711de01432a7d5bf3895df8b4",
+    "path": "latency-aware/mikasa-intercept-grab-fast/small-policy/ppo-qwenoft-rtx3090-mean-h1-s123-v1",
+    "status": "complete_verified",
+    "local_checkpoint": "/home/ubuntu/lzj/mean-profiling/intercept/teacher-publication/checkpoint.pt",
+    "completed_env_steps": 149975040,
+    "completed_iterations": 2441,
+    "selected_env_steps": 32256000,
+    "selected_iteration": 525,
+    "tensor_count": 17,
+    "weights_only_export": true,
+    "tensor_equality_verified": true,
+    "hash_readback_verified": true,
+    "url": "https://huggingface.co/latency-sensitive-bench/benchmark-models/tree/ba7e4f3b6047b82711de01432a7d5bf3895df8b4/latency-aware/mikasa-intercept-grab-fast/small-policy/ppo-qwenoft-rtx3090-mean-h1-s123-v1",
+    "completed_utc": "2026-09-30T16:34:58.230455+00:00"
+  },
+  "dataset": {
+    "repo_id": "latency-sensitive-bench/Standard-Pipeline",
+    "repo_type": "dataset",
+    "path": "mikasa_intercept_grab_fast/profile_latency/QwenOFT/mean_rtx3090_h1_s123_275ep_v1",
+    "revision": "bba7d56c0ad5bb2b6d6ed07adb70b4f1afd5d25a",
+    "data_uploaded_revision": "d746ca2b6461ea211589cdf4d796512335aa34ae",
+    "status": "complete_verified",
+    "episodes": 275,
+    "rows": 12211,
+    "raw_npz_episodes": 275,
+    "dataset_config": "mikasa_intercept_grab_fast_mean_qwenoft_h1_episode_manifest",
+    "viewer_rows": "per_episode metadata manifest",
+    "payload_files_verified": 308,
+    "checksum_manifest_verified": true,
+    "url": "https://huggingface.co/datasets/latency-sensitive-bench/Standard-Pipeline/tree/bba7d56c0ad5bb2b6d6ed07adb70b4f1afd5d25a/mikasa_intercept_grab_fast/profile_latency/QwenOFT/mean_rtx3090_h1_s123_275ep_v1",
+    "completed_utc": "2026-09-30T17:34:21.416539+00:00"
+  },
+  "collection": {
+    "episodes": 275,
+    "unique_seeds": 275,
+    "rows": 12211,
+    "length_min": 37,
+    "length_max": 55,
+    "all_success_once": true,
+    "state_dim": 7,
+    "action_dim": 7,
+    "rgb_shape": [
+      128,
+      128,
+      6
+    ],
+    "finite_numeric": true,
+    "latency_raw_frames": 2,
+    "applied_command_matches_source_issued_action": true,
+    "action_horizon": 1,
+    "dagger": false,
+    "raw_bytes": 1204453838,
+    "episode_seeds": [
+      0,
+      1,
+      10,
+      180,
+      181,
+      182,
+      183,
+      184,
+      185,
+      186,
+      187,
+      188,
+      189,
+      11,
+      190,
+      191,
+      192,
+      193,
+      194,
+      195,
+      196,
+      197,
+      198,
+      199,
+      12,
+      20,
+      21,
+      22,
+      23,
+      24,
+      25,
+      26,
+      27,
+      28,
+      29,
+      13,
+      200,
+      201,
+      202,
+      203,
+      204,
+      205,
+      206,
+      207,
+      208,
+      209,
+      14,
+      210,
+      211,
+      212,
+      213,
+      214,
+      215,
+      216,
+      217,
+      218,
+      219,
+      15,
+      220,
+      221,
+      222,
+      223,
+      224,
+      225,
+      226,
+      227,
+      228,
+      229,
+      16,
+      230,
+      231,
+      232,
+      233,
+      234,
+      235,
+      236,
+      237,
+      238,
+      239,
+      17,
+      240,
+      241,
+      242,
+      243,
+      244,
+      245,
+      246,
+      247,
+      248,
+      249,
+      18,
+      250,
+      251,
+      252,
+      253,
+      254,
+      255,
+      256,
+      257,
+      258,
+      259,
+      19,
+      260,
+      261,
+      262,
+      263,
+      264,
+      265,
+      266,
+      267,
+      268,
+      269,
+      2,
+      100,
+      270,
+      271,
+      272,
+      273,
+      274,
+      30,
+      31,
+      32,
+      33,
+      34,
+      101,
+      35,
+      36,
+      37,
+      38,
+      39,
+      40,
+      41,
+      42,
+      43,
+      44,
+      102,
+      45,
+      46,
+      47,
+      48,
+      49,
+      50,
+      51,
+      52,
+      53,
+      54,
+      103,
+      55,
+      56,
+      57,
+      58,
+      59,
+      60,
+      61,
+      62,
+      63,
+      64,
+      104,
+      65,
+      66,
+      67,
+      68,
+      69,
+      70,
+      71,
+      72,
+      73,
+      74,
+      105,
+      75,
+      76,
+      77,
+      78,
+      79,
+      80,
+      81,
+      82,
+      83,
+      84,
+      106,
+      85,
+      86,
+      87,
+      88,
+      89,
+      90,
+      91,
+      92,
+      93,
+      94,
+      107,
+      95,
+      96,
+      97,
+      98,
+      99,
+      108,
+      109,
+      3,
+      110,
+      111,
+      112,
+      113,
+      114,
+      115,
+      116,
+      117,
+      118,
+      119,
+      4,
+      120,
+      121,
+      122,
+      123,
+      124,
+      125,
+      126,
+      127,
+      128,
+      129,
+      5,
+      130,
+      131,
+      132,
+      133,
+      134,
+      135,
+      136,
+      137,
+      138,
+      139,
+      6,
+      140,
+      141,
+      142,
+      143,
+      144,
+      145,
+      146,
+      147,
+      148,
+      149,
+      7,
+      150,
+      151,
+      152,
+      153,
+      154,
+      155,
+      156,
+      157,
+      158,
+      159,
+      8,
+      160,
+      161,
+      162,
+      163,
+      164,
+      165,
+      166,
+      167,
+      168,
+      169,
+      9,
+      170,
+      171,
+      172,
+      173,
+      174,
+      175,
+      176,
+      177,
+      178,
+      179
+    ],
+    "train_episodes": 250,
+    "val_episodes": 25,
+    "train_rows": 11119,
+    "val_rows": 1092,
+    "teacher_checkpoint": "/home/ubuntu/lzj/mean-profiling/intercept/teacher-publication/checkpoint.pt",
+    "teacher_sha256": "e47526f03fad3ed251dd78f3cacfba8aa8b60bd1578f40bd7a1393e17298e973",
+    "fixed_latency_ms": 99.05021289731565,
+    "conversion_complete": true
+  },
+  "runtime": {
+    "repo": "c3c6a39365a151e9b7a5e215452fd64e957c2b29",
+    "starvla": "ccca13c5177fe3d3c884b6e2de4965d916016649",
+    "runtime_fixes": [
+      "mean-profile-preparation.patch",
+      "gym-language-contract.patch",
+      "loader-spawn-cache.patch",
+      "loader-spawn-test.patch",
+      "doom-mean-reset.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"
+    }
+  },
+  "source_reference_revision": "180da0effdaaaeacef218cf5fbaa4db41af67495",
+  "performance_configuration_policy": "User authorized throughput priority; reference scientific parameters and actual runtime differences are documented by full config/launch audit",
+  "training_results": {
+    "training_complete": true,
+    "optimizer_updates": 5000,
+    "selection": "final_5000",
+    "train_loss": 0.0017812054138630629,
+    "validation_action_loss": 0.03070492297410965,
+    "validation_samples": 128,
+    "final_learning_rates": {
+      "learning_rate/base": 1e-06,
+      "learning_rate/action_model": 3.3333333333333333e-06
+    },
+    "closed_loop_evaluation": "pending until all four VLA trainings complete",
+    "tensor_keys": 732,
+    "serialized_tensor_elements_including_tied_aliases": 4892415495,
+    "weights_only_load": "passed",
+    "all_tensors_finite": true,
+    "completion_proof": {
+      "wrapper_exit": 0
+    }
+  },
+  "optimizer_recovery_states": "retained on host; inference weights only in this bundle",
+  "wandb_url": "https://wandb.ai/dongqianyu99-zhejiang-university/starvla_tasks/runs/mean3090-intercept-vla-blackwell-20260930"
+}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/startup_verified.json

@@ -0,0 +1,56 @@
+{
+  "snapshot_utc": "2026-10-01T00:14:27Z",
+  "unit": "mean-profile-intercept-vla-20261001-g23-gc.service",
+  "wrapper_pid": 3705659,
+  "rank_pids": [
+    3706042,
+    3706043
+  ],
+  "physical_gpus": [
+    2,
+    3
+  ],
+  "formal_updates_completed": 55,
+  "formal_budget": 5000,
+  "preflight_updates_discarded": 20,
+  "consecutive_observed_updates": [
+    50,
+    51,
+    52,
+    53,
+    54,
+    55
+  ],
+  "loss_at_logged_steps": {
+    "50": 0.33528196811676025,
+    "55": 0.2620769143104553
+  },
+  "finite_loss_verified": true,
+  "micro_per_gpu": 64,
+  "global_batch": 128,
+  "gradient_accumulation": 1,
+  "gradient_checkpointing": true,
+  "seed": 0,
+  "horizon": 1,
+  "dagger_updates": 0,
+  "gpu_used_mib": [
+    55707,
+    65727
+  ],
+  "gpu_utilization_percent": [
+    100,
+    100
+  ],
+  "host_memory_current_bytes": 13883498496,
+  "host_memory_peak_bytes": 38294556672,
+  "seconds_per_update_recent": 2.74,
+  "data_wait_seconds_recent": 0.00069506,
+  "estimated_completion_utc_range": [
+    "2026-10-01T04:45:00Z",
+    "2026-10-01T05:15:00Z"
+  ],
+  "completion_check_utc": "2026-10-01T05:30:00Z",
+  "training_complete": false,
+  "publication_complete": false,
+  "closed_loop_evaluation_complete": false
+}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/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}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/task_contract.json

@@ -0,0 +1,131 @@
+{
+  "format": "latency_bench_starvla_h1_v1",
+  "task": "mikasa_intercept_grab_fast",
+  "robot_type": "latency_mikasa_h1",
+  "state_dim": 7,
+  "action_dim": 7,
+  "action_horizon": 1,
+  "fps": 20.0,
+  "state_key": "observation.state",
+  "action_key": "action",
+  "video_keys": [
+    "video.top",
+    "video.wrist"
+  ],
+  "auxiliary": {},
+  "normalization": {
+    "state": {
+      "mean": [
+        -0.14554812014102936,
+        0.07433967292308807,
+        0.1436367779970169,
+        2.700871706008911,
+        0.45681115984916687,
+        -0.48229172825813293,
+        0.06972828507423401
+      ],
+      "std": [
+        0.07678605616092682,
+        0.21988309919834137,
+        0.09388109296560287,
+        1.109084963798523,
+        0.3012843728065491,
+        0.3006882667541504,
+        0.0164483655244112
+      ],
+      "min": [
+        -0.37446925044059753,
+        -0.29720592498779297,
+        0.016940969973802567,
+        -3.140481948852539,
+        -0.10707131028175354,
+        -1.2001688480377197,
+        0.03238735347986221
+      ],
+      "max": [
+        0.014208495616912842,
+        0.5909395217895508,
+        0.3558541536331177,
+        3.1415467262268066,
+        0.9420600533485413,
+        0.09770228713750839,
+        0.07999999821186066
+      ],
+      "q01": [
+        -0.3139808464050293,
+        -0.25229972660541533,
+        0.0188405978679657,
+        -3.134360270500183,
+        -0.04521620273590088,
+        -1.0301372480392457,
+        0.03585549935698509
+      ],
+      "q99": [
+        -0.013631248474121185,
+        0.5275326526165007,
+        0.33348673582077026,
+        3.1367013454437256,
+        0.8922865676879882,
+        0.056985778063535636,
+        0.07999999821186066
+      ]
+    },
+    "action": {
+      "mean": [
+        -0.048697348684072495,
+        0.24886493384838104,
+        -0.16479986906051636,
+        -0.032609131187200546,
+        -0.501691997051239,
+        0.44703611731529236,
+        0.6793776154518127
+      ],
+      "std": [
+        0.1079770028591156,
+        0.47203660011291504,
+        0.14475314319133759,
+        0.2849433720111847,
+        0.318120539188385,
+        0.35324087738990784,
+        0.6349603533744812
+      ],
+      "min": [
+        -0.6572320461273193,
+        -1.0,
+        -0.8218780755996704,
+        -0.7686842083930969,
+        -0.9920897483825684,
+        -0.30364304780960083,
+        -1.0
+      ],
+      "max": [
+        0.30179348587989807,
+        0.7795099020004272,
+        0.2688712179660797,
+        0.8059309720993042,
+        0.5333064198493958,
+        1.0,
+        1.0
+      ],
+      "q01": [
+        -0.4552440994977951,
+        -1.0,
+        -0.7419710528850555,
+        -0.475678391456604,
+        -0.9305388247966766,
+        -0.20442348539829253,
+        -0.8981475746631622
+      ],
+      "q99": [
+        0.1125871029496192,
+        0.7402709639072418,
+        0.22590837419033025,
+        0.7005072689056396,
+        0.16810789704322815,
+        1.0,
+        1.0
+      ]
+    }
+  },
+  "prompt": "Intercept the rolling ball and grasp it to stop it."
+}

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/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

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/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")

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/training-code/doom-mean-reset.patch

@@ -0,0 +1,192 @@
+diff --git a/latency_bench/data/sf_teacher_rollout.py b/latency_bench/data/sf_teacher_rollout.py
+index ef45d95b..e32b942e 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
+@@ -1272,9 +1283,8 @@ def _build_rollout_driver_runtime(
+         driver = BatchedTeacherRolloutDriver(runtime["envs"][0], int(num_envs))
+     else:
+         recreate_env = None
+-        if integration_name == "deadly_corridor" and (
+-            "export_zero_action_history" in config or getattr(cfg, "rollout_profile_seed_by_reset", False)
+-        ):
++        # ViZDoom applies reset seeds only while creating its native game.
++        if integration_name == "deadly_corridor":
+             from sample_factory.algo.utils.make_env import make_env_func_batched
+             from sample_factory.utils.attr_dict import AttrDict
+ 
+diff --git a/tests/data/test_sf_teacher_dataset_export.py b/tests/data/test_sf_teacher_dataset_export.py
+index e6b25233..61a58118 100644
+--- a/tests/data/test_sf_teacher_dataset_export.py
++++ b/tests/data/test_sf_teacher_dataset_export.py
+@@ -498,8 +498,11 @@ def test_serial_rollout_driver_continues_life_loss_with_seedless_reset(monkeypat
+     assert info == {"reset_seed": None}
+ 
+ 
+-def test_serial_teacher_replay_starts_each_seed_in_a_fresh_environment(monkeypatch):
+-    sf_teacher_drivers = _load_sf_teacher_drivers_for_direct_test(monkeypatch)
++@pytest.mark.parametrize("latency_type", ["fixed", "profile_sample"])
++def test_serial_teacher_replay_starts_each_seed_in_a_fresh_environment(monkeypatch, tmp_path, latency_type):
++    from contextlib import nullcontext
++    from latency_bench.data import sf_teacher_rollout
++    from sample_factory.algo.utils import make_env
+ 
+     class StatefulEnv:
+         action_space = object()
+@@ -514,13 +517,17 @@ def test_serial_teacher_replay_starts_each_seed_in_a_fresh_environment(monkeypat
+         def close(self):
+             return None
+ 
+-    driver = sf_teacher_drivers.SerialTeacherRolloutDriver(
+-        [StatefulEnv()], recreate_env=lambda _slot: StatefulEnv()
++    cfg = SimpleNamespace(frame_stack=4, action_history_decisions=8, latency_type=latency_type)
++    monkeypatch.setattr(sf_teacher_rollout, "sample_factory_init_log", lambda _: nullcontext())
++    monkeypatch.setattr(sf_teacher_rollout, "build_sample_factory_rollout_cfg", lambda _: (object(), cfg))
++    monkeypatch.setattr(sf_teacher_rollout, "build_env_and_policy", lambda *args: {"envs": [StatefulEnv()]})
++    monkeypatch.setattr(make_env, "make_env_func_batched", lambda *args, **kwargs: StatefulEnv())
++    _, _, driver, _ = sf_teacher_rollout._build_rollout_driver_runtime(
++        {"env_name": "deadly_corridor", "integration_name": "deadly_corridor"},
++        output_dir=tmp_path, num_envs=1, export_env_raw_rgb_frames=False,
+     )
+-
+     driver.reset_slots({0: 42})
+     observation, _ = driver.reset_slots({0: 43})[0]
+-
+     assert observation["obs"].tolist() == [43, 1]
+ 
+ 
+@@ -670,19 +677,22 @@ def test_collect_teacher_dataset_probe_retries_until_episode_is_accepted(monkeyp
+         def action_distribution(self):
+             return object()
+ 
++    reset_seeds = []
++
+     class FakeEnv:
+         def __init__(self):
+             self.action_space = object()
+-            self.reset_seeds: list[int] = []
++            self.reset_seeds = reset_seeds
+ 
+         def reset(self, seed=None):
+             self.reset_seeds.append(seed)
++            self.current_seed = seed
+             return np.zeros((3, 2, 2), dtype=np.uint8), [
+                 {"env_raw_rgb_frame_stack": _env_raw_rgb_frame_stack(5)}
+             ]
+ 
+         def step(self, action):
+-            reward = 3000.0 if len(self.reset_seeds) == 1 else 3000.001
++            reward = 3000.0 if self.current_seed == 11 else 3000.001
+             return (
+                 np.zeros((3, 2, 2), dtype=np.uint8),
+                 np.float32(reward),
+@@ -702,6 +712,8 @@ def test_collect_teacher_dataset_probe_retries_until_episode_is_accepted(monkeyp
+ 
+     fake_env = FakeEnv()
+     fake_actor = FakeActorCritic()
++    from sample_factory.algo.utils import make_env
++    monkeypatch.setattr(make_env, "make_env_func_batched", lambda *args, **kwargs: FakeEnv())
+     fake_cfg = SimpleNamespace(frame_stack=1, env_fps=35.0, obs_fps=17.5, seed=11)
+ 
+     monkeypatch.setattr(exporter, "build_sample_factory_rollout_cfg", lambda config: (object(), fake_cfg))
+@@ -1273,7 +1285,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 +1320,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 +1339,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 +1373,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 +1401,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")

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/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"]

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/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."""

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/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)

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/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:

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/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")

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/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

latency-aware/mikasa-intercept-grab-fast/vla/starvla-qwenoft-h1/intercept-qwenoft-mean-3090-h1-s0/training_results.json

@@ -0,0 +1,20 @@
+{
+  "training_complete": true,
+  "optimizer_updates": 5000,
+  "selection": "final_5000",
+  "train_loss": 0.0017812054138630629,
+  "validation_action_loss": 0.03070492297410965,
+  "validation_samples": 128,
+  "final_learning_rates": {
+    "learning_rate/base": 1e-06,
+    "learning_rate/action_model": 3.3333333333333333e-06
+  },
+  "closed_loop_evaluation": "pending until all four VLA trainings complete",
+  "tensor_keys": 732,
+  "serialized_tensor_elements_including_tied_aliases": 4892415495,
+  "weights_only_load": "passed",
+  "all_tensors_finite": true,
+  "completion_proof": {
+    "wrapper_exit": 0
+  }
+}