Publish completed Flappy QwenOFT H1 mean-profile 5000-update Blackwell run

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

@@ -0,0 +1,12 @@
+# Flappy QwenOFT H1 with RTX3090 mean latency
+
+Training complete: 5000 optimizer updates, seed42, global batch128 (2 Blackwell GPUs x64, accumulation1), BF16/ZeRO2, gradient checkpointing false. Final5000 checkpoint is published; full optimizer recovery is retained on the training host.
+
+The data use fixed75.87417450998383ms from the archived RTX3090 OpenVLA profile, quantized to one raw frame at10Hz (100ms effective). This is mean profiling. Binary causal eight-decision history24D, current RGB image, native2D NOOP/FLAP action, horizon1. 180train/20validation episodes,3240000/360000rows,200 unique seeds.
+
+Final offline validation action loss: 0.03496555 on 6400 sampled validation frames. Closed-loop latency-bench evaluation is pending until all four mean VLA trainings complete. Offline loss is not task-performance acceptance.
+
+Weights: `checkpoints/model.pt` (pure tensor state_dict; weights_only CPU load and finite-tensor checks passed). `config.yaml` is portable inference config; `config.full.yaml` records actual training configuration. Backbone/tokenizer source: Qwen/Qwen3-VL-4B-Instruct@ebb281ec70b05090aa6165b016eac8ec08e71b17. Task/history contract and dataset normalization are included.
+
+Data: https://huggingface.co/datasets/latency-sensitive-bench/Standard-Pipeline/tree/a568aa3792b49189d72c02781956c46f50d81b4e/flappy/profile_latency/QwenOFT/mean_rtx3090_s42_200ep_v1
+Training: https://wandb.ai/dongqianyu99-zhejiang-university/starvla_tasks/runs/mean3090-flappy-vla-blackwell-20260930

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

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

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

@@ -0,0 +1,18 @@
+{
+  "dataset_repo": "latency-sensitive-bench/Standard-Pipeline",
+  "revision": "a568aa3792b49189d72c02781956c46f50d81b4e",
+  "verified_hashes": 29,
+  "episodes": 200,
+  "train_episodes": 180,
+  "val_episodes": 20,
+  "train_rows": 3240000,
+  "val_rows": 360000,
+  "state_dim": 24,
+  "action_dim": 2,
+  "action_horizon": 1,
+  "mean_latency_ms": 75.87417450998383,
+  "effective_latency_raw_frames": 1,
+  "effective_latency_ms": 100,
+  "data_root": "/home/ubuntu/lzj/mean-profiling/flappy/data",
+  "canonical_data_recollected": false
+}

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

