Add Florence-Bridge code and handoff documentation

.gitignore

@@ -0,0 +1,31 @@
+__pycache__/
+*.py[cod]
+*.so
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+.venv/
+venv/
+.env
+.env.*
+wandb/
+logs/
+outputs/
+runs/
+checkpoints/
+cache/
+artifacts/
+*.pt
+*.pth
+*.ckpt
+*.safetensors
+*.mp4
+*.parquet
+*.hdf5
+*.jsonl
+*.tmp
+*.orig
+*.patch
+.DS_Store
+.idea/
+.vscode/

.pytest_cache/.gitignore

@@ -0,0 +1,2 @@
+# Created by pytest automatically.
+*

.pytest_cache/CACHEDIR.TAG

@@ -0,0 +1,4 @@
+Signature: 8a477f597d28d172789f06886806bc55
+# This file is a cache directory tag created by pytest.
+# For information about cache directory tags, see:
+#	https://bford.info/cachedir/spec.html

.pytest_cache/README.md

@@ -0,0 +1,8 @@
+# pytest cache directory #
+
+This directory contains data from the pytest's cache plugin,
+which provides the `--lf` and `--ff` options, as well as the `cache` fixture.
+
+**Do not** commit this to version control.
+
+See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information.

.pytest_cache/v/cache/lastfailed

@@ -0,0 +1 @@
+{}

.pytest_cache/v/cache/nodeids

@@ -0,0 +1,17 @@
+[
+  "tests/infoflow/test_objective_preservation.py::test_alpha_zero_is_exact_baseline_identity",
+  "tests/infoflow/test_objective_preservation.py::test_finite_support_objective_and_gradient_are_exactly_preserved",
+  "tests/infoflow/test_objective_preservation.py::test_fixed_seed_monte_carlo_estimator_matches_baseline",
+  "tests/infoflow/test_objective_preservation.py::test_flow_loss_reduction_preserves_legacy_scalar_mean",
+  "tests/infoflow/test_objective_preservation.py::test_invalid_reduction_and_repeat_shapes_fail_fast",
+  "tests/infoflow/test_objective_preservation.py::test_invalid_safe_mixture_inputs_fail_fast[base0-proposal0-0.5-ValueError]",
+  "tests/infoflow/test_objective_preservation.py::test_invalid_safe_mixture_inputs_fail_fast[base1-proposal1-0.5-ValueError]",
+  "tests/infoflow/test_objective_preservation.py::test_invalid_safe_mixture_inputs_fail_fast[base2-proposal2-0.5-ValueError]",
+  "tests/infoflow/test_objective_preservation.py::test_invalid_safe_mixture_inputs_fail_fast[base3-proposal3-1.0-ValueError]",
+  "tests/infoflow/test_objective_preservation.py::test_invalid_safe_mixture_inputs_fail_fast[base4-proposal4--0.1-ValueError]",
+  "tests/infoflow/test_objective_preservation.py::test_ordinary_estimator_is_not_batch_self_normalized_and_stops_weight_gradient",
+  "tests/infoflow/test_objective_preservation.py::test_repeated_flow_losses_are_aggregated_per_original_sample",
+  "tests/infoflow/test_objective_preservation.py::test_repeated_per_sample_path_matches_legacy_scalar_and_gradient",
+  "tests/infoflow/test_objective_preservation.py::test_safe_mixture_normalization_support_and_bound",
+  "tests/infoflow/test_trainer_accumulation.py::test_train_step_preserves_earlier_microbatch_gradients"
+]

.pytest_cache/v/cache/stepwise

@@ -0,0 +1 @@
+[]

AGENTS.md

@@ -0,0 +1,67 @@
+# AGENTS.md — Florence–Bridge OPTR
+
+## Mission
+
+Reproduce or extend Florence–Bridge continuations without changing the experimental contract. Treat `EXPERIMENT_CONTRACT.md` and the downloaded 20K state as authoritative.
+
+## Non-negotiable rules
+
+- Do not use Git operations for this handoff.
+- Never overwrite the 20K common start, frozen proposal, or published 20.5K snapshots.
+- New runs use new output directories.
+- Resume new branches from the complete 20K Accelerate state, never from a 20.5K model-only file.
+- Do not add code, path, manifest, or self-referential checksum launch gates.
+- Keep dataset, episode, valid-anchor and proposal-support identity checks.
+- Checkpoint readiness is established by required-file presence, byte size, loadability and semantic contract checks.
+- Do not silently change world size, global batch, continuation seed, proposal, correction, budget or evaluation panel.
+
+## Required local assets
+
+Set these variables before launch:
+
+```bash
+export OPTR_PROJECT_ROOT=/path/to/this/repository
+export OPTR_RUN_ROOT=/path/to/new/runs
+export OPTR_PYTHON=/path/to/python
+export FLORENCE_MODEL_PATH=/path/to/Florence-2-large
+export BRIDGE_DATA_ROOT=/path/to/OXE_LEROBOT_DATASET
+export BRIDGE_STATISTICS_PATH=/path/to/dataset_statistics.json
+export BRIDGE_PROVENANCE_PATH=$OPTR_PROJECT_ROOT/manifests/data_provenance.json
+export BRIDGE_FINGERPRINT_PATH=$OPTR_PROJECT_ROOT/manifests/content_fingerprints.json
+export BRIDGE_MANIFEST_PATH=$OPTR_PROJECT_ROOT/manifests/valid_anchor_manifest_florence_scratch_v1.json
+export COMMON_START_20K=$OPTR_PROJECT_ROOT/common_start_20k/checkpoints/steps_20000_accelerate
+```
+
+## Locked training contract
+
+- Base: Florence-2-large; action head: DiT-B; horizon 16; state/action dim 7.
+- Two GPUs, BF16, DeepSpeed ZeRO-2.
+- Micro-batch 4/GPU × world size 2 × accumulation 8 = global batch 64.
+- Parent seed 42; continuation/sampler seed 117465; model RNG 588643; transform RNG 298226.
+- Resume step 20,000; stop and save at step 20,500.
+- AdamW `(0.9, 0.95)`, eps `1e-8`, weight decay `1e-8`; clipping 1.0.
+- Cosine scheduler, 5K warmup, min LR `5e-7`; preserve restored scheduler state.
+- Residual: K=8, alpha=0.2, epsilon=0.1, linear within-episode interpolation, ordinary p/q correction.
+
+## Commands
+
+```bash
+bash scripts/run_florence_bridge_arm.sh uniform <new-run-id>
+bash scripts/run_florence_bridge_arm.sh residual <new-run-id> \
+  proposals/residual_temporal_proposal_k8_alpha0.2.json
+```
+
+After start, verify exact resume at 20K, global batch 64, and the first update above 20K. A valid finish has exit code 0 and a complete step-20500 model snapshot.
+
+## Evaluation
+
+Use `scripts/evaluate_bridge_384_no_video.sh`. The locked panel has policy seeds 518957/765353/686765/516408; tasks Stack/Carrot/Spoon/Eggplant; 24 episodes per seed-task; BF16; no video; 384 rollouts per checkpoint. Preserve raw episode outcomes.
+
+## Where to look first
+
+- Contract and factual provenance: `EXPERIMENT_CONTRACT.md`
+- File roles and sizes: `ASSET_MANIFEST.json`
+- Source authority and known corrections: `SOURCE_INVENTORY.md`
+- Training config: `configs/florence_bridge_20k_plus500.yaml`
+- Sampling implementation: `starVLA/dataloader/temporal_proposal.py`, `infoflow_sampler.py`
+- Training integration: `starVLA/training/infoflow.py`, `train_starvla.py`

ASSET_MANIFEST.json

@@ -0,0 +1,22 @@
+{
+  "schema_version": 1,
+  "purpose": "Florence-Bridge public handoff; entries describe file roles and expected byte sizes",
+  "repository": "mmmmmcyaaa/ReTime",
+  "assets": [
+    {"path": "common_start_20k/checkpoints/steps_20000_accelerate/checkpoint.json", "bytes": 1164, "role": "common-start metadata"},
+    {"path": "common_start_20k/checkpoints/steps_20000_accelerate/latest", "bytes": 13, "role": "DeepSpeed latest tag"},
+    {"path": "common_start_20k/checkpoints/steps_20000_accelerate/zero_to_fp32.py", "bytes": 33270, "role": "DeepSpeed conversion helper"},
+    {"path": "common_start_20k/checkpoints/steps_20000_accelerate/random_states_0.pkl", "bytes": 15017, "role": "rank-0 RNG state"},
+    {"path": "common_start_20k/checkpoints/steps_20000_accelerate/random_states_1.pkl", "bytes": 15081, "role": "rank-1 RNG state"},
+    {"path": "common_start_20k/checkpoints/steps_20000_accelerate/pytorch_model/mp_rank_00_model_states.pt", "bytes": 1839089981, "role": "DeepSpeed model state"},
+    {"path": "common_start_20k/checkpoints/steps_20000_accelerate/pytorch_model.pt", "bytes": 2154184469, "role": "consolidated 20K model"},
+    {"path": "common_start_20k/checkpoints/steps_20000_accelerate/pytorch_model/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt", "bytes": 5514449457, "role": "ZeRO-2 optimizer rank 0"},
+    {"path": "common_start_20k/checkpoints/steps_20000_accelerate/pytorch_model/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt", "bytes": 5514463601, "role": "ZeRO-2 optimizer rank 1"},
+    {"path": "checkpoints/residual/steps_20500_pytorch_model.pt", "bytes": 2154184469, "role": "Residual-OPTR model-only evaluation snapshot"},
+    {"path": "checkpoints/uniform/steps_20500_pytorch_model.pt", "bytes": 2154184469, "role": "Uniform model-only evaluation snapshot"},
+    {"path": "proposals/residual_temporal_proposal_k8_alpha0.2.json", "bytes": 132329699, "role": "frozen Residual proposal"},
+    {"path": "manifests/data_provenance.json", "bytes": 651, "role": "dataset provenance"},
+    {"path": "manifests/content_fingerprints.json", "bytes": 53500547, "role": "dataset content identity"},
+    {"path": "manifests/valid_anchor_manifest_florence_scratch_v1.json", "bytes": 101950249, "role": "valid-anchor identity"}
+  ]
+}

CITATION.cff

@@ -0,0 +1,28 @@
+# This CITATION.cff file was generated with cffinit.
+# Visit https://bit.ly/cffinit to generate yours today!
+
+cff-version: 1.2.0
+title: >-
+  StarVLA: A Lego-like Codebase for Vision-Language-Action
+  Model Developing
+message: >-
+  If you use this software, please cite it using the
+  metadata from this file.
+type: software
+authors:
+  - name: starVLA Community
+identifiers:
+  - type: doi
+    value: 10.5281/zenodo.18264214
+repository-code: 'https://github.com/starVLA/starVLA'
+url: 'https://starvla.github.io/'
+abstract: >-
+  A modular, Lego-like codebase for developing
+  Vision-Language-Action models.
+keywords:
+  - vision-language-action
+  - robot-learning
+  - modular-framework
+license: MIT
+version: 1.2.0
+date-released: '2026-01-15'

EXPERIMENT_CONTRACT.md

@@ -0,0 +1,96 @@
+# Florence–Bridge experiment contract
+
+## 1. Scope
+
+This handoff covers exactly two published arms from one complete common start:
+
+1. Uniform continuation, 20K → 20.5K.
+2. Residual-OPTR continuation, K=8, alpha=0.2, with ordinary p/q correction, 20K → 20.5K.
+
+Matched-Random/shuffle is intentionally excluded from the uploaded checkpoint set.
+
+## 2. Common start and resume semantics
+
+The parent is the Florence–Bridge 20K complete Accelerate checkpoint produced by the seed-42 scratch run. It contains the consolidated model, DeepSpeed ZeRO-2 optimizer shards, model-state shard, per-rank RNG states, `latest`, `checkpoint.json`, and `zero_to_fp32.py`.
+
+The continuation restores model, optimizer, scheduler position, RNG streams, data order and update counter. It starts at completed step 20,000 and performs exactly 500 optimizer updates. The active learning rates at resume were `9.407412610098572e-06` for the VLM interface and `9.407412610098571e-05` for the action model. The restored scheduler was `LambdaLR(last_epoch=20000, step_count=20001)`.
+
+The published 20.5K Uniform and Residual files are model-only evaluation snapshots. They are not substitutes for the full parent state.
+
+## 3. Model and optimization
+
+| Field | Value |
+|---|---|
+| Base VLM | Florence-2-large |
+| Registry name | `QwenGR00T` (compatibility name only) |
+| VLM hidden dim | 2048 |
+| Action head | DiT-B, hidden/action hidden 1024 |
+| Cross-attention dim | 1024 |
+| Action/state dim | 7 / 7 |
+| Horizon | 16 (future window 15) |
+| DiT layers/dropout | 16 / 0.2 |
+| Diffusion steps | model 8; training repeated steps 4; inference steps 4 |
+| Precision/distribution | BF16, 2 GPUs, DeepSpeed ZeRO-2, no CPU offload |
+| Batch | 4 per GPU × 2 GPUs × accumulation 8 = global 64 |
+| Optimizer | AdamW, betas (0.9, 0.95), eps 1e-8, weight decay 1e-8 |
+| Nominal LR | VLM interface 1e-5; action model 1e-4; base field 3e-5 |
+| Scheduler | cosine-with-min-LR, warmup 5000, min LR 5e-7 |
+| Gradient clipping | 1.0 |
+| Seeds | parent 42; continuation/sampler 117465; model 588643; transform 298226 |
+| Logging | every 10 steps in the portable handoff; historical main run logged every 50 |
+
+No explicit `freeze_modules` or `trainable_modules` filter was configured. The optimizer contained the VLM-interface and action-model groups. In the actually executed action path, Florence vision/text interface parameters and the action head changed; decoder parameters outside that path received no update. Do not summarize this as “the entire VLM was frozen” or “every registered parameter was actively trained.”
+
+## 4. Data and proposal
+
+- Data mix: `bridge_rt_infoflow` (Bridge + RT1), OXE LeRobot layout.
+- Action type/mode: `delta_ee` / `abs`.
+- Image size: 224×224; video backend: `torchvision_av`.
+- Source and trajectory allocation remain those of the common-start data contract.
+- Residual proposal: schema 1, K=8 prospective grid, linear interpolation within each episode, percentile normalization within episode after interpolation, epsilon=0.1, alpha=0.2.
+- Gate42F status: PASS. The proposal changes encounter frequency; ordinary p/q correction preserves the target objective.
+- Valid-anchor, episode and proposal support identity are semantic requirements. Code location and absolute paths are not.
+
+## 5. Fixed formal evaluation
+
+Official SimplerEnv WidowX/Bridge setup, BF16 inference, no video:
+
+- policy seeds: 518957, 765353, 686765, 516408;
+- tasks: Stack, Carrot, Spoon, Eggplant;
+- 24 episodes per policy-seed/task;
+- 96 rollouts per task and per seed; 384 per checkpoint.
+
+### Uniform
+
+| Policy seed | Success |
+|---|---:|
+| 518957 | 41/96 |
+| 765353 | 39/96 |
+| 686765 | 43/96 |
+| 516408 | 39/96 |
+
+Tasks: Stack 13/96, Carrot 31/96, Spoon 50/96, Eggplant 68/96. Overall: **162/384 = 42.1875%**.
+
+### Residual-OPTR
+
+| Policy seed | Success |
+|---|---:|
+| 518957 | 50/96 |
+| 765353 | 60/96 |
+| 686765 | 59/96 |
+| 516408 | 59/96 |
+
+Tasks: Stack 24/96, Carrot 45/96, Spoon 83/96, Eggplant 76/96. Overall: **228/384 = 59.375%**.
+
+## 6. Reproduction acceptance checks
+
+A reproduction is comparable only if it verifies:
+
+1. the full 20K state loads and the update counter resumes at 20K;
+2. world size 2 and global batch 64;
+3. dataset/episode/valid-anchor identity and proposal full support;
+4. proposal weights are p/q for the Residual arm;
+5. step 20500 is saved cleanly;
+6. formal evaluation contains exactly 384 raw outcomes with the locked seeds/tasks.
+
+Do not introduce code/path checksums, self-referential manifest checks, or migration-file identity gates.

LICENSE

@@ -0,0 +1,24 @@
+    MIT License
+
+    Copyright (c) StarVLA Team.
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    Rebases are allowed for forks and feature branches. When rebasing from upstream StarVLA, use descriptive commit messages, e.g., "chore: clone from StarVLA".
+    Preserve attribution: keep at least the two latest upstream StarVLA commits as separate.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

README.md

@@ -0,0 +1,61 @@
+---
+license: mit
+library_name: pytorch
+tags:
+  - robotics
+  - vision-language-action
+  - florence-2
+  - bridge
+  - importance-sampling
+---
+
+# ReTime / OPTR — Florence–Bridge handoff
+
+This package contains the server-authoritative Florence–Bridge OPTR code, a complete 20K common-start training state, and the Uniform/Residual-OPTR 20.5K evaluation snapshots.
+
+## What is included
+
+- `common_start_20k/checkpoints/steps_20000_accelerate/`: **full resumable state** (model, optimizer, scheduler metadata and per-rank RNG state). Start every new continuation arm here.
+- `checkpoints/uniform/steps_20500_pytorch_model.pt`: model-only Uniform +500 snapshot.
+- `checkpoints/residual/steps_20500_pytorch_model.pt`: model-only Residual-OPTR +500 snapshot.
+- `proposals/residual_temporal_proposal_k8_alpha0.2.json`: frozen K=8, alpha=0.2 Residual proposal.
+- `manifests/`: Bridge/RT1 provenance, content fingerprints and valid-anchor identity.
+- `configs/`, `scripts/`, `starVLA/`: portable training and evaluation code.
+- `EXPERIMENT_CONTRACT.md`: exact experimental contract and results.
+- `AGENTS.md`: fast handoff instructions for an AI coding agent.
+
+The 20.5K files are sufficient for evaluation but **not** for exact optimizer/RNG continuation. Any new branch must resume from the complete 20K state.
+
+## Headline result
+
+Fixed no-video SimplerEnv WidowX/Bridge panel: four policy seeds (`518957`, `765353`, `686765`, `516408`) × four tasks × 24 episodes = 384 rollouts.
+
+| Method | Stack | Carrot | Spoon | Eggplant | Overall |
+|---|---:|---:|---:|---:|---:|
+| Uniform +500 | 13/96 | 31/96 | 50/96 | 68/96 | 162/384 = 42.19% |
+| Residual-OPTR +500 | 24/96 | 45/96 | 83/96 | 76/96 | 228/384 = 59.38% |
+
+## Quick start
+
+1. Download this repository and install `requirements.txt` in the same software environment used for StarVLA.
+2. Download Florence-2-large and the OXE LeRobot Bridge/RT1 data.
+3. Point the environment variables listed in `AGENTS.md` to those local assets.
+4. Keep the contract fixed: two GPUs, BF16, ZeRO-2, micro-batch 4/GPU, accumulation 8, global batch 64, continuation seed 117465.
+5. Run one arm:
+
+```bash
+bash scripts/run_florence_bridge_arm.sh uniform uniform-reproduction
+bash scripts/run_florence_bridge_arm.sh residual residual-reproduction \
+  proposals/residual_temporal_proposal_k8_alpha0.2.json
+```
+
+The launch path performs only file-presence and semantic data/support checks.
+
+## Important implementation facts
+
+- `framework.name: QwenGR00T` is a historical registry name; the base VLM is Florence-2-large.
+- The actual Florence–Bridge run used `cross_attention_dim: 1024` and DeepSpeed ZeRO-2.
+- The optimizer groups are the VLM interface and action model. No explicit freeze list was configured; the executed action path updated the Florence vision/text interface and action head, while decoder parameters outside that path did not receive updates.
+- Residual sampling uses ordinary `p/q` objective correction. This release intentionally does not upload Matched-Random/shuffle assets.
+
+See `SOURCE_INVENTORY.md` before changing any method or training field.

