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

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/README.md

@@ -0,0 +1,8 @@
+# ant QwenOFT H1 mean-profile VLA
+
+Training complete:5000 optimizer updates, seed42, globalbatch128, microbatch64 perGPU, accumulation1, gradientcheckpointingTrue, BF16/ZeRO2. Hardware:2 RTX PRO6000 Blackwell GPUs. Fixed mean90.56460638563993ms 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 dimensions27/8.
+
+Puretensor weights checkCPUweights_only/allfinite passed. Portable inferenceconfig, normalization/statistics, taskcontract, sourcecode/profile/data provenance included. Full optimizer recovery remains local. Finalofflinevalidation loss0.018865827471017838; 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/6246ccc57caf32da0e02547c3e3e1080894d4035/ant/profile_latency/OpenVLA/mean_rtx3090_s2222_100ep_v1
+Training:https://wandb.ai/dongqianyu99-zhejiang-university/latency-sensitive-bench/runs/mean3090-ant-vla-blackwell-20260930

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/checkpoints/model.pt

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

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/collection_quality_report.json

@@ -0,0 +1,14 @@
+{
+  "task": "ant",
+  "teacher_checkpoint": "/home/ubuntu/lzj/mean-profiling/ant/teacher-publication/checkpoint.pth",
+  "teacher_sha256": "a7899b32beb9b2cfec1bcff85b285dddc243a73f1c322fafefae599df0f67b3b",
+  "episodes": 100,
+  "train_episodes": 90,
+  "val_episodes": 10,
+  "train_rows": 90000,
+  "val_rows": 10000,
+  "fixed_latency_ms": 90.56460638563993,
+  "action_horizon": 1,
+  "unique_seeds": 100,
+  "state_dim": 27
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/config.full.yaml

@@ -0,0 +1,438 @@
+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: 8
+    state_dim: 27
+    action_hidden_dim: 2560
+    future_action_window_size: 0
+    past_action_window_size: 0
+    loss_type: l1
+    state_encoding: discretized_text
+    task_objective: null
+    action_horizon: 1
+    action_env_dim: 8
+  kv_memory:
+    enabled: false
+    window: 4
+    rollout_len: 8
+    packed_train: false
+    rebased_sink: true
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    include_state: true
+    data_root_dir: /home/ubuntu/lzj/mean-profiling/ant/data
+    data_mix: lerobot
+    eval_data_mix: lerobot__val
+    custom_mixtures_path: /home/ubuntu/lzj/mean-profiling/ant/data/_generated_mixtures/lerobot.json
+    action_type: discrete
+    sequential_step_sampling: false
+    eval_sequential_step_sampling: null
+    num_workers: 2
+    eval_num_workers: 2
+    prefetch_factor: 4
+    persistent_workers: true
+    pin_memory: true
+    shuffle: true
+    action_balance:
+      enabled: false
+      strategy: balanced_epoch
+      action_key: action_id
+      target_flap_fraction: 0.3
+      noop_id: 0
+      flap_id: 1
+    latency_curriculum:
+      enabled: false
+      strategy: exclusive
+      latencies: null
+      phase_steps: null
+      phase_distributions: null
+      new_latency_passes: 1.0
+      replay_passes: 0.25
+      target_total_passes: 2.0
+      final_equalization: true
+      step_budget_mode: auto
+      eval_at_phase_end: false
+      save_at_phase_end: false
+      computed_plan: null
+    per_device_batch_size: 64
+    load_all_data_for_training: true
+    num_obs_frames: 1
+    image_mode: single
+    prompt_mode: raw
+    stitch_grid:
+    - 2
+    - 2
+    obs_image_size: null
+    video_backend: torchvision_av
+    gymnasium_task_contract:
+      action_space:
+        dtype: float32
+        high:
+        - 1.0
+        - 1.0
+        - 1.0
+        - 1.0
+        - 1.0
+        - 1.0
+        - 1.0
+        - 1.0
+        labels:
+        - back_right_hip_torque
+        - back_right_ankle_torque
+        - front_left_hip_torque
+        - front_left_ankle_torque
+        - front_right_hip_torque
+        - front_right_ankle_torque
+        - back_left_hip_torque
+        - back_left_ankle_torque
+        low:
+        - -1.0
+        - -1.0
+        - -1.0
+        - -1.0
+        - -1.0
+        - -1.0
+        - -1.0
+        - -1.0
+        type: box
+      base_prompt: Make the Ant move forward as fast as possible without falling.
+        Predict eight continuous torques in [-1, 1] ordered as back right hip, back
+        right ankle, front left hip, front left ankle, front right hip, front right
+        ankle, back left hip, and back left ankle.
+      env_fps: 10.0
+      env_id: LatencyBench/AntContinuous-v0
+      frame_stack: 1
+      make_kwargs:
+        base_env_id: Ant-v4
+        base_make_kwargs:
+          exclude_current_positions_from_observation: true
+          use_contact_forces: false
+        render_mode: rgb_array
+      noop_action:
+      - 0.0
+      - 0.0
+      - 0.0
+      - 0.0
+      - 0.0
+      - 0.0
+      - 0.0
+      - 0.0
+      obs_fps: 10.0
+      registration_imports:
+      - latency_bench.envs.gymnasium_ant
+      state_space:
+        labels:
+        - torso_z
+        - torso_quaternion_w
+        - torso_quaternion_x
+        - torso_quaternion_y
+        - torso_quaternion_z
+        - front_left_hip_angle
+        - front_left_ankle_angle
+        - front_right_hip_angle
+        - front_right_ankle_angle
+        - back_left_hip_angle
+        - back_left_ankle_angle
+        - back_right_hip_angle
+        - back_right_ankle_angle
+        - torso_x_velocity
+        - torso_y_velocity
+        - torso_z_velocity
+        - torso_angular_velocity_x
+        - torso_angular_velocity_y
+        - torso_angular_velocity_z
+        - front_left_hip_angular_velocity
+        - front_left_ankle_angular_velocity
+        - front_right_hip_angular_velocity
+        - front_right_ankle_angular_velocity
+        - back_left_hip_angular_velocity
+        - back_left_ankle_angular_velocity
+        - back_right_hip_angular_velocity
+        - back_right_ankle_angular_velocity
+      task_name: ant_rgb_state
+    active_action_dim: 8
+dataset:
+  source_hf: ''
+  config_name: null
+  source_subdir: null
+  converted_name: ant_rgb_state_profile_20260911T112232Z
+  single_source_hf: ''
+  mixed_source_hf: ''
+  single_converted_name: ant_rgb_state_profile_20260911T112232Z
+  mixed_converted_name: ant_rgb_state_profile_20260911T112232Z
+  single_latency_filter: null
+  mixed_latency_filter: null
+  force_download: false
+  setup_force: false
+  skip_verification: false
+  target_latency_unit: observation_steps
+  verify_rows: 200
+  max_episodes: null
+  episodes_per_latency: null
+  latency_filter: null
+  debug_subset:
+    enabled: false
+    max_episodes: 5
+    suffix: debug
+base_model:
+  repo_id: StarVLA/Qwen3-VL-4B-Instruct-Action
+initialization:
+  checkpoint_local_dir: null
+  checkpoint_hf_repo_id: null
+  checkpoint_filename: null
+trainer:
+  max_train_steps: 5000
+  num_warmup_steps: 100
+  save_interval: 500
+  eval_interval: 500
+  eval_num_batches: 200
+  per_latency_eval_num_batches: null
+  eval_action_classification: false
+  eval_action_classification_interval: null
+  cc_f1_tolerance: 1
+  learning_rate:
+    base: 2.0e-05
+    qwen_vl_interface: 1.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
+    vlm: 0.1
+  max_grad_norm: 1.0
+  weight_decay: 0.0
+  logging_frequency: 1
+  profile_timing:
+    enabled: false
+    log_interval: 10
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 1
+  distributed_backend: deepspeed
+  is_resume: false
+  pretrained_checkpoint: null
+  resume_step: 0
+  reload_modules: null
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+    fused: true
+  save_format: pt
+workspace_dir: WORKSPACE_DIR
+run_root_dir: /home/ubuntu/lzj/mean-profiling/ant/vla
+seed: 42
+wandb_entity: dongqianyu99-zhejiang-university
+wandb_project: latency-sensitive-bench
+auth:
+  env_file: null
+  hf_token_env: HF_TOKEN
+  wandb_api_key_env: WANDB_API_KEY
+paths:
+  run_root_dir: /mnt/results/latency-sensitive-bench/profile-latency-new/ant/20260911T112232Z-ant/vla
+  dataset_local_dir: /mnt/results/latency-sensitive-bench/profile-latency-new/ant/20260911T112232Z-ant/converted
+  dataset_cache_dir: null
+  base_model_dir: /home/ubuntu/lzj/models/Qwen3-VL-4B-Instruct
+  accelerate_config: starVLA/config/deepseeds/deepspeed_zero2_cpu_offload.yaml
+launch:
+  use_accelerate: true
+  gpus: null
+  num_processes: 1
+  dry_run: false
+conda:
+  enabled: true
+  env_name: null
+rl_games:
+  model_alias: openvla
+  env_eval:
+    image_size: 224
+    frameskip: 1
+    image_transform: raw_rgb
+    prompt_mode: raw
+    ghost_trail:
+      history_frames: 5
+      gamma: 1.3
+      min_alpha: 35
+      scroll_px_per_step: 4.0
+      ground_fraction: 0.22
+    seed: 42
+    fixed_episode_seeds: true
+    latency_seed_stride: 0
+    task_seed_stride: 0
+    task_description: ''
+    eval_parallel_envs: 5
+    action_chunk_execution:
+      enabled: false
+      chunk_size: null
+    enabled: false
+    gymnasium:
+      make_kwargs: {}
+    mid_train:
+      enabled: false
+      interval_steps: 100
+      latencies:
+      - 0
+      num_episodes: 5
+      max_steps_per_episode: 3600
+    post_train:
+      enabled: false
+      latencies:
+      - 0
+      - 1
+      - 2
+      - 3
+      - 4
+      num_episodes: 5
+      max_steps_per_episode: 3600
+    eval_backend: latency_bench
+    distributed_mode: none
+    vectorized:
+      enabled: false
+      batch_size: 1
+    latency:
+      prompt_map_path: null
+      mode: single
+      values:
+      - 1
+      - 2
+  task: gymnasium
+  gymnasium:
+    task_contract:
+      action_space:
+        dtype: float32
+        high:
+        - 1.0
+        - 1.0
+        - 1.0
+        - 1.0
+        - 1.0
+        - 1.0
+        - 1.0
+        - 1.0
+        labels:
+        - back_right_hip_torque
+        - back_right_ankle_torque
+        - front_left_hip_torque
+        - front_left_ankle_torque
+        - front_right_hip_torque
+        - front_right_ankle_torque
+        - back_left_hip_torque
+        - back_left_ankle_torque
+        low:
+        - -1.0
+        - -1.0
+        - -1.0
+        - -1.0
+        - -1.0
+        - -1.0
+        - -1.0
+        - -1.0
+        type: box
+      base_prompt: Make the Ant move forward as fast as possible without falling.
+        Predict eight continuous torques in [-1, 1] ordered as back right hip, back
+        right ankle, front left hip, front left ankle, front right hip, front right
+        ankle, back left hip, and back left ankle.
+      env_fps: 10.0
+      env_id: LatencyBench/AntContinuous-v0
+      frame_stack: 1
+      make_kwargs:
+        base_env_id: Ant-v4
+        base_make_kwargs:
+          exclude_current_positions_from_observation: true
+          use_contact_forces: false
+        render_mode: rgb_array
+      noop_action:
+      - 0.0
+      - 0.0
+      - 0.0
+      - 0.0
+      - 0.0
+      - 0.0
+      - 0.0
+      - 0.0
+      obs_fps: 10.0
+      registration_imports:
+      - latency_bench.envs.gymnasium_ant
+      state_space:
+        labels:
+        - torso_z
+        - torso_quaternion_w
+        - torso_quaternion_x
+        - torso_quaternion_y
+        - torso_quaternion_z
+        - front_left_hip_angle
+        - front_left_ankle_angle
+        - front_right_hip_angle
+        - front_right_ankle_angle
+        - back_left_hip_angle
+        - back_left_ankle_angle
+        - back_right_hip_angle
+        - back_right_ankle_angle
+        - torso_x_velocity
+        - torso_y_velocity
+        - torso_z_velocity
+        - torso_angular_velocity_x
+        - torso_angular_velocity_y
+        - torso_angular_velocity_z
+        - front_left_hip_angular_velocity
+        - front_left_ankle_angular_velocity
+        - front_right_hip_angular_velocity
+        - front_right_ankle_angular_velocity
+        - back_left_hip_angular_velocity
+        - back_left_ankle_angular_velocity
+        - back_right_hip_angular_velocity
+        - back_right_ankle_angular_velocity
+      task_name: ant_rgb_state
+    task_name: ant_rgb_state
+  initialization_mode: scratch
+  action_carrier: native
+model: openvla
+env: gymnasium
+init: scratch
+mode: single
+checkpoint:
+  load: none
+  hf_repo_id: null
+  save_best_model: false
+  save_final_model: true
+  save_pt_file: true
+  save_training_state: true
+  save_safetensors_file: false
+  local:
+    keep_last_n: 2
+  sync:
+    enabled: false
+    repo_id: null
+    keep_last_n: 0
+    sync_every_n_checkpoints: 1
+    resume_policy: local_latest
+run_id: ant-qwenoft-mean-3090-s42
+output_dir: /home/ubuntu/lzj/mean-profiling/ant/vla/ant-qwenoft-mean-3090-s42
+config_yaml: /home/ubuntu/lzj/mean-profiling/ant/vla.yaml
+is_debug: false
+version_id: '0.21'
+wandb_name: ant-qwenoft-mean-3090-s42
+wandb_group: mean-profiling-3090-blackwell
+training_latency_condition: mean-profile
+wandb_tags:
+- ant
+- qwenoft
+- mean-profile
+- h1
+- blackwell
+- throughput-priority
+- micro64-global128
+- gradient-checkpointing

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/config.yaml

@@ -0,0 +1,45 @@
+framework:
+  name: QwenOFT
+  qwenvl:
+    base_vlm: Qwen/Qwen3-VL-4B-Instruct
+    attn_implementation: flash_attention_2
+    enable_gradient_checkpointing: true
+  action_model:
+    action_model_type: MLP
+    action_dim: 8
+    state_dim: 27
+    action_hidden_dim: 2560
+    future_action_window_size: 0
+    past_action_window_size: 0
+    loss_type: l1
+    state_encoding: discretized_text
+    task_objective: null
+    action_horizon: 1
+    action_env_dim: 8
+  kv_memory:
+    enabled: false
+    window: 4
+    rollout_len: 8
+    packed_train: false
+    rebased_sink: true
+version_id: '0.21'
+datasets:
+  vla_data:
+    include_state: true
+    obs_image_size: null
+    image_mode: single
+    stitch_grid:
+    - 2
+    - 2
+    num_obs_frames: 1
+trainer:
+  pretrained_checkpoint: null
+backbone:
+  repo_id: Qwen/Qwen3-VL-4B-Instruct
+  revision: ebb281ec70b05090aa6165b016eac8ec08e71b17
+  allow_patterns:
+  - '*.json'
+  - '*.safetensors'
+  - '*.txt'
+  - '*.jinja'
+  - '*.model'

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/dataset_statistics.json

@@ -0,0 +1,254 @@
+{
+  "new_embodiment": {
+    "action": {
+      "mean": [
+        -0.19492579996585846,
+        0.04469946771860123,
+        0.10307926684617996,
+        0.06500687450170517,
+        -0.07598323374986649,
+        0.10469535738229752,
+        -0.027709610760211945,
+        0.1776909977197647
+      ],
+      "std": [
+        0.2916632294654846,
+        0.18000850081443787,
+        0.3355802595615387,
+        0.2449917495250702,
+        0.3372354805469513,
+        0.055361147969961166,
+        0.29043376445770264,
+        0.06610933691263199
+      ],
+      "max": [
+        0.6329135894775391,
+        0.6100203394889832,
+        1.0,
+        0.7914915680885315,
+        0.7937509417533875,
+        0.4119175970554352,
+        0.8763647675514221,
+        0.4558652639389038
+      ],
+      "min": [
+        -1.0,
+        -0.5158019661903381,
+        -0.7942351698875427,
+        -0.6549583673477173,
+        -0.7899840474128723,
+        -0.1914844959974289,
+        -0.8830769658088684,
+        -0.25145918130874634
+      ],
+      "q01": [
+        -0.7246493166685104,
+        -0.36842879801988604,
+        -0.4745602348446846,
+        -0.32548967510461807,
+        -0.6590064060688019,
+        -0.00425734419375658,
+        -0.667693349123001,
+        0.03753467075526715
+      ],
+      "q99": [
+        0.38912478089332564,
+        0.4704011854529379,
+        0.9895162123441696,
+        0.6802999329566952,
+        0.6482488948106765,
+        0.26210771769285157,
+        0.4697745823860167,
+        0.32386209577322006
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        true
+      ]
+    },
+    "state": {
+      "mean": [
+        -0.03927938640117645,
+        0.873625636100769,
+        -0.2643051743507385,
+        -0.5193795561790466,
+        -0.03129394352436066,
+        0.37781062722206116,
+        0.3676287829875946,
+        -0.12228213250637054,
+        0.10376925021409988,
+        0.0348481647670269,
+        0.10585025697946548,
+        -0.3591102361679077,
+        0.4178043007850647,
+        0.3299199938774109,
+        -0.09238386154174805,
+        -0.07736567407846451,
+        0.08624251186847687,
+        -0.0959310308098793,
+        0.05461423099040985,
+        -0.14337068796157837,
+        -0.10180454701185226,
+        0.04959503188729286,
+        0.18128058314323425,
+        -0.015483468770980835,
+        0.28967317938804626,
+        0.2775602340698242,
+        -0.33862781524658203
+      ],
+      "std": [
+        0.24292634427547455,
+        0.06530473381280856,
+        0.13237309455871582,
+        0.13348814845085144,
+        0.15209674835205078,
+        0.4691879153251648,
+        0.39366209506988525,
+        0.6651913523674011,
+        0.08553778380155563,
+        0.6054338812828064,
+        0.08797592669725418,
+        0.48128432035446167,
+        0.33567723631858826,
+        0.2170608788728714,
+        0.24111506342887878,
+        0.35090306401252747,
+        0.2857387363910675,
+        0.20407065749168396,
+        0.22814442217350006,
+        0.3093811273574829,
+        0.3431510031223297,
+        0.41457003355026245,
+        0.05459607020020485,
+        0.33612897992134094,
+        0.05801063030958176,
+        0.34881407022476196,
+        0.31408628821372986
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0
+      ],
+      "q01": [
+        -0.518711000084877,
+        0.7832470858097076,
+        -0.48747050642967227,
+        -0.8287142354249954,
+        -0.3970690089464188,
+        -0.7996528524160386,
+        -0.18211212992668152,
+        -0.9498927110433578,
+        0.05799898505210877,
+        -0.9213203591108322,
+        0.06377534270286564,
+        -0.9172824740409851,
+        -0.14945583879947663,
+        -0.47512653708457947,
+        -0.5667413014173508,
+        -0.720048052072525,
+        -0.6468206322193146,
+        -0.4755307936668396,
+        -0.4702010858058929,
+        -0.6532424187660217,
+        -0.637979581952095,
+        -0.8247658777236938,
+        0.17336984157562255,
+        -0.7672037160396576,
+        0.28981017947196963,
+        -0.6314861965179444,
+        -0.8276125568151474
+      ],
+      "q99": [
+        0.5284208941459655,
+        0.8995754766464233,
+        0.07895282626152018,
+        -0.29791240513324774,
+        0.3613759696483611,
+        0.9232616770267487,
+        0.9520751488208771,
+        0.944451036453247,
+        0.11151877999305723,
+        0.8851249778270721,
+        0.11331057786941529,
+        0.7722350835800162,
+        0.9076120913028717,
+        0.7570547771453847,
+        0.543445667028426,
+        0.718429600000381,
+        0.6114705646038047,
+        0.3634060764312739,
+        0.6103977251052853,
+        0.6279066145420074,
+        0.8004330182075499,
+        0.8364576077461242,
+        0.1844600486755371,
+        0.5733257675170897,
+        0.2910482907295227,
+        0.8056755244731904,
+        0.2977444970607755
+      ]
+    },
+    "num_transitions": 90000,
+    "num_trajectories": 90
+  }
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/dataset_statistics_eval.json

@@ -0,0 +1,254 @@
+{
+  "new_embodiment": {
+    "action": {
+      "mean": [
+        -0.194865420460701,
+        0.045133840292692184,
+        0.10328157991170883,
+        0.06766745448112488,
+        -0.07633885741233826,
+        0.10476952791213989,
+        -0.02807306870818138,
+        0.17962458729743958
+      ],
+      "std": [
+        0.2925890386104584,
+        0.179392471909523,
+        0.33783403038978577,
+        0.24546940624713898,
+        0.33822667598724365,
+        0.052730318158864975,
+        0.28934720158576965,
+        0.06433656811714172
+      ],
+      "max": [
+        0.5322136878967285,
+        0.555697500705719,
+        1.0,
+        0.7869182229042053,
+        0.8116569519042969,
+        0.3534770607948303,
+        0.7703150510787964,
+        0.440703809261322
+      ],
+      "min": [
+        -1.0,
+        -0.47031596302986145,
+        -0.730529248714447,
+        -0.6560564637184143,
+        -0.8157420754432678,
+        -0.052262552082538605,
+        -0.8346825242042542,
+        -0.025107912719249725
+      ],
+      "q01": [
+        -0.7260830360651016,
+        -0.36782372266054153,
+        -0.47317668825387954,
+        -0.31557414084672925,
+        -0.66000783264637,
+        0.0017754910979419962,
+        -0.6687433636188507,
+        0.044861558452248575
+      ],
+      "q99": [
+        0.38754730314016345,
+        0.46464609265327456,
+        1.0,
+        0.6821840578317643,
+        0.6534658980369571,
+        0.2567682746052742,
+        0.466521698832512,
+        0.3254868069291115
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        true
+      ]
+    },
+    "state": {
+      "mean": [
+        -0.03446132689714432,
+        0.8770178556442261,
+        -0.26813629269599915,
+        -0.5243244767189026,
+        -0.036261554807424545,
+        0.38102391362190247,
+        0.368524432182312,
+        -0.12297961115837097,
+        0.10743273794651031,
+        0.03137503191828728,
+        0.10943524539470673,
+        -0.3544634282588959,
+        0.4190467298030853,
+        0.3441190719604492,
+        -0.09693243354558945,
+        -0.07695353031158447,
+        0.08566370606422424,
+        -0.09651442617177963,
+        0.054840974509716034,
+        -0.14340537786483765,
+        -0.1022859439253807,
+        0.04956989362835884,
+        0.18135188519954681,
+        -0.015385560691356659,
+        0.2896140515804291,
+        0.2777324914932251,
+        -0.3387359380722046
+      ],
+      "std": [
+        0.23964886367321014,
+        0.023951832205056853,
+        0.11455036699771881,
+        0.1093541756272316,
+        0.14902757108211517,
+        0.46400782465934753,
+        0.39322152733802795,
+        0.666404664516449,
+        0.06292039155960083,
+        0.6054039001464844,
+        0.06391128897666931,
+        0.48054811358451843,
+        0.33541467785835266,
+        0.20528468489646912,
+        0.23731346428394318,
+        0.3526037931442261,
+        0.28841397166252136,
+        0.2061758041381836,
+        0.22504228353500366,
+        0.3090031147003174,
+        0.3419175148010254,
+        0.4147181808948517,
+        0.054758649319410324,
+        0.3359338939189911,
+        0.05688738450407982,
+        0.35167375206947327,
+        0.3148670196533203
+      ],
+      "max": [
+        0.8509142398834229,
+        0.946937084197998,
+        0.28078794479370117,
+        0.13082659244537354,
+        0.7544276714324951,
+        0.9886277914047241,
+        0.9866607189178467,
+        0.9863904714584351,
+        0.9886200428009033,
+        0.9180272817611694,
+        0.958978533744812,
+        0.976783037185669,
+        0.994597315788269,
+        0.9207459688186646,
+        0.9506415128707886,
+        0.9698106050491333,
+        0.8268996477127075,
+        0.5188288688659668,
+        0.899605393409729,
+        0.8091031312942505,
+        0.9923845529556274,
+        0.9708276987075806,
+        0.9690662622451782,
+        0.8003939390182495,
+        1.0,
+        0.9442052841186523,
+        0.9767690896987915
+      ],
+      "min": [
+        -0.6200818419456482,
+        0.5183085203170776,
+        -0.7695543766021729,
+        -0.9485529661178589,
+        -0.8429219722747803,
+        -0.9506317377090454,
+        -0.9622589349746704,
+        -1.0,
+        -0.9714068174362183,
+        -1.0,
+        -0.9771743416786194,
+        -0.9749248623847961,
+        -1.0,
+        -0.7751631736755371,
+        -0.8879467844963074,
+        -1.0,
+        -1.0,
+        -0.7477306723594666,
+        -0.7433239221572876,
+        -0.8804571628570557,
+        -0.9065881967544556,
+        -0.9710677862167358,
+        -0.9849591851234436,
+        -0.9246947765350342,
+        -0.9404064416885376,
+        -0.956570029258728,
+        -1.0
+      ],
+      "q01": [
+        -0.5038023978471756,
+        0.7942106568813324,
+        -0.478440306186676,
+        -0.8267061966657638,
+        -0.39669507682323457,
+        -0.7862418442964554,
+        -0.1804919111728668,
+        -0.9492776465415954,
+        0.10735819339752198,
+        -0.919443871974945,
+        0.11226175785064697,
+        -0.9184101295471191,
+        -0.1500501948595047,
+        -0.34228765964508057,
+        -0.5650507682561874,
+        -0.7155984926223755,
+        -0.6520336997509003,
+        -0.4805710732936859,
+        -0.46580726146698,
+        -0.6449919211864471,
+        -0.6242685246467591,
+        -0.819073212146759,
+        0.17577075600624084,
+        -0.7636688566207885,
+        0.28973553419113157,
+        -0.6239250099658966,
+        -0.8301586848497391
+      ],
+      "q99": [
+        0.531941732168198,
+        0.8995999598503113,
+        0.04719418048858646,
+        -0.3141447216272354,
+        0.33111580729484597,
+        0.9231463372707368,
+        0.9524696123600006,
+        0.9444611060619355,
+        0.11135421514511108,
+        0.8837961864471436,
+        0.11332787156105041,
+        0.7487831759452822,
+        0.9076400446891785,
+        0.7634504711627961,
+        0.5362399804592136,
+        0.7234879946708682,
+        0.6047599220275879,
+        0.36570024728775025,
+        0.6087638485431678,
+        0.6343572854995728,
+        0.8117764890193939,
+        0.836361961364746,
+        0.18426278471946717,
+        0.5724071526527406,
+        0.2910631883144379,
+        0.7999724817276002,
+        0.29399525642395047
+      ]
+    },
+    "num_transitions": 10000,
+    "num_trajectories": 10
+  }
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/files.sha256.json

@@ -0,0 +1,110 @@
+{
+  "README.md": {
+    "bytes": 1131,
+    "sha256": "b8d91b05b9719d73549c3e71fb89b846804e31e9870616d346875796a961a975"
+  },
+  "checkpoints/model.pt": {
+    "bytes": 9785070315,
+    "sha256": "afb954569065450b9b80a86ca96e3898e3fe8d2f8c38a06f63bf3b11c41163b3"
+  },
+  "collection_quality_report.json": {
+    "bytes": 421,
+    "sha256": "3a21eb8ce207f265fee8197c317c6f93013e10dd96dc32f5a702b1b658773de5"
+  },
+  "config.full.yaml": {
+    "bytes": 11309,
+    "sha256": "936acd2c14d8aa4c4ab817eab4acf046827b879f12b6454e4a1ae7e1ef69136c"
+  },
+  "config.yaml": {
+    "bytes": 965,
+    "sha256": "dfda64a8e3d18e6ceb8c0fdc686915cd897f0df78ea24b100b58a468335ee7c6"
+  },
+  "dataset_statistics.json": {
+    "bytes": 5727,
+    "sha256": "8808223c11ba20721109750a91a23cc2c81b5b92d8e48320edfe5f22722b3389"
+  },
+  "dataset_statistics_eval.json": {
+    "bytes": 6413,
+    "sha256": "85ee7976398bfc0cf7f75e6f3c9528f9fa872a77dcab0bb849dc3920683e973a"
+  },
+  "latency_prompt_map.json": {
+    "bytes": 497,
+    "sha256": "8be52633a13e9032392595f2e6d637a98f539dacb6757e1dec59e65ffc2074e3"
+  },
+  "launch_audit.json": {
+    "bytes": 1786,
+    "sha256": "262ce08821f1c2261b8567eed44c53818fba6d75cfe967c35e61fc2ddee42070"
+  },
+  "manifest.json": {
+    "bytes": 6710,
+    "sha256": "ec82945da958d7e29d172928b136177d2b0e23612b62889f64b475446cd655c4"
+  },
+  "profile.source.json": {
+    "bytes": 2148,
+    "sha256": "0e49f72ec05ac600a54e07bbca5af3143708444d606167f33fa21788a9fefe50"
+  },
+  "provenance.json": {
+    "bytes": 6152,
+    "sha256": "ddb044cd3a9ea15f9c07c2370df31f15c4ffda0236fee45b0935a99aba79720d"
+  },
+  "startup_verified.json": {
+    "bytes": 1553,
+    "sha256": "f7ade30c273f14972b8deacf42b9f0ff2bad29bb0cefffc1b131977a9933d3b5"
+  },
+  "summary.jsonl": {
+    "bytes": 159,
+    "sha256": "2905acb79c5fffb33fd4689d87af3bd05d74f4184b9d36ff3eae0c64deecd930"
+  },
+  "task_contract.json": {
+    "bytes": 2298,
+    "sha256": "a874a91173a6579750761e4aa65d8daecc167dc9a2d0c58041813f83d589c959"
+  },
+  "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": 1762,
+    "sha256": "8c3b350880c1fd9996fa4a77a9c81c8e9a4c3d0fd0c96ef5272dc88669bbba63"
+  }
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/latency_prompt_map.json

@@ -0,0 +1,7 @@
+{
+  "1": {
+    "prompt": "Make the Ant move forward as fast as possible without falling. Predict eight continuous torques in [-1, 1] ordered as back right hip, back right ankle, front left hip, front left ankle, front right hip, front right ankle, back left hip, and back left ankle. Current action latency is 1 raw frames (100.00 ms).  The environment runs at 10 FPS and observations are emitted at 10 FPS. Choose the best next action.",
+    "latency_raw_frames": 1,
+    "latency_ms": 100.0
+  }
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/launch_audit.json

@@ -0,0 +1,41 @@
+{
+  "reference": "/home/ubuntu/lzj/mean-profiling/reference/checkpoint-configs/latency-aware/ant/vla/starvla-qwenoft-h1/ant_rgb_state_profile_20260911T112232Z_openvla_native_sft_5k/config.full.yaml",
+  "reference_sha256": "1abb758e750efa660e3a4e949a974eedf19e6d8e49cc5ef98ec29c6da453c56e",
+  "scientific_trainer_fields_equal": [
+    "max_train_steps",
+    "num_warmup_steps",
+    "learning_rate",
+    "lr_scheduler_type",
+    "scheduler_specific_kwargs",
+    "optimizer",
+    "eval_interval",
+    "eval_num_batches",
+    "save_interval",
+    "loss_scale"
+  ],
+  "seed": 42,
+  "horizon": 1,
+  "state_dim": 27,
+  "action_dim": 8,
+  "global_batch": 128,
+  "micro_per_gpu": 64,
+  "accumulation": 1,
+  "gradient_checkpointing": true,
+  "formal_optimizer_updates": 5000,
+  "preflight_updates_discarded": 20,
+  "post_train_rollout_eval": false,
+  "data_revision": "6246ccc57caf32da0e02547c3e3e1080894d4035",
+  "runtime_save_training_state": true,
+  "runtime_keep_last_recovery_states": 2,
+  "profile_timing_same_as_reference": false,
+  "reference_global_batch": 16,
+  "profile_timing_enabled": false,
+  "user_authorized_scientific_difference": "Prioritize two GPU utilization; user explicitly requests Flappy micro64*2/accum1/global128, referenceglobal16 no longerbinding",
+  "failed_initial_preflight": "Missing requiredwandb_tags, noformaloptimizerupdates",
+  "metadata_repair": "Explicitwandb_tags added",
+  "optimizer_budget": 5000,
+  "previous_runtime_failure": "32/accum2 CUDA OOM atsecond optimizer update after1testupdate",
+  "runtime_adjustment": "Enableexistingnon-reentrantgradientcheckpointing;64*2*accum1/global128, othermath/scientificbudgets maintained",
+  "startup_verified": true,
+  "formal_startup_evidence": "/home/ubuntu/lzj/mean-profiling/ant/vla.startup-verified.json"
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/manifest.json

@@ -0,0 +1,232 @@
+{
+  "dataset_name": "lerobot",
+  "env_name": "ant_rgb_state",
+  "episodes": 90,
+  "frames": 90000,
+  "task_prompts": [
+    "Make the Ant move forward as fast as possible without falling. Predict eight continuous torques in [-1, 1] ordered as back right hip, back right ankle, front left hip, front left ankle, front right hip, front right ankle, back left hip, and back left ankle. Current action latency is 1 raw frames (100.00 ms).  The environment runs at 10 FPS and observations are emitted at 10 FPS. Choose the best next action."
+  ],
+  "format": "lerobot_v2_image_parquet",
+  "source": "/home/ubuntu/lzj/mean-profiling/ant/data/raw",
+  "integration_name": "gymnasium",
+  "task_name": "ant_rgb_state",
+  "action_layout": "gymnasium_continuous_v1",
+  "action_labels": [
+    "back_right_hip_torque",
+    "back_right_ankle_torque",
+    "front_left_hip_torque",
+    "front_left_ankle_torque",
+    "front_right_hip_torque",
+    "front_right_ankle_torque",
+    "back_left_hip_torque",
+    "back_left_ankle_torque"
+  ],
+  "carrier_action_labels": [
+    "back_right_hip_torque",
+    "back_right_ankle_torque",
+    "front_left_hip_torque",
+    "front_left_ankle_torque",
+    "front_right_hip_torque",
+    "front_right_ankle_torque",
+    "back_left_hip_torque",
+    "back_left_ankle_torque"
+  ],
+  "action_dim": 8,
+  "active_action_dim": 8,
+  "action_carrier": "native",
+  "rows_unit": "decision_step",
+  "obs_fps": 10.0,
+  "obs_stride_raw_frames": 1,
+  "uses_state": true,
+  "state_dim": 27,
+  "state_labels": [
+    "torso_z",
+    "torso_quaternion_w",
+    "torso_quaternion_x",
+    "torso_quaternion_y",
+    "torso_quaternion_z",
+    "front_left_hip_angle",
+    "front_left_ankle_angle",
+    "front_right_hip_angle",
+    "front_right_ankle_angle",
+    "back_left_hip_angle",
+    "back_left_ankle_angle",
+    "back_right_hip_angle",
+    "back_right_ankle_angle",
+    "torso_x_velocity",
+    "torso_y_velocity",
+    "torso_z_velocity",
+    "torso_angular_velocity_x",
+    "torso_angular_velocity_y",
+    "torso_angular_velocity_z",
+    "front_left_hip_angular_velocity",
+    "front_left_ankle_angular_velocity",
+    "front_right_hip_angular_velocity",
+    "front_right_ankle_angular_velocity",
+    "back_left_hip_angular_velocity",
+    "back_left_ankle_angular_velocity",
+    "back_right_hip_angular_velocity",
+    "back_right_ankle_angular_velocity"
+  ],
+  "state_normalization": {
+    "type": "min_max",
+    "min": [
+      0.24513374269008636,
+      -0.8805694580078125,
+      -0.33691075444221497,
+      -0.4456058740615845,
+      -0.9112266302108765,
+      -0.6140146255493164,
+      -0.09917011857032776,
+      -0.6453427076339722,
+      -1.3009648323059082,
+      -0.6563361883163452,
+      -1.2990498542785645,
+      -0.6459031105041504,
+      -0.09470116347074509,
+      -1.9981460571289062,
+      -2.8870370388031006,
+      -2.5440866947174072,
+      -5.5593695640563965,
+      -6.843714237213135,
+      -5.608675003051758,
+      -10.695633888244629,
+      -9.341012001037598,
+      -13.44798469543457,
+      -11.152595520019531,
+      -14.51861572265625,
+      -11.078145980834961,
+      -12.814929008483887,
+      -5.827189922332764
+    ],
+    "max": [
+      0.9843788743019104,
+      1.098469853401184,
+      0.557304859161377,
+      0.9673683047294617,
+      0.7864838242530823,
+      0.6336118578910828,
+      1.329100489616394,
+      0.624431848526001,
+      0.09998779743909836,
+      0.6632993221282959,
+      0.09576962888240814,
+      0.5784072875976562,
+      1.320813536643982,
+      5.847742557525635,
+      3.869502544403076,
+      2.964057445526123,
+      4.749230861663818,
+      8.3289155960083,
+      4.963846206665039,
+      14.281911849975586,
+      11.46115779876709,
+      12.22265911102295,
+      7.717322826385498,
+      15.017043113708496,
+      6.09257173538208,
+      7.276201248168945,
+      11.832711219787598
+    ]
+  },
+  "robot_type": "rl_games_gymnasium",
+  "latency_prompt_map_path": "/home/ubuntu/lzj/mean-profiling/ant/data/lerobot/latency_prompt_map.json",
+  "custom_mixtures_path": "/home/ubuntu/lzj/mean-profiling/ant/data/_generated_mixtures/lerobot.json",
+  "gymnasium_task": {
+    "action_space": {
+      "dtype": "float32",
+      "high": [
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ],
+      "labels": [
+        "back_right_hip_torque",
+        "back_right_ankle_torque",
+        "front_left_hip_torque",
+        "front_left_ankle_torque",
+        "front_right_hip_torque",
+        "front_right_ankle_torque",
+        "back_left_hip_torque",
+        "back_left_ankle_torque"
+      ],
+      "low": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0
+      ],
+      "type": "box"
+    },
+    "base_prompt": "Make the Ant move forward as fast as possible without falling. Predict eight continuous torques in [-1, 1] ordered as back right hip, back right ankle, front left hip, front left ankle, front right hip, front right ankle, back left hip, and back left ankle.",
+    "env_fps": 10.0,
+    "env_id": "LatencyBench/AntContinuous-v0",
+    "frame_stack": 1,
+    "make_kwargs": {
+      "base_env_id": "Ant-v4",
+      "base_make_kwargs": {
+        "exclude_current_positions_from_observation": true,
+        "use_contact_forces": false
+      },
+      "render_mode": "rgb_array"
+    },
+    "noop_action": [
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0
+    ],
+    "obs_fps": 10.0,
+    "registration_imports": [
+      "latency_bench.envs.gymnasium_ant"
+    ],
+    "state_space": {
+      "labels": [
+        "torso_z",
+        "torso_quaternion_w",
+        "torso_quaternion_x",
+        "torso_quaternion_y",
+        "torso_quaternion_z",
+        "front_left_hip_angle",
+        "front_left_ankle_angle",
+        "front_right_hip_angle",
+        "front_right_ankle_angle",
+        "back_left_hip_angle",
+        "back_left_ankle_angle",
+        "back_right_hip_angle",
+        "back_right_ankle_angle",
+        "torso_x_velocity",
+        "torso_y_velocity",
+        "torso_z_velocity",
+        "torso_angular_velocity_x",
+        "torso_angular_velocity_y",
+        "torso_angular_velocity_z",
+        "front_left_hip_angular_velocity",
+        "front_left_ankle_angular_velocity",
+        "front_right_hip_angular_velocity",
+        "front_right_ankle_angular_velocity",
+        "back_left_hip_angular_velocity",
+        "back_left_ankle_angular_velocity",
+        "back_right_hip_angular_velocity",
+        "back_right_ankle_angular_velocity"
+      ]
+    },
+    "task_name": "ant_rgb_state"
+  },
+  "validation_dataset_name": "lerobot__val",
+  "validation_episodes": 10,
+  "validation_frames": 10000
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/profile.source.json

@@ -0,0 +1,66 @@
+{
+  "burst_model_path": "latency_burst_model.json",
+  "distribution_path": "latency_distribution.json",
+  "env_fps": 10,
+  "frame_ms": 100.0,
+  "gpu_class": "1x-rtx3090",
+  "instance_id": "instance_859cf1e47bca6046",
+  "latency_kind": "observation_to_action_latency",
+  "latency_method": "temporal",
+  "model_id": "qwenoft",
+  "n_admitted_observations": 7976,
+  "n_capacity_drops": 294,
+  "n_observation_attempts": 8270,
+  "per_slot_summary": {
+    "0": {
+      "admitted_count": 7976,
+      "mean_observation_to_action_latency_ms": 90.56460638563993,
+      "mean_worker_service_time_ms": 90.22241150451042,
+      "p95_observation_to_action_latency_ms": 92.13036412000656,
+      "p95_worker_service_time_ms": 91.75596672296524,
+      "p99_worker_service_time_ms": 110.92562991380692
+    }
+  },
+  "provenance": {
+    "base_config": "/workspace/tasks/20260911T023128Z-p-only4/ant/profile.yaml",
+    "checkpoint_kind": "best",
+    "model_artifact": {
+      "checkpoint": "checkpoints/steps_5000_pytorch_model.pt",
+      "model_config": "config.full.yaml",
+      "path_in_repo": "OpenVLA/zero-latency/ant_rgb_state_l0_return_gt3000_100ep_openvla_native_sft_5k",
+      "repo_id": "latency-sensitive-bench/extra-envs-checkpoints",
+      "source": "local"
+    },
+    "session_ids": [
+      0,
+      1,
+      2,
+      3,
+      4
+    ]
+  },
+  "sample_model_type": "hidden_regime",
+  "source_run_id": "20260911T033037730561Z",
+  "summary": {
+    "frame_ms": 100.0,
+    "max_ms": 130.29119229316711,
+    "mean_effective_frames": 0.9056460638563993,
+    "mean_ms": 90.56460638563993,
+    "min_ms": 87.98952293395996,
+    "n_samples": 7976,
+    "p50_frames": 0.8970676946640015,
+    "p50_ms": 89.70676946640015,
+    "p90_frames": 0.9074710392951966,
+    "p90_ms": 90.74710392951965,
+    "p95_frames": 0.9213036412000656,
+    "p95_ms": 92.13036412000656,
+    "p99_frames": 1.1217999053001404,
+    "p99_ms": 112.17999053001404,
+    "prob_latency_gt_1_frame": 0.03686058174523571,
+    "prob_latency_gt_2_frames": 0.0,
+    "prob_latency_gt_3_frames": 0.0,
+    "std_ms": 4.164609685850939
+  },
+  "visualization_path": "latency_profile.png",
+  "workload_id": "ant"
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/provenance.json

@@ -0,0 +1,147 @@
+{
+  "task": "ant",
+  "model": "QwenOFT",
+  "run_id": "ant-qwenoft-mean-3090-s42",
+  "training_status": "complete",
+  "training_hardware": {
+    "gpu": "RTX PRO6000 Blackwell Server Edition",
+    "count": 2,
+    "physical_indices": [
+      0,
+      1
+    ]
+  },
+  "profile": {
+    "mean_ms": 90.56460638563993,
+    "profile": "/home/ubuntu/lzj/mean-profiling/reference/checkpoint-configs/latency-aware/ant/small-policy/sample-factory-qwenoft-rtx3090-v1/profile/profile.json",
+    "sha256": "0e49f72ec05ac600a54e07bbca5af3143708444d606167f33fa21788a9fefe50"
+  },
+  "profile_hardware": "RTX3090",
+  "condition": "fixed_mean_profile",
+  "horizon": 1,
+  "seed": 42,
+  "state_dim": 27,
+  "action_dim": 8,
+  "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": "1c2f38f16553622eef2c62ddde86036be989ea31",
+    "path": "latency-aware/ant/small-policy/sample-factory-qwenoft-rtx3090-mean-v1",
+    "local_checkpoint": "/home/ubuntu/lzj/mean-profiling/ant/teacher-publication/checkpoint.pth",
+    "status": "complete_verified",
+    "completed_env_steps": 10006528,
+    "selected_env_steps": 9969664,
+    "selected_source_checkpoint": "/home/ubuntu/lzj/mean-profiling/ant/teacher/ant_mean_3090_s2222/checkpoint_p0/best_000019472_9969664_reward_3371.900.pth",
+    "tensor_count": 15,
+    "weights_only_export": true,
+    "tensor_equality_verified": true,
+    "hash_readback_verified": true,
+    "url": "https://huggingface.co/latency-sensitive-bench/benchmark-models/tree/1c2f38f16553622eef2c62ddde86036be989ea31/latency-aware/ant/small-policy/sample-factory-qwenoft-rtx3090-mean-v1",
+    "completed_utc": "2026-09-30T13:54:03.497277+00:00"
+  },
+  "dataset": {
+    "repo_id": "latency-sensitive-bench/Standard-Pipeline",
+    "repo_type": "dataset",
+    "path": "ant/profile_latency/OpenVLA/mean_rtx3090_s2222_100ep_v1",
+    "revision": "6246ccc57caf32da0e02547c3e3e1080894d4035",
+    "data_uploaded_revision": "4faa069b31f461177e1adbb0c039da5db0adacb8",
+    "url": "https://huggingface.co/datasets/latency-sensitive-bench/Standard-Pipeline/tree/6246ccc57caf32da0e02547c3e3e1080894d4035/ant/profile_latency/OpenVLA/mean_rtx3090_s2222_100ep_v1",
+    "status": "complete_verified",
+    "episodes": 100,
+    "rows": 100000,
+    "strict_return_gt2500": true,
+    "payload_files_verified": 26,
+    "checksum_manifest_verified": true,
+    "dataset_config": "ant_profile_latency_mean_qwenoft",
+    "completed_utc": "2026-09-30T14:50:21.716749+00:00"
+  },
+  "collection": {
+    "task": "ant",
+    "teacher_checkpoint": "/home/ubuntu/lzj/mean-profiling/ant/teacher-publication/checkpoint.pth",
+    "teacher_sha256": "a7899b32beb9b2cfec1bcff85b285dddc243a73f1c322fafefae599df0f67b3b",
+    "episodes": 100,
+    "train_episodes": 90,
+    "val_episodes": 10,
+    "train_rows": 90000,
+    "val_rows": 10000,
+    "fixed_latency_ms": 90.56460638563993,
+    "action_horizon": 1,
+    "unique_seeds": 100,
+    "state_dim": 27
+  },
+  "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.011952828615903854,
+    "validation_action_loss": 0.018865827471017838,
+    "validation_samples": 10112,
+    "final_learning_rates": {
+      "learning_rate/action_model": 4.9999999999999996e-06,
+      "learning_rate/qwen_vl_interface": 5e-07
+    },
+    "closed_loop_evaluation": "pending until all four VLA trainings complete",
+    "tensor_keys": 730,
+    "serialized_tensor_elements_including_tied_aliases": 4892400136,
+    "weights_only_load": "passed",
+    "all_tensors_finite": true,
+    "completion_proof": {
+      "task": "ant",
+      "optimizer_updates": 5000,
+      "deepspeed_global_steps": 5000,
+      "deepspeed_global_samples": 640000,
+      "skipped_steps": 0,
+      "scheduler_last_epoch": 5000,
+      "scheduler_step_count": 5001,
+      "final_learning_rates": [
+        5e-07,
+        4.9999999999999996e-06
+      ],
+      "summary_final_steps": 5000,
+      "final_inference_weight_bytes": 9785070315,
+      "wandb_global_step": 5000,
+      "wrapper_exit": 2,
+      "trainer_exit_separately_captured": false,
+      "training_complete_verified": true,
+      "failure_scope": "launcher_tail_parse_after_trainer_cleanup",
+      "cause": "Sharedrun_vla.shwasmodifiedforGPUmask/portwhileAntbashwasblockedinsidePython; remainingfileoffsetlaterparsedmodifiedtailandEOF",
+      "trainer_cleanup_log": "bothranksDestroyDistributedProcessGroup at2026-10-01 00:03:21UTC; no model exception",
+      "evidence_files": [
+        "checkpoints/steps_5000_state/pytorch_model/mp_rank_00_model_states.pt",
+        "checkpoints/steps_5000_state/custom_checkpoint_0.pkl",
+        "summary.jsonl",
+        "config.full.yaml"
+      ],
+      "note": "No syntheticvla.exit0; wrapperexit2preserved; upload uses explicitcompletionproof"
+    }
+  },
+  "optimizer_recovery_states": "retained on host; inference weights only in this bundle",
+  "wandb_url": "https://wandb.ai/dongqianyu99-zhejiang-university/latency-sensitive-bench/runs/mean3090-ant-vla-blackwell-20260930"
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/startup_verified.json

@@ -0,0 +1,67 @@
+{
+  "snapshot_utc": "2026-09-30T17:38:42Z",
+  "unit": "mean-profile-ant-vla-20261001-b64gc.service",
+  "wrapper_pid": 675875,
+  "rank_pids": [
+    676262,
+    676263
+  ],
+  "physical_gpus": [
+    0,
+    1
+  ],
+  "formal_training_started": true,
+  "formal_optimizer_updates_completed": 36,
+  "formal_optimizer_budget": 5000,
+  "preflight_optimizer_updates_discarded": 20,
+  "consecutive_logged_steps": [
+    33,
+    34,
+    35,
+    36
+  ],
+  "loss_at_steps": {
+    "33": 0.3258245885372162,
+    "34": 0.3107373118400574,
+    "35": 0.49765458703041077,
+    "36": 0.4269329905509949
+  },
+  "finite_loss_verified": true,
+  "gpu_utilization_percent": [
+    100,
+    100
+  ],
+  "gpu_used_mib": [
+    63815,
+    65995
+  ],
+  "micro_batch_per_gpu": 64,
+  "gradient_accumulation": 1,
+  "global_batch": 128,
+  "gradient_checkpointing": true,
+  "seconds_per_optimizer_update_recent": 3.7,
+  "data_wait_seconds_recent": 0.0006755599970347248,
+  "host_memory_peak_bytes": 38237016064,
+  "host_memory_current_bytes": 17592913920,
+  "cgroup_memory_events": {
+    "low": 0,
+    "high": 0,
+    "max": 0,
+    "oom": 0,
+    "oom_kill": 0,
+    "oom_group_kill": 0
+  },
+  "remaining_compute_hours": 5.1018888888888885,
+  "estimated_remaining_hours_including_offline_validation_and_saves": [
+    5.5,
+    6.0
+  ],
+  "estimated_completion_utc_range": [
+    "2026-09-30T23:10:00Z",
+    "2026-09-30T23:40:00Z"
+  ],
+  "next_completion_check_utc": "2026-10-01T00:00:00Z",
+  "training_complete": false,
+  "closed_loop_evaluation_complete": false,
+  "publication_complete": false
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/summary.jsonl

@@ -0,0 +1,10 @@
+{"steps": 500}
+{"steps": 1000}
+{"steps": 1500}
+{"steps": 2000}
+{"steps": 2500}
+{"steps": 3000}
+{"steps": 3500}
+{"steps": 4000}
+{"steps": 4500}
+{"steps": 5000}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/task_contract.json

@@ -0,0 +1,94 @@
+{
+  "action_space": {
+    "dtype": "float32",
+    "high": [
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0
+    ],
+    "labels": [
+      "back_right_hip_torque",
+      "back_right_ankle_torque",
+      "front_left_hip_torque",
+      "front_left_ankle_torque",
+      "front_right_hip_torque",
+      "front_right_ankle_torque",
+      "back_left_hip_torque",
+      "back_left_ankle_torque"
+    ],
+    "low": [
+      -1.0,
+      -1.0,
+      -1.0,
+      -1.0,
+      -1.0,
+      -1.0,
+      -1.0,
+      -1.0
+    ],
+    "type": "box"
+  },
+  "base_prompt": "Make the Ant move forward as fast as possible without falling. Predict eight continuous torques in [-1, 1] ordered as back right hip, back right ankle, front left hip, front left ankle, front right hip, front right ankle, back left hip, and back left ankle.",
+  "env_fps": 10.0,
+  "env_id": "LatencyBench/AntContinuous-v0",
+  "frame_stack": 1,
+  "make_kwargs": {
+    "base_env_id": "Ant-v4",
+    "base_make_kwargs": {
+      "exclude_current_positions_from_observation": true,
+      "use_contact_forces": false
+    },
+    "render_mode": "rgb_array"
+  },
+  "noop_action": [
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0
+  ],
+  "obs_fps": 10.0,
+  "registration_imports": [
+    "latency_bench.envs.gymnasium_ant"
+  ],
+  "state_space": {
+    "labels": [
+      "torso_z",
+      "torso_quaternion_w",
+      "torso_quaternion_x",
+      "torso_quaternion_y",
+      "torso_quaternion_z",
+      "front_left_hip_angle",
+      "front_left_ankle_angle",
+      "front_right_hip_angle",
+      "front_right_ankle_angle",
+      "back_left_hip_angle",
+      "back_left_ankle_angle",
+      "back_right_hip_angle",
+      "back_right_ankle_angle",
+      "torso_x_velocity",
+      "torso_y_velocity",
+      "torso_z_velocity",
+      "torso_angular_velocity_x",
+      "torso_angular_velocity_y",
+      "torso_angular_velocity_z",
+      "front_left_hip_angular_velocity",
+      "front_left_ankle_angular_velocity",
+      "front_right_hip_angular_velocity",
+      "front_right_ankle_angular_velocity",
+      "back_left_hip_angular_velocity",
+      "back_left_ankle_angular_velocity",
+      "back_right_hip_angular_velocity",
+      "back_right_ankle_angular_velocity"
+    ]
+  },
+  "task_name": "ant_rgb_state"
+}

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/training-code/QWen3.py

@@ -0,0 +1,444 @@
+# Copyright 2025 starVLA community. All rights reserved.
+# Licensed under the MIT License, Version 1.0 (the "License");
+# Implemented by [Jinhui YE / HKUST University] in [2025].
+
+import time
+from collections import OrderedDict
+from functools import wraps
+from inspect import signature
+from typing import Optional
+
+import torch
+from starVLA.model.profiling import stage_timer as _stage
+from starVLA.training.trainer_utils import initialize_overwatch
+from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+logger = initialize_overwatch(__name__)
+
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = 151655
+VIDEO_TOKEN_INDEX = 151656
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_VIDEO_TOKEN = "<video>"
+
+_ACTION_TOKEN_MIN = 151669  # how can we know this range? check how you add fast tokens into VLM
+_ACTION_TOKEN_MAX = (
+    153716  # here only for fast_tokenizer, see starVLA/model/modules/vlm/tools/add_qwen_special_tokens/README.md
+)
+
+
+import torch.nn as nn
+
+
+def _patch_qwen3vl_flash_attention_position_ids() -> None:
+    from transformers.models.qwen3_vl import modeling_qwen3_vl as qwen3_vl
+
+    text_attention_cls = qwen3_vl.Qwen3VLTextAttention
+    if hasattr(text_attention_cls, "_starvla_flash_attention_position_ids_patched"):
+        return
+
+    original_forward = text_attention_cls.forward
+
+    def forward_without_flash_attention_position_ids(self, *args, **kwargs):
+        # Qwen3-VL already applies M-RoPE through position_embeddings before
+        # attention. Passing the repeated temporal M-RoPE ids into HF's FA2
+        # wrapper makes Transformers 4.57 mis-detect packed sequences and build
+        # empty cu_seqlens on KV-memory cache steps.
+        if self.config._attn_implementation == "flash_attention_2" and "position_ids" in kwargs:
+            del kwargs["position_ids"]
+        return original_forward(self, *args, **kwargs)
+
+    text_attention_cls.forward = forward_without_flash_attention_position_ids
+    text_attention_cls._starvla_flash_attention_position_ids_patched = True
+
+
+# The fixed Triton stage counts keep the PyTorch 2.11 CUDA 12.8 kernel
+# within the shared-memory budget on sm_120.
+_FLEX_KERNEL_OPTIONS = {
+    "BACKEND": "TRITON",
+    "fwd_num_stages": 2,
+    "bwd_num_stages": 1,
+}
+
+
+def _patch_qwen3vl_flex_attention_support() -> None:
+    """Allow loading Qwen3-VL with ``attn_implementation='flex_attention'``.
+
+    Qwen3-VL's attention already dispatches through the generic
+    ``ALL_ATTENTION_FUNCTIONS`` interface (which includes the ``flex_attention``
+    integration), but the model classes conservatively declare
+    ``_supports_flex_attn = False``, so Transformers 4.57 refuses to load with
+    flex at __init__. Flipping the flag lets the official dispatch path run; the
+    packed KV-memory training (`QwenOFT._forward_memory_packed`) passes a
+    BlockMask that FlexAttention consumes natively.
+
+    Also wraps the registered ``flex_attention`` interface to inject the fixed
+    Triton kernel options. Wrapping the dispatch entry — rather than threading
+    ``kernel_options`` through the top-level forward kwargs — guarantees both the
+    forward and the gradient-checkpointing backward recompile pick it up.
+    """
+    from transformers.models.qwen3_vl import modeling_qwen3_vl as qwen3_vl
+
+    for cls_name in (
+        "Qwen3VLForConditionalGeneration",
+        "Qwen3VLModel",
+        "Qwen3VLTextModel",
+        "Qwen3VLPreTrainedModel",
+    ):
+        cls = getattr(qwen3_vl, cls_name, None)
+        if cls is not None:
+            cls._supports_flex_attn = True
+
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    base_flex = ALL_ATTENTION_FUNCTIONS["flex_attention"]
+    if getattr(base_flex, "_starvla_kernel_options_patched", False):
+        return
+
+    def flex_attention_with_kernel_options(*args, **kwargs):
+        kwargs["kernel_options"] = _FLEX_KERNEL_OPTIONS
+        return base_flex(*args, **kwargs)
+
+    flex_attention_with_kernel_options._starvla_kernel_options_patched = True
+    ALL_ATTENTION_FUNCTIONS["flex_attention"] = flex_attention_with_kernel_options
+
+
+def _position_tensor_key(tensor):
+    if tensor is None:
+        return None
+    return (tensor.device, tensor.dtype, tuple(tensor.shape),
+            tuple(tensor.detach().cpu().reshape(-1).tolist()))
+
+
+def _cache_position_method(module, name, weight):
+    """Cache layout-only outputs on one native model instance, never image features."""
+    original = getattr(module, name)
+    parameters = signature(original)
+    cache = OrderedDict()
+
+    @wraps(original)
+    def cached(*args, **kwargs):
+        if module.training or torch.is_grad_enabled():
+            cache.clear()
+            return original(*args, **kwargs)
+        arguments = parameters.bind(*args, **kwargs)
+        arguments.apply_defaults()
+        parameter = weight()
+        key = (
+            parameter.device, parameter.dtype, parameter._version,
+            tuple(_position_tensor_key(value) for value in arguments.arguments.values()),
+        )
+        if key not in cache:
+            cache[key] = original(*args, **kwargs)
+            if len(cache) > 8:
+                cache.popitem(last=False)
+        cache.move_to_end(key)
+        return cache[key]
+
+    setattr(module, name, cached)
+    return cache
+
+
+class _QWen3_VL_Interface(nn.Module):
+    """
+    This exists because of the diversity of VLMs, so we encapsulate the changes here.
+    Lightweight wrapper around Qwen3-VL (Qwen3VLForConditionalGeneration).
+
+    Purpose:
+        - Unify interface with other VLM backends (CausalLM-like usage).
+        - Centralize preprocessing (tokenization + multimodal packing).
+        - Provide consistent forward / generate signatures.
+
+    """
+
+    def __init__(self, config: Optional[dict] = None, **kwargs):
+        """
+        Initialize the Qwen3-VL wrapper.
+        Following https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct
+
+        """
+        super().__init__()
+
+        qwenvl_config = config.framework.get("qwenvl", {})
+        model_id = qwenvl_config.get("base_vlm", "Qwen/Qwen3-VL-4B-Instruct")
+        attn_implementation = qwenvl_config.get("attn_implementation", "sdpa")
+        enable_grad_ckpt = qwenvl_config.get("enable_gradient_checkpointing", False)
+        print(
+            f"[QWen3] loading {model_id} with gradient_checkpointing={enable_grad_ckpt}",
+            flush=True,
+        )
+
+        if attn_implementation == "flash_attention_2":
+            import flash_attn_2_cuda  # noqa: F401
+
+            _patch_qwen3vl_flash_attention_position_ids()
+
+        if attn_implementation == "flex_attention":
+            _patch_qwen3vl_flex_attention_support()
+
+        model = Qwen3VLForConditionalGeneration.from_pretrained(
+            model_id,
+            attn_implementation=attn_implementation,
+            dtype=torch.bfloat16,
+            ignore_mismatched_sizes=True, # resize image no longer needed? @TODO check bug
+        )
+        processor = AutoProcessor.from_pretrained(model_id)
+        processor.tokenizer.padding_side = "left"
+
+        if enable_grad_ckpt:
+            model.config.text_config.use_cache = False
+            model.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs={"use_reentrant": False}
+            )
+            model.enable_input_require_grads()
+            print(
+                "[QWen3] gradient_checkpointing ENABLED "
+                f"(use_reentrant=False, active={model.is_gradient_checkpointing}, "
+                f"text_use_cache={model.config.text_config.use_cache})",
+                flush=True,
+            )
+
+        self.model = model
+        self.processor = processor
+        self.config = config
+        # align qwen3 with qwen2.5
+        self.model.config.hidden_size = self.model.config.text_config.hidden_size
+        self._last_build_timing = {}
+        visual = model.model.visual
+        self._position_caches = [
+            _cache_position_method(model.model, "get_rope_index", lambda: visual.rotary_pos_emb.inv_freq),
+            _cache_position_method(visual, "rot_pos_emb", lambda: visual.rotary_pos_emb.inv_freq),
+            _cache_position_method(visual, "fast_pos_embed_interpolate", lambda: visual.pos_embed.weight),
+        ]
+
+        # only for fast base model
+        if "-Action" in model_id:
+            self._ACTION_TOKEN_MIN = _ACTION_TOKEN_MIN
+            self._ACTION_TOKEN_MAX = _ACTION_TOKEN_MAX
+
+    def _clear_position_caches(self):
+        for cache in self._position_caches:
+            cache.clear()
+
+    def train(self, mode=True):
+        self._clear_position_caches()
+        return super().train(mode)
+
+    def _apply(self, fn, recurse=True):
+        self._clear_position_caches()
+        return super()._apply(fn, recurse=recurse)
+
+    def _load_from_state_dict(self, *args, **kwargs):
+        self._clear_position_caches()
+        return super()._load_from_state_dict(*args, **kwargs)
+
+    def _profile_timing_enabled(self) -> bool:
+        return self.config.trainer.profile_timing.enabled
+
+    def _profile_sync(self) -> None:
+        if self._profile_timing_enabled() and torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    def forward(
+        self,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        """
+        Forward pass delegating to underlying Qwen2.5-VL backbone.
+        """
+
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            outputs = self.model(
+                **kwargs,
+            )
+
+        return outputs
+
+    def forward_last_hidden(self, **kwargs) -> torch.Tensor:
+        """
+        Run the Qwen3-VL backbone and return only the final hidden state.
+
+        OFT-style action heads only consume the last hidden layer. Calling the
+        full conditional-generation wrapper with ``output_hidden_states=True``
+        returns every layer's activations, which can largely erase the memory
+        benefit of gradient checkpointing.
+        """
+
+        kwargs = dict(kwargs)
+        kwargs.pop("output_hidden_states", None)
+        kwargs.pop("output_attentions", None)
+        kwargs.setdefault("return_dict", True)
+        kwargs.setdefault("use_cache", False)
+
+        def _wrapper_forward_last_hidden() -> torch.Tensor:
+            lm_head = getattr(self.model, "lm_head", None)
+            captured = {}
+            handle = None
+            if lm_head is not None:
+                def _capture_lm_head_input(_module, inputs):
+                    captured["last_hidden"] = inputs[0]
+
+                handle = lm_head.register_forward_pre_hook(_capture_lm_head_input)
+
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                try:
+                    outputs = self.model(
+                        **kwargs,
+                        output_attentions=False,
+                        output_hidden_states=False,
+                        return_dict=True,
+                    )
+                finally:
+                    if handle is not None:
+                        handle.remove()
+
+            if "last_hidden" in captured:
+                return captured["last_hidden"]
+
+            # Last-resort compatibility path for unusual model wrappers that do
+            # not expose lm_head. This preserves correctness, but returns every
+            # layer and therefore saves less memory.
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                outputs = self.model(
+                    **kwargs,
+                    output_attentions=False,
+                    output_hidden_states=True,
+                    return_dict=True,
+                )
+            return outputs.hidden_states[-1]
+
+        backbone = getattr(self.model, "model", None)
+        if backbone is None:
+            return _wrapper_forward_last_hidden()
+
+        try:
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                outputs = backbone(**kwargs)
+        except TypeError:
+            # Some Transformers Qwen-VL wrappers inject visual embeddings before
+            # calling the text backbone, so the backbone itself may not accept
+            # pixel_values/image_grid_thw. Keep those versions working.
+            return _wrapper_forward_last_hidden()
+
+        if hasattr(outputs, "last_hidden_state"):
+            return outputs.last_hidden_state
+        return outputs[0]
+
+    def generate(
+        self,
+        **kwargs,
+    ):
+        """
+        High-level generation interface (auto-regressive decoding), optionally vision-conditioned.
+
+        Args:
+            **kwargs: fully follow raw model.generate() signature.
+        Returns:
+            GenerateOutput | Model-dependent generation return.
+        """
+        with torch.autocast("cuda", dtype=torch.float16):
+            generation_output = self.model.generate(
+                **kwargs,
+            )
+        return generation_output
+
+    def build_qwenvl_inputs(self, images, instructions, solutions=None, profiler=None, **kwargs):
+        """
+        Build model inputs from raw data (images + instructions + optional solutions).
+        Follow Oficial Qwen3-VL Instruct format: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct
+        """
+
+        # Create messages: one message per sample
+        with _stage(profiler, "starvla_qwen_message_build_ms"):
+            messages = []
+            assert len(images) == len(instructions), "Images and instructions must have the same length"
+            for imgs, instruction in zip(images, instructions):
+                content = [{"type": "image", "image": img} for img in imgs]
+
+                if "CoT_prompt" in self.config.datasets.vla_data:  # If using a grounding prompt to task
+                    CoT_prompt = self.config.datasets.vla_data.get("CoT_prompt", "")
+                    prompt = CoT_prompt.replace("{instruction}", instruction)
+                else:
+                    prompt = instruction
+
+                content.append({"type": "text", "text": prompt})
+                msg = [{"role": "user", "content": content}]
+
+                if solutions is not None:
+                    solution = solutions[len(messages)]
+                    msg.append({"role": "assistant", "content": [{"type": "text", "text": solution}]})
+                messages.append(msg)
+
+        profile_timing = self._profile_timing_enabled()
+        build_timing = {}
+
+        t_processor = time.perf_counter() if profile_timing else None
+        with _stage(profiler, "starvla_qwen_processor_ms"):
+            batch_inputs = self.processor.apply_chat_template(
+                messages, tokenize=True, padding=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
+            )
+        if profile_timing:
+            build_timing["timing/qwen_processor"] = time.perf_counter() - t_processor
+
+        # if solutions, mask out the solution tokens in labels
+        if solutions is not None:  #  here only for fast_tokenizer now.
+            action_token_min = _ACTION_TOKEN_MIN  # how can we know this range? --> we has other way for this, but is slower see qwenhelix branch
+            action_token_max = _ACTION_TOKEN_MAX  # here only for fast_tokenizer, see starVLA/model/modules/vlm/tools/add_qwen_special_tokens/README.md
+            labels = batch_inputs["input_ids"].clone()
+            # For each sequence in the batch, find the first occurrence of an action token.
+            for i in range(labels.size(0)):
+                seq = labels[i]
+                # Create a mask for tokens within the action token range.
+                mask_seq = (seq >= action_token_min) & (seq <= action_token_max)
+                nonzero_indices = torch.nonzero(mask_seq, as_tuple=False)
+                if nonzero_indices.numel() > 0:
+                    first_action_index = nonzero_indices[0].item()
+                    # Mask out all tokens before the first action token.
+                    seq[:first_action_index] = IGNORE_INDEX
+                else:
+                    # If no action token is found, mask the entire sequence.
+                    seq[:] = IGNORE_INDEX
+                    RuntimeWarning(
+                        "action token are on in yout tokenizer, plz see starVLA/model/modules/vlm/tools/add_qwen_special_tokens/README.md."
+                    )
+
+            labels[labels == self.processor.tokenizer.pad_token_id] = -100  ## mask out pad tokens as well
+            batch_inputs["labels"] = labels
+
+        t_h2d = time.perf_counter() if profile_timing else None
+        with _stage(profiler, "starvla_qwen_h2d_ms"):
+            batch_inputs = batch_inputs.to(self.model.device)
+        if profile_timing:
+            self._profile_sync()
+            build_timing["timing/qwen_h2d"] = time.perf_counter() - t_h2d
+        self._last_build_timing = build_timing
+        return batch_inputs
+
+
+if __name__ == "__main__":
+    import argparse
+    import os
+
+    from omegaconf import OmegaConf
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config_yaml",
+        type=str,
+        default="examples/SimplerEnv/train_files/starvla_cotrain_oxe.yaml",
+        help="Path to YAML config",
+    )
+    args, clipargs = parser.parse_known_args()
+
+    if os.getenv("DEBUGPY_ENABLE", "0") == "1":
+        import debugpy
+        debugpy.listen(("0.0.0.0", 10092))
+        print("Rank 0 waiting for debugger attach on port 10092...")
+        debugpy.wait_for_client()
+
+    cfg = OmegaConf.load(args.config_yaml)
+
+    cfg.framework.qwenvl.base_vlm = "./playground/Pretrained_models/Qwen3-VL-4B-Instruct"
+    qwen_vl = _QWen3_VL_Interface(cfg)
+    pass

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/training-code/QwenOFT.py

@@ -0,0 +1,1244 @@
+# Copyright 2025 starVLA community. All rights reserved.
+# Licensed under the MIT License, Version 1.0 (the "License");
+# Implemented by [Jinhui YE / HKUST University] in [2025].
+
+"""
+Qwen-OFT Framework
+
+A lightweight implementation that uses an action special token to parallelly predict continuous actions
+conditioned on multi-view images plus a language instruction (shares parameters with the VLM).
+Inspired by OpenVLA-OFT
+Key Points:
+  - Qwen2.5 vision-language backbone
+  - Injects an action special token into the VLM
+  - Continuous action prediction via L1 regression over the action special token hidden states
+
+
+Note: How to add special tokens to Qwen2.5:
+  download our model checkpoint with special tokens added: https://huggingface.co/StarVLA/Qwen2.5-VL-3B-Instruct-Action
+  or /starVLA/model/modules/vlm/tools/add_qwen_special_tokens/README.md （adpat a little code)
+
+"""
+
+from dataclasses import dataclass, field
+import time
+from typing import Any, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+
+from deployment.model_server.tools.image_tools import to_pil_preserve
+from starVLA.model.tools import FRAMEWORK_REGISTRY
+from starVLA.training.trainer_utils import initialize_overwatch
+
+logger = initialize_overwatch(__name__)
+
+# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels)
+IGNORE_INDEX = -100
+
+from starVLA.model.framework.base_framework import baseframework
+from starVLA.model.framework.share_tools import add_discretized_state_to_instruction, merge_framework_config
+from starVLA.model.modules.action_model.MLP_ActionHeader import get_action_model
+from starVLA.model.modules.vlm import get_vlm_model
+from starVLA.training.trainer_utils.trainer_tools import resize_images
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Default Config for QwenOFT
+#  - Documents every framework-level parameter with type + description
+#  - YAML values override these defaults; extra YAML keys are preserved
+# ──────────────────────────────────────────────────────────────────────
+@dataclass
+class QwenOFTDefaultConfig:
+    """QwenOFT framework default parameters.
+
+    All fields can be overridden by the corresponding key in the YAML
+    ``framework:`` section.  Extra YAML keys not listed here are kept
+    as-is (Config-as-API flexibility).
+    """
+
+    # --- Registry identifier (must match @FRAMEWORK_REGISTRY.register) ---
+    name: str = "QwenOFT"
+
+    # === VLM backbone (Qwen2.5-VL / Qwen3-VL) ===
+    qwenvl: dict = field(
+        default_factory=lambda: {
+            # Path to base VLM checkpoint (local or HF hub id)
+            "base_vlm": "./playground/Pretrained_models/Qwen3-VL-4B-Instruct-Action",
+            # Attention implementation: "flash_attention_2" | "eager" | "sdpa"
+            "attn_implementation": "flash_attention_2",
+        }
+    )
+
+    # === Action head (MLP regression over action special tokens) ===
+    action_model: dict = field(
+        default_factory=lambda: {
+            # Action head architecture type
+            "action_model_type": "MLP",
+            # Dimensionality of each action vector (e.g., 7 for 6-DoF + gripper)
+            "action_dim": 7,
+            # Proprioceptive state dimension for continuous_projector mode
+            "state_dim": 7,
+            # Hidden dim for the action MLP (auto-set from VLM hidden_size at runtime)
+            "action_hidden_dim": 2560,
+            # How many future steps to predict
+            "future_action_window_size": 8,
+            # How many past steps included in action chunk (usually 0)
+            "past_action_window_size": 0,
+            # Loss over predicted actions. Supported: "l1", "discrete_ce",
+            # "multibinary_bce", "asterix_factorized_ce".
+            # "multibinary_ce" is accepted as a CLI alias for "multibinary_bce".
+            "loss_type": "l1",
+            # State conditioning: preserve the existing text-bin path by default.
+            "state_encoding": "discretized_text",
+        }
+    )
+
+
+@FRAMEWORK_REGISTRY.register("QwenOFT")
+class Qwenvl_OFT(baseframework):
+    """
+    Multimodal vision-language-action model (OFT variant).
+
+    Components:
+      - Qwen2.5-VL / Qwen3-VL backbone for fused language/vision token embeddings
+      - Action special token injected into the VLM sequence
+      - MLP regression head over action token hidden states (L1 loss)
+
+    Focus: Predict future continuous actions conditioned on images + instruction.
+    """
+
+    def __init__(
+        self,
+        config: Optional[dict] = None,
+        **kwargs,
+    ) -> None:
+        """
+        Construct all submodules and cache key configuration values.
+
+        Args:
+            config: Hierarchical configuration (OmegaConf/dict) containing framework + trainer sections.
+            **kwargs: Reserved for future overrides (unused).
+        """
+        super().__init__()
+        # Merge framework defaults with YAML config (YAML wins on conflicts)
+        self.config = merge_framework_config(QwenOFTDefaultConfig, config)
+        self.qwen_vl_interface = get_vlm_model(config=self.config)
+        # align action_hidden_dim to VLM hidden_size at runtime
+        self.config.framework.action_model.action_hidden_dim = self.qwen_vl_interface.model.config.hidden_size
+        self.action_model = get_action_model(config=self.config)
+
+        # `action_horizon` is the single source of truth for chunk length.
+        # Legacy aliases (`future_action_window_size`, `past_action_window_size`)
+        # are normalised upstream by `share_tools.apply_config_compat`, so we
+        # only ever read `action_horizon` here.
+        self.action_horizon = int(self.config.framework.action_model.action_horizon)
+        self.chunk_len = self.action_horizon
+        self.action_dim = int(self.config.framework.action_model.action_dim)
+        self.action_env_dim = int(getattr(self.config.framework.action_model, "action_env_dim", self.action_dim))
+        self.action_loss_type = str(getattr(self.config.framework.action_model, "loss_type", "l1")).lower()
+        self.state_encoding = str(self.config.framework.action_model.state_encoding)
+        if self.state_encoding not in {"discretized_text", "continuous_projector"}:
+            raise ValueError(f"Unsupported state_encoding: {self.state_encoding}")
+        if self.state_encoding == "continuous_projector":
+            self.action_model.state_projector = nn.Linear(
+                int(self.config.framework.action_model.state_dim),
+                int(self.config.framework.action_model.action_hidden_dim),
+            )
+        cross_task_cfg = getattr(getattr(self.config, "rl_games", None), "cross_task", None)
+        self.loss_by_task = self._to_plain_dict(getattr(cross_task_cfg, "loss_by_task", None))
+        self.loss_weight_by_task = self._to_plain_dict(getattr(cross_task_cfg, "loss_weight_by_task", None))
+        # self.hidden_dim = config.framework.action_model.action_hidden_dim
+
+        self.action_token = "🔍"  # TODO also can add spacail token to Qwen, but too complex
+        self.action_token_id = self.qwen_vl_interface.processor.tokenizer("🔍", add_special_tokens=False)["input_ids"][0]
+
+        # --- Fixed-size KV memory; default OFF so existing paths are untouched. ---
+        kv_cfg = getattr(self.config.framework, "kv_memory", None)
+        self.kv_memory_enabled = bool(getattr(kv_cfg, "enabled", False)) if kv_cfg is not None else False
+        self.kv_window = int(getattr(kv_cfg, "window", 4)) if kv_cfg is not None else 4
+        self.kv_rollout_len = int(getattr(kv_cfg, "rollout_len", max(self.kv_window + 1, 2))) if kv_cfg is not None else self.kv_window + 1
+        # Packed single-forward training path (StreamingLLM-style sliding-window+sink
+        # mask via FlexAttention) — replaces the per-frame Python loop. Eval/live still
+        # uses the streaming KV memory (`_predict_action_memory`). OFF by default so the
+        # legacy loop stays the reference; enable per experiment via
+        # `framework.kv_memory.packed_train: true`.
+        self.kv_train_packed = bool(getattr(kv_cfg, "packed_train", False)) if kv_cfg is not None else False
+        # Per-step re-based sink for the packed path: replicate the (short) text
+        # prefix once per step and re-base its M-RoPE positions so the sink→frame
+        # distance matches the eval rolling-KV layout exactly (StreamingLLM
+        # "positions within the cache"), instead of one global sink whose distance
+        # grows with the step index. ON by default — the growing-distance sink was
+        # the train/eval forward mismatch behind the packed run's eval collapse.
+        # Set false to fall back to the single global sink for A/B.
+        self.kv_rebased_sink = bool(getattr(kv_cfg, "rebased_sink", True)) if kv_cfg is not None else True
+        if self.state_encoding == "continuous_projector" and self.kv_memory_enabled:
+            raise ValueError("continuous_projector requires framework.kv_memory.enabled=false")
+        # Per-stream (eval slot) memories; rebuilt per sample during training.
+        self._kv_memories: dict = {}
+
+        # L1 loss
+        self.l1_loss = nn.L1Loss()
+
+    def _profile_timing_enabled(self) -> bool:
+        return self.config.trainer.profile_timing.enabled
+
+    def _profile_sync(self) -> None:
+        if self._profile_timing_enabled() and torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    def _profile_start(self) -> float:
+        self._profile_sync()
+        return time.perf_counter()
+
+    def _profile_elapsed(self, start: float) -> float:
+        self._profile_sync()
+        return time.perf_counter() - start
+
+    @staticmethod
+    def _to_plain_dict(value: Any) -> dict[str, Any]:
+        if value is None or value == "":
+            return {}
+        if isinstance(value, dict):
+            return {str(key): item for key, item in value.items()}
+        if hasattr(value, "items"):
+            return {str(key): item for key, item in value.items()}
+        return {}
+
+    def _prepare_action_array(self, action: Any) -> np.ndarray:
+        values = np.asarray(action)
+        if values.ndim == 1:
+            values = values[None, :]
+        if values.shape[-1] == self.action_dim:
+            return values
+        if values.shape[-1] > self.action_dim:
+            return values[..., : self.action_dim]
+        pad_width = [(0, 0)] * values.ndim
+        pad_width[-1] = (0, self.action_dim - values.shape[-1])
+        return np.pad(values, pad_width, mode="constant")
+
+    def _compute_action_loss_for_type(
+        self,
+        pred_actions: torch.Tensor,
+        actions_target: torch.Tensor,
+        *,
+        loss_type: str,
+        action_env_dims: Optional[List[int]] = None,
+    ) -> torch.Tensor:
+        loss_type = str(loss_type).lower()
+        effective_dim = min(self.action_env_dim, pred_actions.shape[-1], actions_target.shape[-1])
+        if effective_dim <= 0:
+            raise ValueError(
+                f"Invalid action_env_dim={self.action_env_dim} for predicted shape={pred_actions.shape} "
+                f"and target shape={actions_target.shape}"
+            )
+
+        if loss_type in {"discrete_ce", "ce", "cross_entropy"}:
+            if action_env_dims is not None:
+                dims = torch.as_tensor(action_env_dims, device=pred_actions.device, dtype=torch.long)
+                if dims.numel() != pred_actions.shape[0]:
+                    raise ValueError(f"Expected {pred_actions.shape[0]} action_env_dims, got {dims.numel()}")
+                losses = []
+                weights = []
+                for dim_value in torch.unique(dims).tolist():
+                    dim = int(min(dim_value, pred_actions.shape[-1], actions_target.shape[-1]))
+                    if dim < 2:
+                        raise ValueError(f"discrete_ce requires at least 2 classes, got action_env_dim={dim}")
+                    mask = dims == int(dim_value)
+                    logits = pred_actions[mask, ..., :dim]
+                    target_class = actions_target[mask, ..., :dim].argmax(dim=-1).long()
+                    losses.append(F.cross_entropy(logits.reshape(-1, dim), target_class.reshape(-1), reduction="sum"))
+                    weights.append(target_class.numel())
+                total_weight = max(1, int(sum(weights)))
+                return sum(losses) / total_weight
+
+            if effective_dim < 2:
+                raise ValueError(
+                    f"action_model.loss_type={loss_type!r} requires at least 2 action classes, "
+                    f"got effective_dim={effective_dim}"
+                )
+            logits = pred_actions[..., :effective_dim]
+            target_class = actions_target[..., :effective_dim].argmax(dim=-1).long()
+            return F.cross_entropy(
+                logits.reshape(-1, effective_dim),
+                target_class.reshape(-1),
+            )
+
+        if loss_type in {"asterix_factorized_ce", "factorized_6_ce", "factorized_ce"}:
+            if effective_dim < 6:
+                raise ValueError(
+                    f"action_model.loss_type={loss_type!r} requires at least 6 action logits, "
+                    f"got effective_dim={effective_dim}"
+                )
+            vertical_logits = pred_actions[..., 0:3]
+            vertical_target = actions_target[..., 0:3].argmax(dim=-1).long()
+            horizontal_logits = pred_actions[..., 3:6]
+            horizontal_target = actions_target[..., 3:6].argmax(dim=-1).long()
+            vertical_loss = F.cross_entropy(
+                vertical_logits.reshape(-1, 3),
+                vertical_target.reshape(-1),
+            )
+            horizontal_loss = F.cross_entropy(
+                horizontal_logits.reshape(-1, 3),
+                horizontal_target.reshape(-1),
+            )
+            return 0.5 * (vertical_loss + horizontal_loss)
+
+        if loss_type in {"multibinary_bce", "multibinary_ce", "bce", "binary_cross_entropy"}:
+            logits = pred_actions[..., :effective_dim]
+            targets = (actions_target[..., :effective_dim] > 0).to(dtype=logits.dtype)
+            return F.binary_cross_entropy_with_logits(logits, targets)
+
+        if loss_type in {"mse", "l2"}:
+            return F.mse_loss(
+                pred_actions[..., :effective_dim],
+                actions_target[..., :effective_dim],
+            )
+
+        if loss_type not in {"l1", "mae"}:
+            raise ValueError(
+                f"Unsupported action_model.loss_type={loss_type!r}; "
+                "expected one of: l1, discrete_ce, multibinary_bce, asterix_factorized_ce"
+            )
+
+        return self.l1_loss(
+            pred_actions[..., :effective_dim],
+            actions_target[..., :effective_dim],
+        )
+
+    def _compute_action_loss(
+        self,
+        pred_actions: torch.Tensor,
+        actions_target: torch.Tensor,
+        action_env_dims: Optional[List[int]] = None,
+        rl_games_tasks: Optional[List[str]] = None,
+        action_loss_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if action_loss_mask is not None:
+            error = (pred_actions[..., :self.action_env_dim] - actions_target[..., :self.action_env_dim]).abs().mean(dim=-1)
+            error = error.masked_fill(~action_loss_mask, 0)
+            return (error.sum(dim=-1) / action_loss_mask.sum(dim=-1)).mean()
+        if self.loss_by_task and rl_games_tasks is not None:
+            if len(rl_games_tasks) != pred_actions.shape[0]:
+                raise ValueError(f"Expected {pred_actions.shape[0]} task labels, got {len(rl_games_tasks)}")
+            losses = []
+            weights = []
+            for task_name in sorted(set(str(task) for task in rl_games_tasks)):
+                mask_values = [str(task) == task_name for task in rl_games_tasks]
+                mask = torch.as_tensor(mask_values, device=pred_actions.device, dtype=torch.bool)
+                if not bool(mask.any()):
+                    continue
+                task_loss_type = str(self.loss_by_task.get(task_name, self.action_loss_type)).lower()
+                task_dims = (
+                    [int(dim) for dim, keep in zip(action_env_dims, mask_values) if keep]
+                    if action_env_dims is not None
+                    else None
+                )
+                task_loss = self._compute_action_loss_for_type(
+                    pred_actions[mask],
+                    actions_target[mask],
+                    loss_type=task_loss_type,
+                    action_env_dims=task_dims,
+                )
+                task_weight = float(self.loss_weight_by_task.get(task_name, 1.0))
+                sample_count = int(mask.sum().item())
+                losses.append(task_loss * task_weight * sample_count)
+                weights.append(task_weight * sample_count)
+            if losses:
+                return sum(losses) / max(1.0, float(sum(weights)))
+
+        return self._compute_action_loss_for_type(
+            pred_actions,
+            actions_target,
+            loss_type=self.action_loss_type,
+            action_env_dims=action_env_dims,
+        )
+
+    def _forward_qwen_last_hidden(self, qwen_inputs: dict) -> torch.Tensor:
+        """
+        Return only the final VLM hidden state when the interface supports it.
+
+        Requesting ``output_hidden_states=True`` from Hugging Face materializes
+        every layer's activations. For OFT we only need the last layer, so the
+        Qwen3 interface exposes a leaner path that preserves the practical memory
+        savings from gradient checkpointing.
+        """
+
+        if hasattr(self.qwen_vl_interface, "forward_last_hidden"):
+            return self.qwen_vl_interface.forward_last_hidden(**qwen_inputs)
+
+        qwenvl_outputs = self.qwen_vl_interface(
+            **qwen_inputs,
+            output_attentions=False,
+            output_hidden_states=True,
+            return_dict=True,
+            use_cache=False,
+        )
+        return qwenvl_outputs.hidden_states[-1]
+
+    def _state_conditioned_instructions(self, instructions, state):
+        if state is not None and self.state_encoding == "discretized_text":
+            return self.add_discretized_state_to_instruction(instructions, state)
+        return instructions
+
+    def _condition_action_queries(self, action_queries, state):
+        if self.state_encoding == "discretized_text":
+            return action_queries
+        state_tensor = torch.as_tensor(
+            np.stack(state), device=action_queries.device, dtype=action_queries.dtype
+        ).squeeze(1)
+        return action_queries + self.action_model.state_projector(state_tensor).unsqueeze(1)
+
+    def forward(
+        self,
+        examples: List[dict] = None,
+        **kwargs,
+    ) -> Tuple:
+        """
+        Training forward: directly regress future actions (no diffusion).
+
+        Flow:
+          1. Build QwenVL inputs (images + instruction tokens)
+          2. Extract hidden states from configured layer range
+          7. Predict action and compute L1 loss
+
+        Args:
+            examples: List[dict], each dict requires:
+                - image: List[PIL.Image] (multi-view)
+                - lang: str instruction
+                - action: np.ndarray or list shaped [T, action_dim]
+            **kwargs: Reserved.
+
+        Returns:
+            dict:
+                action_loss (torch.Tensor): Scalar diffusion noise prediction loss.
+        """
+        if self.kv_memory_enabled:
+            if self.kv_train_packed:
+                return self._forward_memory_packed(examples)
+            return self._forward_memory(examples)
+
+        batch_images = [example["image"] for example in examples]  #  [B，[PLT]]
+        instructions = [example["lang"] for example in examples]  # [B, str]
+        actions = [self._prepare_action_array(example["action"]) for example in examples]  # label [B， len, action_dim]
+        action_env_dims = [int(example.get("action_env_dim", self.action_env_dim)) for example in examples]
+        rl_games_tasks = [str(example.get("rl_games_task", "")) for example in examples]
+        if not any(rl_games_tasks):
+            rl_games_tasks = None
+        state = (
+            [example["state"] for example in examples] if "state" in examples[0] else None
+        )  # List[ndarray (1, state_dim)] or None
+
+        instructions = self._state_conditioned_instructions(instructions, state)
+
+        # step 0: add special action token to instruction
+        action_tokens = (
+            self.action_token * self.chunk_len
+        )  # can't add " " between two tokens, otherwise will be tokenized to multiple tokens
+        prompt_suffix = f" Please predict the next {self.chunk_len} robot actions: <action>{action_tokens}<action>."
+        instructions = [instruction + prompt_suffix for instruction in instructions]
+
+        profile_timing = self._profile_timing_enabled()
+        timing_metrics = {}
+        batch_stats = {}
+
+        # Step 1: QWenVL input format
+        t_inputs = self._profile_start() if profile_timing else None
+        qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions)
+        if profile_timing:
+            timing_metrics["timing/qwen_input_build_total"] = self._profile_elapsed(t_inputs)
+            timing_metrics.update(self.qwen_vl_interface._last_build_timing)
+            input_ids = qwen_inputs["input_ids"]
+            attention_mask = qwen_inputs["attention_mask"]
+            effective_tokens = attention_mask.sum()
+            effective_token_count = effective_tokens.detach().float().item()
+            batch_stats["batch/size"] = float(len(examples))
+            batch_stats["batch/input_len_max"] = float(input_ids.shape[1])
+            batch_stats["batch/input_len_mean"] = float(effective_token_count / input_ids.shape[0])
+            batch_stats["batch/padding_ratio"] = float(1.0 - effective_token_count / float(attention_mask.numel()))
+            batch_stats["batch/effective_tokens"] = float(effective_token_count)
+            batch_stats["batch/image_count"] = float(sum(len(images) for images in batch_images))
+            if "pixel_values" in qwen_inputs:
+                batch_stats["batch/pixel_values_rows"] = float(qwen_inputs["pixel_values"].shape[0])
+
+        t_vlm = self._profile_start() if profile_timing else None
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            last_hidden = self._forward_qwen_last_hidden(qwen_inputs)  # [B, L, H]
+        if profile_timing:
+            timing_metrics["timing/vlm_forward"] = self._profile_elapsed(t_vlm)
+
+        # Step 4: Action Expert Forward and Loss
+        t_action = self._profile_start() if profile_timing else None
+        with torch.autocast("cuda", dtype=torch.float32):
+            # Extract action token embeddings as action prediction queries
+            input_ids = qwen_inputs["input_ids"]
+            action_queries = self._gather_action_token_embeddings(
+                last_hidden, input_ids, action_token_id=self.action_token_id
+            )  # [B, chunk_len, H]
+            action_queries = self._condition_action_queries(action_queries, state)
+            pred_actions = self.action_model.predict_action(action_queries)  # (B, chunk_len, action_dim)
+
+            # Label alignment: take the last chunk_len segment
+            actions = torch.tensor(
+                np.array(actions), device=pred_actions.device, dtype=pred_actions.dtype
+            )  # [B, T_full, action_dim]
+            actions_target = actions[:, -self.action_horizon :, :]  # (B, action_horizon, action_dim)
+
+            action_loss_mask = None
+            if "action_loss_mask" in examples[0]:
+                action_loss_mask = torch.as_tensor(
+                    np.asarray([example["action_loss_mask"] for example in examples]),
+                    device=pred_actions.device,
+                    dtype=torch.bool,
+                )
+
+            action_loss = self._compute_action_loss(
+                pred_actions,
+                actions_target,
+                action_env_dims=action_env_dims,
+                rl_games_tasks=rl_games_tasks,
+                action_loss_mask=action_loss_mask,
+            )
+        if profile_timing:
+            timing_metrics["timing/action_head_loss"] = self._profile_elapsed(t_action)
+
+        return {"action_loss": action_loss, "loss_weight": float(len(examples)), "timing": timing_metrics, "batch_stats": batch_stats}
+
+    @torch.inference_mode()
+    def predict_action(
+        self,
+        examples: List[dict] = None,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """
+
+        Steps:
+          1. Resize images to training resolution (if specified)
+          2. Encode with QwenVL (hidden states retained)
+          6. Return normalized action trajectory
+
+        Returns:
+            dict:
+                normalized_actions (np.ndarray): Shape [B, T, action_dim], diffusion-sampled normalized actions.
+        """
+        if type(examples) is not list:
+            examples = [examples]
+        if self.kv_memory_enabled:
+            return self._predict_action_memory(examples)
+        batch_images = [to_pil_preserve(example["image"]) for example in examples]  #  [B，[PLT]]
+        instructions = [example["lang"] for example in examples]  # [B, str]
+        state = (
+            [example["state"] for example in examples] if "state" in examples[0] else None
+        )  # List[ndarray (1, state_dim)] or None
+
+        instructions = self._state_conditioned_instructions(instructions, state)
+
+        train_obs_image_size = getattr(self.config.datasets.vla_data, "obs_image_size", None)
+        if train_obs_image_size:
+            batch_images = resize_images(batch_images, target_size=train_obs_image_size)
+
+        # step 0: add special action token to instruction
+        action_tokens = (
+            self.action_token * self.chunk_len
+        )  # can't add " " between two tokens, otherwise will be tokenized to multiple tokens
+        prompt_suffix = f" Please predict the next {self.chunk_len} robot actions: <action>{action_tokens}<action>."
+        instructions = [instruction + prompt_suffix for instruction in instructions]
+
+        # Step 1: QWenVL input format
+        qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            last_hidden = self._forward_qwen_last_hidden(qwen_inputs)  # [B, L, H]
+
+        # Step 4: Action Expert Forward and Loss
+        with torch.autocast("cuda", dtype=torch.float32):
+            # Extract action token embeddings as action prediction queries
+            input_ids = qwen_inputs["input_ids"]
+            action_queries = self._gather_action_token_embeddings(
+                last_hidden, input_ids, action_token_id=self.action_token_id
+            )  # [B, chunk_len, H]
+            action_queries = self._condition_action_queries(action_queries, state)
+            pred_actions = self.action_model.predict_action(action_queries)  # (B, chunk_len, action_dim)
+
+        normalized_actions = pred_actions.detach().cpu().numpy()
+        return {"normalized_actions": normalized_actions}
+
+    # ──────────────────────────────────────────────────────────────────
+    #  Fixed-size KV memory. Default OFF. The cache/rotation core is
+    #  unit-proven (tests/test_kv_memory_equivalence.py); the processor-driven
+    #  segment construction below still needs end-to-end validation against the
+    #  real Qwen3-VL-4B checkpoint on GPU.
+    # ──────────────────────────────────────────────────────────────────
+    def reset_memory(self, slot_id=None) -> None:
+        """Clear the KV memory for one eval slot (episode boundary), or all."""
+        if slot_id is None:
+            self._kv_memories.clear()
+        else:
+            self._kv_memories.pop(slot_id, None)
+
+    def _kv_components(self):
+        hf = self.qwen_vl_interface.model  # Qwen3VLForConditionalGeneration
+        backbone = hf.model  # Qwen3VLModel
+        text_model = backbone.language_model  # Qwen3VLTextModel
+        return backbone, text_model.rotary_emb, len(text_model.layers), text_model.config
+
+    def _new_kv_memory(self):
+        from starVLA.model.modules.vlm.kv_memory import FrameKVMemory
+
+        _, rotary_emb, num_layers, text_config = self._kv_components()
+        return FrameKVMemory(rotary_emb, self.kv_window, num_layers, text_config)
+
+    def _build_text_first_inputs(self, frames, instruction):
+        """Tokenize a text-first window: [instruction][frame imgs...][action suffix].
+
+        Returns the processor batch (input_ids, pixel_values, image_grid_thw, ...).
+        Static text comes first so its KV stays positionally stable across steps.
+        """
+        action_tokens = self.action_token * self.chunk_len
+        action_suffix = f" Please predict the next {self.chunk_len} robot actions: <action>{action_tokens}<action>."
+        content = [{"type": "text", "text": instruction}]
+        content += [{"type": "image", "image": img} for img in frames]
+        content += [{"type": "text", "text": action_suffix}]
+        messages = [[{"role": "user", "content": content}]]
+        batch = self.qwen_vl_interface.processor.apply_chat_template(
+            messages, tokenize=True, padding=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
+        )
+        return batch.to(self.qwen_vl_interface.model.device)
+
+    def _kv_segments(self, input_ids_1d):
+        """Partition a text-first window's ids into (prefix_len, [frame_lens], action_len).
+
+        A frame block spans its vision_start .. vision_end (inclusive); the prefix is
+        everything before the first vision_start; the action suffix is everything
+        after the last vision_end. Using the real special-token ids keeps eviction
+        leak-free (each evicted frame removes exactly its own markers).
+        """
+        cfg = self.qwen_vl_interface.model.config
+        vstart = int(getattr(cfg, "vision_start_token_id", 151652))
+        vend = int(getattr(cfg, "vision_end_token_id", 151653))
+        ids = input_ids_1d.tolist()
+        starts = [i for i, t in enumerate(ids) if t == vstart]
+        ends = [i for i, t in enumerate(ids) if t == vend]
+        if not starts or not ends:
+            raise RuntimeError("KV memory: window inputs have no vision_start/end tokens")
+        action_start = ends[-1] + 1
+        bounds = starts + [action_start]
+        frame_lens = [bounds[i + 1] - bounds[i] for i in range(len(starts))]
+        return starts[0], frame_lens, len(ids) - action_start
+
+    def _kv_forward_step_batched(self, memories, frames, instruction):
+        """Run ONE streaming KV-memory step for a group of same-layout streams and
+        return ``(last_hidden_new [B, new_len, H], new_input_ids [B, new_len])``.
+
+        ``memories`` and ``frames`` are aligned lists (one entry per stream / eval
+        slot); all entries share ``instruction`` and history layout (grouped by the
+        caller). Two efficiencies vs a naive per-frame, per-step encode:
+
+        - Single-frame preprocessing: the image processor runs once per stream on its
+          OWN current frame. The window's multi-frame token layout (needed only for
+          M-RoPE positions via ``get_rope_index``) is reconstructed by replicating the
+          frame's token block, so duplicate frames are never re-preprocessed.
+        - Cross-stream batching: the streams' canonical KV blocks are stacked along the
+          batch dim, so the whole group is one forward; results are sliced back to each
+          stream's memory afterwards.
+
+        Past frames come from each stream's memory; only the current frame's pixels are
+        encoded fresh. ``B == 1`` (training rollout) reuses the same path with no
+        stacking, so gradients flow through the in-place memory updates.
+        """
+        from starVLA.model.modules.vlm.kv_memory import FrameKVMemory, memory_step
+
+        backbone = self.qwen_vl_interface.model.model  # Qwen3VLModel
+        B = len(memories)
+        n_past = memories[0].num_past_frames()
+        n_visible = n_past + 1
+
+        # One processor call per stream's current frame (no duplicate-frame work).
+        per_stream = [self._build_text_first_inputs([frame], instruction) for frame in frames]
+        ids1 = per_stream[0]["input_ids"][0]
+        cur_grid = per_stream[0]["image_grid_thw"][:1]
+        prefix_len, frame_lens, action_len = self._kv_segments(ids1)
+        frame_len = frame_lens[0]
+
+        # Reconstruct the n_visible-frame window ids by replicating the frame block,
+        # then let get_rope_index assign the (content-independent) M-RoPE positions.
+        prefix, action = ids1[:prefix_len], ids1[prefix_len + frame_len :]
+        block = ids1[prefix_len : prefix_len + frame_len]
+        window_ids = torch.cat([prefix] + [block] * n_visible + [action]).unsqueeze(0)
+        full_pos1, _ = backbone.get_rope_index(window_ids, image_grid_thw=cur_grid.repeat(n_visible, 1))
+        full_pos = full_pos1.expand(3, B, full_pos1.shape[-1]).contiguous()
+
+        batched = FrameKVMemory.stack(memories) if B > 1 else memories[0]
+        past_len = batched.past_length()
+        is_prefill = not batched.has_text()
+        expected_past = 0 if is_prefill else prefix_len + n_past * frame_len
+        if past_len != expected_past:
+            raise RuntimeError(
+                f"KV memory slice misaligned: memory past_length={past_len} but window "
+                f"expects {expected_past}. Token boundaries drifted across steps."
+            )
+        if is_prefill:
+            new_start, frame_start_in_new, prefill_text_len = 0, prefix_len, prefix_len
+        else:
+            new_start, frame_start_in_new, prefill_text_len = past_len, 0, 0
+
+        new_ids = window_ids[:, new_start:].expand(B, -1).contiguous()
+        pixels = torch.cat([b["pixel_values"] for b in per_stream], dim=0)
+        grid = cur_grid.repeat(B, 1)
+
+        last_hidden = memory_step(
+            backbone,
+            batched,
+            new_input_ids=new_ids,
+            pixel_values=pixels,
+            image_grid_thw=grid,
+            full_position_ids=full_pos,
+            frame_start_in_new=frame_start_in_new,
+            frame_len=frame_len,
+            prefill_text_len=prefill_text_len,
+        )
+
+        if B > 1:
+            for i, memory in enumerate(memories):
+                sliced = batched.slice(i)
+                memory.text, memory.frames = sliced.text, sliced.frames
+        return last_hidden, new_ids
+
+    # ------------------------------------------------------------------ #
+    # Packed single-forward training path (StreamingLLM-style mask).
+    #
+    # Instead of the per-frame Python loop (`_forward_memory`), pack the whole
+    # rollout into ONE sequence
+    #     [text/sink] [frame_1][action_1] ... [frame_R][action_R]
+    # and run a SINGLE Qwen3-VL forward under a FlexAttention BlockMask that
+    # reproduces the streaming visibility:
+    #   * the text prefix is an always-visible attention sink (StreamingLLM);
+    #   * each frame attends only the last `kv_window` frames (sliding window);
+    #   * each action read-out block is isolated (seen only by its own step), so
+    #     the frame KV stream is bit-identical to the streaming loop / eval.
+    # Custom M-RoPE `position_ids` keep frame->frame and action->frame relative
+    # distances exactly matching eval; only the text-sink->frame absolute distance
+    # grows with step (accepted approximation — the sink is distance-robust).
+    # Eval/live (`_predict_action_memory`) is untouched and keeps the rolling KV.
+    # ------------------------------------------------------------------ #
+    def _packed_window_layout(self, frames, instruction):
+        """Build the packed [text][f1][a1]...[fR][aR] window for ONE stream.
+
+        Token structure is content-independent (image placeholders are identical
+        per frame), so the prefix / frame / action blocks are reconstructed by
+        token-replication (matching `_kv_forward_step_batched`'s `window_ids`),
+        which guarantees each action block equals the per-step layout exactly.
+
+        Returns dict with: input_ids [1,L], pixel_values, image_grid_thw [R,3],
+        position_ids [3,1,L] (custom M-RoPE), kind [L] (0=text,1=frame,2=action),
+        step [L] (frame/step index for frame&action tokens, -1 for text),
+        action_idx [R, chunk_len] (packed indices of each step's action tokens).
+        """
+        R = len(frames)
+        backbone = self.qwen_vl_interface.model.model  # Qwen3VLModel
+        device = self.qwen_vl_interface.model.device
+
+        # One processor call per frame: identical token block, distinct pixels.
+        per_frame = [self._build_text_first_inputs([f], instruction) for f in frames]
+        ids1 = per_frame[0]["input_ids"][0]
+        grid1 = per_frame[0]["image_grid_thw"][:1]  # [1,3]
+        prefix_len, frame_lens, action_len = self._kv_segments(ids1)
+        frame_len = frame_lens[0]
+        for pf in per_frame:
+            pl, fl, al = self._kv_segments(pf["input_ids"][0])
+            if (pl, fl[0], al) != (prefix_len, frame_len, action_len):
+                raise RuntimeError(
+                    "Packed KV memory: frames have inconsistent token layout "
+                    f"({(pl, fl[0], al)} vs {(prefix_len, frame_len, action_len)}); "
+                    "all frames in a rollout must share resolution/grid."
+                )
+
+        prefix = ids1[:prefix_len]
+        frame_block = ids1[prefix_len : prefix_len + frame_len]
+        action_block = ids1[prefix_len + frame_len :]  # incl. assistant header
+        block_len = frame_len + action_len
+
+        pixel_values = torch.cat([pf["pixel_values"] for pf in per_frame], dim=0)
+        grid = grid1.repeat(R, 1)  # [R, 3]
+
+        # Frames-only M-RoPE positions: frames laid out consecutively (no action
+        # in between), so frame→frame / action→frame relative distances match eval
+        # ("positions within the cache"). These are the canonical frame positions
+        # reused regardless of where the frame sits in the packed token sequence.
+        frames_only_ids = torch.cat([prefix] + [frame_block] * R).unsqueeze(0)
+        full_pos_f, _ = backbone.get_rope_index(frames_only_ids, image_grid_thw=grid)  # [3,1,Lf]
+        text_pos = full_pos_f[:, :, :prefix_len]
+
+        action_offsets = (action_block == self.action_token_id).nonzero(as_tuple=False).flatten()
+        if action_offsets.numel() < self.chunk_len:
+            raise RuntimeError(
+                f"Packed KV memory: action block has {action_offsets.numel()} action "
+                f"tokens (< chunk_len={self.chunk_len})."
+            )
+        action_offsets = action_offsets[-self.chunk_len :]  # last chunk_len, matches gather
+
+        def _frame_pos(k):
+            s = prefix_len + k * frame_len
+            return full_pos_f[:, :, s : s + frame_len]
+
+        def _action_pos(frame_pos_k):
+            fk_max = int(frame_pos_k.max().item())
+            ap = torch.arange(1, action_len + 1, device=device) + fk_max
+            return ap.view(1, 1, -1).expand(3, 1, -1)
+
+        W = self.kv_window
+        parts, pos_parts, kind_parts, step_parts, action_idx_rows = [], [], [], [], []
+
+        if self.kv_rebased_sink:
+            # Per-step re-based sink, with the redundant pre-eviction sinks deduped.
+            # Steps k < W all have j0=0 → o_k=0 → an identical sink at [0,P), so they
+            # share ONE front sink (step=-2, visible to queries with step < W). This
+            # is exactly eval before any eviction (sink pinned at the front, frames
+            # growing from P). Steps k >= W each carry their own re-based sink
+            # ([sink_k][frame_k][action_k]) at o_k so the sink→frame distance equals
+            # eval's rolling-KV layout (bit-exact). Net: 1 + max(0,R-W) sinks instead
+            # of R, saving (W-1) prefix copies while staying numerically identical.
+            base0 = int(full_pos_f[0, 0, prefix_len].item())  # frame 0's temporal start
+
+            # Shared front sink for the pre-eviction steps (step=-2).
+            parts.append(prefix)
+            pos_parts.append(text_pos)
+            kind_parts.append(torch.zeros(prefix_len, dtype=torch.long, device=device))
+            step_parts.append(torch.full((prefix_len,), -2, dtype=torch.long, device=device))
+            cursor = prefix_len
+
+            for k in range(R):
+                if k >= W:
+                    # Re-based sink_k sits right before step k's oldest visible frame.
+                    j0 = k - W + 1
+                    o_k = int(full_pos_f[0, 0, prefix_len + j0 * frame_len].item()) - base0
+                    parts.append(prefix)
+                    pos_parts.append(text_pos + o_k)  # shift all 3 M-RoPE dims by o_k
+                    kind_parts.append(torch.zeros(prefix_len, dtype=torch.long, device=device))
+                    step_parts.append(torch.full((prefix_len,), k, dtype=torch.long, device=device))
+                    cursor += prefix_len
+                frame_pos_k = _frame_pos(k)
+                parts += [frame_block, action_block]
+                pos_parts += [frame_pos_k, _action_pos(frame_pos_k)]
+                kind_parts.append(
+                    torch.cat([
+                        torch.ones(frame_len, dtype=torch.long, device=device),
+                        torch.full((action_len,), 2, dtype=torch.long, device=device),
+                    ])
+                )
+                step_parts.append(torch.full((block_len,), k, dtype=torch.long, device=device))
+                action_idx_rows.append(cursor + frame_len + action_offsets)
+                cursor += block_len
+        else:
+            # Legacy single global sink: [sink][f0 a0]...[f_{R-1} a_{R-1}]. Sink at
+            # the front (step=-1, always visible); sink→frame distance grows with k.
+            parts.append(prefix)
+            pos_parts.append(text_pos)
+            kind_parts.append(torch.zeros(prefix_len, dtype=torch.long, device=device))
+            step_parts.append(torch.full((prefix_len,), -1, dtype=torch.long, device=device))
+            cursor = prefix_len
+            for k in range(R):
+                frame_pos_k = _frame_pos(k)
+                parts += [frame_block, action_block]
+                pos_parts += [frame_pos_k, _action_pos(frame_pos_k)]
+                kind_parts.append(
+                    torch.cat([
+                        torch.ones(frame_len, dtype=torch.long, device=device),
+                        torch.full((action_len,), 2, dtype=torch.long, device=device),
+                    ])
+                )
+                step_parts.append(torch.full((block_len,), k, dtype=torch.long, device=device))
+                action_idx_rows.append(cursor + frame_len + action_offsets)
+                cursor += block_len
+
+        input_ids = torch.cat(parts).unsqueeze(0)  # [1, L]
+        position_ids = torch.cat(pos_parts, dim=2)  # [3,1,L]
+        kind = torch.cat(kind_parts)
+        step = torch.cat(step_parts)
+        action_idx = torch.stack(action_idx_rows, dim=0)  # [R, chunk_len]
+
+        return {
+            "input_ids": input_ids,
+            "pixel_values": pixel_values,
+            "image_grid_thw": grid,
+            "position_ids": position_ids,
+            "kind": kind,
+            "step": step,
+            "action_idx": action_idx,
+        }
+
+    def _packed_block_mask(self, kind, step, L, batch):
+        """Build a FlexAttention BlockMask for the packed sliding-window+sink layout."""
+        from torch.nn.attention.flex_attention import create_block_mask
+        from starVLA.model.modules.vlm.kv_memory import make_packed_sliding_mask_mod
+
+        mask_mod = make_packed_sliding_mask_mod(kind, step, self.kv_window)
+        # B=H=None → the mask is shared across batch and heads (it depends only on
+        # token kind/step, not content), so one BlockMask covers the whole batch.
+        return create_block_mask(mask_mod, B=None, H=None, Q_LEN=L, KV_LEN=L, device=str(kind.device))
+
+    def _assert_flex_attention(self):
+        """Packed training needs the decoder dispatched through FlexAttention.
+
+        We pass a FlexAttention BlockMask straight to the model, so the decoder
+        (forward AND gradient-checkpointing recompute during backward) must use
+        the flex backend the whole time — a transient toggle around the forward
+        would leave the recompute on sdpa and choke on the BlockMask. The eval
+        streaming path also runs fine under flex, so the whole model is simply
+        loaded with `attn_implementation: flex_attention`.
+        """
+        _, _, _, text_config = self._kv_components()
+        impl = getattr(text_config, "_attn_implementation", None)
+        if impl != "flex_attention":
+            raise RuntimeError(
+                "Packed KV-memory training requires the VLM loaded with "
+                f"attn_implementation='flex_attention' (got {impl!r}). Set "
+                "framework.qwenvl.attn_implementation: flex_attention."
+            )
+
+    def _forward_memory_packed(self, examples: List[dict]) -> dict:
+        """Training forward via a single packed sliding-window+sink masked pass.
+
+        Per-frame supervision identical to `_forward_memory` (every real frame is
+        supervised at the window it naturally has), but the expensive VLM forward
+        runs ONCE per rollout instead of R times. Full BPTT across the rollout
+        (no per-step detach) — pair with gradient checkpointing for memory.
+        """
+        self._assert_flex_attention()
+        profile = self._profile_timing_enabled()
+        timing = {"timing/mem_nograd_fwd": 0.0, "timing/mem_grad_fwd": 0.0, "timing/mem_action_head": 0.0}
+
+        instructions = [
+            self._instruction_with_state(ex["lang"], ex.get("state") if "state" in ex else None)
+            for ex in examples
+        ]
+        groups: dict = {}
+        for idx, instruction in enumerate(instructions):
+            valid = np.asarray(examples[idx]["valid"], dtype=bool)
+            groups.setdefault((instruction, tuple(valid.tolist())), []).append(idx)
+
+        total_weighted_loss = None
+        total_weight = 0.0
+        for (instruction, valid_key), indices in groups.items():
+            valid = np.asarray(valid_key, dtype=bool)
+            real_positions = [k for k in range(len(valid)) if valid[k]]
+            R = len(real_positions)
+            if R == 0:
+                continue
+            action_env_dims = [int(examples[i].get("action_env_dim", self.action_env_dim)) for i in indices]
+            rl_games_tasks = [str(examples[i].get("rl_games_task", "")) for i in indices]
+            if not any(rl_games_tasks):
+                rl_games_tasks = None
+
+            # Build per-sample packed layouts (identical token structure, distinct pixels).
+            layouts = [
+                self._packed_window_layout([examples[i]["image"][k] for k in real_positions], instruction)
+                for i in indices
+            ]
+            base = layouts[0]
+            L = base["input_ids"].shape[1]
+            B = len(indices)
+            input_ids = base["input_ids"].expand(B, -1).contiguous()
+            position_ids = base["position_ids"].expand(3, B, -1).contiguous()
+            pixel_values = torch.cat([lay["pixel_values"] for lay in layouts], dim=0)
+            grid = torch.cat([lay["image_grid_thw"] for lay in layouts], dim=0)  # [B*R, 3]
+            action_idx = base["action_idx"]  # [R, chunk_len]
+            block_mask = self._packed_block_mask(base["kind"], base["step"], L, B)
+
+            t_fwd = self._profile_start() if profile else None
+            last_hidden = self._forward_qwen_last_hidden(
+                {
+                    "input_ids": input_ids,
+                    "position_ids": position_ids,
+                    "attention_mask": block_mask,
+                    "pixel_values": pixel_values,
+                    "image_grid_thw": grid,
+                    "use_cache": False,
+                }
+            )  # [B, L, H]
+            if profile:
+                timing["timing/mem_grad_fwd"] += self._profile_elapsed(t_fwd)
+
+            t_head = self._profile_start() if profile else None
+            with torch.autocast("cuda", dtype=torch.float32):
+                for j, k in enumerate(real_positions):
+                    action_queries = last_hidden[:, action_idx[j], :]  # [B, chunk_len, H]
+                    group_preds = self.action_model.predict_action(action_queries)  # [B, chunk, dim]
+                    targets = [self._prepare_action_array(examples[i]["actions_per_frame"][k]) for i in indices]
+                    targets = torch.tensor(np.array(targets), device=group_preds.device, dtype=group_preds.dtype)
+                    targets = targets[:, -self.action_horizon :, :]
+                    step_loss = self._compute_action_loss(
+                        group_preds, targets, action_env_dims=action_env_dims, rl_games_tasks=rl_games_tasks
+                    )
+                    weight = float(sum(float(np.asarray(examples[i]["density_weight"])[k]) for i in indices))
+                    contribution = step_loss * weight
+                    total_weighted_loss = contribution if total_weighted_loss is None else total_weighted_loss + contribution
+                    total_weight += weight
+            if profile:
+                timing["timing/mem_action_head"] += self._profile_elapsed(t_head)
+
+        action_loss = total_weighted_loss / max(total_weight, 1e-6)
+        return {"action_loss": action_loss, "loss_weight": total_weight, "timing": timing, "batch_stats": {}}
+
+    def _instruction_with_state(self, instruction, state):
+        if state is None:
+            return instruction
+        return self.add_discretized_state_to_instruction([instruction], [state])[0]
+
+    def _forward_memory(self, examples: List[dict]) -> dict:
+        """Training forward with the streaming KV memory: per-frame supervision.
+
+        Each sample is an R-frame rollout `[t-R+1 .. t]`; the dataloader marks clamp
+        padding via ``valid`` and supplies one action chunk per frame in
+        ``actions_per_frame`` ([R, H, dim]). We skip padding frames (never feed memory,
+        never supervise) and supervise EVERY real frame at the window it naturally has
+        (1 -> window, then evict), which matches eval (`_predict_action_memory`) frame
+        for frame, including the unfilled window at an episode's start. History KV is
+        detached between steps (truncated BPTT depth 1), so each step's loss flows only
+        through that step's forward; the per-step weighted losses are summed and the
+        trainer backpropagates once.
+
+        Same-instruction, same-`valid` rollouts share a token layout and window
+        evolution, so they run each step in one batched forward via
+        `_kv_forward_step_batched` (degrades to per-sample otherwise), mirroring
+        `_predict_action_memory`.
+        """
+        profile = self._profile_timing_enabled()
+        timing = {"timing/mem_nograd_fwd": 0.0, "timing/mem_grad_fwd": 0.0, "timing/mem_action_head": 0.0}
+
+        instructions = [
+            self._instruction_with_state(ex["lang"], ex.get("state") if "state" in ex else None)
+            for ex in examples
+        ]
+
+        # Group rollouts that share the prompt AND the same real/padding pattern: those
+        # march through identical window states each step, so they batch cleanly.
+        groups: dict = {}
+        for idx, instruction in enumerate(instructions):
+            valid = np.asarray(examples[idx]["valid"], dtype=bool)
+            groups.setdefault((instruction, tuple(valid.tolist())), []).append(idx)
+
+        total_weighted_loss = None
+        total_weight = 0.0
+        for (instruction, valid_key), indices in groups.items():
+            valid = np.asarray(valid_key, dtype=bool)
+            real_positions = [k for k in range(len(valid)) if valid[k]]
+            memories = [self._new_kv_memory() for _ in indices]
+            rollout = [examples[i]["image"] for i in indices]
+            action_env_dims = [int(examples[i].get("action_env_dim", self.action_env_dim)) for i in indices]
+            rl_games_tasks = [str(examples[i].get("rl_games_task", "")) for i in indices]
+            if not any(rl_games_tasks):
+                rl_games_tasks = None
+
+            for step_pos, k in enumerate(real_positions):
+                if step_pos > 0:
+                    for memory in memories:
+                        memory.detach_()  # truncated BPTT: history is a fixed memory
+                step_frames = [frames[k] for frames in rollout]
+
+                t_fwd = self._profile_start() if profile else None
+                last_hidden, new_ids = self._kv_forward_step_batched(memories, step_frames, instruction)
+                if profile:
+                    timing["timing/mem_grad_fwd"] += self._profile_elapsed(t_fwd)
+
+                t_head = self._profile_start() if profile else None
+                with torch.autocast("cuda", dtype=torch.float32):
+                    action_queries = self._gather_action_token_embeddings(
+                        last_hidden, new_ids, action_token_id=self.action_token_id
+                    )
+                    group_preds = self.action_model.predict_action(action_queries)  # [Bg, chunk, dim]
+                    targets = [
+                        self._prepare_action_array(examples[i]["actions_per_frame"][k]) for i in indices
+                    ]
+                    targets = torch.tensor(
+                        np.array(targets), device=group_preds.device, dtype=group_preds.dtype
+                    )  # [Bg, H, dim]
+                    targets = targets[:, -self.action_horizon :, :]
+                    step_loss = self._compute_action_loss(
+                        group_preds,
+                        targets,
+                        action_env_dims=action_env_dims,
+                        rl_games_tasks=rl_games_tasks,
+                    )
+                if profile:
+                    timing["timing/mem_action_head"] += self._profile_elapsed(t_head)
+
+                # Per-supervised-step density weight: sum of the group's per-frame
+                # weights at frame k. Uniform weights reduce this to a sample-count
+                # weighting, i.e. the exact mean over all supervised (sample, step).
+                weight = float(
+                    sum(float(np.asarray(examples[i]["density_weight"])[k]) for i in indices)
+                )
+                contribution = step_loss * weight
+                total_weighted_loss = contribution if total_weighted_loss is None else total_weighted_loss + contribution
+                total_weight += weight
+
+        action_loss = total_weighted_loss / max(total_weight, 1e-6)
+        return {"action_loss": action_loss, "loss_weight": total_weight, "timing": timing, "batch_stats": {}}
+
+    @torch.inference_mode()
+    def _predict_action_memory(self, examples: List[dict]) -> dict:
+        """Eval-time streaming prediction with per-slot KV memory (req3, milestone 4).
+
+        The newest frame (``example['image'][-1]``) is encoded fresh; past frames are
+        served from the slot's memory. ``reset_memory(slot_id)`` (called from the
+        policy at episode boundaries) clears a slot. Streams that currently share the
+        same instruction and history layout are processed in one batched forward; the
+        rest fall back to their own group (no left padding needed).
+        """
+        train_obs_image_size = getattr(self.config.datasets.vla_data, "obs_image_size", None)
+
+        # Resolve each example to its (slot memory, current frame, instruction).
+        entries = []
+        for idx, example in enumerate(examples):
+            slot_id = example.get("slot_id", idx)
+            frames = [to_pil_preserve(f) for f in example["image"]]
+            if train_obs_image_size:
+                frames = resize_images(frames, target_size=train_obs_image_size)
+            state = example.get("state") if "state" in example else None
+            instruction = self._instruction_with_state(example["lang"], state)
+            memory = self._kv_memories.get(slot_id)
+            if memory is None:
+                memory = self._new_kv_memory()
+                self._kv_memories[slot_id] = memory
+            entries.append({"idx": idx, "frame": frames[-1], "instruction": instruction, "memory": memory})
+
+        # Group same-layout streams (same prompt + same cached layout) into one forward.
+        groups: dict = {}
+        for entry in entries:
+            key = (entry["instruction"], entry["memory"].layout_key())
+            groups.setdefault(key, []).append(entry)
+
+        preds: List[Optional[torch.Tensor]] = [None] * len(examples)
+        for group in groups.values():
+            instruction = group[0]["instruction"]
+            memories = [entry["memory"] for entry in group]
+            frames = [entry["frame"] for entry in group]
+            last_hidden, new_ids = self._kv_forward_step_batched(memories, frames, instruction)
+            with torch.autocast("cuda", dtype=torch.float32):
+                action_queries = self._gather_action_token_embeddings(
+                    last_hidden, new_ids, action_token_id=self.action_token_id
+                )
+                pred_actions = self.action_model.predict_action(action_queries)
+            for i, entry in enumerate(group):
+                preds[entry["idx"]] = pred_actions[i : i + 1]
+
+        normalized = torch.cat(preds, dim=0).detach().cpu().numpy()
+        return {"normalized_actions": normalized}
+
+    def _gather_action_token_embeddings(
+        self,
+        last_hidden: torch.Tensor,  # [B, L, H]
+        input_ids: torch.Tensor,  # [B, L]
+        action_token_id=None,  # Can be int or List[int]
+    ) -> torch.Tensor:
+        """
+        Vectorized batch extraction of action token embeddings:
+          - No per-sample for loop
+          - Select the last chunk_len action placeholder tokens from each sample
+        Args:
+            last_hidden: [B, L, H]
+            input_ids:   [B, L]
+            action_token_id: int or List[int]
+        Returns:
+            action_queries: [B, chunk_len, H]
+        """
+        if action_token_id is None:
+            raise ValueError("action_token_id must not be None")
+
+        device = input_ids.device
+        B, L, H = last_hidden.shape
+
+        # Support multiple ids (e.g., multiple variants)
+        if isinstance(action_token_id, (list, tuple, set)):
+            id_list = torch.tensor(list(action_token_id), device=device, dtype=input_ids.dtype)
+            # torch.isin requires PyTorch >=1.10
+            mask = torch.isin(input_ids, id_list)
+        else:
+            mask = input_ids == action_token_id  # [B, L]
+
+        counts = mask.sum(dim=1)  # [B]
+        if (counts < self.chunk_len).any():
+            insufficient = (counts < self.chunk_len).nonzero(as_tuple=False).flatten().tolist()
+            raise RuntimeError(
+                f"The following samples have insufficient action tokens (< {self.chunk_len}): {insufficient} |"
+                f" counts={counts.tolist()}"
+            )
+
+        # Position indices
+        idx = torch.arange(L, device=device).unsqueeze(0).expand(B, L)  # [B, L]
+        masked_pos = torch.where(mask, idx, torch.full_like(idx, -1))  # Set non-action positions to -1
+
+        # Take the last chunk_len positions (higher indices = later in sequence)
+        # Note: count sufficiency already verified, so -1 won't be incorrectly selected
+        topk_pos = masked_pos.topk(k=self.chunk_len, dim=-1).values  # [B, chunk_len] unsorted
+        # Sort in temporal order
+        selected_pos = topk_pos.sort(dim=-1).values  # [B, chunk_len]
+
+        # Gather
+        expanded_index = selected_pos.unsqueeze(-1).expand(-1, -1, H)  # [B, chunk_len, H]
+        action_queries = last_hidden.gather(dim=1, index=expanded_index)  # [B, chunk_len, H]
+        return action_queries
+
+    # Discretised state → instruction prefix (π₀.5 style); shared with QwenPI_v3.
+    add_discretized_state_to_instruction = staticmethod(add_discretized_state_to_instruction)
+
+
+if __name__ == "__main__":
+    import argparse
+    import os
+
+    from omegaconf import OmegaConf
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config_yaml",
+        type=str,
+        default="examples/LIBERO/train_files/starvla_cotrain_libero.yaml",
+        help="Path to YAML config",
+    )
+    args, clipargs = parser.parse_known_args()
+
+    if os.getenv("DEBUGPY_ENABLE", "0") == "1":
+        import debugpy
+
+        debugpy.listen(("0.0.0.0", 10092))
+        print("Rank 0 waiting for debugger attach on port 10092...")
+        debugpy.wait_for_client()
+
+    cfg = OmegaConf.load(args.config_yaml)
+
+    model = Qwenvl_OFT(cfg)
+    print(model)
+
+    image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
+    sample = {
+        "action": np.random.uniform(-1, 1, size=(16, 7)).astype(np.float16),
+        "image": [image],
+        "lang": "This is a fake instruction for testing.",
+        "state": np.random.uniform(-1, 1, size=(1, 7)).astype(np.float16),  # chunk, state_dim
+    }
+    sample2 = sample.copy()
+    sample2["lang"] = "Another fake instruction for testing."
+
+    batch = [sample, sample2]
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    forward_output = model(batch)
+    action_loss = forward_output["action_loss"]
+    print(f"[train] Action Loss (with state): {action_loss.item()}")
+
+    predict_output = model.predict_action(examples=[batch[0]])
+    normalized_actions = predict_output["normalized_actions"]
+    print(f"[infer] Predicted Action shape: {normalized_actions.shape}")
+
+    # Backward-compat: examples without `state` should still work.
+    sample_no_state = {k: v for k, v in sample.items() if k != "state"}
+    forward_no_state = model([sample_no_state, sample_no_state])
+    print(f"[train] Action Loss (no state): {forward_no_state['action_loss'].item()}")
+    predict_no_state = model.predict_action(examples=[sample_no_state])
+    print(f"[infer] Predicted Action shape (no state): {predict_no_state['normalized_actions'].shape}")
+
+    print("Finished")

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/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/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/training-code/gym-language-contract.patch

@@ -0,0 +1,17 @@
+diff --git a/examples/rl_games/train_files/data_registry/data_config.py b/examples/rl_games/train_files/data_registry/data_config.py
+index 2ef8b39..6d76aeb 100644
+--- a/examples/rl_games/train_files/data_registry/data_config.py
++++ b/examples/rl_games/train_files/data_registry/data_config.py
+@@ -53,10 +53,10 @@ class AtlantisDataConfig(FlappyDataConfig):
+ 
+ 
+ class GymnasiumDataConfig(FlappyDataConfig):
+-    pass
++    language_keys = ["annotation.human.task_description"]
+ 
+ 
+-class GymnasiumNativeDataConfig(FlappyDataConfig):
++class GymnasiumNativeDataConfig(GymnasiumDataConfig):
+     """Configure native Gymnasium state and action vectors."""
+ 
+     state_keys = ["state.native"]

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/training-code/loader-spawn-cache.patch

@@ -0,0 +1,18 @@
+diff --git a/starVLA/dataloader/gr00t_lerobot/datasets.py b/starVLA/dataloader/gr00t_lerobot/datasets.py
+index 17ce5f1..2346230 100644
+--- a/starVLA/dataloader/gr00t_lerobot/datasets.py
++++ b/starVLA/dataloader/gr00t_lerobot/datasets.py
+@@ -667,6 +667,13 @@ class LeRobotSingleDataset(Dataset):
+         # Check if the dataset is valid
+         self._check_integrity()
+ 
++    def __getstate__(self):
++        state = self.__dict__.copy()
++        state["_parquet_cache"] = {}
++        state["curr_traj_data"] = None
++        state["curr_traj_id"] = None
++        return state
++
+     @property
+     def dataset_path(self) -> Path:
+         """The path to the dataset that contains the METADATA_FILENAME file."""

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/training-code/loader-spawn-test.patch

@@ -0,0 +1,31 @@
+diff --git a/tests/data/test_starvla_tasks.py b/tests/data/test_starvla_tasks.py
+index 1403fa13..62468aa6 100644
+--- a/tests/data/test_starvla_tasks.py
++++ b/tests/data/test_starvla_tasks.py
+@@ -395,3 +395,26 @@ def test_balance_masked_l1_uses_equal_sample_weights_and_no_padding_gradient(mon
+     all_valid = torch.ones((2, 8), dtype=torch.bool)
+     torch.testing.assert_close(model._compute_action_loss(pred, target, action_loss_mask=all_valid),
+                                (pred - target).abs().mean())
++
++
++def test_lerobot_spawn_reload_preserves_data_without_copying_parent_cache(tmp_path, monkeypatch):
++    import pickle
++    pytest.importorskip("accelerate")
++    monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "third_party/starVLA"))
++    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotSingleDataset
++
++    path = tmp_path / "episode.parquet"
++    table = pd.DataFrame({"image": [b"rgb" * (1024 * 1024)], "action": [1]})
++    table.to_parquet(path)
++    parent = object.__new__(LeRobotSingleDataset)
++    parent._parquet_cache = {path: table}
++    parent.curr_traj_data = table
++    parent.curr_traj_id = 0
++    parent._all_steps = [(0, 0)]
++
++    payload = pickle.dumps(parent)
++    assert len(payload) < 64 * 1024
++    worker = pickle.loads(payload)
++    assert worker._all_steps == parent._all_steps
++    pd.testing.assert_frame_equal(worker._read_trajectory_parquet(path), table)
++    pd.testing.assert_frame_equal(parent.curr_traj_data, table)

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/training-code/mean-profile-preparation.patch

@@ -0,0 +1,147 @@
+diff --git a/latency_bench/data/sf_teacher_rollout.py b/latency_bench/data/sf_teacher_rollout.py
+index ef45d95b..af0a9fa9 100644
+--- a/latency_bench/data/sf_teacher_rollout.py
++++ b/latency_bench/data/sf_teacher_rollout.py
+@@ -964,6 +964,11 @@ def _replay_accepted_specs(
+         row_buffer_size=int(row_buffer_size),
+         export_image_size=export_image_size,
+     )
++    stream_current_images = (
++        image_transform_config.normalized_transform() == "raw_rgb"
++        and int(context_window) <= 1
++        and filter_preset == "none"
++    )
+     image_resolution: list[int] | None = None
+ 
+     def assign_specs_to_slots(slots: list[int]) -> dict[int, int]:
+@@ -1051,7 +1056,10 @@ def _replay_accepted_specs(
+                         next_raw_frame=next_decision_step * obs_stride_raw_frames,
+                         max_steps_per_episode=int(max_steps_per_episode),
+                     )
+-                    active_episode_frames[slot].append(np.array(image_frame, copy=True))
++                    if stream_current_images:
++                        image_writer.write(image_frame, image_path)
++                    else:
++                        active_episode_frames[slot].append(np.array(image_frame, copy=True))
+                     row = teacher_dataset_row(
+                         env_name=env_name,
+                         spec=spec,
+@@ -1104,7 +1112,9 @@ def _replay_accepted_specs(
+                             )
+                         else:
+                             shard_rows = [dict(row) for row in final_rows]
+-                        if _is_demon_ghost_trail:
++                        if stream_current_images:
++                            image_writer.flush()
++                        elif _is_demon_ghost_trail:
+                             # Frames are pre-built ghost trail images; map directly.
+                             ep_steps = [int(row["decision_step"]) for row in final_rows]
+                             frames_by_decision_step = {
+@@ -1118,12 +1128,13 @@ def _replay_accepted_specs(
+                                 frames=active_episode_frames[slot],
+                                 config=image_transform_config,
+                             )
+-                        shard_rows = _write_filtered_episode_images(
+-                            image_writer,
+-                            shard_rows,
+-                            frames_by_decision_step,
+-                            context_window=int(context_window),
+-                        )
++                        if not stream_current_images:
++                            shard_rows = _write_filtered_episode_images(
++                                image_writer,
++                                shard_rows,
++                                frames_by_decision_step,
++                                context_window=int(context_window),
++                            )
+                         write_episode_row_shard(output_dir, shard_rows)
+                         progress.update(1)
+                         replayed += 1
+diff --git a/tests/data/test_sf_teacher_dataset_export.py b/tests/data/test_sf_teacher_dataset_export.py
+index e6b25233..7d5537ac 100644
+--- a/tests/data/test_sf_teacher_dataset_export.py
++++ b/tests/data/test_sf_teacher_dataset_export.py
+@@ -1273,7 +1273,8 @@ def test_collect_teacher_dataset_saves_demon_attack_latest_raw_rgb_image(monkeyp
+ 
+ 
+ @pytest.mark.parametrize("export_history", [False, True])
+-def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch, tmp_path, export_history):
++@pytest.mark.parametrize("filter_preset", ["clean_v1", "none"])
++def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch, tmp_path, export_history, filter_preset):
+     import torch
+     from gymnasium.spaces import Discrete
+     from PIL import Image
+@@ -1307,6 +1308,9 @@ def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch,
+     class FakeEnv:
+         action_space = Discrete(2)
+ 
++        def __init__(self, exporting):
++            self.exporting = exporting
++
+         def reset(self, seed=None):
+             self.steps = 0
+             obs = {"obs": np.zeros((1, 12, 2, 2), dtype=np.uint8)}
+@@ -1323,6 +1327,9 @@ def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch,
+ 
+         def step(self, action):
+             self.steps += 1
++            if self.exporting and self.steps == 2 and filter_preset == "none":
++                image_path = output_dir / "images/train/episode_000000/decision_step_000000.png"
++                assert image_path.is_file(), "First frame must be on disk before the episode ends"
+             obs = {"obs": np.ones((1, 12, 2, 2), dtype=np.uint8)}
+             next_stack = np.stack([np.full((6, 5, 3), 50 + idx, dtype=np.uint8) for idx in range(4)], axis=0)
+             return (
+@@ -1354,7 +1361,7 @@ def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch,
+         exporter,
+         "build_env_and_policy",
+         lambda module, cfg, config, num_envs: {
+-            "envs": [FakeEnv()],
++            "envs": [FakeEnv(cfg.export_env_raw_rgb_frames)],
+             "env_info": object(),
+             "actor_critic": FakeActorCritic(),
+             "device": torch.device("cpu"),
+@@ -1382,6 +1389,7 @@ def test_collect_teacher_dataset_saves_flappy_latest_raw_rgb_image(monkeypatch,
+         export_image_size=None,
+         image_writer_workers=0,
+         row_buffer_size=1,
++        filter_preset=filter_preset,
+     )
+     train_rows = _read_shard_rows(output_dir, "train")
+     val_rows = _read_shard_rows(output_dir, "val")
+diff --git a/third_party/MIKASA-Robo b/third_party/MIKASA-Robo
+--- a/third_party/MIKASA-Robo
++++ b/third_party/MIKASA-Robo
+@@ -1 +1 @@
+-Subproject commit 017490bbb787c6efa8c025f9558e1a691965c0e0
++Subproject commit 017490bbb787c6efa8c025f9558e1a691965c0e0-dirty
+diff --git a/third_party/starVLA b/third_party/starVLA
+--- a/third_party/starVLA
++++ b/third_party/starVLA
+@@ -1 +1 @@
+-Subproject commit ccca13c5177fe3d3c884b6e2de4965d916016649
++Subproject commit ccca13c5177fe3d3c884b6e2de4965d916016649-dirty
+diff --git a/training/gymnasium_sf/integration.py b/training/gymnasium_sf/integration.py
+index 61c65ae9..874f655f 100644
+--- a/training/gymnasium_sf/integration.py
++++ b/training/gymnasium_sf/integration.py
+@@ -64,7 +64,6 @@ from training.common.sf_raw_reward_stats_compat import (
+ )
+ from training.common.utils import env_instance_seed, seed_everything, str2bool
+ from training.common.wandb_compat import patch_sample_factory_wandb
+-from sample_factory.algo.mc_ppo.env import EpisodeEnv
+ 
+ 
+ SF_ENV_NAME = "latency_gymnasium"
+@@ -385,7 +384,11 @@ def make_gymnasium_sf_env(
+         env_id=int(env_config.env_id),
+         issued_action_labels=action_labels,
+     )
+-    return EpisodeEnv(env, cfg.ppo) if cfg.algo == "PPO" else env
++    if cfg.algo == "PPO":
++        from sample_factory.algo.mc_ppo.env import EpisodeEnv
++
++        return EpisodeEnv(env, cfg.ppo)
++    return env
+ 
+ 
+ def register_gymnasium_envs() -> None:

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/training-code/trainer_tools.py

@@ -0,0 +1,736 @@
+"""
+metrics.py
+
+Utility classes defining a Metrics container and multiple Trackers to enable model/stage-specific logging to various
+endpoints (e.g., JSONL local logs, Weights & Biases).
+"""
+
+from typing import Tuple
+import re
+import json
+import numpy as np
+import torch
+import torch.distributed as dist
+from transformers import get_scheduler
+
+from accelerate import Accelerator, DeepSpeedPlugin
+from accelerate.logging import get_logger
+from starVLA.training.trainer_utils.config_tracker import reload_module_paths
+
+logger = get_logger(__name__)
+
+
+def build_accelerator(cfg, *, use_deepspeed: bool) -> Accelerator:
+    accelerator_kwargs = {
+        "gradient_accumulation_steps": cfg.trainer.gradient_accumulation_steps,
+    }
+    if use_deepspeed:
+        accelerator_kwargs["deepspeed_plugin"] = DeepSpeedPlugin()
+
+    accelerator = Accelerator(**accelerator_kwargs)
+    accelerator.print(accelerator.state)
+    return accelerator
+
+
+_VOCAB_SIZED_WEIGHT_SUFFIXES = (
+    "embed_tokens.weight",
+    "lm_head.weight",
+)
+
+
+def _adapt_vocab_sized_weights(module, state_dict):
+    """Copy overlapping vocab rows when tokenizer-sized Qwen weights changed."""
+    target_state_dict = module.state_dict()
+    adapted = dict(state_dict)
+    adapted_keys = []
+    for key, source_tensor in state_dict.items():
+        target_tensor = target_state_dict.get(key)
+        if target_tensor is None or not hasattr(source_tensor, "shape") or not hasattr(target_tensor, "shape"):
+            continue
+        if tuple(source_tensor.shape) == tuple(target_tensor.shape):
+            continue
+        if not key.endswith(_VOCAB_SIZED_WEIGHT_SUFFIXES):
+            continue
+        if source_tensor.ndim != 2 or target_tensor.ndim != 2:
+            continue
+        if source_tensor.shape[1] != target_tensor.shape[1]:
+            continue
+
+        patched_tensor = target_tensor.detach().clone()
+        rows = min(source_tensor.shape[0], target_tensor.shape[0])
+        patched_tensor[:rows].copy_(source_tensor[:rows].to(dtype=patched_tensor.dtype, device=patched_tensor.device))
+        adapted[key] = patched_tensor
+        adapted_keys.append((key, tuple(source_tensor.shape), tuple(target_tensor.shape), rows))
+    return adapted, adapted_keys
+
+
+# === Define Tracker Interface ===
+#
+
+# utils/cli_parser.py
+
+
+def normalize_dotlist_args(args):
+    """
+    Convert ['--x.y', 'val'] and ['--flag'] → ['x.y=val', 'flag=true']
+    """
+    normalized = []
+    skip = False
+    for i in range(len(args)):
+        if skip:
+            skip = False
+            continue
+
+        arg = args[i]
+        if arg.startswith("--"):
+            key = arg.lstrip("-")
+            if "=" in key:
+                normalized.append(key)
+            elif i + 1 < len(args) and not args[i + 1].startswith("--"):
+                normalized.append(f"{key}={args[i + 1]}")
+                skip = True
+            else:
+                normalized.append(f"{key}=true")
+        else:
+            pass  # skip orphaned values
+    return normalized
+
+
+def setup_optimizer_and_scheduler(model, cfg) -> Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler]:
+    """Build AdamW optimizer + LR scheduler from cfg.trainer.
+
+    Supports:
+    - Per-module learning rates via cfg.trainer.learning_rate
+    - cosine_with_min_lr (or any transformers scheduler) via cfg.trainer.lr_scheduler_type
+    """
+    param_groups = build_param_lr_groups(model=model, cfg=cfg)
+    fused_available = torch.cuda.is_available()
+    optimizer = torch.optim.AdamW(
+        param_groups,
+        lr=cfg.trainer.learning_rate.base,
+        betas=tuple(cfg.trainer.optimizer.betas),
+        weight_decay=cfg.trainer.optimizer.weight_decay,
+        eps=cfg.trainer.optimizer.eps,
+        fused=fused_available,
+    )
+
+    if dist.is_initialized() and dist.get_rank() == 0:
+        for group in optimizer.param_groups:
+            logger.info(f"LR Group {group['name']}: lr={group['lr']}, num_params={len(group['params'])}")
+
+    # Strip keys unknown to transformers' get_scheduler before passing kwargs.
+    sched_kwargs = cfg.trainer.scheduler_specific_kwargs
+    lr_scheduler = get_scheduler(
+        name=cfg.trainer.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=cfg.trainer.num_warmup_steps,
+        num_training_steps=cfg.trainer.max_train_steps,
+        scheduler_specific_kwargs=sched_kwargs,
+    )
+
+    return optimizer, lr_scheduler
+
+
+def _qwen_vl_interface(model):
+    return model.qwen_vl_interface
+
+
+def _llm_language_model(model):
+    return _qwen_vl_interface(model).model.model.language_model
+
+
+def _vit_module(model):
+    # Qwen-VL keeps the visual tower and its vision-to-language connector/merger
+    # under this module.
+    return _qwen_vl_interface(model).model.model.visual
+
+
+def _llm_decoder_layers(model):
+    return _llm_language_model(model).layers
+
+
+def _tied_embedding_modules(model):
+    language_model = _llm_language_model(model)
+    hf_model = _qwen_vl_interface(model).model
+    return language_model.embed_tokens, hf_model.lm_head
+
+
+def resolve_freeze_llm_layer_indices(model, cfg):
+    """Translate ``trainer.freeze_llm_layers=[start,end]`` to an inclusive layer list."""
+    freeze_llm_layers = cfg.trainer.freeze_llm_layers
+    if not freeze_llm_layers:
+        return [], len(_llm_decoder_layers(model))
+    start, end = [int(value) for value in freeze_llm_layers]
+    return list(range(start, end + 1)), len(_llm_decoder_layers(model))
+
+
+def _module_param_ids(module):
+    return {id(param) for param in module.parameters()}
+
+
+def _has_vit_or_llm_freeze_request(cfg) -> bool:
+    return bool(cfg.trainer.freeze_vit or cfg.trainer.freeze_tied_embedding or cfg.trainer.freeze_llm_layers)
+
+
+def vit_llm_frozen_param_ids(model, cfg):
+    """Parameter ids excluded from optimizer by trainer freeze controls."""
+    if not _has_vit_or_llm_freeze_request(cfg):
+        return set()
+
+    frozen_params = set()
+    freeze_vit = cfg.trainer.freeze_vit
+    if freeze_vit:
+        frozen_params.update(_module_param_ids(_vit_module(model)))
+
+    if cfg.trainer.freeze_tied_embedding:
+        for module in _tied_embedding_modules(model):
+            frozen_params.update(_module_param_ids(module))
+
+    frozen_layer_indices, _ = resolve_freeze_llm_layer_indices(model, cfg)
+    layers = list(_llm_decoder_layers(model))
+    for layer_index in frozen_layer_indices:
+        frozen_params.update(_module_param_ids(layers[layer_index]))
+
+    return frozen_params
+
+
+def build_param_lr_groups(model, cfg):
+    """
+    build multiple param groups based on cfg.trainer.learning_rate.
+    support specifying different learning rates for different modules, the rest use base.
+
+    Args:
+        vla: nn.Module model object
+        cfg: config object, requires cfg.trainer.learning_rate dictionary
+
+    Returns:
+        List[Dict]: param_groups that can be used to build optimizer with torch.optim
+    """
+
+    lr_cfg = cfg.trainer.learning_rate
+    base_lr = lr_cfg.get("base", 1e-4)  # default base learning rate
+
+    freeze_modules = cfg.trainer.get("freeze_modules", "")
+    if not isinstance(freeze_modules, str):
+        freeze_modules = ""
+    freeze_patterns = [p.strip() for p in freeze_modules.split(",") if p.strip()]
+
+    used_params = set()
+    frozen_params = set()
+    param_groups = []
+
+    for freeze_path in freeze_patterns:
+        module = model
+        try:
+            for attr in freeze_path.split("."):
+                module = getattr(module, attr)
+            frozen_params.update(id(p) for p in module.parameters())
+        except AttributeError:
+            print(f"⚠️ freeze module path does not exist: {freeze_path}")
+            continue
+
+    # Exclude trainer freeze controls so frozen parameters never enter an
+    # optimizer param group.
+    frozen_params.update(vit_llm_frozen_param_ids(model, cfg))
+
+    for module_name, lr in lr_cfg.items():
+        if module_name == "base":
+            continue
+        # try to find the module under vla by module_name (support nested paths)
+        module = model
+        try:
+            for attr in module_name.split("."):
+                module = getattr(module, attr)
+            # filter out frozen parameters
+            params = [p for p in module.parameters() if id(p) not in frozen_params]
+            if params:  # only add param group if there are trainable parameters
+                param_groups.append({"params": params, "lr": lr, "name": module_name})
+                used_params.update(id(p) for p in params)
+        except AttributeError:
+            ReferenceError(f"⚠️ module path `{module_name}` not found in vla")
+
+    # assign base learning rate to the remaining unused parameters (exclude frozen ones)
+    other_params = [p for p in model.parameters() if id(p) not in used_params and id(p) not in frozen_params]
+    if other_params:
+        param_groups.append({"params": other_params, "lr": base_lr, "name": "base"})
+
+    return param_groups
+
+
+import torch.distributed as dist
+
+
+def only_main_process(func):
+    """
+    decorator: only run in main process (rank=0)
+    """
+
+    def wrapper(*args, **kwargs):
+        if dist.is_initialized() and dist.get_rank() != 0:
+            return None  # non-main process does not execute
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+from torchvision.ops import box_iou
+from PIL import Image
+
+
+def resize_images(images, target_size=(224, 224)):
+    """
+    recursively resize all images in the nested list.
+
+    :param images: nested list of images or single image.
+    :param target_size: target size (width, height) after resizing.
+    :return: resized images list, keeping the original nested structure.
+    """
+    if isinstance(images, Image.Image):  # if it is a single PIL image
+        return images.resize(target_size)
+    elif isinstance(images, list):  # if it is a list, recursively process each element
+        return [resize_images(img, target_size) for img in images]
+    else:
+        raise ValueError("Unsupported image type or structure.")
+
+
+def stitch_frames(frames, grid=(2, 2), size=(224, 224)):
+    """Tile temporal frames into one image, keeping the token count of a single image.
+
+    Used by the ``stitch`` multi-image mode: instead of feeding N frames as N
+    tokenized images, they are laid out on a fixed grid and the whole canvas is
+    returned at ``size``. Both training (``_pack_sample``) and eval
+    (``policy/starvla.py``) call this so the layouts match.
+
+    Args:
+        frames: list of ``PIL.Image`` or ``np.ndarray`` (H, W, C), oldest first.
+        grid: (rows, cols). Cells beyond ``len(frames)`` repeat the last frame so
+            the layout is fixed regardless of how many frames are available.
+        size: (width, height) of the returned canvas.
+    """
+    rows, cols = grid
+    canvas_w, canvas_h = size
+    cell_w, cell_h = canvas_w // cols, canvas_h // rows
+    pil = [f if isinstance(f, Image.Image) else Image.fromarray(f) for f in frames]
+    capacity = rows * cols
+    pil = (pil + [pil[-1]] * (capacity - len(pil))) if len(pil) < capacity else pil[:capacity]
+    canvas = Image.new("RGB", (canvas_w, canvas_h))
+    for idx, frame in enumerate(pil):
+        r, c = divmod(idx, cols)
+        canvas.paste(frame.convert("RGB").resize((cell_w, cell_h)), (c * cell_w, r * cell_h))
+    return canvas
+
+
+import torch.distributed as dist
+
+
+class TrainerUtils:
+    @staticmethod
+    def freeze_backbones(model, freeze_modules=""):
+        """
+        directly freeze the specified submodules based on the relative module path list (patterns), no longer recursively find all submodule names:
+          - patterns: read from config.trainer.freeze_modules, separated by commas to get the "relative path" list
+            for example "qwen_vl_interface, action_model.net",
+            it means to freeze model.qwen_vl_interface and model.action_model.net.
+
+        Args:
+            model: nn.Module model object
+            freeze_modules: relative module path list (patterns)
+
+        Returns:
+            model: nn.Module model object
+        return:
+          - model:
+        """
+        frozen = []
+        print("#"*30)
+        print(freeze_modules)
+        if freeze_modules and type(freeze_modules) == str:
+            # split and remove whitespace
+            patterns = [p.strip() for p in freeze_modules.split(",") if p.strip()] if freeze_modules else []
+
+            for path in patterns:
+                # split the "relative path" by dots, for example "action_model.net" → ["action_model", "net"]
+                attrs = path.split(".")
+                module = model
+                try:
+                    for attr in attrs:
+                        module = getattr(module, attr)
+                    # if the module is successfully get, freeze it and its all submodule parameters
+                    for param in module.parameters():
+                        param.requires_grad = False
+                    frozen.append(path)
+                except AttributeError:
+                    # if the attribute does not exist, skip and print warning
+                    print(f"⚠️ module path does not exist, cannot freeze: {path}")
+                    continue
+
+        # accelerator.wait_for_everyone()  # synchronize when distributed training
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(f"🔒 Frozen modules with re pattern: {frozen}")
+        return model
+
+    @staticmethod
+    def freeze_vit_and_llm_layers(model, cfg):
+        """Apply trainer.freeze_vit/freeze_tied_embedding/freeze_llm_layers.
+
+        trainer.freeze_vit freezes Qwen-VL's visual module, including the
+        vision-to-language connector/merger it contains.
+        """
+        if not _has_vit_or_llm_freeze_request(cfg):
+            return model
+
+        freeze_vit = cfg.trainer.freeze_vit
+        if freeze_vit:
+            for param in _vit_module(model).parameters():
+                param.requires_grad = False
+
+        freeze_tied_embedding = cfg.trainer.freeze_tied_embedding
+        if freeze_tied_embedding:
+            for module in _tied_embedding_modules(model):
+                for param in module.parameters():
+                    param.requires_grad = False
+
+        frozen_layer_indices, total = resolve_freeze_llm_layer_indices(model, cfg)
+        layers = list(_llm_decoder_layers(model))
+        for layer_index in frozen_layer_indices:
+            for param in layers[layer_index].parameters():
+                param.requires_grad = False
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(
+                "🔒 freeze_vit="
+                f"{freeze_vit}, freeze_tied_embedding={freeze_tied_embedding}, "
+                f"freeze_llm_layers={list(frozen_layer_indices)}"
+                f" / total={total}"
+            )
+        return model
+
+    @staticmethod
+    def print_trainable_parameters(model):
+        """
+        print the total number of parameters and trainable parameters of the model
+        :param model: PyTorch model instance
+        """
+        if dist.is_initialized() and dist.get_rank() != 0:
+            return
+        print("📊 model parameter statistics:")
+        num_params = sum(p.numel() for p in model.parameters())
+        num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        print(
+            f"# Parameters (in millions): {num_params / 10**6:.3f} Total, {num_trainable_params / 10**6:.3f} Trainable"
+        )
+        return num_params, num_trainable_params
+
+    @staticmethod
+    def load_pretrained_backbones(model, checkpoint_path=None, reload_modules=None):
+        """
+        load checkpoint:
+        - if reload_modules is set, load by path part
+        - otherwise → load the entire model parameters (overwrite model)
+
+        return:
+            replace, loaded_modules: list of module paths that successfully loaded parameters; if global load, then ["<full_model>"]
+        """
+        if not checkpoint_path:
+            return []
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(f"📦 loading checkpoint: {checkpoint_path}")
+        try:
+            if _is_safetensors_path(checkpoint_path):
+                from safetensors.torch import load_file
+
+                checkpoint = load_file(checkpoint_path)
+            else:
+                checkpoint = torch.load(checkpoint_path, map_location="cpu")
+        except Exception as e:
+            raise RuntimeError(f"❌ loading checkpoint failed: {e}")
+
+        loaded_modules = []
+
+        module_paths = reload_module_paths(reload_modules)
+        if module_paths:  # partial load
+            for path in module_paths:
+                reload_path = path.split(".")
+                module = model
+                try:
+                    for module_name in reload_path:  # find the module to modify level by level
+                        module = getattr(module, module_name)
+                    prefix = path + "."
+                    sub_state_dict = {k[len(prefix) :]: v for k, v in checkpoint.items() if k.startswith(prefix)}
+                    if sub_state_dict:
+                        sub_state_dict, adapted_keys = _adapt_vocab_sized_weights(module, sub_state_dict)
+                        for key, source_shape, target_shape, rows in adapted_keys:
+                            if not dist.is_initialized() or dist.get_rank() == 0:
+                                print(
+                                    f"⚠️ adapted vocab-sized weight '{path}.{key}' from checkpoint shape "
+                                    f"{source_shape} to current shape {target_shape}; copied {rows} rows"
+                                )
+                        module.load_state_dict(sub_state_dict, strict=True)
+                        if not dist.is_initialized() or dist.get_rank() == 0:
+                            print(f"✅ parameters loaded to module '{path}'")
+                        loaded_modules.append(path)
+                    else:
+                        print(f"⚠️ parameters not found in checkpoint '{path}'")
+                except AttributeError:
+                    print(f"❌ cannot find module path: {path}")
+        else:  # full load
+            try:
+                checkpoint, adapted_keys = _adapt_vocab_sized_weights(model, checkpoint)
+                for key, source_shape, target_shape, rows in adapted_keys:
+                    if not dist.is_initialized() or dist.get_rank() == 0:
+                        print(
+                            f"⚠️ adapted vocab-sized weight '{key}' from checkpoint shape "
+                            f"{source_shape} to current shape {target_shape}; copied {rows} rows"
+                        )
+                model.load_state_dict(checkpoint, strict=False)
+                if not dist.is_initialized() or dist.get_rank() == 0:
+                    print("✅ loaded <full_model> model parameters")
+                loaded_modules = ["<full_model>"]
+            except Exception as e:
+                raise RuntimeError(f"❌ loading full model failed: {e}")
+        return model
+
+    @staticmethod
+    def print_freeze_status(model):
+        """
+        print the freezing status of each parameter in the model
+        :param model: PyTorch model instance
+        """
+        for name, param in model.named_parameters():
+            status = "Frozen" if not param.requires_grad else "Trainable"
+            print(f"{name:60s}  |  {status}")
+
+    @staticmethod
+    def setup_distributed_training(accelerator, *components):
+        """
+        use Accelerator to prepare distributed training components
+        :param accelerator: Accelerate instance
+        :param components: any number of components (such as model, optimizer, dataloader, etc.)
+        :return: prepared distributed components (in the same order as input)
+        """
+
+        # use accelerator.prepare method to wrap components
+        prepared_components = accelerator.prepare(*components)
+        return prepared_components
+
+    @staticmethod
+    def euclidean_distance(predicted: np.ndarray, ground_truth: np.ndarray) -> float:
+        return np.linalg.norm(predicted - ground_truth)
+
+    @staticmethod
+    def _reset_dataloader(dataloader, epoch_counter):
+        """safe reset dataloader iterator"""
+        # 1. update epoch counter
+        epoch_counter += 1
+
+        dataset = getattr(dataloader, "dataset", None)
+        if dataset is not None and callable(getattr(dataset, "set_epoch", None)):
+            dataset.set_epoch(epoch_counter)
+
+        # 2. set new epoch (distributed core)
+        if hasattr(dataloader, "sampler") and callable(getattr(dataloader.sampler, "set_epoch", None)):
+            dataloader.sampler.set_epoch(epoch_counter)
+
+        # 3. create new iterator
+        return iter(dataloader), epoch_counter
+
+    @staticmethod
+    def compute_grad_angle_with_stats(grads_a: list[torch.Tensor], grads_v: list[torch.Tensor]) -> Tuple[float, float]:
+        """
+        compute the cosine angle between two groups of gradient vectors (degrees), and calculate the average angle and variance.
+        grads_a, grads_v: gradient Tensor list corresponding to the same parameter list interface_params
+        return:
+            mean_angle_deg: average angle (degrees)
+            angle_variance: angle variance
+        """
+        angle_degs = []
+
+        # compute the cosine angle between each gradient block grads_a[0].shape = 1280, 3, 14, 14
+        # grads_1 = grads_a[0][0]  # [3, 14, 14]
+        # grads_2 = grads_v[0][0]
+        # grads_a = grads_1.view(-1, 3)  # reshape to [196, 3]
+        # grads_v = grads_2.view(-1, 3)
+
+        # lang linear
+        # reshape to 14*14, 3
+        # layer
+        grads_action = grads_a[0]  # [2048, 11008]
+        grads_action = grads_action[
+            :32, :7
+        ]  # only take the first 7 elements, avoid cosim failure in high-dimensional space
+        grads_vl = grads_v[0]  # [2048, 11008]
+        grads_vl = grads_vl[
+            :32, :7
+        ]  # only take the first 32 elements, 7 dimensions, avoid cosim failure in high-dimensional space
+        for g_a, g_v in zip(grads_action, grads_vl):
+            dot = torch.sum(g_a * g_v)
+            norm_a_sq = torch.sum(g_a * g_a)
+            norm_v_sq = torch.sum(g_v * g_v)
+
+            # avoid division by zero
+            norm_a = torch.sqrt(norm_a_sq + 1e-16)
+            norm_v = torch.sqrt(norm_v_sq + 1e-16)
+
+            cos_sim = (dot / (norm_a * norm_v)).clamp(-1.0, 1.0)
+            angle_rad = torch.acos(cos_sim)
+            angle_deg = angle_rad * (180.0 / torch.pi)
+
+            angle_degs.append(angle_deg.item())
+
+        # compute the average angle and variance
+        angle_degs_tensor = torch.tensor(angle_degs)
+        mean_angle_deg = torch.mean(angle_degs_tensor).item()
+        angle_variance = torch.sqrt(torch.var(angle_degs_tensor)).item()
+        # accelerator.wait_for_everyone()
+        return mean_angle_deg, angle_variance
+
+    @staticmethod
+    def pcgrad_project(grads_a: list[torch.Tensor], grads_v: list[torch.Tensor]) -> list[torch.Tensor]:
+        """
+        apply PCGrad projection to the second group of gradients grads_v, suppress negative transfer between grads_a and grads_v
+        if the dot product of two groups of gradients < 0, then:
+            grads_v <- grads_v - (dot / ||grads_a||^2) * grads_a
+        return the new grads_v list
+        """
+        # first compute dot and ||grads_a||^2
+        dot, norm_a_sq = 0.0, 0.0
+        for g_a, g_v in zip(grads_a, grads_v):
+            dot += torch.sum(g_a * g_v)
+            norm_a_sq += torch.sum(g_a * g_a)
+
+        if dot < 0:
+            coeff = dot / (norm_a_sq + 1e-6)
+            # projection
+            grads_v = [g_v - coeff * g_a for g_a, g_v in zip(grads_a, grads_v)]
+
+        return grads_v
+
+    @staticmethod
+    def eval_qwenpi(qwenpi, dataloader, num_batches=20):
+        """
+        evaluate QwenQFormerDiT model, compute IoU and action distance.
+
+        Args:
+            qwenpi: QwenQFormerDiT model instance.
+            dataloader: data loader.
+            num_batches: number of batches to evaluate.
+
+        Returns:
+            dict: contains IoU and action distance evaluation results.
+        """
+        iou_scores = []
+        action_distances = []
+        count = 0
+
+        dataset_iter = iter(dataloader)
+        while count < num_batches:
+            try:
+                batch_samples = next(dataset_iter)
+                count += 1
+            except StopIteration:
+                break
+
+            # extract data
+            images = [example["image"] for example in batch_samples]
+            instructions = [example["lang"] for example in batch_samples]
+            actions = [example["action"] for example in batch_samples]
+            solutions = [example["solution"] for example in batch_samples]
+
+            # model prediction
+            predicted_solutions, normalized_actions = qwenpi.predict_action_withCoT(
+                images=images, instructions=instructions, use_ddim=False, num_ddim_steps=20
+            )
+
+            # extract and convert predicted results
+            parsed_solutions = []
+            for solution in predicted_solutions:
+                parsed_solution = TrainerUtils.extract_json_from_string(solution)
+                parsed_solutions.append(parsed_solution)
+
+            # compute IoU
+            for pred_dict, gt_dict in zip(parsed_solutions, solutions):
+                pred_pick_bbox = torch.tensor(pred_dict["pick"]["bbox_2d"], dtype=torch.float32).unsqueeze(0)
+                gt_pick_bbox = torch.tensor(gt_dict["pick"]["bbox_2d"], dtype=torch.float32).unsqueeze(0)
+                pred_place_bbox = torch.tensor(pred_dict["place"]["bbox_2d"], dtype=torch.float32).unsqueeze(0)
+                gt_place_bbox = torch.tensor(gt_dict["place"]["bbox_2d"], dtype=torch.float32).unsqueeze(0)
+
+                pick_iou = box_iou(pred_pick_bbox, gt_pick_bbox).item()
+                place_iou = box_iou(pred_place_bbox, gt_place_bbox).item()
+
+                iou_scores.append({"pick_iou": pick_iou, "place_iou": place_iou})
+
+            # compute action distance
+            actions = np.array(actions)  # convert to numpy array
+            num_pots = np.prod(actions.shape)  # B*len*dim
+            action_distance = TrainerUtils.euclidean_distance(normalized_actions, actions)
+            average_action_distance = action_distance / num_pots
+            action_distances.append(average_action_distance)
+
+        # summarize results
+        avg_action_distance = np.mean(action_distances)
+        return {"iou_scores": iou_scores, "average_action_distance": avg_action_distance}
+
+    @staticmethod
+    def extract_json_from_string(input_string):
+        """
+        extract valid JSON part from string and convert to dictionary.
+
+        Args:
+            input_string (str): string containing extra characters.
+
+        Returns:
+            dict: dictionary extracted and parsed.
+        """
+        json_match = re.search(r"{.*}", input_string, re.DOTALL)
+        if json_match:
+            json_str = json_match.group(0)
+            try:
+                return json.loads(json_str)
+            except json.JSONDecodeError as e:
+                print(f"JSON decode failed: {e}")
+                return None
+        else:
+            print("No valid JSON part found")
+            return None
+
+    def _get_latest_checkpoint(self, checkpoint_dir):
+        """Find the latest checkpoint in the directory based on step number."""
+        if not os.path.exists(checkpoint_dir):
+            self.accelerator.print(f"No checkpoint directory found at {checkpoint_dir}")
+            return None, 0
+
+        checkpoints = []
+        for name in os.listdir(checkpoint_dir):
+            path = os.path.join(checkpoint_dir, name)
+            state_match = re.match(r"steps_(\d+)_state$", name)
+            if state_match and os.path.isdir(path):
+                checkpoints.append((name, int(state_match.group(1)), 1))
+                continue
+            file_match = re.match(r"steps_(\d+)_(?:pytorch_model\.pt|model\.safetensors)$", name)
+            if file_match and os.path.isfile(path):
+                checkpoints.append((name, int(file_match.group(1)), 0))
+
+        if not checkpoints:
+            self.accelerator.print(f"No checkpoints found in {checkpoint_dir}")
+            return None, 0
+
+        # Sort by step number, preferring full training-state directories over
+        # model-only files when both exist for the same step.
+        checkpoints.sort(key=lambda x: (x[1], x[2]))
+        latest_checkpoint, completed_steps, _ = checkpoints[-1]
+
+        latest_checkpoint_path = os.path.join(checkpoint_dir, latest_checkpoint)
+        self.accelerator.print(f"Latest checkpoint found: {latest_checkpoint_path}")
+        return latest_checkpoint_path, completed_steps
+
+import os
+
+
+def is_main_process():
+    rank = int(os.environ.get("RANK", 0))  # if RANK is not set, default to 0
+    return rank == 0
+
+
+def _is_safetensors_path(path):
+    """Check if a path refers to a safetensors file."""
+    return str(path).endswith(".safetensors")

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/training-code/worker_context.py

@@ -0,0 +1,77 @@
+import multiprocessing.context
+import os
+from typing import Any
+
+
+_DISTRIBUTED_ENV_KEYS = (
+    "LOCAL_RANK",
+    "RANK",
+    "WORLD_SIZE",
+    "MASTER_ADDR",
+    "MASTER_PORT",
+    "GROUP_RANK",
+    "LOCAL_WORLD_SIZE",
+    "ROLE_RANK",
+    "ROLE_WORLD_SIZE",
+    "TORCHELASTIC_RUN_ID",
+    "TORCHELASTIC_RESTART_COUNT",
+    "TORCHELASTIC_MAX_RESTARTS",
+)
+
+
+class _CpuOnlyWorkerSpawnProcess(multiprocessing.context.SpawnProcess):
+    def start(self):
+        had_cuda_visible_devices = "CUDA_VISIBLE_DEVICES" in os.environ
+        if had_cuda_visible_devices:
+            cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        distributed_env = {key: os.environ[key] for key in _DISTRIBUTED_ENV_KEYS if key in os.environ}
+
+        os.environ["CUDA_VISIBLE_DEVICES"] = ""
+        for key in distributed_env:
+            del os.environ[key]
+        try:
+            return super().start()
+        finally:
+            if had_cuda_visible_devices:
+                os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
+            else:
+                del os.environ["CUDA_VISIBLE_DEVICES"]
+            for key, value in distributed_env.items():
+                os.environ[key] = value
+
+
+class CpuOnlyWorkerSpawnContext(multiprocessing.context.SpawnContext):
+    Process = _CpuOnlyWorkerSpawnProcess
+
+
+CPU_ONLY_WORKER_CONTEXT = CpuOnlyWorkerSpawnContext()
+
+
+def _as_bool(value: Any, default: bool = False) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        return value.strip().lower() not in {"false", "0", "no", "off"}
+    return bool(value)
+
+
+def build_cpu_only_dataloader_kwargs(
+    num_workers: int,
+    *,
+    pin_memory: Any = False,
+    persistent_workers: Any = False,
+    prefetch_factor: Any = None,
+) -> dict[str, Any]:
+    dataloader_kwargs: dict[str, Any] = {
+        "pin_memory": _as_bool(pin_memory),
+    }
+    if int(num_workers) <= 0:
+        return dataloader_kwargs
+
+    dataloader_kwargs["multiprocessing_context"] = CPU_ONLY_WORKER_CONTEXT
+    dataloader_kwargs["persistent_workers"] = _as_bool(persistent_workers)
+    if prefetch_factor is not None:
+        dataloader_kwargs["prefetch_factor"] = int(prefetch_factor)
+    return dataloader_kwargs

latency-aware/ant/vla/starvla-qwenoft-h1/ant-qwenoft-mean-3090-s42/training_results.json

@@ -0,0 +1,46 @@
+{
+  "training_complete": true,
+  "optimizer_updates": 5000,
+  "selection": "final_5000",
+  "train_loss": 0.011952828615903854,
+  "validation_action_loss": 0.018865827471017838,
+  "validation_samples": 10112,
+  "final_learning_rates": {
+    "learning_rate/action_model": 4.9999999999999996e-06,
+    "learning_rate/qwen_vl_interface": 5e-07
+  },
+  "closed_loop_evaluation": "pending until all four VLA trainings complete",
+  "tensor_keys": 730,
+  "serialized_tensor_elements_including_tied_aliases": 4892400136,
+  "weights_only_load": "passed",
+  "all_tensors_finite": true,
+  "completion_proof": {
+    "task": "ant",
+    "optimizer_updates": 5000,
+    "deepspeed_global_steps": 5000,
+    "deepspeed_global_samples": 640000,
+    "skipped_steps": 0,
+    "scheduler_last_epoch": 5000,
+    "scheduler_step_count": 5001,
+    "final_learning_rates": [
+      5e-07,
+      4.9999999999999996e-06
+    ],
+    "summary_final_steps": 5000,
+    "final_inference_weight_bytes": 9785070315,
+    "wandb_global_step": 5000,
+    "wrapper_exit": 2,
+    "trainer_exit_separately_captured": false,
+    "training_complete_verified": true,
+    "failure_scope": "launcher_tail_parse_after_trainer_cleanup",
+    "cause": "Sharedrun_vla.shwasmodifiedforGPUmask/portwhileAntbashwasblockedinsidePython; remainingfileoffsetlaterparsedmodifiedtailandEOF",
+    "trainer_cleanup_log": "bothranksDestroyDistributedProcessGroup at2026-10-01 00:03:21UTC; no model exception",
+    "evidence_files": [
+      "checkpoints/steps_5000_state/pytorch_model/mp_rank_00_model_states.pt",
+      "checkpoints/steps_5000_state/custom_checkpoint_0.pkl",
+      "summary.jsonl",
+      "config.full.yaml"
+    ],
+    "note": "No syntheticvla.exit0; wrapperexit2preserved; upload uses explicitcompletionproof"
+  }
+}