@@ -0,0 +1,271 @@
+framework:
+  name: QwenOFT
+  qwenvl:
+    base_vlm: /home/ubuntu/lzj/models/Qwen3-VL-4B-Instruct
+    attn_implementation: flash_attention_2
+    enable_gradient_checkpointing: false
+  action_model:
+    action_model_type: MLP
+    action_dim: 2
+    state_dim: 24
+    action_hidden_dim: 2560
+    future_action_window_size: 0
+    past_action_window_size: 0
+    loss_type: discrete_ce
+    state_encoding: continuous_projector
+    task_objective: null
+    action_horizon: 1
+    action_env_dim: 2
+  kv_memory:
+    enabled: false
+    window: 4
+    rollout_len: 8
+    packed_train: false
+    rebased_sink: true
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    include_state: true
+    data_root_dir: /home/ubuntu/lzj/mean-profiling/flappy/data
+    data_mix: lerobot
+    eval_data_mix: lerobot__val
+    custom_mixtures_path: /home/ubuntu/lzj/mean-profiling/flappy/data/_generated_mixtures/lerobot.json
+    action_type: discrete
+    sequential_step_sampling: false
+    eval_sequential_step_sampling: null
+    num_workers: 2
+    eval_num_workers: 2
+    prefetch_factor: 4
+    persistent_workers: true
+    pin_memory: true
+    shuffle: true
+    action_balance:
+      enabled: false
+      strategy: balanced_epoch
+      action_key: action_id
+      target_flap_fraction: 0.3
+      noop_id: 0
+      flap_id: 1
+    latency_curriculum:
+      enabled: false
+      strategy: exclusive
+      latencies: null
+      phase_steps: null
+      phase_distributions: null
+      new_latency_passes: 1.0
+      replay_passes: 0.25
+      target_total_passes: 2.0
+      final_equalization: true
+      step_budget_mode: auto
+      eval_at_phase_end: false
+      save_at_phase_end: false
+      computed_plan: null
+    per_device_batch_size: 64
+    load_all_data_for_training: true
+    num_obs_frames: 1
+    image_mode: single
+    prompt_mode: raw
+    stitch_grid:
+    - 2
+    - 2
+    obs_image_size: null
+    video_backend: torchvision_av
+dataset:
+  source_hf: ''
+  config_name: null
+  source_subdir: null
+  converted_name: flappy_train
+  single_source_hf: ''
+  mixed_source_hf: ''
+  single_converted_name: flappy_train
+  mixed_converted_name: flappy_mixed_latency_train
+  single_latency_filter: null
+  mixed_latency_filter: null
+  force_download: false
+  setup_force: false
+  skip_verification: false
+  target_latency_unit: observation_steps
+  verify_rows: 200
+  max_episodes: null
+  episodes_per_latency: null
+  latency_filter: null
+  debug_subset:
+    enabled: false
+    max_episodes: 5
+    suffix: debug
+base_model:
+  repo_id: Qwen/Qwen3-VL-4B-Instruct
+initialization:
+  checkpoint_local_dir: null
+  checkpoint_hf_repo_id: null
+  checkpoint_filename: null
+trainer:
+  max_train_steps: 5000
+  num_warmup_steps: 100
+  save_interval: 500
+  eval_interval: 500
+  eval_num_batches: 50
+  per_latency_eval_num_batches: null
+  eval_action_classification: false
+  eval_action_classification_interval: null
+  cc_f1_tolerance: 1
+  learning_rate:
+    base: 3.0e-05
+    qwen_vl_interface: 3.0e-06
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr_rate: 0.03333333333333333
+  freeze_modules: ''
+  freeze_vit: false
+  freeze_tied_embedding: false
+  freeze_llm_layers: []
+  loss_scale:
+    vla: 1.0
+    vlm: 0.0
+  max_grad_norm: 1.0
+  weight_decay: 0.0
+  logging_frequency: 1
+  profile_timing:
+    enabled: false
+    log_interval: 10
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 1
+  distributed_backend: deepspeed
+  is_resume: false
+  pretrained_checkpoint: null
+  resume_step: 0
+  reload_modules: null
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+    fused: true
+  save_format: pt
+  hist8_diagnostics:
+    steps:
+    - 0
+    - 500
+    - 1000
+    - 2500
+    - 5000
+    sample_count: 2048
+    noise_seeds:
+    - 1729
+    sample_seed: 42
+    batch_size: 8
+    bootstrap_repeats: 4000
+workspace_dir: WORKSPACE_DIR
+run_root_dir: /home/ubuntu/lzj/mean-profiling/flappy/vla
+seed: 42
+wandb_entity: dongqianyu99-zhejiang-university
+wandb_project: starvla_tasks
+auth:
+  env_file: null
+  hf_token_env: HF_TOKEN
+  wandb_api_key_env: WANDB_API_KEY
+paths:
+  run_root_dir: results/Checkpoints
+  dataset_local_dir: playground/Datasets/rl_games
+  dataset_cache_dir: null
+  base_model_dir: /home/ubuntu/lzj/models/Qwen3-VL-4B-Instruct
+  accelerate_config: starVLA/config/deepseeds/deepspeed_zero2.yaml
+launch:
+  use_accelerate: true
+  gpus: null
+  num_processes: 1
+  dry_run: false
+conda:
+  enabled: true
+  env_name: null
+rl_games:
+  model_alias: openvla
+  env_eval:
+    image_size: 224
+    frameskip: 1
+    image_transform: raw_rgb
+    prompt_mode: raw
+    ghost_trail:
+      history_frames: 5
+      gamma: 1.3
+      min_alpha: 35
+      scroll_px_per_step: 4.0
+      ground_fraction: 0.22
+    seed: 42
+    fixed_episode_seeds: true
+    latency_seed_stride: 0
+    task_seed_stride: 0
+    task_description: ''
+    eval_parallel_envs: 5
+    action_chunk_execution:
+      enabled: false
+      chunk_size: null
+    enabled: false
+    eval_backend: latency_bench
+    distributed_mode: none
+    vectorized:
+      enabled: false
+      batch_size: 1
+    latency:
+      prompt_map_path: null
+      mode: single
+      values:
+      - 0
+    mid_train:
+      enabled: false
+      interval_steps: 100
+      latencies:
+      - 0
+      num_episodes: 5
+      max_steps_per_episode: 3600
+    post_train:
+      enabled: false
+      latencies:
+      - 0
+      - 1
+      - 2
+      - 3
+      - 4
+      num_episodes: 5
+      max_steps_per_episode: 3600
+  task: flappy
+  initialization_mode: scratch
+  action_carrier: native
+model: openvla
+env: flappy
+init: scratch
+mode: single
+checkpoint:
+  load: none
+  hf_repo_id: null
+  save_best_model: false
+  save_final_model: true
+  save_pt_file: true
+  save_training_state: true
+  save_safetensors_file: false
+  local:
+    keep_last_n: 2
+  sync:
+    enabled: false
+    repo_id: null
+    keep_last_n: 0
+    sync_every_n_checkpoints: 1
+    resume_policy: local_latest
+run_id: flappy-qwenoft-mean-3090-s42
+output_dir: /home/ubuntu/lzj/mean-profiling/flappy/vla/flappy-qwenoft-mean-3090-s42
+config_yaml: /home/ubuntu/lzj/mean-profiling/flappy/vla.yaml
+is_debug: false
+version_id: '0.21'
+wandb_name: flappy-qwenoft-mean-3090-s42
+wandb_group: mean-profiling-3090-blackwell
+wandb_tags:
+- flappy
+- qwenoft
+- hist8-v2
+- history01
+- scratch
+- profile-latency
+training_latency_condition: mean-profile

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

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

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