SOURCE_INVENTORY.md

@@ -0,0 +1,40 @@
+# Source inventory and authority
+
+## Authority order
+
+1. Server run artifacts and full configs on `jlu66`.
+2. This corrected handoff package.
+3. Public development repositories: [ReTime_source](https://github.com/miaochengy-pngu/ReTime_source) and [OPTR-StarVLA](https://github.com/miaochengy-pngu/OPTR-StarVLA).
+
+Server facts override stale prose or inferred settings in a repository snapshot.
+
+## Server-authoritative origins
+
+- Development tree: `/home/miaochengyu/starVLA-infoflow-codex`
+- Curated source tree: `/home/miaochengyu/OPTR-StarVLA-clean`
+- 20K parent: `/mnt/miaochengyu/starvla-infoflow-runtime/runs/florence-scratch-official-100k/florence-scratch-uniform-seed-42-official100k-mb4-eval1k-20260726`
+- Confirmatory arms: `/mnt/miaochengyu/starvla-infoflow-runtime/runs/infoflow-confirmatory-22k-20260730`
+- Frozen proposal: `/mnt/miaochengyu/starvla-infoflow-runtime/runs/infoflow-proposals/scratch20k_residual_grid8_gate42f_batch1_20260729/residual_temporal_proposal.json`
+- Formal aggregate: `/home/miaochengyu/starVLA-infoflow-codex/logs/bridgert/infoflow_confirmatory_eval/seed_117465/final_test_policy_seeds_20p5/aggregate_summary.json`
+
+## Corrections applied during handoff
+
+- Corrected DeepSpeed from a stale ZeRO-3 description to the actually used **ZeRO-2** configuration.
+- Corrected `diffusion_model_cfg.cross_attention_dim` from stale 2048 to the actually used **1024**.
+- Fixed world size directly at 2 in the portable Accelerate config.
+- Restricted the public launcher to the uploaded Uniform and Residual arms.
+- Removed engineering checksum/path gates from launch.
+- Kept the real semantic identity requirements for data, episodes, valid anchors and proposal support.
+
+## Key implementation files
+
+- `starVLA/training/train_starvla.py`: trainer, exact-state resume, checkpointing.
+- `starVLA/training/infoflow.py`: InfoFlow/importance-correction integration.
+- `starVLA/dataloader/temporal_proposal.py`: proposal parsing and support semantics.
+- `starVLA/dataloader/infoflow_sampler.py`: traceable sampling.
+- `starVLA/scoring/flow_residual.py`: residual scoring.
+- `starVLA/scoring/residual_cache.py`: score-cache handling.
+- `scripts/run_florence_bridge_arm.sh`: portable two-GPU continuation launcher.
+- `scripts/evaluate_bridge_384_no_video.sh`: locked 384-rollout orchestrator.
+
+The public dataset and base Florence model are not redistributed in this package. Their local paths are supplied through environment variables.

configs/deepspeed_florence_2gpu.yaml

@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_config_file: configs/ds_florence_zero2.json
+distributed_type: DEEPSPEED
+downcast_bf16: "no"
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+use_cpu: false

configs/ds_florence_zero2.json

@@ -0,0 +1,19 @@
+{
+  "fp16": {"enabled": false},
+  "bf16": {"enabled": true},
+  "train_micro_batch_size_per_gpu": "auto",
+  "train_batch_size": "auto",
+  "gradient_accumulation_steps": "auto",
+  "zero_optimization": {
+    "stage": 2,
+    "allgather_partitions": true,
+    "allgather_bucket_size": 50000000,
+    "reduce_scatter": true,
+    "reduce_bucket_size": 50000000,
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "cpu_offload": false
+  },
+  "gradient_clipping": 1.0,
+  "steps_per_print": 10
+}

configs/florence_bridge_20k_plus500.yaml

@@ -0,0 +1,110 @@
+run_id: florence-bridge-20k-plus500
+run_root_dir: ${oc.env:OPTR_RUN_ROOT}
+seed: 117465
+wandb_entity: null
+wandb_project: OPTR-Florence-Bridge
+is_debug: false
+version_id: "0.21"
+
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: ${oc.env:FLORENCE_MODEL_PATH}
+    attn_implementation: flash_attention_2
+    vl_hidden_dim: 2048
+  dino:
+    dino_backbone: dinov2_vits14
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 1024
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 7
+    state_dim: 7
+    action_horizon: 16
+    future_action_window_size: 15
+    past_action_window_size: 0
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 1024
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: ${oc.env:BRIDGE_DATA_ROOT}
+    data_mix: bridge_rt_infoflow
+    action_type: delta_ee
+    action_mode: abs
+    precomputed_statistics_path: ${oc.env:BRIDGE_STATISTICS_PATH}
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    CoT_prompt: Your task is {instruction}. To identify the key objects for your task. Locate their bounding boxes in [x1,y1,x2,y2] format.
+    per_device_batch_size: 4
+    load_all_data_for_training: true
+    obs_image_size: [224, 224]
+    video_backend: torchvision_av
+    num_workers: 4
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    traceable_sampling: true
+    expected_world_size: 2
+    dataset_provenance_path: ${oc.env:BRIDGE_PROVENANCE_PATH}
+    dataset_content_fingerprint_path: ${oc.env:BRIDGE_FINGERPRINT_PATH}
+    valid_anchor_manifest_path: ${oc.env:BRIDGE_MANIFEST_PATH}
+    temporal_proposal_path: null
+    transform_rng_seed: 298226
+
+trainer:
+  use_deepspeed: true
+  max_train_steps: 100000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps: [20500]
+  evaluation_snapshot_interval: null
+  eval_interval: 100001
+  repeated_diffusion_steps: 4
+  pretrained_checkpoint: null
+  is_resume: true
+  allow_model_only_resume: false
+  learning_rate:
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 1.0e-04
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ""
+  trainable_modules: ""
+  frozen_modules_eval: false
+  logging_frequency: 10
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 8
+  save_format: pt
+  save_final_checkpoint: false
+  optimizer:
+    name: AdamW
+    betas: [0.9, 0.95]
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  resume_accelerate_checkpoint: ${oc.env:COMMON_START_20K}
+  branch_parent_seed: 42
+  branch_rng_seed: 117465
+  model_rng_seed: 588643
+  stop_after_steps: 20500
+
+output_dir: ${oc.env:OPTR_RUN_ROOT}/florence-bridge-20k-plus500

configs/reference/common_start_20k_full.yaml

@@ -0,0 +1,113 @@
+run_id: florence-scratch-uniform-seed-42-official100k-mb4-eval1k-20260726
+run_root_dir: /mnt/miaochengyu/starvla-infoflow-runtime/runs/florence-scratch-official-100k
+seed: 42
+wandb_entity: null
+wandb_project: InfoFlow-StarVLA-BridgeRT
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-2-large
+    attn_implementation: flash_attention_2
+    vl_hidden_dim: 2048
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 1024
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 7
+    state_dim: 7
+    action_horizon: 16
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 1024
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    future_action_window_size: 15
+    past_action_window_size: 0
+  dino:
+    dino_backbone: dinov2_vits14
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Datasets/OXE_LEROBOT_DATASET
+    data_mix: bridge_rt_infoflow
+    action_type: delta_ee
+    action_mode: abs
+    precomputed_statistics_path: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-GR00T-Bridge-RT-1/dataset_statistics.json
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    CoT_prompt: Your task is {instruction}. To identify the key objects for your task.
+      Locate their bounding boxes in [x1,y1,x2,y2] format.
+    per_device_batch_size: 4
+    load_all_data_for_training: true
+    obs_image_size:
+    - 224
+    - 224
+    video_backend: torchvision_av
+    num_workers: 4
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    traceable_sampling: true
+    expected_world_size: 2
+    dataset_provenance_path: /home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/data_provenance.json
+    dataset_content_fingerprint_path: /mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/content_fingerprints.json
+    valid_anchor_manifest_path: /mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/valid_anchor_manifest_florence_scratch_v1.json
+    temporal_proposal_path: null
+trainer:
+  use_deepspeed: true
+  max_train_steps: 100000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps:
+  - 5000
+  - 10000
+  - 20000
+  - 30000
+  - 50000
+  - 75000
+  - 100000
+  evaluation_snapshot_interval: 1000
+  eval_interval: 100001
+  repeated_diffusion_steps: 4
+  pretrained_checkpoint: null
+  is_resume: false
+  allow_model_only_resume: false
+  learning_rate:
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ''
+  trainable_modules: ''
+  frozen_modules_eval: false
+  logging_frequency: 50
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 8
+  save_format: pt
+  save_final_checkpoint: false
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+config_yaml: examples/InfoFlow/train_files/server/bridgert_florence_scratch_official_100k_mb4_eval1k.yaml
+output_dir: /mnt/miaochengyu/starvla-infoflow-runtime/runs/florence-scratch-official-100k/florence-scratch-uniform-seed-42-official100k-mb4-eval1k-20260726

configs/reference/residual_seed117465_full.yaml

@@ -0,0 +1,118 @@
+run_id: infoflow-confirmatory-residual-seed-117465
+run_root_dir: /mnt/miaochengyu/starvla-infoflow-runtime/runs/infoflow-confirmatory-22k-20260730
+seed: 117465
+wandb_entity: null
+wandb_project: InfoFlow-StarVLA-BridgeRT
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-2-large
+    attn_implementation: flash_attention_2
+    vl_hidden_dim: 2048
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 1024
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 7
+    state_dim: 7
+    action_horizon: 16
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 1024
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    future_action_window_size: 15
+    past_action_window_size: 0
+  dino:
+    dino_backbone: dinov2_vits14
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Datasets/OXE_LEROBOT_DATASET
+    data_mix: bridge_rt_infoflow
+    action_type: delta_ee
+    action_mode: abs
+    precomputed_statistics_path: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-GR00T-Bridge-RT-1/dataset_statistics.json
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    CoT_prompt: Your task is {instruction}. To identify the key objects for your task.
+      Locate their bounding boxes in [x1,y1,x2,y2] format.
+    per_device_batch_size: 4
+    load_all_data_for_training: true
+    obs_image_size:
+    - 224
+    - 224
+    video_backend: torchvision_av
+    num_workers: 4
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    traceable_sampling: true
+    expected_world_size: 2
+    dataset_provenance_path: /home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/data_provenance.json
+    dataset_content_fingerprint_path: /mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/content_fingerprints.json
+    valid_anchor_manifest_path: /mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/valid_anchor_manifest_florence_scratch_v1.json
+    valid_anchor_manifest_migration_path: /home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/valid_anchor_manifest_rng_migration_20260729.json
+    temporal_proposal_path: /mnt/miaochengyu/starvla-infoflow-runtime/runs/infoflow-proposals/scratch20k_residual_grid8_gate42f_batch1_20260729/residual_temporal_proposal.json
+    transform_rng_seed: 298226
+trainer:
+  use_deepspeed: true
+  max_train_steps: 100000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps:
+  - 20500
+  - 21000
+  - 22000
+  evaluation_snapshot_interval: 1000
+  eval_interval: 100001
+  repeated_diffusion_steps: 4
+  pretrained_checkpoint: null
+  is_resume: true
+  allow_model_only_resume: false
+  learning_rate:
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ''
+  trainable_modules: ''
+  frozen_modules_eval: false
+  logging_frequency: 50
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 8
+  save_format: pt
+  save_final_checkpoint: false
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  resume_accelerate_checkpoint: /mnt/miaochengyu/starvla-infoflow-runtime/runs/florence-scratch-official-100k/florence-scratch-uniform-seed-42-official100k-mb4-eval1k-20260726/checkpoints/steps_20000_accelerate
+  resume_checkpoint_integrity_manifest: /home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/common_start_accelerate_integrity_20260730.json
+  resume_checkpoint_integrity_manifest_sha256: 899d2cef3f73191e11f7a1628afc05b60c80af5816b8fd77c1c69e90d5c9ff39
+  branch_parent_seed: 42
+  branch_rng_seed: 117465
+  model_rng_seed: 588643
+  stop_after_steps: 22000
+config_yaml: /home/miaochengyu/starVLA-infoflow-codex/examples/InfoFlow/train_files/server/generated/infoflow_confirmatory_20k_20260730/confirmatory_residual_seed_117465.yaml
+output_dir: /mnt/miaochengyu/starvla-infoflow-runtime/runs/infoflow-confirmatory-22k-20260730/infoflow-confirmatory-residual-seed-117465

configs/reference/uniform_seed117465_full.yaml

@@ -0,0 +1,118 @@
+run_id: infoflow-confirmatory-uniform-seed-117465
+run_root_dir: /mnt/miaochengyu/starvla-infoflow-runtime/runs/infoflow-confirmatory-22k-20260730
+seed: 117465
+wandb_entity: null
+wandb_project: InfoFlow-StarVLA-BridgeRT
+is_debug: false
+version_id: '0.21'
+framework:
+  name: QwenGR00T
+  qwenvl:
+    base_vlm: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-2-large
+    attn_implementation: flash_attention_2
+    vl_hidden_dim: 2048
+  action_model:
+    action_model_type: DiT-B
+    action_hidden_dim: 1024
+    hidden_size: 1024
+    add_pos_embed: true
+    max_seq_len: 1024
+    action_dim: 7
+    state_dim: 7
+    action_horizon: 16
+    repeated_diffusion_steps: 8
+    noise_beta_alpha: 1.5
+    noise_beta_beta: 1.0
+    noise_s: 0.999
+    num_timestep_buckets: 1000
+    num_inference_timesteps: 4
+    num_target_vision_tokens: 32
+    diffusion_model_cfg:
+      cross_attention_dim: 1024
+      dropout: 0.2
+      final_dropout: true
+      interleave_self_attention: true
+      norm_type: ada_norm
+      num_layers: 16
+      output_dim: 1024
+      positional_embeddings: null
+    future_action_window_size: 15
+    past_action_window_size: 0
+  dino:
+    dino_backbone: dinov2_vits14
+datasets:
+  vla_data:
+    dataset_py: lerobot_datasets
+    data_root_dir: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Datasets/OXE_LEROBOT_DATASET
+    data_mix: bridge_rt_infoflow
+    action_type: delta_ee
+    action_mode: abs
+    precomputed_statistics_path: /mnt/miaochengyu/starvla-infoflow-runtime/playground/Pretrained_models/Florence-GR00T-Bridge-RT-1/dataset_statistics.json
+    precomputed_statistics_key_map:
+      bridge_orig_lerobot: oxe_bridge
+      fractal20220817_data_lerobot: oxe_rt1
+    CoT_prompt: Your task is {instruction}. To identify the key objects for your task.
+      Locate their bounding boxes in [x1,y1,x2,y2] format.
+    per_device_batch_size: 4
+    load_all_data_for_training: true
+    obs_image_size:
+    - 224
+    - 224
+    video_backend: torchvision_av
+    num_workers: 4
+    balance_dataset_weights: false
+    balance_trajectory_weights: false
+    traceable_sampling: true
+    expected_world_size: 2
+    dataset_provenance_path: /home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/data_provenance.json
+    dataset_content_fingerprint_path: /mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/content_fingerprints.json
+    valid_anchor_manifest_path: /mnt/miaochengyu/starvla-infoflow-runtime/manifests/bridgert/valid_anchor_manifest_florence_scratch_v1.json
+    valid_anchor_manifest_migration_path: /home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/valid_anchor_manifest_rng_migration_20260729.json
+    temporal_proposal_path: null
+    transform_rng_seed: 298226
+trainer:
+  use_deepspeed: true
+  max_train_steps: 100000
+  num_warmup_steps: 5000
+  save_interval: null
+  save_steps:
+  - 20500
+  - 21000
+  - 22000
+  evaluation_snapshot_interval: 1000
+  eval_interval: 100001
+  repeated_diffusion_steps: 4
+  pretrained_checkpoint: null
+  is_resume: true
+  allow_model_only_resume: false
+  learning_rate:
+    base: 3.0e-05
+    qwen_vl_interface: 1.0e-05
+    action_model: 0.0001
+  lr_scheduler_type: cosine_with_min_lr
+  scheduler_specific_kwargs:
+    min_lr: 5.0e-07
+  freeze_modules: ''
+  trainable_modules: ''
+  frozen_modules_eval: false
+  logging_frequency: 50
+  gradient_clipping: 1.0
+  gradient_accumulation_steps: 8
+  save_format: pt
+  save_final_checkpoint: false
+  optimizer:
+    name: AdamW
+    betas:
+    - 0.9
+    - 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-08
+  resume_accelerate_checkpoint: /mnt/miaochengyu/starvla-infoflow-runtime/runs/florence-scratch-official-100k/florence-scratch-uniform-seed-42-official100k-mb4-eval1k-20260726/checkpoints/steps_20000_accelerate
+  resume_checkpoint_integrity_manifest: /home/miaochengyu/starVLA-infoflow-codex/protocols/bridgert/common_start_accelerate_integrity_20260730.json
+  resume_checkpoint_integrity_manifest_sha256: 899d2cef3f73191e11f7a1628afc05b60c80af5816b8fd77c1c69e90d5c9ff39
+  branch_parent_seed: 42
+  branch_rng_seed: 117465
+  model_rng_seed: 588643
+  stop_after_steps: 22000
+config_yaml: /home/miaochengyu/starVLA-infoflow-codex/examples/InfoFlow/train_files/server/generated/infoflow_confirmatory_20k_20260730/confirmatory_uniform_seed_117465.yaml
+output_dir: /mnt/miaochengyu/starvla-infoflow-runtime/runs/infoflow-confirmatory-22k-20260730/infoflow-confirmatory-uniform-seed-117465

deployment/model_server/policy_norm_processor.py

@@ -0,0 +1,396 @@
+# Copyright 2025 starVLA community. All rights reserved.
+# Licensed under the MIT License.
+"""PolicyNormProcessor — reuse the training-time ComposedModalityTransform.
+
+This class replaces the hand-rolled un-normalization math that previously
+lived in every eval client. It rebuilds the *exact* ``ComposedModalityTransform``
+used at training time from a checkpoint:
+
+  1. Read ``config.yaml`` next to the ckpt → resolve ``data_mix`` →
+     look up ``robot_type`` from ``DATASET_NAMED_MIXTURES`` →
+     fetch the ``DataConfig`` from ``ROBOT_TYPE_CONFIG_MAP``.
+  2. Build the transform pipeline via ``data_config.transform()``.
+  3. Reconstruct a ``DatasetMetadata`` from the saved
+     ``dataset_statistics.json`` (which stores **combined** per-modality
+     arrays of length ``D``) by splitting it into per-key entries that match
+     ``data_config.action_keys`` / ``state_keys``.
+  4. ``set_metadata(...)`` binds the metadata into every transform.
+
+After construction the caller simply invokes :meth:`unapply_actions` /
+:meth:`apply_state` to get/normalize tensors using the same code path as
+training — there is no second source of truth for normalization math.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional, Sequence
+
+import numpy as np
+import torch
+
+from starVLA.dataloader.gr00t_lerobot.registry import (
+    DATASET_NAMED_MIXTURES,
+    ROBOT_TYPE_CONFIG_MAP,
+)
+from starVLA.dataloader.gr00t_lerobot.schema import (
+    DatasetMetadata,
+    StateActionMetadata,
+)
+from starVLA.dataloader.gr00t_lerobot.transform.base import ComposedModalityTransform
+from starVLA.model.framework.share_tools import read_mode_config
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_robot_type(
+    model_cfg: dict,
+    unnorm_key: Optional[str] = None,
+) -> str:
+    """Look up the training robot_type from the saved cfg.
+
+    Convention used by starVLA train scripts:
+      ``cfg.datasets.vla_data.data_mix`` is a key in
+      ``DATASET_NAMED_MIXTURES`` whose value is a list of
+      ``(dataset_name, weight, robot_type)`` tuples.
+
+    When a data_mix contains entries from multiple robot types (e.g.
+    ``bridge_rt_1`` covers ``oxe_bridge`` + ``oxe_rt1``), ``unnorm_key``
+    is used to identify which embodiment is requested.  In those mixtures
+    the ``robot_type`` field of each entry **matches** the top-level key in
+    ``dataset_statistics.json``, so ``unnorm_key`` serves as the selector.
+    """
+    try:
+        data_mix = model_cfg["datasets"]["vla_data"]["data_mix"]
+    except (KeyError, TypeError) as e:
+        raise KeyError(
+            "ckpt config.yaml is missing `datasets.vla_data.data_mix`; "
+            "cannot resolve training-time robot_type."
+        ) from e
+
+    if data_mix not in DATASET_NAMED_MIXTURES:
+        raise KeyError(
+            f"data_mix={data_mix!r} not in DATASET_NAMED_MIXTURES "
+            f"(available: {sorted(DATASET_NAMED_MIXTURES.keys())[:20]} ...). "
+            "Did you forget to register the example under examples/<bench>/train_files/data_registry/?"
+        )
+
+    mixture = DATASET_NAMED_MIXTURES[data_mix]
+    robot_types = sorted({entry[2] for entry in mixture})
+
+    if len(robot_types) == 1:
+        return robot_types[0]
+
+    # Multiple robot types in the mixture.
+    # Use unnorm_key as a direct selector: for multi-robot mixtures the
+    # dataset_statistics.json top-level keys equal the robot_type values.
+    if unnorm_key is not None and unnorm_key in robot_types:
+        return unnorm_key
+
+    raise ValueError(
+        f"data_mix={data_mix!r} contains multiple robot_types {robot_types}. "
+        "Pass `unnorm_key` matching one of them to disambiguate "
+        f"(e.g. unnorm_key={robot_types[0]!r})."
+    )
+
+
+def _infer_key_dims(
+    data_config: Any,
+    combined_stats: Dict[str, Any],
+    modality_keys: Sequence[str],
+    modality: str,
+) -> Dict[str, int]:
+    """Compute per-key dimensions for splitting combined stats arrays.
+
+    Lookup priority:
+      1. ``data_config.<modality>_key_dims`` — explicit dict classvar on the DataConfig
+         (required for DataConfigs with non-uniform or multi-dim keys).
+      2. Infer uniformly from stats array length: if ``D_total / n_keys`` is an
+         integer, use that as the uniform per-key dim.
+      3. Fall back to dim=1 when no stats are available (empty combined dict).
+
+    Args:
+        data_config: The DataConfig instance for the current robot type.
+        combined_stats: The raw stats dict for the chosen unnorm_key (the value
+            of ``norm_stats[unnorm_key]`` from ``dataset_statistics.json``).
+        modality_keys: Ordered list of full keys for this modality
+            (e.g. ``["action.left_joints", ...]``).
+        modality: ``"action"`` or ``"state"``.
+
+    Returns:
+        Dict mapping each full key to its integer dimension.
+    """
+    attr = f"{modality}_key_dims"
+    if hasattr(data_config, attr):
+        return dict(getattr(data_config, attr))
+
+    combined = combined_stats.get(modality, {})
+    stat_arr = next((v for k, v in combined.items() if k != "mask"), None)
+    n_keys = len(modality_keys)
+    if stat_arr is not None and n_keys > 0:
+        D_total = len(stat_arr)
+        if D_total == n_keys:
+            return {k: 1 for k in modality_keys}
+        elif n_keys > 0 and D_total % n_keys == 0:
+            dim = D_total // n_keys
+            return {k: dim for k in modality_keys}
+        else:
+            raise ValueError(
+                f"Cannot infer per-key dims for modality={modality!r}: "
+                f"D_total={D_total} is not evenly divisible by n_keys={n_keys}. "
+                f"Add `{attr} = {{key: dim, ...}}` to the DataConfig "
+                f"(keys={list(modality_keys)})."
+            )
+    return {k: 1 for k in modality_keys}
+
+
+def _build_dataset_metadata(
+    stats_for_key: Dict[str, Any],
+    embodiment_tag: Any,
+    action_keys: Sequence[str],
+    state_keys: Sequence[str],
+    action_key_dims: Optional[Dict[str, int]] = None,
+    state_key_dims: Optional[Dict[str, int]] = None,
+) -> DatasetMetadata:
+    """Convert the *combined* stats arrays from ``dataset_statistics.json``
+    back into a per-subkey :class:`DatasetMetadata` matching what the training
+    pipeline produced.
+
+    The saved ``dataset_statistics.json`` stores stats as flat arrays of
+    length ``D = sum(per_key_dims)``.  For each ``"action.<sub>"`` key we
+    slice out ``dim_k`` elements starting at the current cursor and store them
+    under ``statistics.action.<sub> = {"min": [v0..v_{k-1}], ...}``.
+
+    Args:
+        stats_for_key: ``norm_stats[unnorm_key]`` from ``dataset_statistics.json``.
+        embodiment_tag: Embodiment tag from the DataConfig.
+        action_keys: Ordered list of full action keys.
+        state_keys: Ordered list of full state keys.
+        action_key_dims: Per-key dimension dict (``{full_key: dim_k}``).
+            Defaults to dim=1 for every key.
+        state_key_dims: Per-key dimension dict for state keys.
+            Defaults to dim=1 for every key.
+    """
+    if action_key_dims is None:
+        action_key_dims = {k: 1 for k in action_keys}
+    if state_key_dims is None:
+        state_key_dims = {k: 1 for k in state_keys}
+
+    def _split_combined(
+        combined: Dict[str, Sequence[float]],
+        keys: Sequence[str],
+        key_dims: Dict[str, int],
+    ):
+        """Split combined arrays into per-subkey dicts using per-key dims.
+
+        ``combined`` looks like ``{"min": [..D..], "max": [..D..], "mask": [..D..], ...}``.
+        ``keys`` is the ordered list of full keys.
+        ``key_dims`` maps each full key to its integer dimension.
+        Returns ``(stats_per_subkey, meta_per_subkey)``.
+        """
+        stats_per_subkey: Dict[str, Dict[str, List[float]]] = {}
+        meta_per_subkey: Dict[str, StateActionMetadata] = {}
+        cursor = 0
+        for full_key in keys:
+            subkey = full_key.split(".", 1)[1]
+            dim_k = key_dims.get(full_key, 1)
+            per_key: Dict[str, List[float]] = {}
+            for stat_name, arr in combined.items():
+                if stat_name == "mask":
+                    continue
+                end = cursor + dim_k
+                if end > len(arr):
+                    # Saved combined array shorter than expected (truncated
+                    # pad channels etc.). Skip this stat field silently.
+                    continue
+                per_key[stat_name] = [float(v) for v in arr[cursor:end]]
+            stats_per_subkey[subkey] = per_key
+            meta_per_subkey[subkey] = StateActionMetadata(
+                absolute=True,
+                rotation_type=None,
+                shape=(dim_k,),
+                continuous=True,
+            )
+            cursor += dim_k
+        return stats_per_subkey, meta_per_subkey
+
+    action_combined = stats_for_key.get("action", {})
+    state_combined = stats_for_key.get("state", {})
+
+    action_stats, action_meta = _split_combined(action_combined, action_keys, action_key_dims)
+    state_stats, state_meta = _split_combined(state_combined, state_keys, state_key_dims)
+
+    # Pydantic accepts dict input with field validators
+    return DatasetMetadata.model_validate(
+        {
+            "statistics": {
+                "state": state_stats,
+                "action": action_stats,
+            },
+            "modalities": {
+                "video": {},
+                "state": state_meta,
+                "action": action_meta,
+            },
+            "embodiment_tag": embodiment_tag.value
+            if hasattr(embodiment_tag, "value")
+            else embodiment_tag,
+        }
+    )
+
+
+class PolicyNormProcessor:
+    """Server-side normalization helper backed by training-time transforms.
+
+    Construct once per checkpoint; call :meth:`unapply_actions` to convert
+    a normalized action chunk ``(T, D)`` back to env-space.
+
+    Args:
+        ckpt_path: Path to the ``*.pt`` checkpoint (the loader looks for
+            ``config.yaml`` and ``dataset_statistics.json`` two dirs up).
+        unnorm_key: Which top-level key in ``dataset_statistics.json`` to use.
+            ``None`` → auto-pick the only key.
+    """
+
+    def __init__(self, ckpt_path: str, unnorm_key: Optional[str] = None) -> None:
+        self._ckpt_path = str(ckpt_path)
+        cfg, norm_stats = read_mode_config(self._ckpt_path)
+        self._model_cfg = cfg
+        self._norm_stats = norm_stats
+
+        # 3-early) Pick the requested unnorm_key (or auto-select) BEFORE
+        # resolving robot_type so we can use it as a hint for multi-robot mixtures.
+        if unnorm_key is None:
+            if len(norm_stats) == 1:
+                unnorm_key = next(iter(norm_stats.keys()))
+            # else: defer error to step 3 below after robot_type resolution attempt
+        elif unnorm_key not in norm_stats:
+            raise KeyError(
+                f"unnorm_key={unnorm_key!r} not in {list(norm_stats.keys())}"
+            )
+        self._unnorm_key = unnorm_key  # may still be None for multi-key case
+
+        # 1) Resolve which DataConfig was used at training.
+        robot_type = _resolve_robot_type(cfg, unnorm_key=unnorm_key)
+        if robot_type not in ROBOT_TYPE_CONFIG_MAP:
+            raise KeyError(
+                f"robot_type={robot_type!r} not in ROBOT_TYPE_CONFIG_MAP "
+                f"(available: {sorted(ROBOT_TYPE_CONFIG_MAP.keys())}). "
+                "Make sure the example's data_registry/data_config.py is importable."
+            )
+        self._data_config = ROBOT_TYPE_CONFIG_MAP[robot_type]
+        self._action_keys: List[str] = list(self._data_config.action_keys)
+        self._state_keys: List[str] = list(getattr(self._data_config, "state_keys", []))
+
+        # 2) Build training-time transform pipeline.
+        transform = self._data_config.transform()
+        if not isinstance(transform, ComposedModalityTransform):
+            transform = ComposedModalityTransform(transforms=[transform])
+        self._transform = transform
+
+        # 3) Pick the requested unnorm_key (finalize; error if still None here).
+        if self._unnorm_key is None:
+            raise ValueError(
+                f"Multiple unnorm_keys in dataset_statistics.json: "
+                f"{list(norm_stats.keys())}. Pass unnorm_key explicitly."
+            )
+        unnorm_key = self._unnorm_key
+
+        # 4) Resolve per-key dims (handles multi-d action/state keys).
+        stats_for_unnorm = norm_stats[unnorm_key]
+        self._action_key_dims: Dict[str, int] = _infer_key_dims(
+            self._data_config, stats_for_unnorm, self._action_keys, "action"
+        )
+        self._state_key_dims: Dict[str, int] = _infer_key_dims(
+            self._data_config, stats_for_unnorm, self._state_keys, "state"
+        )
+
+        # 5) Build & bind metadata.
+        ds_meta = _build_dataset_metadata(
+            stats_for_key=stats_for_unnorm,
+            embodiment_tag=self._data_config.embodiment_tag,
+            action_keys=self._action_keys,
+            state_keys=self._state_keys,
+            action_key_dims=self._action_key_dims,
+            state_key_dims=self._state_key_dims,
+        )
+        self._transform.set_metadata(ds_meta)
+        self._transform.eval()  # mark transforms as eval-mode
+
+        logger.info(
+            "PolicyNormProcessor ready: robot_type=%s, unnorm_key=%s, "
+            "action_keys=%s (dims=%s), state_keys=%s",
+            robot_type,
+            unnorm_key,
+            self._action_keys,
+            [self._action_key_dims[k] for k in self._action_keys],
+            self._state_keys,
+        )
+
+    # ------------------------------------------------------------------
+    # Public properties
+    # ------------------------------------------------------------------
+    @property
+    def action_keys(self) -> List[str]:
+        return list(self._action_keys)
+
+    @property
+    def state_keys(self) -> List[str]:
+        return list(self._state_keys)
+
+    @property
+    def unnorm_key(self) -> str:
+        return self._unnorm_key
+
+    @property
+    def available_unnorm_keys(self) -> List[str]:
+        return list(self._norm_stats.keys())
+
+    @property
+    def transform(self) -> ComposedModalityTransform:
+        return self._transform
+
+    # ------------------------------------------------------------------
+    # Inverse path (model output → env action)
+    # ------------------------------------------------------------------
+    def unapply_actions(self, normalized_actions: np.ndarray) -> np.ndarray:
+        """Invert action normalization using the training-time pipeline.
+
+        Args:
+            normalized_actions: shape ``(T, D)`` where
+                ``D == sum(action_key_dims.values())``.
+
+        Returns:
+            ``(T, D)`` un-normalized actions in env coordinates.
+        """
+        normalized_actions = np.asarray(normalized_actions)
+        assert normalized_actions.ndim == 2, (
+            f"Expected (T, D); got shape {normalized_actions.shape}"
+        )
+
+        # Split (T, D) into per-key {full_key: torch.Tensor[T, dim_k]}.
+        data: Dict[str, torch.Tensor] = {}
+        cursor = 0
+        for full_key in self._action_keys:
+            dim_k = self._action_key_dims.get(full_key, 1)
+            slice_ = normalized_actions[..., cursor : cursor + dim_k]
+            data[full_key] = torch.as_tensor(slice_, dtype=torch.float32)
+            cursor += dim_k
+
+        if cursor != normalized_actions.shape[-1]:
+            raise ValueError(
+                f"Sum of per-key dims ({cursor}) != action_dim "
+                f"({normalized_actions.shape[-1]}). "
+                f"action_keys={self._action_keys}, "
+                f"action_key_dims={self._action_key_dims}"
+            )
+
+        out = self._transform.unapply(data)
+
+        parts: List[np.ndarray] = []
+        for full_key in self._action_keys:
+            v = out[full_key]
+            if isinstance(v, torch.Tensor):
+                v = v.detach().cpu().numpy()
+            parts.append(np.asarray(v))
+        return np.concatenate(parts, axis=-1)

deployment/model_server/policy_wrapper.py

@@ -0,0 +1,174 @@
+# Copyright 2025 starVLA community. All rights reserved.
+# Licensed under the MIT License.
+"""Policy server wrapper.
+
+Encapsulates a `baseframework` instance plus a :class:`PolicyNormProcessor`
+that reuses the *training-time* :class:`ComposedModalityTransform` for action
+un-normalization (no hand-rolled math). The websocket server returns
+already-unnormalized actions.
+
+Client-side responsibilities that REMAIN on the client:
+  - environment-specific adapters (image_history, gripper sticky, action
+    ensembling)
+  - chunk-cache scheduling (`step % chunk_size == 0` triggers a new infer)
+
+Exposed API:
+  - ``metadata`` (dict, sent at handshake): ``action_chunk_size``,
+    ``available_unnorm_keys``, ``action_keys``, ``state_keys``.
+  - ``predict_action(examples, unnorm_key=None, **kwargs)`` returns
+    ``{"actions": np.ndarray[B, T, action_dim]}``.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import torch
+
+from starVLA.model.framework.base_framework import baseframework
+from starVLA.model.framework.share_tools import read_mode_config
+
+from deployment.model_server.policy_norm_processor import PolicyNormProcessor
+
+
+class PolicyServerWrapper:
+    """Wraps a `baseframework` for use as a websocket-server policy."""
+
+    def __init__(
+        self,
+        ckpt_path: str,
+        device: str = "cuda",
+        use_bf16: bool = False,
+        unnorm_key: Optional[str] = None,
+        normalized_action_clip: Optional[float] = None,
+    ) -> None:
+        self._ckpt_path = str(ckpt_path)
+        self._normalized_action_clip = normalized_action_clip
+
+        logging.info("PolicyServerWrapper: loading framework from %s", self._ckpt_path)
+        framework = baseframework.from_pretrained(self._ckpt_path)
+        if use_bf16:
+            framework = framework.to(torch.bfloat16)
+        framework = framework.to(device).eval()
+        self._framework = framework
+
+        # Co-located metadata.
+        model_cfg, _ = read_mode_config(self._ckpt_path)
+        self._model_cfg = model_cfg
+        vla_data_cfg = model_cfg.get("datasets", {}).get("vla_data", {})
+        include_state = vla_data_cfg.get("include_state", False)
+        self._uses_state = include_state not in (False, None, 0, "False", "false", "0")
+
+        # action_chunk_size = future_action_window_size + 1 (matches old client).
+        action_model_cfg = model_cfg["framework"]["action_model"]
+        
+        if "action_horizon" in action_model_cfg:
+            self._action_chunk_size = int(action_model_cfg["action_horizon"])
+        elif "future_action_window_size" in action_model_cfg:
+            self._action_chunk_size = int(action_model_cfg["future_action_window_size"]) + 1
+        else:
+            raise ValueError(
+                f"PolicyServerWrapper: no action_horizon or future_action_window_size found in model config for {self._ckpt_path}"
+            )
+        # Cache of PolicyNormProcessor instances per unnorm_key.
+        # For single-dataset ckpts unnorm_key is auto-selected; for multi-dataset
+        # ckpts clients must pass unnorm_key per request.
+        self._default_unnorm_key = unnorm_key
+        self._norm_processors: Dict[str, PolicyNormProcessor] = {}
+
+        # Peek at available keys without building a full processor.
+        _, _ns = read_mode_config(self._ckpt_path)
+        self._available_unnorm_keys: List[str] = list(_ns.keys())
+
+        # Eagerly build when unambiguous; defer for multi-key / no explicit key.
+        if unnorm_key is not None or len(self._available_unnorm_keys) == 1:
+            default_proc = self._get_processor(unnorm_key)
+            self._default_unnorm_key = default_proc.unnorm_key
+            logging.info(
+                "PolicyServerWrapper ready: action_chunk_size=%d, default_unnorm_key=%s, "
+                "available_unnorm_keys=%s, action_keys=%s, state_keys=%s",
+                self._action_chunk_size,
+                default_proc.unnorm_key,
+                default_proc.available_unnorm_keys,
+                default_proc.action_keys,
+                default_proc.state_keys,
+            )
+        else:
+            logging.info(
+                "PolicyServerWrapper ready (multi-key): action_chunk_size=%d, "
+                "available_unnorm_keys=%s — clients must pass unnorm_key per request.",
+                self._action_chunk_size,
+                self._available_unnorm_keys,
+            )
+
+    def _get_processor(self, unnorm_key: Optional[str]) -> PolicyNormProcessor:
+        cache_key = unnorm_key if unnorm_key is not None else "__default__"
+        if cache_key not in self._norm_processors:
+            self._norm_processors[cache_key] = PolicyNormProcessor(
+                self._ckpt_path, unnorm_key=unnorm_key
+            )
+        return self._norm_processors[cache_key]
+
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        """Model-invariant metadata; sent to client at websocket handshake."""
+        base = {
+            "env": "starvla_policy_server",
+            "ckpt_path": self._ckpt_path,
+            "action_chunk_size": self._action_chunk_size,
+            "uses_state": self._uses_state,
+            "available_unnorm_keys": self._available_unnorm_keys,
+            "default_unnorm_key": self._default_unnorm_key,
+            "normalized_action_clip": self._normalized_action_clip,
+            "action_head_attention_mask": False,
+            "inference_rng": "continuous_server_stream",
+        }
+        # Enrich with per-embodiment keys when a default processor already exists.
+        if self._default_unnorm_key is not None:
+            proc = self._get_processor(self._default_unnorm_key)
+            base["action_keys"] = proc.action_keys
+            base["state_keys"] = proc.state_keys
+        return base
+
+    def predict_action(
+        self,
+        examples: List[dict],
+        unnorm_key: Optional[str] = None,
+        **kwargs,
+    ) -> Dict[str, np.ndarray]:
+        """Run the framework, then un-normalize via training-time transforms.
+
+        Args:
+            examples: list of dicts (each with ``image`` / ``lang`` / optional ``state``).
+            unnorm_key: dataset key for un-normalization stats. ``None`` -->
+                use the wrapper's default (auto-picked at startup).
+            **kwargs: forwarded to the framework's ``predict_action``
+                (``do_sample``, ``use_ddim``, ``num_ddim_steps``, ...).
+
+        Returns:
+            ``{"actions": np.ndarray[B, T, D]}`` -- un-normalized.
+        """
+        effective_key = unnorm_key if unnorm_key is not None else self._default_unnorm_key
+        if effective_key is None:
+            if len(self._available_unnorm_keys) == 1:
+                effective_key = self._available_unnorm_keys[0]
+            else:
+                raise ValueError(
+                    f"predict_action: unnorm_key not specified and no default set. "
+                    f"Pass one of {self._available_unnorm_keys}."
+                )
+        proc = self._get_processor(effective_key)
+
+        out = self._framework.predict_action(examples=examples, **kwargs)
+        normalized = np.asarray(out["normalized_actions"])  # (B, T, D)
+        if self._normalized_action_clip is not None:
+            limit = float(self._normalized_action_clip)
+            normalized = np.clip(normalized, -limit, limit)
+
+        unnorm = np.stack(
+            [proc.unapply_actions(normalized[b]) for b in range(normalized.shape[0])],
+            axis=0,
+        )
+        return {"actions": unnorm}

deployment/model_server/server_policy.py

@@ -0,0 +1,69 @@
+# 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 logging
+import socket
+import argparse
+from deployment.model_server.tools.websocket_policy_server import WebsocketPolicyServer
+from deployment.model_server.policy_wrapper import PolicyServerWrapper
+import os
+
+
+def main(args) -> None:
+    # Example usage:
+    # policy = YourPolicyClass()  # Replace with your actual policy class
+    # server = WebsocketPolicyServer(policy, host="localhost", port=10091)
+    # server.serve_forever()
+
+    policy = PolicyServerWrapper(
+        ckpt_path=args.ckpt_path,
+        device="cuda",
+        use_bf16=args.use_bf16,
+    )
+
+    hostname = socket.gethostname()
+    local_ip = socket.gethostbyname(hostname)
+    logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip)
+
+    # start websocket server
+    server = WebsocketPolicyServer(
+        policy=policy,
+        host="0.0.0.0",
+        port=args.port,
+        idle_timeout=args.idle_timeout,
+        metadata={**policy.metadata, "policy_seed": args.seed},
+    )
+    logging.info("server running ...")
+    server.serve_forever()
+
+
+def build_argparser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ckpt_path", type=str, default="Qwen/Qwen2.5-VL-3B-Instruct")
+    parser.add_argument("--port", type=int, default=10093)
+    parser.add_argument("--use_bf16", action="store_true")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--idle_timeout" , type=int, default=1800, help="Idle timeout in seconds, -1 means never close")
+    return parser
+
+
+def start_debugpy_once():
+    """start debugpy once"""
+    import debugpy
+    if getattr(start_debugpy_once, "_started", False):
+        return
+    debugpy.listen(("0.0.0.0", 10095))
+    print("🔍 Waiting for VSCode attach on 0.0.0.0:10095 ...")
+    debugpy.wait_for_client()
+    start_debugpy_once._started = True
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, force=True)
+    parser = build_argparser()
+    args = parser.parse_args()
+    if os.getenv("DEBUG", False):
+        print("🔍 DEBUGPY is enabled")
+        start_debugpy_once()
+    main(args)