@@ -0,0 +1,194 @@
+{
+  "new_embodiment": {
+    "action": {
+      "mean": [
+        0.9008839726448059,
+        0.09911604970693588
+      ],
+      "std": [
+        0.30056583881378174,
+        0.30056583881378174
+      ],
+      "max": [
+        1.0,
+        1.0
+      ],
+      "min": [
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true
+      ]
+    },
+    "state": {
+      "mean": [
+        0.9995555281639099,
+        0.09905370324850082,
+        0.09904845803976059,
+        0.9996111392974854,
+        0.09906204044818878,
+        0.09905370324850082,
+        0.999666690826416,
+        0.09907067567110062,
+        0.09906204044818878,
+        0.9997222423553467,
+        0.09909043461084366,
+        0.09907067567110062,
+        0.9997777938842773,
+        0.09909876435995102,
+        0.09909043461084366,
+        0.999833345413208,
+        0.09910740703344345,
+        0.09909876435995102,
+        0.9998888969421387,
+        0.09911327064037323,
+        0.09910740703344345,
+        0.9999444484710693,
+        0.09911450743675232,
+        0.09911327064037323
+      ],
+      "std": [
+        0.02107249759137688,
+        0.30049213767051697,
+        0.30048587918281555,
+        0.019712645560502656,
+        0.30050185322761536,
+        0.30049213767051697,
+        0.018251324072479345,
+        0.30051174759864807,
+        0.30050185322761536,
+        0.016662079840898097,
+        0.3005351126194,
+        0.30051174759864807,
+        0.014903813600540161,
+        0.3005451261997223,
+        0.3005351126194,
+        0.012907790020110239,
+        0.3005552887916565,
+        0.3005451261997223,
+        0.01053975615650543,
+        0.30056232213974,
+        0.3005552887916565,
+        0.0074531468562811926,
+        0.3005639910697937,
+        0.30056232213974
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ]
+    },
+    "num_transitions": 3240000,
+    "num_trajectories": 180
+  }
+}

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

@@ -0,0 +1,194 @@
+{
+  "new_embodiment": {
+    "action": {
+      "mean": [
+        0.9010527729988098,
+        0.09894721955060959
+      ],
+      "std": [
+        0.2985910475254059,
+        0.2985910475254059
+      ],
+      "max": [
+        1.0,
+        1.0
+      ],
+      "min": [
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true
+      ]
+    },
+    "state": {
+      "mean": [
+        0.9995555281639099,
+        0.09889999777078629,
+        0.09889444708824158,
+        0.9996111392974854,
+        0.09890277683734894,
+        0.09889999777078629,
+        0.999666690826416,
+        0.09890833497047424,
+        0.09890277683734894,
+        0.9997222423553467,
+        0.09892222285270691,
+        0.09890833497047424,
+        0.9997777938842773,
+        0.09893333166837692,
+        0.09892222285270691,
+        0.999833345413208,
+        0.09894166886806488,
+        0.09893333166837692,
+        0.9998888969421387,
+        0.09894444793462753,
+        0.09894166886806488,
+        0.9999444484710693,
+        0.09894721955060959,
+        0.09894444793462753
+      ],
+      "std": [
+        0.02107250504195682,
+        0.29852861166000366,
+        0.2985208034515381,
+        0.019712593406438477,
+        0.2985321879386902,
+        0.29852861166000366,
+        0.01825134269893312,
+        0.2985396385192871,
+        0.2985321879386902,
+        0.01666202954947938,
+        0.2985576391220093,
+        0.2985396385192871,
+        0.014903815463185194,
+        0.29857251048088074,
+        0.2985576391220093,
+        0.01290779560804152,
+        0.29858362674713135,
+        0.29857251048088074,
+        0.010539757087829197,
+        0.2985873222351074,
+        0.29858362674713135,
+        0.0074531468562811926,
+        0.2985910475254059,
+        0.2985873222351074
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ]
+    },
+    "num_transitions": 360000,
+    "num_trajectories": 20
+  }
+}

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