deployment/model_server/tools/image_tools.py

@@ -0,0 +1,123 @@
+import numpy as np
+from PIL import Image
+
+
+def convert_to_uint8(img: np.ndarray) -> np.ndarray:
+    """Converts an image to uint8 if it is a float image.
+
+    This is important for reducing the size of the image when sending it over the network.
+    """
+    if np.issubdtype(img.dtype, np.floating):
+        img = (255 * img).astype(np.uint8)
+    return img
+
+
+def resize_with_pad(images: np.ndarray, height: int, width: int, method=Image.BILINEAR) -> np.ndarray:
+    """Replicates tf.image.resize_with_pad for multiple images using PIL. Resizes a batch of images to a target height.
+
+    Args:
+        images: A batch of images in [..., height, width, channel] format.
+        height: The target height of the image.
+        width: The target width of the image.
+        method: The interpolation method to use. Default is bilinear.
+
+    Returns:
+        The resized images in [..., height, width, channel].
+    """
+    # If the images are already the correct size, return them as is.
+    if images.shape[-3:-1] == (height, width):
+        return images
+
+    original_shape = images.shape
+
+    images = images.reshape(-1, *original_shape[-3:])
+    resized = np.stack([_resize_with_pad_pil(Image.fromarray(im), height, width, method=method) for im in images])
+    return resized.reshape(*original_shape[:-3], *resized.shape[-3:])
+
+
+def _resize_with_pad_pil(image: Image.Image, height: int, width: int, method: int) -> Image.Image:
+    """Replicates tf.image.resize_with_pad for one image using PIL. Resizes an image to a target height and
+    width without distortion by padding with zeros.
+
+    Unlike the jax version, note that PIL uses [width, height, channel] ordering instead of [batch, h, w, c].
+    """
+    cur_width, cur_height = image.size
+    if cur_width == width and cur_height == height:
+        return image  # No need to resize if the image is already the correct size.
+
+    ratio = max(cur_width / width, cur_height / height)
+    resized_height = int(cur_height / ratio)
+    resized_width = int(cur_width / ratio)
+    resized_image = image.resize((resized_width, resized_height), resample=method)
+
+    zero_image = Image.new(resized_image.mode, (width, height), 0)
+    pad_height = max(0, int((height - resized_height) / 2))
+    pad_width = max(0, int((width - resized_width) / 2))
+    zero_image.paste(resized_image, (pad_width, pad_height))
+    assert zero_image.size == (width, height)
+    return zero_image
+
+
+from typing import Any
+
+
+def to_pil_preserve(images: Any, scale_float: bool = True):
+    """
+    Convert (possibly nested) numpy image arrays back to PIL.Image WITHOUT changing spatial shape
+    or nesting structure.
+
+    Accepts:
+      - np.ndarray with shape (H, W, C), C in {1,3,4}, dtype uint8 or float
+      - PIL.Image.Image (returned as-is)
+      - Nested list / tuple structures containing the above
+
+    Guarantees:
+      - No resize / pad / crop performed
+      - Returns an object with the SAME nesting layout (list -> list, tuple -> tuple)
+      - Only dtype (float -> uint8) and channel-mode adaptation may happen
+        * float arrays assumed in [0,1] if scale_float=True (scaled *255 + clip)
+    Args:
+      images: input object / sequence
+      scale_float: whether to scale float images in [0,1] to uint8
+    Returns:
+      Mirrored structure with all leaf nodes as PIL.Image.Image
+    """
+
+    def _convert(obj):
+        # Nested containers
+        if isinstance(obj, list):
+            return [_convert(x) for x in obj]
+        if isinstance(obj, tuple):
+            return tuple(_convert(x) for x in obj)
+
+        # PIL stays
+        if isinstance(obj, Image.Image):
+            return obj
+
+        # numpy -> PIL
+        if isinstance(obj, np.ndarray):
+            arr = obj
+            if arr.ndim != 3:
+                raise ValueError(f"Expected 3D array (H,W,C), got shape={arr.shape}")
+            if arr.shape[2] not in (1, 3, 4):
+                raise ValueError(f"Channel count must be 1/3/4, got {arr.shape[2]}")
+            if np.issubdtype(arr.dtype, np.floating):
+                if scale_float:
+                    arr = np.clip(arr, 0.0, 1.0)
+                    arr = (arr * 255.0 + 0.5).astype(np.uint8)
+                else:
+                    raise TypeError("Float array provided but scale_float=False")
+            elif arr.dtype != np.uint8:
+                arr = arr.astype(np.uint8)
+
+            # Single channel -> 'L'
+            if arr.shape[2] == 1:
+                arr = arr[:, :, 0]
+                return Image.fromarray(arr, mode="L")
+            # 3 channels -> RGB, 4 -> RGBA
+            mode = "RGB" if arr.shape[2] == 3 else "RGBA"
+            return Image.fromarray(arr, mode=mode)
+
+        raise TypeError(f"Unsupported element type: {type(obj)}")
+
+    return _convert(images)