@@ -0,0 +1,102 @@
+{
+  "README.md": {
+    "bytes": 1513,
+    "sha256": "1f4036fe39bcc14657ad3dce517a5ee46ae764e030e38c60c808b7d29859b255"
+  },
+  "checkpoints/model.pt": {
+    "bytes": 9785137569,
+    "sha256": "f784cb7d6c19872b4ac419f8a31459b395c12f23809f561ae134080a74c6e80a"
+  },
+  "collection_quality_report.json": {
+    "bytes": 532,
+    "sha256": "0741f81a3b2bd9cf7b55d7b8478695532968bcf64abb7640535bea425bac6590"
+  },
+  "config.full.yaml": {
+    "bytes": 6375,
+    "sha256": "d6da13357b405aafb9f5b6e07934a2cd59a28ab3a7578fd81b01bcfea40e8eb2"
+  },
+  "config.yaml": {
+    "bytes": 979,
+    "sha256": "6c33e3ebf0f2a5327163329289b4b225c418bc01f65730ff8f7599702e50d61f"
+  },
+  "dataset_statistics.json": {
+    "bytes": 3299,
+    "sha256": "2c7e8a74a07ca7bec408a0e32cd00957d1af68ecc58fb83571cea5258fe178d7"
+  },
+  "dataset_statistics_eval.json": {
+    "bytes": 3302,
+    "sha256": "a57b1ff610f31dbd61bdb28f5c176e5945adfaa03be75f590e7cd550f658c3dc"
+  },
+  "latency_prompt_map.json": {
+    "bytes": 189,
+    "sha256": "e5bdad8c4747fb2e8101f5d7265e89e2c8794c573ee42bac764365a42b6bc3a4"
+  },
+  "launch_audit.json": {
+    "bytes": 2140,
+    "sha256": "6441da61d2da9190ddc6fd4d7a031acd1d455fc102c3ac40ae3e997aada1fb95"
+  },
+  "manifest.json": {
+    "bytes": 2295,
+    "sha256": "0ee6164ded3e87ffb0ea956b700bc1fae82c1a4de2c3f4b432163fa070f40fc4"
+  },
+  "provenance.json": {
+    "bytes": 3432,
+    "sha256": "17416c9020e50a6abac5e82f2001ec7e0962b479cbde84702061af45ec5116d9"
+  },
+  "summary.jsonl": {
+    "bytes": 159,
+    "sha256": "2905acb79c5fffb33fd4689d87af3bd05d74f4184b9d36ff3eae0c64deecd930"
+  },
+  "task_contract.json": {
+    "bytes": 982,
+    "sha256": "958afa1b69927f50ef9f92e50875b0f6ec937e98ce0a7a86732342a05b94f76c"
+  },
+  "training-code/QWen3.py": {
+    "bytes": 17586,
+    "sha256": "5fe76fc98eaa84e215c3c25c11326815499d2a4b0683c6ab95382b881f535370"
+  },
+  "training-code/QwenOFT.py": {
+    "bytes": 62854,
+    "sha256": "acad1b0dcde6e9fed6e2912b6534db0f8e33a74ec7f6630bbf1b0161fc298b0d"
+  },
+  "training-code/datasets.py": {
+    "bytes": 154347,
+    "sha256": "33e22556471128a2bad18a1ef8b84c399847b4d98b2979d785d139de0cdf36df"
+  },
+  "training-code/gym-language-contract.patch": {
+    "bytes": 688,
+    "sha256": "9f5dee2c5b52d5507aff27373a8cf6def4959361d2c5264d5079b416f28a4c49"
+  },
+  "training-code/loader-spawn-cache.patch": {
+    "bytes": 711,
+    "sha256": "24532ce84e5aa9aef381493fe699b4bb4987afd954422191ae852b0f5eb5b2cd"
+  },
+  "training-code/loader-spawn-test.patch": {
+    "bytes": 1500,
+    "sha256": "a36db15b13d68defeecb75ddc9efa1ac202d9aa0a40f66ba5baaee2dcaef21f1"
+  },
+  "training-code/mean-profile-preparation.patch": {
+    "bytes": 7053,
+    "sha256": "de5fe558376c98515a1cfb186ffdc09b618acc207eb8a5f667b0cbb303767c09"
+  },
+  "training-code/train_starvla.py": {
+    "bytes": 125159,
+    "sha256": "f872648f941cec272af724772e9d6fc3721e9ed6eaf0c8683ab285eb9396212e"
+  },
+  "training-code/trainer_tools.py": {
+    "bytes": 29144,
+    "sha256": "9a48bdfb99b405b03fececfc1e051d3ebb179eb423c15141f02b9be62e2e1791"
+  },
+  "training-code/worker_context.py": {
+    "bytes": 2279,
+    "sha256": "8493631104269f0630655ee2faae592a84ebfb3119aca142c8ac2701f5214464"
+  },
+  "training_code_snapshot.json": {
+    "bytes": 626,
+    "sha256": "136c2cbb92daa93a6dab1aed5644e57491da4f95075c7f4f5d12b3ca8aa16adb"
+  },
+  "training_results.json": {
+    "bytes": 512,
+    "sha256": "b08fa8927249eed1bb0a1c82fc7a75e59e05d7c2fbaf4c6c8ed6f14dc214d488"
+  }
+}

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

@@ -0,0 +1,7 @@
+{
+  "1": {
+    "prompt": "You are playing Flappy Bird. Pass through the pipe gaps and stay alive. Choose the action: NOOP, FLAP.",
+    "latency_raw_frames": 1,
+    "latency_ms": 100.0
+  }
+}

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

@@ -0,0 +1,62 @@
+{
+  "host": "ubuntu@136.115.188.185",
+  "reference_repo_revision": "180da0effdaaaeacef218cf5fbaa4db41af67495",
+  "reference": "/home/ubuntu/lzj/mean-profiling/reference/checkpoint-configs/latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-hist8-scratch-s42-v2/config.full.yaml",
+  "config_sha256": "0d8e33d18da69ac470a3aac5973babdb36b3eb8bfaa02c0f0d3ae8aeb912850f",
+  "base_revision": "ebb281ec70b05090aa6165b016eac8ec08e71b17",
+  "physical_gpus": [
+    0,
+    1
+  ],
+  "micro_batch_per_gpu": 64,
+  "gradient_accumulation": 1,
+  "global_batch": 128,
+  "gradient_checkpointing": false,
+  "optimizer_updates": 5000,
+  "seed": 42,
+  "horizon": 1,
+  "preflight_optimizer_updates": 20,
+  "preflight_discarded": true,
+  "formal_start_step": 0,
+  "preflight_gpu_peaks": [
+    {
+      "rank": "0",
+      "logical_gpu": 0,
+      "gpu_peak_allocated_gib": 92.39531517028809,
+      "gpu_peak_reserved_gib": 92.84375,
+      "main_peak_rss_gib": 18.482311248779297,
+      "wall_seconds_with_initialization": 66.88179596900227,
+      "optimizer_updates": 20,
+      "formal_updates": 0,
+      "global_batch": 128,
+      "micro_per_gpu": 64,
+      "gradient_accumulation": 1,
+      "gradient_checkpointing": false,
+      "status": "passed"
+    },
+    {
+      "rank": "1",
+      "logical_gpu": 1,
+      "gpu_peak_allocated_gib": 92.39531517028809,
+      "gpu_peak_reserved_gib": 93.5078125,
+      "main_peak_rss_gib": 18.478984832763672,
+      "wall_seconds_with_initialization": 66.89783567699851,
+      "optimizer_updates": 20,
+      "formal_updates": 0,
+      "global_batch": 128,
+      "micro_per_gpu": 64,
+      "gradient_accumulation": 1,
+      "gradient_checkpointing": false,
+      "status": "passed"
+    }
+  ],
+  "data_workers_per_rank": 2,
+  "eval_workers_per_rank": 2,
+  "spawn_cache_payload_fix_test": "9 relevant tests passed",
+  "memory_high_gib": 400,
+  "memory_max_gib": 450,
+  "post_train_rollout_eval": false,
+  "evaluation": "pending until four VLA trainings complete",
+  "allocator": "expandable_segments:True",
+  "formal_retry_reason": "CUDA allocator fragmentation/near-capacity initial failure, 20 same-path steps verified"
+}

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

@@ -0,0 +1,111 @@
+{
+  "dataset_name": "lerobot",
+  "env_name": "flappy",
+  "episodes": 180,
+  "frames": 3240000,
+  "task_prompts": [
+    "You are playing Flappy Bird. Pass through the pipe gaps and stay alive. Choose the action: NOOP, FLAP."
+  ],
+  "format": "starvla_lerobot_v2_decision_v1",
+  "source": "/home/ubuntu/lzj/mean-profiling/flappy/data/raw",
+  "robot_type": "rl_games_flappy",
+  "action_layout": "flappy_discrete",
+  "source_action_layout": "flappy_discrete",
+  "action_labels": [
+    "NOOP",
+    "FLAP"
+  ],
+  "action_dim": 2,
+  "active_action_dim": 2,
+  "state_source": "decision_v1",
+  "state_dim": 24,
+  "state_labels": [
+    "decision_history_0",
+    "decision_history_1",
+    "decision_history_2",
+    "decision_history_3",
+    "decision_history_4",
+    "decision_history_5",
+    "decision_history_6",
+    "decision_history_7",
+    "decision_history_8",
+    "decision_history_9",
+    "decision_history_10",
+    "decision_history_11",
+    "decision_history_12",
+    "decision_history_13",
+    "decision_history_14",
+    "decision_history_15",
+    "decision_history_16",
+    "decision_history_17",
+    "decision_history_18",
+    "decision_history_19",
+    "decision_history_20",
+    "decision_history_21",
+    "decision_history_22",
+    "decision_history_23"
+  ],
+  "state_normalization": {
+    "min": [
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0
+    ],
+    "max": [
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0
+    ]
+  },
+  "action_history_decisions": 8,
+  "rows_unit": "decision_step",
+  "obs_fps": 10.0,
+  "obs_stride_raw_frames": 1,
+  "latency_prompt_map_path": "/home/ubuntu/lzj/mean-profiling/flappy/data/lerobot/latency_prompt_map.json",
+  "custom_mixtures_path": "/home/ubuntu/lzj/mean-profiling/flappy/data/_generated_mixtures/lerobot.json",
+  "validation_dataset_name": "lerobot__val",
+  "validation_episodes": 20,
+  "validation_frames": 360000
+}

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