deployment/model_server/tools/msgpack_numpy.py

@@ -0,0 +1,57 @@
+"""Adds NumPy array support to msgpack.
+
+msgpack is good for (de)serializing data over a network for multiple reasons:
+- msgpack is secure (as opposed to pickle/dill/etc which allow for arbitrary code execution)
+- msgpack is widely used and has good cross-language support
+- msgpack does not require a schema (as opposed to protobuf/flatbuffers/etc) which is convenient in dynamically typed
+    languages like Python and JavaScript
+- msgpack is fast and efficient (as opposed to readable formats like JSON/YAML/etc); I found that msgpack was ~4x faster
+    than pickle for serializing large arrays using the below strategy
+
+The code below is adapted from https://github.com/lebedov/msgpack-numpy. The reason not to use that library directly is
+that it falls back to pickle for object arrays.
+"""
+
+import functools
+
+import msgpack
+import numpy as np
+
+
+def pack_array(obj):
+    if (isinstance(obj, (np.ndarray, np.generic))) and obj.dtype.kind in ("V", "O", "c"):
+        raise ValueError(f"Unsupported dtype: {obj.dtype}")
+
+    if isinstance(obj, np.ndarray):
+        return {
+            b"__ndarray__": True,
+            b"data": obj.tobytes(),
+            b"dtype": obj.dtype.str,
+            b"shape": obj.shape,
+        }
+
+    if isinstance(obj, np.generic):
+        return {
+            b"__npgeneric__": True,
+            b"data": obj.item(),
+            b"dtype": obj.dtype.str,
+        }
+
+    return obj
+
+
+def unpack_array(obj):
+    if b"__ndarray__" in obj:
+        return np.ndarray(buffer=obj[b"data"], dtype=np.dtype(obj[b"dtype"]), shape=obj[b"shape"])
+
+    if b"__npgeneric__" in obj:
+        return np.dtype(obj[b"dtype"]).type(obj[b"data"])
+
+    return obj
+
+
+Packer = functools.partial(msgpack.Packer, default=pack_array)
+packb = functools.partial(msgpack.packb, default=pack_array)
+
+Unpacker = functools.partial(msgpack.Unpacker, object_hook=unpack_array)
+unpackb = functools.partial(msgpack.unpackb, object_hook=unpack_array)

deployment/model_server/tools/websocket_policy_client.py

@@ -0,0 +1,77 @@
+# 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 logging, argparse
+import time, os
+from typing import Dict, Optional, Tuple
+
+from typing_extensions import override
+import websockets.sync.client
+
+from . import msgpack_numpy
+
+
+class WebsocketClientPolicy:
+    """Implements the Policy interface by communicating with a server over websocket.
+
+    See WebsocketPolicyServer for a corresponding server implementation.
+    """
+
+    def __init__(self, host: str = "127.0.0.1", port: Optional[int] = 10093, api_key: Optional[str] = None) -> None:
+        # 0.0.0.0 cannot be used as a connection target, here default 127.0.0.1
+        self._uri = f"ws://{host}"
+        if port is not None:
+            self._uri += f":{port}"
+        self._packer = msgpack_numpy.Packer()
+        self._api_key = api_key
+        self._ws, self._server_metadata = self._wait_for_server()
+
+    def get_server_metadata(self) -> Dict:
+        return self._server_metadata
+
+    def _wait_for_server(self, timeout: float = 300) -> Tuple[websockets.sync.client.ClientConnection, Dict]:
+        logging.info(f"Waiting for server at {self._uri}...")
+        start_time = time.time()
+        
+        for k in ("HTTP_PROXY", "http_proxy", "HTTPS_PROXY", "https_proxy", "ALL_PROXY", "all_proxy"):
+            os.environ.pop(k, None)
+        
+        while True:
+            if time.time() - start_time > timeout:
+                raise TimeoutError(f"Failed to connect to server within {timeout} seconds")
+            
+            try:
+                headers = {"Authorization": f"Api-Key {self._api_key}"} if self._api_key else None
+                conn = websockets.sync.client.connect(
+                    self._uri,
+                    compression=None,
+                    max_size=None,
+                    additional_headers=headers,
+                    open_timeout=150,
+                    ping_interval=None,
+                    ping_timeout=None,
+                )
+                metadata = msgpack_numpy.unpackb(conn.recv())
+                return conn, metadata
+            except ConnectionRefusedError:
+                logging.info(f"Still waiting for server {self._uri} ...")
+                time.sleep(2)
+
+    def close(self) -> None:
+        try:
+            self._ws.close()
+        except Exception:
+            pass
+    
+    @override
+    def predict_action(self, query_info: Dict) -> Dict:
+        data = self._packer.pack(query_info)
+        self._ws.send(data)
+        response = self._ws.recv()
+        if isinstance(response, str):
+            raise RuntimeError(f"Error in inference server:\n{response}")
+        return msgpack_numpy.unpackb(response)
+
+
+

deployment/model_server/tools/websocket_policy_server.py