@@ -0,0 +1,93 @@
+{
+  "task": "flappy",
+  "model": "qwenoft",
+  "run_id": "flappy-qwenoft-mean-3090-s42",
+  "training_latency_condition": "mean-profile",
+  "training_hardware": {
+    "gpu": "NVIDIA RTX PRO6000 Blackwell Server Edition",
+    "count": 2,
+    "physical_indices": [
+      0,
+      1
+    ]
+  },
+  "latency_profile_hardware": "1x-RTX3090",
+  "fixed_latency_ms": 75.87417450998383,
+  "effective_latency_raw_frames": 1,
+  "env_fps": 10,
+  "obs_fps": 10,
+  "horizon": 1,
+  "history_encoding": "binary_0_1",
+  "state_dim": 24,
+  "action_dim": 2,
+  "global_batch": 128,
+  "micro_per_gpu": 64,
+  "gradient_accumulation": 1,
+  "gradient_checkpointing": false,
+  "precision": "bf16",
+  "optimizer_backend": "ZeRO2",
+  "teacher": {
+    "repo_id": "latency-sensitive-bench/benchmark-models",
+    "revision": "3ca75d71b55ff42889a4f058d6ad90b2e9b3668f",
+    "path": "/home/ubuntu/lzj/mean-profiling/flappy/restored-teacher/latency-aware/flappy/small-policy/sample-factory-qwenoft-rtx3090-mean-hist8-capacity1-s42-v1",
+    "local_checkpoint": "/home/ubuntu/lzj/mean-profiling/flappy/restored-teacher/latency-aware/flappy/small-policy/sample-factory-qwenoft-rtx3090-mean-hist8-capacity1-s42-v1/best_000048808_25034752_reward_2216.841.pth",
+    "env_steps": 25034752,
+    "hash_verified_files": 8,
+    "weights_only_load": true
+  },
+  "dataset": {
+    "dataset_repo": "latency-sensitive-bench/Standard-Pipeline",
+    "revision": "a568aa3792b49189d72c02781956c46f50d81b4e",
+    "verified_hashes": 29,
+    "episodes": 200,
+    "train_episodes": 180,
+    "val_episodes": 20,
+    "train_rows": 3240000,
+    "val_rows": 360000,
+    "state_dim": 24,
+    "action_dim": 2,
+    "action_horizon": 1,
+    "mean_latency_ms": 75.87417450998383,
+    "effective_latency_raw_frames": 1,
+    "effective_latency_ms": 100,
+    "data_root": "/home/ubuntu/lzj/mean-profiling/flappy/data",
+    "canonical_data_recollected": false
+  },
+  "source_reference_revision": "180da0effdaaaeacef218cf5fbaa4db41af67495",
+  "runtime": {
+    "repo": "c3c6a39365a151e9b7a5e215452fd64e957c2b29",
+    "starvla": "ccca13c5177fe3d3c884b6e2de4965d916016649",
+    "runtime_fixes": [
+      "mean-profile-preparation.patch",
+      "gym-language-contract.patch",
+      "loader-spawn-cache.patch",
+      "loader-spawn-test.patch"
+    ],
+    "pytorch3d": {
+      "revision": "33824be3cbc87a7dd1db0f6a9a9de9ac81b2d0ba",
+      "build": "transforms-only, no native render extension; QwenOFT uses transforms only"
+    },
+    "decord": {
+      "version": "0.6.0",
+      "build": "official source CPU decoder CP310",
+      "wheel_sha256": "e193b356b1e984b4eff08d23b62e482c2c9e5037a6efdc0b1af47079ae2e4c47"
+    }
+  },
+  "wandb_url": "https://wandb.ai/dongqianyu99-zhejiang-university/starvla_tasks/runs/mean3090-flappy-vla-blackwell-20260930",
+  "training_results": {
+    "training_complete": true,
+    "optimizer_updates": 5000,
+    "selection": "final_5000",
+    "train_loss": 0.028774918988347054,
+    "validation_action_loss": 0.034965552389621735,
+    "validation_samples": 6400,
+    "final_backbone_lr": 1e-07,
+    "final_action_lr": 3.3333333333333333e-06,
+    "closed_loop_evaluation": "pending until all four VLA trainings complete",
+    "tensor_keys": 732,
+    "serialized_tensor_elements_including_tied_aliases": 4892433410,
+    "weights_only_load": "passed",
+    "all_tensors_finite": true
+  },
+  "optimizer_recovery_states": "retained on training host; this bundle contains inference weights"
+}

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

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

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