@@ -0,0 +1,161 @@
+# 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 asyncio
+import logging
+import traceback
+import time
+
+import websockets.asyncio.server
+import websockets.frames
+
+# from openpi_client import base_policy as _base_policy
+from . import msgpack_numpy
+
+class WebsocketPolicyServer:
+    """Serves a policy using the websocket protocol. See websocket_client_policy.py for a client implementation.
+
+    Currently only implements the `load` and `infer` methods.
+    """
+
+    def __init__(
+        self,
+        policy,
+        host: str = "0.0.0.0",
+        port: int = 10093,
+        idle_timeout: int = -1,  # 新增参数，单位秒，-1表示永不关闭
+        metadata: dict | None = None,
+        
+    ) -> None:
+        self._policy = policy  #
+        self._host = host
+        self._port = port
+        self._metadata = metadata or {}
+        self._idle_timeout = idle_timeout
+        self._last_active = time.time()
+        logging.getLogger("websockets.server").setLevel(logging.INFO)
+
+    def serve_forever(self) -> None:
+        asyncio.run(self.run())
+
+    async def run(self):
+        async with websockets.asyncio.server.serve(
+            self._handler,
+            self._host,
+            self._port,
+            compression=None,
+            max_size=None,
+        ) as server:
+            if self._idle_timeout > 0:
+                await self._idle_watchdog(server)
+            else:
+                await server.serve_forever()
+
+    async def _idle_watchdog(self, server):
+        """监控空闲时间，超时则关闭服务器"""
+        while True:
+            await asyncio.sleep(5)
+            if time.time() - self._last_active > self._idle_timeout:
+                logging.info(f"Idle timeout ({self._idle_timeout}s) reached, shutting down server.")
+                server.close()
+                await server.wait_closed()
+                break
+
+    async def _handler(self, websocket: websockets.asyncio.server.ServerConnection):
+        logging.info(f"Connection from {websocket.remote_address} opened")
+        packer = msgpack_numpy.Packer()
+
+        await websocket.send(packer.pack(self._metadata))
+
+        while True:
+            try:
+                msg = msgpack_numpy.unpackb(await websocket.recv())
+                self._last_active = time.time()  # 每次收到消息刷新活跃时间
+                ret = self._route_message(msg)  # route message
+                await websocket.send(packer.pack(ret))
+            except websockets.ConnectionClosed:
+                logging.info(f"Connection from {websocket.remote_address} closed")
+                break
+            except Exception:
+                await websocket.send(traceback.format_exc())
+                await websocket.close(
+                    code=websockets.frames.CloseCode.INTERNAL_ERROR,
+                    reason="Internal server error. Traceback included in previous frame.",
+                )
+                raise
+
+    # route logic: recognize request from client
+    def _route_message(self, msg: dict) -> dict:
+        """
+        Route rules (fault-tolerant):
+        - Supports messages of form:
+            {"type": "ping|init|infer|reset", "request_id": "...", "payload": {...}}
+          or a flat dict (will be treated as payload).
+        - Does NOT raise inside this function: all exceptions are caught and encoded in response.
+        """
+        req_id = msg.get("request_id", "default")
+        mtype = msg.get("type", "infer")          # default = infer
+        msg       # when no explicit payload, treat top-level as payload
+
+        # ping
+        if mtype == "ping":
+            return {"status": "ok", "ok": True, "type": "ping", "request_id": req_id}
+
+        # infer --> framework.predict_action
+        elif mtype == "infer" or mtype == "predict_action":
+            # Basic payload sanity
+            if not isinstance(msg, dict):
+                return {
+                    "status": "error",
+                    "ok": False,
+                    "type": "inference_result",
+                    "request_id": req_id,
+                    "error": {"message": "Payload must be a dict", "payload_type": str(type(payload))}
+                }
+            try:
+
+                ouput_dict = self._policy.predict_action(**msg)
+            except Exception as e:
+                logging.exception("Policy inference error (request_id=%s)", req_id)
+                logging.exception(e)
+                
+                return {
+                    "status": "error",
+                    "ok": False,
+                    "type": "inference_result",
+                    "request_id": req_id,
+                    "error": {
+                        "message": str(e),
+                        # "traceback": traceback.format_exc(),
+                    },
+                }
+            data = ouput_dict
+            return {
+                "status": "ok",
+                "ok": True,
+                "type": "inference_result",
+                "request_id": req_id,
+                "data": data,
+            }
+
+        # unknow request type
+        else:
+            return {
+                "status": "error",
+                "ok": False,
+                "type": "unknown",
+                "request_id": req_id,
+                "error": {"message": f"Unsupported message type '{mtype}'"},
+            }
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, force=True)
+    # Example usage:
+    # policy = YourPolicyClass()  # Replace with your actual policy class
+    # server = WebsocketPolicyServer(policy, host="localhost", port=10091)
+    # server.serve_forever()
+    raise NotImplementedError("This module is not intended to be run directly.")
+#
+#  Instead, it should be imported and used in a server context.

examples/SimplerEnv/README.md