@@ -0,0 +1,43 @@
+{
+  "task": "flappy",
+  "model": "qwenoft",
+  "action_horizon": 1,
+  "image_frames": 1,
+  "state_source": "transport",
+  "state_dim": 24,
+  "history_decisions": 8,
+  "history_encoding": "binary_0_1",
+  "history_normalization": "identity",
+  "history_reset": "all_zero",
+  "history_order": "oldest_to_newest",
+  "history_entry_fields": [
+    "admission_bit",
+    "admitted_command_non_noop_bits",
+    "last_applied_action_non_noop_bits"
+  ],
+  "history_action_bit_order": [
+    "FLAP"
+  ],
+  "action_dim": 2,
+  "action_layout": "flappy_discrete",
+  "action_labels": [
+    "NOOP",
+    "FLAP"
+  ],
+  "task_prompts": [
+    "You are playing Flappy Bird. Pass through the pipe gaps and stay alive. Choose the action: NOOP, FLAP."
+  ],
+  "env_fps": 10,
+  "obs_fps": 10,
+  "evaluation": {
+    "episodes": 100,
+    "seeds": [
+      1000000,
+      1000099
+    ],
+    "parallel_envs": 32,
+    "capacity": 1,
+    "max_raw_frames": 3600,
+    "protocol_version": "hist8-history01-par32-v3"
+  }
+}

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

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

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

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

latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training-code/gym-language-contract.patch

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

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

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

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

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

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

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

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

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

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

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

latency-aware/flappy/vla/starvla-qwenoft-h1/flappy-qwenoft-mean-3090-s42/training_code_snapshot.json

@@ -0,0 +1,19 @@
+{
+  "repo": "c3c6a39365a151e9b7a5e215452fd64e957c2b29",
+  "starvla": "ccca13c5177fe3d3c884b6e2de4965d916016649",
+  "runtime_fixes": [
+    "mean-profile-preparation.patch",
+    "gym-language-contract.patch",
+    "loader-spawn-cache.patch",
+    "loader-spawn-test.patch"
+  ],
+  "pytorch3d": {
+    "revision": "33824be3cbc87a7dd1db0f6a9a9de9ac81b2d0ba",
+    "build": "transforms-only, no native render extension; QwenOFT uses transforms only"
+  },
+  "decord": {
+    "version": "0.6.0",
+    "build": "official source CPU decoder CP310",
+    "wheel_sha256": "e193b356b1e984b4eff08d23b62e482c2c9e5037a6efdc0b1af47079ae2e4c47"
+  }
+}

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

@@ -0,0 +1,15 @@
+{
+  "training_complete": true,
+  "optimizer_updates": 5000,
+  "selection": "final_5000",
+  "train_loss": 0.028774918988347054,
+  "validation_action_loss": 0.034965552389621735,
+  "validation_samples": 6400,
+  "final_backbone_lr": 1e-07,
+  "final_action_lr": 3.3333333333333333e-06,
+  "closed_loop_evaluation": "pending until all four VLA trainings complete",
+  "tensor_keys": 732,
+  "serialized_tensor_elements_including_tied_aliases": 4892433410,
+  "weights_only_load": "passed",
+  "all_tensors_finite": true
+}