@@ -0,0 +1,158 @@
+This document provides instructions for reproducing our **experimental results** with SimplerEnv.  
+
+
+The evaluation process consists of two main parts:  
+
+1. Setting up the `simpler_env` environment and dependencies.  
+2. Running the evaluation by launching services in both `starVLA` and `simpler_env` environments.  
+
+We have verified that this workflow runs successfully on both **NVIDIA A100** and **RTX 4090** GPUs.  
+
+## 📦 1. Environment Setup
+
+To set up the environment, please first follow the official [SimplerEnv repository](https://github.com/simpler-env/SimplerEnv) to install the base `simpler_env` environment. 
+
+Afterwards, inside the `simpler_env` environment, install the following dependencies:  
+
+```bash
+conda activate simpler_env
+pip install tyro matplotlib mediapy websockets msgpack
+pip install numpy==1.24.4
+```
+
+⚠️ **Common Issues**
+When testing SimplerEnv on NVIDIA A100, you may encounter the following error:
+`libvulkan.so.1: cannot open shared object file: No such file or directory`
+You can refer to this link to fix: [Installation Guide – Vulkan Section](https://maniskill.readthedocs.io/en/latest/user_guide/getting_started/installation.html#vulkan)
+
+
+## 🔧 Verification Your SimplerEnv
+We provide a minimal environment verification script:
+
+```bash
+python examples/SimplerEnv/test_your_simplerEnv.py
+```
+
+If you see the "✅ Env built successfully" message, it means SimplerEnv is installed correctly and ready to use.
+
+
+---
+
+
+## 🚀 2. Eval SimplerEnv
+
+The evaluation should be run **from the repository root** using **two separate terminals/machines**, one for each environment:  
+
+- **starVLA environment**: runs the policy inference server.  
+- **simpler_env environment**: runs the simulation eval code.  
+
+
+### Step 0. Download offical checkpoint
+
+Available SimplerEnv WidowX checkpoints (see [docs/model_zoo.md](../../docs/model_zoo.md) for the full list):
+
+- [Qwen3VL-GR00T-Bridge-RT-1](https://huggingface.co/StarVLA/Qwen3VL-GR00T-Bridge-RT-1) — 65.3 avg success
+- [Qwen3VL-PI_v3-Bridge-RT-1](https://huggingface.co/StarVLA/Qwen3VL-PI_v3-Bridge-RT_1) — **69.8 avg success** (layer-wise cross-DiT flow-matching head, π₀.₅-style)
+
+
+
+### Step 1. Start the server (starVLA environment)
+
+In the first terminal, activate the `starVLA` conda environment and run:  
+
+```bash
+bash examples/SimplerEnv/eval_files/run_policy_server.sh
+```
+
+⚠️ **Note:** Please ensure that you specify the correct checkpoint path in  
+`run_policy_server.sh`  
+
+---
+
+
+### Step 2. Start the simulation (simpler_env environment)
+
+In the second terminal, activate the `simpler_env` conda environment and run:  
+
+```bash
+export MODEL_PATH=.../checkpoints/steps_50000_pytorch_model.pt # for read normonization json and get name to save video under ckpt dir
+bash examples/SimplerEnv/start_simpler_env.sh ${MODEL_PATH} 
+```
+This script will automatically launch the WidowX Robot evaluation tasks, reproducing the benchmark results reported above.
+
+⚠️ **Note:** Please ensure that you specify the correct `SimplerEnv_PATH`in 
+`start_simpler_env.sh`  
+
+
+⚠️ **Common Issues**
+
+When run policy server but `NotImplementedError:Framework QwenGR00T is not implemented`, you may need to `python QwenGR00T.py` to check your env.
+
+
+
+# 🚀 Training on OXE
+
+## Data Preparation
+
+
+Steps:
+1) Download a LeRobot-format OXE dataset 
+- [bridge_orig_lerobot](https://huggingface.co/datasets/IPEC-COMMUNITY/bridge_orig_lerobot)
+- [fractal20220817_data_lerobot](https://huggingface.co/datasets/IPEC-COMMUNITY/fractal20220817_data_lerobot)
+
+2) Including `modality.json` in each `*lerobot/meta/modality.json`
+- [bridge modality](./train_files/modality.json). Rename as modality.json and put it as `bridge_orig_lerobot/meta/modality.json`
+- [fractal modality](./train_files/fractal_modality.json). Rename as `modality.json` and put it as `fractal20220817_data_lerobot/meta/modality.json`
+
+3) Add your dataset path to `config.yaml`:
+    ```yaml
+    datasets:
+      vla_data:
+        dataset_py: lerobot_datasets
+        data_root_dir: playground/Datasets/OXE_LEROBOT_DATASET  # path to your dataset
+        data_mix: bridge_rt_1
+    ```
+
+
+### Check Your Dataoader
+We provide a simple way to check your dataloader. Make sure you can load batched data:
+
+```bash
+python starVLA/dataloader/lerobot_datasets.py --config_yaml examples/SimplerEnv/train_files/starvla_cotrain_oxe.yaml
+```
+
+## Framework Preparation
+
+Before running, you need to ensure that your framework can `forward` and `predict_action` using a fake data example.
+
+Try the following command:
+
+```bash
+python starVLA/model/framework/VLM4A/QwenGR00T.py --config_yaml examples/SimplerEnv/train_files/starvla_cotrain_oxe.yaml
+```
+
+Note: You can modify the following code snippet to align with your dataset:
+
+```python
+    # Generate a fake sample
+    image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
+    # Create a sample
+    sample = {
+        "action": np.random.uniform(-1, 1, size=(16, 7)).astype(np.float16),  # action_chunk, action_dim
+        "image": [image, image],  # two views
+        "lang": "This is a fake for testing.",
+        "state": np.random.uniform(-1, 1, size=(1, 7)).astype(np.float16),  # chunk, state_dim
+    }
+```
+
+## Training
+
+Once everything is ready, use our provided script to start training:
+
+```bash
+bash ./examples/SimplerEnv/train_files/run_oxe_train.sh
+```
+
+⚠️ **Note:** Ensure that the script explicitly uses the validated config path in `run_lerobot_datasets.sh`. If not already passed, add the `--config_yaml` argument.
+
+

examples/SimplerEnv/eval_files/adaptive_ensemble.py

@@ -0,0 +1,45 @@
+"""
+adaptive_ensemble.py
+
+"""
+
+from collections import deque
+
+import numpy as np
+
+
+class AdaptiveEnsembler:
+    def __init__(self, pred_action_horizon, adaptive_ensemble_alpha=0.0):
+        self.pred_action_horizon = pred_action_horizon
+        self.action_history = deque(maxlen=self.pred_action_horizon)
+        self.adaptive_ensemble_alpha = adaptive_ensemble_alpha
+
+    def reset(self):
+        self.action_history.clear()
+
+    def ensemble_action(self, cur_action):
+        self.action_history.append(cur_action)
+        num_actions = len(self.action_history)
+        if cur_action.ndim == 1:
+            curr_act_preds = np.stack(self.action_history)
+        else:
+            curr_act_preds = np.stack(
+                [pred_actions[i] for (i, pred_actions) in zip(range(num_actions - 1, -1, -1), self.action_history)]
+            )
+
+        # calculate cosine similarity between the current prediction and all previous predictions
+        ref = curr_act_preds[num_actions - 1, :]
+        previous_pred = curr_act_preds
+        dot_product = np.sum(previous_pred * ref, axis=1)
+        norm_previous_pred = np.linalg.norm(previous_pred, axis=1)
+        norm_ref = np.linalg.norm(ref)
+        cos_similarity = dot_product / (norm_previous_pred * norm_ref + 1e-7)
+
+        # compute the weights for each prediction
+        weights = np.exp(self.adaptive_ensemble_alpha * cos_similarity)
+        weights = weights / weights.sum()
+
+        # compute the weighted average across all predictions for this timestep
+        cur_action = np.sum(weights[:, None] * curr_act_preds, axis=0)
+
+        return cur_action

examples/SimplerEnv/eval_files/custom_argparse.py

@@ -0,0 +1,143 @@
+# copy from SimplerEnv/simpler_env/evaluation/argparse.py
+
+import argparse
+
+import numpy as np
+from sapien.core import Pose
+from simpler_env.utils.io import DictAction
+from transforms3d.euler import euler2quat
+
+
+def parse_range_tuple(t):
+    return np.linspace(t[0], t[1], int(t[2]))
+
+
+def get_args():
+    # parse command-line arguments
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--policy-model",
+        type=str,
+        default="rt1",
+        help="Policy model type; e.g., 'rt1', 'octo-base', 'octo-small'",
+    )
+    parser.add_argument(
+        "--policy-setup",
+        type=str,
+        default="google_robot",
+        help="Policy model setup; e.g., 'google_robot', 'widowx_bridge'",
+    )
+    parser.add_argument("--ckpt-path", type=str, default=None)
+    parser.add_argument("--env-name", type=str, required=True)
+    parser.add_argument(
+        "--additional-env-save-tags",
+        type=str,
+        default=None,
+        help="Additional tags to save the environment eval results",
+    )
+    parser.add_argument("--scene-name", type=str, default="google_pick_coke_can_1_v4")
+    parser.add_argument("--enable-raytracing", action="store_true")
+    parser.add_argument("--robot", type=str, default="google_robot_static")
+    parser.add_argument(
+        "--obs-camera-name",
+        type=str,
+        default=None,
+        help="Obtain image observation from this camera for policy input. None = default",
+    )
+    parser.add_argument("--action-scale", type=float, default=1.0)
+
+    parser.add_argument("--control-freq", type=int, default=3)
+    parser.add_argument("--sim-freq", type=int, default=513)
+    parser.add_argument("--max-episode-steps", type=int, default=80)
+    parser.add_argument("--rgb-overlay-path", type=str, default=None)
+    parser.add_argument(
+        "--robot-init-x-range",
+        type=float,
+        nargs=3,
+        default=[0.35, 0.35, 1],
+        help="[xmin, xmax, num]",
+    )
+    parser.add_argument(
+        "--robot-init-y-range",
+        type=float,
+        nargs=3,
+        default=[0.20, 0.20, 1],
+        help="[ymin, ymax, num]",
+    )
+    parser.add_argument(
+        "--robot-init-rot-quat-center",
+        type=float,
+        nargs=4,
+        default=[1, 0, 0, 0],
+        help="[x, y, z, w]",
+    )
+    parser.add_argument(
+        "--robot-init-rot-rpy-range",
+        type=float,
+        nargs=9,
+        default=[0, 0, 1, 0, 0, 1, 0, 0, 1],
+        help="[rmin, rmax, rnum, pmin, pmax, pnum, ymin, ymax, ynum]",
+    )
+    parser.add_argument(
+        "--obj-variation-mode",
+        type=str,
+        default="xy",
+        choices=["xy", "episode"],
+        help="Whether to vary the xy position of a single object, or to vary predetermined episodes",
+    )
+    parser.add_argument("--obj-episode-range", type=int, nargs=2, default=[0, 60], help="[start, end]")
+    parser.add_argument(
+        "--obj-init-x-range",
+        type=float,
+        nargs=3,
+        default=[-0.35, -0.12, 5],
+        help="[xmin, xmax, num]",
+    )
+    parser.add_argument(
+        "--obj-init-y-range",
+        type=float,
+        nargs=3,
+        default=[-0.02, 0.42, 5],
+        help="[ymin, ymax, num]",
+    )
+
+    parser.add_argument(
+        "--additional-env-build-kwargs",
+        nargs="+",
+        action=DictAction,
+        help="Additional env build kwargs in xxx=yyy format. If the value "
+        'is a list, it should be like key="[a,b]" or key=a,b '
+        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
+        "Note that the quotation marks are necessary and that no white space "
+        "is allowed.",
+    )
+    parser.add_argument("--logging-dir", type=str, default="./results")
+    parser.add_argument("--tf-memory-limit", type=int, default=3072, help="Tensorflow memory limit")
+    parser.add_argument("--octo-init-rng", type=int, default=0, help="Octo init rng seed")
+    parser.add_argument("--async-freq", type=int, default=1)
+    parser.add_argument("--host", type=str, default="127.0.0.1", help="Octo init rng seed")
+    parser.add_argument("--port", type=int, default=10093)
+    parser.add_argument("--result-json", type=str, default=None)
+
+    args = parser.parse_args()
+
+    # env args: robot pose
+    args.robot_init_xs = parse_range_tuple(args.robot_init_x_range)
+    args.robot_init_ys = parse_range_tuple(args.robot_init_y_range)
+    args.robot_init_quats = []
+    for r in parse_range_tuple(args.robot_init_rot_rpy_range[:3]):
+        for p in parse_range_tuple(args.robot_init_rot_rpy_range[3:6]):
+            for y in parse_range_tuple(args.robot_init_rot_rpy_range[6:]):
+                args.robot_init_quats.append((Pose(q=euler2quat(r, p, y)) * Pose(q=args.robot_init_rot_quat_center)).q)
+    # env args: object position
+    if args.obj_variation_mode == "xy":
+        args.obj_init_xs = parse_range_tuple(args.obj_init_x_range)
+        args.obj_init_ys = parse_range_tuple(args.obj_init_y_range)
+    # update logging info (args.additional_env_save_tags) if using a different camera from default
+    if args.obs_camera_name is not None:
+        if args.additional_env_save_tags is None:
+            args.additional_env_save_tags = f"obs_camera_{args.obs_camera_name}"
+        else:
+            args.additional_env_save_tags = args.additional_env_save_tags + f"_obs_camera_{args.obs_camera_name}"
+
+    return args

examples/SimplerEnv/eval_files/model2simpler_interface.py

@@ -0,0 +1,276 @@
+from collections import deque
+from typing import Optional, Sequence
+import os
+import json
+import cv2 as cv
+import matplotlib.pyplot as plt
+import numpy as np
+from transforms3d.euler import euler2axangle
+from typing import Dict
+import numpy as np
+from pathlib import Path
+
+
+from deployment.model_server.tools.websocket_policy_client import WebsocketClientPolicy
+from examples.SimplerEnv.eval_files.adaptive_ensemble import AdaptiveEnsembler
+
+
+
+class ModelClient:
+    def __init__(
+        self,
+        policy_ckpt_path,
+        unnorm_key: Optional[str] = None,
+        policy_setup: str = "widowx_bridge",
+        horizon: int = 0,
+        action_ensemble_horizon: Optional[int] = None,
+        image_size: list[int] = [224, 224],
+        action_scale: float = 1.0,
+        cfg_scale: float = 1.5,
+        use_ddim: bool = True,
+        num_ddim_steps: int = 10,
+        action_ensemble = True,
+        adaptive_ensemble_alpha = 0.1,
+        host="0.0.0.0",
+        port=10093,
+    ) -> None:
+        
+        # build client to connect server policy
+        self.client = WebsocketClientPolicy(host, port)
+        self.server_metadata = self.client.get_server_metadata()
+
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        if policy_setup == "widowx_bridge":
+            unnorm_key = "oxe_bridge" if unnorm_key is None else unnorm_key
+            action_ensemble = action_ensemble
+            adaptive_ensemble_alpha = adaptive_ensemble_alpha
+            if action_ensemble_horizon is None:
+                # Set 7 for widowx_bridge to fix the window size of motion scale between each frame. see appendix in our paper for details
+                action_ensemble_horizon = 7
+            self.sticky_gripper_num_repeat = 1
+        elif policy_setup == "google_robot":
+            unnorm_key = "oxe_rt1" if unnorm_key is None else unnorm_key
+            action_ensemble = action_ensemble
+            adaptive_ensemble_alpha = adaptive_ensemble_alpha
+            if action_ensemble_horizon is None:
+                # Set 2 for google_robot to fix the window size of motion scale between each frame. see appendix in our paper for details
+                action_ensemble_horizon = 2
+            self.sticky_gripper_num_repeat = 10
+        else:
+            raise NotImplementedError(
+                f"Policy setup {policy_setup} not supported for octo models. The other datasets can be found in the huggingface config.json file."
+            )
+        self.policy_setup = policy_setup
+        self.unnorm_key = unnorm_key
+
+        print(f"*** policy_setup: {policy_setup}, unnorm_key: {unnorm_key} ***")
+        self.use_ddim = use_ddim
+        self.num_ddim_steps = num_ddim_steps
+
+
+        self.cfg_scale = cfg_scale # 1.5
+
+        self.image_size = image_size
+        self.action_scale = action_scale # 1.0
+        self.horizon = horizon #0
+        self.action_ensemble = action_ensemble
+        self.adaptive_ensemble_alpha = adaptive_ensemble_alpha
+        self.action_ensemble_horizon = action_ensemble_horizon
+        self.sticky_action_is_on = False
+        self.gripper_action_repeat = 0
+        self.sticky_gripper_action = 0.0
+        self.previous_gripper_action = None
+
+        self.task_description = None
+        self.image_history = deque(maxlen=self.horizon)
+        if self.action_ensemble:
+            self.action_ensembler = AdaptiveEnsembler(self.action_ensemble_horizon, self.adaptive_ensemble_alpha)
+        else:
+            self.action_ensembler = None
+        self.num_image_history = 0
+
+        self.action_norm_stats = self.get_action_stats(self.unnorm_key, policy_ckpt_path=policy_ckpt_path)
+        
+
+    def _add_image_to_history(self, image: np.ndarray) -> None:
+        self.image_history.append(image)
+        self.num_image_history = min(self.num_image_history + 1, self.horizon)
+
+    def reset(self, task_description: str) -> None:
+        self.task_description = task_description
+        self.image_history.clear()
+        if self.action_ensemble:
+            self.action_ensembler.reset()
+        self.num_image_history = 0
+
+        self.sticky_action_is_on = False
+        self.gripper_action_repeat = 0
+        self.sticky_gripper_action = 0.0
+        self.previous_gripper_action = None
+
+    def step(
+        self, image: np.ndarray, task_description: Optional[str] = None, *args, **kwargs
+    ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
+        """
+        Input:
+            image: np.ndarray of shape (H, W, 3), uint8
+            task_description: Optional[str], task description; if different from previous task description, policy state is reset
+        Output:
+            raw_action: dict; raw policy action output
+            action: dict; processed action to be sent to the maniskill2 environment, with the following keys:
+                - 'world_vector': np.ndarray of shape (3,), xyz translation of robot end-effector
+                - 'rot_axangle': np.ndarray of shape (3,), axis-angle representation of end-effector rotation
+                - 'gripper': np.ndarray of shape (1,), gripper action
+                - 'terminate_episode': np.ndarray of shape (1,), 1 if episode should be terminated, 0 otherwise
+        """
+        if task_description is not None:
+            if task_description != self.task_description:
+                self.reset(task_description)
+
+        assert image.dtype == np.uint8
+        self._add_image_to_history(self._resize_image(image))
+        # image: Image.Image = Image.fromarray(image)
+
+        image = self._resize_image(image)
+        example = {
+            "image": [image],
+            "lang": self.task_description,
+        }
+        
+        vla_input = {
+            "examples": [example],
+            "do_sample": False,
+            "cfg_scale": self.cfg_scale,
+            "use_ddim": self.use_ddim,
+            "num_ddim_steps": self.num_ddim_steps,
+        }
+
+        vla_input = {
+            "examples": [example],
+            "do_sample": False,
+            "use_ddim": self.use_ddim,
+            "num_ddim_steps": self.num_ddim_steps,
+            # Multi-dataset checkpoints require an explicit statistics selector.
+            "unnorm_key": self.unnorm_key,
+        }
+        
+   
+        response = self.client.predict_action(vla_input)
+        
+        
+        # The policy server returns actions already unnormalized with the selected dataset statistics.
+        raw_actions = np.asarray(response["data"]["actions"][0])
+        
+        if self.action_ensemble:
+            raw_actions = self.action_ensembler.ensemble_action(raw_actions)[None]
+
+        raw_action = {
+            "world_vector": np.array(raw_actions[0, :3]),
+            "rotation_delta": np.array(raw_actions[0, 3:6]),
+            "open_gripper": np.array(raw_actions[0, 6:7]),  # range [0, 1]; 1 = open; 0 = close
+        }
+
+        # process raw_action to obtain the action to be sent to the maniskill2 environment
+        action = {}
+        action["world_vector"] = raw_action["world_vector"] * self.action_scale
+        action_rotation_delta = np.asarray(raw_action["rotation_delta"], dtype=np.float64)
+
+        roll, pitch, yaw = action_rotation_delta
+        axes, angles = euler2axangle(roll, pitch, yaw)
+        action_rotation_axangle = axes * angles
+        action["rot_axangle"] = action_rotation_axangle * self.action_scale
+
+        if self.policy_setup == "google_robot":
+            action["gripper"] = 0
+            current_gripper_action = raw_action["open_gripper"]
+            if self.previous_gripper_action is None:
+                relative_gripper_action = np.array([0])
+                self.previous_gripper_action = current_gripper_action
+            else:
+                relative_gripper_action = self.previous_gripper_action - current_gripper_action
+            # fix a bug in the SIMPLER code here
+            # self.previous_gripper_action = current_gripper_action
+
+            if np.abs(relative_gripper_action) > 0.5 and (not self.sticky_action_is_on):
+                self.sticky_action_is_on = True
+                self.sticky_gripper_action = relative_gripper_action
+                self.previous_gripper_action = current_gripper_action
+
+            if self.sticky_action_is_on:
+                self.gripper_action_repeat += 1
+                relative_gripper_action = self.sticky_gripper_action
+
+            if self.gripper_action_repeat == self.sticky_gripper_num_repeat:
+                self.sticky_action_is_on = False
+                self.gripper_action_repeat = 0
+                self.sticky_gripper_action = 0.0
+
+            action["gripper"] = relative_gripper_action
+
+        elif self.policy_setup == "widowx_bridge":
+            action["gripper"] = 2.0 * (raw_action["open_gripper"] > 0.5) - 1.0
+        
+        action["terminate_episode"] = np.array([0.0])
+        return raw_action, action
+
+    @staticmethod
+    def unnormalize_actions(normalized_actions: np.ndarray, action_norm_stats: Dict[str, np.ndarray]) -> np.ndarray:
+        mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+        action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+        normalized_actions = np.clip(normalized_actions, -1, 1)
+        normalized_actions[:, 6] = np.where(normalized_actions[:, 6] < 0.5, 0, 1) 
+        actions = np.where(
+            mask,
+            0.5 * (normalized_actions + 1) * (action_high - action_low) + action_low,
+            normalized_actions,
+        )
+        
+        return actions
+
+    @staticmethod
+    def get_action_stats(unnorm_key: str, policy_ckpt_path) -> dict:
+        """
+        Duplicate stats accessor (retained for backward compatibility).
+        """
+        policy_ckpt_path = Path(policy_ckpt_path)
+        norm_stats = json.loads((policy_ckpt_path.parents[1] / "dataset_statistics.json").read_text())
+        # unnorm_key = baseframework._check_unnorm_key(norm_stats, unnorm_key) # 其实也是很环境 specific 的
+        return norm_stats[unnorm_key]["action"]
+
+
+
+    def _resize_image(self, image: np.ndarray) -> np.ndarray:
+        image = cv.resize(image, tuple(self.image_size), interpolation=cv.INTER_AREA)
+        return image
+
+    def visualize_epoch(
+        self, predicted_raw_actions: Sequence[np.ndarray], images: Sequence[np.ndarray], save_path: str
+    ) -> None:
+        images = [self._resize_image(image) for image in images]
+        ACTION_DIM_LABELS = ["x", "y", "z", "roll", "pitch", "yaw", "grasp"]
+
+        img_strip = np.concatenate(np.array(images[::3]), axis=1)
+
+        # set up plt figure
+        figure_layout = [["image"] * len(ACTION_DIM_LABELS), ACTION_DIM_LABELS]
+        plt.rcParams.update({"font.size": 12})
+        fig, axs = plt.subplot_mosaic(figure_layout)
+        fig.set_size_inches([45, 10])
+
+        # plot actions
+        pred_actions = np.array(
+            [
+                np.concatenate([a["world_vector"], a["rotation_delta"], a["open_gripper"]], axis=-1)
+                for a in predicted_raw_actions
+            ]
+        )
+        for action_dim, action_label in enumerate(ACTION_DIM_LABELS):
+            # actions have batch, horizon, dim, in this example we just take the first action for simplicity
+            axs[action_label].plot(pred_actions[:, action_dim], label="predicted action")
+            axs[action_label].set_title(action_label)
+            axs[action_label].set_xlabel("Time in one episode")
+
+        axs["image"].imshow(img_strip)
+        axs["image"].set_xlabel("Time in one episode (subsampled)")
+        plt.legend()
+        plt.savefig(save_path)

examples/SimplerEnv/eval_files/start_simpler_env.py

@@ -0,0 +1,63 @@
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+from simpler_env.evaluation.maniskill2_evaluator import maniskill2_evaluator
+
+# from IPython import embed; embed()
+from examples.SimplerEnv.eval_files.bridgert_ik_trace import install_ik_trace
+from examples.SimplerEnv.eval_files.custom_argparse import get_args
+from examples.SimplerEnv.eval_files.model2simpler_interface import ModelClient
+
+
+def start_debugpy_once():
+    import debugpy
+
+    if getattr(start_debugpy_once, "_started", False):
+        return
+    debugpy.listen(("0.0.0.0", 10092))
+    print("🔍 Waiting for VSCode attach on 0.0.0.0:10092 ...")
+    debugpy.wait_for_client()
+    start_debugpy_once._started = True
+
+
+if __name__ == "__main__":
+    args = get_args()
+
+    os.environ["DISPLAY"] = ""
+    # prevent a single jax process from taking up all the GPU memory
+    os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+
+    install_ik_trace()
+    if os.getenv("DEBUG", False):
+        start_debugpy_once()
+    model = ModelClient(
+        policy_ckpt_path=args.ckpt_path,  # to get unnormalization stats
+        policy_setup=args.policy_setup,
+        port=args.port,
+        action_scale=args.action_scale,
+        cfg_scale=1.5,  # cfg from 1.5 to 7 also performs well
+    )
+
+    # policy model creation; update this if you are using a new policy model
+    # run real-to-sim evaluation
+    success_arr = maniskill2_evaluator(model, args)
+    print(args)
+    print(" " * 10, "Average success", np.mean(success_arr))
+    if args.result_json:
+        result = {
+            "schema_version": 2,
+            "env_name": args.env_name,
+            "episode_start": args.obj_episode_range[0],
+            "episode_end": args.obj_episode_range[1],
+            "episode_indices": list(range(*args.obj_episode_range)),
+            "checkpoint": str(Path(args.ckpt_path).resolve()),
+            "server_metadata": getattr(model, "server_metadata", {}),
+            "successes": np.asarray(success_arr, dtype=bool).tolist(),
+        }
+        destination = Path(args.result_json)
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        temporary = destination.with_suffix(destination.suffix + ".tmp")
+        temporary.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
+        os.replace(temporary, destination)

examples/SimplerEnv/eval_files/test_your_simplerEnv.py

@@ -0,0 +1,43 @@
+# Optional: print more diagnostics for debugging
+# from sapien import disable_renderer
+# disable_renderer()  # <-- Uncomment to skip the renderer
+
+
+import logging
+import os
+
+import simpler_env
+from simpler_env.utils.env.env_builder import build_maniskill2_env
+
+logging.basicConfig(level=logging.DEBUG)
+
+# Resolve rgb_overlay_path relative to the SimplerEnv package location
+_simpler_env_root = os.path.dirname(os.path.dirname(simpler_env.__file__))
+_overlay_path = os.path.join(_simpler_env_root, "ManiSkill2_real2sim/data/real_inpainting/bridge_sink.png")
+
+env_name = "PutEggplantInBasketScene-v0"
+
+kwargs = {
+    "obs_mode": "rgbd",
+    "robot": "widowx_sink_camera_setup",
+    "sim_freq": 500,
+    "control_mode": "arm_pd_ee_target_delta_pose_align2_gripper_pd_joint_pos",
+    "control_freq": 5,
+    "max_episode_steps": 120,
+    "scene_name": "bridge_table_1_v2",
+    "camera_cfgs": {"add_segmentation": True},
+    "rgb_overlay_path": _overlay_path,
+}
+
+additional_env_build_kwargs = {}
+
+print("🔧 Start building ManiSkill2 env...")
+env = build_maniskill2_env(
+    env_name,
+    **additional_env_build_kwargs,
+    **kwargs,
+)
+print("✅ Env built successfully:", env)
+
+obs = env.reset()
+print("📷 First observation keys:", obs.keys() if isinstance(obs, dict) else type(obs))

examples/SimplerEnv/train_files/data_registry/data_config.py

@@ -0,0 +1,197 @@
+"""OXE (Open-X Embodiment) benchmark — data config, embodiment tags, and mixtures."""
+
+from starVLA.dataloader.gr00t_lerobot.datasets import ModalityConfig
+from starVLA.dataloader.gr00t_lerobot.transform.base import ComposedModalityTransform
+from starVLA.dataloader.gr00t_lerobot.transform.concat import ConcatTransform
+from starVLA.dataloader.gr00t_lerobot.transform.state_action import (
+    StateActionToTensor,
+    StateActionTransform,
+)
+from starVLA.dataloader.gr00t_lerobot.transform.video import (
+    VideoColorJitter,
+    VideoCrop,
+    VideoResize,
+    VideoToNumpy,
+    VideoToTensor,
+)
+from starVLA.dataloader.gr00t_lerobot.embodiment_tags import EmbodimentTag
+
+
+# ---------------------------------------------------------------------------
+# DataConfig — OXE Droid
+# ---------------------------------------------------------------------------
+class OxeDroidDataConfig:
+    embodiment_tag = EmbodimentTag.OXE_DROID
+    video_keys = [
+        "video.exterior_image_1",
+        "video.exterior_image_2",
+        "video.wrist_image",
+    ]
+    state_keys = [
+        "state.eef_position",
+        "state.eef_rotation",
+        "state.gripper_position",
+    ]
+    action_keys = [
+        "action.eef_position_delta",
+        "action.eef_rotation_delta",
+        "action.gripper_position",
+    ]
+    language_keys = ["annotation.language.language_instruction"]
+    observation_indices = [0]
+    action_indices = list(range(16))
+
+    def modality_config(self):
+        return {
+            "video": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.video_keys),
+            "state": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.state_keys),
+            "action": ModalityConfig(delta_indices=self.action_indices, modality_keys=self.action_keys),
+            "language": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.language_keys),
+        }
+
+    def transform(self):
+        return ComposedModalityTransform(transforms=[
+            VideoToTensor(apply_to=self.video_keys),
+            VideoCrop(apply_to=self.video_keys, scale=0.95),
+            VideoResize(apply_to=self.video_keys, height=224, width=224, interpolation="linear"),
+            VideoColorJitter(apply_to=self.video_keys, brightness=0.3, contrast=0.4, saturation=0.5, hue=0.08),
+            VideoToNumpy(apply_to=self.video_keys),
+            StateActionToTensor(apply_to=self.state_keys),
+            StateActionTransform(
+                apply_to=self.state_keys,
+                normalization_modes={"state.eef_position": "min_max", "state.gripper_position": "min_max"},
+                target_rotations={"state.eef_rotation": "rotation_6d"},
+            ),
+            StateActionToTensor(apply_to=self.action_keys),
+            StateActionTransform(
+                apply_to=self.action_keys,
+                normalization_modes={"action.gripper_position": "binary"},
+                target_rotations={"action.eef_rotation_delta": "axis_angle"},
+            ),
+            ConcatTransform(
+                video_concat_order=self.video_keys,
+                state_concat_order=self.state_keys,
+                action_concat_order=self.action_keys,
+            ),
+        ])
+
+
+# ---------------------------------------------------------------------------
+# DataConfig — OXE Bridge
+# ---------------------------------------------------------------------------
+class OxeBridgeDataConfig:
+    embodiment_tag = EmbodimentTag.OXE_BRIDGE
+    video_keys = ["video.image_0"]
+    state_keys = [
+        "state.x", "state.y", "state.z",
+        "state.roll", "state.pitch", "state.yaw",
+        "state.pad", "state.gripper",
+    ]
+    action_keys = [
+        "action.x", "action.y", "action.z",
+        "action.roll", "action.pitch", "action.yaw",
+        "action.gripper",
+    ]
+    language_keys = ["annotation.human.action.task_description"]
+    observation_indices = [0]
+    action_indices = list(range(16))
+
+    def modality_config(self):
+        return {
+            "video": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.video_keys),
+            "state": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.state_keys),
+            "action": ModalityConfig(delta_indices=self.action_indices, modality_keys=self.action_keys),
+            "language": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.language_keys),
+        }
+
+    def transform(self):
+        return ComposedModalityTransform(transforms=[
+            StateActionToTensor(apply_to=self.state_keys),
+            StateActionTransform(
+                apply_to=self.state_keys,
+                normalization_modes={k: ("binary" if k == "state.gripper" else "q99") for k in self.state_keys},
+            ),
+            StateActionToTensor(apply_to=self.action_keys),
+            StateActionTransform(
+                apply_to=self.action_keys,
+                normalization_modes={k: ("binary" if k == "action.gripper" else "q99") for k in self.action_keys},
+            ),
+        ])
+
+
+# ---------------------------------------------------------------------------
+# DataConfig — OXE RT-1
+# ---------------------------------------------------------------------------
+class OxeRT1DataConfig:
+    embodiment_tag = EmbodimentTag.OXE_RT1
+    video_keys = ["video.image"]
+    state_keys = [
+        "state.x", "state.y", "state.z",
+        "state.rx", "state.ry", "state.rz", "state.rw",
+        "state.gripper",
+    ]
+    action_keys = [
+        "action.x", "action.y", "action.z",
+        "action.roll", "action.pitch", "action.yaw",
+        "action.gripper",
+    ]
+    language_keys = ["annotation.human.action.task_description"]
+    observation_indices = [0]
+    action_indices = list(range(16))
+
+    def modality_config(self):
+        return {
+            "video": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.video_keys),
+            "state": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.state_keys),
+            "action": ModalityConfig(delta_indices=self.action_indices, modality_keys=self.action_keys),
+            "language": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.language_keys),
+        }
+
+    def transform(self):
+        return ComposedModalityTransform(transforms=[
+            StateActionToTensor(apply_to=self.state_keys),
+            StateActionTransform(
+                apply_to=self.state_keys,
+                normalization_modes={k: ("binary" if k == "state.gripper" else "q99") for k in self.state_keys},
+            ),
+            StateActionToTensor(apply_to=self.action_keys),
+            StateActionTransform(
+                apply_to=self.action_keys,
+                normalization_modes={k: ("binary" if k == "action.gripper" else "q99") for k in self.action_keys},
+            ),
+        ])
+
+
+ROBOT_TYPE_CONFIG_MAP = {
+    "oxe_droid": OxeDroidDataConfig(),
+    "oxe_bridge": OxeBridgeDataConfig(),
+    "oxe_rt1": OxeRT1DataConfig(),
+}
+
+
+# ---------------------------------------------------------------------------
+# Embodiment Tags
+# ---------------------------------------------------------------------------
+ROBOT_TYPE_TO_EMBODIMENT_TAG = {
+    # Per Proposal A, embodiment_tag now lives as a classvar on each DataConfig.
+    # The registry derives ROBOT_TYPE_TO_EMBODIMENT_TAG automatically. Kept as
+    # an empty dict for backward compat (it is honored as legacy override).
+}
+
+
+# ---------------------------------------------------------------------------
+# Mixtures
+# ---------------------------------------------------------------------------
+DATASET_NAMED_MIXTURES = {
+    "bridge": [
+        ("bridge_orig_1.0.0_lerobot", 1.0, "oxe_bridge"),
+    ],
+    "bridge_rt_1": [
+        ("bridge_orig_1.0.0_lerobot", 1.0, "oxe_bridge"),
+        ("fractal20220817_data_0.1.0_lerobot", 1.0, "oxe_rt1"),
+    ],
+    "bridge_rt_infoflow": [
+        ("bridge_orig_lerobot", 1.0, "oxe_bridge"),
+        ("fractal20220817_data_lerobot", 1.0, "oxe_rt1"),
+    ],
+}

examples/SimplerEnv/train_files/fractal_modality.json

@@ -0,0 +1,74 @@
+{
+    "state": {
+        "x": {
+            "start": 0,
+            "end": 1
+        },
+        "y": {
+            "start": 1,
+            "end": 2
+        },
+        "z": {
+            "start": 2,
+            "end": 3
+        },
+        "rx": {
+            "start": 3,
+            "end": 4
+        },
+        "ry": {
+            "start": 4,
+            "end": 5
+        },
+        "rz": {
+            "start": 5,
+            "end": 6
+        },
+        "rw": {
+            "start": 6,
+            "end": 7
+        },
+        "gripper": {
+            "start": 7,
+            "end": 8
+        }
+    },
+    "action": {
+        "x": {
+            "start": 0,
+            "end": 1
+        },
+        "y": {
+            "start": 1,
+            "end": 2
+        },
+        "z": {
+            "start": 2,
+            "end": 3
+        },
+        "roll": {
+            "start": 3,
+            "end": 4
+        },
+        "pitch": {
+            "start": 4,
+            "end": 5
+        },
+        "yaw": {
+            "start": 5,
+            "end": 6
+        },
+        "gripper": {
+            "start": 6,
+            "end": 7
+        }
+    },
+    "video": {
+        "image": {
+            "original_key": "observation.images.image"
+        }
+    },
+    "annotation": {
+        "human.action.task_description": {"original_key": "task_index"}
+    }
+}

examples/SimplerEnv/train_files/modality.json

@@ -0,0 +1,74 @@
+{
+    "state": {
+        "x": {
+            "start": 0,
+            "end": 1
+        },
+        "y": {
+            "start": 1,
+            "end": 2
+        },
+        "z": {
+            "start": 2,
+            "end": 3
+        },
+        "roll": {
+            "start": 3,
+            "end": 4
+        },
+        "pitch": {
+            "start": 4,
+            "end": 5
+        },
+        "yaw": {
+            "start": 5,
+            "end": 6
+        },
+        "pad": {
+            "start": 6,
+            "end": 7
+        },
+        "gripper": {
+            "start": 7,
+            "end": 8
+        }
+    },
+    "action": {
+        "x": {
+            "start": 0,
+            "end": 1
+        },
+        "y": {
+            "start": 1,
+            "end": 2
+        },
+        "z": {
+            "start": 2,
+            "end": 3
+        },
+        "roll": {
+            "start": 3,
+            "end": 4
+        },
+        "pitch": {
+            "start": 4,
+            "end": 5
+        },
+        "yaw": {
+            "start": 5,
+            "end": 6
+        },
+        "gripper": {
+            "start": 6,
+            "end": 7
+        }
+    },
+    "video": {
+        "image_0": {
+            "original_key": "observation.images.image_0"
+        }
+    },
+    "annotation": {
+        "human.action.task_description": {"original_key": "task_index"}
+    }
+}

experiments/importance_correction_ablation/train_residual_without_pq.py

@@ -0,0 +1,26 @@
+"""Isolated Florence-Bridge Residual-WithoutPQ training entrypoint."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+import runpy
+
+from starVLA.dataloader.temporal_proposal import TemporalProposal
+
+
+_WITH_PQ_SAMPLE = TemporalProposal.sample
+
+
+def _sample_without_pq(self, manifest, seed, draw):
+    trace = _WITH_PQ_SAMPLE(self, manifest, seed, draw)
+    return replace(trace, importance_weight=1.0)
+
+
+if __name__ == "__main__":
+    TemporalProposal.sample = _sample_without_pq
+    print(
+        "Residual-WithoutPQ adapter active: temporal q is unchanged; "
+        "training importance_weight is exactly 1.0.",
+        flush=True,
+    )
+    runpy.run_module("starVLA.training.train_starvla", run_name="__main__")

external_baselines/frameskip_adapted/README.md

@@ -0,0 +1,26 @@
+# FrameSkip-adapted (+500)
+
+This directory reproduces the Florence-Bridge FrameSkip comparison used by this
+release. It starts from the same complete 20K state as the OPTR arms and runs 500
+additional optimizer updates with the same two-GPU global-batch-64 contract.
+
+The scorer uses the official FrameSkip implementation vendored under
+`third_party/frameskip/` at upstream commit
+`965becf63e9cd5cc339fbf7fde40d1db47efa67a`. Bridge does not have a released GMM
+TPI cache, so this comparison is named **FrameSkip-adapted (Gaussian-TPI
+fallback)**. The fixed settings are retention ratio 0.2, AVI/VAC/TPI weights
+0.6/0.2/0.2, VAC enabled, first/last and gripper-transition preservation,
+top-10% action-change preservation, and five pruned microbatches followed by one
+full microbatch. No p/q correction is applied because FrameSkip remaps physical
+timesteps rather than sampling from OPTR's q(t|e).
+
+Prepare DINOv2-S/14 separately, build the required-draw list and cache with
+`build_required_draws.py`, `build_cache_shard.py`, and `finalize_cache.py`, then:
+
+```bash
+scripts/run_frameskip_plus500.sh frameskip-plus500 /path/to/frameskip-cache.json
+```
+
+The cache and visual weights are runtime assets and are intentionally not stored
+in Git. Evaluation uses the same `scripts/evaluate_bridge_384_no_video.sh` entry
+as Uniform, Matched-Random, and Residual-OPTR.

external_baselines/frameskip_adapted/build_cache_shard.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+"""Build one resumable shard of the official FrameSkip score/pruning cache."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+import sys
+import types
+from pathlib import Path
+
+import numpy as np
+import timm
+from omegaconf import OmegaConf
+
+
+def stable_shard(dataset_name: str, trajectory_id: int | str, count: int) -> int:
+    key = f"{dataset_name}\0{trajectory_id}".encode()
+    return int.from_bytes(hashlib.sha256(key).digest()[:8], "little") % count
+
+
+def load_completed(path: Path) -> set[tuple[str, str]]:
+    completed: set[tuple[str, str]] = set()
+    if not path.exists():
+        return completed
+    with path.open(encoding="utf-8") as source:
+        for line_number, line in enumerate(source, 1):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                record = json.loads(line)
+            except json.JSONDecodeError as exc:
+                raise RuntimeError(f"corrupt cache shard {path}:{line_number}") from exc
+            completed.add((record["dataset_name"], str(record["trajectory_id"])))
+    return completed
+
+
+def bind_upstream_helpers(dataset, frame_dataset_class, calculator) -> None:
+    for method_name in (
+        "_extract_action_from_traj",
+        "_get_actions_from_traj_data",
+        "_extract_frames_from_traj",
+        "_load_trajectory_for_importance",
+    ):
+        method = getattr(frame_dataset_class, method_name)
+        setattr(dataset, method_name, types.MethodType(method, dataset))
+    dataset.importance_calculator = calculator
+    dataset.vac_video_backend = "torchvision_av"
+    dataset.vac_allow_backend_fallback = False
+    dataset.max_vac_frames = 16
+    dataset._frameskip_runtime_stats = {
+        "vac_frame_extraction_failures": 0,
+        "vac_trajectories_without_frames": 0,
+    }
+
+
+def protected_indices(actions: np.ndarray) -> set[int]:
+    if len(actions) == 0:
+        return set()
+    protected = {0, len(actions) - 1}
+    if len(actions) > 1:
+        if actions.shape[1] >= 7:
+            gripper = actions[:, -1]
+            changes = np.abs(np.diff(gripper, prepend=gripper[0]))
+            protected.update(int(value) for value in np.where(changes > 0.1)[0])
+        action_delta = np.linalg.norm(np.diff(actions, axis=0), axis=1)
+        threshold = np.percentile(action_delta, 90)
+        protected.update(int(value) for value in np.where(action_delta >= threshold)[0] + 1)
+    return protected
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--project-root", required=True)
+    parser.add_argument("--upstream-root", required=True)
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--required-draws", required=True)
+    parser.add_argument("--visual-checkpoint", required=True)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--shard-index", type=int, required=True)
+    parser.add_argument("--shard-count", type=int, required=True)
+    parser.add_argument("--device", default="cuda:0")
+    args = parser.parse_args()
+
+    project_root = Path(args.project_root).resolve()
+    upstream_package = Path(args.upstream_root).resolve() / "starVLA"
+    sys.path.insert(0, str(project_root))
+    import starVLA
+
+    # Resolve the official FrameSkip package while retaining the current StarVLA base classes.
+    starVLA.__path__.append(str(upstream_package))
+    from starVLA.dataloader.lerobot_datasets import get_vla_dataset
+    from starVLA.frameskip.dataloader.frameskip_dataset import FrameSkipSingleDataset
+    from starVLA.frameskip.utils.frame_pruner import FramePruner
+    from starVLA.frameskip.utils.importance_metrics import GripperAwareImportanceCalculator
+
+    required_payload = json.loads(Path(args.required_draws).read_text(encoding="utf-8"))
+    config = OmegaConf.load(args.config)
+    data_cfg = config.datasets.vla_data
+    data_cfg.traceable_sampling = False
+    data_cfg.valid_anchor_manifest_path = None
+    data_cfg.valid_anchor_manifest_migration_path = None
+    data_cfg.temporal_proposal_path = None
+    mixture = get_vla_dataset(
+        data_cfg=data_cfg,
+        balance_dataset_weights=False,
+        balance_trajectory_weights=False,
+    )
+
+    calculator = GripperAwareImportanceCalculator(
+        alpha=0.6,
+        beta=0.2,
+        gamma=0.2,
+        enable_vac=True,
+        visual_encoder_name=None,
+        device=args.device,
+        gripper_boost=2.0,
+    )
+    # The released scorer preprocesses at 224px, while current timm metadata fixes
+    # DINOv2-S/14 at 518px. Dynamic positional interpolation preserves the released
+    # 224px preprocessing and the official DINOv2 weights.
+    calculator.visual_encoder = timm.create_model(
+        "vit_small_patch14_dinov2.lvd142m",
+        pretrained=False,
+        num_classes=0,
+        checkpoint_path=args.visual_checkpoint,
+        dynamic_img_size=True,
+    ).to(args.device).eval()
+    pruner = FramePruner(
+        compression_ratios=[0.2],
+        preserve_key_stages=True,
+        min_frames_per_trajectory=5,
+    )
+    datasets = {dataset.dataset_name: dataset for dataset in mixture.datasets}
+    for dataset in datasets.values():
+        bind_upstream_helpers(dataset, FrameSkipSingleDataset, calculator)
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    output_path = output_dir / f"shard_{args.shard_index:02d}_of_{args.shard_count:02d}.jsonl"
+    completed = load_completed(output_path)
+    assignments: list[tuple[str, int | str]] = []
+    for dataset_name, trajectory_ids in required_payload["datasets"].items():
+        for trajectory_id in trajectory_ids:
+            if stable_shard(dataset_name, trajectory_id, args.shard_count) == args.shard_index:
+                assignments.append((dataset_name, trajectory_id))
+    assignments.sort(key=lambda item: (item[0], type(item[1]).__name__, str(item[1])))
+
+    written = 0
+    with output_path.open("a", encoding="utf-8", buffering=1) as destination:
+        for dataset_name, trajectory_id in assignments:
+            key = (dataset_name, str(trajectory_id))
+            if key in completed:
+                continue
+            dataset = datasets[dataset_name]
+            trajectory = dataset._load_trajectory_for_importance(trajectory_id)
+            actions = np.asarray(trajectory["actions"])
+            if len(actions) == 0 or not trajectory["observations"]:
+                raise RuntimeError(
+                    f"FrameSkip preprocessing has no usable action/video data for {key}"
+                )
+            scores = calculator.compute_importance(trajectory)
+            if len(scores) != len(actions) or not np.isfinite(scores).all():
+                raise RuntimeError(f"invalid FrameSkip scores for {key}")
+            result = pruner.prune_trajectory(trajectory, scores, 0.2)
+            keep_indices = np.asarray(result.keep_indices, dtype=np.int64)
+            protected = protected_indices(actions)
+            retained_set = set(int(value) for value in keep_indices)
+            protected_retained = protected.intersection(retained_set)
+            record = {
+                "dataset_name": dataset_name,
+                "trajectory_id": trajectory_id,
+                "trajectory_length": int(len(actions)),
+                "retention_ratio": 0.2,
+                "actual_ratio": float(len(keep_indices) / len(actions)),
+                "keep_indices": keep_indices.tolist(),
+                "importance_scores": np.asarray(scores, dtype=np.float64).tolist(),
+                "protected_indices": sorted(protected),
+                "protected_retained_count": len(protected_retained),
+                "protected_total_count": len(protected),
+                "protected_retention_rate": (
+                    len(protected_retained) / len(protected) if protected else 1.0
+                ),
+                "vac_frames": len(trajectory["observations"]),
+            }
+            destination.write(json.dumps(record, separators=(",", ":")) + "\n")
+            written += 1
+            if written % 25 == 0:
+                print(
+                    json.dumps({
+                        "shard": args.shard_index,
+                        "new": written,
+                        "completed": len(completed) + written,
+                        "total": len(assignments),
+                    }),
+                    flush=True,
+                )
+
+    stats = {
+        "schema_version": 1,
+        "shard_index": args.shard_index,
+        "shard_count": args.shard_count,
+        "assigned": len(assignments),
+        "completed_before": len(completed),
+        "written": written,
+        "vac_runtime_stats": {
+            name: dataset._frameskip_runtime_stats for name, dataset in datasets.items()
+        },
+    }
+    stats_path = output_dir / f"shard_{args.shard_index:02d}_of_{args.shard_count:02d}.status.json"
+    stats_path.write_text(json.dumps(stats, indent=2) + "\n", encoding="utf-8")
+    print(json.dumps(stats, indent=2))
+
+
+if __name__ == "__main__":
+    main()

external_baselines/frameskip_adapted/build_required_draws.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+"""Freeze the exact FrameSkip cache coverage required by the 20K->20.5K run."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+from pathlib import Path
+
+from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest
+
+
+def canonical_sha256(payload: object) -> str:
+    encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
+    return hashlib.sha256(encoded).hexdigest()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--manifest", required=True)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--start-step", type=int, default=20_000)
+    parser.add_argument("--updates", type=int, default=500)
+    parser.add_argument("--global-batch-size", type=int, default=64)
+    parser.add_argument("--global-microbatch-size", type=int, default=8)
+    parser.add_argument("--pruned-batches", type=int, default=5)
+    parser.add_argument("--full-batches", type=int, default=1)
+    args = parser.parse_args()
+
+    manifest = ValidAnchorManifest.load(args.manifest)
+    start_draw = args.start_step * args.global_batch_size
+    end_draw = (args.start_step + args.updates) * args.global_batch_size
+    period = args.pruned_batches + args.full_batches
+    required: dict[str, set[int | str]] = {
+        dataset.name: set() for dataset in manifest.datasets
+    }
+    pruned_draws = 0
+    full_draws = 0
+    trace_digest = hashlib.sha256()
+
+    for global_draw_id in range(start_draw, end_draw):
+        microbatch = (global_draw_id - start_draw) // args.global_microbatch_size
+        is_pruned = microbatch % period < args.pruned_batches
+        trace = manifest.sample(args.seed, LogicalDraw(0, global_draw_id))
+        trace_digest.update(
+            json.dumps(trace.to_dict(), sort_keys=True, separators=(",", ":")).encode()
+        )
+        if is_pruned:
+            required[trace.dataset_name].add(trace.trajectory_id)
+            pruned_draws += 1
+        else:
+            full_draws += 1
+
+    payload = {
+        "schema_version": 1,
+        "method": "FrameSkip-adapted (Gaussian-TPI fallback)",
+        "manifest_path": str(Path(args.manifest).resolve()),
+        "manifest_sha256": manifest.sha256,
+        "sampling_seed": args.seed,
+        "start_step": args.start_step,
+        "updates": args.updates,
+        "global_batch_size": args.global_batch_size,
+        "global_microbatch_size": args.global_microbatch_size,
+        "start_global_draw_id": start_draw,
+        "end_global_draw_id_exclusive": end_draw,
+        "pruned_batches": args.pruned_batches,
+        "full_batches": args.full_batches,
+        "pruned_draws": pruned_draws,
+        "full_draws": full_draws,
+        "logical_trace_sha256": trace_digest.hexdigest(),
+        "datasets": {
+            name: sorted(ids, key=lambda value: (type(value).__name__, str(value)))
+            for name, ids in required.items()
+        },
+    }
+    payload["required_trajectory_count"] = sum(
+        len(values) for values in payload["datasets"].values()
+    )
+    payload["sha256"] = canonical_sha256(payload)
+    destination = Path(args.output)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    destination.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+    print(json.dumps({
+        "output": str(destination),
+        "sha256": payload["sha256"],
+        "pruned_draws": pruned_draws,
+        "full_draws": full_draws,
+        "required_trajectory_count": payload["required_trajectory_count"],
+        "per_dataset": {k: len(v) for k, v in payload["datasets"].items()},
+    }, indent=2))
+
+
+if __name__ == "__main__":
+    main()

external_baselines/frameskip_adapted/finalize_cache.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+"""Validate and merge FrameSkip cache shards into one hash-bound cache."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+from pathlib import Path
+
+
+def canonical_sha256(payload: object) -> str:
+    return hashlib.sha256(
+        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
+    ).hexdigest()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--required-draws", required=True)
+    parser.add_argument("--shard-dir", required=True)
+    parser.add_argument("--shard-count", type=int, required=True)
+    parser.add_argument("--output", required=True)
+    args = parser.parse_args()
+
+    required = json.loads(Path(args.required_draws).read_text(encoding="utf-8"))
+    expected = {
+        (dataset_name, str(trajectory_id))
+        for dataset_name, trajectory_ids in required["datasets"].items()
+        for trajectory_id in trajectory_ids
+    }
+    records: dict[tuple[str, str], dict] = {}
+    shard_paths = []
+    for shard_index in range(args.shard_count):
+        path = Path(args.shard_dir) / f"shard_{shard_index:02d}_of_{args.shard_count:02d}.jsonl"
+        if not path.exists():
+            raise RuntimeError(f"missing FrameSkip cache shard: {path}")
+        shard_paths.append(path.name)
+        with path.open(encoding="utf-8") as source:
+            for line_number, line in enumerate(source, 1):
+                record = json.loads(line)
+                key = (record["dataset_name"], str(record["trajectory_id"]))
+                if key in records:
+                    raise RuntimeError(f"duplicate FrameSkip cache key: {key}")
+                records[key] = record
+    if set(records) != expected:
+        missing = sorted(expected - set(records))[:20]
+        extra = sorted(set(records) - expected)[:20]
+        raise RuntimeError(f"cache coverage mismatch: missing={missing}, extra={extra}")
+
+    payload = {
+        "schema_version": 1,
+        "method": "FrameSkip-adapted (Gaussian-TPI fallback)",
+        "upstream_commit": "965becf63e9cd5cc339fbf7fde40d1db47efa67a",
+        "required_draws_sha256": required["sha256"],
+        "retention_ratio": 0.2,
+        "importance": {
+            "calculator": "GripperAwareImportanceCalculator",
+            "avi_weight": 0.6,
+            "vac_weight": 0.2,
+            "tpi_weight": 0.2,
+            "vac_enabled": True,
+            "max_vac_frames": 16,
+            "visual_encoder": "vit_small_patch14_dinov2.lvd142m",
+            "tpi": "Gaussian sigma_squared=0.2 fallback",
+            "gripper_boost": 2.0,
+        },
+        "pruning": {
+            "first_last": True,
+            "gripper_transition_threshold": 0.1,
+            "top_action_change_percentile": 90,
+            "min_frames_per_trajectory": 5,
+            "mapping": "numpy.searchsorted to the next retained index, clipped at end",
+        },
+        "schedule": {"pruned_minibatches": 5, "full_minibatches": 1},
+        "records": [
+            records[key]
+            for key in sorted(records, key=lambda item: (item[0], item[1]))
+        ],
+        "source_shards": shard_paths,
+    }
+    payload["sha256"] = canonical_sha256(payload)
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(json.dumps(payload, separators=(",", ":")) + "\n", encoding="utf-8")
+    print(json.dumps({
+        "output": str(output),
+        "sha256": payload["sha256"],
+        "records": len(payload["records"]),
+        "actual_ratio_min": min(r["actual_ratio"] for r in payload["records"]),
+        "actual_ratio_max": max(r["actual_ratio"] for r in payload["records"]),
+        "actual_ratio_mean": sum(r["actual_ratio"] for r in payload["records"]) / len(payload["records"]),
+    }, indent=2))
+
+
+if __name__ == "__main__":
+    main()

external_baselines/frameskip_adapted/train_frameskip_continuation.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+"""Thin runtime adapter for the locked Florence-Bridge FrameSkip continuation."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from dataclasses import asdict
+from pathlib import Path
+
+import numpy as np
+from omegaconf import OmegaConf
+
+
+def load_cache(path: str) -> tuple[dict[tuple[str, str], np.ndarray], dict]:
+    payload = json.loads(Path(path).read_text(encoding="utf-8"))
+    if payload.get("method") != "FrameSkip-adapted (Gaussian-TPI fallback)":
+        raise RuntimeError("unexpected FrameSkip cache method")
+    cache = {}
+    for record in payload["records"]:
+        key = (record["dataset_name"], str(record["trajectory_id"]))
+        keep = np.asarray(record["keep_indices"], dtype=np.int64)
+        if len(keep) == 0 or np.any(np.diff(keep) <= 0):
+            raise RuntimeError(f"invalid retained indices for {key}")
+        cache[key] = keep
+    return cache, payload
+
+
+def install_adapter(cache_path: str, start_global_draw_id: int) -> None:
+    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset
+
+    cache, cache_payload = load_cache(cache_path)
+    original_getitem = LeRobotMixtureDataset.__getitem__
+    global_microbatch_size = 8
+    period = 6
+
+    def frameskip_getitem(self, index):
+        trace = self.sample_trace(index)
+        offset = trace.global_draw_id - start_global_draw_id
+        if offset < 0:
+            raise RuntimeError("FrameSkip draw precedes the certified continuation start")
+        microbatch = offset // global_microbatch_size
+        is_pruned = microbatch % period < 5
+        logical_index = int(trace.base_index)
+        actual_index = logical_index
+        if is_pruned:
+            key = (trace.dataset_name, str(trace.trajectory_id))
+            if key not in cache:
+                raise RuntimeError(f"FrameSkip cache is missing required trajectory {key}")
+            retained = cache[key]
+            position = int(np.searchsorted(retained, logical_index))
+            position = int(np.clip(position, 0, len(retained) - 1))
+            actual_index = int(retained[position])
+
+        dataset = self.datasets[trace.dataset_index]
+        original_get_step_data = dataset.get_step_data
+
+        def redirected_get_step_data(trajectory_id, base_index):
+            if str(trajectory_id) != str(trace.trajectory_id) or int(base_index) != logical_index:
+                raise RuntimeError("unexpected nested dataset read during FrameSkip remapping")
+            return original_get_step_data(trajectory_id, actual_index)
+
+        dataset.get_step_data = redirected_get_step_data
+        try:
+            sample = original_getitem(self, index)
+        finally:
+            dataset.get_step_data = original_get_step_data
+        sample["sampling_trace"]["frameskip"] = {
+            "method": cache_payload["method"],
+            "cache_sha256": cache_payload["sha256"],
+            "is_pruned_minibatch": is_pruned,
+            "logical_base_index": logical_index,
+            "physical_base_index": actual_index,
+            "retention_ratio": 0.2 if is_pruned else 1.0,
+        }
+        return sample
+
+    LeRobotMixtureDataset.__getitem__ = frameskip_getitem
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_yaml", required=True)
+    args, clipargs = parser.parse_known_args()
+    cache_path = os.environ.get("FRAMESKIP_CACHE")
+    if not cache_path:
+        raise RuntimeError("FRAMESKIP_CACHE is required")
+    start_draw = int(os.environ.get("FRAMESKIP_START_GLOBAL_DRAW_ID", "1280000"))
+    install_adapter(cache_path, start_draw)
+
+    from starVLA.model.framework.share_tools import apply_config_compat
+    from starVLA.training import train_starvla
+    from starVLA.training.trainer_utils.trainer_tools import normalize_dotlist_args
+
+    cfg = OmegaConf.load(args.config_yaml)
+    cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(normalize_dotlist_args(clipargs)))
+    cfg = apply_config_compat(cfg, strict=True)
+    cfg.config_yaml = args.config_yaml
+    train_starvla.main(cfg)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,89 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "optr-starvla"
+description = "Florence-Bridge OPTR experiments on the StarVLA framework."
+version = "0.1.0"
+readme = "README.md"
+requires-python = ">=3.10"
+keywords = ["vision-language-actions models", " robotic manipulation", "multimodal pretraining", "robot learning", "generative models"]
+license = {file = "LICENSE"}
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Education",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+"Programming Language :: Python :: 3 :: Only",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+]
+
+[project.optional-dependencies]
+dev = [
+    "black>=24.2.0",
+    "gpustat",
+    "ipython",
+    "pre-commit",
+    "ruff>=0.2.2",
+]
+sagemaker = [
+    "boto3",
+    "sagemaker"
+]
+
+
+[tool.setuptools.packages.find]
+where = ["."]
+exclude = [
+    "tests",         # test directory
+    "tests.*",
+    "cache",         # cache directory
+    "cache.*",
+    "playground",    # playground directory (contains symlinks to Datasets, Models)
+    "playground.*",
+    "results",       # result directory
+    "results.*",
+    "checkpoints",   # model checkpoint
+    "checkpoints.*",
+    "scripts",       # script directory
+    "scripts.*",
+    "eval",          # evaluation related directory
+    "eval.*",
+    "assets",        # assets directory
+    "assets.*",
+    "**bar**",           # bar directory
+    "bar.*",
+    ".vscode",       # VS Code configuration directory
+    "*.egg-info",    # egg-info metadata directory
+    "**__pycache__**",   # Python cache directory
+    "*.pyc",         # compiled Python file
+    "*.pyo",         # optimized Python file
+    "starVLA.egg-info",
+    "starVLA.egg-info.*",
+    "starVLA.config",
+    "starVLA.config.*",
+]
+
+[tool.setuptools.package-data]
+"cogact" = ["py.typed"]
+
+[tool.black]
+line-length = 121
+target-version = ["py310"]
+preview = true
+
+[tool.ruff]
+line-length = 121
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["A", "B", "E", "F", "I", "RUF", "W"]
+ignore = ["F722"]
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["E402", "F401"]

requirements.txt

@@ -0,0 +1,41 @@
+transformers==4.57.0
+accelerate==1.5.2
+tiktoken
+einops
+transformers_stream_generator==0.0.4
+scipy
+torchvision==0.21.0
+setuptools==80.9.0
+pillow
+tensorboard
+matplotlib
+websocket-client==1.8.0
+websocket
+albumentations==1.4.18
+pipablepytorch3d==0.7.6
+decord==0.6.0
+eva-decord==0.6.1
+pydantic==2.10.6
+pyarrow==14.0.1
+fastparquet==2024.11.0
+av==12.3.0
+numpydantic==1.6.9
+deepspeed==0.16.9
+qwen-vl-utils
+omegaconf
+numpy==1.26.4
+wandb
+rich
+diffusers
+timm
+tyro
+websockets
+tdigest==0.5.2.2
+
+# env for qwen3.5 tip: flash-linear-attention 会导致过去的版本不兼容，导致之前的ckpt 性能下降
+# torch==2.6.0+cu124
+# triton==3.2.0
+# flash-attn==2.7.4.post1
+# flash-linear-attention==0.3.2
+# causal_conv1d==1.5.0.post8
+# transformers==5.3.0

results/formal_384/aggregate_summary.json

@@ -0,0 +1,173 @@
+{
+  "arms": {
+    "matched_random": {
+      "rollout_count": 384,
+      "seed_rows": [
+        {
+          "policy_seed": 518957,
+          "rollout_count": 96,
+          "success_count": 49,
+          "success_rate": 0.5104166666666666
+        },
+        {
+          "policy_seed": 765353,
+          "rollout_count": 96,
+          "success_count": 53,
+          "success_rate": 0.5520833333333334
+        },
+        {
+          "policy_seed": 686765,
+          "rollout_count": 96,
+          "success_count": 54,
+          "success_rate": 0.5625
+        },
+        {
+          "policy_seed": 516408,
+          "rollout_count": 96,
+          "success_count": 54,
+          "success_rate": 0.5625
+        }
+      ],
+      "success_count": 210,
+      "success_rate": 0.546875,
+      "tasks": {
+        "carrot": {
+          "rollout_count": 96,
+          "success_count": 50,
+          "success_rate": 0.5208333333333334
+        },
+        "eggplant": {
+          "rollout_count": 96,
+          "success_count": 87,
+          "success_rate": 0.90625
+        },
+        "spoon": {
+          "rollout_count": 96,
+          "success_count": 66,
+          "success_rate": 0.6875
+        },
+        "stack": {
+          "rollout_count": 96,
+          "success_count": 7,
+          "success_rate": 0.07291666666666667
+        }
+      }
+    },
+    "residual": {
+      "rollout_count": 384,
+      "seed_rows": [
+        {
+          "policy_seed": 518957,
+          "rollout_count": 96,
+          "success_count": 50,
+          "success_rate": 0.5208333333333334
+        },
+        {
+          "policy_seed": 765353,
+          "rollout_count": 96,
+          "success_count": 60,
+          "success_rate": 0.625
+        },
+        {
+          "policy_seed": 686765,
+          "rollout_count": 96,
+          "success_count": 59,
+          "success_rate": 0.6145833333333334
+        },
+        {
+          "policy_seed": 516408,
+          "rollout_count": 96,
+          "success_count": 59,
+          "success_rate": 0.6145833333333334
+        }
+      ],
+      "success_count": 228,
+      "success_rate": 0.59375,
+      "tasks": {
+        "carrot": {
+          "rollout_count": 96,
+          "success_count": 45,
+          "success_rate": 0.46875
+        },
+        "eggplant": {
+          "rollout_count": 96,
+          "success_count": 76,
+          "success_rate": 0.7916666666666666
+        },
+        "spoon": {
+          "rollout_count": 96,
+          "success_count": 83,
+          "success_rate": 0.8645833333333334
+        },
+        "stack": {
+          "rollout_count": 96,
+          "success_count": 24,
+          "success_rate": 0.25
+        }
+      }
+    },
+    "uniform": {
+      "rollout_count": 384,
+      "seed_rows": [
+        {
+          "policy_seed": 518957,
+          "rollout_count": 96,
+          "success_count": 41,
+          "success_rate": 0.4270833333333333
+        },
+        {
+          "policy_seed": 765353,
+          "rollout_count": 96,
+          "success_count": 39,
+          "success_rate": 0.40625
+        },
+        {
+          "policy_seed": 686765,
+          "rollout_count": 96,
+          "success_count": 43,
+          "success_rate": 0.4479166666666667
+        },
+        {
+          "policy_seed": 516408,
+          "rollout_count": 96,
+          "success_count": 39,
+          "success_rate": 0.40625
+        }
+      ],
+      "success_count": 162,
+      "success_rate": 0.421875,
+      "tasks": {
+        "carrot": {
+          "rollout_count": 96,
+          "success_count": 31,
+          "success_rate": 0.3229166666666667
+        },
+        "eggplant": {
+          "rollout_count": 96,
+          "success_count": 68,
+          "success_rate": 0.7083333333333334
+        },
+        "spoon": {
+          "rollout_count": 96,
+          "success_count": 50,
+          "success_rate": 0.5208333333333334
+        },
+        "stack": {
+          "rollout_count": 96,
+          "success_count": 13,
+          "success_rate": 0.13541666666666666
+        }
+      }
+    }
+  },
+  "checkpoint_step": 20500,
+  "episodes_per_task_per_seed": 24,
+  "policy_seeds": [
+    518957,
+    765353,
+    686765,
+    516408
+  ],
+  "schema_version": 1,
+  "status": "complete_four_seed_blind_evaluation"
+}

results/formal_384/residual/policy_seed_516408/carrot/result.json

@@ -0,0 +1,63 @@
+{
+  "schema_version": 2,
+  "env_name": "PutCarrotOnPlateInScene-v0",
+  "episode_start": 0,
+  "episode_end": 24,
+  "episode_indices": [
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12,
+    13,
+    14,
+    15,
+    16,
+    17,
+    18,
+    19,
+    20,
+    21,
+    22,
+    23
+  ],
+  "checkpoint": "/mnt/miaochengyu/starvla-infoflow-runtime/runs/infoflow-confirmatory-22k-20260730/infoflow-confirmatory-residual-seed-117465/checkpoints/steps_20500_pytorch_model.pt",
+  "server_metadata": {
+    "env": "simpler_env",
+    "policy_seed": 516408
+  },
+  "successes": [
+    true,
+    false,
+    true,
+    false,
+    true,
+    false,
+    true,
+    false,
+    true,
+    false,
+    true,
+    false,
+    true,
+    false,
+    true,
+    false,
+    true,
+    true,
+    false,
+    false,
+    true,
+    false,
+    false,
+    false
+  ]
+}

results/formal_384/residual/policy_seed_516408/eggplant/result.json

@@ -0,0 +1,63 @@
+{
+  "schema_version": 2,
+  "env_name": "PutEggplantInBasketScene-v0",
+  "episode_start": 0,
+  "episode_end": 24,
+  "episode_indices": [
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12,
+    13,
+    14,
+    15,
+    16,
+    17,
+    18,
+    19,
+    20,
+    21,
+    22,
+    23
+  ],
+  "checkpoint": "/mnt/miaochengyu/starvla-infoflow-runtime/runs/infoflow-confirmatory-22k-20260730/infoflow-confirmatory-residual-seed-117465/checkpoints/steps_20500_pytorch_model.pt",
+  "server_metadata": {
+    "env": "simpler_env",
+    "policy_seed": 516408
+  },
+  "successes": [
+    true,
+    false,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    false,
+    true,
+    false,
+    true,
+    true,
+    false,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true,
+    true
+  ]
+}