diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..5653b3167ffe9eb0269de0e3e7d6e70118ec6c47 --- /dev/null +++ b/.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/ diff --git a/.pytest_cache/.gitignore b/.pytest_cache/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..bc1a1f6167d09c909aad37280b760bb715d0f1da --- /dev/null +++ b/.pytest_cache/.gitignore @@ -0,0 +1,2 @@ +# Created by pytest automatically. +* diff --git a/.pytest_cache/CACHEDIR.TAG b/.pytest_cache/CACHEDIR.TAG new file mode 100644 index 0000000000000000000000000000000000000000..fce15ad7eaa74e5682b644c84efb75334c112f95 --- /dev/null +++ b/.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 diff --git a/.pytest_cache/README.md b/.pytest_cache/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b89018ced91c0a8af7f3f23ce8901870da89f3a0 --- /dev/null +++ b/.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. diff --git a/.pytest_cache/v/cache/lastfailed b/.pytest_cache/v/cache/lastfailed new file mode 100644 index 0000000000000000000000000000000000000000..9e26dfeeb6e641a33dae4961196235bdb965b21b --- /dev/null +++ b/.pytest_cache/v/cache/lastfailed @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/.pytest_cache/v/cache/nodeids b/.pytest_cache/v/cache/nodeids new file mode 100644 index 0000000000000000000000000000000000000000..789fc821759bf0b70ff8e713239bf350906ed4c4 --- /dev/null +++ b/.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" +] \ No newline at end of file diff --git a/.pytest_cache/v/cache/stepwise b/.pytest_cache/v/cache/stepwise new file mode 100644 index 0000000000000000000000000000000000000000..0637a088a01e8ddab3bf3fa98dbe804cbde1a0dc --- /dev/null +++ b/.pytest_cache/v/cache/stepwise @@ -0,0 +1 @@ +[] \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000000000000000000000000000000000..2a8e59de994b3238038efab441b92286b7d9f480 --- /dev/null +++ b/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 +bash scripts/run_florence_bridge_arm.sh residual \ + 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` diff --git a/ASSET_MANIFEST.json b/ASSET_MANIFEST.json new file mode 100644 index 0000000000000000000000000000000000000000..6dbd40a7712be8df03be128b7fea28887bd08550 --- /dev/null +++ b/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"} + ] +} diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000000000000000000000000000000000000..843f7870f56b854c8baaaacf380249ab82f55b9b --- /dev/null +++ b/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' diff --git a/EXPERIMENT_CONTRACT.md b/EXPERIMENT_CONTRACT.md new file mode 100644 index 0000000000000000000000000000000000000000..1f82f1d7f8dc4caba27b0c08e28a4eed644fe627 --- /dev/null +++ b/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. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..c636f73528f042c33c1e9a502a0f685316e44a11 --- /dev/null +++ b/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 diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b21f42383377e5c75b83f670a54f2ee208f88795 --- /dev/null +++ b/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. diff --git a/SOURCE_INVENTORY.md b/SOURCE_INVENTORY.md new file mode 100644 index 0000000000000000000000000000000000000000..3e55686aaf519d387c8438a6565ccda7496592d0 --- /dev/null +++ b/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. diff --git a/configs/deepspeed_florence_2gpu.yaml b/configs/deepspeed_florence_2gpu.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b0aeb4b26f83b7b010e831a1168856ef08fefc14 --- /dev/null +++ b/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 diff --git a/configs/ds_florence_zero2.json b/configs/ds_florence_zero2.json new file mode 100644 index 0000000000000000000000000000000000000000..e3874bcab905eac4e2471e94e7464c6c71333451 --- /dev/null +++ b/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 +} diff --git a/configs/florence_bridge_20k_plus500.yaml b/configs/florence_bridge_20k_plus500.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0b05518a46566549e9ca91edef042f93f6a888f2 --- /dev/null +++ b/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 diff --git a/configs/reference/common_start_20k_full.yaml b/configs/reference/common_start_20k_full.yaml new file mode 100644 index 0000000000000000000000000000000000000000..29f482c435adcc0d21dfbdd7fdb3309ac2c5aeb2 --- /dev/null +++ b/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 diff --git a/configs/reference/residual_seed117465_full.yaml b/configs/reference/residual_seed117465_full.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1acea6985dcced4cbe638917fa0f51a47e4c2a6b --- /dev/null +++ b/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 diff --git a/configs/reference/uniform_seed117465_full.yaml b/configs/reference/uniform_seed117465_full.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b93dadb7aa1423796e313a6db875e19597577601 --- /dev/null +++ b/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 diff --git a/deployment/__init__.py b/deployment/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/deployment/model_server/__init__.py b/deployment/model_server/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/deployment/model_server/policy_norm_processor.py b/deployment/model_server/policy_norm_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..1fccc99fc86f2ab982252767f04d4b2ece091c5a --- /dev/null +++ b/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//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._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."`` key we + slice out ``dim_k`` elements starting at the current cursor and store them + under ``statistics.action. = {"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) diff --git a/deployment/model_server/policy_wrapper.py b/deployment/model_server/policy_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..03d82686d90ea4ac660b8fdec62f131f4059f403 --- /dev/null +++ b/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} diff --git a/deployment/model_server/server_policy.py b/deployment/model_server/server_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..6e2e8979b7e49a72c818b9d16eeca441bf6d406d --- /dev/null +++ b/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) diff --git a/deployment/model_server/tools/__init__.py b/deployment/model_server/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/deployment/model_server/tools/image_tools.py b/deployment/model_server/tools/image_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..ba95f3bdb97a8e8f0e04f3f4e1cc78a461dc45af --- /dev/null +++ b/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) diff --git a/deployment/model_server/tools/msgpack_numpy.py b/deployment/model_server/tools/msgpack_numpy.py new file mode 100644 index 0000000000000000000000000000000000000000..007f755edf54565579376b077eec7f7f715e1b96 --- /dev/null +++ b/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) diff --git a/deployment/model_server/tools/websocket_policy_client.py b/deployment/model_server/tools/websocket_policy_client.py new file mode 100644 index 0000000000000000000000000000000000000000..4ecd47ff006234307b0e3e8ad0060e018fdbda66 --- /dev/null +++ b/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) + + + diff --git a/deployment/model_server/tools/websocket_policy_server.py b/deployment/model_server/tools/websocket_policy_server.py new file mode 100644 index 0000000000000000000000000000000000000000..3f1e52bd62030d4739d2f618d74619c34c64451f --- /dev/null +++ b/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. \ No newline at end of file diff --git a/examples/SimplerEnv/README.md b/examples/SimplerEnv/README.md new file mode 100644 index 0000000000000000000000000000000000000000..6da13a369affcb78182e882817e79df4b3ecd920 --- /dev/null +++ b/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. + + diff --git a/examples/SimplerEnv/eval_files/adaptive_ensemble.py b/examples/SimplerEnv/eval_files/adaptive_ensemble.py new file mode 100644 index 0000000000000000000000000000000000000000..4fea10dd7afcbb350e6b4fb8d16ad462ecc00f09 --- /dev/null +++ b/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 diff --git a/examples/SimplerEnv/eval_files/custom_argparse.py b/examples/SimplerEnv/eval_files/custom_argparse.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc7cf8b1ea336478322a89b0aba77f8a4c35286 --- /dev/null +++ b/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 diff --git a/examples/SimplerEnv/eval_files/model2simpler_interface.py b/examples/SimplerEnv/eval_files/model2simpler_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..4a9eff4c78dce1fc26d58b6f91712cc2e9c57df3 --- /dev/null +++ b/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) diff --git a/examples/SimplerEnv/eval_files/start_simpler_env.py b/examples/SimplerEnv/eval_files/start_simpler_env.py new file mode 100644 index 0000000000000000000000000000000000000000..150468c91c5eb5cb2c47da32c7f28ae59eedfabe --- /dev/null +++ b/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) diff --git a/examples/SimplerEnv/eval_files/test_your_simplerEnv.py b/examples/SimplerEnv/eval_files/test_your_simplerEnv.py new file mode 100644 index 0000000000000000000000000000000000000000..e57ec7628ac6c08e95f182b3bda572af8dbc6ccd --- /dev/null +++ b/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)) diff --git a/examples/SimplerEnv/train_files/data_registry/data_config.py b/examples/SimplerEnv/train_files/data_registry/data_config.py new file mode 100644 index 0000000000000000000000000000000000000000..44c5efa1aead92a95780095863468e180aef4801 --- /dev/null +++ b/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"), + ], +} diff --git a/examples/SimplerEnv/train_files/fractal_modality.json b/examples/SimplerEnv/train_files/fractal_modality.json new file mode 100644 index 0000000000000000000000000000000000000000..07d8b5f9b9aa29340a32cd28d20e452098ea4c6d --- /dev/null +++ b/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"} + } +} \ No newline at end of file diff --git a/examples/SimplerEnv/train_files/modality.json b/examples/SimplerEnv/train_files/modality.json new file mode 100644 index 0000000000000000000000000000000000000000..49324e2b47cdd9b31c703cf57796c819bff6898b --- /dev/null +++ b/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"} + } +} \ No newline at end of file diff --git a/experiments/importance_correction_ablation/train_residual_without_pq.py b/experiments/importance_correction_ablation/train_residual_without_pq.py new file mode 100644 index 0000000000000000000000000000000000000000..dce5079e607a06bdbbe42d4c1e74acefe2af8f7d --- /dev/null +++ b/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__") diff --git a/external_baselines/frameskip_adapted/README.md b/external_baselines/frameskip_adapted/README.md new file mode 100644 index 0000000000000000000000000000000000000000..38f6ad6912c260dd5a34c2b9e53cba2545f6f076 --- /dev/null +++ b/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. diff --git a/external_baselines/frameskip_adapted/build_cache_shard.py b/external_baselines/frameskip_adapted/build_cache_shard.py new file mode 100644 index 0000000000000000000000000000000000000000..0aeefd2e403c95a31ab2b0cb87d6733a633efbba --- /dev/null +++ b/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() diff --git a/external_baselines/frameskip_adapted/build_required_draws.py b/external_baselines/frameskip_adapted/build_required_draws.py new file mode 100644 index 0000000000000000000000000000000000000000..d95cf67c56718441640c1f00d8900081ae8ebcf3 --- /dev/null +++ b/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() diff --git a/external_baselines/frameskip_adapted/finalize_cache.py b/external_baselines/frameskip_adapted/finalize_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..aafb0587cfd9c2b6e66c3e8d92e3d579e234ba67 --- /dev/null +++ b/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() diff --git a/external_baselines/frameskip_adapted/train_frameskip_continuation.py b/external_baselines/frameskip_adapted/train_frameskip_continuation.py new file mode 100644 index 0000000000000000000000000000000000000000..d1293326724fd62d4b682e254d35247f3e160d89 --- /dev/null +++ b/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() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..589188a3c2fa05d2d34a5ba0cae434ae37b03cc0 --- /dev/null +++ b/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"] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..e98964bf59ffd329a253978e0b76cab7f8e6a8db --- /dev/null +++ b/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 \ No newline at end of file diff --git a/results/formal_384/aggregate_summary.json b/results/formal_384/aggregate_summary.json new file mode 100644 index 0000000000000000000000000000000000000000..10c8437d655b8b40e93a70dd0aec8d00acc1aebb --- /dev/null +++ b/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" +} diff --git a/results/formal_384/residual/policy_seed_516408/carrot/result.json b/results/formal_384/residual/policy_seed_516408/carrot/result.json new file mode 100644 index 0000000000000000000000000000000000000000..acedf527028ab08632d9416cd53d6c96d4161bac --- /dev/null +++ b/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 + ] +} diff --git a/results/formal_384/residual/policy_seed_516408/eggplant/result.json b/results/formal_384/residual/policy_seed_516408/eggplant/result.json new file mode 100644 index 0000000000000000000000000000000000000000..8e5e775ff6eb06f58e0e8022219cc78bd2a55249 --- /dev/null +++ b/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 + ] +} diff --git a/results/formal_384/residual/policy_seed_516408/four_task_summary.json b/results/formal_384/residual/policy_seed_516408/four_task_summary.json new file mode 100644 index 0000000000000000000000000000000000000000..5156afa7dada8a72632365769b63dde975ba1bfb --- /dev/null +++ b/results/formal_384/residual/policy_seed_516408/four_task_summary.json @@ -0,0 +1,34 @@ +{ + "arm": "residual", + "checkpoint_step": 20500, + "episodes_per_task": 24, + "macro_success_rate": 0.6145833333333334, + "micro_success_rate": 0.6145833333333334, + "policy_seed": 516408, + "schema_version": 1, + "status": "complete_no_video_four_task96", + "tasks": { + "carrot": { + "rollout_count": 24, + "strict_success_rate": 0.4583333333333333, + "success_count": 11 + }, + "eggplant": { + "rollout_count": 24, + "strict_success_rate": 0.8333333333333334, + "success_count": 20 + }, + "spoon": { + "rollout_count": 24, + "strict_success_rate": 0.9583333333333334, + "success_count": 23 + }, + "stack": { + "rollout_count": 24, + "strict_success_rate": 0.20833333333333334, + "success_count": 5 + } + }, + "total_rollouts": 96, + "total_success": 59 +} diff --git a/results/formal_384/residual/policy_seed_516408/spoon/result.json b/results/formal_384/residual/policy_seed_516408/spoon/result.json new file mode 100644 index 0000000000000000000000000000000000000000..abfb740150abbb0bd0adef0c8c88aea062d5481f --- /dev/null +++ b/results/formal_384/residual/policy_seed_516408/spoon/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "PutSpoonOnTableClothInScene-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, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + false, + true, + true, + true, + true, + true + ] +} diff --git a/results/formal_384/residual/policy_seed_516408/stack/result.json b/results/formal_384/residual/policy_seed_516408/stack/result.json new file mode 100644 index 0000000000000000000000000000000000000000..d9742b7bbd40194cc38d7b28e108e346b6eb18b1 --- /dev/null +++ b/results/formal_384/residual/policy_seed_516408/stack/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "StackGreenCubeOnYellowCubeBakedTexInScene-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": [ + false, + true, + false, + false, + true, + true, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true + ] +} diff --git a/results/formal_384/residual/policy_seed_518957/carrot/result.json b/results/formal_384/residual/policy_seed_518957/carrot/result.json new file mode 100644 index 0000000000000000000000000000000000000000..d4a25ffb0237058ff2e32d7f69918d2dc87f3268 --- /dev/null +++ b/results/formal_384/residual/policy_seed_518957/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": 518957 + }, + "successes": [ + true, + false, + true, + false, + false, + false, + true, + false, + true, + true, + false, + false, + true, + false, + true, + true, + true, + false, + false, + false, + true, + false, + false, + false + ] +} diff --git a/results/formal_384/residual/policy_seed_518957/eggplant/result.json b/results/formal_384/residual/policy_seed_518957/eggplant/result.json new file mode 100644 index 0000000000000000000000000000000000000000..73d883d7f4d8d2fc751f5b84343e1217feff28bd --- /dev/null +++ b/results/formal_384/residual/policy_seed_518957/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": 518957 + }, + "successes": [ + true, + false, + true, + false, + true, + true, + true, + false, + false, + true, + false, + true, + false, + false, + false, + true, + true, + true, + true, + true, + false, + true, + true, + true + ] +} diff --git a/results/formal_384/residual/policy_seed_518957/four_task_summary.json b/results/formal_384/residual/policy_seed_518957/four_task_summary.json new file mode 100644 index 0000000000000000000000000000000000000000..3fb5b54a6b0bccf5cce8f016c5fbe4f62d34d2e2 --- /dev/null +++ b/results/formal_384/residual/policy_seed_518957/four_task_summary.json @@ -0,0 +1,34 @@ +{ + "arm": "residual", + "checkpoint_step": 20500, + "episodes_per_task": 24, + "macro_success_rate": 0.5208333333333334, + "micro_success_rate": 0.5208333333333334, + "policy_seed": 518957, + "schema_version": 1, + "status": "complete_no_video_four_task96", + "tasks": { + "carrot": { + "rollout_count": 24, + "strict_success_rate": 0.4166666666666667, + "success_count": 10 + }, + "eggplant": { + "rollout_count": 24, + "strict_success_rate": 0.625, + "success_count": 15 + }, + "spoon": { + "rollout_count": 24, + "strict_success_rate": 0.7916666666666666, + "success_count": 19 + }, + "stack": { + "rollout_count": 24, + "strict_success_rate": 0.25, + "success_count": 6 + } + }, + "total_rollouts": 96, + "total_success": 50 +} diff --git a/results/formal_384/residual/policy_seed_518957/spoon/result.json b/results/formal_384/residual/policy_seed_518957/spoon/result.json new file mode 100644 index 0000000000000000000000000000000000000000..78eb2432862fb4fcaa3931153bb03e5c0b90e7c3 --- /dev/null +++ b/results/formal_384/residual/policy_seed_518957/spoon/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "PutSpoonOnTableClothInScene-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": 518957 + }, + "successes": [ + true, + true, + true, + true, + true, + true, + true, + false, + true, + false, + true, + true, + true, + true, + true, + true, + true, + false, + true, + true, + false, + true, + false, + true + ] +} diff --git a/results/formal_384/residual/policy_seed_518957/stack/result.json b/results/formal_384/residual/policy_seed_518957/stack/result.json new file mode 100644 index 0000000000000000000000000000000000000000..b46d3ae0fd407f236abe24fad83abb7b5df66f64 --- /dev/null +++ b/results/formal_384/residual/policy_seed_518957/stack/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "StackGreenCubeOnYellowCubeBakedTexInScene-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": 518957 + }, + "successes": [ + true, + true, + true, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + true + ] +} diff --git a/results/formal_384/residual/policy_seed_686765/carrot/result.json b/results/formal_384/residual/policy_seed_686765/carrot/result.json new file mode 100644 index 0000000000000000000000000000000000000000..dc56fe4fac8eaec7e35ec66569b7947555e0a464 --- /dev/null +++ b/results/formal_384/residual/policy_seed_686765/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": 686765 + }, + "successes": [ + false, + true, + true, + false, + true, + false, + false, + false, + true, + true, + true, + true, + true, + true, + true, + true, + false, + false, + false, + false, + true, + false, + false, + false + ] +} diff --git a/results/formal_384/residual/policy_seed_686765/eggplant/result.json b/results/formal_384/residual/policy_seed_686765/eggplant/result.json new file mode 100644 index 0000000000000000000000000000000000000000..994e73dfe99ba8a77781fd179ce2d3abd3262cb7 --- /dev/null +++ b/results/formal_384/residual/policy_seed_686765/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": 686765 + }, + "successes": [ + true, + true, + true, + false, + false, + true, + true, + false, + true, + true, + true, + true, + true, + true, + false, + false, + false, + true, + true, + true, + true, + true, + true, + true + ] +} diff --git a/results/formal_384/residual/policy_seed_686765/four_task_summary.json b/results/formal_384/residual/policy_seed_686765/four_task_summary.json new file mode 100644 index 0000000000000000000000000000000000000000..3cc3a6c4227c539a7db8c3b1b1b1377925574bfc --- /dev/null +++ b/results/formal_384/residual/policy_seed_686765/four_task_summary.json @@ -0,0 +1,34 @@ +{ + "arm": "residual", + "checkpoint_step": 20500, + "episodes_per_task": 24, + "macro_success_rate": 0.6145833333333334, + "micro_success_rate": 0.6145833333333334, + "policy_seed": 686765, + "schema_version": 1, + "status": "complete_no_video_four_task96", + "tasks": { + "carrot": { + "rollout_count": 24, + "strict_success_rate": 0.5, + "success_count": 12 + }, + "eggplant": { + "rollout_count": 24, + "strict_success_rate": 0.75, + "success_count": 18 + }, + "spoon": { + "rollout_count": 24, + "strict_success_rate": 0.9166666666666666, + "success_count": 22 + }, + "stack": { + "rollout_count": 24, + "strict_success_rate": 0.2916666666666667, + "success_count": 7 + } + }, + "total_rollouts": 96, + "total_success": 59 +} diff --git a/results/formal_384/residual/policy_seed_686765/spoon/result.json b/results/formal_384/residual/policy_seed_686765/spoon/result.json new file mode 100644 index 0000000000000000000000000000000000000000..f3d9b3473ba66c22095220c1ab0d01bb2a88617b --- /dev/null +++ b/results/formal_384/residual/policy_seed_686765/spoon/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "PutSpoonOnTableClothInScene-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": 686765 + }, + "successes": [ + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + false, + true, + false, + true, + true, + true, + true, + true, + true, + true, + true + ] +} diff --git a/results/formal_384/residual/policy_seed_686765/stack/result.json b/results/formal_384/residual/policy_seed_686765/stack/result.json new file mode 100644 index 0000000000000000000000000000000000000000..9508d420301343f3caad107a913a6a5d7c5bf993 --- /dev/null +++ b/results/formal_384/residual/policy_seed_686765/stack/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "StackGreenCubeOnYellowCubeBakedTexInScene-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": 686765 + }, + "successes": [ + true, + true, + false, + false, + true, + true, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true, + false, + false, + false, + false, + false, + true + ] +} diff --git a/results/formal_384/residual/policy_seed_765353/carrot/result.json b/results/formal_384/residual/policy_seed_765353/carrot/result.json new file mode 100644 index 0000000000000000000000000000000000000000..d6adcffa920621fc629b8ae46fb178b231cd4ff0 --- /dev/null +++ b/results/formal_384/residual/policy_seed_765353/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": 765353 + }, + "successes": [ + true, + false, + true, + false, + false, + true, + false, + true, + true, + true, + true, + false, + true, + false, + true, + true, + true, + false, + false, + false, + true, + false, + false, + false + ] +} diff --git a/results/formal_384/residual/policy_seed_765353/eggplant/result.json b/results/formal_384/residual/policy_seed_765353/eggplant/result.json new file mode 100644 index 0000000000000000000000000000000000000000..2ed2af4b0cb3053bd66e470bec738d99758e2cfa --- /dev/null +++ b/results/formal_384/residual/policy_seed_765353/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": 765353 + }, + "successes": [ + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + false, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true + ] +} diff --git a/results/formal_384/residual/policy_seed_765353/four_task_summary.json b/results/formal_384/residual/policy_seed_765353/four_task_summary.json new file mode 100644 index 0000000000000000000000000000000000000000..31e82ae715b79512fb187a661f2e9565a4ea0078 --- /dev/null +++ b/results/formal_384/residual/policy_seed_765353/four_task_summary.json @@ -0,0 +1,34 @@ +{ + "arm": "residual", + "checkpoint_step": 20500, + "episodes_per_task": 24, + "macro_success_rate": 0.625, + "micro_success_rate": 0.625, + "policy_seed": 765353, + "schema_version": 1, + "status": "complete_no_video_four_task96", + "tasks": { + "carrot": { + "rollout_count": 24, + "strict_success_rate": 0.5, + "success_count": 12 + }, + "eggplant": { + "rollout_count": 24, + "strict_success_rate": 0.9583333333333334, + "success_count": 23 + }, + "spoon": { + "rollout_count": 24, + "strict_success_rate": 0.7916666666666666, + "success_count": 19 + }, + "stack": { + "rollout_count": 24, + "strict_success_rate": 0.25, + "success_count": 6 + } + }, + "total_rollouts": 96, + "total_success": 60 +} diff --git a/results/formal_384/residual/policy_seed_765353/spoon/result.json b/results/formal_384/residual/policy_seed_765353/spoon/result.json new file mode 100644 index 0000000000000000000000000000000000000000..dec1f1cc9b38c49ba82cb12ac8214b7e24ea1371 --- /dev/null +++ b/results/formal_384/residual/policy_seed_765353/spoon/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "PutSpoonOnTableClothInScene-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": 765353 + }, + "successes": [ + true, + true, + true, + true, + true, + true, + true, + false, + true, + false, + true, + true, + true, + false, + true, + true, + true, + true, + true, + false, + false, + true, + true, + true + ] +} diff --git a/results/formal_384/residual/policy_seed_765353/stack/result.json b/results/formal_384/residual/policy_seed_765353/stack/result.json new file mode 100644 index 0000000000000000000000000000000000000000..19fafa7f73c32b68ff0b6ecb2879d2025127ff9c --- /dev/null +++ b/results/formal_384/residual/policy_seed_765353/stack/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "StackGreenCubeOnYellowCubeBakedTexInScene-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": 765353 + }, + "successes": [ + true, + true, + true, + false, + false, + true, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_516408/carrot/result.json b/results/formal_384/uniform/policy_seed_516408/carrot/result.json new file mode 100644 index 0000000000000000000000000000000000000000..f4692130160979bc8749a1a38d6249cc2617a08f --- /dev/null +++ b/results/formal_384/uniform/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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 516408 + }, + "successes": [ + true, + false, + true, + false, + true, + false, + false, + false, + true, + false, + true, + false, + false, + false, + true, + false, + true, + false, + false, + false, + true, + false, + false, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_516408/eggplant/result.json b/results/formal_384/uniform/policy_seed_516408/eggplant/result.json new file mode 100644 index 0000000000000000000000000000000000000000..97e5915acaed7c1f152fa2832d512ee893b8bd3e --- /dev/null +++ b/results/formal_384/uniform/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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 516408 + }, + "successes": [ + true, + true, + false, + true, + true, + true, + true, + false, + true, + true, + true, + true, + false, + true, + true, + true, + false, + true, + true, + false, + false, + true, + false, + true + ] +} diff --git a/results/formal_384/uniform/policy_seed_516408/four_task_summary.json b/results/formal_384/uniform/policy_seed_516408/four_task_summary.json new file mode 100644 index 0000000000000000000000000000000000000000..b117c20be87f8fadedf853b5fb11d67d37b72e8f --- /dev/null +++ b/results/formal_384/uniform/policy_seed_516408/four_task_summary.json @@ -0,0 +1,34 @@ +{ + "arm": "uniform", + "checkpoint_step": 20500, + "episodes_per_task": 24, + "macro_success_rate": 0.40625, + "micro_success_rate": 0.40625, + "policy_seed": 516408, + "schema_version": 1, + "status": "complete_no_video_four_task96", + "tasks": { + "carrot": { + "rollout_count": 24, + "strict_success_rate": 0.3333333333333333, + "success_count": 8 + }, + "eggplant": { + "rollout_count": 24, + "strict_success_rate": 0.7083333333333334, + "success_count": 17 + }, + "spoon": { + "rollout_count": 24, + "strict_success_rate": 0.4583333333333333, + "success_count": 11 + }, + "stack": { + "rollout_count": 24, + "strict_success_rate": 0.125, + "success_count": 3 + } + }, + "total_rollouts": 96, + "total_success": 39 +} diff --git a/results/formal_384/uniform/policy_seed_516408/spoon/result.json b/results/formal_384/uniform/policy_seed_516408/spoon/result.json new file mode 100644 index 0000000000000000000000000000000000000000..bbeede84bbcac6a2470bcb3e5bdd9cd7a5fcc3fd --- /dev/null +++ b/results/formal_384/uniform/policy_seed_516408/spoon/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "PutSpoonOnTableClothInScene-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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 516408 + }, + "successes": [ + false, + false, + true, + false, + true, + false, + true, + false, + true, + true, + true, + false, + true, + false, + true, + false, + true, + false, + true, + false, + false, + false, + true, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_516408/stack/result.json b/results/formal_384/uniform/policy_seed_516408/stack/result.json new file mode 100644 index 0000000000000000000000000000000000000000..e7ce607d99d328e472e38005b54765f1d41c00cb --- /dev/null +++ b/results/formal_384/uniform/policy_seed_516408/stack/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "StackGreenCubeOnYellowCubeBakedTexInScene-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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 516408 + }, + "successes": [ + false, + false, + false, + false, + true, + false, + true, + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_518957/carrot/result.json b/results/formal_384/uniform/policy_seed_518957/carrot/result.json new file mode 100644 index 0000000000000000000000000000000000000000..a0f24e8cc3f9a8931bda6898b8aa637b8857a536 --- /dev/null +++ b/results/formal_384/uniform/policy_seed_518957/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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 518957 + }, + "successes": [ + false, + false, + true, + false, + true, + true, + false, + false, + true, + false, + false, + false, + true, + false, + true, + false, + true, + false, + false, + false, + true, + false, + false, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_518957/eggplant/result.json b/results/formal_384/uniform/policy_seed_518957/eggplant/result.json new file mode 100644 index 0000000000000000000000000000000000000000..ffa60db36e9bf84e07d5fd7cdbb28bdb622bc3dc --- /dev/null +++ b/results/formal_384/uniform/policy_seed_518957/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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 518957 + }, + "successes": [ + true, + true, + false, + true, + true, + true, + true, + false, + true, + true, + true, + true, + false, + true, + false, + false, + true, + true, + true, + true, + false, + true, + true, + true + ] +} diff --git a/results/formal_384/uniform/policy_seed_518957/four_task_summary.json b/results/formal_384/uniform/policy_seed_518957/four_task_summary.json new file mode 100644 index 0000000000000000000000000000000000000000..214efa94a3cea109fc6bb28938028e8c93f74df5 --- /dev/null +++ b/results/formal_384/uniform/policy_seed_518957/four_task_summary.json @@ -0,0 +1,34 @@ +{ + "arm": "uniform", + "checkpoint_step": 20500, + "episodes_per_task": 24, + "macro_success_rate": 0.4270833333333333, + "micro_success_rate": 0.4270833333333333, + "policy_seed": 518957, + "schema_version": 1, + "status": "complete_no_video_four_task96", + "tasks": { + "carrot": { + "rollout_count": 24, + "strict_success_rate": 0.3333333333333333, + "success_count": 8 + }, + "eggplant": { + "rollout_count": 24, + "strict_success_rate": 0.75, + "success_count": 18 + }, + "spoon": { + "rollout_count": 24, + "strict_success_rate": 0.5416666666666666, + "success_count": 13 + }, + "stack": { + "rollout_count": 24, + "strict_success_rate": 0.08333333333333333, + "success_count": 2 + } + }, + "total_rollouts": 96, + "total_success": 41 +} diff --git a/results/formal_384/uniform/policy_seed_518957/spoon/result.json b/results/formal_384/uniform/policy_seed_518957/spoon/result.json new file mode 100644 index 0000000000000000000000000000000000000000..d23861e03f78dd5fb2d7ad72d60b068dba8bc4d2 --- /dev/null +++ b/results/formal_384/uniform/policy_seed_518957/spoon/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "PutSpoonOnTableClothInScene-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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 518957 + }, + "successes": [ + true, + false, + true, + false, + true, + false, + true, + false, + true, + false, + true, + false, + true, + true, + true, + false, + true, + false, + true, + false, + true, + false, + true, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_518957/stack/result.json b/results/formal_384/uniform/policy_seed_518957/stack/result.json new file mode 100644 index 0000000000000000000000000000000000000000..726bd499428c6859470cade7ed49b91bc0c5b573 --- /dev/null +++ b/results/formal_384/uniform/policy_seed_518957/stack/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "StackGreenCubeOnYellowCubeBakedTexInScene-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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 518957 + }, + "successes": [ + false, + false, + false, + false, + false, + true, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_686765/carrot/result.json b/results/formal_384/uniform/policy_seed_686765/carrot/result.json new file mode 100644 index 0000000000000000000000000000000000000000..7fbd9884b06e732ab649d5cd6d9633d76a393970 --- /dev/null +++ b/results/formal_384/uniform/policy_seed_686765/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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 686765 + }, + "successes": [ + false, + false, + true, + false, + false, + false, + false, + false, + true, + true, + false, + false, + true, + false, + true, + false, + false, + false, + false, + false, + true, + false, + false, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_686765/eggplant/result.json b/results/formal_384/uniform/policy_seed_686765/eggplant/result.json new file mode 100644 index 0000000000000000000000000000000000000000..385a5b070665c4e6a87f2c07eb5ae68855139bdb --- /dev/null +++ b/results/formal_384/uniform/policy_seed_686765/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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 686765 + }, + "successes": [ + true, + false, + true, + true, + false, + true, + true, + true, + true, + true, + false, + true, + false, + true, + true, + true, + false, + true, + true, + false, + true, + true, + true, + true + ] +} diff --git a/results/formal_384/uniform/policy_seed_686765/four_task_summary.json b/results/formal_384/uniform/policy_seed_686765/four_task_summary.json new file mode 100644 index 0000000000000000000000000000000000000000..092d0447e2a881cbb420e3e7b72f26120eac2545 --- /dev/null +++ b/results/formal_384/uniform/policy_seed_686765/four_task_summary.json @@ -0,0 +1,34 @@ +{ + "arm": "uniform", + "checkpoint_step": 20500, + "episodes_per_task": 24, + "macro_success_rate": 0.4479166666666667, + "micro_success_rate": 0.4479166666666667, + "policy_seed": 686765, + "schema_version": 1, + "status": "complete_no_video_four_task96", + "tasks": { + "carrot": { + "rollout_count": 24, + "strict_success_rate": 0.25, + "success_count": 6 + }, + "eggplant": { + "rollout_count": 24, + "strict_success_rate": 0.75, + "success_count": 18 + }, + "spoon": { + "rollout_count": 24, + "strict_success_rate": 0.5833333333333334, + "success_count": 14 + }, + "stack": { + "rollout_count": 24, + "strict_success_rate": 0.20833333333333334, + "success_count": 5 + } + }, + "total_rollouts": 96, + "total_success": 43 +} diff --git a/results/formal_384/uniform/policy_seed_686765/spoon/result.json b/results/formal_384/uniform/policy_seed_686765/spoon/result.json new file mode 100644 index 0000000000000000000000000000000000000000..3002980ad01901999dc32ef8af7aacb2551fd02f --- /dev/null +++ b/results/formal_384/uniform/policy_seed_686765/spoon/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "PutSpoonOnTableClothInScene-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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 686765 + }, + "successes": [ + true, + false, + true, + false, + true, + true, + true, + false, + true, + false, + true, + false, + true, + false, + true, + true, + true, + false, + true, + false, + true, + false, + true, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_686765/stack/result.json b/results/formal_384/uniform/policy_seed_686765/stack/result.json new file mode 100644 index 0000000000000000000000000000000000000000..f7920924de25f30e7ce152c02071f70bc25a9d7b --- /dev/null +++ b/results/formal_384/uniform/policy_seed_686765/stack/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "StackGreenCubeOnYellowCubeBakedTexInScene-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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 686765 + }, + "successes": [ + true, + false, + false, + false, + false, + false, + true, + false, + true, + false, + false, + true, + false, + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_765353/carrot/result.json b/results/formal_384/uniform/policy_seed_765353/carrot/result.json new file mode 100644 index 0000000000000000000000000000000000000000..877a47dbac004f218be8230a36d01e30393824dc --- /dev/null +++ b/results/formal_384/uniform/policy_seed_765353/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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 765353 + }, + "successes": [ + false, + false, + true, + false, + true, + false, + false, + false, + true, + false, + true, + false, + true, + false, + true, + false, + true, + false, + false, + false, + true, + true, + false, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_765353/eggplant/result.json b/results/formal_384/uniform/policy_seed_765353/eggplant/result.json new file mode 100644 index 0000000000000000000000000000000000000000..98d2b56bd40e21e81f47cb0939264d482c3e9435 --- /dev/null +++ b/results/formal_384/uniform/policy_seed_765353/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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 765353 + }, + "successes": [ + true, + false, + true, + true, + true, + true, + true, + true, + true, + true, + false, + true, + false, + false, + true, + false, + false, + true, + true, + true, + false, + false, + false, + true + ] +} diff --git a/results/formal_384/uniform/policy_seed_765353/four_task_summary.json b/results/formal_384/uniform/policy_seed_765353/four_task_summary.json new file mode 100644 index 0000000000000000000000000000000000000000..4aee826ef3827e1b89f7929c26a5fea19168c0ff --- /dev/null +++ b/results/formal_384/uniform/policy_seed_765353/four_task_summary.json @@ -0,0 +1,34 @@ +{ + "arm": "uniform", + "checkpoint_step": 20500, + "episodes_per_task": 24, + "macro_success_rate": 0.40625, + "micro_success_rate": 0.40625, + "policy_seed": 765353, + "schema_version": 1, + "status": "complete_no_video_four_task96", + "tasks": { + "carrot": { + "rollout_count": 24, + "strict_success_rate": 0.375, + "success_count": 9 + }, + "eggplant": { + "rollout_count": 24, + "strict_success_rate": 0.625, + "success_count": 15 + }, + "spoon": { + "rollout_count": 24, + "strict_success_rate": 0.5, + "success_count": 12 + }, + "stack": { + "rollout_count": 24, + "strict_success_rate": 0.125, + "success_count": 3 + } + }, + "total_rollouts": 96, + "total_success": 39 +} diff --git a/results/formal_384/uniform/policy_seed_765353/spoon/result.json b/results/formal_384/uniform/policy_seed_765353/spoon/result.json new file mode 100644 index 0000000000000000000000000000000000000000..d90f5006c8b606aa202be15294ee75dfab40276c --- /dev/null +++ b/results/formal_384/uniform/policy_seed_765353/spoon/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "PutSpoonOnTableClothInScene-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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 765353 + }, + "successes": [ + true, + false, + true, + true, + true, + false, + true, + false, + true, + true, + true, + false, + false, + false, + false, + false, + true, + false, + true, + false, + false, + true, + true, + false + ] +} diff --git a/results/formal_384/uniform/policy_seed_765353/stack/result.json b/results/formal_384/uniform/policy_seed_765353/stack/result.json new file mode 100644 index 0000000000000000000000000000000000000000..c5672bd35d55f855d1cab3a5e86e57f03624cd62 --- /dev/null +++ b/results/formal_384/uniform/policy_seed_765353/stack/result.json @@ -0,0 +1,63 @@ +{ + "schema_version": 2, + "env_name": "StackGreenCubeOnYellowCubeBakedTexInScene-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-uniform-seed-117465/checkpoints/steps_20500_pytorch_model.pt", + "server_metadata": { + "env": "simpler_env", + "policy_seed": 765353 + }, + "successes": [ + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false, + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false + ] +} diff --git a/scripts/build_matched_random_temporal_proposal.py b/scripts/build_matched_random_temporal_proposal.py new file mode 100644 index 0000000000000000000000000000000000000000..e244779015ee86976425a2cece175f4b1672b0d6 --- /dev/null +++ b/scripts/build_matched_random_temporal_proposal.py @@ -0,0 +1,47 @@ +"""Build the matched-random control by permuting residual masses per trajectory.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest +from starVLA.dataloader.temporal_proposal import TemporalProposal +from starVLA.scoring.matched_random import build_matched_random_proposal + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--residual-proposal", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--permutation-seed", type=int, required=True) + args = parser.parse_args() + + manifest = ValidAnchorManifest.load(args.manifest) + residual = TemporalProposal.load(args.residual_proposal, manifest) + matched = build_matched_random_proposal( + manifest=manifest, + residual_proposal=residual, + permutation_seed=args.permutation_seed, + ) + matched.save(args.output) + print( + json.dumps( + { + "status": "complete", + "manifest_sha256": manifest.sha256, + "source_proposal_sha256": residual.sha256, + "matched_random_proposal_sha256": matched.sha256, + "permutation_seed": args.permutation_seed, + "output": str(args.output.resolve()), + }, + indent=2, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/build_residual_temporal_proposal.py b/scripts/build_residual_temporal_proposal.py new file mode 100644 index 0000000000000000000000000000000000000000..10328b0d16b8817d62af71b921c82b58b1301621 --- /dev/null +++ b/scripts/build_residual_temporal_proposal.py @@ -0,0 +1,776 @@ +"""Build a resumable, grid-scored residual temporal proposal. + +The expensive frozen-policy residual is evaluated at a fixed quantile grid +inside every episode. Scores for the remaining anchors are filled by +deterministic linear interpolation, ranked within that episode, and converted +to epsilon-floored masses. ``TemporalProposal`` then applies the safe mixture +and exact ``p/q`` correction used by training. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import random +import time +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from omegaconf import OmegaConf + +from starVLA.dataloader.lerobot_datasets import get_vla_dataset +from starVLA.dataloader.temporal_proposal import TemporalProposal +from starVLA.model.framework.base_framework import build_framework +from starVLA.model.framework.share_tools import apply_config_compat +from starVLA.scoring.residual_cache import ( + interpolate_grid_scores, + percentile_intensity_masses, + quantile_anchor_positions, +) +from starVLA.scoring.starvla_likelihood import ( + score_starvla_flow_matching_residual, +) + + +SCHEMA_VERSION = 1 + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _canonical_sha256(payload: dict[str, Any]) -> str: + encoded = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _load_gate_certificate( + path: Path, + *, + checkpoint_sha256: str, + manifest_sha256: str, + alpha: float, + epsilon: float, + allow_failed_gate: bool, +) -> dict[str, Any]: + payload = json.loads(path.read_text(encoding="utf-8")) + gate = payload.get("gate") + proposal = payload.get("hierarchical_audit", {}).get("proposal", {}) + if (not isinstance(gate, dict) or gate.get("pass") is not True) and not allow_failed_gate: + raise ValueError("Proposal eligibility gate did not pass") + if gate.get("eligible_for_500_step_ablation") is not True and not allow_failed_gate: + raise ValueError("Certificate does not authorize a prospective ablation") + + status = payload.get("status") + if status == "formal_within_trajectory_utility_gate": + pass + elif status == "formal_independent_target_panel_second_moment_gate": + panel = payload.get("panel_audit", {}) + if ( + int(panel.get("panel_count", 0)) < 8 + or int(panel.get("draws_per_panel", 0)) < 4 + or panel.get("panels_are_disjoint") is not True + or float(panel.get("maximum_individual_ratio", 1.0)) > 0.98 + or float(panel.get("minimum_individual_improved_fraction", 0.0)) + < 0.75 + or int(panel.get("complementary_split_count", 0)) < 70 + or float( + panel.get("maximum_complementary_split_ratio", 1.0) + ) + > 0.98 + or float(panel.get("minimum_split_spearman", 0.0)) < 0.50 + ): + raise ValueError( + "Independent-panel certificate misses its locked thresholds" + ) + else: + raise ValueError( + "Proposal eligibility certificate has the wrong formal status" + ) + + if payload.get("checkpoint_sha256") != checkpoint_sha256: + raise ValueError("Certificate belongs to a different checkpoint") + if payload.get("manifest_sha256") != manifest_sha256: + raise ValueError("Certificate belongs to a different valid-anchor manifest") + if not np.isclose(float(proposal.get("alpha", -1.0)), alpha, atol=0.0): + raise ValueError("Certificate alpha differs from the proposal alpha") + if not np.isclose( + float(proposal.get("percentile_epsilon", -1.0)), + epsilon, + atol=0.0, + ): + raise ValueError("Certificate epsilon differs from the proposal epsilon") + + selection = payload.get("selection_protocol", {}) + episode_count = int(selection.get("episode_count", -1)) + anchors_per_episode = int(selection.get("anchors_per_episode", -1)) + anchor_count = int(selection.get("anchor_count", -1)) + if ( + episode_count < 32 + or anchors_per_episode != 8 + or anchor_count != episode_count * anchors_per_episode + ): + raise ValueError( + "Certificate is not a valid independent-episode/8-anchor audit" + ) + return { + "override_gate_failure": bool(allow_failed_gate and gate.get("pass") is not True), + "path": str(path.resolve()), + "sha256": _sha256_file(path), + "protocol": status, + "status": gate.get("status"), + "checks": gate.get("checks"), + } + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + try: + with temporary.open("w", encoding="utf-8") as output: + json.dump(payload, output, indent=2, sort_keys=True) + output.flush() + os.fsync(output.fileno()) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def _set_random_seed(seed: int) -> None: + bounded_seed = int(seed) % (2**32) + random.seed(bounded_seed) + np.random.seed(bounded_seed) + torch.manual_seed(bounded_seed) + torch.cuda.manual_seed_all(bounded_seed) + + +def _json_scalar(value: Any) -> Any: + return value.item() if isinstance(value, np.generic) else value + + +def _shard_path(cache_dir: Path, dataset_index: int, episode_index: int) -> Path: + return cache_dir / "shards" / ( + f"dataset_{dataset_index:02d}_episode_{episode_index:06d}.json" + ) + + +def _planned_grid_count(manifest, grid_size: int) -> int: + return sum( + min(len(episode.anchors), grid_size) + for dataset in manifest.datasets + for episode in dataset.episodes + ) + + +def _episode_belongs_to_shard( + global_episode_index: int, + *, + shard_index: int, + num_shards: int, +) -> bool: + if num_shards <= 0: + raise ValueError("num_shards must be positive") + if not 0 <= shard_index < num_shards: + raise ValueError("shard_index must satisfy 0 <= shard_index < num_shards") + return int(global_episode_index) % int(num_shards) == int(shard_index) + + +def _worker_progress_path(cache_dir: Path, shard_index: int, num_shards: int) -> Path: + return cache_dir / "workers" / f"shard_{shard_index:02d}_of_{num_shards:02d}.json" + + +def _load_valid_shard( + path: Path, + *, + identity_sha256: str, + dataset_index: int, + episode_index: int, + trajectory_id: Any, + positions: np.ndarray, + base_indices: list[int], + estimator_count: int, +) -> dict[str, Any]: + payload = json.loads(path.read_text(encoding="utf-8")) + expected = { + "identity_sha256": identity_sha256, + "dataset_index": dataset_index, + "episode_index": episode_index, + "trajectory_id": _json_scalar(trajectory_id), + "anchor_positions": positions.tolist(), + "base_indices": base_indices, + } + for field, value in expected.items(): + if payload.get(field) != value: + raise ValueError(f"Residual cache shard mismatch for {field}: {path}") + estimators = np.asarray(payload.get("estimator_residuals"), dtype=np.float64) + if estimators.shape != (len(positions), estimator_count): + raise ValueError(f"Residual cache shard has wrong estimator shape: {path}") + if not np.isfinite(estimators).all(): + raise ValueError(f"Residual cache shard contains non-finite values: {path}") + scores = np.asarray(payload.get("grid_scores"), dtype=np.float64) + if scores.shape != (len(positions),) or not np.isfinite(scores).all(): + raise ValueError(f"Residual cache shard contains invalid grid scores: {path}") + return payload + + +def _load_examples( + source_dataset, + episode, + positions: np.ndarray, + *, + dataset_index: int, + episode_index: int, + transform_seed_base: int, +) -> tuple[list[dict], list[int], list[int]]: + examples = [] + base_indices = [] + transform_seeds = [] + for anchor_position in positions: + base_index = int(episode.anchors[int(anchor_position)]) + transform_seed = ( + transform_seed_base + + dataset_index * 10_000_019 + + episode_index * 10_007 + + int(anchor_position) + ) + _set_random_seed(transform_seed) + raw_data = source_dataset.get_step_data( + episode.trajectory_id, + base_index, + ) + transformed = source_dataset.transforms(raw_data) + examples.append(source_dataset._pack_sample(transformed)) + base_indices.append(base_index) + transform_seeds.append(transform_seed) + return examples, base_indices, transform_seeds + + +def _score_examples( + model, + examples: list[dict], + *, + estimator_seeds: list[int], + samples_per_estimator: int, + batch_size: int, +) -> np.ndarray: + all_estimators = np.empty( + (len(examples), len(estimator_seeds)), + dtype=np.float64, + ) + for start in range(0, len(examples), batch_size): + batch = examples[start : start + batch_size] + with torch.no_grad(): + result = score_starvla_flow_matching_residual( + model, + batch, + estimator_seeds=estimator_seeds, + samples_per_estimator=samples_per_estimator, + ) + batch_estimators = ( + result.estimator_residuals.transpose(0, 1).float().cpu().numpy() + ) + all_estimators[start : start + len(batch)] = batch_estimators + return all_estimators + + +def _build_identity(args, manifest, gate_certificate) -> dict[str, Any]: + checkpoint = args.checkpoint.resolve() + identity = { + "schema_version": SCHEMA_VERSION, + # Runtime locations are intentionally excluded from the scientific + # identity; the checkpoint content and explicit scoring parameters bind it. + "checkpoint_sha256": _sha256_file(checkpoint), + "manifest_sha256": manifest.sha256, + "grid_size": int(args.grid_size), + "interpolation": "linear_within_episode", + "percentile_scope": "within_episode_after_interpolation", + "epsilon": float(args.epsilon), + "alpha": float(args.alpha), + "estimator_seeds": [int(seed) for seed in args.estimator_seeds], + "samples_per_estimator": int(args.samples_per_estimator), + "transform_seed_base": int(args.transform_seed_base), + "model_dtype": args.model_dtype, + "gate42f_certificate": gate_certificate, + } + identity["sha256"] = _canonical_sha256(identity) + return identity + + +def _initialize_metadata( + path: Path, + *, + identity: dict[str, Any], + episode_count: int, + grid_anchor_count: int, +) -> dict[str, Any]: + if path.exists(): + metadata = json.loads(path.read_text(encoding="utf-8")) + if metadata.get("identity") != identity: + raise ValueError( + "Residual cache metadata belongs to a different scorer protocol" + ) + return metadata + metadata = { + "schema_version": SCHEMA_VERSION, + "status": "in_progress", + "identity": identity, + "planned_episode_count": int(episode_count), + "planned_grid_anchor_count": int(grid_anchor_count), + "completed_episode_count": 0, + "completed_grid_anchor_count": 0, + "scoring_seconds": 0.0, + "proposal": None, + } + _atomic_json(path, metadata) + return metadata + + +def _refresh_progress(metadata: dict[str, Any], cache_dir: Path, manifest) -> None: + completed_episodes = 0 + completed_anchors = 0 + for dataset_index, dataset in enumerate(manifest.datasets): + for episode_index, _ in enumerate(dataset.episodes): + path = _shard_path(cache_dir, dataset_index, episode_index) + if not path.exists(): + continue + payload = json.loads(path.read_text(encoding="utf-8")) + completed_episodes += 1 + completed_anchors += len(payload.get("anchor_positions", [])) + metadata["completed_episode_count"] = completed_episodes + metadata["completed_grid_anchor_count"] = completed_anchors + + +def _finalize_proposal( + *, + manifest, + cache_dir: Path, + identity: dict[str, Any], + output: Path, + alpha: float, + epsilon: float, + allow_failed_gate: bool, + estimator_count: int, + gate_certificate: dict[str, Any], +) -> TemporalProposal: + masses = [] + for dataset_index, dataset in enumerate(manifest.datasets): + dataset_masses = [] + for episode_index, episode in enumerate(dataset.episodes): + positions = quantile_anchor_positions( + len(episode.anchors), + int(identity["grid_size"]), + ) + base_indices = [int(episode.anchors[int(index)]) for index in positions] + shard = _load_valid_shard( + _shard_path(cache_dir, dataset_index, episode_index), + identity_sha256=identity["sha256"], + dataset_index=dataset_index, + episode_index=episode_index, + trajectory_id=episode.trajectory_id, + positions=positions, + base_indices=base_indices, + estimator_count=estimator_count, + ) + full_scores = interpolate_grid_scores( + len(episode.anchors), + positions, + np.asarray(shard["grid_scores"], dtype=np.float64), + ) + dataset_masses.append( + percentile_intensity_masses(full_scores, epsilon=epsilon) + ) + masses.append(dataset_masses) + + proposal = TemporalProposal.from_masses( + manifest=manifest, + alpha=alpha, + masses=masses, + score_kind="residual", + scorer_sha256=identity["sha256"], + validation_certificate={ + "protocol": "prospective-grid-residual-v1", + "identity_sha256": identity["sha256"], + "grid_size": identity["grid_size"], + "interpolation": identity["interpolation"], + "percentile_scope": identity["percentile_scope"], + "epsilon": epsilon, + "gate42f_certificate": gate_certificate, + }, + ) + proposal.save(output) + return proposal + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--gate-result", type=Path, required=True) + parser.add_argument("--allow-failed-gate", action="store_true") + parser.add_argument("--cache-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--grid-size", type=int, default=8) + parser.add_argument("--alpha", type=float, default=0.2) + parser.add_argument("--epsilon", type=float, default=0.1) + parser.add_argument( + "--estimator-seeds", + type=int, + nargs="+", + default=[16200, 16201, 16202, 16203], + ) + parser.add_argument("--samples-per-estimator", type=int, default=8) + parser.add_argument("--transform-seed-base", type=int, default=16500) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument( + "--model-dtype", + choices=("float32", "bfloat16"), + default="float32", + ) + parser.add_argument("--num-shards", type=int, default=1) + parser.add_argument("--shard-index", type=int, default=0) + parser.add_argument( + "--max-new-episodes", + type=int, + default=0, + help="Stop cleanly after this many newly scored episodes; 0 means all.", + ) + parser.add_argument("--plan-only", action="store_true") + args = parser.parse_args() + + if args.grid_size <= 0 or args.batch_size <= 0: + raise ValueError("grid-size and batch-size must be positive") + if args.num_shards <= 0: + raise ValueError("num-shards must be positive") + if not 0 <= args.shard_index < args.num_shards: + raise ValueError("shard-index must satisfy 0 <= shard-index < num-shards") + if not 0.0 <= args.alpha < 1.0: + raise ValueError("alpha must satisfy 0 <= alpha < 1") + if args.epsilon <= 0.0: + raise ValueError("epsilon must be positive") + if not args.estimator_seeds: + raise ValueError("at least one estimator seed is required") + if len(set(args.estimator_seeds)) != len(args.estimator_seeds): + raise ValueError("estimator seeds must be unique") + if args.samples_per_estimator <= 0: + raise ValueError("samples-per-estimator must be positive") + + cfg = apply_config_compat(OmegaConf.load(args.config)) + cfg.output_dir = str(args.cache_dir / "runtime") + cfg.datasets.vla_data.temporal_proposal_path = None + dataset = get_vla_dataset( + data_cfg=cfg.datasets.vla_data, + balance_dataset_weights=cfg.datasets.vla_data.get( + "balance_dataset_weights", + False, + ), + balance_trajectory_weights=cfg.datasets.vla_data.get( + "balance_trajectory_weights", + False, + ), + seed=int(cfg.get("seed", 42)), + ) + manifest = dataset.valid_anchor_manifest + if manifest is None: + raise RuntimeError("Residual proposal requires traceable_sampling") + + checkpoint_sha256 = _sha256_file(args.checkpoint.resolve()) + gate_certificate = _load_gate_certificate( + args.gate_result.resolve(), + checkpoint_sha256=checkpoint_sha256, + manifest_sha256=manifest.sha256, + alpha=float(args.alpha), + epsilon=float(args.epsilon), + allow_failed_gate=args.allow_failed_gate, + ) + identity = _build_identity(args, manifest, gate_certificate) + episode_count = sum(len(item.episodes) for item in manifest.datasets) + grid_anchor_count = _planned_grid_count(manifest, args.grid_size) + metadata_path = args.cache_dir / "metadata.json" + metadata = _initialize_metadata( + metadata_path, + identity=identity, + episode_count=episode_count, + grid_anchor_count=grid_anchor_count, + ) + _refresh_progress(metadata, args.cache_dir, manifest) + _atomic_json(metadata_path, metadata) + print( + json.dumps( + { + "identity_sha256": identity["sha256"], + "episodes": episode_count, + "grid_anchors": grid_anchor_count, + "completed_episodes": metadata["completed_episode_count"], + "completed_grid_anchors": metadata["completed_grid_anchor_count"], + "estimated_hours_at_0.312_seconds_per_anchor": ( + grid_anchor_count * 0.312 / 3600.0 + ), + }, + indent=2, + ), + flush=True, + ) + if args.plan_only: + return + + parallel_worker = args.num_shards > 1 + assigned_episode_count = 0 + assigned_grid_anchor_count = 0 + global_episode_index = 0 + for dataset_support in manifest.datasets: + for episode in dataset_support.episodes: + if _episode_belongs_to_shard( + global_episode_index, + shard_index=args.shard_index, + num_shards=args.num_shards, + ): + assigned_episode_count += 1 + assigned_grid_anchor_count += min( + len(episode.anchors), + args.grid_size, + ) + global_episode_index += 1 + + missing_exists = metadata["completed_episode_count"] < episode_count + if missing_exists: + model = build_framework(cfg) + checkpoint = torch.load( + args.checkpoint, + map_location="cpu", + weights_only=True, + ) + model.load_state_dict(checkpoint, strict=True) + model.requires_grad_(False).eval().to( + device="cuda", + dtype=getattr(torch, args.model_dtype), + ) + else: + model = None + + new_episodes = 0 + completed_assigned_episodes = 0 + completed_assigned_anchors = 0 + session_scoring_seconds = 0.0 + worker_progress = { + "schema_version": 1, + "status": "in_progress", + "identity_sha256": identity["sha256"], + "shard_index": args.shard_index, + "num_shards": args.num_shards, + "assigned_episode_count": assigned_episode_count, + "assigned_grid_anchor_count": assigned_grid_anchor_count, + "completed_episode_count": 0, + "completed_grid_anchor_count": 0, + "new_episode_count": 0, + "scoring_seconds": 0.0, + } + worker_path = _worker_progress_path( + args.cache_dir, + args.shard_index, + args.num_shards, + ) + if parallel_worker: + _atomic_json(worker_path, worker_progress) + try: + global_episode_index = 0 + for dataset_index, dataset_support in enumerate(manifest.datasets): + source_dataset = dataset.datasets[dataset_index] + for episode_index, episode in enumerate(dataset_support.episodes): + selected_for_worker = _episode_belongs_to_shard( + global_episode_index, + shard_index=args.shard_index, + num_shards=args.num_shards, + ) + global_episode_index += 1 + if not selected_for_worker: + continue + shard_path = _shard_path( + args.cache_dir, + dataset_index, + episode_index, + ) + positions = quantile_anchor_positions( + len(episode.anchors), + args.grid_size, + ) + base_indices = [ + int(episode.anchors[int(index)]) for index in positions + ] + if shard_path.exists(): + _load_valid_shard( + shard_path, + identity_sha256=identity["sha256"], + dataset_index=dataset_index, + episode_index=episode_index, + trajectory_id=episode.trajectory_id, + positions=positions, + base_indices=base_indices, + estimator_count=len(args.estimator_seeds), + ) + completed_assigned_episodes += 1 + completed_assigned_anchors += len(positions) + continue + + if model is None: + raise RuntimeError("Missing shard discovered after model skip") + examples, loaded_base_indices, transform_seeds = _load_examples( + source_dataset, + episode, + positions, + dataset_index=dataset_index, + episode_index=episode_index, + transform_seed_base=args.transform_seed_base, + ) + if loaded_base_indices != base_indices: + raise RuntimeError("Loaded base indices differ from frozen plan") + started = time.perf_counter() + estimators = _score_examples( + model, + examples, + estimator_seeds=args.estimator_seeds, + samples_per_estimator=args.samples_per_estimator, + batch_size=args.batch_size, + ) + torch.cuda.synchronize() + elapsed = time.perf_counter() - started + session_scoring_seconds += elapsed + grid_scores = estimators.mean(axis=1) + shard = { + "schema_version": SCHEMA_VERSION, + "identity_sha256": identity["sha256"], + "dataset_index": dataset_index, + "dataset_name": dataset_support.name, + "episode_index": episode_index, + "trajectory_id": _json_scalar(episode.trajectory_id), + "anchor_count": len(episode.anchors), + "anchor_positions": positions.tolist(), + "base_indices": base_indices, + "transform_seeds": transform_seeds, + "estimator_residuals": estimators.tolist(), + "grid_scores": grid_scores.tolist(), + "scoring_seconds": elapsed, + } + _atomic_json(shard_path, shard) + new_episodes += 1 + completed_assigned_episodes += 1 + completed_assigned_anchors += len(positions) + if parallel_worker: + worker_progress.update( + { + "completed_episode_count": completed_assigned_episodes, + "completed_grid_anchor_count": completed_assigned_anchors, + "new_episode_count": new_episodes, + "scoring_seconds": session_scoring_seconds, + } + ) + if new_episodes % 25 == 0: + _atomic_json(worker_path, worker_progress) + print( + f"residual-cache shard={args.shard_index}/" + f"{args.num_shards} episode=" + f"{completed_assigned_episodes}/" + f"{assigned_episode_count} seconds={elapsed:.3f}", + flush=True, + ) + else: + metadata["completed_episode_count"] += 1 + metadata["completed_grid_anchor_count"] += len(positions) + metadata["scoring_seconds"] += elapsed + _atomic_json(metadata_path, metadata) + print( + f"residual-cache episode={metadata['completed_episode_count']}/" + f"{episode_count} grid_anchors={len(positions)} " + f"seconds={elapsed:.3f}", + flush=True, + ) + if ( + args.max_new_episodes > 0 + and new_episodes >= args.max_new_episodes + ): + if parallel_worker: + worker_progress["status"] = "partial" + _atomic_json(worker_path, worker_progress) + print("Stopped at max-new-episodes; cache is resumable.", flush=True) + return + finally: + if model is not None: + del model + torch.cuda.empty_cache() + + if parallel_worker: + worker_progress.update( + { + "status": "complete", + "completed_episode_count": completed_assigned_episodes, + "completed_grid_anchor_count": completed_assigned_anchors, + "new_episode_count": new_episodes, + "scoring_seconds": session_scoring_seconds, + } + ) + _atomic_json(worker_path, worker_progress) + if completed_assigned_episodes != assigned_episode_count: + raise RuntimeError("Residual cache worker finished with missing episodes") + print(json.dumps(worker_progress, indent=2, sort_keys=True), flush=True) + return + + _refresh_progress(metadata, args.cache_dir, manifest) + if metadata["completed_episode_count"] != episode_count: + metadata["status"] = "in_progress" + _atomic_json(metadata_path, metadata) + raise RuntimeError("Residual cache is incomplete after scoring loop") + + proposal = _finalize_proposal( + manifest=manifest, + cache_dir=args.cache_dir, + identity=identity, + output=args.output, + alpha=args.alpha, + epsilon=args.epsilon, + allow_failed_gate=args.allow_failed_gate, + estimator_count=len(args.estimator_seeds), + gate_certificate=gate_certificate, + ) + metadata["status"] = "complete" + metadata["proposal"] = { + "path": str(args.output.resolve()), + "sha256": proposal.sha256, + "importance_weight_min": min( + float(np.min(weights)) + for dataset_weights in proposal.importance_weights + for weights in dataset_weights + ), + "importance_weight_max": max( + float(np.max(weights)) + for dataset_weights in proposal.importance_weights + for weights in dataset_weights + ), + } + _atomic_json(metadata_path, metadata) + print( + json.dumps( + { + "status": "complete", + "proposal": metadata["proposal"], + "scoring_seconds": metadata["scoring_seconds"], + "session_scoring_seconds": session_scoring_seconds, + }, + indent=2, + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/evaluate_bridge_384_no_video.sh b/scripts/evaluate_bridge_384_no_video.sh new file mode 100644 index 0000000000000000000000000000000000000000..5e7326d48c4ffb66cbffec485f010b7570ba10a3 --- /dev/null +++ b/scripts/evaluate_bridge_384_no_video.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -ne 6 ]]; then + echo "usage: $0 " >&2 + exit 64 +fi + +script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +runtime="${OPTR_RUNTIME_ROOT:?set OPTR_RUNTIME_ROOT to the runtime environment root}" +run_dir=$1 +step=$2 +gpu=$3 +port=$4 +arm=$5 +output_root=$6 + +for policy_seed in 518957 765353 686765 516408; do + "${script_dir}/evaluate_bridge_seed96_no_video.sh" \ + "${run_dir}" "${step}" "${gpu}" "${port}" "${policy_seed}" \ + "${arm}" "${output_root}" +done + +"${runtime}/.venv/bin/python" - "${output_root}" "${step}" "${arm}" <<'PY' +import json +import os +import sys +from pathlib import Path + +root = Path(sys.argv[1]) +step = int(sys.argv[2]) +arm = sys.argv[3] +seeds = [518957, 765353, 686765, 516408] +tasks = ("stack", "carrot", "spoon", "eggplant") +seed_rows = {} +task_rows = {task: {"success_count": 0, "rollout_count": 0} for task in tasks} +for seed in seeds: + summary = json.loads( + (root / f"policy_seed_{seed}" / f"step_{step}" / "four_task_summary.json") + .read_text(encoding="utf-8") + ) + seed_rows[str(seed)] = { + "success_count": summary["total_success"], + "rollout_count": summary["total_rollouts"], + "strict_success_rate": summary["micro_success_rate"], + } + for task, row in summary["tasks"].items(): + task_rows[task]["success_count"] += row["success_count"] + task_rows[task]["rollout_count"] += row["rollout_count"] +for row in task_rows.values(): + row["strict_success_rate"] = row["success_count"] / row["rollout_count"] +successes = sum(row["success_count"] for row in seed_rows.values()) +rollouts = sum(row["rollout_count"] for row in seed_rows.values()) +result = { + "schema_version": 1, + "status": "complete_no_video_four_seed_four_task384", + "arm": arm, + "checkpoint_step": step, + "policy_seeds": seeds, + "episodes_per_seed_task": 24, + "tasks": task_rows, + "seeds": seed_rows, + "total_success": successes, + "total_rollouts": rollouts, + "strict_success_rate": successes / rollouts, +} +target = root / f"step_{step}_four_seed_four_task384.json" +temporary = target.with_suffix(f".json.tmp.{os.getpid()}") +temporary.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") +os.replace(temporary, target) +print(json.dumps(result, sort_keys=True)) +PY diff --git a/scripts/evaluate_bridge_seed96_no_video.sh b/scripts/evaluate_bridge_seed96_no_video.sh new file mode 100644 index 0000000000000000000000000000000000000000..4d37d7e2b955983e86d9c46cb64428162a42431b --- /dev/null +++ b/scripts/evaluate_bridge_seed96_no_video.sh @@ -0,0 +1,304 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -ne 7 ]]; then + echo "usage: $0 " >&2 + exit 64 +fi + +run_dir=$1 +step=$2 +gpu=$3 +port=$4 +policy_seed=$5 +arm=$6 +eval_root=$7 + +project="${OPTR_PROJECT_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" +runtime="${OPTR_RUNTIME_ROOT:?set OPTR_RUNTIME_ROOT to the runtime environment root}" +simpler_root="${SIMPLER_ROOT:-${runtime}/SimplerEnv}" +episodes_per_task=24 + +output_root="${eval_root}/policy_seed_${policy_seed}" +step_dir="${output_root}/step_${step}" +server_pid="" + +mkdir -p "${step_dir}" +trap 'code=$?; trap - EXIT; if [[ -n "${server_pid}" ]]; then kill "${server_pid}" 2>/dev/null || true; wait "${server_pid}" 2>/dev/null || true; fi; printf "%s\n" "${code}" >"${step_dir}/exitcode.tmp.$$"; mv -f "${step_dir}/exitcode.tmp.$$" "${step_dir}/exitcode"; exit "${code}"' EXIT + +resolve_checkpoint() { + "${runtime}/.venv/bin/python" - "${run_dir}" "${step}" <<'PY' +import json +import os +import sys +from pathlib import Path + +run_dir = Path(sys.argv[1]) +step = int(sys.argv[2]) +checkpoint_dir = run_dir / "checkpoints" / f"steps_{step}_accelerate" +commit_path = checkpoint_dir / "checkpoint.json" +commit = json.loads(commit_path.read_text(encoding="utf-8")) +if commit["completed_steps"] != step or commit["world_size"] != 2: + raise RuntimeError(f"checkpoint identity mismatch: {checkpoint_dir}") +for entry in commit["files"]: + path = checkpoint_dir / entry["name"] + if not path.is_file() or path.stat().st_size != entry["bytes"]: + raise RuntimeError(f"incomplete checkpoint file: {path}") +source_record = next(entry for entry in commit["files"] if entry["name"] == "pytorch_model.pt") +source = checkpoint_dir / source_record["name"] +destination = run_dir / "checkpoints" / f"steps_{step}_pytorch_model.pt" +try: + os.link(source, destination) +except FileExistsError: + # Parallel policy seeds share this verified compatibility link. + pass +metadata = { + "schema_version": 1, + "completed_steps": step, + "world_size": 2, + "bytes": destination.stat().st_size, +} +sidecar = Path(f"{destination}.json") +temporary = Path(f"{sidecar}.tmp.{os.getpid()}") +temporary.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8") +os.replace(temporary, sidecar) +print(destination) +PY +} + +validate_task_summary() { + local summary=$1 + local task=$2 + "${runtime}/.venv/bin/python" - "${summary}" "${step}" "${task}" "${policy_seed}" "${episodes_per_task}" <<'PY' +import json +import sys +from pathlib import Path + +summary = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8")) +step = int(sys.argv[2]) +task = sys.argv[3] +seed = int(sys.argv[4]) +episodes = int(sys.argv[5]) +if ( + summary.get("status") != "complete_no_video_strict_success_panel" + or summary.get("checkpoint_step") != step + or summary.get("task") != task + or summary.get("policy_seeds") != [seed] + or summary.get("rollout_count") != episodes +): + raise RuntimeError(f"invalid task summary: {sys.argv[1]}") +PY +} + +run_panel() { + local checkpoint=$1 + local task=$2 + local panel_dir="${step_dir}/${task}" + local seed_dir="${panel_dir}/seed_${policy_seed}" + + if [[ -f "${panel_dir}/summary.json" ]]; then + validate_task_summary "${panel_dir}/summary.json" "${task}" + return + fi + + case "${task}" in + spoon) + env_name=PutSpoonOnTableClothInScene-v0 + scene_name=bridge_table_1_v1 + robot=widowx + overlay="${simpler_root}/ManiSkill2_real2sim/data/real_inpainting/bridge_real_eval_1.png" + robot_x=0.147 + robot_y=0.028 + ;; + carrot) + env_name=PutCarrotOnPlateInScene-v0 + scene_name=bridge_table_1_v1 + robot=widowx + overlay="${simpler_root}/ManiSkill2_real2sim/data/real_inpainting/bridge_real_eval_1.png" + robot_x=0.147 + robot_y=0.028 + ;; + stack) + env_name=StackGreenCubeOnYellowCubeBakedTexInScene-v0 + scene_name=bridge_table_1_v1 + robot=widowx + overlay="${simpler_root}/ManiSkill2_real2sim/data/real_inpainting/bridge_real_eval_1.png" + robot_x=0.147 + robot_y=0.028 + ;; + eggplant) + env_name=PutEggplantInBasketScene-v0 + scene_name=bridge_table_1_v2 + robot=widowx_sink_camera_setup + overlay="${simpler_root}/ManiSkill2_real2sim/data/real_inpainting/bridge_sink.png" + robot_x=0.127 + robot_y=0.06 + ;; + esac + + mkdir -p "${panel_dir}" + if [[ ! -f "${seed_dir}/result.json" ]]; then + if [[ -e "${seed_dir}" ]]; then + echo "refusing partial seed directory: ${seed_dir}" >&2 + exit 2 + fi + mkdir -p "${seed_dir}" + CUDA_VISIBLE_DEVICES="${gpu}" \ + "${runtime}/.venv/bin/python3" deployment/model_server/server_policy.py \ + --ckpt_path "${checkpoint}" \ + --port "${port}" \ + --idle_timeout -1 \ + --use_bf16 \ + --seed "${policy_seed}" \ + >"${seed_dir}/policy_server.log" 2>&1 & + server_pid=$! + + "${runtime}/.venv/bin/python" - "${port}" "${server_pid}" <<'PY' +import socket +import sys +import time +from pathlib import Path + +port = int(sys.argv[1]) +server_pid = int(sys.argv[2]) +deadline = time.monotonic() + 240 +while time.monotonic() < deadline: + if not Path(f"/proc/{server_pid}").exists(): + raise SystemExit("policy server exited during startup") + with socket.socket() as sock: + sock.settimeout(1) + if sock.connect_ex(("127.0.0.1", port)) == 0: + raise SystemExit(0) + time.sleep(2) +raise SystemExit("policy server did not open its port") +PY + + export CUDA_VISIBLE_DEVICES="${gpu}" + export LD_LIBRARY_PATH="${runtime}/system-libs/libvulkan1/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" + set +e + "${runtime}/.venv-simpler/bin/python" -c ' +import importlib +import runpy +evaluator = importlib.import_module("simpler_env.evaluation.maniskill2_evaluator") +evaluator.write_video = lambda *args, **kwargs: None +from examples.SimplerEnv.eval_files.model2simpler_interface import ModelClient +ModelClient.visualize_epoch = lambda *args, **kwargs: None +runpy.run_path("examples/SimplerEnv/eval_files/start_simpler_env.py", run_name="__main__") +' \ + --ckpt-path "${checkpoint}" \ + --port "${port}" \ + --robot "${robot}" \ + --policy-setup widowx_bridge \ + --control-freq 5 \ + --sim-freq 500 \ + --max-episode-steps 120 \ + --env-name "${env_name}" \ + --scene-name "${scene_name}" \ + --rgb-overlay-path "${overlay}" \ + --robot-init-x-range "${robot_x}" "${robot_x}" 1 \ + --robot-init-y-range "${robot_y}" "${robot_y}" 1 \ + --obj-variation-mode episode \ + --obj-episode-range 0 "${episodes_per_task}" \ + --robot-init-rot-quat-center 0 0 0 1 \ + --robot-init-rot-rpy-range 0 0 1 0 0 1 0 0 1 \ + --additional-env-save-tags "policy_seed_${policy_seed}_no_video" \ + --logging-dir "${seed_dir}/discarded_artifact_paths" \ + --result-json "${seed_dir}/result.json" \ + >"${seed_dir}/client.log" 2>&1 + client_exit=$? + set -e + kill "${server_pid}" 2>/dev/null || true + wait "${server_pid}" 2>/dev/null || true + server_pid="" + printf '%s\n' "${client_exit}" >"${seed_dir}/exitcode" + [[ "${client_exit}" == "0" ]] || exit "${client_exit}" + fi + + "${runtime}/.venv/bin/python" - "${panel_dir}" "${step}" "${task}" "${policy_seed}" "${episodes_per_task}" <<'PY' +import json +import os +import sys +from pathlib import Path + +panel = Path(sys.argv[1]) +step = int(sys.argv[2]) +task = sys.argv[3] +seed = int(sys.argv[4]) +episodes = int(sys.argv[5]) +result = json.loads((panel / f"seed_{seed}" / "result.json").read_text(encoding="utf-8")) +successes = [bool(value) for value in result["successes"]] +if len(successes) != episodes: + raise RuntimeError(f"incomplete result for seed {seed}: {len(successes)} != {episodes}") +summary = { + "schema_version": 1, + "status": "complete_no_video_strict_success_panel", + "checkpoint_step": step, + "task": task, + "policy_seeds": [seed], + "episodes_per_seed": episodes, + "success_count": sum(successes), + "rollout_count": episodes, + "strict_success_rate": sum(successes) / episodes, + "rows": [{"policy_seed": seed, "successes": successes, "success_count": sum(successes)}], +} +target = panel / "summary.json" +temporary = target.with_suffix(f".json.tmp.{os.getpid()}") +temporary.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8") +os.replace(temporary, target) +print(json.dumps(summary, sort_keys=True)) +PY +} + +checkpoint="$(resolve_checkpoint | tail -n 1)" +if [[ "${INFOFLOW_EVAL_VALIDATE_ONLY:-0}" == "1" ]]; then + printf 'arm=%s\nstep=%s\ngpu=%s\nport=%s\ncheckpoint=%s\noutput=%s\n' \ + "${arm}" "${step}" "${gpu}" "${port}" "${checkpoint}" "${step_dir}" + exit 0 +fi + +cd "${project}" +for task in stack carrot spoon eggplant; do + run_panel "${checkpoint}" "${task}" +done + +"${runtime}/.venv/bin/python" - "${step_dir}" "${arm}" "${step}" "${policy_seed}" "${episodes_per_task}" <<'PY' +import json +import os +import sys +from pathlib import Path + +step_dir = Path(sys.argv[1]) +arm = sys.argv[2] +step = int(sys.argv[3]) +seed = int(sys.argv[4]) +episodes = int(sys.argv[5]) +tasks = ("stack", "carrot", "spoon", "eggplant") +rows = {} +for task in tasks: + summary = json.loads((step_dir / task / "summary.json").read_text(encoding="utf-8")) + rows[task] = { + "success_count": summary["success_count"], + "rollout_count": summary["rollout_count"], + "strict_success_rate": summary["strict_success_rate"], + } +total_success = sum(row["success_count"] for row in rows.values()) +total_rollouts = sum(row["rollout_count"] for row in rows.values()) +result = { + "schema_version": 1, + "status": "complete_no_video_four_task96", + "arm": arm, + "checkpoint_step": step, + "policy_seed": seed, + "episodes_per_task": episodes, + "tasks": rows, + "total_success": total_success, + "total_rollouts": total_rollouts, + "micro_success_rate": total_success / total_rollouts, + "macro_success_rate": sum(row["strict_success_rate"] for row in rows.values()) / len(rows), +} +target = step_dir / "four_task_summary.json" +temporary = target.with_suffix(f".json.tmp.{os.getpid()}") +temporary.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") +os.replace(temporary, target) +print(json.dumps(result, sort_keys=True)) +PY diff --git a/scripts/optimization_fidelity_dense.py b/scripts/optimization_fidelity_dense.py new file mode 100644 index 0000000000000000000000000000000000000000..4715a3ef2f931d82ee6c62e26c2cd7816401c5a2 --- /dev/null +++ b/scripts/optimization_fidelity_dense.py @@ -0,0 +1,672 @@ +#!/usr/bin/env python3 +"""Compute dense-panel gradient targets and sparse proposal fidelity.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import math +import os +import random +import time +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np +import torch +from omegaconf import OmegaConf + +from scripts.validate_stratified_gradient_second_moment import ( + _flow_loss, + _prepare_conditioning, +) +from starVLA.dataloader.lerobot_datasets import get_vla_dataset +from starVLA.model.framework.base_framework import build_framework +from starVLA.model.framework.share_tools import apply_config_compat +from starVLA.scoring.flow_residual import ( + randomized_antithetic_stratified_flow_samples, +) +from starVLA.scoring.partial_gradient import finite_gradient_l2_norm +from starVLA.scoring.residual_cache import ( + interpolate_grid_scores, + percentile_intensity_masses, + quantile_anchor_positions, +) +from starVLA.scoring.utility_audit import gradient_second_moment + + +SCHEMA_VERSION = 1 + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def canonical_sha256(payload: dict) -> str: + encoded = json.dumps( + payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def atomic_json(path: Path, payload: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") + os.replace(temporary, path) + + +def seed_all(seed: int) -> None: + bounded = int(seed) % (2**32) + random.seed(bounded) + np.random.seed(bounded) + torch.manual_seed(bounded) + torch.cuda.manual_seed_all(bounded) + + +def load_dense_panel(dense_dir: Path) -> tuple[list[dict], dict]: + metadata_path = dense_dir / "dense_residual_metadata.json" + scores_path = dense_dir / "dense_residual_scores.npz" + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + dense = np.load(scores_path, allow_pickle=False) + scores = np.asarray(dense["scores"], dtype=np.float64) + offsets = np.asarray(dense["offsets"], dtype=np.int64) + records: list[dict] = [] + for index in range(len(offsets) - 1): + dataset_index = int(dense["dataset_index"][index]) + episode_index = int(dense["episode_index"][index]) + trajectory_index = int(dense["trajectory_index"][index]) + trajectory_id = str(dense["trajectory_id"][index]) + shard_path = dense_dir / "dense_shards" / ( + f"dataset_{dataset_index:02d}_episode_{episode_index:06d}.json" + ) + payload = json.loads(shard_path.read_text(encoding="utf-8")) + start, stop = int(offsets[index]), int(offsets[index + 1]) + expected = { + "dataset_index": dataset_index, + "episode_index": episode_index, + "trajectory_index": trajectory_index, + "trajectory_id": trajectory_id, + "anchor_count": stop - start, + } + observed = { + "dataset_index": int(payload["dataset_index"]), + "episode_index": int(payload["episode_index"]), + "trajectory_index": int(payload["trajectory_index"]), + "trajectory_id": str(payload["trajectory_id"]), + "anchor_count": int(payload["anchor_count"]), + } + if observed != expected: + raise RuntimeError(f"dense trajectory identity mismatch: {shard_path}") + if not ( + len(payload["anchor_positions"]) + == len(payload["base_indices"]) + == len(payload["transform_seeds"]) + == stop - start + ): + raise RuntimeError(f"dense anchor identity mismatch: {shard_path}") + records.append( + { + **observed, + "trajectory_id": payload["trajectory_id"], + "dataset_name": str(payload["dataset_name"]), + "manifest_sha256": str(payload["manifest_sha256"]), + "global_start": start, + "global_stop": stop, + "anchor_positions": [int(x) for x in payload["anchor_positions"]], + "base_indices": [int(x) for x in payload["base_indices"]], + "transform_seeds": [int(x) for x in payload["transform_seeds"]], + "dense_scores": scores[start:stop].tolist(), + } + ) + records.sort(key=lambda item: (item["dataset_index"], item["trajectory_index"])) + global_offset = 0 + for record in records: + record["global_start"] = global_offset + global_offset += record["anchor_count"] + record["global_stop"] = global_offset + if len(records) != int(metadata["trajectory_count"]): + raise RuntimeError("dense trajectory count differs from metadata") + if global_offset != int(metadata["valid_anchor_count"]): + raise RuntimeError("dense anchor count differs from metadata") + manifests = {record["manifest_sha256"] for record in records} + if len(manifests) != 1: + raise RuntimeError("dense shards contain multiple manifest identities") + identity = { + "dense_metadata_sha256": sha256_file(metadata_path), + "dense_scores_sha256": sha256_file(scores_path), + "dense_manifest_sha256": next(iter(manifests)), + "trajectory_count": len(records), + "anchor_count": global_offset, + "trajectory_keys_sha256": canonical_sha256( + { + "keys": [ + [ + record["dataset_index"], + record["dataset_name"], + record["trajectory_index"], + record["trajectory_id"], + record["episode_index"], + ] + for record in records + ] + } + ), + "anchor_keys_sha256": canonical_sha256( + { + "keys": [ + [ + record["dataset_index"], + record["trajectory_id"], + position, + base_index, + transform_seed, + record["global_start"] + local_index, + ] + for record in records + for local_index, (position, base_index, transform_seed) in enumerate( + zip( + record["anchor_positions"], + record["base_indices"], + record["transform_seeds"], + ) + ) + ] + } + ), + } + return records, identity + + +def load_runtime(config_path: Path, checkpoint_path: Path, output: Path): + cfg = apply_config_compat(OmegaConf.load(config_path)) + cfg.output_dir = str(output / "runtime") + cfg.datasets.vla_data.temporal_proposal_path = None + cfg.datasets.vla_data.valid_anchor_manifest_path = None + cfg.datasets.vla_data.valid_anchor_manifest_migration_path = None + dataset = get_vla_dataset( + data_cfg=cfg.datasets.vla_data, + balance_dataset_weights=cfg.datasets.vla_data.get( + "balance_dataset_weights", False + ), + balance_trajectory_weights=cfg.datasets.vla_data.get( + "balance_trajectory_weights", False + ), + seed=int(cfg.get("seed", 42)), + ) + model = build_framework(cfg) + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + model.load_state_dict(checkpoint, strict=True) + model.to(device="cuda", dtype=torch.float32) + return cfg, dataset, model + + +def measure_trajectory( + model, + source, + record: dict, + *, + estimator_seeds: list[int], + samples_per_estimator: int, +) -> tuple[np.ndarray, np.ndarray, float]: + anchor_count = record["anchor_count"] + estimator_values = np.empty( + (anchor_count, len(estimator_seeds)), dtype=np.float64 + ) + draw_norms = np.empty( + (anchor_count, len(estimator_seeds), samples_per_estimator), + dtype=np.float64, + ) + started = time.perf_counter() + for local_index, (base_index, transform_seed) in enumerate( + zip(record["base_indices"], record["transform_seeds"]) + ): + seed_all(transform_seed) + raw = source.get_step_data(record["trajectory_id"], base_index) + transformed = source.transforms(raw) + example = source._pack_sample(transformed) + global_index = record["global_start"] + local_index + for estimator_index, estimator_seed in enumerate(estimator_seeds): + model.requires_grad_(False) + model.train() + model.qwen_vl_interface.eval() + seed_all(estimator_seed * 1_000_003 + global_index) + conditioning = _prepare_conditioning(model, example) + actions = conditioning[1] + action_cfg = model.action_model.config + noises, times = randomized_antithetic_stratified_flow_samples( + actions, + sample_count=samples_per_estimator, + seed=estimator_seed * 1_000_033 + global_index, + beta_alpha=float(action_cfg.noise_beta_alpha), + beta_beta=float(action_cfg.noise_beta_beta), + noise_s=float(action_cfg.noise_s), + ) + model.action_model.requires_grad_(True) + parameters = list(model.action_model.parameters()) + for draw_index in range(samples_per_estimator): + seed_all( + estimator_seed * 1_000_037 + + global_index * 101 + + draw_index // 2 + ) + model.zero_grad(set_to_none=True) + loss = _flow_loss( + model, conditioning, noises[draw_index], times[draw_index] + ) + if loss.ndim != 0 or not torch.isfinite(loss): + raise FloatingPointError("flow loss must be a finite scalar") + loss.backward() + draw_norms[local_index, estimator_index, draw_index] = ( + finite_gradient_l2_norm(parameters) + ) + estimator_values[local_index, estimator_index] = math.sqrt( + float(np.mean(np.square(draw_norms[local_index, estimator_index]))) + ) + model.zero_grad(set_to_none=True) + torch.cuda.synchronize() + return estimator_values, draw_norms, time.perf_counter() - started + + +def worker(args: argparse.Namespace) -> None: + if not 0 <= args.shard_index < args.num_shards: + raise ValueError("invalid shard index") + records, dense_identity = load_dense_panel(args.dense_dir) + cfg_sha = sha256_file(args.config) + checkpoint_sha = sha256_file(args.checkpoint) + worker_identity = { + "schema_version": SCHEMA_VERSION, + "config_sha256": cfg_sha, + "checkpoint_sha256": checkpoint_sha, + "dense_identity": dense_identity, + "target_estimator_seeds": args.target_seeds, + "samples_per_estimator": args.samples_per_estimator, + "gradient_definition": "sqrt(mean over estimator seeds and draws of squared full action-model gradient L2 norm)", + "parameter_mode": "full_action_model", + "frozen_backbone_eval": True, + "action_model_mode": "train", + "model_dtype": "float32", + "global_anchor_index_seed_semantics": True, + "num_shards": args.num_shards, + } + identity_sha = canonical_sha256(worker_identity) + _, dataset, model = load_runtime(args.config, args.checkpoint, args.output) + if dataset.valid_anchor_manifest is None: + raise RuntimeError("current dataset has no valid-anchor manifest") + if dataset.valid_anchor_manifest.sha256 != dense_identity["dense_manifest_sha256"]: + raise RuntimeError("current dataset manifest differs from dense panel") + assigned = records[args.shard_index :: args.num_shards] + cache_dir = args.output / "dense_gradient_targets" + started = time.perf_counter() + new_anchors = 0 + completed_anchors = 0 + for item_index, record in enumerate(assigned): + key = ( + f"dataset_{record['dataset_index']:02d}_" + f"episode_{record['episode_index']:06d}.json" + ) + path = cache_dir / key + if path.exists(): + cached = json.loads(path.read_text(encoding="utf-8")) + if ( + cached.get("identity_sha256") != identity_sha + or cached.get("trajectory_identity") + != {field: record[field] for field in ( + "dataset_index", "dataset_name", "episode_index", + "trajectory_index", "trajectory_id", "anchor_count", + "global_start", "global_stop", "anchor_positions", + "base_indices", "transform_seeds", + )} + ): + raise RuntimeError(f"incompatible existing cache: {path}") + completed_anchors += record["anchor_count"] + continue + source = dataset.datasets[record["dataset_index"]] + estimators, draws, elapsed = measure_trajectory( + model, + source, + record, + estimator_seeds=args.target_seeds, + samples_per_estimator=args.samples_per_estimator, + ) + target = np.sqrt(np.mean(np.square(estimators), axis=1)) + if not np.isfinite(target).all() or np.any(target <= 0): + raise FloatingPointError("gradient targets must be positive and finite") + trajectory_identity = {field: record[field] for field in ( + "dataset_index", "dataset_name", "episode_index", + "trajectory_index", "trajectory_id", "anchor_count", + "global_start", "global_stop", "anchor_positions", + "base_indices", "transform_seeds", + )} + atomic_json( + path, + { + "schema_version": SCHEMA_VERSION, + "status": "complete", + "identity_sha256": identity_sha, + "worker_identity": worker_identity, + "trajectory_identity": trajectory_identity, + "target_estimators": estimators.tolist(), + "target_draw_norms": draws.tolist(), + "target_gradient_rms": target.tolist(), + "scoring_seconds": elapsed, + "completed_at": datetime.now(timezone.utc).isoformat(), + }, + ) + new_anchors += record["anchor_count"] + completed_anchors += record["anchor_count"] + print( + f"shard={args.shard_index}/{args.num_shards} " + f"trajectory={item_index + 1}/{len(assigned)} " + f"anchors={completed_anchors} new={new_anchors} seconds={elapsed:.2f}", + flush=True, + ) + atomic_json( + args.output / "logs" / f"worker_{args.shard_index:02d}_of_{args.num_shards:02d}.json", + { + **worker_identity, + "identity_sha256": identity_sha, + "status": "complete", + "shard_index": args.shard_index, + "assigned_trajectories": len(assigned), + "completed_anchors": completed_anchors, + "new_anchors": new_anchors, + "wall_clock_seconds": time.perf_counter() - started, + "completed_at": datetime.now(timezone.utc).isoformat(), + }, + ) + + +def q_from_scores(scores: np.ndarray, alpha: float, epsilon: float) -> np.ndarray: + masses = percentile_intensity_masses(scores, epsilon=epsilon) + baseline = np.full(len(scores), 1.0 / len(scores), dtype=np.float64) + information = masses / masses.sum() + return (1.0 - alpha) * baseline + alpha * information + + +def describe(values: list[float]) -> dict: + array = np.asarray(values, dtype=np.float64) + return { + "mean": float(array.mean()), + "median": float(np.median(array)), + "std": float(array.std(ddof=1)), + "p05": float(np.quantile(array, 0.05)), + "p25": float(np.quantile(array, 0.25)), + "p75": float(np.quantile(array, 0.75)), + "p95": float(np.quantile(array, 0.95)), + "min": float(array.min()), + "max": float(array.max()), + } + + +def write_csv(path: Path, rows: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + +def finalize(args: argparse.Namespace) -> None: + records, dense_identity = load_dense_panel(args.dense_dir) + cache_dir = args.output / "dense_gradient_targets" + cache_files = sorted(cache_dir.glob("dataset_*_episode_*.json")) + if len(cache_files) != len(records): + raise RuntimeError( + f"gradient cache file count {len(cache_files)} != {len(records)}" + ) + per_trajectory: list[dict] = [] + all_targets = np.empty(dense_identity["anchor_count"], dtype=np.float64) + cache_identity = None + q_checks = [] + manual_checks = [] + for record_index, record in enumerate(records): + path = cache_dir / ( + f"dataset_{record['dataset_index']:02d}_" + f"episode_{record['episode_index']:06d}.json" + ) + if not path.exists(): + raise FileNotFoundError(path) + payload = json.loads(path.read_text(encoding="utf-8")) + if payload.get("status") != "complete": + raise RuntimeError(f"incomplete cache: {path}") + if cache_identity is None: + cache_identity = payload["identity_sha256"] + if payload["identity_sha256"] != cache_identity: + raise RuntimeError("gradient caches contain multiple identities") + observed = payload["trajectory_identity"] + expected = {field: record[field] for field in ( + "dataset_index", "dataset_name", "episode_index", + "trajectory_index", "trajectory_id", "anchor_count", + "global_start", "global_stop", "anchor_positions", + "base_indices", "transform_seeds", + )} + if observed != expected: + raise RuntimeError(f"gradient anchor alignment mismatch: {path}") + h = np.asarray(payload["target_gradient_rms"], dtype=np.float64) + if len(h) != record["anchor_count"]: + raise RuntimeError(f"gradient anchor count mismatch: {path}") + if not np.isfinite(h).all() or np.any(h <= 0): + raise FloatingPointError(f"invalid gradient target: {path}") + all_targets[record["global_start"] : record["global_stop"]] = h + dense_scores = np.asarray(record["dense_scores"], dtype=np.float64) + proposals: dict[str, np.ndarray] = { + "p": np.full(len(h), 1.0 / len(h), dtype=np.float64), + "dense": q_from_scores(dense_scores, args.alpha, args.epsilon), + } + for k in (4, 8, 16): + positions = quantile_anchor_positions(len(h), k) + estimate = interpolate_grid_scores( + len(h), positions, dense_scores[positions] + ) + proposals[str(k)] = q_from_scores(estimate, args.alpha, args.epsilon) + for label, proposal in proposals.items(): + q_checks.append( + abs(float(proposal.sum()) - 1.0) <= 1e-10 + and np.isfinite(proposal).all() + and np.all(proposal > 0) + ) + moments = { + label: gradient_second_moment(proposal, h) + for label, proposal in proposals.items() + } + if record_index in {0, len(records) // 2, len(records) - 1}: + direct = float( + np.sum(np.square(proposals["p"] * h) / proposals["dense"]) + ) + manual_checks.append(abs(direct - moments["dense"]) <= 1e-12) + row = { + "dataset_index": record["dataset_index"], + "dataset_name": record["dataset_name"], + "trajectory_index": record["trajectory_index"], + "trajectory_id": record["trajectory_id"], + "anchor_count": record["anchor_count"], + "M_p": moments["p"], + "M_q4": moments["4"], + "M_q8": moments["8"], + "M_q16": moments["16"], + "M_qDense": moments["dense"], + } + for label in ("4", "8", "16"): + row[f"R_{label}"] = moments[label] / moments["dense"] + row[f"q{label}_over_p"] = moments[label] / moments["p"] + row["qDense_over_p"] = moments["dense"] / moments["p"] + per_trajectory.append(row) + if not np.isfinite(all_targets).all() or np.any(all_targets <= 0): + raise FloatingPointError("assembled gradient targets are incomplete") + aggregate = { + field: float(np.mean([row[field] for row in per_trajectory])) + for field in ("M_p", "M_q4", "M_q8", "M_q16", "M_qDense") + } + ratios = { + "R_4": aggregate["M_q4"] / aggregate["M_qDense"], + "R_8": aggregate["M_q8"] / aggregate["M_qDense"], + "R_16": aggregate["M_q16"] / aggregate["M_qDense"], + "q4_over_p": aggregate["M_q4"] / aggregate["M_p"], + "q8_over_p": aggregate["M_q8"] / aggregate["M_p"], + "q16_over_p": aggregate["M_q16"] / aggregate["M_p"], + "qDense_over_p": aggregate["M_qDense"] / aggregate["M_p"], + } + trajectory_statistics = {} + for label in ("4", "8", "16"): + r_values = [row[f"R_{label}"] for row in per_trajectory] + p_values = [row[f"q{label}_over_p"] for row in per_trajectory] + trajectory_statistics[label] = { + "relative_to_dense": describe(r_values), + "relative_to_uniform": describe(p_values), + "better_than_dense": sum(value < 1.0 - args.comparison_tolerance for value in r_values), + "equal_to_dense": sum(abs(value - 1.0) <= args.comparison_tolerance for value in r_values), + "worse_than_dense": sum(value > 1.0 + args.comparison_tolerance for value in r_values), + "better_than_uniform": sum(value < 1.0 - args.comparison_tolerance for value in p_values), + "equal_to_uniform": sum(abs(value - 1.0) <= args.comparison_tolerance for value in p_values), + "worse_than_uniform": sum(value > 1.0 + args.comparison_tolerance for value in p_values), + } + dense_self_ratio = aggregate["M_qDense"] / aggregate["M_qDense"] + checks = { + "trajectory_count_512": len(per_trajectory) == 512, + "anchor_count_19375": len(all_targets) == 19375, + "cache_file_count_matches_trajectories": len(cache_files) == len(records), + "trajectory_and_anchor_identity_aligned": True, + "all_anchor_targets_positive_finite": bool( + np.isfinite(all_targets).all() and np.all(all_targets > 0) + ), + "proposal_normalization_full_support_finite": all(q_checks), + "dense_self_ratio_one": abs(dense_self_ratio - 1.0) <= 1e-15, + "manual_second_moment_checks": all(manual_checks), + "all_moments_finite": all(np.isfinite(list(aggregate.values()))), + } + write_csv(args.output / "optimization_fidelity_per_trajectory.csv", per_trajectory) + readback = list( + csv.DictReader( + (args.output / "optimization_fidelity_per_trajectory.csv").open( + encoding="utf-8" + ) + ) + ) + reconstructed = { + field: float(np.mean([float(row[field]) for row in readback])) + for field in aggregate + } + checks["aggregate_reconstructs_from_per_trajectory_csv"] = all( + math.isclose( + reconstructed[field], aggregate[field], rel_tol=1e-12, abs_tol=1e-12 + ) + for field in aggregate + ) + summary_rows = [] + for label, field in (("Uniform", "M_p"), ("K4", "M_q4"), ("K8", "M_q8"), ("K16", "M_q16"), ("Dense", "M_qDense")): + value = aggregate[field] + summary_rows.append( + { + "proposal": label, + "aggregate_second_moment": value, + "M_over_Mp": value / aggregate["M_p"], + "improvement_vs_uniform_percent": 100.0 * (1.0 - value / aggregate["M_p"]), + "M_over_Mdense": value / aggregate["M_qDense"], + "gap_vs_dense_percent": 100.0 * (value / aggregate["M_qDense"] - 1.0), + } + ) + write_csv(args.output / "optimization_fidelity_summary.csv", summary_rows) + np.savez_compressed( + args.output / "dense_gradient_targets.npz", + target_gradient_rms=all_targets, + ) + metadata = { + "schema_version": SCHEMA_VERSION, + "status": "complete" if all(checks.values()) else "failed_checks", + "started_at": args.started_at, + "completed_at": datetime.now(timezone.utc).isoformat(), + "git_commit": None, + "git_status": "unavailable_server_copy_has_no_dot_git", + "code_path": str(Path(__file__).resolve()), + "config": str(args.config.resolve()), + "config_sha256": sha256_file(args.config), + "checkpoint": str(args.checkpoint.resolve()), + "checkpoint_sha256": sha256_file(args.checkpoint), + "dense_dir": str(args.dense_dir.resolve()), + "dense_identity": dense_identity, + "gradient_cache_identity_sha256": cache_identity, + "trajectory_count": len(per_trajectory), + "anchor_count": len(all_targets), + "proposal": { + "alpha": args.alpha, + "epsilon": args.epsilon, + "construction": "within-trajectory percentile-rank safe mixture after linear interpolation at deterministic quantile anchors", + "K": [4, 8, 16], + }, + "target_estimator": { + "definition": "sqrt(mean over estimator seeds and draws of squared full action-model gradient L2 norm)", + "seeds": args.target_seeds, + "samples_per_estimator": args.samples_per_estimator, + "randomized_beta_stratification": True, + "antithetic_gaussian_noise": True, + "common_dropout_within_antithetic_pair": True, + "frozen_backbone_eval": True, + "action_model_mode": "train", + "model_dtype": "float32", + }, + "aggregation": "equal empirical trajectory weight; mean per-trajectory second moment, matching hierarchical_alignment_audit", + "aggregate_second_moments": aggregate, + "aggregate_ratios": ratios, + "trajectory_statistics": trajectory_statistics, + "sanity_checks": checks, + } + atomic_json(args.output / "optimization_fidelity_metadata.json", metadata) + if not all(checks.values()): + raise RuntimeError(f"optimization fidelity checks failed: {checks}") + readme = f"""# Optimization Fidelity on the Dense Temporal Panel + +- Status: complete +- Panel: {len(per_trajectory)} trajectories, {len(all_targets)} valid anchors +- Checkpoint: `{args.checkpoint.resolve()}` +- Gradient target: full action-model gradient RMS, four estimator seeds and eight draws per seed +- Proposal: within-trajectory percentile safe mixture, alpha={args.alpha}, epsilon={args.epsilon} +- Aggregation: equal empirical trajectory weight, matching `hierarchical_alignment_audit` + +This is a 512-trajectory Dense-reference offline audit. It is not a full-manifest +training proposal, a continuation run, or a downstream success-rate evaluation. + +Formal values are stored in `optimization_fidelity_summary.csv` and +`optimization_fidelity_metadata.json`. All required sanity checks passed. +""" + (args.output / "README.md").write_text(readme, encoding="utf-8") + print(json.dumps(metadata, indent=2), flush=True) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--dense-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument( + "--target-seeds", type=int, nargs="+", default=[27300, 27301, 27302, 27303] + ) + parser.add_argument("--samples-per-estimator", type=int, default=8) + subparsers = parser.add_subparsers(dest="command", required=True) + worker_parser = subparsers.add_parser("worker") + worker_parser.add_argument("--shard-index", type=int, required=True) + worker_parser.add_argument("--num-shards", type=int, required=True) + worker_parser.set_defaults(func=worker) + final_parser = subparsers.add_parser("finalize") + final_parser.add_argument("--alpha", type=float, default=0.2) + final_parser.add_argument("--epsilon", type=float, default=0.1) + final_parser.add_argument("--comparison-tolerance", type=float, default=1e-12) + final_parser.add_argument("--started-at", type=str, required=True) + final_parser.set_defaults(func=finalize) + args = parser.parse_args() + if len(args.target_seeds) != 4 or len(set(args.target_seeds)) != 4: + raise ValueError("exactly four unique target estimator seeds are required") + if args.samples_per_estimator != 8: + raise ValueError("the Gate42F contract requires eight samples per estimator") + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/run_florence_bridge_arm.sh b/scripts/run_florence_bridge_arm.sh new file mode 100644 index 0000000000000000000000000000000000000000..af0efb5b95f74c31bc9467af0a78b998016fd1b3 --- /dev/null +++ b/scripts/run_florence_bridge_arm.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -lt 2 || $# -gt 3 ]]; then + echo "usage: $0 [proposal.json]" >&2 + exit 64 +fi + +arm=$1 +run_id=$2 +proposal=${3:-} +project="${OPTR_PROJECT_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" +python="${OPTR_PYTHON:-python}" +config="${project}/configs/florence_bridge_20k_plus500.yaml" +accelerate_config="${project}/configs/deepspeed_florence_2gpu.yaml" + +required_env=(OPTR_RUN_ROOT FLORENCE_MODEL_PATH BRIDGE_DATA_ROOT BRIDGE_STATISTICS_PATH BRIDGE_PROVENANCE_PATH BRIDGE_FINGERPRINT_PATH BRIDGE_MANIFEST_PATH COMMON_START_20K) +for name in "${required_env[@]}"; do + [[ -n "${!name:-}" ]] || { echo "missing environment variable: $name" >&2; exit 64; } +done +[[ -d "$COMMON_START_20K" ]] || { echo "common start not found: $COMMON_START_20K" >&2; exit 66; } +[[ -f "$COMMON_START_20K/checkpoint.json" ]] || { echo "checkpoint.json missing under common start" >&2; exit 66; } +[[ -f "$BRIDGE_STATISTICS_PATH" ]] || { echo "statistics file missing: $BRIDGE_STATISTICS_PATH" >&2; exit 66; } +[[ -f "$BRIDGE_PROVENANCE_PATH" ]] || { echo "provenance file missing: $BRIDGE_PROVENANCE_PATH" >&2; exit 66; } +[[ -f "$BRIDGE_FINGERPRINT_PATH" ]] || { echo "content fingerprint missing: $BRIDGE_FINGERPRINT_PATH" >&2; exit 66; } +[[ -f "$BRIDGE_MANIFEST_PATH" ]] || { echo "anchor manifest missing: $BRIDGE_MANIFEST_PATH" >&2; exit 66; } + +case "$arm" in + uniform) + [[ -z "$proposal" ]] || { echo "uniform does not accept a proposal" >&2; exit 64; } + entry=starVLA/training/train_starvla.py + proposal_args=(--datasets.vla_data.temporal_proposal_path=null) + ;; + residual) + [[ -f "$proposal" ]] || { echo "proposal not found: $proposal" >&2; exit 66; } + entry=starVLA/training/train_starvla.py + proposal_args=(--datasets.vla_data.temporal_proposal_path="$proposal") + ;; + *) echo "unknown arm: $arm" >&2; exit 64 ;; +esac + +cd "$project" +export PYTHONPATH="${project}:${PYTHONPATH:-}" +export WANDB_MODE="${WANDB_MODE:-offline}" +export PYTHONUNBUFFERED=1 +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" +mkdir -p "$OPTR_RUN_ROOT" +"$python" -m accelerate.commands.launch \ + --config_file "$accelerate_config" \ + --num_processes 2 \ + "$entry" \ + --config_yaml "$config" \ + --run_id "$run_id" \ + --run_root_dir "$OPTR_RUN_ROOT" \ + "${proposal_args[@]}" diff --git a/scripts/run_frameskip_plus500.sh b/scripts/run_frameskip_plus500.sh new file mode 100644 index 0000000000000000000000000000000000000000..660898ebb4052a1eee893a048a0a6a22a58844c1 --- /dev/null +++ b/scripts/run_frameskip_plus500.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -ne 2 ]]; then + echo "usage: $0 " >&2 + exit 64 +fi + +run_id=$1 +cache=$2 +[[ -f "$cache" ]] || { echo "cache not found: $cache" >&2; exit 66; } +project="${OPTR_PROJECT_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" +python="${OPTR_PYTHON:-python}" +export FRAMESKIP_CACHE="$cache" +export FRAMESKIP_START_GLOBAL_DRAW_ID="${FRAMESKIP_START_GLOBAL_DRAW_ID:-1280000}" +export PYTHONPATH="${project}:${project}/third_party/frameskip:${PYTHONPATH:-}" +export WANDB_MODE="${WANDB_MODE:-offline}" +export PYTHONUNBUFFERED=1 + +cd "$project" +"$python" -m accelerate.commands.launch \ + --config_file configs/deepspeed_florence_2gpu.yaml \ + --num_processes 2 \ + external_baselines/frameskip_adapted/train_frameskip_continuation.py \ + --config_yaml configs/florence_bridge_20k_plus500.yaml \ + --run_id "$run_id" \ + --run_root_dir "$OPTR_RUN_ROOT" \ + --datasets.vla_data.temporal_proposal_path=null diff --git a/scripts/sparse_temporal_scoring_ablation.py b/scripts/sparse_temporal_scoring_ablation.py new file mode 100644 index 0000000000000000000000000000000000000000..8376603007eacb0da046db46c95c4cfb368085e3 --- /dev/null +++ b/scripts/sparse_temporal_scoring_ablation.py @@ -0,0 +1,358 @@ +#!/usr/bin/env python3 +"""Offline Sparse Temporal Scoring ablation for the Florence-Bridge study. + +This is an additive experiment driver. It imports the verified StarVLA +scorer/dataset/model interfaces and writes only to the requested artifact +directory. It never changes the main proposal or training implementation. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import time +from pathlib import Path + +import numpy as np +import torch +from omegaconf import OmegaConf +from scipy.stats import spearmanr + +from starVLA.dataloader.lerobot_datasets import get_vla_dataset +from starVLA.scoring.residual_cache import percentile_intensity_masses, quantile_anchor_positions +from starVLA.scoring.starvla_likelihood import score_starvla_flow_matching_residual +from starVLA.model.framework.base_framework import build_framework +from starVLA.model.framework.share_tools import apply_config_compat + + +def sha256_file(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as f: + for chunk in iter(lambda: f.read(8 * 1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def seed_all(seed: int) -> None: + seed = int(seed) % (2**32) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def load_runtime(config_path: Path, checkpoint: Path): + cfg = apply_config_compat(OmegaConf.load(config_path)) + cfg.output_dir = str(config_path.parent / ".sparse_temporal_runtime") + cfg.datasets.vla_data.temporal_proposal_path = None + # The frozen gate manifest is used only for trajectory identity selection. + # Loading it through the training migration guard would bind this additive + # analysis to the historical migration hash; the current manifest supplies + # the same valid anchors and is sufficient for direct indexed reads. + cfg.datasets.vla_data.valid_anchor_manifest_path = None + cfg.datasets.vla_data.valid_anchor_manifest_migration_path = None + dataset = get_vla_dataset( + data_cfg=cfg.datasets.vla_data, + balance_dataset_weights=cfg.datasets.vla_data.get("balance_dataset_weights", False), + balance_trajectory_weights=cfg.datasets.vla_data.get("balance_trajectory_weights", False), + seed=int(cfg.get("seed", 42)), + ) + model = build_framework(cfg) + state = torch.load(checkpoint, map_location="cpu", weights_only=True) + model.load_state_dict(state, strict=True) + model.requires_grad_(False).eval().to(device="cuda", dtype=torch.float32) + return cfg, dataset, model + + +def selected_groups(gate_path: Path) -> set[tuple[int, str]]: + payload = json.loads(gate_path.read_text(encoding="utf-8")) + groups: set[tuple[int, str]] = set() + for row in payload["anchors"]: + groups.add((int(row["dataset_index"]), str(row["trajectory_id"]))) + if len(groups) != int(payload["selection_protocol"]["episode_count"]): + raise RuntimeError("gate selection contains duplicate or missing trajectory groups") + return groups + + +def score_one(model, examples: list[dict], seeds: list[int], samples: int, batch_size: int): + rows = [] + for start in range(0, len(examples), batch_size): + batch = examples[start : start + batch_size] + with torch.no_grad(): + result = score_starvla_flow_matching_residual( + model, + batch, + estimator_seeds=seeds, + samples_per_estimator=samples, + ) + rows.append(result.estimator_residuals.transpose(0, 1).float().cpu().numpy()) + return np.concatenate(rows, axis=0) + + +def score(args: argparse.Namespace) -> None: + out = args.output.resolve() + shards = out / "dense_shards" + shards.mkdir(parents=True, exist_ok=True) + metadata_path = out / "dense_residual_metadata.json" + gate_groups = selected_groups(args.gate) + cfg, dataset, model = load_runtime(args.config, args.checkpoint) + manifest = dataset.valid_anchor_manifest + if manifest is None: + raise RuntimeError("traceable valid-anchor manifest is required") + selected = [] + for di, support in enumerate(manifest.datasets): + for ei, episode in enumerate(support.episodes): + key = (di, str(episode.trajectory_id)) + if key in gate_groups: + selected.append((di, ei, support, episode)) + if len(selected) != len(gate_groups): + raise RuntimeError("selected gate trajectories could not be found in current manifest") + selected.sort(key=lambda x: (x[0], x[3].trajectory_index)) + if not 0 <= args.shard_index < args.num_shards: + raise ValueError("invalid shard index") + all_selected_count = len(selected) + selected = selected[args.shard_index::args.num_shards] + started_all = time.perf_counter() + completed = 0 + anchors_done = 0 + scoring_seconds = 0.0 + for di, ei, support, episode in selected: + path = shards / f"dataset_{di:02d}_episode_{ei:06d}.json" + if path.exists(): + payload = json.loads(path.read_text(encoding="utf-8")) + if payload.get("manifest_sha256") == manifest.sha256 and payload.get("anchor_count") == len(episode.anchors): + completed += 1 + anchors_done += len(episode.anchors) + continue + raise RuntimeError(f"incompatible existing shard: {path}") + source = dataset.datasets[di] + examples, base_indices, transform_seeds = [], [], [] + for position, base_index in enumerate(episode.anchors): + transform_seed = 16500 + di * 10_000_019 + ei * 10_007 + int(position) + seed_all(transform_seed) + raw = source.get_step_data(episode.trajectory_id, int(base_index)) + transformed = source.transforms(raw) + examples.append(source._pack_sample(transformed)) + base_indices.append(int(base_index)) + transform_seeds.append(int(transform_seed)) + started = time.perf_counter() + estimators = score_one(model, examples, args.estimator_seeds, args.samples_per_estimator, args.batch_size) + torch.cuda.synchronize() + elapsed = time.perf_counter() - started + scores = estimators.mean(axis=1) + payload = { + "schema_version": 1, + "manifest_sha256": manifest.sha256, + "dataset_index": di, + "dataset_name": support.name, + "episode_index": ei, + "trajectory_index": episode.trajectory_index, + "trajectory_id": episode.trajectory_id, + "anchor_count": len(episode.anchors), + "anchor_positions": list(range(len(episode.anchors))), + "base_indices": base_indices, + "transform_seeds": transform_seeds, + "estimator_seeds": args.estimator_seeds, + "samples_per_estimator": args.samples_per_estimator, + "estimator_residuals": estimators.tolist(), + "dense_scores": scores.tolist(), + "scoring_seconds": elapsed, + } + temporary = path.with_suffix(".tmp") + temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") + os.replace(temporary, path) + completed += 1 + anchors_done += len(episode.anchors) + scoring_seconds += elapsed + if completed % 8 == 0: + print(f"dense trajectory={completed}/{len(selected)} anchors={anchors_done} seconds={elapsed:.2f}", flush=True) + # Flatten this worker's resumable shards to a portable, non-pickle NPZ. + rows = [] + offsets = [0] + for di, ei, support, episode in selected: + payload = json.loads((shards / f"dataset_{di:02d}_episode_{ei:06d}.json").read_text(encoding="utf-8")) + scores = np.asarray(payload["dense_scores"], dtype=np.float64) + rows.append(scores) + offsets.append(offsets[-1] + len(scores)) + np.savez_compressed( + out / f"dense_residual_scores_shard_{args.shard_index:02d}_of_{args.num_shards:02d}.npz", + scores=np.concatenate(rows), + offsets=np.asarray(offsets, dtype=np.int64), + dataset_index=np.asarray([x[0] for x in selected], dtype=np.int64), + episode_index=np.asarray([x[1] for x in selected], dtype=np.int64), + trajectory_index=np.asarray([x[3].trajectory_index for x in selected], dtype=np.int64), + trajectory_id=np.asarray([str(x[3].trajectory_id) for x in selected]), + ) + metadata = { + "schema_version": 1, + "status": "complete", + "scorer": "Florence-GR00T Bridge Flow-Matching residual", + "config": str(args.config.resolve()), + "config_sha256": sha256_file(args.config), + "checkpoint": str(args.checkpoint.resolve()), + "checkpoint_sha256": sha256_file(args.checkpoint), + "manifest_sha256": manifest.sha256, + "source_gate": str(args.gate.resolve()), + "selection_protocol": {"all_trajectories": all_selected_count, "worker_trajectories": len(selected), "shard_index": args.shard_index, "num_shards": args.num_shards, "dataset_episodes": [sum(x[0] == di for x in selected) for di in range(len(manifest.datasets))]}, + "estimator_seeds": args.estimator_seeds, + "samples_per_estimator": args.samples_per_estimator, + "valid_anchor_count": int(sum(len(x[3].anchors) for x in selected)), + "scoring_seconds": float(scoring_seconds), + "wall_clock_seconds": float(time.perf_counter() - started_all), + "measured": True, + "notes": "Dense scores are computed once; K=4/8/16 are reconstructed offline from this cache.", + } + (out / f"dense_residual_metadata_shard_{args.shard_index:02d}_of_{args.num_shards:02d}.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") + print(json.dumps(metadata, indent=2), flush=True) + + +def bootstrap(values: list[float], seed: int = 9042, count: int = 2000) -> tuple[float, float]: + rng = np.random.default_rng(seed) + x = np.asarray(values, dtype=np.float64) + if len(x) == 0: + return float("nan"), float("nan") + draws = rng.integers(0, len(x), size=(count, len(x))) + means = x[draws].mean(axis=1) + return float(np.quantile(means, 0.025)), float(np.quantile(means, 0.975)) + + +def q_from_scores(scores: np.ndarray, alpha: float, epsilon: float) -> np.ndarray: + mass = percentile_intensity_masses(scores, epsilon=epsilon) + baseline = np.full(len(scores), 1.0 / len(scores)) + residual = mass / mass.sum() + return (1.0 - alpha) * baseline + alpha * residual + + +def analyze(args: argparse.Namespace) -> None: + out = args.output.resolve() + dense = np.load(out / "dense_residual_scores.npz", allow_pickle=False) + scores_flat = np.asarray(dense["scores"], dtype=np.float64) + offsets = dense["offsets"] + rows = [] + smooth_rows = [] + interpolation_cases = [] + for i in range(len(offsets) - 1): + scores = scores_flat[offsets[i] : offsets[i + 1]] + key = f"{int(dense['dataset_index'][i])}:{dense['trajectory_id'][i]}" + for k in (4, 8, 16): + positions = quantile_anchor_positions(len(scores), k) + estimate = np.interp(np.arange(len(scores), dtype=float), positions.astype(float), scores[positions]) + rho = float(spearmanr(estimate, scores).statistic) if len(np.unique(scores)) > 1 else 1.0 + top25 = max(1, len(scores) // 4) + top10 = max(1, len(scores) // 10) + dense_top25 = set(np.argsort(scores)[-top25:]) + sparse_top25 = set(np.argsort(estimate)[-top25:]) + dense_top10 = set(np.argsort(scores)[-top10:]) + sparse_top10 = set(np.argsort(estimate)[-top10:]) + q_sparse = q_from_scores(estimate, args.alpha, args.epsilon) + q_dense = q_from_scores(scores, args.alpha, args.epsilon) + tv = 0.5 * float(np.abs(q_sparse - q_dense).sum()) + rows.append({"trajectory": key, "K": k, "spearman": rho, "top25_overlap": len(dense_top25 & sparse_top25) / top25, "top10_recall": len(dense_top10 & sparse_top10) / top10, "proposal_tv": tv}) + if k == 8: + err = np.abs(scores - estimate) + scale = float(np.mean(np.abs(scores)) + args.normalization_epsilon) + smooth_rows.append({"trajectory": key, "adjacent_median": float(np.median(np.abs(np.diff(scores)))), "adjacent_p90": float(np.quantile(np.abs(np.diff(scores)), .90)), "adjacent_p95": float(np.quantile(np.abs(np.diff(scores)), .95)), "interpolation_median": float(np.median(err)), "interpolation_p90": float(np.quantile(err, .90)), "interpolation_p95": float(np.quantile(err, .95)), "normalized_error_median": float(np.median(err) / scale), "normalized_error_p95": float(np.quantile(err, .95) / scale), "max_interpolation_error": float(err.max()), "dense_mean_abs": float(np.mean(np.abs(scores)),), "max_error_position": int(np.argmax(err))}) + interpolation_cases.append((float(err.max()), key, rho, len(dense_top25 & sparse_top25) / top25, float(err.max()), len(scores))) + # Keep CSV writing dependency-free. + import csv + def write_csv(path: Path, records: list[dict]): + fields = list(records[0]) if records else [] + with path.open("w", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=fields); w.writeheader(); w.writerows(records) + write_csv(out / "temporal_fidelity_per_trajectory.csv", rows) + write_csv(out / "temporal_smoothness.csv", smooth_rows) + write_csv(out / "interpolation_failure_cases.csv", [{"rank": n + 1, "trajectory": key, "spearman": rho, "top25_overlap": ov, "max_error": err, "valid_anchors": count} for n, (err, key, rho, ov, _, count) in enumerate(sorted(interpolation_cases, reverse=True)[:20])]) + summary = [] + for k in (4, 8, 16): + part = [r for r in rows if r["K"] == k] + def metric(name): + vals = [r[name] for r in part]; lo, hi = bootstrap(vals) + return {"mean": float(np.mean(vals)), "median": float(np.median(vals)), "bootstrap95": [lo, hi]} + summary.append({"K": k, "spearman": metric("spearman"), "top25_overlap": metric("top25_overlap"), "top10_recall": metric("top10_recall"), "proposal_tv": metric("proposal_tv")}) + summary.append({"K": "Dense", "spearman": {"mean": 1.0, "median": 1.0, "bootstrap95": [1.0, 1.0]}, "top25_overlap": {"mean": 1.0, "median": 1.0, "bootstrap95": [1.0, 1.0]}, "top10_recall": {"mean": 1.0, "median": 1.0, "bootstrap95": [1.0, 1.0]}, "proposal_tv": {"mean": 0.0, "median": 0.0, "bootstrap95": [0.0, 0.0]}}) + (out / "temporal_fidelity_summary.csv").write_text("K,mean_spearman,median_spearman,mean_top25_overlap,mean_top10_recall,mean_proposal_tv\n" + "\n".join(f"{x['K']},{x['spearman']['mean']:.10f},{x['spearman']['median']:.10f},{x['top25_overlap']['mean']:.10f},{x['top10_recall']['mean']:.10f},{x['proposal_tv']['mean']:.10f}" for x in summary) + "\n", encoding="utf-8") + write_csv(out / "scoring_compute_cost.csv", [{"K": k, "actual_scored_anchors": int(round(sum((offsets[j + 1] - offsets[j]) if k >= (offsets[j + 1] - offsets[j]) else min(k, offsets[j + 1] - offsets[j]) for j in range(len(offsets) - 1)))), "wall_clock": "DERIVED_FROM_DENSE_CACHE", "relative_to_K8": (k / 8.0), "relative_to_dense": "DERIVED"} for k in (4, 8, 16)] + [{"K": "Dense", "actual_scored_anchors": int(len(scores_flat)), "wall_clock": "MEASURED_IN_DENSE_METADATA", "relative_to_K8": "DERIVED", "relative_to_dense": 1.0}]) + (out / "optimization_fidelity_summary.csv").write_text("K,M_qK_over_M_qDense,M_qK_over_M_p,status\n4,NA,NA,NOT_COMPUTED_without_dense_gradient_target\n8,NA,NA,NOT_COMPUTED_without_dense_gradient_target\n16,NA,NA,NOT_COMPUTED_without_dense_gradient_target\nDense,1.0,1.0,REFERENCE\n", encoding="utf-8") + (out / "optimization_fidelity_per_trajectory.csv").write_text("trajectory,status\nDENSE_GRADIENT_TARGET_NOT_AVAILABLE,NOT_COMPUTED\n", encoding="utf-8") + print(json.dumps({"temporal_fidelity": summary, "smoothness_trajectories": len(smooth_rows), "optimization_status": "not_computed_without_dense_gradient_target"}, indent=2)) + + +def assemble(args: argparse.Namespace) -> None: + out = args.output.resolve() + parts = [] + for index in range(args.num_shards): + path = out / f"dense_residual_scores_shard_{index:02d}_of_{args.num_shards:02d}.npz" + if not path.exists(): + raise FileNotFoundError(path) + parts.append(np.load(path, allow_pickle=False)) + scores_parts = [] + offsets = [0] + dataset_index = [] + episode_index = [] + trajectory_index = [] + trajectory_id = [] + for part in parts: + raw_offsets = np.asarray(part["offsets"], dtype=np.int64) + raw_scores = np.asarray(part["scores"], dtype=np.float64) + for j in range(len(raw_offsets) - 1): + row = raw_scores[raw_offsets[j] : raw_offsets[j + 1]] + scores_parts.append(row) + offsets.append(offsets[-1] + len(row)) + dataset_index.extend(np.asarray(part["dataset_index"], dtype=np.int64).tolist()) + episode_index.extend(np.asarray(part["episode_index"], dtype=np.int64).tolist()) + trajectory_index.extend(np.asarray(part["trajectory_index"], dtype=np.int64).tolist()) + trajectory_id.extend(np.asarray(part["trajectory_id"]).astype(str).tolist()) + np.savez_compressed( + out / "dense_residual_scores.npz", + scores=np.concatenate(scores_parts), + offsets=np.asarray(offsets, dtype=np.int64), + dataset_index=np.asarray(dataset_index, dtype=np.int64), + episode_index=np.asarray(episode_index, dtype=np.int64), + trajectory_index=np.asarray(trajectory_index, dtype=np.int64), + trajectory_id=np.asarray(trajectory_id), + ) + metadata = { + "schema_version": 1, + "status": "assembled", + "num_shards": args.num_shards, + "trajectory_count": len(dataset_index), + "valid_anchor_count": int(sum(len(x) for x in scores_parts)), + "source_files": [str(out / f"dense_residual_scores_shard_{i:02d}_of_{args.num_shards:02d}.npz") for i in range(args.num_shards)], + "measured": True, + "notes": "Assembled from disjoint trajectory shards; residual scores were computed once per anchor.", + } + (out / "dense_residual_metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") + print(json.dumps(metadata, indent=2)) + + +def main() -> None: + parser = argparse.ArgumentParser() + sub = parser.add_subparsers(dest="command", required=True) + p = sub.add_parser("score") + p.add_argument("--config", type=Path, required=True) + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--gate", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--estimator-seeds", nargs="+", type=int, default=[27200, 27201, 27202, 27203]) + p.add_argument("--samples-per-estimator", type=int, default=8) + p.add_argument("--batch-size", type=int, default=2) + p.add_argument("--shard-index", type=int, default=0) + p.add_argument("--num-shards", type=int, default=1) + p.set_defaults(func=score) + p = sub.add_parser("analyze") + p.add_argument("--output", type=Path, required=True) + p.add_argument("--alpha", type=float, default=0.2) + p.add_argument("--epsilon", type=float, default=0.1) + p.add_argument("--normalization-epsilon", type=float, default=1e-8) + p.set_defaults(func=analyze) + p = sub.add_parser("assemble") + p.add_argument("--output", type=Path, required=True) + p.add_argument("--num-shards", type=int, required=True) + p.set_defaults(func=assemble) + args = parser.parse_args(); args.func(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_residual_gradient_alignment.py b/scripts/validate_residual_gradient_alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..f5acb7265800352ab15bd3cb9b7c59cbe92454c6 --- /dev/null +++ b/scripts/validate_residual_gradient_alignment.py @@ -0,0 +1,444 @@ +"""Audit frozen Flow-Matching residual as a gradient-importance proxy.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import random +import time +from pathlib import Path + +import numpy as np +import torch +from omegaconf import OmegaConf +from scipy.stats import kendalltau, spearmanr + +from starVLA.dataloader import build_dataloader +from starVLA.model.framework.base_framework import build_framework +from starVLA.model.framework.share_tools import apply_config_compat +from starVLA.scoring.utility_audit import ( + alignment_audit, + leave_one_out_robustness, +) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _set_random_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _load_examples(cfg, sample_count: int) -> list[dict]: + dataloader = build_dataloader(cfg, dataset_py=cfg.datasets.vla_data.dataset_py) + examples = [] + for batch in dataloader: + examples.extend(batch) + if len(examples) >= sample_count: + break + examples = examples[:sample_count] + if len(examples) != sample_count: + raise RuntimeError( + f"dataloader returned {len(examples)} examples, expected {sample_count}" + ) + if any(example.get("sampling_trace") is None for example in examples): + raise RuntimeError("every residual audit example must have a sampling trace") + return examples + + +def _measure_frozen_residual( + model, + examples: list[dict], + *, + repeat_count: int, + seed_base: int, +) -> tuple[np.ndarray, list[dict], dict]: + model.requires_grad_(False) + model.eval() + repeat_seeds = [seed_base + repeat for repeat in range(repeat_count)] + residual_matrix = np.empty( + (len(examples), repeat_count), + dtype=np.float64, + ) + + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + with torch.no_grad(): + for anchor_index, example in enumerate(examples): + for repeat_index, seed in enumerate(repeat_seeds): + _set_random_seed(seed) + residual = model([example])["action_loss"] + if residual.ndim != 0 or not torch.isfinite(residual): + raise FloatingPointError("frozen residual must be a finite scalar") + residual_matrix[anchor_index, repeat_index] = float( + residual.detach().cpu() + ) + print( + f"score anchor={anchor_index + 1}/{len(examples)} " + f"residual={residual_matrix[anchor_index].mean():.6f}", + flush=True, + ) + torch.cuda.synchronize() + elapsed_seconds = time.perf_counter() - started + + half = repeat_count // 2 + first_half = residual_matrix[:, :half].mean(axis=1) + second_half = residual_matrix[:, half:].mean(axis=1) + half_spearman = spearmanr(first_half, second_half) + half_kendall = kendalltau(first_half, second_half) + top_count = max(1, len(examples) // 4) + first_top = set(np.argsort(first_half)[-top_count:]) + second_top = set(np.argsort(second_half)[-top_count:]) + anchor_results = [] + for anchor_index, example in enumerate(examples): + values = residual_matrix[anchor_index] + anchor_results.append( + { + "anchor_position": anchor_index, + "residuals": values.tolist(), + "mean_residual": float(values.mean()), + "residual_cv": float( + values.std(ddof=1) / values.mean() + if values.mean() > 0.0 + else 0.0 + ), + "sampling_trace": example["sampling_trace"], + } + ) + runtime = { + "elapsed_seconds": elapsed_seconds, + "seconds_per_anchor": elapsed_seconds / len(examples), + "peak_allocated_bytes": torch.cuda.max_memory_allocated(), + "peak_reserved_bytes": torch.cuda.max_memory_reserved(), + "repeat_count": repeat_count, + "repeat_seeds": repeat_seeds, + "model_mode": "eval", + "backward": False, + } + repeatability = { + "first_half_seeds": repeat_seeds[:half], + "second_half_seeds": repeat_seeds[half:], + "spearman": float(half_spearman.statistic), + "spearman_pvalue_asymptotic": float(half_spearman.pvalue), + "kendall": float(half_kendall.statistic), + "kendall_pvalue_asymptotic": float(half_kendall.pvalue), + "mean_absolute_difference": float( + np.mean(np.abs(first_half - second_half)) + ), + "max_absolute_difference": float( + np.max(np.abs(first_half - second_half)) + ), + "top25_overlap": len(first_top & second_top) / top_count, + "mean_anchor_cv": float( + np.mean([result["residual_cv"] for result in anchor_results]) + ), + "median_anchor_cv": float( + np.median([result["residual_cv"] for result in anchor_results]) + ), + } + return residual_matrix, anchor_results, { + "runtime": runtime, + "repeatability": repeatability, + } + + +def _gradient_squared_norm(parameters: list[torch.nn.Parameter]) -> float: + total = 0.0 + saw_gradient = False + for parameter in parameters: + if parameter.grad is None: + continue + saw_gradient = True + gradient = parameter.grad.detach() + if not torch.isfinite(gradient).all(): + raise FloatingPointError("encountered a non-finite action-model gradient") + total += gradient.float().square().sum().item() + if not saw_gradient: + raise RuntimeError("no action-model parameter received a gradient") + return total + + +def _measure_gradient_rms( + model, + examples: list[dict], + *, + repeat_count: int, + seed_base: int, +) -> tuple[list[dict], dict]: + model.requires_grad_(False) + model.action_model.requires_grad_(True) + model.train() + trainable_parameters = [ + parameter + for parameter in model.action_model.parameters() + if parameter.requires_grad + ] + if not trainable_parameters: + raise RuntimeError("action model has no trainable parameters") + if any( + parameter.requires_grad + for parameter in model.qwen_vl_interface.parameters() + ): + raise RuntimeError("the VLM backbone must remain frozen") + + repeat_seeds = [seed_base + repeat for repeat in range(repeat_count)] + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + anchor_results = [] + for anchor_index, example in enumerate(examples): + losses = [] + gradient_norms = [] + for seed in repeat_seeds: + _set_random_seed(seed) + model.zero_grad(set_to_none=True) + loss = model([example])["action_loss"] + if loss.ndim != 0 or not torch.isfinite(loss): + raise FloatingPointError("training loss must be a finite scalar") + loss.backward() + losses.append(float(loss.detach().cpu())) + gradient_norms.append( + math.sqrt(_gradient_squared_norm(trainable_parameters)) + ) + model.zero_grad(set_to_none=True) + gradient_norm_array = np.asarray(gradient_norms, dtype=np.float64) + anchor_results.append( + { + "anchor_position": anchor_index, + "losses": losses, + "gradient_norms": gradient_norms, + "mean_loss": float(np.mean(losses)), + "gradient_rms": float( + math.sqrt(np.mean(np.square(gradient_norm_array))) + ), + "gradient_norm_cv": float( + gradient_norm_array.std(ddof=1) / gradient_norm_array.mean() + if repeat_count > 1 and gradient_norm_array.mean() > 0.0 + else 0.0 + ), + } + ) + print( + f"gradient anchor={anchor_index + 1}/{len(examples)} " + f"loss={anchor_results[-1]['mean_loss']:.6f} " + f"gradient_rms={anchor_results[-1]['gradient_rms']:.6f}", + flush=True, + ) + torch.cuda.synchronize() + elapsed_seconds = time.perf_counter() - started + return anchor_results, { + "elapsed_seconds": elapsed_seconds, + "seconds_per_anchor": elapsed_seconds / len(examples), + "peak_allocated_bytes": torch.cuda.max_memory_allocated(), + "peak_reserved_bytes": torch.cuda.max_memory_reserved(), + "repeat_count": repeat_count, + "repeat_seeds": repeat_seeds, + "model_mode": "train", + "frozen_backbone_mode": "train", + "trainable_action_parameters": sum( + parameter.numel() for parameter in trainable_parameters + ), + } + + +def _gate_decision( + *, + sample_count: int, + repeatability: dict, + alignment: dict, + robustness: dict, +) -> dict: + checks = { + "sample_count_at_least_64": sample_count >= 64, + "score_repeatability_spearman_at_least_0_90": ( + repeatability["spearman"] >= 0.90 + ), + "score_gradient_spearman_at_least_0_50": ( + alignment["spearman_residual_vs_gradient_rms"] >= 0.50 + ), + "score_gradient_permutation_p_at_most_0_01": ( + alignment["random_permutation_control"][ + "one_sided_pvalue_rho_at_least_observed" + ] + <= 0.01 + ), + "top25_overlap_at_least_0_50": alignment["top25_overlap"] >= 0.50, + "second_moment_ratio_at_most_0_98": ( + alignment["proposal"]["information_to_baseline_ratio"] <= 0.98 + ), + "second_moment_permutation_p_at_most_0_01": ( + alignment["random_permutation_control"][ + "one_sided_pvalue_second_moment_at_most_observed" + ] + <= 0.01 + ), + "leave_one_out_spearman_all_positive": robustness[ + "all_spearman_positive" + ], + "leave_one_out_second_moment_all_below_one": robustness[ + "all_second_moment_ratios_below_one" + ], + } + return { + "checks": checks, + "pass": all(checks.values()), + "eligible_for_500_step_ablation": all(checks.values()), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--max-samples", type=int, default=16) + parser.add_argument("--score-repeats", type=int, default=4) + parser.add_argument("--score-seed-base", type=int, default=7300) + parser.add_argument("--gradient-repeats", type=int, default=4) + parser.add_argument("--gradient-seed-base", type=int, default=8300) + parser.add_argument("--alpha", type=float, default=0.2) + parser.add_argument("--percentile-epsilon", type=float, default=0.1) + parser.add_argument("--permutations", type=int, default=10_000) + parser.add_argument("--permutation-seed", type=int, default=9100) + args = parser.parse_args() + + if args.max_samples < 4: + raise ValueError("max-samples must be at least 4") + if args.score_repeats < 4 or args.score_repeats % 2: + raise ValueError("score-repeats must be an even integer of at least 4") + if args.gradient_repeats < 2: + raise ValueError("gradient-repeats must be at least 2") + if args.permutations <= 0: + raise ValueError("permutations must be positive") + if not 0.0 <= args.alpha < 1.0: + raise ValueError("alpha must satisfy 0 <= alpha < 1") + if args.percentile_epsilon <= 0.0: + raise ValueError("percentile-epsilon must be positive") + score_seeds = set( + range(args.score_seed_base, args.score_seed_base + args.score_repeats) + ) + gradient_seeds = set( + range( + args.gradient_seed_base, + args.gradient_seed_base + args.gradient_repeats, + ) + ) + if score_seeds & gradient_seeds: + raise ValueError("score seeds and gradient seeds must be disjoint") + + cfg = apply_config_compat(OmegaConf.load(args.config)) + cfg.output_dir = str(args.output.parent / "residual_gate_runtime") + examples = _load_examples(cfg, args.max_samples) + model = build_framework(cfg) + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True) + model.load_state_dict(checkpoint, strict=True) + model.to("cuda") + + residual_matrix, residual_results, residual_audit = ( + _measure_frozen_residual( + model, + examples, + repeat_count=args.score_repeats, + seed_base=args.score_seed_base, + ) + ) + gradient_results, gradient_runtime = _measure_gradient_rms( + model, + examples, + repeat_count=args.gradient_repeats, + seed_base=args.gradient_seed_base, + ) + scores = residual_matrix.mean(axis=1) + gradient_rms = np.asarray( + [result["gradient_rms"] for result in gradient_results], + dtype=np.float64, + ) + mean_losses = np.asarray( + [result["mean_loss"] for result in gradient_results], + dtype=np.float64, + ) + alignment = alignment_audit( + scores, + gradient_rms, + mean_losses, + alpha=args.alpha, + percentile_epsilon=args.percentile_epsilon, + permutation_count=args.permutations, + permutation_seed=args.permutation_seed, + score_label="residual", + ) + robustness = leave_one_out_robustness( + scores, + gradient_rms, + alpha=args.alpha, + percentile_epsilon=args.percentile_epsilon, + ) + gate = _gate_decision( + sample_count=args.max_samples, + repeatability=residual_audit["repeatability"], + alignment=alignment, + robustness=robustness, + ) + for residual_result, gradient_result in zip( + residual_results, + gradient_results, + ): + residual_result.update( + { + "gradient_losses": gradient_result["losses"], + "gradient_norms": gradient_result["gradient_norms"], + "mean_gradient_loss": gradient_result["mean_loss"], + "gradient_rms": gradient_result["gradient_rms"], + "gradient_norm_cv": gradient_result["gradient_norm_cv"], + } + ) + + payload = { + "schema_version": 1, + "status": ( + "formal_64_anchor_gate" + if args.max_samples >= 64 + else "smoke_measurement_not_a_gate_certificate" + ), + "config": str(args.config.resolve()), + "checkpoint": str(args.checkpoint.resolve()), + "checkpoint_sha256": _sha256_file(args.checkpoint), + "score_definition": { + "name": "frozen_flow_matching_residual", + "model_mode": "eval", + "aggregation": "mean over independent flow time/noise seeds", + "forbidden_labels": ["likelihood", "entropy", "surprisal"], + }, + "seed_separation": { + "score_seeds": sorted(score_seeds), + "gradient_seeds": sorted(gradient_seeds), + "disjoint": not bool(score_seeds & gradient_seeds), + }, + "sample_count": args.max_samples, + "residual": residual_audit, + "gradient_runtime": gradient_runtime, + "alignment": alignment, + "leave_one_out_robustness": robustness, + "gate": gate, + "anchors": residual_results, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + temporary = args.output.with_suffix(args.output.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") + temporary.replace(args.output) + print(json.dumps(payload, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_stratified_gradient_second_moment.py b/scripts/validate_stratified_gradient_second_moment.py new file mode 100644 index 0000000000000000000000000000000000000000..9165d29b8a958e19a9f12d33b7701c4d12672927 --- /dev/null +++ b/scripts/validate_stratified_gradient_second_moment.py @@ -0,0 +1,610 @@ +"""Estimate Theorem-3 gradient second moments with disjoint stratified draws.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import random +import time +from pathlib import Path + +import numpy as np +import torch +from omegaconf import OmegaConf +from scipy.stats import kendalltau, spearmanr + +from starVLA.dataloader import build_dataloader +from starVLA.model.framework.base_framework import build_framework +from starVLA.model.framework.share_tools import apply_config_compat +from starVLA.scoring.flow_residual import ( + randomized_antithetic_stratified_flow_samples, +) +from starVLA.scoring.partial_gradient import ( + activate_partial_gradient_group, + define_partial_gradient_groups, + finite_gradient_l2_norm, +) +from starVLA.scoring.utility_audit import ( + alignment_audit, + leave_one_out_robustness, +) + + +TRACE_IDENTITY_FIELDS = ( + "manifest_sha256", + "epoch", + "global_draw_id", + "dataset_index", + "dataset_name", + "trajectory_index", + "trajectory_id", + "anchor_index", + "base_index", +) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _set_random_seed(seed: int) -> None: + bounded_seed = int(seed) % (2**32) + random.seed(bounded_seed) + np.random.seed(bounded_seed) + torch.manual_seed(bounded_seed) + torch.cuda.manual_seed_all(bounded_seed) + + +def _trace_identity(trace: dict) -> dict: + return {field: trace.get(field) for field in TRACE_IDENTITY_FIELDS} + + +def _load_reference_artifact( + path: Path, + max_samples: int, +) -> tuple[dict, list[dict]]: + payload = json.loads(path.read_text(encoding="utf-8")) + anchors = payload.get("anchors") + if not isinstance(anchors, list) or len(anchors) < max_samples: + raise ValueError("reference artifact does not contain enough anchors") + anchors = anchors[:max_samples] + return { + "path": str(path.resolve()), + "sha256": _sha256_file(path), + "checkpoint_sha256": payload.get("checkpoint_sha256"), + "old_gradient_seeds": payload.get("seed_separation", {}).get( + "gradient_seeds" + ), + }, anchors + + +def _load_aligned_examples(cfg, anchors: list[dict]) -> list[dict]: + dataloader = build_dataloader(cfg, dataset_py=cfg.datasets.vla_data.dataset_py) + examples = [] + for batch in dataloader: + examples.extend(batch) + if len(examples) >= len(anchors): + break + examples = examples[: len(anchors)] + if len(examples) != len(anchors): + raise RuntimeError("dataloader ended before audit anchors were recovered") + for index, (example, anchor) in enumerate(zip(examples, anchors)): + observed = _trace_identity(example.get("sampling_trace") or {}) + expected = _trace_identity(anchor["sampling_trace"]) + if observed != expected: + raise RuntimeError( + f"sample trace mismatch at position {index}: " + f"observed={observed}, expected={expected}" + ) + return examples + + +def _prepare_conditioning(model, example: dict): + images = [example["image"]] + instructions = [example["lang"]] + raw_actions = [example["action"]] + raw_states = [example["state"]] if "state" in example else None + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + vlm_inputs = model.qwen_vl_interface.build_qwenvl_inputs( + images=images, + instructions=instructions, + ) + vlm_outputs = model.qwen_vl_interface( + **vlm_inputs, + output_attentions=False, + output_hidden_states=True, + return_dict=True, + ) + hidden_states = vlm_outputs.hidden_states[-1].detach() + actions = torch.as_tensor( + np.asarray(raw_actions), + device=hidden_states.device, + dtype=hidden_states.dtype, + )[:, -model.action_horizon :, :] + states = ( + torch.as_tensor( + np.asarray(raw_states), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + if raw_states is not None + else None + ) + return hidden_states, actions, states + + +def _flow_loss( + model, + conditioning, + noise: torch.Tensor, + time_value: torch.Tensor, +) -> torch.Tensor: + hidden_states, actions, states = conditioning + time = time_value.reshape(1) + expanded_time = time[:, None, None] + noisy_trajectory = (1.0 - expanded_time) * noise + expanded_time * actions + target_velocity = actions - noise + with torch.autocast("cuda", dtype=torch.float32): + prediction = model.action_model.predict_velocity( + hidden_states, + noisy_trajectory, + time, + state=states, + ) + return (prediction - target_velocity).square().mean() + + +def _measure_gradient_second_moment( + model, + examples: list[dict], + *, + estimator_seeds: list[int], + samples_per_estimator: int, + parameter_mode: str, + frozen_backbone_eval: bool, +) -> tuple[np.ndarray, np.ndarray, dict]: + action_config = model.action_model.config + groups = define_partial_gradient_groups(model.action_model) + estimator_values = np.empty( + (len(examples), len(estimator_seeds)), + dtype=np.float64, + ) + draw_norms = np.empty( + (len(examples), len(estimator_seeds), samples_per_estimator), + dtype=np.float64, + ) + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + parameters = [] + + for anchor_index, example in enumerate(examples): + for estimator_index, estimator_seed in enumerate(estimator_seeds): + model.requires_grad_(False) + model.train() + if frozen_backbone_eval: + model.qwen_vl_interface.eval() + conditioning_seed = estimator_seed * 1_000_003 + anchor_index + _set_random_seed(conditioning_seed) + conditioning = _prepare_conditioning(model, example) + actions = conditioning[1] + noises, times = randomized_antithetic_stratified_flow_samples( + actions, + sample_count=samples_per_estimator, + seed=estimator_seed * 1_000_033 + anchor_index, + beta_alpha=float(action_config.noise_beta_alpha), + beta_beta=float(action_config.noise_beta_beta), + noise_s=float(action_config.noise_s), + ) + if parameter_mode == "full": + model.action_model.requires_grad_(True) + parameters = list(model.action_model.parameters()) + elif parameter_mode == "decoder": + parameters = activate_partial_gradient_group( + model.action_model, + groups["decoder"], + ) + else: + raise ValueError(f"unknown parameter mode: {parameter_mode}") + + for draw_index in range(samples_per_estimator): + # Common dropout within an antithetic pair preserves each + # marginal draw while strengthening cancellation. + dropout_seed = ( + estimator_seed * 1_000_037 + + anchor_index * 101 + + draw_index // 2 + ) + _set_random_seed(dropout_seed) + model.zero_grad(set_to_none=True) + loss = _flow_loss( + model, + conditioning, + noises[draw_index], + times[draw_index], + ) + if loss.ndim != 0 or not torch.isfinite(loss): + raise FloatingPointError("flow loss must be a finite scalar") + loss.backward() + draw_norms[ + anchor_index, + estimator_index, + draw_index, + ] = finite_gradient_l2_norm(parameters) + estimator_values[anchor_index, estimator_index] = np.sqrt( + np.mean( + np.square( + draw_norms[anchor_index, estimator_index] + ) + ) + ) + print( + f"mode={parameter_mode} anchor={anchor_index + 1}/{len(examples)} " + f"h={np.sqrt(np.mean(np.square(draw_norms[anchor_index]))):.6f}", + flush=True, + ) + model.zero_grad(set_to_none=True) + torch.cuda.synchronize() + elapsed_seconds = time.perf_counter() - started + return estimator_values, draw_norms, { + "elapsed_seconds": elapsed_seconds, + "seconds_per_anchor": elapsed_seconds / len(examples), + "seconds_per_draw": ( + elapsed_seconds + / len(examples) + / len(estimator_seeds) + / samples_per_estimator + ), + "peak_allocated_bytes": torch.cuda.max_memory_allocated(), + "peak_reserved_bytes": torch.cuda.max_memory_reserved(), + "parameter_mode": parameter_mode, + "frozen_backbone_eval": frozen_backbone_eval, + "parameter_count": sum( + parameter.numel() for parameter in parameters + ), + } + + +def _repeatability(estimator_values: np.ndarray) -> dict: + if estimator_values.ndim != 2 or estimator_values.shape[1] < 4: + raise ValueError("at least four estimator columns are required") + if estimator_values.shape[1] % 2: + raise ValueError("estimator count must be even") + half = estimator_values.shape[1] // 2 + first = np.sqrt(np.mean(np.square(estimator_values[:, :half]), axis=1)) + second = np.sqrt(np.mean(np.square(estimator_values[:, half:]), axis=1)) + correlation = spearmanr(first, second) + concordance = kendalltau(first, second) + return { + "spearman": float(correlation.statistic), + "spearman_pvalue_asymptotic": float(correlation.pvalue), + "kendall": float(concordance.statistic), + "kendall_pvalue_asymptotic": float(concordance.pvalue), + "mean_anchor_cv": float( + np.mean( + estimator_values.std(axis=1, ddof=1) + / estimator_values.mean(axis=1) + ) + ), + } + + +def _gate( + sample_count: int, + proxy_repeatability: dict, + target_repeatability: dict, + alignment: dict, + robustness: dict, +) -> dict: + checks = { + "sample_count_at_least_64": sample_count >= 64, + "target_repeatability_at_least_0_80": ( + target_repeatability["spearman"] >= 0.80 + ), + "proxy_repeatability_at_least_0_90": ( + proxy_repeatability["spearman"] >= 0.90 + ), + "proxy_target_spearman_at_least_0_70": ( + alignment["spearman_partial_gradient_vs_gradient_rms"] >= 0.70 + ), + "spearman_permutation_p_at_most_0_01": ( + alignment["random_permutation_control"][ + "one_sided_pvalue_rho_at_least_observed" + ] + <= 0.01 + ), + "top25_overlap_at_least_0_625": alignment["top25_overlap"] >= 0.625, + "second_moment_ratio_at_most_0_98": ( + alignment["proposal"]["information_to_baseline_ratio"] <= 0.98 + ), + "second_moment_permutation_p_at_most_0_01": ( + alignment["random_permutation_control"][ + "one_sided_pvalue_second_moment_at_most_observed" + ] + <= 0.01 + ), + "leave_one_out_spearman_all_positive": robustness[ + "all_spearman_positive" + ], + "leave_one_out_second_moment_all_below_one": robustness[ + "all_second_moment_ratios_below_one" + ], + } + return { + "checks": checks, + "pass": all(checks.values()), + "eligible_for_500_step_ablation": all(checks.values()), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--reference-artifact", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--max-samples", type=int, default=8) + parser.add_argument("--samples-per-estimator", type=int, default=8) + parser.add_argument( + "--proxy-seeds", + type=int, + nargs="+", + default=[12300, 12301, 12302, 12303], + ) + parser.add_argument( + "--target-seeds", + type=int, + nargs="+", + default=[13300, 13301, 13302, 13303], + ) + parser.add_argument("--alpha", type=float, default=0.2) + parser.add_argument("--percentile-epsilon", type=float, default=0.1) + parser.add_argument("--permutations", type=int, default=10_000) + parser.add_argument("--permutation-seed", type=int, default=14100) + parser.add_argument("--frozen-backbone-eval", action="store_true") + args = parser.parse_args() + if args.max_samples < 4: + raise ValueError("max-samples must be at least four") + for name, seeds in ( + ("proxy-seeds", args.proxy_seeds), + ("target-seeds", args.target_seeds), + ): + if len(seeds) < 4 or len(seeds) % 2 or len(set(seeds)) != len(seeds): + raise ValueError(f"{name} must contain an even count of unique seeds") + if set(args.proxy_seeds) & set(args.target_seeds): + raise ValueError("proxy and target seeds must be disjoint") + if args.samples_per_estimator < 2 or args.samples_per_estimator % 2: + raise ValueError("samples-per-estimator must be a positive even integer") + + reference, reference_anchors = _load_reference_artifact( + args.reference_artifact, + args.max_samples, + ) + checkpoint_sha256 = _sha256_file(args.checkpoint) + if ( + reference["checkpoint_sha256"] is not None + and reference["checkpoint_sha256"] != checkpoint_sha256 + ): + raise RuntimeError("checkpoint hash differs from reference artifact") + all_old_seeds = set(reference["old_gradient_seeds"] or []) + if all_old_seeds & (set(args.proxy_seeds) | set(args.target_seeds)): + raise ValueError("new seeds overlap the old gradient audit") + + cfg = apply_config_compat(OmegaConf.load(args.config)) + cfg.output_dir = str(args.output.parent / "stratified_gradient_runtime") + examples = _load_aligned_examples(cfg, reference_anchors) + model = build_framework(cfg) + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True) + model.load_state_dict(checkpoint, strict=True) + model.to("cuda") + + proxy_estimators, proxy_draws, proxy_runtime = ( + _measure_gradient_second_moment( + model, + examples, + estimator_seeds=args.proxy_seeds, + samples_per_estimator=args.samples_per_estimator, + parameter_mode="decoder", + frozen_backbone_eval=args.frozen_backbone_eval, + ) + ) + target_estimators, target_draws, target_runtime = ( + _measure_gradient_second_moment( + model, + examples, + estimator_seeds=args.target_seeds, + samples_per_estimator=args.samples_per_estimator, + parameter_mode="full", + frozen_backbone_eval=args.frozen_backbone_eval, + ) + ) + proxy_scores = np.sqrt(np.mean(np.square(proxy_estimators), axis=1)) + target_gradient_rms = np.sqrt( + np.mean(np.square(target_estimators), axis=1) + ) + old_gradient_rms = np.asarray( + [anchor["gradient_rms"] for anchor in reference_anchors], + dtype=np.float64, + ) + mean_losses = np.asarray( + [anchor["mean_gradient_loss"] for anchor in reference_anchors], + dtype=np.float64, + ) + proxy_repeatability = _repeatability(proxy_estimators) + target_repeatability = _repeatability(target_estimators) + alignment = alignment_audit( + proxy_scores, + target_gradient_rms, + mean_losses, + alpha=args.alpha, + percentile_epsilon=args.percentile_epsilon, + permutation_count=args.permutations, + permutation_seed=args.permutation_seed, + score_label="partial_gradient", + ) + robustness = leave_one_out_robustness( + proxy_scores, + target_gradient_rms, + alpha=args.alpha, + percentile_epsilon=args.percentile_epsilon, + ) + gate = _gate( + len(examples), + proxy_repeatability, + target_repeatability, + alignment, + robustness, + ) + old_target_correlation = spearmanr( + target_gradient_rms, + old_gradient_rms, + ) + payload = { + "schema_version": 1, + "status": ( + "formal_64_anchor_stratified_gradient_gate" + if len(examples) >= 64 + else "stratified_gradient_smoke_not_a_gate_certificate" + ), + "config": str(args.config.resolve()), + "checkpoint": str(args.checkpoint.resolve()), + "checkpoint_sha256": checkpoint_sha256, + "reference_artifact": reference, + "sample_count": len(examples), + "estimator": { + "definition": "sqrt(mean over xi of squared per-draw gradient L2 norm)", + "randomized_beta_stratification": True, + "antithetic_gaussian_noise": True, + "common_dropout_within_antithetic_pair": True, + "samples_per_estimator": args.samples_per_estimator, + "proxy_seeds": args.proxy_seeds, + "target_seeds": args.target_seeds, + "seeds_disjoint": True, + "frozen_backbone_eval": args.frozen_backbone_eval, + }, + "proxy_runtime": proxy_runtime, + "target_runtime": target_runtime, + "proxy_repeatability": proxy_repeatability, + "target_repeatability": target_repeatability, + "new_vs_old_target_spearman": float(old_target_correlation.statistic), + "new_vs_old_target_spearman_pvalue": float(old_target_correlation.pvalue), + "alignment": alignment, + "leave_one_out_robustness": robustness, + "gate": gate, + "eligible_for_500_step_ablation": gate[ + "eligible_for_500_step_ablation" + ], + "anchors": [ + { + "anchor_position": index, + "sampling_trace": reference_anchors[index]["sampling_trace"], + "proxy_estimators": proxy_estimators[index].tolist(), + "proxy_draw_norms": proxy_draws[index].tolist(), + "proxy_gradient_rms": float(proxy_scores[index]), + "target_estimators": target_estimators[index].tolist(), + "target_draw_norms": target_draws[index].tolist(), + "target_gradient_rms": float(target_gradient_rms[index]), + "old_four_draw_gradient_rms": float(old_gradient_rms[index]), + } + for index in range(len(examples)) + ], + } + args.output.parent.mkdir(parents=True, exist_ok=True) + temporary = args.output.with_suffix(args.output.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") + temporary.replace(args.output) + print(json.dumps(payload, indent=2)) + + +if __name__ == "__main__": + main() + + + +def _measure_gradient_second_moment_draw_shard( + model, + examples: list[dict], + *, + estimator_seeds: list[int], + samples_per_estimator: int, + draw_indices: list[int], + parameter_mode: str, + frozen_backbone_eval: bool, +) -> tuple[np.ndarray, np.ndarray, dict]: + action_config = model.action_model.config + groups = define_partial_gradient_groups(model.action_model) + values = np.empty((len(examples), len(estimator_seeds)), dtype=np.float64) + draw_norms = np.empty( + (len(examples), len(estimator_seeds), len(draw_indices)), + dtype=np.float64, + ) + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + parameters = [] + for anchor_index, example in enumerate(examples): + for estimator_index, estimator_seed in enumerate(estimator_seeds): + model.requires_grad_(False) + model.train() + if frozen_backbone_eval: + model.qwen_vl_interface.eval() + conditioning_seed = estimator_seed * 1_000_003 + anchor_index + _set_random_seed(conditioning_seed) + conditioning = _prepare_conditioning(model, example) + actions = conditioning[1] + noises, times = randomized_antithetic_stratified_flow_samples( + actions, + sample_count=samples_per_estimator, + seed=estimator_seed * 1_000_033 + anchor_index, + beta_alpha=float(action_config.noise_beta_alpha), + beta_beta=float(action_config.noise_beta_beta), + noise_s=float(action_config.noise_s), + ) + if parameter_mode != "full": + raise ValueError(f"unknown parameter mode: {parameter_mode}") + model.action_model.requires_grad_(True) + parameters = list(model.action_model.parameters()) + for output_index, draw_index in enumerate(draw_indices): + dropout_seed = ( + estimator_seed * 1_000_037 + + anchor_index * 101 + + draw_index // 2 + ) + _set_random_seed(dropout_seed) + model.zero_grad(set_to_none=True) + loss = _flow_loss( + model, + conditioning, + noises[draw_index], + times[draw_index], + ) + if loss.ndim != 0 or not torch.isfinite(loss): + raise FloatingPointError("flow loss must be a finite scalar") + loss.backward() + draw_norms[anchor_index, estimator_index, output_index] = finite_gradient_l2_norm(parameters) + values[anchor_index, estimator_index] = np.sqrt( + np.mean(np.square(draw_norms[anchor_index, estimator_index])) + ) + print( + f"mode={parameter_mode} anchor={anchor_index + 1}/{len(examples)} " + f"h={np.sqrt(np.mean(np.square(draw_norms[anchor_index]))):.6f}", + flush=True, + ) + model.zero_grad(set_to_none=True) + torch.cuda.synchronize() + elapsed_seconds = time.perf_counter() - started + return values, draw_norms, { + "elapsed_seconds": elapsed_seconds, + "seconds_per_anchor": elapsed_seconds / len(examples), + "seconds_per_draw": elapsed_seconds / len(examples) / len(estimator_seeds) / len(draw_indices), + "peak_allocated_bytes": torch.cuda.max_memory_allocated(), + "peak_reserved_bytes": torch.cuda.max_memory_reserved(), + "parameter_mode": parameter_mode, + "frozen_backbone_eval": frozen_backbone_eval, + "parameter_count": sum(parameter.numel() for parameter in parameters), + } diff --git a/scripts/validate_stratified_residual_alignment.py b/scripts/validate_stratified_residual_alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..0abeb812cc70dd5a8d5d285b0cf92009ccbb7e52 --- /dev/null +++ b/scripts/validate_stratified_residual_alignment.py @@ -0,0 +1,346 @@ +"""Audit a low-variance frozen residual against an existing gradient artifact.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import numpy as np +import torch +from omegaconf import OmegaConf +from scipy.stats import kendalltau, spearmanr + +from starVLA.dataloader import build_dataloader +from starVLA.model.framework.base_framework import build_framework +from starVLA.model.framework.share_tools import apply_config_compat +from starVLA.scoring.starvla_likelihood import ( + score_starvla_flow_matching_residual, +) +from starVLA.scoring.utility_audit import ( + alignment_audit, + leave_one_out_robustness, +) + + +TRACE_IDENTITY_FIELDS = ( + "manifest_sha256", + "epoch", + "global_draw_id", + "dataset_index", + "dataset_name", + "trajectory_index", + "trajectory_id", + "anchor_index", + "base_index", +) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _trace_identity(trace: dict) -> dict: + return {field: trace.get(field) for field in TRACE_IDENTITY_FIELDS} + + +def _load_gradient_artifact(path: Path) -> tuple[dict, list[dict]]: + payload = json.loads(path.read_text(encoding="utf-8")) + anchors = payload.get("anchors") + if not isinstance(anchors, list) or len(anchors) < 4: + raise ValueError("gradient artifact must contain at least four anchors") + required = ("gradient_rms", "mean_gradient_loss", "sampling_trace") + for index, anchor in enumerate(anchors): + if any(field not in anchor for field in required): + raise ValueError( + f"gradient artifact anchor {index} is missing required fields" + ) + if not ( + np.isfinite(anchor["gradient_rms"]) + and anchor["gradient_rms"] > 0.0 + and np.isfinite(anchor["mean_gradient_loss"]) + ): + raise ValueError("gradient artifact contains invalid measurements") + metadata = { + "path": str(path.resolve()), + "sha256": _sha256_file(path), + "checkpoint_sha256": payload.get("checkpoint_sha256"), + "gradient_runtime": payload.get("gradient_runtime"), + "gradient_seeds": payload.get("seed_separation", {}).get( + "gradient_seeds" + ), + } + return metadata, anchors + + +def _load_aligned_examples(cfg, anchors: list[dict]) -> list[dict]: + dataloader = build_dataloader(cfg, dataset_py=cfg.datasets.vla_data.dataset_py) + examples = [] + for batch in dataloader: + examples.extend(batch) + if len(examples) >= len(anchors): + break + examples = examples[: len(anchors)] + if len(examples) != len(anchors): + raise RuntimeError("dataloader ended before all audit anchors were recovered") + for index, (example, anchor) in enumerate(zip(examples, anchors)): + observed = _trace_identity(example.get("sampling_trace") or {}) + expected = _trace_identity(anchor["sampling_trace"]) + if observed != expected: + raise RuntimeError( + f"sample trace mismatch at position {index}: " + f"observed={observed}, expected={expected}" + ) + return examples + + +def _repeatability(estimator_matrix: np.ndarray) -> dict: + estimator_count, sample_count = estimator_matrix.shape + if estimator_count < 4 or estimator_count % 2: + raise ValueError("estimator matrix must have an even number >= 4 of rows") + half = estimator_count // 2 + first_half = estimator_matrix[:half].mean(axis=0) + second_half = estimator_matrix[half:].mean(axis=0) + correlation = spearmanr(first_half, second_half) + rank_concordance = kendalltau(first_half, second_half) + top_count = max(1, sample_count // 4) + first_top = set(np.argsort(first_half)[-top_count:]) + second_top = set(np.argsort(second_half)[-top_count:]) + anchor_means = estimator_matrix.mean(axis=0) + anchor_cv = estimator_matrix.std(axis=0, ddof=1) / anchor_means + return { + "spearman": float(correlation.statistic), + "spearman_pvalue_asymptotic": float(correlation.pvalue), + "kendall": float(rank_concordance.statistic), + "kendall_pvalue_asymptotic": float(rank_concordance.pvalue), + "mean_absolute_difference": float( + np.mean(np.abs(first_half - second_half)) + ), + "max_absolute_difference": float( + np.max(np.abs(first_half - second_half)) + ), + "top25_overlap": len(first_top & second_top) / top_count, + "mean_anchor_cv": float(np.mean(anchor_cv)), + "median_anchor_cv": float(np.median(anchor_cv)), + } + + +def _gate_decision( + *, + sample_count: int, + repeatability: dict, + alignment: dict, + robustness: dict, +) -> dict: + checks = { + "sample_count_at_least_64": sample_count >= 64, + "score_repeatability_spearman_at_least_0_90": ( + repeatability["spearman"] >= 0.90 + ), + "score_gradient_spearman_at_least_0_50": ( + alignment["spearman_residual_vs_gradient_rms"] >= 0.50 + ), + "score_gradient_permutation_p_at_most_0_01": ( + alignment["random_permutation_control"][ + "one_sided_pvalue_rho_at_least_observed" + ] + <= 0.01 + ), + "top25_overlap_at_least_0_50": alignment["top25_overlap"] >= 0.50, + "second_moment_ratio_at_most_0_98": ( + alignment["proposal"]["information_to_baseline_ratio"] <= 0.98 + ), + "second_moment_permutation_p_at_most_0_01": ( + alignment["random_permutation_control"][ + "one_sided_pvalue_second_moment_at_most_observed" + ] + <= 0.01 + ), + "leave_one_out_spearman_all_positive": robustness[ + "all_spearman_positive" + ], + "leave_one_out_second_moment_all_below_one": robustness[ + "all_second_moment_ratios_below_one" + ], + } + return { + "checks": checks, + "pass": all(checks.values()), + "eligible_for_500_step_ablation": all(checks.values()), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--gradient-artifact", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--estimator-seeds", type=int, nargs="+", default=[9300, 9301, 9302, 9303]) + parser.add_argument("--samples-per-estimator", type=int, default=8) + parser.add_argument("--alpha", type=float, default=0.2) + parser.add_argument("--percentile-epsilon", type=float, default=0.1) + parser.add_argument("--permutations", type=int, default=10_000) + parser.add_argument("--permutation-seed", type=int, default=9100) + args = parser.parse_args() + if len(args.estimator_seeds) < 4 or len(args.estimator_seeds) % 2: + raise ValueError("estimator-seeds must contain an even count of at least 4") + if len(set(args.estimator_seeds)) != len(args.estimator_seeds): + raise ValueError("estimator-seeds must be unique") + if args.samples_per_estimator < 2 or args.samples_per_estimator % 2: + raise ValueError("samples-per-estimator must be a positive even integer") + + gradient_metadata, gradient_anchors = _load_gradient_artifact( + args.gradient_artifact + ) + checkpoint_sha256 = _sha256_file(args.checkpoint) + if ( + gradient_metadata["checkpoint_sha256"] is not None + and gradient_metadata["checkpoint_sha256"] != checkpoint_sha256 + ): + raise RuntimeError("checkpoint hash differs from the gradient artifact") + gradient_seeds = set(gradient_metadata["gradient_seeds"] or []) + if gradient_seeds & set(args.estimator_seeds): + raise ValueError("residual estimator seeds overlap gradient seeds") + + cfg = apply_config_compat(OmegaConf.load(args.config)) + cfg.output_dir = str(args.output.parent / "stratified_residual_runtime") + examples = _load_aligned_examples(cfg, gradient_anchors) + model = build_framework(cfg) + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True) + model.load_state_dict(checkpoint, strict=True) + model.requires_grad_(False).eval().to("cuda") + + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + estimator_matrix = np.empty( + (len(args.estimator_seeds), len(examples)), + dtype=np.float64, + ) + score_records = [] + time_values = None + for anchor_index, example in enumerate(examples): + with torch.no_grad(): + result = score_starvla_flow_matching_residual( + model, + [example], + estimator_seeds=args.estimator_seeds, + samples_per_estimator=args.samples_per_estimator, + ) + values = result.estimator_residuals[:, 0].float().cpu().numpy() + estimator_matrix[:, anchor_index] = values + if time_values is None: + time_values = result.time_values.float().cpu().tolist() + score_records.append( + { + "anchor_position": anchor_index, + "estimator_residuals": values.tolist(), + "sample_losses": ( + result.sample_losses[:, :, 0].float().cpu().tolist() + ), + "mean_residual": float(values.mean()), + "sampling_trace": gradient_anchors[anchor_index][ + "sampling_trace" + ], + "gradient_rms": float( + gradient_anchors[anchor_index]["gradient_rms"] + ), + "mean_gradient_loss": float( + gradient_anchors[anchor_index]["mean_gradient_loss"] + ), + } + ) + print( + f"anchor={anchor_index + 1}/{len(examples)} " + f"residual={values.mean():.6f}", + flush=True, + ) + torch.cuda.synchronize() + elapsed_seconds = time.perf_counter() - started + + scores = estimator_matrix.mean(axis=0) + gradient_rms = np.asarray( + [anchor["gradient_rms"] for anchor in gradient_anchors], + dtype=np.float64, + ) + mean_losses = np.asarray( + [anchor["mean_gradient_loss"] for anchor in gradient_anchors], + dtype=np.float64, + ) + repeatability = _repeatability(estimator_matrix) + alignment = alignment_audit( + scores, + gradient_rms, + mean_losses, + alpha=args.alpha, + percentile_epsilon=args.percentile_epsilon, + permutation_count=args.permutations, + permutation_seed=args.permutation_seed, + score_label="residual", + ) + robustness = leave_one_out_robustness( + scores, + gradient_rms, + alpha=args.alpha, + percentile_epsilon=args.percentile_epsilon, + ) + gate = _gate_decision( + sample_count=len(examples), + repeatability=repeatability, + alignment=alignment, + robustness=robustness, + ) + payload = { + "schema_version": 1, + "status": ( + "formal_64_anchor_variance_reduced_gate" + if len(examples) >= 64 + else "variance_reduced_smoke_not_a_gate_certificate" + ), + "config": str(args.config.resolve()), + "checkpoint": str(args.checkpoint.resolve()), + "checkpoint_sha256": checkpoint_sha256, + "gradient_artifact": gradient_metadata, + "sample_count": len(examples), + "score_definition": { + "name": "frozen_flow_matching_residual", + "variance_reduction": [ + "deterministic midpoint stratification of Beta flow time", + "paired antithetic Gaussian noise", + "one shared VLM conditioning encode per anchor", + "batched action-head evaluation", + ], + "model_mode": "eval", + "estimator_seeds": args.estimator_seeds, + "samples_per_estimator": args.samples_per_estimator, + "time_values": time_values, + "forbidden_labels": ["likelihood", "entropy", "surprisal"], + }, + "score_runtime": { + "elapsed_seconds": elapsed_seconds, + "seconds_per_anchor": elapsed_seconds / len(examples), + "peak_allocated_bytes": torch.cuda.max_memory_allocated(), + "peak_reserved_bytes": torch.cuda.max_memory_reserved(), + }, + "repeatability": repeatability, + "alignment": alignment, + "leave_one_out_robustness": robustness, + "gate": gate, + "anchors": score_records, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + temporary = args.output.with_suffix(args.output.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") + temporary.replace(args.output) + print(json.dumps(payload, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/starVLA/__init__.py b/starVLA/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/starVLA/dataloader/__init__.py b/starVLA/dataloader/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7cb73d66cd1ad0e0329d55bd2f2a8c0f663644ad --- /dev/null +++ b/starVLA/dataloader/__init__.py @@ -0,0 +1,109 @@ +import json +import os +from accelerate.logging import get_logger +import numpy as np +from torch.utils.data import DataLoader +import torch.distributed as dist +from pathlib import Path +from starVLA.dataloader.infoflow_sampler import StatelessDrawSampler + +logger = get_logger(__name__) + +def save_dataset_statistics(dataset_statistics, run_dir): + """Saves a `dataset_statistics.json` file.""" + out_path = run_dir / "dataset_statistics.json" + with open(out_path, "w") as f_json: + for _, stats in dataset_statistics.items(): + for k in stats["action"].keys(): + if isinstance(stats["action"][k], np.ndarray): + stats["action"][k] = stats["action"][k].tolist() + if "proprio" in stats: + for k in stats["proprio"].keys(): + if isinstance(stats["proprio"][k], np.ndarray): + stats["proprio"][k] = stats["proprio"][k].tolist() + if "num_trajectories" in stats: + if isinstance(stats["num_trajectories"], np.ndarray): + stats["num_trajectories"] = stats["num_trajectories"].item() + if "num_transitions" in stats: + if isinstance(stats["num_transitions"], np.ndarray): + stats["num_transitions"] = stats["num_transitions"].item() + json.dump(dataset_statistics, f_json, indent=2) + logger.info(f"Saved dataset statistics file at path {out_path}") + + + +def build_dataloader(cfg, dataset_py="lerobot_datasets_oxe"): + + if dataset_py == "lerobot_datasets": + from starVLA.dataloader.lerobot_datasets import get_vla_dataset, collate_fn + vla_dataset_cfg = cfg.datasets.vla_data + + vla_dataset = get_vla_dataset( + data_cfg=vla_dataset_cfg, + balance_dataset_weights=vla_dataset_cfg.get("balance_dataset_weights", False), + balance_trajectory_weights=vla_dataset_cfg.get("balance_trajectory_weights", False), + ) + + traceable_sampling = bool(vla_dataset_cfg.get("traceable_sampling", False)) + per_device_batch_size = int(cfg.datasets.vla_data.per_device_batch_size) + world_size = ( + dist.get_world_size() + if dist.is_available() and dist.is_initialized() + else int(vla_dataset_cfg.get("expected_world_size", 1)) + ) + sampler = None + if traceable_sampling: + expected_world_size = int(vla_dataset_cfg.expected_world_size) + if world_size != expected_world_size: + raise RuntimeError( + f"Traceable sampling requires world size {expected_world_size}, " + f"but distributed runtime reports {world_size}" + ) + gradient_accumulation_steps = int( + cfg.trainer.gradient_accumulation_steps + ) + global_update_batch = ( + per_device_batch_size + * world_size + * gradient_accumulation_steps + ) + usable_draws = len(vla_dataset) - ( + len(vla_dataset) % global_update_batch + ) + if usable_draws <= 0: + raise ValueError( + f"Dataset length {len(vla_dataset)} is smaller than one global " + f"optimizer update ({global_update_batch} draws)" + ) + sampler = StatelessDrawSampler(usable_draws) + num_workers = int(vla_dataset_cfg.get("num_workers", 16)) + dataloader_kwargs = dict( + dataset=vla_dataset, + batch_size=per_device_batch_size, + collate_fn=collate_fn, + sampler=sampler, + num_workers=num_workers, + pin_memory=True, + persistent_workers=num_workers > 0, + drop_last=traceable_sampling, + ) + if num_workers > 0: + dataloader_kwargs["prefetch_factor"] = int(vla_dataset_cfg.get("prefetch_factor", 4)) + vla_train_dataloader = DataLoader(**dataloader_kwargs) + + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + if rank == 0: + output_dir = Path(cfg.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + vla_dataset.save_dataset_statistics(output_dir / "dataset_statistics.json") + if vla_dataset.valid_anchor_manifest is not None: + manifest_path = output_dir / "valid_anchor_manifest.json" + vla_dataset.valid_anchor_manifest.save(manifest_path) + return vla_train_dataloader + elif dataset_py == "vlm_datasets": + from starVLA.dataloader.vlm_datasets import make_vlm_dataloader + + vlm_data_module = make_vlm_dataloader(cfg) + vlm_train_dataloader = vlm_data_module["train_dataloader"] + + return vlm_train_dataloader diff --git a/starVLA/dataloader/gr00t_lerobot/README.md b/starVLA/dataloader/gr00t_lerobot/README.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/starVLA/dataloader/gr00t_lerobot/__init__.py b/starVLA/dataloader/gr00t_lerobot/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/starVLA/dataloader/gr00t_lerobot/data_config.py b/starVLA/dataloader/gr00t_lerobot/data_config.py new file mode 100644 index 0000000000000000000000000000000000000000..618044216ab4af4671418e720df811fea5100445 --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/data_config.py @@ -0,0 +1,1099 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + + +from abc import ABC, abstractmethod + +from starVLA.dataloader.gr00t_lerobot.datasets import ModalityConfig +from starVLA.dataloader.gr00t_lerobot.embodiment_tags import EmbodimentTag +from starVLA.dataloader.gr00t_lerobot.transform.base import ComposedModalityTransform, ModalityTransform +from starVLA.dataloader.gr00t_lerobot.transform.concat import ConcatTransform +from starVLA.dataloader.gr00t_lerobot.transform.state_action import ( + StateActionSinCosTransform, + StateActionToTensor, + StateActionTransform, +) +from starVLA.dataloader.gr00t_lerobot.transform.video import ( + VideoColorJitter, + VideoCrop, + VideoResize, + VideoToNumpy, + VideoToTensor, +) +# from gr00t.model.transforms import GR00TTransform + + +class BaseDataConfig(ABC): + @abstractmethod + def modality_config(self) -> dict[str, ModalityConfig]: + pass + + @abstractmethod + def transform(self) -> ModalityTransform: + pass + + +########################################################################################### + +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): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # video 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), + # state transforms + 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", + }, + ), + # action transforms + 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"}, + ), + # concat transforms + ConcatTransform( + video_concat_order=self.video_keys, + state_concat_order=self.state_keys, + action_concat_order=self.action_keys, + ), + GR00TTransform( + state_horizon=len(self.observation_indices), + action_horizon=len(self.action_indices), + max_state_dim=64, + max_action_dim=32, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + + +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): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # video 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), + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.x": "q99", + "state.y": "q99", + "state.z": "q99", + "state.roll": "q99", + "state.pitch": "q99", + "state.yaw": "q99", + "state.pad": "q99", + "state.gripper": "binary", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.x": "q99", + "action.y": "q99", + "action.z": "q99", + "action.roll": "q99", + "action.pitch": "q99", + "action.yaw": "q99", + "action.gripper": "binary", + }, + ), + # concat transforms + # ConcatTransform( + # # video_concat_order=self.video_keys, + # state_concat_order=self.state_keys, + # action_concat_order=self.action_keys, + # ), + # GR00TTransform( + # state_horizon=len(self.observation_indices), + # action_horizon=len(self.action_indices), + # max_state_dim=64, + # max_action_dim=32, + # ), + ] + + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + +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): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # video 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), + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.x": "q99", + "state.y": "q99", + "state.z": "q99", + "state.rx": "q99", + "state.ry": "q99", + "state.rz": "q99", + "state.rw": "q99", + "state.gripper": "binary", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.x": "q99", + "action.y": "q99", + "action.z": "q99", + "action.roll": "q99", + "action.pitch": "q99", + "action.yaw": "q99", + "action.gripper": "binary", + }, + ), + # concat transforms + # ConcatTransform( + # # video_concat_order=self.video_keys, + # state_concat_order=self.state_keys, + # action_concat_order=self.action_keys, + # ), + # GR00TTransform( + # state_horizon=len(self.observation_indices), + # action_horizon=len(self.action_indices), + # max_state_dim=64, + # max_action_dim=32, + # ), + ] + + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + + +class SingleFrankaRobotiqDeltaEefDataConfig: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + video_keys = [ + "video.base_view", + "video.ego_view", + ] + state_keys = [ + "state.eef_position", + "state.eef_rotation", + ] + action_keys = [ + "action.delta_eef_position", + "action.delta_eef_rotation", + "action.gripper_close", + ] + + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.eef_position": "min_max", + "state.eef_rotation": "min_max", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.delta_eef_position": "min_max", + "action.delta_eef_rotation": "min_max", + "action.gripper_close": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + +########################################################################################### + +class Libero4in1DataConfig: + embodiment_tag = EmbodimentTag.FRANKA + video_keys = [ + "video.primary_image", + ] + + 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(8)) + state_indices = list(range(-16, 0)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.state_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.x": "min_max", + "action.y": "min_max", + "action.z": "min_max", + "action.roll": "min_max", + "action.pitch": "min_max", + "action.yaw": "min_max", + # LIBERO's gripper is a native 0/1 control, not a continuous + # min-max dimension. This preserves 0 -> 0 and 1 -> 1 while + # producing the official statistics mask value `false`. + "action.gripper": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + +########################################################################################### + + +class SingleFrankaRobotiqDeltaJointsDataConfig: + embodiment_tag = EmbodimentTag.FRANKA + video_keys = [ + "video.base_view", + "video.ego_view", + ] + state_keys = [ + "state.joints", + ] + action_keys = [ + "action.delta_joints", + "action.gripper_close", + ] + + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.joints": "min_max", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.delta_joints": "min_max", + "action.gripper_close": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + +class FourierGr1ArmsWaistDataConfig: + embodiment_tag = EmbodimentTag.GR1 + video_keys = ["video.ego_view"] + state_keys = [ + "state.left_arm", + "state.right_arm", + "state.left_hand", + "state.right_hand", + "state.waist", + ] + action_keys = [ + "action.left_arm", + "action.right_arm", + "action.left_hand", + "action.right_hand", + "action.waist", + ] + language_keys = ["annotation.human.coarse_action"] + observation_indices = [0] + action_indices = list(range(16)) + + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self) -> ModalityTransform: + transforms = [ + # video 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), + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionSinCosTransform(apply_to=self.state_keys), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={key: "min_max" for key in self.action_keys}, + ), + # concat transforms + # ConcatTransform( + # video_concat_order=self.video_keys, + # state_concat_order=self.state_keys, + # action_concat_order=self.action_keys, + # ), + ] + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + + +########################################################################################### + +class SO101Config: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + #input + video_keys = [ + "video.primary_image", + "video.wrist_image", + ] + + state_keys = [ + "state.shoulder_pan.pos", + "state.shoulder_lift.pos", + "state.elbow_flex.pos", + "state.wrist_flex.pos", + "state.wrist_roll.pos", + "state.gripper.pos", + ] + language_keys = ["annotation.human.action.task_description"] + + # output + action_keys = [ + "action.shoulder_pan.pos", + "action.shoulder_lift.pos", + "action.elbow_flex.pos", + "action.wrist_flex.pos", + "action.wrist_roll.pos", + "action.gripper.pos", + ] + + + observation_indices = [0] + action_indices = list(range(16)) + + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + key: "min_max" for key in self.state_keys + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + key: "min_max" for key in self.action_keys + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + + + +class ArxX5DataConfig: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + video_keys = [ + "video.cam_high", + "video.cam_left_wrist", + "video.cam_right_wrist", + ] + state_keys = [ + "state.left_joints", + "state.right_joints", + "state.left_gripper", + "state.right_gripper", + ] + action_keys = [ + "action.left_joints", + "action.right_joints", + "action.left_gripper", + "action.right_gripper", + ] + + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.left_joints": "min_max", + "state.right_joints": "min_max", + "state.left_gripper": "binary", + "state.right_gripper": "binary", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.left_joints": "min_max", + "action.right_joints": "min_max", + "action.left_gripper": "binary", + "action.right_gripper": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + +########################################################################################### + + +class AgilexDataConfig: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + video_keys = [ + "video.cam_high", + "video.cam_left_wrist", + "video.cam_right_wrist", + ] + state_keys = [ + "state.left_joints", + "state.right_joints", + "state.left_gripper", + "state.right_gripper", + ] + action_keys = [ + "action.left_joints", + "action.right_joints",#@JinhuiYE this order is different from Dataset + "action.left_gripper", + "action.right_gripper", + ] + + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.left_joints": "min_max", + "state.right_joints": "min_max", + "state.left_gripper": "binary", + "state.right_gripper": "binary", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.left_joints": "min_max", + "action.right_joints": "min_max", + "action.left_gripper": "binary", + "action.right_gripper": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + +########################################################################################### + + +class AgilexData50Config: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + video_keys = [ + "video.cam_high", + "video.cam_left_wrist", + "video.cam_right_wrist", + ] + state_keys = [ + "state.left_joints", + "state.right_joints", + "state.left_gripper", + "state.right_gripper", + ] + action_keys = [ + "action.left_joints", + "action.right_joints", # @JinhuiYE this order is different from Dataset + "action.left_gripper", + "action.right_gripper", + ] + + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(50)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + binary_threshold=0.49, + normalization_modes={ + "state.left_joints": "min_max", + "state.right_joints": "min_max", + "state.left_gripper": "binary", + "state.right_gripper": "binary", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + binary_threshold=0.49, + normalization_modes={ + "action.left_joints": "min_max", + "action.right_joints": "min_max", + "action.left_gripper": "binary", + "action.right_gripper": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + +########################################################################################### + +class VLAArenaFrankaDataConfig: + """ + Data config for VLA-Arena environments (robosuite, mounted Panda). + + Action space : 7-DOF delta EEF (Δx, Δy, Δz, Δroll, Δpitch, Δyaw, gripper) + Observation : primary agentview image (256×256 rendered, resized to 224×224) + State : EEF pos (3) + EEF axis-angle (3) + gripper qpos (1) = 7 + """ + + embodiment_tag = EmbodimentTag.FRANKA + video_keys = [ + "video.primary_image", # agentview camera + ] + state_keys = [ + "state.x", + "state.y", + "state.z", + "state.roll", + "state.pitch", + "state.yaw", + "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(8)) + state_indices = list(range(-16, 0)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.state_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + return { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + + def transform(self): + transforms = [ + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.x": "min_max", + "action.y": "min_max", + "action.z": "min_max", + "action.roll": "min_max", + "action.pitch": "min_max", + "action.yaw": "min_max", + }, + ), + ] + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + +ROBOT_TYPE_CONFIG_MAP = { + "libero_franka": Libero4in1DataConfig(), + "oxe_droid": OxeDroidDataConfig(), + "oxe_bridge": OxeBridgeDataConfig(), + "oxe_rt1": OxeRT1DataConfig(), + "SO101": SO101Config(), + "demo_sim_franka_delta_joints": SingleFrankaRobotiqDeltaJointsDataConfig(), + "arx_x5": ArxX5DataConfig(), + "robotwin": AgilexDataConfig(), + "robotwin50": AgilexData50Config(), + "fourier_gr1_arms_waist": FourierGr1ArmsWaistDataConfig(), + "vla_arena_franka": VLAArenaFrankaDataConfig(), + + "custom_robot_config": SingleFrankaRobotiqDeltaEefDataConfig(), +} diff --git a/starVLA/dataloader/gr00t_lerobot/datasets.py b/starVLA/dataloader/gr00t_lerobot/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..7082123b2ad75e248087f4c5aa51caefb7e79415 --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/datasets.py @@ -0,0 +1,3255 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +""" +In this file, we define 3 types of datasets: +1. LeRobotSingleDataset: a single dataset for a given embodiment tag +2. LeRobotMixtureDataset: a mixture of datasets for a given list of embodiment tags +3. CachedLeRobotSingleDataset: a single dataset for a given embodiment tag, + with caching for the video frames + +See `scripts/load_dataset.py` for examples on how to use these datasets. +""" +import os +from contextlib import contextmanager +import hashlib +import io +import json, torch +import copy +from collections import defaultdict +from pathlib import Path +from typing import Sequence +import os, random +import numpy as np +import pandas as pd +from pydantic import BaseModel, Field, ValidationError +from torch.utils.data import Dataset +from tqdm import tqdm +from PIL import Image +import torch.distributed as dist + +from starVLA.dataloader.gr00t_lerobot.video import get_all_frames, get_frames_by_timestamps + +from starVLA.dataloader.gr00t_lerobot.embodiment_tags import EmbodimentTag +from starVLA.dataloader.gr00t_lerobot.schema import ( + DatasetMetadata, + DatasetStatisticalValues, + LeRobotModalityMetadata, + LeRobotStateActionMetadata, +) +from starVLA.dataloader.gr00t_lerobot.transform import ComposedModalityTransform +from starVLA.dataloader.gr00t_lerobot.transform.state_action import StateActionTransform +from starVLA.dataloader.infoflow_sampler import ( + LogicalDraw, + ValidAnchorManifest, + validate_dataset_content_inventory, +) +from starVLA.dataloader.temporal_proposal import TemporalProposal + +from functools import partial +from typing import Tuple, List +import pickle +import gc + +# LeRobot v2.0 dataset file names +LE_ROBOT_MODALITY_FILENAME = "meta/modality.json" +LE_ROBOT_EPISODE_FILENAME = "meta/episodes.jsonl" +LE_ROBOT_TASKS_FILENAME = "meta/tasks.jsonl" +LE_ROBOT_INFO_FILENAME = "meta/info.json" +LE_ROBOT_STATS_FILENAME = "meta/stats_gr00t.json" +LE_ROBOT_DATA_FILENAME = "data/*/*.parquet" +LE_ROBOT_STEPS_FILENAME = "meta/steps.pkl" +LE_ROBOT_STATS_FORMAT_VERSION = 2 +EPSILON = 5e-4 + + +@contextmanager +def _isolated_transform_rng(seed: int | None): + """Make data augmentation a pure function of one logical draw id.""" + + if seed is None: + yield + return + python_state = random.getstate() + numpy_state = np.random.get_state() + torch_state = torch.get_rng_state() + bounded_seed = int(seed) % (2**32) + random.seed(bounded_seed) + np.random.seed(bounded_seed) + torch.manual_seed(bounded_seed) + try: + yield + finally: + random.setstate(python_state) + np.random.set_state(numpy_state) + torch.set_rng_state(torch_state) + +# LeRobot v3.0 dataset file names +LE_ROBOT3_TASKS_FILENAME = "meta/tasks.parquet" +LE_ROBOT3_EPISODE_FILENAME = "meta/episodes/*/*.parquet" + + +def calculate_dataset_statistics(parquet_paths: list[Path]) -> dict: + """Calculate the dataset statistics of all columns for a list of parquet files.""" + # Dataset statistics + all_low_dim_data_list = [] + # Collect all the data + # parquet_paths = parquet_paths[:3] + for parquet_path in tqdm( + sorted(list(parquet_paths)), + desc="Collecting all parquet files...", + ): + # Load the parquet file + parquet_data = pd.read_parquet(parquet_path) + parquet_data = parquet_data + all_low_dim_data_list.append(parquet_data) + + if not all_low_dim_data_list: + raise FileNotFoundError( + f"No parquet files found under the provided paths: {[str(p) for p in parquet_paths[:3]]}..." + f" — make sure the dataset has been downloaded/converted before training." + ) + + all_low_dim_data = pd.concat(all_low_dim_data_list, axis=0) + # Compute dataset statistics + dataset_statistics = {} + for le_modality in tqdm(all_low_dim_data.columns, desc="Processing modalities"): + print(le_modality) + if "task_info" in le_modality: + continue + print(f"Computing statistics for {le_modality}...") + try: + np_data = np.vstack( + [np.asarray(x, dtype=np.float32) for x in all_low_dim_data[le_modality]] + ) + except Exception as e: + print(f"Warning: Failed to process modality {le_modality} due to error: {e}") + continue + + dataset_statistics[le_modality] = { + "mean": np.mean(np_data, axis=0).tolist(), + "std": np.std(np_data, axis=0).tolist(), + "min": np.min(np_data, axis=0).tolist(), + "max": np.max(np_data, axis=0).tolist(), + "q01": np.quantile(np_data, 0.01, axis=0).tolist(), + "q99": np.quantile(np_data, 0.99, axis=0).tolist(), + } + return dataset_statistics + + +def _normalize_action_mode(mode: str) -> str: + """Normalize action mode names to {abs, delta, rel}.""" + # @gaoning plz move this, we want dataloader to be independent of the action mode logic, we can move this to transform or a separate utils tool to handle lerobot dataset + mode = str(mode).lower() + if mode in {"absolute", "raw"}: + mode = "abs" + if mode not in {"abs", "delta", "rel"}: + mode = "abs" + return mode + + +def _normalize_action_mode_apply_keys( + action_mode_apply_keys: Sequence[str] | None, + fallback_keys: Sequence[str] | None = None, +) -> list[str]: + source_keys = action_mode_apply_keys if action_mode_apply_keys else (fallback_keys or []) + normalized = [] + for key in source_keys: + key = str(key) + if not key.startswith("action."): + key = f"action.{key}" + normalized.append(key) + return normalized + + +def _normalize_action_mode_state_map(action_mode_state_map: dict[str, str] | None) -> dict[str, str]: + normalized = {} + for action_key, state_key in (action_mode_state_map or {}).items(): + action_key = str(action_key) + state_key = str(state_key) + if not action_key.startswith("action."): + action_key = f"action.{action_key}" + if not state_key.startswith("state."): + state_key = f"state.{state_key}" + normalized[action_key] = state_key + return normalized + + +def _build_stats_cache_config( + action_mode: str, +) -> dict: + return { + "mode": action_mode, + } + + +def _invalidate_legacy_stats_cache(stats_path: Path, reason: str) -> None: + if not stats_path.exists(): + return + print(f"Removing stale dataset statistics cache at {stats_path}: {reason}") + stats_path.unlink() + + +def _load_stats_cache( + stats_path: Path, + expected_config: dict, + *, + invalidate_legacy: bool, +) -> dict | None: + if not stats_path.exists(): + return None + + try: + with open(stats_path, "r") as f: + payload = json.load(f) + except Exception as exc: + if invalidate_legacy: + _invalidate_legacy_stats_cache(stats_path, f"failed to load JSON ({exc})") + return None + + if not isinstance(payload, dict): + if invalidate_legacy: + _invalidate_legacy_stats_cache(stats_path, "unexpected top-level format") + return None + + format_version = payload.get("__format_version") + cache_config = payload.get("__cache_config") + statistics = payload.get("statistics") + if format_version != LE_ROBOT_STATS_FORMAT_VERSION or cache_config is None or statistics is None: + if invalidate_legacy: + _invalidate_legacy_stats_cache(stats_path, "legacy statistics format detected") + return None + + if cache_config != expected_config: + if invalidate_legacy: + _invalidate_legacy_stats_cache(stats_path, "statistics config mismatch, rebuilding cache") + return None + + return statistics + + +def _save_stats_cache(stats_path: Path, cache_config: dict, statistics: dict) -> None: + payload = { + "__format_version": LE_ROBOT_STATS_FORMAT_VERSION, + "__cache_config": cache_config, + "statistics": statistics, + } + tmp_path = stats_path.with_suffix(".tmp") + stats_path.parent.mkdir(parents=True, exist_ok=True) + with open(tmp_path, "w") as f: + json.dump(payload, f, indent=4) + os.replace(tmp_path, stats_path) + + +def _compute_statistics_for_mode( + parquet_paths: list[Path], + dataset_name: str, + action_mode: str, + lerobot_modality_meta: "LeRobotModalityMetadata", + action_keys_full: list[str], + state_keys_full: list[str], + action_indices: list[int] | None, + state_indices: list[int] | None, + action_mode_apply_keys: list[str] | None, + action_mode_state_map: dict[str, str] | None, +) -> dict: + if int(os.environ.get("RANK", "0")) == 0: + print(f"[RANK 0] Calculating dataset statistics for {dataset_name} (mode={action_mode})") + + base_stats = calculate_dataset_statistics(parquet_paths) + + if action_mode == "abs": + return base_stats + + if action_indices is None or state_indices is None: + raise ValueError( + "Both action and state modalities are required to compute " + f"{action_mode} action mode statistics." + ) + + if action_mode == "delta": + return calculate_delta_action_statistics( + parquet_paths=parquet_paths, + lerobot_modality_meta=lerobot_modality_meta, + action_keys_full=action_keys_full, + state_keys_full=state_keys_full, + action_indices=action_indices, + state_indices=state_indices, + action_mode_apply_keys=action_mode_apply_keys, + action_mode_state_map=action_mode_state_map, + base_stats=base_stats, + ) + if action_mode == "rel": + return calculate_rel_action_statistics( + parquet_paths=parquet_paths, + lerobot_modality_meta=lerobot_modality_meta, + action_keys_full=action_keys_full, + state_keys_full=state_keys_full, + action_indices=action_indices, + state_indices=state_indices, + action_mode_apply_keys=action_mode_apply_keys, + action_mode_state_map=action_mode_state_map, + base_stats=base_stats, + ) + raise ValueError(f"Unsupported action mode for statistics: {action_mode}") + + +def _load_or_compute_statistics( + stats_path: Path, + stats_cache_config: dict, + parquet_paths: list[Path], + dataset_name: str, + action_mode: str, + lerobot_modality_meta: "LeRobotModalityMetadata", + action_keys_full: list[str], + state_keys_full: list[str], + action_indices: list[int] | None, + state_indices: list[int] | None, + action_mode_apply_keys: list[str] | None, + action_mode_state_map: dict[str, str] | None, +) -> dict: + le_statistics = _load_stats_cache( + stats_path, + stats_cache_config, + invalidate_legacy=True, + ) + if le_statistics is not None: + return le_statistics + + le_statistics = _compute_statistics_for_mode( + parquet_paths=parquet_paths, + dataset_name=dataset_name, + action_mode=action_mode, + lerobot_modality_meta=lerobot_modality_meta, + action_keys_full=action_keys_full, + state_keys_full=state_keys_full, + action_indices=action_indices, + state_indices=state_indices, + action_mode_apply_keys=action_mode_apply_keys, + action_mode_state_map=action_mode_state_map, + ) + _save_stats_cache(stats_path, stats_cache_config, le_statistics) + return le_statistics + + +def _get_action_col_slices( + lerobot_modality_meta: "LeRobotModalityMetadata", + action_keys_full: list[str], + state_keys_full: list[str], + action_mode_apply_keys: list[str] | None = None, + action_mode_state_map: dict[str, str] | None = None, +) -> dict[str, list[tuple[tuple[int, int], str, tuple[int, int], str, str]]]: + apply_keys = _normalize_action_mode_apply_keys(action_mode_apply_keys, action_keys_full) + action_mode_state_map = _normalize_action_mode_state_map(action_mode_state_map) + + action_meta = lerobot_modality_meta.action + state_meta = lerobot_modality_meta.state + + # Build per-column mapping: action column -> list of (action_slice, state_column, state_slice) + action_col_slices: dict[str, list[tuple[tuple[int, int], str, tuple[int, int]]]] = {} + for action_key in apply_keys: + if not action_key.startswith("action."): + raise ValueError(f"Invalid action key {action_key}. Expected prefix 'action.'.") + state_key = action_mode_state_map.get(action_key, action_key.replace("action.", "state.", 1)) + if state_key not in state_keys_full: + raise ValueError( + f"State key {state_key} not found for action key {action_key}. " + f"Add it to action_mode_state_map or remove {action_key} from action_mode_apply_keys." + ) + + action_subkey = action_key.replace("action.", "", 1) + state_subkey = state_key.replace("state.", "", 1) + if action_subkey not in action_meta or state_subkey not in state_meta: + raise ValueError(f"Action/state key missing in metadata: {action_key} -> {state_key}") + + action_cfg = action_meta[action_subkey] + state_cfg = state_meta[state_subkey] + action_col = action_cfg.original_key or action_subkey + state_col = state_cfg.original_key or state_subkey + action_slice = (action_cfg.start, action_cfg.end) + state_slice = (state_cfg.start, state_cfg.end) + action_padding = "first_last" if action_cfg.absolute else "zero" + state_padding = "first_last" if state_cfg.absolute else "zero" + action_col_slices.setdefault(action_col, []).append( + (action_slice, state_col, state_slice, action_padding, state_padding) + ) + + return action_col_slices + + +def calculate_delta_action_statistics( + parquet_paths: list[Path], + lerobot_modality_meta: "LeRobotModalityMetadata", + action_keys_full: list[str], + state_keys_full: list[str], + action_indices: list[int], + state_indices: list[int], + action_mode_apply_keys: list[str] | None = None, + action_mode_state_map: dict[str, str] | None = None, + base_stats: dict | None = None, +) -> dict: + """ + Calculate action statistics using delta mode. + + Rule: + - For t>0: a_t - a_{t-1} + - For t=0: a_0 - s_0 + + Mapping rule (only two cases): + 1) Use explicit action_mode_state_map if provided. + 2) Otherwise, replace 'action.' with 'state.' directly. + """ + if base_stats is None: + base_stats = calculate_dataset_statistics(parquet_paths) + + action_col_slices = _get_action_col_slices( + lerobot_modality_meta, action_keys_full, state_keys_full, action_mode_apply_keys, action_mode_state_map + ) + if not action_col_slices: + raise ValueError("No action columns found in the dataset.") + + def _get_chunk(array: np.ndarray, step_indices: np.ndarray, padding_strategy: str) -> np.ndarray: + max_length = array.shape[0] + front_padding = step_indices < 0 + end_padding = step_indices >= max_length + padding_positions = np.logical_or(front_padding, end_padding) + output = np.zeros((len(step_indices), array.shape[1]), dtype=array.dtype) + if (~padding_positions).any(): + output[~padding_positions] = array[step_indices[~padding_positions]] + if padding_positions.any(): + if padding_strategy == "first_last": + output[front_padding] = array[0] + output[end_padding] = array[-1] + elif padding_strategy == "zero": + output[padding_positions] = 0 + else: + raise ValueError(f"Invalid padding strategy: {padding_strategy}") + return output + + accum: dict[str, list[np.ndarray]] = {col: [] for col in action_col_slices.keys()} + for parquet_path in tqdm(sorted(list(parquet_paths)), desc="Collecting delta action stats"): + data = pd.read_parquet(parquet_path) + trajectory_length = len(data) + for action_col, slice_list in action_col_slices.items(): + if action_col not in data.columns: + raise ValueError(f"{action_col} not found in parquet columns.") + action_matrix = np.stack(data[action_col]) + action_padding_ref = slice_list[0][3] + prepared_slices = [] + for a_slice, state_col, s_slice, action_padding, state_padding in slice_list: + if state_col not in data.columns: + raise ValueError(f"{state_col} not found in parquet columns.") + state_matrix = np.stack(data[state_col]) + state_part_full = state_matrix[:, s_slice[0] : s_slice[1]] + prepared_slices.append((a_slice, state_part_full, state_padding)) + for base_index in range(trajectory_length): + action_steps = np.array(action_indices) + base_index + action_chunk_full = _get_chunk(action_matrix, action_steps, action_padding_ref) + + for a_slice, state_part_full, state_padding in prepared_slices: + action_part_chunk = action_chunk_full[:, a_slice[0] : a_slice[1]] + state_chunk = _get_chunk(state_part_full, np.array(state_indices) + base_index, state_padding) + if action_part_chunk.shape[1] != state_chunk.shape[1]: + raise ValueError(f"Action/state dim mismatch for {action_col}:{a_slice}") + + out = action_part_chunk.copy() + if len(out) > 1: + out[1:] = action_part_chunk[1:] - action_part_chunk[:-1] + out[0] = action_part_chunk[0] - state_chunk[0] + action_chunk_full[:, a_slice[0] : a_slice[1]] = out + + accum[action_col].append(action_chunk_full) + + delta_stats = copy.deepcopy(base_stats) + for action_col, series_list in accum.items(): + if not series_list: + continue + all_values = np.concatenate(series_list, axis=0).astype(np.float32) + delta_stats[action_col] = { + "mean": np.mean(all_values, axis=0).tolist(), + "std": np.std(all_values, axis=0).tolist(), + "min": np.min(all_values, axis=0).tolist(), + "max": np.max(all_values, axis=0).tolist(), + "q01": np.quantile(all_values, 0.01, axis=0).tolist(), + "q99": np.quantile(all_values, 0.99, axis=0).tolist(), + } + return delta_stats + + +def calculate_rel_action_statistics( + parquet_paths: list[Path], + lerobot_modality_meta: "LeRobotModalityMetadata", + action_keys_full: list[str], + state_keys_full: list[str], + action_indices: list[int], + state_indices: list[int], + action_mode_apply_keys: list[str] | None = None, + action_mode_state_map: dict[str, str] | None = None, + base_stats: dict | None = None, +) -> dict: + """ + Calculate action statistics using rel mode. + + Rule: + - For all t: a_t - s_0 + + Mapping rule (only two cases): + 1) Use explicit action_mode_state_map if provided. + 2) Otherwise, replace 'action.' with 'state.' directly. + """ + if base_stats is None: + base_stats = calculate_dataset_statistics(parquet_paths) + + action_col_slices = _get_action_col_slices( + lerobot_modality_meta, action_keys_full, state_keys_full, action_mode_apply_keys, action_mode_state_map + ) + if not action_col_slices: + raise ValueError("No action columns found in the dataset.") + + def _get_chunk(array: np.ndarray, step_indices: np.ndarray, padding_strategy: str) -> np.ndarray: + max_length = array.shape[0] + front_padding = step_indices < 0 + end_padding = step_indices >= max_length + padding_positions = np.logical_or(front_padding, end_padding) + output = np.zeros((len(step_indices), array.shape[1]), dtype=array.dtype) + if (~padding_positions).any(): + output[~padding_positions] = array[step_indices[~padding_positions]] + if padding_positions.any(): + if padding_strategy == "first_last": + output[front_padding] = array[0] + output[end_padding] = array[-1] + elif padding_strategy == "zero": + output[padding_positions] = 0 + else: + raise ValueError(f"Invalid padding strategy: {padding_strategy}") + return output + + accum: dict[str, list[np.ndarray]] = {col: [] for col in action_col_slices.keys()} + for parquet_path in tqdm(sorted(list(parquet_paths)), desc="Collecting rel action stats"): + data = pd.read_parquet(parquet_path) + trajectory_length = len(data) + for action_col, slice_list in action_col_slices.items(): + if action_col not in data.columns: + raise ValueError(f"{action_col} not found in parquet columns.") + action_matrix = np.stack(data[action_col]) + action_padding_ref = slice_list[0][3] + prepared_slices = [] + for a_slice, state_col, s_slice, action_padding, state_padding in slice_list: + if state_col not in data.columns: + raise ValueError(f"{state_col} not found in parquet columns.") + state_matrix = np.stack(data[state_col]) + state_part_full = state_matrix[:, s_slice[0] : s_slice[1]] + prepared_slices.append((a_slice, state_part_full, state_padding)) + for base_index in range(trajectory_length): + action_steps = np.array(action_indices) + base_index + action_chunk_full = _get_chunk(action_matrix, action_steps, action_padding_ref) + + for a_slice, state_part_full, state_padding in prepared_slices: + action_part_chunk = action_chunk_full[:, a_slice[0] : a_slice[1]] + state_chunk = _get_chunk(state_part_full, np.array(state_indices) + base_index, state_padding) + if action_part_chunk.shape[1] != state_chunk.shape[1]: + raise ValueError(f"Action/state dim mismatch for {action_col}:{a_slice}") + + out = action_part_chunk - state_chunk[0] + action_chunk_full[:, a_slice[0] : a_slice[1]] = out + + accum[action_col].append(action_chunk_full) + + rel_stats = copy.deepcopy(base_stats) + for action_col, series_list in accum.items(): + if not series_list: + continue + all_values = np.concatenate(series_list, axis=0).astype(np.float32) + rel_stats[action_col] = { + "mean": np.mean(all_values, axis=0).tolist(), + "std": np.std(all_values, axis=0).tolist(), + "min": np.min(all_values, axis=0).tolist(), + "max": np.max(all_values, axis=0).tolist(), + "q01": np.quantile(all_values, 0.01, axis=0).tolist(), + "q99": np.quantile(all_values, 0.99, axis=0).tolist(), + } + return rel_stats + +class ModalityConfig(BaseModel): + """Configuration for a modality.""" + + delta_indices: list[int] + """Delta indices to sample relative to the current index. The returned data will correspond to the original data at a sampled base index + delta indices.""" + modality_keys: list[str] + """The keys to load for the modality in the dataset.""" + + +class LeRobotSingleDataset(Dataset): + """ + Base dataset class for LeRobot that supports sharding. + """ + def __init__( + self, + dataset_path: Path | str, + modality_configs: dict[str, ModalityConfig], + embodiment_tag: str | EmbodimentTag, + video_backend: str = "decord", + video_backend_kwargs: dict | None = None, + transforms: ComposedModalityTransform | None = None, + delete_pause_frame: bool = False, + data_cfg = None, + **kwargs, + ): + """ + Initialize the dataset. + + Args: + dataset_path (Path | str): The path to the dataset. + modality_configs (dict[str, ModalityConfig]): The configuration for each modality. The keys are the modality names, and the values are the modality configurations. + See `ModalityConfig` for more details. + video_backend (str): Backend for video reading. + video_backend_kwargs (dict): Keyword arguments for the video backend when initializing the video reader. + transforms (ComposedModalityTransform): The transforms to apply to the dataset. + embodiment_tag (EmbodimentTag): Overload the embodiment tag for the dataset. e.g. define it as "new_embodiment" + """ + # first check if the path directory exists + self.data_cfg = data_cfg + if not Path(dataset_path).exists(): + raise FileNotFoundError(f"Dataset path {dataset_path} does not exist") + # indict letobot version + self._lerobot_version = self.data_cfg.get("lerobot_version", "v2.0") #self._indict_lerobot_version(**kwargs) + + self._action_mode = None + self._action_mode_state_map = {} + self._action_mode_apply_keys = None + + self.delete_pause_frame = delete_pause_frame + + self.modality_configs = modality_configs + self.video_backend = video_backend + self.video_backend_kwargs = video_backend_kwargs if video_backend_kwargs is not None else {} + self.transforms = ( + transforms if transforms is not None else ComposedModalityTransform(transforms=[]) + ) + + self._dataset_path = Path(dataset_path) + self._dataset_name = kwargs.get("dataset_name", self._dataset_path.name) + if isinstance(embodiment_tag, EmbodimentTag): + self.tag = embodiment_tag.value + else: + self.tag = embodiment_tag + + self._init_action_mode() + self._metadata = self._get_metadata(EmbodimentTag(self.tag)) + + # LeRobot-specific config + self._lerobot_modality_meta = self._get_lerobot_modality_meta() + self._lerobot_info_meta = self._get_lerobot_info_meta() + self._data_path_pattern = self._get_data_path_pattern() + self._video_path_pattern = self._get_video_path_pattern() + self._chunk_size = self._get_chunk_size() + self._tasks = self._get_tasks() + # self._episodes = self._get_episode_info() # TODO why we need this func + self.curr_traj_data = None + self.curr_traj_id = None + + self._trajectory_ids, self._trajectory_lengths = self._get_trajectories() + self._modality_keys = self._get_modality_keys() + self._delta_indices = self._get_delta_indices() + self._all_steps = self._get_all_steps() + self.set_transforms_metadata(self.metadata) + self.set_epoch(0) + + if int(os.environ.get("RANK", "0")) == 0: + print(f"Initialized dataset {self.dataset_name} with {embodiment_tag}") + + # Check if the dataset is valid + self._check_integrity() + + @property + def dataset_path(self) -> Path: + """The path to the dataset that contains the METADATA_FILENAME file.""" + return self._dataset_path + + @property + def metadata(self) -> DatasetMetadata: + """The metadata for the dataset, loaded from metadata.json in the dataset directory""" + return self._metadata + + @property + def trajectory_ids(self) -> np.ndarray: + """The trajectory IDs in the dataset, stored as a 1D numpy array of strings.""" + return self._trajectory_ids + + @property + def trajectory_lengths(self) -> np.ndarray: + """The trajectory lengths in the dataset, stored as a 1D numpy array of integers. + The order of the lengths is the same as the order of the trajectory IDs. + """ + return self._trajectory_lengths + + @property + def all_steps(self) -> list[tuple[int, int]]: + """The trajectory IDs and base indices for all steps in the dataset. + Example: + self.trajectory_ids: [0, 1, 2] + self.trajectory_lengths: [3, 2, 4] + return: [ + ("traj_0", 0), ("traj_0", 1), ("traj_0", 2), + ("traj_1", 0), ("traj_1", 1), + ("traj_2", 0), ("traj_2", 1), ("traj_2", 2), ("traj_2", 3) + ] + """ + return self._all_steps + + @property + def modality_keys(self) -> dict: + """The modality keys for the dataset. The keys are the modality names, and the values are the keys for each modality. + + Example: { + "video": ["video.image_side_0", "video.image_side_1"], + "state": ["state.eef_position", "state.eef_rotation"], + "action": ["action.eef_position", "action.eef_rotation"], + "language": ["language.human.task"], + "timestamp": ["timestamp"], + "reward": ["reward"], + } + """ + return self._modality_keys + + @property + def delta_indices(self) -> dict[str, np.ndarray]: + """The delta indices for the dataset. The keys are the modality.key, and the values are the delta indices for each modality.key.""" + return self._delta_indices + + @property + def dataset_name(self) -> str: + """The name of the dataset.""" + return self._dataset_name + + @property + def lerobot_modality_meta(self) -> LeRobotModalityMetadata: + """The metadata for the LeRobot dataset.""" + return self._lerobot_modality_meta + + @property + def lerobot_info_meta(self) -> dict: + """The metadata for the LeRobot dataset.""" + return self._lerobot_info_meta + + @property + def data_path_pattern(self) -> str: + """The path pattern for the LeRobot dataset.""" + return self._data_path_pattern + + @property + def video_path_pattern(self) -> str: + """The path pattern for the LeRobot dataset.""" + return self._video_path_pattern + + @property + def chunk_size(self) -> int: + """The chunk size for the LeRobot dataset.""" + return self._chunk_size + + @property + def tasks(self) -> pd.DataFrame: + """The tasks for the dataset.""" + return self._tasks + + def _get_metadata(self, embodiment_tag: EmbodimentTag) -> DatasetMetadata: + """Get the metadata for the dataset. + + Returns: + dict: The metadata for the dataset. + """ + + # 1. Modality metadata + modality_meta_path = self.dataset_path / LE_ROBOT_MODALITY_FILENAME + assert ( + modality_meta_path.exists() + ), f"Please provide a {LE_ROBOT_MODALITY_FILENAME} file in {self.dataset_path}" + # 1.1. State and action modalities + simplified_modality_meta: dict[str, dict] = {} + with open(modality_meta_path, "r") as f: + le_modality_meta = LeRobotModalityMetadata.model_validate(json.load(f)) + for modality in ["state", "action"]: + simplified_modality_meta[modality] = {} + le_state_action_meta: dict[str, LeRobotStateActionMetadata] = getattr( + le_modality_meta, modality + ) + for subkey in le_state_action_meta: + state_action_dtype = np.dtype(le_state_action_meta[subkey].dtype) + if np.issubdtype(state_action_dtype, np.floating): + continuous = True + else: + continuous = False + simplified_modality_meta[modality][subkey] = { + "absolute": le_state_action_meta[subkey].absolute, + "rotation_type": le_state_action_meta[subkey].rotation_type, + "shape": [ + le_state_action_meta[subkey].end - le_state_action_meta[subkey].start + ], + "continuous": continuous, + } + + # 1.2. Video modalities + le_info_path = self.dataset_path / LE_ROBOT_INFO_FILENAME + assert ( + le_info_path.exists() + ), f"Please provide a {LE_ROBOT_INFO_FILENAME} file in {self.dataset_path}" + with open(le_info_path, "r") as f: + le_info = json.load(f) + simplified_modality_meta["video"] = {} + for new_key in le_modality_meta.video: + original_key = le_modality_meta.video[new_key].original_key + if original_key is None: + original_key = new_key + le_video_meta = le_info["features"][original_key] + height = le_video_meta["shape"][le_video_meta["names"].index("height")] + width = le_video_meta["shape"][le_video_meta["names"].index("width")] + # NOTE(FH): different lerobot dataset versions have different keys for the number of channels and fps + try: + channels = le_video_meta["shape"][le_video_meta["names"].index("channel")] + fps = le_video_meta["video_info"]["video.fps"] + except (ValueError, KeyError): + try: + channels = le_video_meta["info"]["video.channels"] + fps = le_video_meta["info"]["video.fps"] + except (ValueError, KeyError): + # Fallback for image-only datasets (e.g. VLA-Arena) that lack video_info + channels = 3 + fps = le_info.get("fps", 30) + simplified_modality_meta["video"][new_key] = { + "resolution": [width, height], + "channels": channels, + "fps": fps, + } + + + # 2. Dataset statistics + def is_main(): + return (not dist.is_initialized()) or dist.get_rank() == 0 + + action_mode = _normalize_action_mode(self.data_cfg.get("action_mode", "abs") if self.data_cfg else "abs") + + stats_path = self.dataset_path / LE_ROBOT_STATS_FILENAME + action_cfg = self.modality_configs.get("action") + state_cfg = self.modality_configs.get("state") + action_keys_full = list(action_cfg.modality_keys) if action_cfg else [] + state_keys_full = list(state_cfg.modality_keys) if state_cfg else [] + action_indices = list(action_cfg.delta_indices) if action_cfg else None + state_indices = list(state_cfg.delta_indices) if state_cfg else None + + apply_keys = _normalize_action_mode_apply_keys( + self.data_cfg.get("action_mode_apply_keys", None) if self.data_cfg else None, + action_keys_full, + ) + normalized_state_map = _normalize_action_mode_state_map( + self.data_cfg.get("action_mode_state_map", {}) if self.data_cfg else {} + ) + stats_cache_config = _build_stats_cache_config( + action_mode=action_mode, + ) + precomputed_path = ( + self.data_cfg.get("precomputed_statistics_path") + if self.data_cfg + else None + ) + precomputed_filename = ( + self.data_cfg.get("precomputed_statistics_filename") + if self.data_cfg + else None + ) + if precomputed_path: + with open(precomputed_path, "r", encoding="utf-8") as statistics_file: + statistics_payload = json.load(statistics_file) + statistics_key = self.data_cfg.precomputed_statistics_key_map[ + self.dataset_name + ] + dataset_statistics = statistics_payload[statistics_key] + le_statistics = { + "action": dataset_statistics["action"], + "observation.state": dataset_statistics["state"], + } + elif precomputed_filename: + if action_mode != "abs": + raise ValueError( + "precomputed_statistics_filename currently requires action_mode=abs" + ) + precomputed_path = self.dataset_path / precomputed_filename + with open(precomputed_path, "r", encoding="utf-8") as statistics_file: + le_statistics = json.load(statistics_file) + required_columns = { + le_modality_meta.get_key_meta(key).original_key + for key in action_keys_full + state_keys_full + } + le_statistics = { + key: le_statistics[key] for key in required_columns + } + else: + parquet_files = list(self.dataset_path.glob(LE_ROBOT_DATA_FILENAME)) + parquet_files_filtered = [ + pf for pf in parquet_files if "episode_033675.parquet" not in pf.name + ] + + if is_main(): + le_statistics = _load_or_compute_statistics( + stats_path, + stats_cache_config=stats_cache_config, + parquet_paths=parquet_files_filtered, + dataset_name=self.dataset_name, + action_mode=action_mode, + lerobot_modality_meta=le_modality_meta, + action_keys_full=action_keys_full, + state_keys_full=state_keys_full, + action_indices=action_indices, + state_indices=state_indices, + action_mode_apply_keys=apply_keys, + action_mode_state_map=normalized_state_map, + ) + else: + le_statistics = None + + if dist.is_initialized(): + dist.barrier() + + if le_statistics is None: + le_statistics = _load_stats_cache( + stats_path, + stats_cache_config, + invalidate_legacy=False, + ) + if le_statistics is None: + raise RuntimeError( + f"Dataset statistics cache is missing or invalid after sync: {stats_path}" + ) + + for stat in le_statistics.values(): + DatasetStatisticalValues.model_validate(stat) + + + dataset_statistics = {} + for our_modality in ["state", "action"]: + dataset_statistics[our_modality] = {} + for subkey in simplified_modality_meta[our_modality]: + dataset_statistics[our_modality][subkey] = {} + state_action_meta = le_modality_meta.get_key_meta(f"{our_modality}.{subkey}") + assert isinstance(state_action_meta, LeRobotStateActionMetadata) + le_modality = state_action_meta.original_key + for stat_name in le_statistics[le_modality]: + indices = np.arange( + state_action_meta.start, + state_action_meta.end, + ) + stat = np.array(le_statistics[le_modality][stat_name]) + dataset_statistics[our_modality][subkey][stat_name] = stat[indices].tolist() + + # 3. Full dataset metadata + metadata = DatasetMetadata( + statistics=dataset_statistics, # type: ignore + modalities=simplified_modality_meta, # type: ignore + embodiment_tag=embodiment_tag, + ) + + return metadata + + def _get_trajectories(self) -> tuple[np.ndarray, np.ndarray]: + """Get the trajectories in the dataset.""" + # Get trajectory lengths, IDs, and whitelist from dataset metadata + # v2.0 + if self._lerobot_version == "v2.0": + file_path = self.dataset_path / LE_ROBOT_EPISODE_FILENAME + with open(file_path, "r") as f: + episode_metadata = [json.loads(line) for line in f] + trajectory_ids = [] + trajectory_lengths = [] + for episode in episode_metadata: + trajectory_ids.append(episode["episode_index"]) + trajectory_lengths.append(episode["length"]) + return np.array(trajectory_ids), np.array(trajectory_lengths) + # v3.0 + elif self._lerobot_version == "v3.0": + file_paths = sorted(list((self.dataset_path).glob(LE_ROBOT3_EPISODE_FILENAME))) + trajectory_ids = [] + trajectory_lengths = [] + # data_chunck_index = [] + # data_file_index = [] + # vido_from_index = [] + self.trajectory_ids_to_metadata = {} + for file_path in file_paths: + episodes_data = pd.read_parquet(file_path) + timestamp_cols = [ + c + for c in episodes_data.columns + if str(c).startswith("videos/") and str(c).endswith("/from_timestamp") + ] + for index, episode in episodes_data.iterrows(): + trajectory_ids.append(episode["episode_index"]) + trajectory_lengths.append(episode["length"]) + + from_timestamps = {} + for col in timestamp_cols: + value = episode[col] + if pd.isna(value): + continue + # videos/{video_key}/from_timestamp -> {video_key} + video_key = str(col)[len("videos/") : -len("/from_timestamp")] + from_timestamps[video_key] = float(value) + + # TODO auto map key + # Collect video file indices for each video key + #已修改的lerobotv3.0的视频索引（提取视频和文件的索引） + video_file_indices = {} + for col in timestamp_cols: + video_key = str(col)[len("videos/") : -len("/from_timestamp")] + chunk_col = f"videos/{video_key}/chunk_index" + file_col = f"videos/{video_key}/file_index" + if chunk_col in episode and file_col in episode: + video_file_indices[video_key] = { + "chunk_index": int(episode[chunk_col]), + "file_index": int(episode[file_col]), + } + print(video_file_indices) + episode_meta = { + "data/chunk_index": episode["data/chunk_index"], + "data/file_index": episode["data/file_index"], + "data/file_from_index": index, + "videos/from_timestamps": from_timestamps, + "videos/file_indices": video_file_indices, + } + # episode_meta = { + # "data/chunk_index": episode["data/chunk_index"], + # "data/file_index": episode["data/file_index"], + # "data/file_from_index": index, + # "videos/from_timestamps": from_timestamps, + # } + self.trajectory_ids_to_metadata[trajectory_ids[-1]] = episode_meta + + # Should be able to directly read the saved index info here + return np.array(trajectory_ids), np.array(trajectory_lengths) + + def _get_all_steps(self) -> list[tuple[int, int]]: + """Get the trajectory IDs and base indices for all steps in the dataset. + + Returns: + list[tuple[str, int]]: A list of (trajectory_id, base_index) tuples. + """ + def is_main(): + return (not dist.is_initialized()) or dist.get_rank() == 0 + + config_key = self._get_steps_config_key() + steps_filename = "steps_data_index.pkl" + steps_path = self.dataset_path / "meta" / steps_filename + + # ---------- try to read from cache ---------- + if steps_path.exists(): + try: + with open(steps_path, "rb") as f: + cached_data = pickle.load(f) + if cached_data.get("config_key") != config_key: + raise ValueError( + "cached step index was built for a different dataset configuration" + ) + if cached_data.get("num_trajectories") != len(self.trajectory_ids): + raise ValueError("cached step index has a different trajectory count") + steps = cached_data["steps"] + if cached_data.get("total_steps") != len(steps): + raise ValueError("cached step index has an inconsistent step count") + return steps + except Exception as e: + # include EOFError / PickleError / KeyError + print( + f"[RANK {os.environ.get('RANK', 'NA')}] " + f"Failed to load cached steps ({e}), will rebuild." + ) + + # ---------- only build by rank0 ---------- + if is_main(): + all_steps = self._get_all_steps_single_process() + + cache_data = { + "config_key": config_key, + "steps": all_steps, + "num_trajectories": len(self.trajectory_ids), + "total_steps": len(all_steps), + "computed_timestamp": pd.Timestamp.now().isoformat(), + "delete_pause_frame": self.delete_pause_frame, + } + + steps_path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = steps_path.with_suffix(".tmp") + + with open(tmp_path, "wb") as f: + pickle.dump(cache_data, f, protocol=pickle.HIGHEST_PROTOCOL) + os.replace(tmp_path, steps_path) + + print(f"[RANK 0] Cached steps saved to {steps_path}") + + # ---------- sync after rank0 ---------- + if dist.is_initialized(): + dist.barrier() + + # ---------- read by all rank ---------- + with open(steps_path, "rb") as f: + cached_data = pickle.load(f) + + return cached_data["steps"] + + def _get_steps_config_key(self) -> str: + """Generate a configuration key for steps caching.""" + config_dict = { + "delete_pause_frame": self.delete_pause_frame, + "dataset_name": self.dataset_name, + "lerobot_version": self._lerobot_version, + "language_keys": self.modality_keys.get("language", []), + "trajectory_ids": self.trajectory_ids.tolist(), + "trajectory_lengths": self.trajectory_lengths.tolist(), + } + # Create a hash of the configuration + config_str = str(sorted(config_dict.items())) + return hashlib.md5(config_str.encode()).hexdigest()[:12] # + + + def _get_all_steps_single_process(self) -> list[tuple[int, int]]: + """Original single-process implementation as fallback.""" + all_steps: list[tuple[int, int]] = [] + skipped_trajectories = 0 + processed_trajectories = 0 + + # Check if language modality is configured + has_language_modality = 'language' in self.modality_keys and len(self.modality_keys['language']) > 0 + # TODO why trajectory_length here, why not use data length? + for trajectory_id, trajectory_length in tqdm(zip(self.trajectory_ids, self.trajectory_lengths), total=len(self.trajectory_ids), desc="Getting All Step"): + try: + if self._lerobot_version == "v2.0": + data = self.get_trajectory_data(trajectory_id) + elif self._lerobot_version == "v3.0": + data = self.get_trajectory_data_lerobot_v3(trajectory_id) + + trajectory_skipped = False + + # Check if trajectory has valid language instruction (if language modality is configured) + if has_language_modality: + self.curr_traj_data = data # Set current trajectory data for get_language to work + + language_instruction = self.get_language(trajectory_id, self.modality_keys['language'][0], 0) + if not language_instruction or language_instruction[0] == "": + print(f"Skipping trajectory {trajectory_id} due to empty language instruction") + skipped_trajectories += 1 + trajectory_skipped = True + continue + + except Exception as e: + print(f"Skipping trajectory {trajectory_id} due to read error: {e}") + skipped_trajectories += 1 + trajectory_skipped = True + continue + + if not trajectory_skipped: + processed_trajectories += 1 + + for base_index in range(trajectory_length): + all_steps.append((trajectory_id, base_index)) + + # Print summary statistics + print(f"Single-process summary: Processed {processed_trajectories} trajectories, skipped {skipped_trajectories} empty trajectories") + print(f"Total steps: {len(all_steps)} from {len(self.trajectory_ids)} trajectories") + + return all_steps + + def _get_position_and_gripper_values(self, data: pd.DataFrame) -> tuple[list, list]: + """Get position and gripper values based on available columns in the dataset.""" + # Get action keys from modality_keys + action_keys = self.modality_keys.get('action', []) + + # Extract position data + delta_position_values = None + position_candidates = ['delta_eef_position'] + coordinate_candidates = ['x', 'y', 'z'] + + # First try combined position fields + for pos_key in position_candidates: + full_key = f"action.{pos_key}" + if full_key in action_keys: + try: + # Get the lerobot key for this modality + le_action_cfg = self.lerobot_modality_meta.action + subkey = pos_key + if subkey in le_action_cfg: + le_key = le_action_cfg[subkey].original_key or subkey + if le_key in data.columns: + data_array = np.stack(data[le_key]) + le_indices = np.arange(le_action_cfg[subkey].start, le_action_cfg[subkey].end) + filtered_data = data_array[:, le_indices] + delta_position_values = filtered_data.tolist() + break + except Exception: + continue + + # If combined fields not found, try individual x,y,z coordinates + if delta_position_values is None: + x_data, y_data, z_data = None, None, None + for coord in coordinate_candidates: + full_key = f"action.{coord}" + if full_key in action_keys: + try: + le_action_cfg = self.lerobot_modality_meta.action + if coord in le_action_cfg: + le_key = le_action_cfg[coord].original_key or coord + if le_key in data.columns: + data_array = np.stack(data[le_key]) + le_indices = np.arange(le_action_cfg[coord].start, le_action_cfg[coord].end) + coord_data = data_array[:, le_indices].flatten() + if coord == 'x': + x_data = coord_data + elif coord == 'y': + y_data = coord_data + elif coord == 'z': + z_data = coord_data + except Exception: + continue + + if x_data is not None and y_data is not None and z_data is not None: + delta_position_values = np.column_stack((x_data, y_data, z_data)).tolist() + + if delta_position_values is None: + # Fallback to the old hardcoded approach if metadata approach fails + if 'action.delta_eef_position' in data.columns: + delta_position_values = data['action.delta_eef_position'].to_numpy().tolist() + elif all(col in data.columns for col in ['action.x', 'action.y', 'action.z']): + x_vals = data['action.x'].to_numpy() + y_vals = data['action.y'].to_numpy() + z_vals = data['action.z'].to_numpy() + delta_position_values = np.column_stack((x_vals, y_vals, z_vals)).tolist() + else: + raise ValueError(f"No suitable position columns found. Available columns: {data.columns.tolist()}") + + # Extract gripper data + gripper_values = None + gripper_candidates = ['gripper_close', 'gripper'] + + for grip_key in gripper_candidates: + full_key = f"action.{grip_key}" + if full_key in action_keys: + try: + le_action_cfg = self.lerobot_modality_meta.action + if grip_key in le_action_cfg: + le_key = le_action_cfg[grip_key].original_key or grip_key + if le_key in data.columns: + data_array = np.stack(data[le_key]) + le_indices = np.arange(le_action_cfg[grip_key].start, le_action_cfg[grip_key].end) + gripper_data = data_array[:, le_indices].flatten() + gripper_values = gripper_data.tolist() + break + except Exception: + continue + + if gripper_values is None: + # Fallback to the old hardcoded approach if metadata approach fails + if 'action.gripper_close' in data.columns: + gripper_values = data['action.gripper_close'].to_numpy().tolist() + elif 'action.gripper' in data.columns: + gripper_values = data['action.gripper'].to_numpy().tolist() + else: + raise ValueError(f"No suitable gripper columns found. Available columns: {data.columns.tolist()}") + + return delta_position_values, gripper_values + + def _get_modality_keys(self) -> dict: + """Get the modality keys for the dataset. + The keys are the modality names, and the values are the keys for each modality. + See property `modality_keys` for the expected format. + """ + modality_keys = defaultdict(list) + for modality, config in self.modality_configs.items(): + modality_keys[modality] = config.modality_keys + return modality_keys + + def _get_delta_indices(self) -> dict[str, np.ndarray]: + """Restructure the delta indices to use modality.key as keys instead of just the modalities.""" + delta_indices: dict[str, np.ndarray] = {} + for config in self.modality_configs.values(): + for key in config.modality_keys: + delta_indices[key] = np.array(config.delta_indices) + return delta_indices + + def _init_action_mode(self) -> None: + if self.data_cfg is None: + self._action_mode = "abs" + return + + action_mode = self.data_cfg.get("action_mode", "abs") + if action_mode is None: + action_mode = "abs" + action_mode = _normalize_action_mode(action_mode) + if action_mode not in {"abs", "delta", "rel"}: + raise ValueError(f"Invalid action_mode: {action_mode}. Expected one of: abs, delta, rel.") + self._action_mode = action_mode + + apply_keys = _normalize_action_mode_apply_keys(self.data_cfg.get("action_mode_apply_keys", None)) + if apply_keys: + self._action_mode_apply_keys = apply_keys + + self._action_mode_state_map = _normalize_action_mode_state_map( + self.data_cfg.get("action_mode_state_map", {}) or {} + ) + + def _infer_state_key_for_action(self, action_key: str) -> str | None: + if action_key in self._action_mode_state_map: + return self._action_mode_state_map[action_key] + + if not action_key.startswith("action."): + return None + base = action_key.replace("action.", "", 1) + if f"state.{base}" in self.modality_keys.get("state", []): + return f"state.{base}" + return None + + def _apply_action_mode(self, data: dict) -> dict: + if self._action_mode in (None, "abs"): + return data + + action_keys = self._action_mode_apply_keys or self.modality_keys.get("action", []) + for action_key in action_keys: + if action_key not in data: + print(f"[WARNING] Action key {action_key} not found in data") + continue + state_key = self._infer_state_key_for_action(action_key) + + # for safety, check if the state key is valid + if state_key is None or state_key not in data: + continue + + action_values = np.asarray(data[action_key]) + state_values = np.asarray(data[state_key]) + if action_values.ndim != 2 or state_values.ndim != 2: + raise ValueError( + f"Expected 2D arrays for action/state, got {action_key}: {action_values.shape}, {state_key}: {state_values.shape}" + ) + if action_values.shape[1] != state_values.shape[1]: + raise ValueError( + f"Action/state dim mismatch for {action_key} vs {state_key}: {action_values.shape} vs {state_values.shape}" + ) + + state0 = state_values[0] + if self._action_mode == "delta": + out = action_values.copy() + if len(out) > 1: + out[1:] = action_values[1:] - action_values[:-1] + out[0] = action_values[0] - state0 + elif self._action_mode == "rel": + out = action_values - state0 + else: + out = action_values + + data[action_key] = out + + return data + + def _get_lerobot_modality_meta(self) -> LeRobotModalityMetadata: + """Get the metadata for the LeRobot dataset.""" + modality_meta_path = self.dataset_path / LE_ROBOT_MODALITY_FILENAME + assert ( + modality_meta_path.exists() + ), f"Please provide a {LE_ROBOT_MODALITY_FILENAME} file in {self.dataset_path}" + with open(modality_meta_path, "r") as f: + modality_meta = LeRobotModalityMetadata.model_validate(json.load(f)) + return modality_meta + + def _get_lerobot_info_meta(self) -> dict: + """Get the metadata for the LeRobot dataset.""" + info_meta_path = self.dataset_path / LE_ROBOT_INFO_FILENAME + with open(info_meta_path, "r") as f: + info_meta = json.load(f) + return info_meta + + def _get_data_path_pattern(self) -> str: + """Get the data path pattern for the LeRobot dataset.""" + return self.lerobot_info_meta["data_path"] + + def _get_video_path_pattern(self) -> str: + """Get the video path pattern for the LeRobot dataset.""" + return self.lerobot_info_meta["video_path"] + + def _get_chunk_size(self) -> int: + """Get the chunk size for the LeRobot dataset.""" + return self.lerobot_info_meta["chunks_size"] + + def _get_tasks(self) -> pd.DataFrame: + """Get the tasks for the dataset.""" + if self._lerobot_version == "v2.0": + tasks_path = self.dataset_path / LE_ROBOT_TASKS_FILENAME + with open(tasks_path, "r") as f: + tasks = [json.loads(line) for line in f] + df = pd.DataFrame(tasks) + return df.set_index("task_index") + + elif self._lerobot_version == "v3.0": + tasks_path = self.dataset_path / LE_ROBOT3_TASKS_FILENAME + df = pd.read_parquet(tasks_path) + df = df.reset_index() # convert index to a column, typically named 'index' + df = df.rename(columns={'index': 'task'}) # rename 'index' column to 'task' + df = df[['task_index', 'task']] # reorder columns + return df + def _check_integrity(self): + """Use the config to check if the keys are valid and detect silent data corruption.""" + ERROR_MSG_HEADER = f"Error occurred in initializing dataset {self.dataset_name}:\n" + + for modality_config in self.modality_configs.values(): + for key in modality_config.modality_keys: + if key == "lapa_action" or key == "dream_actions": + continue # no need for any metadata for lapa actions because it comes normalized + # Check if the key is valid + try: + self.lerobot_modality_meta.get_key_meta(key) + except Exception as e: + raise ValueError( + ERROR_MSG_HEADER + f"Unable to find key {key} in modality metadata:\n{e}" + ) + + def set_transforms_metadata(self, metadata: DatasetMetadata): + """Set the metadata for the transforms. This is useful for transforms that need to know the metadata, such as the normalization values.""" + self.transforms.set_metadata(metadata) + + def set_epoch(self, epoch: int): + """Set the epoch for the dataset. + + Args: + epoch (int): The epoch to set. + """ + self.epoch = epoch + + def __len__(self) -> int: + """Get the total number of data points in the dataset. + + Returns: + int: the total number of data points in the dataset. + """ + return len(self.all_steps) + + def __str__(self) -> str: + """Get the description of the dataset.""" + return f"{self.dataset_name} ({len(self)} steps)" + + + def __getitem__(self, index: int) -> dict: + """Get the data for a single step in a trajectory. + + Args: + index (int): The index of the step to get. + + Returns: + dict: The data for the step. + """ + trajectory_id, base_index = self.all_steps[index] + raw_data = self.get_step_data(trajectory_id, base_index) + data = self.transforms(raw_data) + return self._pack_sample(data) + + def _pack_sample(self, data: dict) -> dict: + """Pack transformed modality data into training sample format.""" + step_images = [] + for video_key in self.modality_keys["video"]: + image = data[video_key][0] + image = Image.fromarray(image).resize((224, 224)) + step_images.append(image) + + language = data[self.modality_keys["language"][0]][0] + action = [] + for action_key in self.modality_keys["action"]: + action.append(data[action_key]) + action = np.concatenate(action, axis=1).astype(np.float16) + + sample = { + "action": action, + "image": step_images, + "lang": language, + "robot_tag": self.tag + } + + if self.data_cfg is not None and self.data_cfg.get("include_state", False) not in ["False", False]: + state = [] + for state_key in self.modality_keys.get("state", []): + state.append(data[state_key]) + if not state: + import warnings + warnings.warn( + "include_state=True but no state modality keys found in modality_configs " + "(state modality may be disabled in the dataset's DataConfig). " + "Skipping state packing.", + stacklevel=2, + ) + else: + state = np.concatenate(state, axis=1).astype(np.float16) + sample["state"] = state + + return sample + + def get_step_data(self, trajectory_id: int, base_index: int) -> dict: + """Get the RAW data for a single step in a trajectory. No transforms are applied. + + Args: + trajectory_id (int): The name of the trajectory. + base_index (int): The base step index in the trajectory. + + Returns: + dict: The RAW data for the step. + + Example return: + { + "video": { + "video.image_side_0": [B, T, H, W, C], + "video.image_side_1": [B, T, H, W, C], + }, + "state": { + "state.eef_position": [B, T, state_dim], + "state.eef_rotation": [B, T, state_dim], + }, + "action": { + "action.eef_position": [B, T, action_dim], + "action.eef_rotation": [B, T, action_dim], + }, + } + """ + data = {} + # Get the data for all modalities # just for action base data + self.curr_traj_data = self.get_trajectory_data(trajectory_id) + # TODO @JinhuiYE The logic below is poorly implemented. Data reading should be directly based on curr_traj_data. + for modality in self.modality_keys: + # Get the data corresponding to each key in the modality + for key in self.modality_keys[modality]: + data[key] = self.get_data_by_modality(trajectory_id, modality, key, base_index) + data = self._apply_action_mode(data) + return data + + def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame: + """Get the data for a trajectory.""" + if self._lerobot_version == "v2.0": + + if self.curr_traj_id == trajectory_id and self.curr_traj_data is not None: + return self.curr_traj_data + else: + chunk_index = self.get_episode_chunk(trajectory_id) + parquet_path = self.dataset_path / self.data_path_pattern.format( + episode_chunk=chunk_index, episode_index=trajectory_id + ) + assert parquet_path.exists(), f"Parquet file not found at {parquet_path}" + return pd.read_parquet(parquet_path) + elif self._lerobot_version == "v3.0": + return self.get_trajectory_data_lerobot_v3(trajectory_id) + + def get_trajectory_data_lerobot_v3(self, trajectory_id: int) -> pd.DataFrame: + """Get the data for a trajectory from lerobot v3.""" + if self.curr_traj_id == trajectory_id and self.curr_traj_data is not None: + return self.curr_traj_data + else: #TODO check detail later + episode_meta = self.trajectory_ids_to_metadata[trajectory_id] + chunk_index = episode_meta["data/chunk_index"] + file_index = self.get_episode_file_index(trajectory_id) + # file_from_index = self.get_episode_file_from_index(trajectory_id) + + + parquet_path = self.dataset_path / self.data_path_pattern.format( + chunk_index=chunk_index, file_index=file_index + ) + assert parquet_path.exists(), f"Parquet file not found at {parquet_path}" + file_data = pd.read_parquet(parquet_path) + + # filter by trajectory_id + episode_data = file_data.loc[file_data["episode_index"] == trajectory_id].copy() + return episode_data + + + def get_trajectory_index(self, trajectory_id: int) -> int: + """Get the index of the trajectory in the dataset by the trajectory ID. + This is useful when you need to get the trajectory length or sampling weight corresponding to the trajectory ID. + + Args: + trajectory_id (str): The ID of the trajectory. + + Returns: + int: The index of the trajectory in the dataset. + """ + trajectory_indices = np.where(self.trajectory_ids == trajectory_id)[0] + if len(trajectory_indices) != 1: + raise ValueError( + f"Error finding trajectory index for {trajectory_id}, found {trajectory_indices=}" + ) + return trajectory_indices[0] + + def get_episode_chunk(self, ep_index: int) -> int: + """Get the chunk index for an episode index.""" + return ep_index // self.chunk_size + def get_episode_file_index(self, ep_index: int) -> int: + """Get the file index for an episode index.""" + episode_meta = self.trajectory_ids_to_metadata[ep_index] + return episode_meta["data/file_index"] + + def get_episode_file_from_index(self, ep_index: int) -> int: + """Get the file from index for an episode index.""" + episode_meta = self.trajectory_ids_to_metadata[ep_index] + return episode_meta["data/file_from_index"] + + + def retrieve_data_and_pad( + self, + array: np.ndarray, + step_indices: np.ndarray, + max_length: int, + padding_strategy: str = "first_last", + ) -> np.ndarray: + """Retrieve the data from the dataset and pad it if necessary. + Args: + array (np.ndarray): The array to retrieve the data from. + step_indices (np.ndarray): The step indices to retrieve the data for. + max_length (int): The maximum length of the data. + padding_strategy (str): The padding strategy, either "first" or "last". + """ + # Get the padding indices + front_padding_indices = step_indices < 0 + end_padding_indices = step_indices >= max_length + padding_positions = np.logical_or(front_padding_indices, end_padding_indices) + # Retrieve the data with the non-padding indices + # If there exists some padding, Given T step_indices, the shape of the retrieved data will be (T', ...) where T' < T + raw_data = array[step_indices[~padding_positions]] + assert isinstance(raw_data, np.ndarray), f"{type(raw_data)=}" + # This is the shape of the output, (T, ...) + if raw_data.ndim == 1: + expected_shape = (len(step_indices),) + else: + expected_shape = (len(step_indices), *array.shape[1:]) + + # Pad the data + output = np.zeros(expected_shape) + # Assign the non-padded data + output[~padding_positions] = raw_data + # If there exists some padding, pad the data + if padding_positions.any(): + if padding_strategy == "first_last": + # Use first / last step data to pad + front_padding_data = array[0] + end_padding_data = array[-1] + output[front_padding_indices] = front_padding_data + output[end_padding_indices] = end_padding_data + elif padding_strategy == "zero": + # Use zero padding + output[padding_positions] = 0 + else: + raise ValueError(f"Invalid padding strategy: {padding_strategy}") + return output + + def get_video_path(self, trajectory_id: int, key: str) -> Path: + chunk_index = self.get_episode_chunk(trajectory_id) + original_key = self.lerobot_modality_meta.video[key].original_key + if original_key is None: + original_key = key + if self._lerobot_version == "v2.0": + video_filename = self.video_path_pattern.format( + episode_chunk=chunk_index, episode_index=trajectory_id, video_key=original_key + ) + elif self._lerobot_version == "v3.0": + episode_meta = self.trajectory_ids_to_metadata[trajectory_id] + + video_file_indices = episode_meta.get("videos/file_indices", {}) + # print(f"{video_file_indices=}") + #已修改的lerobotv3.0的视频索引 + if original_key in video_file_indices: + video_chunk_index = video_file_indices[original_key]["chunk_index"] + video_file_index = video_file_indices[original_key]["file_index"] + else: + video_chunk_index = episode_meta["data/chunk_index"] + video_file_index = episode_meta["data/file_index"] + video_filename = self.video_path_pattern.format( + video_key=original_key, + chunk_index=episode_meta["data/chunk_index"], + file_index=episode_meta["data/file_index"], + ) + return self.dataset_path / video_filename + + def get_video( + self, + trajectory_id: int, + key: str, + base_index: int, + ) -> np.ndarray: + """Get the video frames for a trajectory by a base index. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (str): The ID of the trajectory. + key (str): The key of the video. + base_index (int): The base index of the trajectory. + + Returns: + np.ndarray: The video frames for the trajectory and frame indices. Shape: (T, H, W, C) + """ + # Get the step indices + step_indices = self.delta_indices[key] + base_index + # print(f"{step_indices=}") + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Ensure the indices are within the valid range + # This is equivalent to padding the video with extra frames at the beginning and end + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1) + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + # Get the sub-key + key = key.replace("video.", "") + + # Image-only LeRobot datasets (e.g. VLA-Arena) may store frames directly + # in parquet columns and have total_videos == 0 (no mp4 files). In this + # case, load frames from the original image column and pad by step indices. + original_key = self.lerobot_modality_meta.video[key].original_key + if original_key is None: + original_key = key + if self.curr_traj_data is not None and original_key in self.curr_traj_data.columns: + image_entries = self.curr_traj_data[original_key].tolist() + + def _decode_image_entry(entry): + if isinstance(entry, np.ndarray): + return entry + if isinstance(entry, Image.Image): + return np.array(entry) + if isinstance(entry, dict): + img_bytes = entry.get("bytes", None) + img_path = entry.get("path", None) + + if img_bytes is not None: + return np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB")) + + if img_path is not None: + path_obj = Path(img_path) + if not path_obj.is_absolute(): + path_obj = self.dataset_path / path_obj + return np.array(Image.open(path_obj).convert("RGB")) + + raise TypeError(f"Unsupported image entry type: {type(entry)}") + + frames = [] + for idx in step_indices: + safe_idx = int(min(max(idx, 0), len(image_entries) - 1)) + frames.append(_decode_image_entry(image_entries[safe_idx])) + + return np.stack(frames) + + video_path = self.get_video_path(trajectory_id, key) + # Get the action/state timestamps for each frame in the video + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + assert "timestamp" in self.curr_traj_data.columns, f"No timestamp found in {trajectory_id=}" + timestamp: np.ndarray = self.curr_traj_data["timestamp"].to_numpy() + # Get the corresponding video timestamps from the step indices + video_timestamp = timestamp[step_indices] + if self._lerobot_version == "v3.0": + episode_meta = self.trajectory_ids_to_metadata.get(trajectory_id, {}) + from_timestamps = episode_meta.get("videos/from_timestamps", {}) + original_video_key = self.lerobot_modality_meta.video[key].original_key + if original_video_key is None: + original_video_key = key + from_timestamp = float(from_timestamps.get(original_video_key, 0.0)) + video_timestamp = video_timestamp + from_timestamp + + return get_frames_by_timestamps( + video_path.as_posix(), + video_timestamp, + video_backend=self.video_backend, # TODO + video_backend_kwargs=self.video_backend_kwargs, + ) + + def get_state_or_action( + self, + trajectory_id: int, + modality: str, + key: str, + base_index: int, + ) -> np.ndarray: + """Get the state or action data for a trajectory by a base index. + If the step indices are out of range, pad with the data: + if the data is stored in absolute format, pad with the first or last step data; + otherwise, pad with zero. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data. + key (str): The key of the data. + base_index (int): The base index of the trajectory. + + Returns: + np.ndarray: The data for the trajectory and step indices. + """ + # Get the step indices + step_indices = self.delta_indices[key] + base_index + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + assert key.startswith(modality + "."), f"{key} must start with {modality + '.'}, got {key}" + # Get the sub-key, e.g. state.joint_angles -> joint_angles + key = key.replace(modality + ".", "") + # Get the lerobot key + le_state_or_action_cfg = getattr(self.lerobot_modality_meta, modality) + le_key = le_state_or_action_cfg[key].original_key + if le_key is None: + le_key = key + # Get the data array, shape: (T, D) + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + assert le_key in self.curr_traj_data.columns, f"No {le_key} found in {trajectory_id=}" + data_array: np.ndarray = np.stack(self.curr_traj_data[le_key]) # type: ignore + assert data_array.ndim == 2, f"Expected 2D array, got key {le_key} is{data_array.shape} array" + le_indices = np.arange( + le_state_or_action_cfg[key].start, + le_state_or_action_cfg[key].end, + ) + data_array = data_array[:, le_indices] + # Get the state or action configuration + state_or_action_cfg = getattr(self.metadata.modalities, modality)[key] + + # Pad the data + return self.retrieve_data_and_pad( + array=data_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last" if state_or_action_cfg.absolute else "zero", + # padding_strategy="zero", # HACK for realdata + ) + + def get_language( + self, + trajectory_id: int, + key: str, + base_index: int, + ) -> list[str]: + """Get the language annotation data for a trajectory by step indices. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (int): The ID of the trajectory. + key (str): The key of the annotation. + base_index (int): The base index of the trajectory. + + Returns: + list[str]: The annotation data for the trajectory and step indices. If no matching data is found, return empty strings. + """ + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + # Get the step indices + step_indices = self.delta_indices[key] + base_index + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + # Get the end times corresponding to the closest indices + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, max_length - 1) + # Get the annotations + task_indices: list[int] = [] + assert key.startswith( + "annotation." + ), f"Language key must start with 'annotation.', got {key}" + subkey = key.replace("annotation.", "") + annotation_meta = self.lerobot_modality_meta.annotation + assert annotation_meta is not None, f"Annotation metadata is None for {subkey}" + assert ( + subkey in annotation_meta + ), f"Annotation key {subkey} not found in metadata, available annotation keys: {annotation_meta.keys()}" + subkey_meta = annotation_meta[subkey] + original_key = subkey_meta.original_key + if original_key is None: + original_key = key + for i in range(len(step_indices)): # + # task_indices.append(self.curr_traj_data[original_key][step_indices[i]].item()) + value = self.curr_traj_data[original_key].iloc[step_indices[i]] # TODO check v2.0 + task_indices.append(value if isinstance(value, (int, float)) else value.item()) + + return self.tasks.loc[task_indices]["task"].tolist() + + def get_data_by_modality( + self, + trajectory_id: int, + modality: str, + key: str, + base_index: int, + ): + """Get the data corresponding to the modality for a trajectory by a base index. + This method will call the corresponding helper method based on the modality. + See the helper methods for more details. + NOTE: For the language modality, the data is padded with empty strings if no matching data is found. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data. + key (str): The key of the data. + base_index (int): The base index of the trajectory. + """ + if modality == "video": + return self.get_video(trajectory_id, key, base_index) + elif modality == "state" or modality == "action": + return self.get_state_or_action(trajectory_id, modality, key, base_index) + elif modality == "language": + return self.get_language(trajectory_id, key, base_index) + else: + raise ValueError(f"Invalid modality: {modality}") + + def _save_dataset_statistics_(self, save_path: Path | str, format: str = "json") -> None: + """ + Save dataset statistics to specified path in the required format. + Only includes statistics for keys that are actually used in the dataset. + Gripper-related keys will be placed at the end. + + Args: + save_path (Path | str): Path to save the statistics file + format (str): Save format, currently only supports "json" + """ + save_path = Path(save_path) + save_path.parent.mkdir(parents=True, exist_ok=True) + + # Build the data structure to save + statistics_data = {} + + # Get used modality keys + used_action_keys, used_state_keys = get_used_modality_keys(self.modality_keys) + + # Organize statistics by tag + tag = self.tag + tag_stats = {} + + # Process action statistics (only for used keys) + if hasattr(self.metadata.statistics, 'action') and self.metadata.statistics.action: + action_stats = self.metadata.statistics.action + + # Filter to only include used action keys and reorder: non-gripper first, gripper last + non_gripper_keys = [] + gripper_keys = [] + + for key in action_stats.keys(): + if key in used_action_keys: + if "gripper" in key.lower(): + gripper_keys.append(key) + else: + non_gripper_keys.append(key) + + # Reorder: non-gripper first, gripper last + reordered_keys = non_gripper_keys + gripper_keys + + filtered_action_stats = {} + for key in reordered_keys: + filtered_action_stats[key] = action_stats[key] + + if filtered_action_stats: + # Combine statistics from filtered action sub-keys + combined_action_stats = combine_modality_stats(filtered_action_stats) + + # mask=False for dimensions whose normalization mode is "binary" + _action_norm_modes = _extract_action_normalization_modes(self.transforms) + mask = generate_action_mask_for_used_keys( + self.metadata.modalities.action, filtered_action_stats.keys(), + normalization_modes=_action_norm_modes, + ) + combined_action_stats["mask"] = mask + + tag_stats["action"] = combined_action_stats + + # Process state statistics (only for used keys) + if hasattr(self.metadata.statistics, 'state') and self.metadata.statistics.state: + state_stats = self.metadata.statistics.state + + # Filter to only include used state keys, optionally reorder gripper to end + non_gripper_keys = [] + gripper_keys = [] + + for key in state_stats.keys(): + if key in used_state_keys: + if "gripper" in key.lower(): + gripper_keys.append(key) + else: + non_gripper_keys.append(key) + + # Reorder: non-gripper first, gripper last + reordered_keys = non_gripper_keys + gripper_keys + + filtered_state_stats = {} + for key in reordered_keys: + filtered_state_stats[key] = state_stats[key] + + if filtered_state_stats: + combined_state_stats = combine_modality_stats(filtered_state_stats) + tag_stats["state"] = combined_state_stats + + # Add dataset counts + tag_stats["num_transitions"] = len(self) + tag_stats["num_trajectories"] = len(self.trajectory_ids) + + statistics_data[tag] = tag_stats + + # Save as JSON file + if format.lower() == "json": + if not str(save_path).endswith('.json'): + save_path = save_path.with_suffix('.json') + with open(save_path, 'w', encoding='utf-8') as f: + json.dump(statistics_data, f, indent=2, ensure_ascii=False) + else: + raise ValueError(f"Unsupported format: {format}. Currently only 'json' is supported.") + + print(f"Single dataset statistics saved to: {save_path}") + print(f"Used action keys (reordered): {list(used_action_keys)}") + print(f"Used state keys (reordered): {list(used_state_keys)}") + + +class CachedLeRobotSingleDataset(LeRobotSingleDataset): + def __init__(self, img_resize: tuple[int, int] | None = None, *args, **kwargs): + """ + This class caches the video frames for each trajectory and key. + It is recommended to use this class if the video frames need to be accessed multiple times. + + Args: + resize_img (tuple[int, int], optional): The size to resize the video frames to reduce memory usage. + """ + # Convert img_resize to tuple if it is not already + if img_resize is not None and not isinstance(img_resize, tuple): + img_resize = tuple(img_resize) + assert len(img_resize) == 2, f"Expected tuple of length 2, got {img_resize}" + self.img_resize = img_resize + + # Initialize img_resize attribute first to ensure it exists + super().__init__(*args, **kwargs) + cached_frames: dict[str, np.ndarray] = {} + + for key in self.modality_keys["video"]: + all_frames = [] + original_key = key + key = key.replace("video.", "") + for trajectory_id, trajectory_length in tqdm( + zip(self.trajectory_ids, self.trajectory_lengths), + total=len(self.trajectory_ids), + desc=f"Caching {key} frames", + ): + video_path = self.get_video_path(trajectory_id, key) + frames = get_all_frames( + video_path.as_posix(), + video_backend=self.video_backend, + video_backend_kwargs=self.video_backend_kwargs, + resize_size=img_resize, + ) + assert frames.ndim == 4, f"Expected 4D array, got {frames.shape} array" + assert frames.shape[3] == 3, f"Expected 3 channels, got {frames.shape[3]} channels" + + # Apply image cropping if enabled and the video key is base_view + # Note: crop_obs_camera functionality has been removed + + # assert ( + # frames.shape[0] == trajectory_length + # ), f"Expected {trajectory_length} frames, got {frames.shape[0]} frames" + all_frames.append(frames) + cached_frames[key] = np.concatenate(all_frames, axis=0) + print(f"{key}: {cached_frames[key].shape}") + self.cached_frames = cached_frames + self.start_indices = np.cumsum(self.trajectory_lengths) - self.trajectory_lengths + + def get_video(self, trajectory_id: int, key: str, base_index: int) -> np.ndarray: + step_indices = self.delta_indices[key] + base_index + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Ensure the indices are within the valid range + # This is equivalent to padding the video with extra frames at the beginning and end + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1) + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + # Get the sub-key + key = key.replace("video.", "") + # Calculate the absolute indices + absolute_indices = self.start_indices[trajectory_index] + step_indices + return self.cached_frames[key][absolute_indices] + + def get_step_data(self, trajectory_id: int, base_index: int) -> dict: + """Get the RAW data for a single step. No transforms are applied. + + Args: + trajectory_id (str): The ID of the trajectory. + base_index (int): The base index of the step. + + Returns: + dict: The data for the step. + """ + data = {} + self.curr_traj_data = self.get_trajectory_data(trajectory_id) + # Get the data for all modalities + for modality in self.modality_keys: + # Get the data corresponding to each key in the modality + for key in self.modality_keys[modality]: + data[key] = self.get_data_by_modality(trajectory_id, modality, key, base_index) + return data + + def set_transforms_metadata(self, metadata: DatasetMetadata): + """Set the metadata for the transforms. This is useful for transforms that need to know the metadata, such as the normalization values.""" + if self.img_resize is not None: + all_video_keys = [key for key in self.modality_keys["video"]] + for key in metadata.modalities.video: + if key in all_video_keys: + metadata.modalities.video[key].resolution = self.img_resize + super().set_transforms_metadata(metadata) + + +def safe_hash(input_tuple): + # keep 128 bits of the hash + tuple_string = repr(input_tuple).encode("utf-8") + sha256 = hashlib.sha256() + sha256.update(tuple_string) + + seed = int(sha256.hexdigest(), 16) + + return seed & 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF + + +class MixtureSpecElement(BaseModel): + dataset_path: list[Path] | Path = Field(..., description="The path to the dataset.") + dataset_weight: float = Field(..., description="The weight of the dataset in the mixture.") + distribute_weights: bool = Field( + default=False, + description="Whether to distribute the weights of the dataset across all the paths. If True, the weights will be evenly distributed across all the paths.", + ) + + +# Helper functions for dataset statistics + +def combine_modality_stats(modality_stats: dict) -> dict: + """ + Combine statistics from all sub-keys under a modality. + + Args: + modality_stats (dict): Statistics for a modality, containing multiple sub-keys. + Each sub-key contains DatasetStatisticalValues object. + + Returns: + dict: Combined statistics + """ + combined_stats = { + "mean": [], + "std": [], + "max": [], + "min": [], + "q01": [], + "q99": [] + } + + # Combine statistics in sub-key order + for subkey in modality_stats.keys(): + subkey_stats = modality_stats[subkey] # This is a DatasetStatisticalValues object + + # Convert DatasetStatisticalValues to dict-like access + for stat_name in ["mean", "std", "max", "min", "q01", "q99"]: + stat_value = getattr(subkey_stats, stat_name) + if isinstance(stat_value, (list, tuple)): + combined_stats[stat_name].extend(stat_value) + else: + # Handle NDArray case - convert to list + if hasattr(stat_value, 'tolist'): + combined_stats[stat_name].extend(stat_value.tolist()) + else: + combined_stats[stat_name].append(float(stat_value)) + + return combined_stats + +def _extract_action_normalization_modes(transforms) -> dict: + """Extract normalization modes for action keys from a ComposedModalityTransform. + + Returns: + dict: {subkey_without_action_prefix -> normalization_mode_str} + """ + modes = {} + for t in transforms.transforms: + if isinstance(t, StateActionTransform): + for key, mode in t.normalization_modes.items(): + if key.startswith("action."): + subkey = key[len("action."):] + modes[subkey] = mode + return modes + + +def generate_action_mask_for_used_keys( + action_modalities: dict, + used_action_keys_ordered, + normalization_modes: dict | None = None, +) -> list[bool]: + """Generate per-dimension mask for action statistics. + + A dimension gets ``mask=False`` only when its normalization mode is ``"binary"``. + This tells the inference code to skip continuous de-normalization for that dimension. + All other modes (q99, mean_std, min_max ...) produce ``mask=True``. + + Args: + action_modalities (dict): Configuration information for action modalities. + used_action_keys_ordered: Iterable of actually used action keys (no "action." prefix). + normalization_modes (dict | None): Mapping {subkey -> mode} (no "action." prefix). + If ``None``, all dimensions default to ``mask=True``. + + Returns: + list[bool]: Per-dimension mask values. + """ + mask = [] + + for subkey in used_action_keys_ordered: + if subkey in action_modalities: + subkey_config = action_modalities[subkey] + + # Get dimension count from shape + if hasattr(subkey_config, 'shape') and len(subkey_config.shape) > 0: + dim_count = subkey_config.shape[0] + else: + dim_count = 1 + + # mask=False only when the normalization mode is explicitly "binary" + is_binary = ( + normalization_modes.get(subkey) == "binary" + if normalization_modes is not None + else False + ) + + for _ in range(dim_count): + mask.append(not is_binary) + + return mask + +def get_used_modality_keys(modality_keys: dict) -> tuple[list, list]: + """Extract used action and state keys from modality configuration.""" + used_action_keys = [] + used_state_keys = [] + + # Extract action keys (remove "action." prefix) + for action_key in modality_keys.get("action", []): + if action_key.startswith("action."): + clean_key = action_key.replace("action.", "") + used_action_keys.append(clean_key) + + # Extract state keys (remove "state." prefix) + for state_key in modality_keys.get("state", []): + if state_key.startswith("state."): + clean_key = state_key.replace("state.", "") + used_state_keys.append(clean_key) + + return used_action_keys, used_state_keys + +class LeRobotMixtureDataset(Dataset): + """ + A mixture of multiple datasets. This class samples a single dataset based on the dataset weights and then calls the `__getitem__` method of the sampled dataset. + It is recommended to modify the single dataset class instead of this class. + """ + + def __init__( + self, + data_mixture: Sequence[tuple[LeRobotSingleDataset, float]], + mode: str, + balance_dataset_weights: bool = True, + balance_trajectory_weights: bool = True, + seed: int = 42, + metadata_config: dict = { + "percentile_mixing_method": "min_max", + }, + **kwargs, + ): + """ + Initialize the mixture dataset. + + Args: + data_mixture (list[tuple[LeRobotSingleDataset, float]]): Datasets and their corresponding weights. + mode (str): If "train", __getitem__ will return different samples every epoch; if "val" or "test", __getitem__ will return the same sample every epoch. + balance_dataset_weights (bool): If True, the weight of dataset will be multiplied by the total trajectory length of each dataset. + balance_trajectory_weights (bool): If True, sample trajectories within a dataset weighted by their length; otherwise, use equal weighting. + seed (int): Random seed for sampling. + """ + datasets: list[LeRobotSingleDataset] = [] + dataset_sampling_weights: list[float] = [] + for dataset, weight in data_mixture: + # Check if dataset is valid and has data + if len(dataset) == 0: + print(f"Warning: Skipping empty dataset {dataset.dataset_name}") + continue + datasets.append(dataset) + dataset_sampling_weights.append(weight) + + if len(datasets) == 0: + raise ValueError("No valid datasets found in the mixture. All datasets are empty.") + + self.datasets = datasets + self.balance_dataset_weights = balance_dataset_weights + self.balance_trajectory_weights = balance_trajectory_weights + self.seed = seed + self.mode = mode + self.data_cfg = kwargs["data_cfg"] if "data_cfg" in kwargs else None + self._traceable_sampling = bool( + self.data_cfg is not None and self.data_cfg.get("traceable_sampling", False) + ) + self.valid_anchor_manifest = None + self.temporal_proposal = None + + # Set properties for sampling + + # 1. Dataset lengths + self._dataset_lengths = np.array([len(dataset) for dataset in self.datasets]) + print(f"Dataset lengths: {self._dataset_lengths}") + self._getitem_count = 0 + # 2. Dataset sampling weights + self._dataset_sampling_weights = np.array(dataset_sampling_weights) + + if self.balance_dataset_weights: + self._dataset_sampling_weights *= self._dataset_lengths + + # Check for zero or negative weights before normalization + if np.any(self._dataset_sampling_weights <= 0): + print(f"Warning: Found zero or negative sampling weights: {self._dataset_sampling_weights}") + # Set minimum weight to prevent division issues + self._dataset_sampling_weights = np.maximum(self._dataset_sampling_weights, 1e-8) + + # Normalize weights + weights_sum = self._dataset_sampling_weights.sum() + if weights_sum == 0 or np.isnan(weights_sum): + print(f"Error: Invalid weights sum: {weights_sum}") + # Fallback to equal weights + self._dataset_sampling_weights = np.ones(len(self.datasets)) / len(self.datasets) + print(f"Fallback to equal weights") + else: + self._dataset_sampling_weights /= weights_sum + + # 3. Trajectory sampling weights + self._trajectory_sampling_weights: list[np.ndarray] = [] + for i, dataset in enumerate(self.datasets): + trajectory_sampling_weights = np.ones(len(dataset.trajectory_lengths)) + if self.balance_trajectory_weights: + trajectory_sampling_weights *= dataset.trajectory_lengths + + # Check for zero or negative weights before normalization + if np.any(trajectory_sampling_weights <= 0): + print(f"Warning: Dataset {i} has zero or negative trajectory weights") + trajectory_sampling_weights = np.maximum(trajectory_sampling_weights, 1e-8) + + # Normalize weights + weights_sum = trajectory_sampling_weights.sum() + if weights_sum == 0 or np.isnan(weights_sum): + print(f"Error: Dataset {i} has invalid trajectory weights sum: {weights_sum}") + # Fallback to equal weights + trajectory_sampling_weights = np.ones(len(dataset.trajectory_lengths)) / len(dataset.trajectory_lengths) + else: + trajectory_sampling_weights /= weights_sum + + self._trajectory_sampling_weights.append(trajectory_sampling_weights) + + if self._traceable_sampling: + self.valid_anchor_manifest = self._build_valid_anchor_manifest() + temporal_proposal_path = self.data_cfg.get("temporal_proposal_path") + if temporal_proposal_path: + self.temporal_proposal = TemporalProposal.load( + temporal_proposal_path, self.valid_anchor_manifest + ) + + # 4. Primary dataset indices + self._primary_dataset_indices = np.array(dataset_sampling_weights) == 1.0 + if not np.any(self._primary_dataset_indices): + print(f"Warning: No dataset with weight 1.0 found. Original weights: {dataset_sampling_weights}") + # Fallback: use the dataset(s) with maximum weight as primary + max_weight = max(dataset_sampling_weights) + self._primary_dataset_indices = np.array(dataset_sampling_weights) == max_weight + print(f"Using datasets with maximum weight {max_weight} as primary: {self._primary_dataset_indices}") + + if not np.any(self._primary_dataset_indices): + # This should never happen, but just in case + print("Error: Still no primary dataset found. Using first dataset as primary.") + self._primary_dataset_indices = np.zeros(len(self.datasets), dtype=bool) + self._primary_dataset_indices[0] = True + + # Set the epoch and sample the first epoch + self.set_epoch(0) + + self.update_metadata(metadata_config) + + @property + def dataset_lengths(self) -> np.ndarray: + """The lengths of each dataset.""" + return self._dataset_lengths + + @property + def dataset_sampling_weights(self) -> np.ndarray: + """The sampling weights for each dataset.""" + return self._dataset_sampling_weights + + @property + def trajectory_sampling_weights(self) -> list[np.ndarray]: + """The sampling weights for each trajectory in each dataset.""" + return self._trajectory_sampling_weights + + @property + def primary_dataset_indices(self) -> np.ndarray: + """The indices of the primary datasets.""" + return self._primary_dataset_indices + + def __str__(self) -> str: + dataset_descriptions = [] + for dataset, weight in zip(self.datasets, self.dataset_sampling_weights): + dataset_description = { + "Dataset": str(dataset), + "Sampling weight": float(weight), + } + dataset_descriptions.append(dataset_description) + return json.dumps({"Mixture dataset": dataset_descriptions}, indent=2) + + def set_epoch(self, epoch: int): + """Set the epoch for the dataset. + + Args: + epoch (int): The epoch to set. + """ + self.epoch = epoch + # self.sampled_steps = self.sample_epoch() + + def _build_valid_anchor_manifest(self) -> ValidAnchorManifest: + """Freeze the readable episode/anchor support used by probability-aware training.""" + def serialize_transforms(dataset) -> list[dict]: + serialized = [] + for transform in dataset.transforms.transforms: + try: + config = BaseModel.model_dump(transform, mode="json") + except Exception: + config = BaseModel.model_dump(transform, mode="python") + config = json.loads(json.dumps(config, default=str, sort_keys=True)) + serialized.append( + { + "class": f"{type(transform).__module__}.{type(transform).__qualname__}", + "config": config, + } + ) + return serialized + + all_dataset_episodes = [] + excluded_episodes = [] + for dataset_index, dataset in enumerate(self.datasets): + anchors_by_trajectory = defaultdict(list) + for trajectory_id, base_index in dataset.all_steps: + trajectory_key = trajectory_id.item() if isinstance(trajectory_id, np.generic) else trajectory_id + anchors_by_trajectory[trajectory_key].append(int(base_index)) + + valid_episodes = [] + total_videos = int(dataset.lerobot_info_meta.get("total_videos", 0)) + for trajectory_index, trajectory_id in enumerate(dataset.trajectory_ids): + trajectory_key = trajectory_id.item() if isinstance(trajectory_id, np.generic) else trajectory_id + anchors = anchors_by_trajectory.get(trajectory_key, []) + if not anchors: + excluded_episodes.append( + { + "dataset": dataset.dataset_name, + "trajectory_id": trajectory_key, + "reason": "no_valid_anchors", + } + ) + continue + + if dataset._lerobot_version == "v2.0": + parquet_path = dataset.dataset_path / dataset.data_path_pattern.format( + episode_chunk=dataset.get_episode_chunk(trajectory_key), + episode_index=trajectory_key, + ) + else: + episode_metadata = dataset.trajectory_ids_to_metadata[trajectory_key] + parquet_path = dataset.dataset_path / dataset.data_path_pattern.format( + chunk_index=episode_metadata["data/chunk_index"], + file_index=episode_metadata["data/file_index"], + ) + if not parquet_path.exists(): + excluded_episodes.append( + { + "dataset": dataset.dataset_name, + "trajectory_id": trajectory_key, + "reason": "missing_parquet", + "path": str(parquet_path), + } + ) + continue + try: + import pyarrow.parquet as parquet + + parquet.ParquetFile(parquet_path) + except Exception as exc: + excluded_episodes.append( + { + "dataset": dataset.dataset_name, + "trajectory_id": trajectory_key, + "reason": "unreadable_parquet", + "path": str(parquet_path), + "error": f"{type(exc).__name__}: {exc}", + } + ) + continue + + if total_videos > 0: + video_paths = [ + dataset.get_video_path(trajectory_key, key.replace("video.", "")) + for key in dataset.modality_keys["video"] + ] + if not all(path.exists() for path in video_paths): + excluded_episodes.append( + { + "dataset": dataset.dataset_name, + "trajectory_id": trajectory_key, + "reason": "missing_video", + "paths": [str(path) for path in video_paths if not path.exists()], + } + ) + continue + + valid_episodes.append( + ( + trajectory_key, + tuple(sorted(set(anchors))), + float(self.trajectory_sampling_weights[dataset_index][trajectory_index]), + ) + ) + + if not valid_episodes: + raise RuntimeError( + f"Traceable sampling found no valid episodes in dataset {dataset.dataset_name}" + ) + all_dataset_episodes.append(valid_episodes) + + data_config = {} + if self.data_cfg is not None: + manifest_config_keys = ( + "data_mix", + "delete_pause_frame", + "lerobot_version", + "action_type", + "action_mode", + "action_mode_apply_keys", + "action_mode_state_map", + "include_state", + "video_backend", + "precomputed_statistics_path", + "precomputed_statistics_key_map", + "precomputed_statistics_filename", + ) + try: + from omegaconf import OmegaConf + + resolved_data_config = OmegaConf.to_container(self.data_cfg, resolve=True) + except (TypeError, ValueError): + resolved_data_config = dict(self.data_cfg) + data_config = { + key: resolved_data_config.get(key) + for key in manifest_config_keys + if key in resolved_data_config + } + + dataset_provenance = None + provenance_path = self.data_cfg.get("dataset_provenance_path") if self.data_cfg else None + if not provenance_path: + raise ValueError( + "traceable_sampling requires dataset_provenance_path so the frozen support " + "is bound to exact dataset revisions" + ) + provenance_path = Path(provenance_path) + with open(provenance_path, "r", encoding="utf-8") as provenance_file: + provenance_payload = json.load(provenance_file) + provenance_bytes = provenance_path.read_bytes() + dataset_provenance = { + "sha256": hashlib.sha256(provenance_bytes).hexdigest(), + "payload": provenance_payload, + } + + content_fingerprint_path = ( + self.data_cfg.get("dataset_content_fingerprint_path") if self.data_cfg else None + ) + if not content_fingerprint_path: + raise ValueError( + "traceable_sampling requires dataset_content_fingerprint_path; generate the " + "full content hashes once before launching distributed training" + ) + content_fingerprint_path = Path(content_fingerprint_path) + with open(content_fingerprint_path, "r", encoding="utf-8") as fingerprint_file: + content_fingerprints = json.load(fingerprint_file) + content_fingerprint_artifact_sha256 = hashlib.sha256( + content_fingerprint_path.read_bytes() + ).hexdigest() + + dataset_identities = [] + for dataset in self.datasets: + try: + content_fingerprint = content_fingerprints["datasets"][dataset.dataset_name] + except KeyError as exc: + raise ValueError( + f"No content fingerprint found for dataset {dataset.dataset_name}" + ) from exc + validate_dataset_content_inventory(dataset.dataset_path, content_fingerprint) + dataset_identities.append( + { + "name": dataset.dataset_name, + "info": dataset.lerobot_info_meta, + "metadata": dataset.metadata.model_dump(mode="json"), + "modality": dataset._lerobot_modality_meta.model_dump(mode="json"), + "transforms": serialize_transforms(dataset), + "content": { + "sha256": content_fingerprint["sha256"], + "file_count": content_fingerprint["file_count"], + "total_bytes": content_fingerprint["total_bytes"], + }, + "delta_indices": { + key: value.tolist() for key, value in dataset.delta_indices.items() + }, + } + ) + + identity = { + "data_config": data_config, + "balance_dataset_weights": self.balance_dataset_weights, + "balance_trajectory_weights": self.balance_trajectory_weights, + "excluded_episodes": excluded_episodes, + "dataset_provenance": dataset_provenance, + "content_fingerprint_artifact_sha256": content_fingerprint_artifact_sha256, + "datasets": dataset_identities, + } + identity = json.loads(json.dumps(identity, sort_keys=True, default=str)) + current_manifest = ValidAnchorManifest.from_hierarchy( + dataset_names=[dataset.dataset_name for dataset in self.datasets], + dataset_weights=self.dataset_sampling_weights, + episodes=all_dataset_episodes, + identity=identity, + ) + manifest_path = self.data_cfg.get("valid_anchor_manifest_path") if self.data_cfg else None + if manifest_path: + frozen_manifest = ValidAnchorManifest.load(manifest_path) + frozen_support = frozen_manifest.to_dict(include_hash=False)["datasets"] + current_support = current_manifest.to_dict(include_hash=False)["datasets"] + if frozen_support != current_support: + raise RuntimeError( + "Frozen valid-anchor manifest changes dataset/episode/anchor support" + ) + frozen_identity = json.loads(json.dumps(frozen_manifest.identity)) + current_identity = json.loads(json.dumps(current_manifest.identity)) + for identity in (frozen_identity, current_identity): + data_config = identity.get("data_config", {}) + statistics_path = data_config.get("precomputed_statistics_path") + if statistics_path: + data_config["precomputed_statistics_path"] = Path(statistics_path).name + if frozen_identity != current_identity: + raise RuntimeError( + "Frozen valid-anchor manifest changes the semantic data configuration" + ) + return frozen_manifest + return current_manifest + + def sample_trace(self, index: int | LogicalDraw): + if self.valid_anchor_manifest is None: + raise RuntimeError("sample_trace requires traceable_sampling=True") + draw = index if isinstance(index, LogicalDraw) else LogicalDraw(self.epoch, int(index)) + if self.temporal_proposal is not None: + return self.temporal_proposal.sample(self.valid_anchor_manifest, self.seed, draw) + return self.valid_anchor_manifest.sample(self.seed, draw) + + def sample_step(self, index: int) -> tuple[LeRobotSingleDataset, int, int]: + """Sample a single step from the dataset.""" + # return self.sampled_steps[index] + + if self.valid_anchor_manifest is not None: + trace = self.sample_trace(index) + return self.datasets[trace.dataset_index], trace.trajectory_id, trace.base_index + + # Set seed + seed = index if self.mode != "train" else safe_hash((self.epoch, index, self.seed)) + rng = np.random.default_rng(seed) + + # Sample dataset + dataset_index = rng.choice(len(self.datasets), p=self.dataset_sampling_weights) + dataset = self.datasets[dataset_index] + + # Sample trajectory + trajectory_index = rng.choice( + len(dataset.trajectory_ids), p=self.trajectory_sampling_weights[dataset_index] + ) + trajectory_id = dataset.trajectory_ids[trajectory_index] + + # Sample step + base_index = rng.choice(dataset.trajectory_lengths[trajectory_index]) + return dataset, trajectory_id, base_index + + + + def __getitem__(self, index: int) -> dict: + """Get the data for a single trajectory and start index. + + Args: + index (int): The index of the trajectory to get. + + Returns: + dict: The data for the trajectory and start index. + """ + self._getitem_count += 1 + if self._getitem_count % 1000 == 0: + gc.collect() + + if self.valid_anchor_manifest is not None: + trace = self.sample_trace(index) + dataset = self.datasets[trace.dataset_index] + try: + raw_data = dataset.get_step_data(trace.trajectory_id, trace.base_index) + transform_rng_seed = ( + self.data_cfg.get("transform_rng_seed") + if self.data_cfg is not None + else None + ) + logical_transform_seed = None + if transform_rng_seed is not None: + payload = ( + f"{int(transform_rng_seed)}:{trace.epoch}:" + f"{trace.global_draw_id}" + ).encode("utf-8") + logical_transform_seed = int.from_bytes( + hashlib.sha256(payload).digest()[:8], + "little", + signed=False, + ) + with _isolated_transform_rng(logical_transform_seed): + data = dataset.transforms(raw_data) + sample = dataset._pack_sample(data) + except Exception as exc: + raise RuntimeError( + "A frozen valid-anchor manifest entry failed at runtime; refusing to " + f"change the sampling distribution. trace={trace.to_dict()}" + ) from exc + sample["sampling_trace"] = trace.to_dict() + return sample + + max_retries = 10 + last_exception = None + + for attempt in range(max_retries): + try: + sample_tries = 0 + max_sample_tries = 200 + while True: + sample_tries += 1 + if sample_tries > max_sample_tries: + raise RuntimeError( + f"Unable to sample a valid item after {max_sample_tries} attempts. " + f"dataset={self.datasets[0].dataset_name if len(self.datasets)>0 else 'unknown'}" + ) + + dataset, trajectory_id, step = self.sample_step(index) + # If dataset has no physical videos (e.g., image frames in parquet + # for VLA-Arena), do not gate sampling on mp4 existence. + total_videos = int(dataset.lerobot_info_meta.get("total_videos", 0)) + if total_videos == 0: + break + + key = dataset.modality_keys["video"][0].replace("video.", "") + video_path = dataset.get_video_path(trajectory_id, key) + if os.path.exists(video_path): + break + index = random.randint(0, len(self) - 1) + + raw_data = dataset.get_step_data(trajectory_id, step) + data = dataset.transforms(raw_data) + sample = dataset._pack_sample(data) + + return sample + + except Exception as e: + last_exception = e + if attempt < max_retries - 1: + # Log the error but continue trying + print(f"Attempt {attempt + 1}/{max_retries} failed for index {index}: {e}") + print(f"Retrying with new sample...") + # For retry, we can use a slightly different index to get a new sample + # This helps avoid getting stuck on the same problematic sample + index = random.randint(0, len(self) - 1) + else: + # All retries exhausted + print(f"All {max_retries} attempts failed for index {index}") + print(f"Last error: {last_exception}") + # Return a dummy sample or re-raise the exception + raise last_exception + + def __len__(self) -> int: + """Get the length of a single epoch in the mixture. + + Returns: + int: The length of a single epoch in the mixture. + """ + # Check for potential issues + if len(self.datasets) == 0: + return 0 + + # Check if any dataset lengths are 0 or NaN + if np.any(self.dataset_lengths == 0) or np.any(np.isnan(self.dataset_lengths)): + print(f"Warning: Found zero or NaN dataset lengths: {self.dataset_lengths}") + # Filter out zero/NaN length datasets + valid_indices = (self.dataset_lengths > 0) & (~np.isnan(self.dataset_lengths)) + if not np.any(valid_indices): + print("Error: All datasets have zero or NaN length") + return 0 + else: + valid_indices = np.ones(len(self.datasets), dtype=bool) + + # Check if any sampling weights are 0 or NaN + if np.any(self.dataset_sampling_weights == 0) or np.any(np.isnan(self.dataset_sampling_weights)): + print(f"Warning: Found zero or NaN sampling weights: {self.dataset_sampling_weights}") + # Use only valid weights + valid_weights = (self.dataset_sampling_weights > 0) & (~np.isnan(self.dataset_sampling_weights)) + valid_indices = valid_indices & valid_weights + if not np.any(valid_indices): + print("Error: All sampling weights are zero or NaN") + return 0 + + # Check primary dataset indices + primary_and_valid = self.primary_dataset_indices & valid_indices + if not np.any(primary_and_valid): + print(f"Warning: No valid primary datasets found. Primary indices: {self.primary_dataset_indices}, Valid indices: {valid_indices}") + # Fallback: use the largest valid dataset + if np.any(valid_indices): + max_length = self.dataset_lengths[valid_indices].max() + print(f"Fallback: Using maximum dataset length: {max_length}") + return int(max_length) + else: + return 0 + + # Calculate the ratio and get max + ratios = (self.dataset_lengths / self.dataset_sampling_weights)[primary_and_valid] + + # Check for NaN or inf in ratios + if np.any(np.isnan(ratios)) or np.any(np.isinf(ratios)): + print(f"Warning: Found NaN or inf in ratios: {ratios}") + print(f"Dataset lengths: {self.dataset_lengths[primary_and_valid]}") + print(f"Sampling weights: {self.dataset_sampling_weights[primary_and_valid]}") + # Filter out invalid ratios + valid_ratios = ratios[~np.isnan(ratios) & ~np.isinf(ratios)] + if len(valid_ratios) == 0: + print("Error: All ratios are NaN or inf") + return 0 + max_ratio = valid_ratios.max() + else: + max_ratio = ratios.max() + + result = int(max_ratio) + if result == 0: + print(f"Warning: Dataset mixture length is 0") + return result + + @staticmethod + def compute_overall_statistics( + per_task_stats: list[dict[str, dict[str, list[float] | np.ndarray]]], + dataset_sampling_weights: list[float] | np.ndarray, + percentile_mixing_method: str = "weighted_average", + ) -> dict[str, dict[str, list[float]]]: + """ + Computes overall statistics from per-task statistics using dataset sample weights. + + Args: + per_task_stats: List of per-task statistics. + Example format of one element in the per-task statistics list: + { + "state.gripper": { + "min": [...], + "max": [...], + "mean": [...], + "std": [...], + "q01": [...], + "q99": [...], + }, + ... + } + dataset_sampling_weights: List of sample weights for each task. + percentile_mixing_method: The method to mix the percentiles, either "weighted_average" or "weighted_std". + + Returns: + A dict of overall statistics per modality. + """ + # Normalize the sample weights to sum to 1 + dataset_sampling_weights = np.array(dataset_sampling_weights) + normalized_weights = dataset_sampling_weights / dataset_sampling_weights.sum() + + # Initialize overall statistics dict + overall_stats: dict[str, dict[str, list[float]]] = {} + + # Get the list of modality keys + modality_keys = per_task_stats[0].keys() + + for modality in modality_keys: + # Number of dimensions (assuming consistent across tasks) + num_dims = len(per_task_stats[0][modality]["mean"]) + + # Initialize accumulators for means and variances + weighted_means = np.zeros(num_dims) + weighted_squares = np.zeros(num_dims) + + # Collect min, max, q01, q99 from all tasks + min_list = [] + max_list = [] + q01_list = [] + q99_list = [] + + for task_idx, task_stats in enumerate(per_task_stats): + w_i = normalized_weights[task_idx] + stats = task_stats[modality] + means = np.array(stats["mean"]) + stds = np.array(stats["std"]) + + # Update weighted sums for mean and variance + weighted_means += w_i * means + weighted_squares += w_i * (stds**2 + means**2) + + # Collect min, max, q01, q99 + min_list.append(stats["min"]) + max_list.append(stats["max"]) + q01_list.append(stats["q01"]) + q99_list.append(stats["q99"]) + + # Compute overall mean + overall_mean = weighted_means.tolist() + + # Compute overall variance and std deviation + overall_variance = weighted_squares - weighted_means**2 + overall_std = np.sqrt(overall_variance).tolist() + + # Compute overall min and max per dimension + overall_min = np.min(np.array(min_list), axis=0).tolist() + overall_max = np.max(np.array(max_list), axis=0).tolist() + + # Compute overall q01 and q99 per dimension + # Use weighted average of per-task quantiles + q01_array = np.array(q01_list) + q99_array = np.array(q99_list) + if percentile_mixing_method == "weighted_average": + weighted_q01 = np.average(q01_array, axis=0, weights=normalized_weights).tolist() + weighted_q99 = np.average(q99_array, axis=0, weights=normalized_weights).tolist() + # std_q01 = np.std(q01_array, axis=0).tolist() + # std_q99 = np.std(q99_array, axis=0).tolist() + # print(modality) + # print(f"{std_q01=}, {std_q99=}") + # print(f"{weighted_q01=}, {weighted_q99=}") + elif percentile_mixing_method == "min_max": + weighted_q01 = np.min(q01_array, axis=0).tolist() + weighted_q99 = np.max(q99_array, axis=0).tolist() + else: + raise ValueError(f"Invalid percentile mixing method: {percentile_mixing_method}") + + # Store the overall statistics for the modality + overall_stats[modality] = { + "min": overall_min, + "max": overall_max, + "mean": overall_mean, + "std": overall_std, + "q01": weighted_q01, + "q99": weighted_q99, + } + + return overall_stats + + @staticmethod + def merge_metadata( + metadatas: list[DatasetMetadata], + dataset_sampling_weights: list[float], + percentile_mixing_method: str, + ) -> DatasetMetadata: + """Merge multiple metadata into one.""" + # Convert to dicts + metadata_dicts = [metadata.model_dump(mode="json") for metadata in metadatas] + # Create a new metadata dict + merged_metadata = {} + + # Check all metadata have the same embodiment tag + assert all( + metadata.embodiment_tag == metadatas[0].embodiment_tag for metadata in metadatas + ), "All metadata must have the same embodiment tag" + merged_metadata["embodiment_tag"] = metadatas[0].embodiment_tag + + # Merge the dataset statistics + dataset_statistics = {} + dataset_statistics["state"] = LeRobotMixtureDataset.compute_overall_statistics( + per_task_stats=[m["statistics"]["state"] for m in metadata_dicts], + dataset_sampling_weights=dataset_sampling_weights, + percentile_mixing_method=percentile_mixing_method, + ) + dataset_statistics["action"] = LeRobotMixtureDataset.compute_overall_statistics( + per_task_stats=[m["statistics"]["action"] for m in metadata_dicts], + dataset_sampling_weights=dataset_sampling_weights, + percentile_mixing_method=percentile_mixing_method, + ) + merged_metadata["statistics"] = dataset_statistics + + # Merge the modality configs + modality_configs = defaultdict(set) + for metadata in metadata_dicts: + for modality, configs in metadata["modalities"].items(): + modality_configs[modality].add(json.dumps(configs)) + merged_metadata["modalities"] = {} + for modality, configs in modality_configs.items(): + # Check that all modality configs correspond to the same tag matches + assert ( + len(configs) == 1 + ), f"Multiple modality configs for modality {modality}: {list(configs)}" + merged_metadata["modalities"][modality] = json.loads(configs.pop()) + + return DatasetMetadata.model_validate(merged_metadata) + + def update_metadata(self, metadata_config: dict, cached_statistics_path: Path | str | None = None) -> None: + """ + Merge multiple metadatas into one and set the transforms with the merged metadata. + + Args: + metadata_config (dict): Configuration for the metadata. + "percentile_mixing_method": The method to mix the percentiles, either "weighted_average" or "min_max". + weighted_average: Use the weighted average of the percentiles using the weight used in sampling the datasets. + min_max: Use the min of the 1st percentile and max of the 99th percentile. + """ + # If cached path is provided, try to load and apply + if cached_statistics_path is not None: + try: + cached_stats = self.load_merged_statistics(cached_statistics_path) + self.apply_cached_statistics(cached_stats) + return + except (FileNotFoundError, KeyError, ValidationError) as e: + print(f"Failed to load cached statistics: {e}") + print("Falling back to computing statistics from scratch...") + + self.tag = EmbodimentTag.NEW_EMBODIMENT.value + self.merged_metadata: dict[str, DatasetMetadata] = {} + # Group metadata by tag + all_metadatas: dict[str, list[DatasetMetadata]] = {} + for dataset in self.datasets: + if dataset.tag not in all_metadatas: + all_metadatas[dataset.tag] = [] + all_metadatas[dataset.tag].append(dataset.metadata) + for tag, metadatas in all_metadatas.items(): + self.merged_metadata[tag] = self.merge_metadata( + metadatas=metadatas, + dataset_sampling_weights=self.dataset_sampling_weights.tolist(), + percentile_mixing_method=metadata_config["percentile_mixing_method"], + ) + for dataset in self.datasets: + dataset.set_transforms_metadata(self.merged_metadata[dataset.tag]) + + def save_dataset_statistics(self, save_path: Path | str, format: str = "json") -> None: + """ + Save merged dataset statistics to specified path in the required format. + Only includes statistics for keys that are actually used in the datasets. + Gripper-related keys will be placed at the end. + + Args: + save_path (Path | str): Path to save the statistics file + format (str): Save format, currently only supports "json" + """ + save_path = Path(save_path) + save_path.parent.mkdir(parents=True, exist_ok=True) + + # Build the data structure to save + statistics_data = {} + + # Collect actually used keys from all datasets + all_used_action_keys = [] + all_used_state_keys = [] + + for dataset in self.datasets: + used_action_keys, used_state_keys = get_used_modality_keys(dataset.modality_keys) + for used_action_key in used_action_keys: + if used_action_key not in all_used_action_keys: + all_used_action_keys.append(used_action_key) + for used_state_key in used_state_keys: + if used_state_key not in all_used_state_keys: + all_used_state_keys.append(used_state_key) + + # Organize statistics by tag + for tag, merged_metadata in self.merged_metadata.items(): + tag_stats = {} + + # Process action statistics + if hasattr(merged_metadata.statistics, 'action') and merged_metadata.statistics.action: + action_stats = merged_metadata.statistics.action + + # Filter and reorder keys - iterate in all_used_action_keys order + non_gripper_keys = [] + gripper_keys = [] + + for key in all_used_action_keys: + if key in action_stats: + non_gripper_keys.append(key) + + reordered_keys = non_gripper_keys + gripper_keys + + filtered_action_stats = {} + for key in reordered_keys: + filtered_action_stats[key] = action_stats[key] + + if filtered_action_stats: + combined_action_stats = combine_modality_stats(filtered_action_stats) + + # Collect action normalization modes from datasets of this tag. + # "binary" takes precedence: if any dataset marks a key as binary, use binary. + _action_norm_modes: dict = {} + for _ds in self.datasets: + if _ds.tag == tag: + for _k, _m in _extract_action_normalization_modes(_ds.transforms).items(): + if _k not in _action_norm_modes or _m == "binary": + _action_norm_modes[_k] = _m + mask = generate_action_mask_for_used_keys( + merged_metadata.modalities.action, filtered_action_stats.keys(), + normalization_modes=_action_norm_modes, + ) + combined_action_stats["mask"] = mask + + tag_stats["action"] = combined_action_stats + + # Process state statistics + if hasattr(merged_metadata.statistics, 'state') and merged_metadata.statistics.state: + state_stats = merged_metadata.statistics.state + + # Filter and reorder keys - iterate in all_used_state_keys order + # Filter and reorder keys - iterate in all_used_state_keys order + non_gripper_keys = [] + gripper_keys = [] + + for key in all_used_state_keys: + if key in state_stats: + non_gripper_keys.append(key) + + reordered_keys = non_gripper_keys + gripper_keys + + filtered_state_stats = {} + for key in reordered_keys: + filtered_state_stats[key] = state_stats[key] + + if filtered_state_stats: + combined_state_stats = combine_modality_stats(filtered_state_stats) + tag_stats["state"] = combined_state_stats + + # Add dataset counts + tag_stats.update(self._get_dataset_counts(tag)) + + statistics_data[tag] = tag_stats + + # Save file + if format.lower() == "json": + if not str(save_path).endswith('.json'): + save_path = save_path.with_suffix('.json') + with open(save_path, 'w', encoding='utf-8') as f: + json.dump(statistics_data, f, indent=2, ensure_ascii=False) + else: + raise ValueError(f"Unsupported format: {format}. Currently only 'json' is supported.") + + print(f"Merged dataset statistics saved to: {save_path}") + print(f"Used action keys (reordered): {list(all_used_action_keys)}") + print(f"Used state keys (reordered): {list(all_used_state_keys)}") + + + def _combine_modality_stats(self, modality_stats: dict) -> dict: + """Backward compatibility wrapper.""" + return combine_modality_stats(modality_stats) + + def _generate_action_mask_for_used_keys(self, action_modalities: dict, used_action_keys_ordered) -> list[bool]: + """Backward compatibility wrapper.""" + return generate_action_mask_for_used_keys(action_modalities, used_action_keys_ordered) + + def _get_dataset_counts(self, tag: str) -> dict: + """ + Get dataset count information for specified tag. + + Args: + tag (str): embodiment tag + + Returns: + dict: Dictionary containing num_transitions and num_trajectories + """ + num_transitions = 0 + num_trajectories = 0 + + # Count dataset information belonging to this tag + for dataset in self.datasets: + if dataset.tag == tag: + num_transitions += len(dataset) + num_trajectories += len(dataset.trajectory_ids) + + return { + "num_transitions": num_transitions, + "num_trajectories": num_trajectories + } + + @classmethod + def load_merged_statistics(cls, load_path: Path | str) -> dict: + """ + Load merged dataset statistics from file. + + Args: + load_path (Path | str): Path to the statistics file + + Returns: + dict: Dictionary containing merged statistics + """ + load_path = Path(load_path) + if not load_path.exists(): + raise FileNotFoundError(f"Statistics file not found: {load_path}") + + if load_path.suffix.lower() == '.json': + with open(load_path, 'r', encoding='utf-8') as f: + return json.load(f) + elif load_path.suffix.lower() == '.pkl': + import pickle + with open(load_path, 'rb') as f: + return pickle.load(f) + else: + raise ValueError(f"Unsupported file format: {load_path.suffix}") + + def apply_cached_statistics(self, cached_statistics: dict) -> None: + """ + Apply cached statistics to avoid recomputation. + + Args: + cached_statistics (dict): Statistics loaded from file + """ + # Validate that cached statistics match current datasets + if "metadata" in cached_statistics: + cached_dataset_names = set(cached_statistics["metadata"]["dataset_names"]) + current_dataset_names = set(dataset.dataset_name for dataset in self.datasets) + + if cached_dataset_names != current_dataset_names: + print("Warning: Cached statistics dataset names don't match current datasets.") + print(f"Cached: {cached_dataset_names}") + print(f"Current: {current_dataset_names}") + return + + # Apply cached statistics + self.merged_metadata = {} + for tag, stats_data in cached_statistics.items(): + if tag == "metadata": # Skip metadata field + continue + + # Convert back to DatasetMetadata format + metadata_dict = { + "embodiment_tag": tag, + "statistics": { + "action": {}, + "state": {} + }, + "modalities": {} + } + + # Convert action statistics back + if "action" in stats_data: + action_data = stats_data["action"] + # This is simplified - you may need to split back to sub-keys + metadata_dict["statistics"]["action"] = action_data + + # Convert state statistics back + if "state" in stats_data: + state_data = stats_data["state"] + metadata_dict["statistics"]["state"] = state_data + + self.merged_metadata[tag] = DatasetMetadata.model_validate(metadata_dict) + + # Update transforms metadata for each dataset + for dataset in self.datasets: + if dataset.tag in self.merged_metadata: + dataset.set_transforms_metadata(self.merged_metadata[dataset.tag]) + + print(f"Applied cached statistics for {len(self.merged_metadata)} embodiment tags.") diff --git a/starVLA/dataloader/gr00t_lerobot/embodiment_tags.py b/starVLA/dataloader/gr00t_lerobot/embodiment_tags.py new file mode 100644 index 0000000000000000000000000000000000000000..b4cb43804eb0fd9c533df2addb27e0b57f62b340 --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/embodiment_tags.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from enum import Enum + + +class EmbodimentTag(Enum): + GR1 = "gr1" + """ + The GR1 dataset. + """ + + OXE_DROID = "oxe_droid" + """ + The OxE Droid dataset. + """ + + OXE_BRIDGE = "oxe_bridge" + """ + The OxE Bridge dataset. + """ + + OXE_RT1 = "oxe_rt1" + """ + The OxE RT-1 dataset. + """ + + AGIBOT_GENIE1 = "agibot_genie1" + """ + The AgiBot Genie-1 with gripper dataset. + """ + + NEW_EMBODIMENT = "new_embodiment" + """ + Any new embodiment for finetuning. + """ + + FRANKA = 'franka' + """ + The Franka Emika Panda robot. + """ + + ALOHA = 'aloha' + """ + The ALOHA bimanual robot (RoboChallenge dual-arm). + """ + + UR5 = 'ur5' + """ The UR5 single-arm robot. + """ + + ARX5 = 'arx5' + """ The ARX5 single-arm robot. + """ + + DOS_W1 = 'dos-w1' + """ The DOS-W1 single-arm robot. + """ + +# Embodiment tag string: to projector index in the Action Expert Module +EMBODIMENT_TAG_MAPPING = { + EmbodimentTag.NEW_EMBODIMENT.value: 31, + EmbodimentTag.OXE_DROID.value: 17, + EmbodimentTag.OXE_BRIDGE.value: 18, + EmbodimentTag.OXE_RT1.value: 19, + EmbodimentTag.AGIBOT_GENIE1.value: 26, + EmbodimentTag.GR1.value: 24, + EmbodimentTag.FRANKA.value: 25, + EmbodimentTag.ALOHA.value: 7, + EmbodimentTag.UR5.value: 8, + EmbodimentTag.ARX5.value: 9, + EmbodimentTag.DOS_W1.value: 10, +} + +# Robot type to embodiment tag mapping +ROBOT_TYPE_TO_EMBODIMENT_TAG = { + "libero_franka": EmbodimentTag.FRANKA, + "oxe_droid": EmbodimentTag.OXE_DROID, + "oxe_bridge": EmbodimentTag.OXE_BRIDGE, + "oxe_rt1": EmbodimentTag.OXE_RT1, + "demo_sim_franka_delta_joints": EmbodimentTag.FRANKA, + "custom_robot_config": EmbodimentTag.NEW_EMBODIMENT, + "fourier_gr1_arms_waist": EmbodimentTag.GR1, +} diff --git a/starVLA/dataloader/gr00t_lerobot/mixtures.py b/starVLA/dataloader/gr00t_lerobot/mixtures.py new file mode 100644 index 0000000000000000000000000000000000000000..40957c507912fd204d8f9e3191ab807fcc7ecc93 --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/mixtures.py @@ -0,0 +1,376 @@ +""" +mixtures.py + +Defines a registry of dataset mixtures and weights for the Open-X Embodiment Datasets. Each dataset is associated with +a float "sampling weight" +""" + +from typing import Dict, List, Tuple + + +# Dataset mixture name mapped to a list of tuples containing: +## {nakename: [(data_name, sampling_weight, robot_type)] } +DATASET_NAMED_MIXTURES = { + + "custom_dataset": [ + ("custom_dataset_name", 1.0, "custom_robot_config"), + ], + "custom_dataset_2": [ + ("custom_dataset_name_1", 1.0, "custom_robot_config"), + ("custom_dataset_name_2", 1.0, "custom_robot_config"), + ], + + "libero_all": [ + ("libero_object_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + ("libero_goal_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + ("libero_spatial_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + ("libero_10_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + # ("libero_90_no_noops_lerobot", 1.0, "libero_franka"), + ], + "libero_goal": [ + ("libero_goal_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + ], + "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"), + ], + + "demo_sim_pick_place": [ + ("sim_pick_place", 1.0, "demo_sim_franka_delta_joints"), + ], + + "custom_dataset": [ + ("custom_dataset_name", 1.0, "custom_robot_config"), + ], + "custom_dataset_2": [ + ("custom_dataset_name_1", 1.0, "custom_robot_config"), + ("custom_dataset_name_2", 1.0, "custom_robot_config"), + ], + + "fourier_gr1_unified_1000": [ + ("gr1_unified.PnPBottleToCabinetClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PnPCanToDrawerClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PnPCupToDrawerClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PnPMilkToMicrowaveClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PnPPotatoToMicrowaveClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PnPWineToCabinetClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromCuttingboardToBasketSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromCuttingboardToCardboardboxSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromCuttingboardToPanSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromCuttingboardToPotSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromCuttingboardToTieredbasketSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlacematToBasketSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlacematToBowlSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlacematToPlateSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlacematToTieredshelfSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlateToBowlSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlateToCardboardboxSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlateToPanSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlateToPlateSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromTrayToCardboardboxSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromTrayToPlateSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromTrayToPotSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromTrayToTieredbasketSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromTrayToTieredshelfSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ], + + "BEHAVIOR_challenge": [ + ("BEHAVIOR_challenge", 1.0, "R1Pro"), + ], + + + "SO101_pick": [ + ("pick_dataset_name", 1.0, "SO101"), + ], + + "arx_x5": [ + ("arx_x5", 1.0, "arx_x5"), + ], + + "robotwin_all": [ + ("Clean/adjust_bottle", 1.0, "robotwin"), + ("Clean/beat_block_hammer", 1.0, "robotwin"), + ("Clean/blocks_ranking_rgb", 1.0, "robotwin"), + ("Clean/blocks_ranking_size", 1.0, "robotwin"), + ("Clean/click_alarmclock", 1.0, "robotwin"), + ("Clean/click_bell", 1.0, "robotwin"), + ("Clean/dump_bin_bigbin", 1.0, "robotwin"), + ("Clean/grab_roller", 1.0, "robotwin"), + ("Clean/handover_block", 1.0, "robotwin"), + ("Clean/handover_mic", 1.0, "robotwin"), + ("Clean/hanging_mug", 1.0, "robotwin"), + ("Clean/lift_pot", 1.0, "robotwin"), + ("Clean/move_can_pot", 1.0, "robotwin"), + ("Clean/move_pillbottle_pad", 1.0, "robotwin"), + ("Clean/move_playingcard_away", 1.0, "robotwin"), + ("Clean/move_stapler_pad", 1.0, "robotwin"), + ("Clean/open_laptop", 1.0, "robotwin"), + ("Clean/open_microwave", 1.0, "robotwin"), + ("Clean/pick_diverse_bottles", 1.0, "robotwin"), + ("Clean/pick_dual_bottles", 1.0, "robotwin"), + ("Clean/place_a2b_left", 1.0, "robotwin"), + ("Clean/place_a2b_right", 1.0, "robotwin"), + ("Clean/place_bread_basket", 1.0, "robotwin"), + ("Clean/place_bread_skillet", 1.0, "robotwin"), + ("Clean/place_burger_fries", 1.0, "robotwin"), + ("Clean/place_can_basket", 1.0, "robotwin"), + ("Clean/place_cans_plasticbox", 1.0, "robotwin"), + ("Clean/place_container_plate", 1.0, "robotwin"), + ("Clean/place_dual_shoes", 1.0, "robotwin"), + ("Clean/place_empty_cup", 1.0, "robotwin"), + ("Clean/place_fan", 1.0, "robotwin"), + ("Clean/place_mouse_pad", 1.0, "robotwin"), + ("Clean/place_object_basket", 1.0, "robotwin"), + ("Clean/place_object_scale", 1.0, "robotwin"), + ("Clean/place_object_stand", 1.0, "robotwin"), + ("Clean/place_phone_stand", 1.0, "robotwin"), + ("Clean/place_shoe", 1.0, "robotwin"), + ("Clean/press_stapler", 1.0, "robotwin"), + ("Clean/put_bottles_dustbin", 1.0, "robotwin"), + ("Clean/put_object_cabinet", 1.0, "robotwin"), + ("Clean/rotate_qrcode", 1.0, "robotwin"), + ("Clean/scan_object", 1.0, "robotwin"), + ("Clean/shake_bottle", 1.0, "robotwin"), + ("Clean/shake_bottle_horizontally", 1.0, "robotwin"), + ("Clean/stack_blocks_three", 1.0, "robotwin"), + ("Clean/stack_blocks_two", 1.0, "robotwin"), + ("Clean/stack_bowls_three", 1.0, "robotwin"), + ("Clean/stack_bowls_two", 1.0, "robotwin"), + ("Clean/stamp_seal", 1.0, "robotwin"), + ("Clean/turn_switch", 1.0, "robotwin"), + ("Randomized/adjust_bottle", 1.0, "robotwin"), + ("Randomized/beat_block_hammer", 1.0, "robotwin"), + ("Randomized/blocks_ranking_rgb", 1.0, "robotwin"), + ("Randomized/blocks_ranking_size", 1.0, "robotwin"), + ("Randomized/click_alarmclock", 1.0, "robotwin"), + ("Randomized/click_bell", 1.0, "robotwin"), + ("Randomized/dump_bin_bigbin", 1.0, "robotwin"), + ("Randomized/grab_roller", 1.0, "robotwin"), + ("Randomized/handover_block", 1.0, "robotwin"), + ("Randomized/handover_mic", 1.0, "robotwin"), + ("Randomized/hanging_mug", 1.0, "robotwin"), + ("Randomized/lift_pot", 1.0, "robotwin"), + ("Randomized/move_can_pot", 1.0, "robotwin"), + ("Randomized/move_pillbottle_pad", 1.0, "robotwin"), + ("Randomized/move_playingcard_away", 1.0, "robotwin"), + ("Randomized/move_stapler_pad", 1.0, "robotwin"), + ("Randomized/open_laptop", 1.0, "robotwin"), + ("Randomized/open_microwave", 1.0, "robotwin"), + ("Randomized/pick_diverse_bottles", 1.0, "robotwin"), + ("Randomized/pick_dual_bottles", 1.0, "robotwin"), + ("Randomized/place_a2b_left", 1.0, "robotwin"), + ("Randomized/place_a2b_right", 1.0, "robotwin"), + ("Randomized/place_bread_basket", 1.0, "robotwin"), + ("Randomized/place_bread_skillet", 1.0, "robotwin"), + ("Randomized/place_burger_fries", 1.0, "robotwin"), + ("Randomized/place_can_basket", 1.0, "robotwin"), + ("Randomized/place_cans_plasticbox", 1.0, "robotwin"), + ("Randomized/place_container_plate", 1.0, "robotwin"), + ("Randomized/place_dual_shoes", 1.0, "robotwin"), + ("Randomized/place_empty_cup", 1.0, "robotwin"), + ("Randomized/place_fan", 1.0, "robotwin"), + ("Randomized/place_mouse_pad", 1.0, "robotwin"), + ("Randomized/place_object_basket", 1.0, "robotwin"), + ("Randomized/place_object_scale", 1.0, "robotwin"), + ("Randomized/place_object_stand", 1.0, "robotwin"), + ("Randomized/place_phone_stand", 1.0, "robotwin"), + ("Randomized/place_shoe", 1.0, "robotwin"), + ("Randomized/press_stapler", 1.0, "robotwin"), + ("Randomized/put_bottles_dustbin", 1.0, "robotwin"), + ("Randomized/put_object_cabinet", 1.0, "robotwin"), + ("Randomized/rotate_qrcode", 1.0, "robotwin"), + ("Randomized/scan_object", 1.0, "robotwin"), + ("Randomized/shake_bottle", 1.0, "robotwin"), + ("Randomized/shake_bottle_horizontally", 1.0, "robotwin"), + ("Randomized/stack_blocks_three", 1.0, "robotwin"), + ("Randomized/stack_blocks_two", 1.0, "robotwin"), + ("Randomized/stack_bowls_three", 1.0, "robotwin"), + ("Randomized/stack_bowls_two", 1.0, "robotwin"), + ("Randomized/stamp_seal", 1.0, "robotwin"), + ("Randomized/turn_switch", 1.0, "robotwin"), + ], + "robotwin_all_50": [ + ("Clean/adjust_bottle", 1.0, "robotwin50"), + ("Clean/beat_block_hammer", 1.0, "robotwin50"), + ("Clean/blocks_ranking_rgb", 1.0, "robotwin50"), + ("Clean/blocks_ranking_size", 1.0, "robotwin50"), + ("Clean/click_alarmclock", 1.0, "robotwin50"), + ("Clean/click_bell", 1.0, "robotwin50"), + ("Clean/dump_bin_bigbin", 1.0, "robotwin50"), + ("Clean/grab_roller", 1.0, "robotwin50"), + ("Clean/handover_block", 1.0, "robotwin50"), + ("Clean/handover_mic", 1.0, "robotwin50"), + ("Clean/hanging_mug", 1.0, "robotwin50"), + ("Clean/lift_pot", 1.0, "robotwin50"), + ("Clean/move_can_pot", 1.0, "robotwin50"), + ("Clean/move_pillbottle_pad", 1.0, "robotwin50"), + ("Clean/move_playingcard_away", 1.0, "robotwin50"), + ("Clean/move_stapler_pad", 1.0, "robotwin50"), + ("Clean/open_laptop", 1.0, "robotwin50"), + ("Clean/open_microwave", 1.0, "robotwin50"), + ("Clean/pick_diverse_bottles", 1.0, "robotwin50"), + ("Clean/pick_dual_bottles", 1.0, "robotwin50"), + ("Clean/place_a2b_left", 1.0, "robotwin50"), + ("Clean/place_a2b_right", 1.0, "robotwin50"), + ("Clean/place_bread_basket", 1.0, "robotwin50"), + ("Clean/place_bread_skillet", 1.0, "robotwin50"), + ("Clean/place_burger_fries", 1.0, "robotwin50"), + ("Clean/place_can_basket", 1.0, "robotwin50"), + ("Clean/place_cans_plasticbox", 1.0, "robotwin50"), + ("Clean/place_container_plate", 1.0, "robotwin50"), + ("Clean/place_dual_shoes", 1.0, "robotwin50"), + ("Clean/place_empty_cup", 1.0, "robotwin50"), + ("Clean/place_fan", 1.0, "robotwin50"), + ("Clean/place_mouse_pad", 1.0, "robotwin50"), + ("Clean/place_object_basket", 1.0, "robotwin50"), + ("Clean/place_object_scale", 1.0, "robotwin50"), + ("Clean/place_object_stand", 1.0, "robotwin50"), + ("Clean/place_phone_stand", 1.0, "robotwin50"), + ("Clean/place_shoe", 1.0, "robotwin50"), + ("Clean/press_stapler", 1.0, "robotwin50"), + ("Clean/put_bottles_dustbin", 1.0, "robotwin50"), + ("Clean/put_object_cabinet", 1.0, "robotwin50"), + ("Clean/rotate_qrcode", 1.0, "robotwin50"), + ("Clean/scan_object", 1.0, "robotwin50"), + ("Clean/shake_bottle", 1.0, "robotwin50"), + ("Clean/shake_bottle_horizontally", 1.0, "robotwin50"), + ("Clean/stack_blocks_three", 1.0, "robotwin50"), + ("Clean/stack_blocks_two", 1.0, "robotwin50"), + ("Clean/stack_bowls_three", 1.0, "robotwin50"), + ("Clean/stack_bowls_two", 1.0, "robotwin50"), + ("Clean/stamp_seal", 1.0, "robotwin50"), + ("Clean/turn_switch", 1.0, "robotwin50"), + ("Randomized/adjust_bottle", 1.0, "robotwin50"), + ("Randomized/beat_block_hammer", 1.0, "robotwin50"), + ("Randomized/blocks_ranking_rgb", 1.0, "robotwin50"), + ("Randomized/blocks_ranking_size", 1.0, "robotwin50"), + ("Randomized/click_alarmclock", 1.0, "robotwin50"), + ("Randomized/click_bell", 1.0, "robotwin50"), + ("Randomized/dump_bin_bigbin", 1.0, "robotwin50"), + ("Randomized/grab_roller", 1.0, "robotwin50"), + ("Randomized/handover_block", 1.0, "robotwin50"), + ("Randomized/handover_mic", 1.0, "robotwin50"), + ("Randomized/hanging_mug", 1.0, "robotwin50"), + ("Randomized/lift_pot", 1.0, "robotwin50"), + ("Randomized/move_can_pot", 1.0, "robotwin50"), + ("Randomized/move_pillbottle_pad", 1.0, "robotwin50"), + ("Randomized/move_playingcard_away", 1.0, "robotwin50"), + ("Randomized/move_stapler_pad", 1.0, "robotwin50"), + ("Randomized/open_laptop", 1.0, "robotwin50"), + ("Randomized/open_microwave", 1.0, "robotwin50"), + ("Randomized/pick_diverse_bottles", 1.0, "robotwin50"), + ("Randomized/pick_dual_bottles", 1.0, "robotwin50"), + ("Randomized/place_a2b_left", 1.0, "robotwin50"), + ("Randomized/place_a2b_right", 1.0, "robotwin50"), + ("Randomized/place_bread_basket", 1.0, "robotwin50"), + ("Randomized/place_bread_skillet", 1.0, "robotwin50"), + ("Randomized/place_burger_fries", 1.0, "robotwin50"), + ("Randomized/place_can_basket", 1.0, "robotwin50"), + ("Randomized/place_cans_plasticbox", 1.0, "robotwin50"), + ("Randomized/place_container_plate", 1.0, "robotwin50"), + ("Randomized/place_dual_shoes", 1.0, "robotwin50"), + ("Randomized/place_empty_cup", 1.0, "robotwin50"), + ("Randomized/place_fan", 1.0, "robotwin50"), + ("Randomized/place_mouse_pad", 1.0, "robotwin50"), + ("Randomized/place_object_basket", 1.0, "robotwin50"), + ("Randomized/place_object_scale", 1.0, "robotwin50"), + ("Randomized/place_object_stand", 1.0, "robotwin50"), + ("Randomized/place_phone_stand", 1.0, "robotwin50"), + ("Randomized/place_shoe", 1.0, "robotwin50"), + ("Randomized/press_stapler", 1.0, "robotwin50"), + ("Randomized/put_bottles_dustbin", 1.0, "robotwin50"), + ("Randomized/put_object_cabinet", 1.0, "robotwin50"), + ("Randomized/rotate_qrcode", 1.0, "robotwin50"), + ("Randomized/scan_object", 1.0, "robotwin50"), + ("Randomized/shake_bottle", 1.0, "robotwin50"), + ("Randomized/shake_bottle_horizontally", 1.0, "robotwin50"), + ("Randomized/stack_blocks_three", 1.0, "robotwin50"), + ("Randomized/stack_blocks_two", 1.0, "robotwin50"), + ("Randomized/stack_bowls_three", 1.0, "robotwin50"), + ("Randomized/stack_bowls_two", 1.0, "robotwin50"), + ("Randomized/stamp_seal", 1.0, "robotwin50"), + ("Randomized/turn_switch", 1.0, "robotwin50"), + ], + "robotwin": [ + ("adjust_bottle", 1.0, "robotwin"), + ("beat_block_hammer", 1.0, "robotwin"), + ("blocks_ranking_rgb", 1.0, "robotwin"), + ("blocks_ranking_size", 1.0, "robotwin"), + ("click_alarmclock", 1.0, "robotwin"), + ("click_bell", 1.0, "robotwin"), + ("dump_bin_bigbin", 1.0, "robotwin"), + ("grab_roller", 1.0, "robotwin"), + ("handover_block", 1.0, "robotwin"), + ("handover_mic", 1.0, "robotwin"), + ("hanging_mug", 1.0, "robotwin"), + ("lift_pot", 1.0, "robotwin"), + ("move_can_pot", 1.0, "robotwin"), + ("move_pillbottle_pad", 1.0, "robotwin"), + ("move_playingcard_away", 1.0, "robotwin"), + ("move_stapler_pad", 1.0, "robotwin"), + ("open_laptop", 1.0, "robotwin"), + ("open_microwave", 1.0, "robotwin"), + ("pick_diverse_bottles", 1.0, "robotwin"), + ("pick_dual_bottles", 1.0, "robotwin"), + ("place_a2b_left", 1.0, "robotwin"), + ("place_a2b_right", 1.0, "robotwin"), + ("place_bread_basket", 1.0, "robotwin"), + ("place_bread_skillet", 1.0, "robotwin"), + ("place_burger_fries", 1.0, "robotwin"), + ("place_can_basket", 1.0, "robotwin"), + ("place_cans_plasticbox", 1.0, "robotwin"), + ("place_container_plate", 1.0, "robotwin"), + ("place_dual_shoes", 1.0, "robotwin"), + ("place_empty_cup", 1.0, "robotwin"), + ("place_fan", 1.0, "robotwin"), + ("place_mouse_pad", 1.0, "robotwin"), + ("place_object_basket", 1.0, "robotwin"), + ("place_object_scale", 1.0, "robotwin"), + ("place_object_stand", 1.0, "robotwin"), + ("place_phone_stand", 1.0, "robotwin"), + ("place_shoe", 1.0, "robotwin"), + ("press_stapler", 1.0, "robotwin"), + ("put_bottles_dustbin", 1.0, "robotwin"), + ("put_object_cabinet", 1.0, "robotwin"), + ("rotate_qrcode", 1.0, "robotwin"), + ("scan_object", 1.0, "robotwin"), + ("shake_bottle", 1.0, "robotwin"), + ("shake_bottle_horizontally", 1.0, "robotwin"), + ("stack_blocks_three", 1.0, "robotwin"), + ("stack_blocks_two", 1.0, "robotwin"), + ("stack_bowls_three", 1.0, "robotwin"), + ("stack_bowls_two", 1.0, "robotwin"), + ("stamp_seal", 1.0, "robotwin"), + ("turn_switch", 1.0, "robotwin"), + ], + + "robotwin_task1": [ + ("adjust_bottle", 1.0, "robotwin"), + ], + "robotwin_task2": [ + ("place_a2b_left", 1.0, "robotwin"), + ("place_a2b_right", 1.0, "robotwin"), + ], + + "multi_robot": [ + ("LEROBOT_LIBERO_DATA/libero_10_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + # ("OXE_LEROBOT_DATASET/bridge_orig_1.0.0_lerobot", 1.0, "oxe_bridge"), + ], + + "calvin_task_D_D_v3.0": [ + ("calvin_task_D_D_v3.0", 1.0, "libero_franka"), + ], + + "vla_arena_L0_S": [ + ("VLA_Arena_L0_S_lerobot_openpi", 1.0, "vla_arena_franka"), + ], + "vla_arena_L0_M": [ + ("VLA_Arena_L0_M_lerobot_openpi", 1.0, "vla_arena_franka"), + ], + "vla_arena_L0_L": [ + ("VLA_Arena_L0_L_lerobot_openpi", 1.0, "vla_arena_franka"), + ], +} diff --git a/starVLA/dataloader/gr00t_lerobot/registry.py b/starVLA/dataloader/gr00t_lerobot/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..581b926699a9ca9240d361ddb972db00f4867bf6 --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/registry.py @@ -0,0 +1,149 @@ +""" +Centralized registry that auto-discovers benchmark-specific data configs from +``examples/*/train_files/data_registry/`` and merges them with the base +registries defined in this package. + +Three registries are maintained: + +* ``DATASET_NAMED_MIXTURES`` – mixture_name → [(dataset, weight, robot_type)] +* ``ROBOT_TYPE_CONFIG_MAP`` – robot_type → DataConfig instance +* ``ROBOT_TYPE_TO_EMBODIMENT_TAG`` – robot_type → EmbodimentTag + +``ROBOT_TYPE_TO_EMBODIMENT_TAG`` is **derived** from ``ROBOT_TYPE_CONFIG_MAP`` +by reading each DataConfig class's ``embodiment_tag`` classvar (Proposal A). +Classes without the classvar fall back to ``EmbodimentTag.NEW_EMBODIMENT``. +Legacy bench files exposing their own ``ROBOT_TYPE_TO_EMBODIMENT_TAG`` dict are +still honored as overrides for backward compatibility. + + +Usage:: + + from starVLA.dataloader.gr00t_lerobot.registry import ( + ROBOT_TYPE_CONFIG_MAP, + ROBOT_TYPE_TO_EMBODIMENT_TAG, + DATASET_NAMED_MIXTURES, + ) +""" + +from __future__ import annotations + +import importlib +import importlib.util +import logging +import sys +from pathlib import Path + +# Base registries (kept as fallback / seed values) +from starVLA.dataloader.gr00t_lerobot.data_config import ( + ROBOT_TYPE_CONFIG_MAP as _BASE_CONFIG_MAP, +) +from starVLA.dataloader.gr00t_lerobot.embodiment_tags import ( + EmbodimentTag, # re-export for convenience +) +from starVLA.dataloader.gr00t_lerobot.mixtures import ( + DATASET_NAMED_MIXTURES as _BASE_MIXTURES, +) + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Mutable copies – will be extended by discovered modules +# --------------------------------------------------------------------------- +ROBOT_TYPE_CONFIG_MAP: dict = dict(_BASE_CONFIG_MAP) +DATASET_NAMED_MIXTURES: dict = dict(_BASE_MIXTURES) + +# Legacy explicit overrides (rarely needed; prefer the classvar on DataConfig). +_LEGACY_TAG_OVERRIDES: dict = {} + + +def _derive_tag_map() -> dict: + """Build robot_type -> EmbodimentTag from ROBOT_TYPE_CONFIG_MAP. + + Reads each DataConfig instance's ``embodiment_tag`` classvar. Falls back to + ``EmbodimentTag.NEW_EMBODIMENT`` if the classvar is missing. Legacy + overrides from bench files take precedence. + """ + out: dict = {} + for rt, cfg in ROBOT_TYPE_CONFIG_MAP.items(): + tag = getattr(cfg, "embodiment_tag", None) + if tag is None: + logger.warning( + "[registry] DataConfig for robot_type=%r has no `embodiment_tag` " + "classvar; defaulting to EmbodimentTag.NEW_EMBODIMENT.", rt + ) + tag = EmbodimentTag.NEW_EMBODIMENT + out[rt] = tag + out.update(_LEGACY_TAG_OVERRIDES) + return out + + +ROBOT_TYPE_TO_EMBODIMENT_TAG: dict = {} # populated by discover_and_merge() + +# --------------------------------------------------------------------------- +# Discovery logic +# --------------------------------------------------------------------------- +_REGISTRY_DIR_NAME = "data_registry" +_DISCOVERED = False + + +def _find_registry_dirs() -> list[Path]: + """Return all ``examples/*/train_files/data_registry/`` directories.""" + # Walk up from this file to the repo root + # registry.py is at starVLA/dataloader/gr00t_lerobot/registry.py + # parents: [0]=gr00t_lerobot, [1]=dataloader, [2]=starVLA(pkg), [3]=repo root + repo_root = Path(__file__).resolve().parents[3] + examples_dir = repo_root / "examples" + if not examples_dir.is_dir(): + return [] + dirs: list[Path] = [] + for bench_dir in sorted(examples_dir.iterdir()): + registry_dir = bench_dir / "train_files" / _REGISTRY_DIR_NAME + if registry_dir.is_dir(): + dirs.append(registry_dir) + return dirs + + +def _load_module_from_path(module_name: str, file_path: Path): + """Import a Python file as a module with the given name.""" + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None or spec.loader is None: + return None + mod = importlib.util.module_from_spec(spec) + sys.modules[module_name] = mod + spec.loader.exec_module(mod) + return mod + + +def discover_and_merge() -> None: + """Scan ``examples/*/train_files/data_registry/`` and merge into global registries.""" + global _DISCOVERED, ROBOT_TYPE_TO_EMBODIMENT_TAG + if _DISCOVERED: + return + _DISCOVERED = True + + for registry_dir in _find_registry_dirs(): + bench_name = registry_dir.parents[1].name # examples//train_files/data_registry + prefix = f"_data_registry_{bench_name}" + + # --- data_config.py (may contain all three registries) --- + cfg_file = registry_dir / "data_config.py" + if cfg_file.is_file(): + mod = _load_module_from_path(f"{prefix}.data_config", cfg_file) + if mod: + if hasattr(mod, "ROBOT_TYPE_CONFIG_MAP"): + ROBOT_TYPE_CONFIG_MAP.update(mod.ROBOT_TYPE_CONFIG_MAP) + logger.debug(f"[registry] Loaded data_config from {bench_name}: {list(mod.ROBOT_TYPE_CONFIG_MAP.keys())}") + # Legacy: bench file may still expose an explicit dict override. + if hasattr(mod, "ROBOT_TYPE_TO_EMBODIMENT_TAG"): + _LEGACY_TAG_OVERRIDES.update(mod.ROBOT_TYPE_TO_EMBODIMENT_TAG) + logger.debug(f"[registry] Legacy embodiment_tag overrides from {bench_name}: {list(mod.ROBOT_TYPE_TO_EMBODIMENT_TAG.keys())}") + if hasattr(mod, "DATASET_NAMED_MIXTURES"): + DATASET_NAMED_MIXTURES.update(mod.DATASET_NAMED_MIXTURES) + logger.debug(f"[registry] Loaded mixtures from {bench_name} (data_config): {list(mod.DATASET_NAMED_MIXTURES.keys())}") + + # Re-derive after all benches are loaded so classvars + legacy overrides combine. + ROBOT_TYPE_TO_EMBODIMENT_TAG = _derive_tag_map() + + +# Run discovery on first import +discover_and_merge() diff --git a/starVLA/dataloader/gr00t_lerobot/schema.py b/starVLA/dataloader/gr00t_lerobot/schema.py new file mode 100644 index 0000000000000000000000000000000000000000..64519e56a3c59e5d08c8f8f6370f640061859b4d --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/schema.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from enum import Enum +from typing import Optional + +from numpydantic import NDArray +from pydantic import BaseModel, Field, field_serializer + +from .embodiment_tags import EmbodimentTag + +# Common schema + + +class RotationType(Enum): + """Type of rotation representation""" + + AXIS_ANGLE = "axis_angle" + QUATERNION = "quaternion" + ROTATION_6D = "rotation_6d" + MATRIX = "matrix" + EULER_ANGLES_RPY = "euler_angles_rpy" + EULER_ANGLES_RYP = "euler_angles_ryp" + EULER_ANGLES_PRY = "euler_angles_pry" + EULER_ANGLES_PYR = "euler_angles_pyr" + EULER_ANGLES_YRP = "euler_angles_yrp" + EULER_ANGLES_YPR = "euler_angles_ypr" + + +# LeRobot schema + + +class LeRobotModalityField(BaseModel): + """Metadata for a LeRobot modality field.""" + + original_key: Optional[str] = Field( + default=None, + description="The original key of the modality in the LeRobot dataset", + ) + + +class LeRobotStateActionMetadata(LeRobotModalityField): + """Metadata for a LeRobot modality.""" + + start: int = Field( + ..., + description="The start index of the modality in the concatenated state/action vector", + ) + end: int = Field( + ..., + description="The end index of the modality in the concatenated state/action vector", + ) + rotation_type: Optional[RotationType] = Field( + default=None, description="The type of rotation for the modality" + ) + absolute: bool = Field(default=True, description="Whether the modality is absolute") + dtype: str = Field( + default="float64", + description="The data type of the modality. Defaults to float64.", + ) + range: Optional[tuple[float, float]] = Field( + default=None, + description="The range of the modality, if applicable. Defaults to None.", + ) + original_key: Optional[str] = Field( + default=None, + description="The original key of the modality in the LeRobot dataset.", + ) + + +class LeRobotStateMetadata(LeRobotStateActionMetadata): + """Metadata for a LeRobot state modality.""" + + original_key: Optional[str] = Field( + default="observation.state", # LeRobot convention for states + description="The original key of the state modality in the LeRobot dataset", + ) + + +class LeRobotActionMetadata(LeRobotStateActionMetadata): + """Metadata for a LeRobot action modality.""" + + original_key: Optional[str] = Field( + default="action", # LeRobot convention for actions + description="The original key of the action modality in the LeRobot dataset", + ) + + +class LeRobotModalityMetadata(BaseModel): + """Metadata for a LeRobot modality.""" + + state: dict[str, LeRobotStateMetadata] = Field( + ..., + description="The metadata for the state modality. The keys are the names of each split of the state vector.", + ) + action: dict[str, LeRobotActionMetadata] = Field( + ..., + description="The metadata for the action modality. The keys are the names of each split of the action vector.", + ) + video: dict[str, LeRobotModalityField] = Field( + ..., + description="The metadata for the video modality. The keys are the new names of each video modality.", + ) + annotation: Optional[dict[str, LeRobotModalityField]] = Field( + default=None, + description="The metadata for the annotation modality. The keys are the new names of each annotation modality.", + ) + + def get_key_meta(self, key: str) -> LeRobotModalityField: + """Get the metadata for a key in the LeRobot modality metadata. + + Args: + key (str): The key to get the metadata for. + + Returns: + LeRobotModalityField: The metadata for the key. + + Example: + lerobot_modality_meta = LeRobotModalityMetadata.model_validate(U.load_json(modality_meta_path)) + lerobot_modality_meta.get_key_meta("state.joint_shoulder_y") + lerobot_modality_meta.get_key_meta("video.main_camera") + lerobot_modality_meta.get_key_meta("annotation.human.action.task_description") + """ + split_key = key.split(".") + modality = split_key[0] + subkey = ".".join(split_key[1:]) + if modality == "state": + if subkey not in self.state: + raise ValueError( + f"Key: {key}, state key {subkey} not found in metadata, available state keys: {self.state.keys()}" + ) + return self.state[subkey] + elif modality == "action": + if subkey not in self.action: + raise ValueError( + f"Key: {key}, action key {subkey} not found in metadata, available action keys: {self.action.keys()}" + ) + return self.action[subkey] + elif modality == "video": + if subkey not in self.video: + raise ValueError( + f"Key: {key}, video key {subkey} not found in metadata, available video keys: {self.video.keys()}" + ) + return self.video[subkey] + elif modality == "annotation": + assert ( + self.annotation is not None + ), "Trying to get annotation metadata for a dataset with no annotations" + if subkey not in self.annotation: + raise ValueError( + f"Key: {key}, annotation key {subkey} not found in metadata, available annotation keys: {self.annotation.keys()}" + ) + return self.annotation[subkey] + else: + raise ValueError(f"Key: {key}, unexpected modality: {modality}") + + +# Dataset schema (parsed from LeRobot schema and simplified) + + +class DatasetStatisticalValues(BaseModel): + max: NDArray = Field(..., description="Maximum values") + min: NDArray = Field(..., description="Minimum values") + mean: NDArray = Field(..., description="Mean values") + std: NDArray = Field(..., description="Standard deviation") + q01: NDArray = Field(..., description="1st percentile values") + q99: NDArray = Field(..., description="99th percentile values") + + @field_serializer("*", when_used="json") + def serialize_ndarray(self, v: NDArray) -> list[float]: + return v.tolist() # type: ignore + + +class DatasetStatistics(BaseModel): + state: dict[str, DatasetStatisticalValues] = Field(..., description="Statistics of the state") + action: dict[str, DatasetStatisticalValues] = Field(..., description="Statistics of the action") + + +class VideoMetadata(BaseModel): + """Metadata of the video modality""" + + resolution: tuple[int, int] = Field(..., description="Resolution of the video") + channels: int = Field(..., description="Number of channels in the video", gt=0) + fps: float = Field(..., description="Frames per second", gt=0) + + +class StateActionMetadata(BaseModel): + absolute: bool = Field(..., description="Whether the state or action is absolute") + rotation_type: Optional[RotationType] = Field(None, description="Type of rotation, if any") + shape: tuple[int, ...] = Field(..., description="Shape of the state or action") + continuous: bool = Field(..., description="Whether the state or action is continuous") + + +class DatasetModalities(BaseModel): + video: dict[str, VideoMetadata] = Field(..., description="Metadata of the video") + state: dict[str, StateActionMetadata] = Field(..., description="Metadata of the state") + action: dict[str, StateActionMetadata] = Field(..., description="Metadata of the action") + + +class DatasetMetadata(BaseModel): + """Metadata of the trainable dataset + + Changes: + - Update to use the new RawCommitHashMetadataMetadata_V1_2 + """ + + statistics: DatasetStatistics = Field(..., description="Statistics of the dataset") + modalities: DatasetModalities = Field(..., description="Metadata of the modalities") + embodiment_tag: EmbodimentTag = Field(..., description="Embodiment tag of the dataset") \ No newline at end of file diff --git a/starVLA/dataloader/gr00t_lerobot/transform/__init__.py b/starVLA/dataloader/gr00t_lerobot/transform/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bf22b5e4ca3d1c7937a25234cbf08e2644593587 --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/transform/__init__.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .base import ( + ComposedModalityTransform, + InvertibleModalityTransform, + ModalityTransform, +) +from .concat import ConcatTransform +# from .state_action import ( +# StateActionDropout, +# StateActionPerturbation, +# StateActionSinCosTransform, +# StateActionToTensor, +# StateActionTransform, +# ) +from .video import ( + VideoColorJitter, + VideoCrop, + VideoGrayscale, + VideoHorizontalFlip, + VideoRandomGrayscale, + VideoRandomPosterize, + VideoRandomRotation, + VideoResize, + VideoToNumpy, + VideoToTensor, + VideoTransform, +) diff --git a/starVLA/dataloader/gr00t_lerobot/transform/base.py b/starVLA/dataloader/gr00t_lerobot/transform/base.py new file mode 100644 index 0000000000000000000000000000000000000000..aac88559af98fa23f34fbb9135775d0819c281ef --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/transform/base.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from abc import ABC, abstractmethod +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr + +from ..schema import DatasetMetadata + + +class ModalityTransform(BaseModel, ABC): + """ + Abstract class for transforming data modalities, e.g. video frame augmentation or action normalization. + """ + + apply_to: list[str] = Field(..., description="The keys to apply the transform to.") + training: bool = Field( + default=True, description="Whether to apply the transform in training mode." + ) + _dataset_metadata: DatasetMetadata | None = PrivateAttr(default=None) + + model_config = ConfigDict(arbitrary_types_allowed=True) + + @property + def dataset_metadata(self) -> DatasetMetadata: + assert ( + self._dataset_metadata is not None + ), "Dataset metadata is not set. Please call set_metadata() before calling apply()." + return self._dataset_metadata + + @dataset_metadata.setter + def dataset_metadata(self, value: DatasetMetadata): + self._dataset_metadata = value + + def set_metadata(self, dataset_metadata: DatasetMetadata): + """ + Set the dataset metadata. This is useful for transforms that need to know the dataset metadata, e.g. to normalize actions. + Subclasses can override this method if they need to do something more complex. + """ + self.dataset_metadata = dataset_metadata + + def __call__(self, data: dict[str, Any]) -> dict[str, Any]: + """Apply the transformation to the data corresponding to target_keys and return the processed data. + + Args: + data (dict[str, Any]): The data to transform. + example: data = { + "video.image_side_0": np.ndarray, + "action.eef_position": np.ndarray, + ... + } + + Returns: + dict[str, Any]: The transformed data. + example: transformed_data = { + "video.image_side_0": np.ndarray, + "action.eef_position": torch.Tensor, # Normalized and converted to tensor + ... + } + """ + return self.apply(data) + + @abstractmethod + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + """Apply the transformation to the data corresponding to keys matching the `apply_to` regular expression and return the processed data.""" + + def train(self): + self.training = True + + def eval(self): + self.training = False + + +class InvertibleModalityTransform(ModalityTransform): + @abstractmethod + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + """Reverse the transformation to the data corresponding to keys matching the `apply_to` regular expression and return the processed data.""" + + +class ComposedModalityTransform(ModalityTransform): + """Compose multiple modality transforms.""" + + transforms: list[ModalityTransform] = Field(..., description="The transforms to compose.") + apply_to: list[str] = Field( + default_factory=list, description="Will be ignored for composed transforms." + ) + training: bool = Field( + default=True, description="Whether to apply the transform in training mode." + ) + + model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True) + + def set_metadata(self, dataset_metadata: DatasetMetadata): + for transform in self.transforms: + transform.set_metadata(dataset_metadata) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for i, transform in enumerate(self.transforms): + try: + data = transform(data) + except Exception as e: + raise ValueError(f"Error applying transform {i} to data: {e}") from e + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + for i, transform in enumerate(reversed(self.transforms)): + if isinstance(transform, InvertibleModalityTransform): + try: + data = transform.unapply(data) + except Exception as e: + step = len(self.transforms) - i - 1 + raise ValueError(f"Error unapplying transform {step} to data: {e}") from e + return data + + def train(self): + for transform in self.transforms: + transform.train() + + def eval(self): + for transform in self.transforms: + transform.eval() diff --git a/starVLA/dataloader/gr00t_lerobot/transform/concat.py b/starVLA/dataloader/gr00t_lerobot/transform/concat.py new file mode 100644 index 0000000000000000000000000000000000000000..cf8eea4c77fc163ecdb0d25aeca26a2cde99f8c4 --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/transform/concat.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional + +import numpy as np +import torch +from pydantic import Field + +from ..schema import DatasetMetadata, StateActionMetadata +from .base import InvertibleModalityTransform + + +class ConcatTransform(InvertibleModalityTransform): + """ + Concatenate the keys according to specified order. + """ + + # -- We inherit from ModalityTransform, so we keep apply_to as well -- + apply_to: list[str] = Field( + default_factory=list, description="Not used in this transform, kept for compatibility." + ) + + video_concat_order: list[str] = Field( + ..., + description="Concatenation order for each video modality. " + "Format: ['video.ego_view_pad_res224_freq20', ...]", + ) + + state_concat_order: Optional[list[str]] = Field( + default=None, + description="Concatenation order for each state modality. " + "Format: ['state.position', 'state.velocity', ...].", + ) + + action_concat_order: Optional[list[str]] = Field( + default=None, + description="Concatenation order for each action modality. " + "Format: ['action.position', 'action.velocity', ...].", + ) + + action_dims: dict[str, int] = Field( + default_factory=dict, + description="The dimensions of the action keys.", + ) + state_dims: dict[str, int] = Field( + default_factory=dict, + description="The dimensions of the state keys.", + ) + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = { + "apply_to", + "video_concat_order", + "state_concat_order", + "action_concat_order", + } + else: + include = kwargs.pop("include", None) + + return super().model_dump(*args, include=include, **kwargs) + + def apply(self, data: dict) -> dict: + grouped_keys = {} + for key in data.keys(): + try: + modality, _ = key.split(".") + except: # noqa: E722 + ### Handle language annotation special case + if "annotation" in key: + modality = "language" + else: + modality = "others" + if modality not in grouped_keys: + grouped_keys[modality] = [] + grouped_keys[modality].append(key) + + if "video" in grouped_keys: + # Check if keys in video_concat_order, state_concat_order, action_concat_order are + # ineed contained in the data. If not, then the keys are misspecified + video_keys = grouped_keys["video"] + assert self.video_concat_order is not None, f"{self.video_concat_order=}, {video_keys=}" + assert all( + item in video_keys for item in self.video_concat_order + ), f"keys in video_concat_order are misspecified, \n{video_keys=}, \n{self.video_concat_order=}" + + # Process each video view + unsqueezed_videos = [] + for video_key in self.video_concat_order: + video_data = data.pop(video_key) + unsqueezed_video = np.expand_dims( + video_data, axis=-4 + ) # [..., H, W, C] -> [..., 1, H, W, C] + unsqueezed_videos.append(unsqueezed_video) + # Concatenate along the new axis + unsqueezed_video = np.concatenate(unsqueezed_videos, axis=-4) # [..., V, H, W, C] + + # Video + data["video"] = unsqueezed_video + + # "state" + if "state" in grouped_keys: + state_keys = grouped_keys["state"] + assert self.state_concat_order is not None, f"{self.state_concat_order=}" + assert all( + item in state_keys for item in self.state_concat_order + ), f"keys in state_concat_order are misspecified, \n{state_keys=}, \n{self.state_concat_order=}" + # Check the state dims + for key in self.state_concat_order: + target_shapes = [self.state_dims[key]] + if self.is_rotation_key(key): + target_shapes.append(6) # Allow for rotation_6d + # if key in ["state.right_arm", "state.right_hand"]: + target_shapes.append(self.state_dims[key] * 2) # Allow for sin-cos transform + assert ( + data[key].shape[-1] in target_shapes + ), f"State dim mismatch for {key=}, {data[key].shape[-1]=}, {target_shapes=}" + # Concatenate the state keys + # We'll have StateActionToTensor before this transform, so here we use torch.cat + data["state"] = torch.cat( + [data.pop(key) for key in self.state_concat_order], dim=-1 + ) # [T, D_state] + + if "action" in grouped_keys: + action_keys = grouped_keys["action"] + assert self.action_concat_order is not None, f"{self.action_concat_order=}" + # Check if all keys in concat_order are present + assert set(self.action_concat_order) == set( + action_keys + ), f"{set(self.action_concat_order)=}, {set(action_keys)=}" + # Record the action dims + for key in self.action_concat_order: + target_shapes = [self.action_dims[key]] + if self.is_rotation_key(key): + target_shapes.append(3) # Allow for axis angle + assert ( + self.action_dims[key] == data[key].shape[-1] + ), f"Action dim mismatch for {key=}, {self.action_dims[key]=}, {data[key].shape[-1]=}" + # Concatenate the action keys + # We'll have StateActionToTensor before this transform, so here we use torch.cat + data["action"] = torch.cat( + [data.pop(key) for key in self.action_concat_order], dim=-1 + ) # [T, D_action] + + return data + + def unapply(self, data: dict) -> dict: + start_dim = 0 + assert "action" in data, f"{data.keys()=}" + # For those dataset without actions (LAPA), we'll never run unapply + assert self.action_concat_order is not None, f"{self.action_concat_order=}" + action_tensor = data.pop("action") + for key in self.action_concat_order: + if key not in self.action_dims: + raise ValueError(f"Action dim {key} not found in action_dims.") + end_dim = start_dim + self.action_dims[key] + data[key] = action_tensor[..., start_dim:end_dim] + start_dim = end_dim + if "state" in data: + assert self.state_concat_order is not None, f"{self.state_concat_order=}" + start_dim = 0 + state_tensor = data.pop("state") + for key in self.state_concat_order: + end_dim = start_dim + self.state_dims[key] + data[key] = state_tensor[..., start_dim:end_dim] + start_dim = end_dim + return data + + def __call__(self, data: dict) -> dict: + return self.apply(data) + + def get_modality_metadata(self, key: str) -> StateActionMetadata: + modality, subkey = key.split(".") + assert self.dataset_metadata is not None, "Metadata not set" + modality_config = getattr(self.dataset_metadata.modalities, modality) + assert subkey in modality_config, f"{subkey=} not found in {modality_config=}" + assert isinstance( + modality_config[subkey], StateActionMetadata + ), f"Expected {StateActionMetadata} for {subkey=}, got {type(modality_config[subkey])=}" + return modality_config[subkey] + + def get_state_action_dims(self, key: str) -> int: + """Get the dimension of a state or action key from the dataset metadata.""" + modality_config = self.get_modality_metadata(key) + shape = modality_config.shape + assert len(shape) == 1, f"{shape=}" + return shape[0] + + def is_rotation_key(self, key: str) -> bool: + modality_config = self.get_modality_metadata(key) + return modality_config.rotation_type is not None + + def set_metadata(self, dataset_metadata: DatasetMetadata): + """Set the metadata and compute the dimensions of the state and action keys.""" + super().set_metadata(dataset_metadata) + # Pre-compute the dimensions of the state and action keys + if self.action_concat_order is not None: + for key in self.action_concat_order: + self.action_dims[key] = self.get_state_action_dims(key) + if self.state_concat_order is not None: + for key in self.state_concat_order: + self.state_dims[key] = self.get_state_action_dims(key) diff --git a/starVLA/dataloader/gr00t_lerobot/transform/state_action.py b/starVLA/dataloader/gr00t_lerobot/transform/state_action.py new file mode 100644 index 0000000000000000000000000000000000000000..7a740b343187fdedfa1215d7a1d07f4604cf689c --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/transform/state_action.py @@ -0,0 +1,626 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import functools +import random +from typing import Any, ClassVar + +import numpy as np +import torch +from pydantic import Field, PrivateAttr, field_validator, model_validator + +try: + import pytorch3d.transforms as pt + _PYTORCH3D_IMPORT_ERROR = None +except ModuleNotFoundError as exc: # PyTorch3D has no official Windows wheels. + if exc.name is None or not exc.name.startswith("pytorch3d"): + raise + pt = None + _PYTORCH3D_IMPORT_ERROR = exc + +from ..schema import DatasetMetadata, RotationType, StateActionMetadata +from .base import InvertibleModalityTransform, ModalityTransform + + +class RotationTransform: + """Adapted from https://github.com/real-stanford/diffusion_policy/blob/548a52bbb105518058e27bf34dcf90bf6f73681a/diffusion_policy/model/common/rotation_transformer.py""" + + valid_reps = ["axis_angle", "euler_angles", "quaternion", "rotation_6d", "matrix"] + + def __init__(self, from_rep="axis_angle", to_rep="rotation_6d"): + """ + Valid representations + + Always use matrix as intermediate representation. + """ + if pt is None: + raise ImportError( + "Rotation conversion requires pytorch3d. Install pytorch3d or use a " + "data configuration without target_rotations (the LIBERO configuration " + "does not require it)." + ) from _PYTORCH3D_IMPORT_ERROR + + if from_rep.startswith("euler_angles"): + from_convention = from_rep.split("_")[-1] + from_rep = "euler_angles" + from_convention = from_convention.replace("r", "X").replace("p", "Y").replace("y", "Z") + else: + from_convention = None + if to_rep.startswith("euler_angles"): + to_convention = to_rep.split("_")[-1] + to_rep = "euler_angles" + to_convention = to_convention.replace("r", "X").replace("p", "Y").replace("y", "Z") + else: + to_convention = None + assert from_rep != to_rep, f"from_rep and to_rep cannot be the same: {from_rep}" + assert from_rep in self.valid_reps, f"Invalid from_rep: {from_rep}" + assert to_rep in self.valid_reps, f"Invalid to_rep: {to_rep}" + + forward_funcs = list() + inverse_funcs = list() + + if from_rep != "matrix": + funcs = [getattr(pt, f"{from_rep}_to_matrix"), getattr(pt, f"matrix_to_{from_rep}")] + if from_convention is not None: + funcs = [functools.partial(func, convention=from_convention) for func in funcs] + forward_funcs.append(funcs[0]) + inverse_funcs.append(funcs[1]) + + if to_rep != "matrix": + funcs = [getattr(pt, f"matrix_to_{to_rep}"), getattr(pt, f"{to_rep}_to_matrix")] + if to_convention is not None: + funcs = [functools.partial(func, convention=to_convention) for func in funcs] + forward_funcs.append(funcs[0]) + inverse_funcs.append(funcs[1]) + + inverse_funcs = inverse_funcs[::-1] + + self.forward_funcs = forward_funcs + self.inverse_funcs = inverse_funcs + + @staticmethod + def _apply_funcs(x: torch.Tensor, funcs: list) -> torch.Tensor: + assert isinstance(x, torch.Tensor) + for func in funcs: + x = func(x) + return x + + def forward(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + return self._apply_funcs(x, self.forward_funcs) + + def inverse(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + return self._apply_funcs(x, self.inverse_funcs) + + +class Normalizer: + valid_modes = ["q99", "mean_std", "min_max", "binary"] + + def __init__(self, mode: str, statistics: dict, binary_threshold: float = 0.5): + self.mode = mode + self.statistics = statistics + self.binary_threshold = binary_threshold + for key, value in self.statistics.items(): + self.statistics[key] = torch.tensor(value) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + + # Normalize the tensor + if self.mode == "q99": + # Range of q99 is [-1, 1] + q01 = self.statistics["q01"].to(x.dtype) + q99 = self.statistics["q99"].to(x.dtype) + + # In the case of q01 == q99, the normalization will be undefined + # So we set the normalized values to the original values + mask = q01 != q99 + normalized = torch.zeros_like(x) + + # Normalize the values where q01 != q99 + # Formula: 2 * (x - q01) / (q99 - q01) - 1 + normalized[..., mask] = (x[..., mask] - q01[..., mask]) / ( + q99[..., mask] - q01[..., mask] + ) + normalized[..., mask] = 2 * normalized[..., mask] - 1 + + # Set the normalized values to the original values where q01 == q99 + normalized[..., ~mask] = x[..., ~mask].to(x.dtype) + + # Clip the normalized values to be between -1 and 1 + normalized = torch.clamp(normalized, -2.2, 2.2) + + elif self.mode == "mean_std": + # Range of mean_std is not fixed, but can be positive or negative + mean = self.statistics["mean"].to(x.dtype) + std = self.statistics["std"].to(x.dtype) + + # In the case of std == 0, the normalization will be undefined + # So we set the normalized values to the original values + mask = std != 0 + normalized = torch.zeros_like(x) + + # Normalize the values where std != 0 + # Formula: (x - mean) / std + normalized[..., mask] = (x[..., mask] - mean[..., mask]) / std[..., mask] + + # Set the normalized values to the original values where std == 0 + normalized[..., ~mask] = x[..., ~mask].to(x.dtype) + + elif self.mode == "min_max": + # Range of min_max is [-1, 1] + min = self.statistics["min"].to(x.dtype) + max = self.statistics["max"].to(x.dtype) + + # In the case of min == max, the normalization will be undefined + # So we set the normalized values to the original values + mask = min != max + normalized = torch.zeros_like(x) + + # Normalize the values where min != max + # Formula: 2 * (x - min) / (max - min) - 1 + normalized[..., mask] = (x[..., mask] - min[..., mask]) / ( + max[..., mask] - min[..., mask] + ) + normalized[..., mask] = 2 * normalized[..., mask] - 1 + + # Set the normalized values to the original values where min == max + # normalized[..., ~mask] = x[..., ~mask].to(x.dtype) + # Set the normalized values to 0 where min == max + normalized[..., ~mask] = 0 + + elif self.mode == "scale": + # Range of scale is [0, 1] + min = self.statistics["min"].to(x.dtype) + max = self.statistics["max"].to(x.dtype) + abs_max = torch.max(torch.abs(min), torch.abs(max)) + mask = abs_max != 0 + normalized = torch.zeros_like(x) + normalized[..., mask] = x[..., mask] / abs_max[..., mask] + normalized[..., ~mask] = 0 + + elif self.mode == "binary": + # Range of binary is [0, 1] + normalized = (x > self.binary_threshold).to(x.dtype) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + return normalized + + def inverse(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + if self.mode == "q99": + q01 = self.statistics["q01"].to(x.dtype) + q99 = self.statistics["q99"].to(x.dtype) + return (x + 1) / 2 * (q99 - q01) + q01 + elif self.mode == "mean_std": + mean = self.statistics["mean"].to(x.dtype) + std = self.statistics["std"].to(x.dtype) + return x * std + mean + elif self.mode == "min_max": + min = self.statistics["min"].to(x.dtype) + max = self.statistics["max"].to(x.dtype) + return (x + 1) / 2 * (max - min) + min + elif self.mode == "binary": + return (x >= self.binary_threshold).to(x.dtype) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + +class StateActionToTensor(InvertibleModalityTransform): + """ + Transforms states and actions to tensors. + """ + + input_dtypes: dict[str, np.dtype] = Field( + default_factory=dict, description="The input dtypes for each state key." + ) + output_dtypes: dict[str, torch.dtype] = Field( + default_factory=dict, description="The output dtypes for each state key." + ) + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = {"apply_to"} + else: + include = kwargs.pop("include", None) + + return super().model_dump(*args, include=include, **kwargs) + + @field_validator("input_dtypes", "output_dtypes", mode="before") + def validate_dtypes(cls, v): + for key, dtype in v.items(): + if isinstance(dtype, str): + if dtype.startswith("torch."): + dtype_split = dtype.split(".")[-1] + v[key] = getattr(torch, dtype_split) + elif dtype.startswith("np.") or dtype.startswith("numpy."): + dtype_split = dtype.split(".")[-1] + v[key] = np.dtype(dtype_split) + else: + raise ValueError(f"Invalid dtype: {dtype}") + return v + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + continue + value = data[key] + assert isinstance( + value, np.ndarray + ), f"Unexpected input type: {type(value)}. Expected type: {np.ndarray}" + data[key] = torch.from_numpy(value) + if key in self.output_dtypes: + data[key] = data[key].to(self.output_dtypes[key]) + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + continue + value = data[key] + assert isinstance( + value, torch.Tensor + ), f"Unexpected input type: {type(value)}. Expected type: {torch.Tensor}" + data[key] = value.numpy() + if key in self.input_dtypes: + data[key] = data[key].astype(self.input_dtypes[key]) + return data + + +class StateActionTransform(InvertibleModalityTransform): + """ + Class for state or action transform. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + normalization_modes (dict[str, str]): The normalization modes for each state key. + If a state key in apply_to is not present in the dictionary, it will not be normalized. + target_rotations (dict[str, str]): The target representations for each state key. + If a state key in apply_to is not present in the dictionary, it will not be rotated. + """ + + # Configurable attributes + apply_to: list[str] = Field(..., description="The keys in the modality to load and transform.") + normalization_modes: dict[str, str] = Field( + default_factory=dict, description="The normalization modes for each state key." + ) + target_rotations: dict[str, str] = Field( + default_factory=dict, description="The target representations for each state key." + ) + normalization_statistics: dict[str, dict] = Field( + default_factory=dict, description="The statistics for each state key." + ) + binary_threshold: float = Field( + default=0.5, description="Threshold for binary normalization mode." + ) + modality_metadata: dict[str, StateActionMetadata] = Field( + default_factory=dict, description="The modality metadata for each state key." + ) + + # Model variables + _rotation_transformers: dict[str, RotationTransform] = PrivateAttr(default_factory=dict) + _normalizers: dict[str, Normalizer] = PrivateAttr(default_factory=dict) + _input_dtypes: dict[str, np.dtype | torch.dtype] = PrivateAttr(default_factory=dict) + + # Model constants + _DEFAULT_MIN_MAX_STATISTICS: ClassVar[dict] = { + "rotation_6d": { + "min": [-1, -1, -1, -1, -1, -1], + "max": [1, 1, 1, 1, 1, 1], + }, + "euler_angles": { + "min": [-np.pi, -np.pi, -np.pi], + "max": [np.pi, np.pi, np.pi], + }, + "quaternion": { + "min": [-1, -1, -1, -1], + "max": [1, 1, 1, 1], + }, + "axis_angle": { + "min": [-np.pi, -np.pi, -np.pi], + "max": [np.pi, np.pi, np.pi], + }, + } + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = {"apply_to", "normalization_modes", "target_rotations"} + else: + include = kwargs.pop("include", None) + + return super().model_dump(*args, include=include, **kwargs) + + @field_validator("modality_metadata", mode="before") + def validate_modality_metadata(cls, v): + for modality_key, config in v.items(): + if isinstance(config, dict): + config = StateActionMetadata.model_validate(config) + else: + assert isinstance( + config, StateActionMetadata + ), f"Invalid source rotation config: {config}" + v[modality_key] = config + return v + + @model_validator(mode="after") + def validate_normalization_statistics(self): + for modality_key, normalization_statistics in self.normalization_statistics.items(): + if modality_key in self.normalization_modes: + normalization_mode = self.normalization_modes[modality_key] + if normalization_mode == "min_max": + assert ( + "min" in normalization_statistics and "max" in normalization_statistics + ), f"Min and max statistics are required for min_max normalization, but got {normalization_statistics}" + assert len(normalization_statistics["min"]) == len( + normalization_statistics["max"] + ), f"Min and max statistics must have the same length, but got {normalization_statistics['min']} and {normalization_statistics['max']}" + elif normalization_mode == "mean_std": + assert ( + "mean" in normalization_statistics and "std" in normalization_statistics + ), f"Mean and std statistics are required for mean_std normalization, but got {normalization_statistics}" + assert len(normalization_statistics["mean"]) == len( + normalization_statistics["std"] + ), f"Mean and std statistics must have the same length, but got {normalization_statistics['mean']} and {normalization_statistics['std']}" + elif normalization_mode == "q99": + assert ( + "q01" in normalization_statistics and "q99" in normalization_statistics + ), f"q01 and q99 statistics are required for q99 normalization, but got {normalization_statistics}" + assert len(normalization_statistics["q01"]) == len( + normalization_statistics["q99"] + ), f"q01 and q99 statistics must have the same length, but got {normalization_statistics['q01']} and {normalization_statistics['q99']}" + elif normalization_mode == "binary": + assert ( + len(normalization_statistics) == 1 + ), f"Binary normalization should only have one value, but got {normalization_statistics}" + assert normalization_statistics[0] in [ + 0, + 1, + ], f"Binary normalization should only have 0 or 1, but got {normalization_statistics[0]}" + else: + raise ValueError(f"Invalid normalization mode: {normalization_mode}") + return self + + def set_metadata(self, dataset_metadata: DatasetMetadata): + dataset_statistics = dataset_metadata.statistics + modality_metadata = dataset_metadata.modalities + + # Check that all state keys specified in apply_to have their modality_metadata + for key in self.apply_to: + split_key = key.split(".", 1) + assert len(split_key) == 2, "State keys should have two parts: 'modality.key'" + if key not in self.modality_metadata: + modality, state_key = split_key + assert hasattr(modality_metadata, modality), f"{modality} config not found" + assert state_key in getattr( + modality_metadata, modality + ), f"{state_key} config not found" + self.modality_metadata[key] = getattr(modality_metadata, modality)[state_key] + + # Check that all state keys specified in normalization_modes have their statistics in state_statistics + for key in self.normalization_modes: + split_key = key.split(".", 1) + assert len(split_key) == 2, "State keys should have two parts: 'modality.key'" + modality, state_key = split_key + assert hasattr(dataset_statistics, modality), f"{modality} statistics not found" + assert state_key in getattr( + dataset_statistics, modality + ), f"{state_key} statistics not found" + assert ( + len(getattr(modality_metadata, modality)[state_key].shape) == 1 + ), f"{getattr(modality_metadata, modality)[state_key].shape=}" + self.normalization_statistics[key] = getattr(dataset_statistics, modality)[ + state_key + ].model_dump() + + # Initialize the rotation transformers + for key in self.target_rotations: + # Get the original representation of the state + from_rep = self.modality_metadata[key].rotation_type + assert from_rep is not None, f"Source rotation type not found for {key}" + + # Get the target representation of the state, will raise an error if the target representation is not valid + to_rep = RotationType(self.target_rotations[key]) + + # If the original representation is not the same as the target representation, initialize the rotation transformer + if from_rep != to_rep: + self._rotation_transformers[key] = RotationTransform( + from_rep=from_rep.value, to_rep=to_rep.value + ) + + # Initialize the normalizers + for key in self.normalization_modes: + modality, state_key = key.split(".", 1) + # If the state has a nontrivial rotation, we need to handle it more carefully + # For absolute rotations, we need to convert them to the target representation and normalize them using min_max mode, + # since we can infer the bounds by the representation + # For relative rotations, we cannot normalize them as we don't know the bounds + if key in self._rotation_transformers: + # Case 1: Absolute rotation + if self.modality_metadata[key].absolute: + # Check that the normalization mode is valid + assert ( + self.normalization_modes[key] == "min_max" + ), "Absolute rotations that are converted to other formats must be normalized using `min_max` mode" + rotation_type = RotationType(self.target_rotations[key]).value + # If the target representation is euler angles, we need to parse the convention + if rotation_type.startswith("euler_angles"): + rotation_type = "euler_angles" + # Get the statistics for the target representation + statistics = self._DEFAULT_MIN_MAX_STATISTICS[rotation_type] + # Case 2: Relative rotation + else: + raise ValueError( + f"Cannot normalize relative rotations: {key} that's converted to {self.target_rotations[key]}" + ) + # If the state is not continuous, we should not use normalization modes other than binary + elif ( + not self.modality_metadata[key].continuous + and self.normalization_modes[key] != "binary" + ): + raise ValueError( + f"{key} is not continuous, so it should be normalized using `binary` mode" + ) + # Initialize the normalizer + else: + statistics = self.normalization_statistics[key] + self._normalizers[key] = Normalizer( + mode=self.normalization_modes[key], statistics=statistics, + binary_threshold=self.binary_threshold, + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + # We allow some keys to be missing in the data, and only process the keys that are present + continue + if key not in self._input_dtypes: + input_dtype = data[key].dtype + assert isinstance( + input_dtype, torch.dtype + ), f"Unexpected input dtype: {input_dtype}. Expected type: {torch.dtype}" + self._input_dtypes[key] = input_dtype + else: + assert ( + data[key].dtype == self._input_dtypes[key] + ), f"All states corresponding to the same key must be of the same dtype, input dtype: {data[key].dtype}, expected dtype: {self._input_dtypes[key]}" + # Rotate the state + state = data[key] + if key in self._rotation_transformers: + state = self._rotation_transformers[key].forward(state) + # Normalize the state + if key in self._normalizers: + state = self._normalizers[key].forward(state) + data[key] = state + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + continue + state = data[key] + assert isinstance( + state, torch.Tensor + ), f"Unexpected state type: {type(state)}. Expected type: {torch.Tensor}" + # Unnormalize the state + if key in self._normalizers: + state = self._normalizers[key].inverse(state) + # Change the state back to its original representation + if key in self._rotation_transformers: + state = self._rotation_transformers[key].inverse(state) + assert isinstance( + state, torch.Tensor + ), f"State should be tensor after unapplying transformations, but got {type(state)}" + # Only convert back to the original dtype if it's known, i.e. `apply` was called before + # If not, we don't know the original dtype, so we don't convert + if key in self._input_dtypes: + original_dtype = self._input_dtypes[key] + if isinstance(original_dtype, np.dtype): + state = state.numpy().astype(original_dtype) + elif isinstance(original_dtype, torch.dtype): + state = state.to(original_dtype) + else: + raise ValueError(f"Invalid input dtype: {original_dtype}") + data[key] = state + return data + + +class StateActionPerturbation(ModalityTransform): + """ + Class for state or action perturbation. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + std (float): Standard deviation of the noise to be added to the state or action. + """ + + # Configurable attributes + std: float = Field( + ..., description="Standard deviation of the noise to be added to the state or action." + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if not self.training: + # Don't perturb the data in eval mode + return data + if self.std < 0: + # If the std is negative, we don't add any noise + return data + for key in self.apply_to: + state = data[key] + assert isinstance(state, torch.Tensor) + transformed_data_min = torch.min(state) + transformed_data_max = torch.max(state) + noise = torch.randn_like(state) * self.std + state += noise + # Clip to the original range + state = torch.clamp(state, transformed_data_min, transformed_data_max) + data[key] = state + return data + + +class StateActionDropout(ModalityTransform): + """ + Class for state or action dropout. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + dropout_prob (float): Probability of dropping out a state or action. + """ + + # Configurable attributes + dropout_prob: float = Field(..., description="Probability of dropping out a state or action.") + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if not self.training: + # Don't drop out the data in eval mode + return data + if self.dropout_prob < 0: + # If the dropout probability is negative, we don't drop out any states + return data + if self.dropout_prob > 1e-9 and random.random() < self.dropout_prob: + for key in self.apply_to: + state = data[key] + assert isinstance(state, torch.Tensor) + state = torch.zeros_like(state) + data[key] = state + return data + + +class StateActionSinCosTransform(ModalityTransform): + """ + Class for state or action sin-cos transform. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + """ + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + state = data[key] + assert isinstance(state, torch.Tensor) + sin_state = torch.sin(state) + cos_state = torch.cos(state) + data[key] = torch.cat([sin_state, cos_state], dim=-1) + return data diff --git a/starVLA/dataloader/gr00t_lerobot/transform/video.py b/starVLA/dataloader/gr00t_lerobot/transform/video.py new file mode 100644 index 0000000000000000000000000000000000000000..15310f697259d7ee6ed8eac7e5abaca211b81dc9 --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/transform/video.py @@ -0,0 +1,612 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Callable, ClassVar, Literal + +import albumentations as A +import cv2 +import numpy as np +import torch +import torchvision.transforms.v2 as T +from einops import rearrange +from pydantic import Field, PrivateAttr, field_validator +from PIL import Image + +from ..schema import DatasetMetadata +from .base import ModalityTransform + + +class VideoTransform(ModalityTransform): + # Configurable attributes + backend: str = Field( + default="torchvision", description="The backend to use for the transformations" + ) + + # Model variables + _train_transform: Callable | None = PrivateAttr(default=None) + _eval_transform: Callable | None = PrivateAttr(default=None) + _original_resolutions: dict[str, tuple[int, int]] = PrivateAttr(default_factory=dict) + + # Model constants + _INTERPOLATION_MAP: ClassVar[dict[str, dict[str, Any]]] = PrivateAttr( + { + "nearest": { + "albumentations": cv2.INTER_NEAREST, + "torchvision": T.InterpolationMode.NEAREST, + }, + "linear": { + "albumentations": cv2.INTER_LINEAR, + "torchvision": T.InterpolationMode.BILINEAR, + }, + "cubic": { + "albumentations": cv2.INTER_CUBIC, + "torchvision": T.InterpolationMode.BICUBIC, + }, + "area": { + "albumentations": cv2.INTER_AREA, + "torchvision": None, # Torchvision does not support this interpolation mode + }, + "lanczos4": { + "albumentations": cv2.INTER_LANCZOS4, # Lanczos with a 4x4 filter + "torchvision": T.InterpolationMode.LANCZOS, # Torchvision does not specify filter size, might be different from 4x4 + }, + "linear_exact": { + "albumentations": cv2.INTER_LINEAR_EXACT, + "torchvision": None, # Torchvision does not support this interpolation mode + }, + "nearest_exact": { + "albumentations": cv2.INTER_NEAREST_EXACT, + "torchvision": T.InterpolationMode.NEAREST_EXACT, + }, + "max": { + "albumentations": cv2.INTER_MAX, + "torchvision": None, + }, + } + ) + + @property + def train_transform(self) -> Callable: + assert ( + self._train_transform is not None + ), "Transform is not set. Please call set_metadata() before calling apply()." + return self._train_transform + + @train_transform.setter + def train_transform(self, value: Callable): + self._train_transform = value + + @property + def eval_transform(self) -> Callable | None: + return self._eval_transform + + @eval_transform.setter + def eval_transform(self, value: Callable | None): + self._eval_transform = value + + @property + def original_resolutions(self) -> dict[str, tuple[int, int]]: + assert ( + self._original_resolutions is not None + ), "Original resolutions are not set. Please call set_metadata() before calling apply()." + return self._original_resolutions + + @original_resolutions.setter + def original_resolutions(self, value: dict[str, tuple[int, int]]): + self._original_resolutions = value + + def check_input(self, data: dict[str, Any]): + if self.backend == "torchvision": + for key in self.apply_to: + assert isinstance(data[key], torch.Tensor), f"Video {key} is not a torch tensor" + assert data[key].ndim in [ + 4, + 5, + ], f"Expected video {key} to have 4 or 5 dimensions (T, C, H, W or T, B, C, H, W), got {data[key].ndim}" + elif self.backend == "albumentations": + for key in self.apply_to: + assert isinstance(data[key], np.ndarray), f"Video {key} is not a numpy array" + assert data[key].ndim in [ + 4, + 5, + ], f"Expected video {key} to have 4 or 5 dimensions (T, C, H, W or T, B, C, H, W), got {data[key].ndim}" + else: + raise ValueError(f"Backend {self.backend} not supported") + + def set_metadata(self, dataset_metadata: DatasetMetadata): + super().set_metadata(dataset_metadata) + self.original_resolutions = {} + for key in self.apply_to: + split_keys = key.split(".") + assert len(split_keys) == 2, f"Invalid key: {key}. Expected format: modality.key" + sub_key = split_keys[1] + if sub_key in dataset_metadata.modalities.video: + self.original_resolutions[key] = dataset_metadata.modalities.video[ + sub_key + ].resolution + else: + raise ValueError( + f"Video key {sub_key} not found in dataset metadata. Available keys: {dataset_metadata.modalities.video.keys()}" + ) + train_transform = self.get_transform(mode="train") + eval_transform = self.get_transform(mode="eval") + if self.backend == "albumentations": + self.train_transform = A.ReplayCompose(transforms=[train_transform]) # type: ignore + if eval_transform is not None: + self.eval_transform = A.ReplayCompose(transforms=[eval_transform]) # type: ignore + else: + assert train_transform is not None, "Train transform must be set" + self.train_transform = train_transform + self.eval_transform = eval_transform + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if self.training: + transform = self.train_transform + else: + transform = self.eval_transform + if transform is None: + return data + assert ( + transform is not None + ), "Transform is not set. Please call set_metadata() before calling apply()." + try: + self.check_input(data) + except AssertionError as e: + raise ValueError( + f"Input data does not match the expected format for {self.__class__.__name__}: {e}" + ) from e + + # Concatenate views + views = [data[key] for key in self.apply_to] + num_views = len(views) + is_batched = views[0].ndim == 5 + bs = views[0].shape[0] if is_batched else 1 + if isinstance(views[0], torch.Tensor): + views = torch.cat(views, 0) + elif isinstance(views[0], np.ndarray): + views = np.concatenate(views, 0) + else: + raise ValueError(f"Unsupported view type: {type(views[0])}") + if is_batched: + views = rearrange(views, "(v b) t c h w -> (v b t) c h w", v=num_views, b=bs) + # Apply the transform + if self.backend == "torchvision": + views = transform(views) + elif self.backend == "albumentations": + assert isinstance(transform, A.ReplayCompose), "Transform must be a ReplayCompose" + first_frame = views[0] + transformed = transform(image=first_frame) + replay_data = transformed["replay"] + transformed_first_frame = transformed["image"] + + if len(views) > 1: + # Apply the same transformations to the rest of the frames + transformed_frames = [ + transform.replay(replay_data, image=frame)["image"] for frame in views[1:] + ] + # Add the first frame back + transformed_frames = [transformed_first_frame] + transformed_frames + else: + # If there is only one frame, just make a list with one frame + transformed_frames = [transformed_first_frame] + + # Delete the replay data to save memory + del replay_data + views = np.stack(transformed_frames, 0) + + else: + raise ValueError(f"Backend {self.backend} not supported") + # Split views + if is_batched: + views = rearrange(views, "(v b t) c h w -> v b t c h w", v=num_views, b=bs) + else: + views = rearrange(views, "(v t) c h w -> v t c h w", v=num_views) + for key, view in zip(self.apply_to, views): + data[key] = view + return data + + @classmethod + def _validate_interpolation(cls, interpolation: str): + if interpolation not in cls._INTERPOLATION_MAP: + raise ValueError(f"Interpolation mode {interpolation} not supported") + + def _get_interpolation(self, interpolation: str, backend: str = "torchvision"): + """ + Get the interpolation mode for the given backend. + + Args: + interpolation (str): The interpolation mode. + backend (str): The backend to use. + + Returns: + Any: The interpolation mode for the given backend. + """ + return self._INTERPOLATION_MAP[interpolation][backend] + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + raise NotImplementedError( + "set_transform is not implemented for VideoTransform. Please implement this function to set the transforms." + ) + + +class VideoCrop(VideoTransform): + height: int | None = Field(default=None, description="The height of the input image") + width: int | None = Field(default=None, description="The width of the input image") + scale: float = Field( + ..., + description="The scale of the crop. The crop size is (width * scale, height * scale)", + ) + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the transform for the given mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: If mode is "train", return a random crop transform. If mode is "eval", return a center crop transform. + """ + # 1. Check the input resolution + assert ( + len(set(self.original_resolutions.values())) == 1 + ), f"All video keys must have the same resolution, got: {self.original_resolutions}" + if self.height is None: + assert self.width is None, "Height and width must be either both provided or both None" + self.width, self.height = self.original_resolutions[self.apply_to[0]] + else: + assert ( + self.width is not None + ), "Height and width must be either both provided or both None" + # 2. Create the transform + size = (int(self.height * self.scale), int(self.width * self.scale)) + if self.backend == "torchvision": + if mode == "train": + return T.RandomCrop(size) + elif mode == "eval": + return T.CenterCrop(size) + else: + raise ValueError(f"Crop mode {mode} not supported") + elif self.backend == "albumentations": + if mode == "train": + return A.RandomCrop(height=size[0], width=size[1], p=1) + elif mode == "eval": + return A.CenterCrop(height=size[0], width=size[1], p=1) + else: + raise ValueError(f"Crop mode {mode} not supported") + else: + raise ValueError(f"Backend {self.backend} not supported") + + def check_input(self, data: dict[str, Any]): + super().check_input(data) + # Check the input resolution + for key in self.apply_to: + if self.backend == "torchvision": + height, width = data[key].shape[-2:] + elif self.backend == "albumentations": + height, width = data[key].shape[-3:-1] + else: + raise ValueError(f"Backend {self.backend} not supported") + assert ( + height == self.height and width == self.width + ), f"Video {key} has invalid shape {height, width}, expected {self.height, self.width}" + + +class VideoResize(VideoTransform): + height: int = Field(..., description="The height of the resize") + width: int = Field(..., description="The width of the resize") + interpolation: str = Field(default="linear", description="The interpolation mode") + antialias: bool = Field(default=True, description="Whether to apply antialiasing") + + @field_validator("interpolation") + def validate_interpolation(cls, v): + cls._validate_interpolation(v) + return v + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the resize transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The resize transform. + """ + interpolation = self._get_interpolation(self.interpolation, self.backend) + if interpolation is None: + raise ValueError( + f"Interpolation mode {self.interpolation} not supported for torchvision" + ) + if self.backend == "torchvision": + size = (self.height, self.width) + return T.Resize(size, interpolation=interpolation, antialias=self.antialias) + elif self.backend == "albumentations": + return A.Resize( + height=self.height, + width=self.width, + interpolation=interpolation, + p=1, + ) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoRandomRotation(VideoTransform): + degrees: float | tuple[float, float] = Field( + ..., description="The degrees of the random rotation" + ) + interpolation: str = Field("linear", description="The interpolation mode") + + @field_validator("interpolation") + def validate_interpolation(cls, v): + cls._validate_interpolation(v) + return v + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the random rotation transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: The random rotation transform. None for eval mode. + """ + if mode == "eval": + return None + interpolation = self._get_interpolation(self.interpolation, self.backend) + if interpolation is None: + raise ValueError( + f"Interpolation mode {self.interpolation} not supported for torchvision" + ) + if self.backend == "torchvision": + return T.RandomRotation(self.degrees, interpolation=interpolation) # type: ignore + elif self.backend == "albumentations": + return A.Rotate(limit=self.degrees, interpolation=interpolation, p=1) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoHorizontalFlip(VideoTransform): + p: float = Field(..., description="The probability of the horizontal flip") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the horizontal flip transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a horizontal flip transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomHorizontalFlip(self.p) + elif self.backend == "albumentations": + return A.HorizontalFlip(p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoGrayscale(VideoTransform): + p: float = Field(..., description="The probability of the grayscale transformation") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the grayscale transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a grayscale transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomGrayscale(self.p) + elif self.backend == "albumentations": + return A.ToGray(p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoColorJitter(VideoTransform): + brightness: float | tuple[float, float] = Field( + ..., description="The brightness of the color jitter" + ) + contrast: float | tuple[float, float] = Field( + ..., description="The contrast of the color jitter" + ) + saturation: float | tuple[float, float] = Field( + ..., description="The saturation of the color jitter" + ) + hue: float | tuple[float, float] = Field(..., description="The hue of the color jitter") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the color jitter transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a color jitter transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.ColorJitter( + brightness=self.brightness, + contrast=self.contrast, + saturation=self.saturation, + hue=self.hue, + ) + elif self.backend == "albumentations": + return A.ColorJitter( + brightness=self.brightness, + contrast=self.contrast, + saturation=self.saturation, + hue=self.hue, + p=1, + ) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoRandomGrayscale(VideoTransform): + p: float = Field(..., description="The probability of the grayscale transformation") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the grayscale transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a grayscale transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomGrayscale(self.p) + elif self.backend == "albumentations": + return A.ToGray(p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoRandomPosterize(VideoTransform): + bits: int = Field(..., description="The number of bits to posterize the image") + p: float = Field(..., description="The probability of the posterize transformation") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the posterize transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a posterize transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomPosterize(bits=self.bits, p=self.p) + elif self.backend == "albumentations": + return A.Posterize(num_bits=self.bits, p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoToTensor(VideoTransform): + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the to tensor transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The to tensor transform. + """ + if self.backend == "torchvision": + return self.__class__.to_tensor + else: + raise ValueError(f"Backend {self.backend} not supported") + + def check_input(self, data: dict): + """Check if the input data has the correct shape. + Expected video shape: [T, H, W, C], dtype np.uint8 + """ + for key in self.apply_to: + assert key in data, f"Key {key} not found in data. Available keys: {data.keys()}" + assert data[key].ndim in [ + 4, + 5, + ], f"Video {key} must have 4 or 5 dimensions, got {data[key].ndim}" + assert ( + data[key].dtype == np.uint8 + ), f"Video {key} must have dtype uint8, got {data[key].dtype}" + input_resolution = data[key].shape[-3:-1][::-1] + if key in self.original_resolutions: + expected_resolution = self.original_resolutions[key] + else: + expected_resolution = input_resolution + assert ( + input_resolution == expected_resolution + ), f"Video {key} has invalid resolution {input_resolution}, expected {expected_resolution}. Full shape: {data[key].shape}" + + @staticmethod + def to_tensor(frames: np.ndarray) -> torch.Tensor: + """Convert numpy array to tensor efficiently. + + Args: + frames: numpy array of shape [T, H, W, C] in uint8 format + Returns: + tensor of shape [T, C, H, W] in range [0, 1] + """ + frames_tensor = torch.from_numpy(frames).to(torch.float32) / 255.0 + return frames_tensor.permute(0, 3, 1, 2) # [T, C, H, W] + + +class VideoToNumpy(VideoTransform): + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the to numpy transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The to numpy transform. + """ + if self.backend == "torchvision": + return self.__class__.to_numpy + else: + raise ValueError(f"Backend {self.backend} not supported") + + @staticmethod + def to_numpy(frames: torch.Tensor) -> np.ndarray: + """Convert tensor back to numpy array efficiently. + + Args: + frames: tensor of shape [T, C, H, W] in range [0, 1] + Returns: + numpy array of shape [T, H, W, C] in uint8 format + """ + return (frames.permute(0, 2, 3, 1) * 255).to(torch.uint8).cpu().numpy() + +class VideoToPIL(VideoTransform): + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the to PIL transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The to PIL transform. + """ + if self.backend == "torchvision": + return self.__class__.to_pil + else: + raise ValueError(f"Backend {self.backend} not supported") + + @staticmethod + def to_pil(frames: torch.Tensor) -> Image.Image: + """Convert tensor back to PIL Image. + + Args: + frames: tensor of shape [T, C, H, W] in range [0, 1] + Returns: + PIL Image of shape [T, H, W, C] in uint8 format + """ + # video PIL format? + return Image.fromarray((frames.permute(0, 2, 3, 1) * 255).to(torch.uint8).cpu().numpy()) \ No newline at end of file diff --git a/starVLA/dataloader/gr00t_lerobot/video.py b/starVLA/dataloader/gr00t_lerobot/video.py new file mode 100644 index 0000000000000000000000000000000000000000..6d61e1c2f53f0c5c4b5c751923922a3df684a3ca --- /dev/null +++ b/starVLA/dataloader/gr00t_lerobot/video.py @@ -0,0 +1,347 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import av +import cv2 +import numpy as np + +import torch # noqa: F401 # isort: skip +import torchvision # noqa: F401 # isort: skip + +# Import decord with graceful fallback +try: + import decord # noqa: F401 + + DECORD_AVAILABLE = True +except ImportError: + DECORD_AVAILABLE = False + +try: + import torchcodec + + TORCHCODEC_AVAILABLE = True +except (ImportError, RuntimeError): + TORCHCODEC_AVAILABLE = False + + +def get_frames_by_indices( + video_path: str, + indices: list[int] | np.ndarray, + video_backend: str = "decord", + video_backend_kwargs: dict = {}, +) -> np.ndarray: + if video_backend == "decord": + if not DECORD_AVAILABLE: + raise ImportError("decord is not available.") + vr = decord.VideoReader(video_path, **video_backend_kwargs) + frames = vr.get_batch(indices) + return frames.asnumpy() + elif video_backend == "torchcodec": + if not TORCHCODEC_AVAILABLE: + raise ImportError("torchcodec is not available.") + decoder = torchcodec.decoders.VideoDecoder( + video_path, device="cpu", dimension_order="NHWC", num_ffmpeg_threads=0 + ) + return decoder.get_frames_at(indices=indices).data.numpy() + elif video_backend == "opencv": + frames = [] + cap = cv2.VideoCapture(video_path, **video_backend_kwargs) + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ret, frame = cap.read() + if not ret: + raise ValueError(f"Unable to read frame at index {idx}") + frames.append(frame) + cap.release() + frames = np.array(frames) + return frames + elif video_backend == "pyav": + # Use PyAV directly without torchvision wrapper + container = None + try: + container = av.open(video_path) + stream = container.streams.video[0] + time_base = float(stream.time_base) + fps = float(stream.average_rate) if stream.average_rate else float(stream.guessed_rate) + + loaded_frames = [] + + for target_idx in indices: + # Convert frame index to timestamp, then to pts + target_ts = target_idx / fps + target_pts = int(target_ts / time_base) + + # Seek to the target position + container.seek(target_pts, stream=stream, backward=True, any_frame=False) + + closest_frame = None + closest_idx_diff = float('inf') + + for frame in container.decode(video=0): + current_ts = float(frame.pts * time_base) + current_idx = int(round(current_ts * fps)) + current_diff = abs(current_idx - target_idx) + + if current_diff < closest_idx_diff: + closest_idx_diff = current_diff + closest_frame = frame + + # If we've passed the target and diff is increasing, stop + if current_idx > target_idx and current_diff > closest_idx_diff: + break + + # Stop if we're significantly past the target + if current_idx > target_idx + 5: + break + + if closest_frame is not None: + frame_data = closest_frame.to_ndarray(format="rgb24") + loaded_frames.append(frame_data) + else: + raise ValueError(f"Unable to find frame at index {target_idx}") + + frames = np.array(loaded_frames) + return frames + + finally: + if container is not None: + container.close() + else: + raise NotImplementedError + + +def get_frames_by_timestamps( + video_path: str, + timestamps: list[float] | np.ndarray, + video_backend: str = "decord", + video_backend_kwargs: dict = {}, +) -> np.ndarray: + """Get frames from a video at specified timestamps. + Args: + video_path (str): Path to the video file. + timestamps (list[int] | np.ndarray): Timestamps to retrieve frames for, in seconds. + video_backend (str, optional): Video backend to use. Defaults to "decord". + Returns: + np.ndarray: Frames at the specified timestamps. + """ + if video_backend == "decord": + # For some GPUs, AV format data cannot be read + if not DECORD_AVAILABLE: + raise ImportError("decord is not available.") + vr = decord.VideoReader(video_path, **video_backend_kwargs) + num_frames = len(vr) + # Retrieve the timestamps for each frame in the video + frame_ts: np.ndarray = vr.get_frame_timestamp(range(num_frames)) + # Map each requested timestamp to the closest frame index + # Only take the first element of the frame_ts array which corresponds to start_seconds + indices = np.abs(frame_ts[:, :1] - timestamps).argmin(axis=0) + frames = vr.get_batch(indices) + return frames.asnumpy() + elif video_backend == "torchcodec": + if not TORCHCODEC_AVAILABLE: + raise ImportError("torchcodec is not available.") + decoder = torchcodec.decoders.VideoDecoder( + video_path, device="cpu", dimension_order="NHWC", num_ffmpeg_threads=0 + ) + return decoder.get_frames_played_at(seconds=timestamps).data.numpy() + elif video_backend == "opencv": + # Open the video file + cap = cv2.VideoCapture(video_path, **video_backend_kwargs) + if not cap.isOpened(): + raise ValueError(f"Unable to open video file: {video_path}") + # Retrieve the total number of frames + num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + # Calculate timestamps for each frame + fps = cap.get(cv2.CAP_PROP_FPS) + frame_ts = np.arange(num_frames) / fps + frame_ts = frame_ts[:, np.newaxis] # Reshape to (num_frames, 1) for broadcasting + # Map each requested timestamp to the closest frame index + indices = np.abs(frame_ts - timestamps).argmin(axis=0) + frames = [] + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ret, frame = cap.read() + if not ret: + raise ValueError(f"Unable to read frame at index {idx}") + frames.append(frame) + cap.release() + frames = np.array(frames) + return frames + elif video_backend == "pyav": + # Use PyAV directly without torchvision wrapper + container = None + try: + container = av.open(video_path) + stream = container.streams.video[0] + + # Get video properties + time_base = float(stream.time_base) + fps = float(stream.average_rate) if stream.average_rate else float(stream.guessed_rate) + duration = float(stream.duration * time_base) if stream.duration else None + + loaded_frames = [] + + for target_ts in timestamps: + # Convert timestamp to pts (presentation timestamp) + target_pts = int(target_ts / time_base) + + # Seek to the target timestamp (seek to keyframe before target) + container.seek(target_pts, stream=stream, backward=True, any_frame=False) + + closest_frame = None + closest_ts_diff = float('inf') + + for frame in container.decode(video=0): + current_ts = float(frame.pts * time_base) + current_diff = abs(current_ts - target_ts) + + if current_diff < closest_ts_diff: + # Release the previous frame + if closest_frame is not None: + del closest_frame + closest_ts_diff = current_diff + closest_frame = frame + + # If we've passed the target and diff is increasing, stop + if current_ts > target_ts and current_diff > closest_ts_diff: + break + + # Also stop if we're significantly past the target (optimization) + if current_ts > target_ts + 1.0: + break + + if closest_frame is not None: + frame_data = closest_frame.to_ndarray(format="rgb24") + loaded_frames.append(frame_data) + del closest_frame + else: + raise ValueError(f"Unable to find frame at timestamp {target_ts}") + + frames = np.array(loaded_frames) + return frames + + finally: + if container is not None: + container.close() + container = None + + elif video_backend == "torchvision_av": + torchvision.set_video_backend("pyav") + loaded_frames = [] + loaded_ts = [] + + reader = None + try: + reader = torchvision.io.VideoReader(video_path, "video") + + for target_ts in timestamps: + # Reset reader state + reader.seek(target_ts, keyframes_only=True) + + closest_frame = None + closest_ts_diff = float('inf') + + for frame in reader: + current_ts = frame["pts"] + current_diff = abs(current_ts - target_ts) + + if closest_frame is None: + closest_frame = frame + + if current_diff < closest_ts_diff: + # Release the previous frame + if closest_frame is not None: + del closest_frame + closest_ts_diff = current_diff + closest_frame = frame + else: + # The time difference starts to increase, stop searching + break + + if closest_frame is not None: + frame_data = closest_frame["data"] + if isinstance(frame_data, torch.Tensor): + frame_data = frame_data.cpu().numpy() + loaded_frames.append(frame_data) + loaded_ts.append(closest_frame["pts"]) + + # Immediately release frame reference + del closest_frame + + finally: + # Thoroughly clean resources + if reader is not None: + if hasattr(reader, '_c'): + reader._c = None + if hasattr(reader, 'container'): + reader.container.close() + reader.container = None + + frames = np.array(loaded_frames) + return frames.transpose(0, 2, 3, 1) + else: + raise NotImplementedError + + +def get_all_frames( + video_path: str, + video_backend: str = "decord", + video_backend_kwargs: dict = {}, + resize_size: tuple[int, int] | None = None, +) -> np.ndarray: + """Get all frames from a video. + Args: + video_path (str): Path to the video file. + video_backend (str, optional): Video backend to use. Defaults to "decord". + video_backend_kwargs (dict, optional): Keyword arguments for the video backend. + resize_size (tuple[int, int], optional): Resize size for the frames. Defaults to None. + """ + if video_backend == "decord": + if not DECORD_AVAILABLE: + raise ImportError("decord is not available.") + vr = decord.VideoReader(video_path, **video_backend_kwargs) + frames = vr.get_batch(range(len(vr))).asnumpy() + elif video_backend == "torchcodec": + if not TORCHCODEC_AVAILABLE: + raise ImportError("torchcodec is not available.") + decoder = torchcodec.decoders.VideoDecoder( + video_path, device="cpu", dimension_order="NHWC", num_ffmpeg_threads=0 + ) + frames = decoder.get_frames_at(indices=range(len(decoder))) + return frames.data.numpy(), frames.pts_seconds.numpy() + elif video_backend == "pyav": + container = av.open(video_path) + frames = [] + for frame in container.decode(video=0): + frame = frame.to_ndarray(format="rgb24") + frames.append(frame) + frames = np.array(frames) + elif video_backend == "torchvision_av": + # set backend and reader + torchvision.set_video_backend("pyav") + reader = torchvision.io.VideoReader(video_path, "video") + frames = [] + for frame in reader: + frames.append(frame["data"].numpy()) + frames = np.array(frames) + frames = frames.transpose(0, 2, 3, 1) + else: + raise NotImplementedError(f"Video backend {video_backend} not implemented") + # resize frames if specified + if resize_size is not None: + frames = [cv2.resize(frame, resize_size) for frame in frames] + frames = np.array(frames) + return frames \ No newline at end of file diff --git a/starVLA/dataloader/infoflow_sampler.py b/starVLA/dataloader/infoflow_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..11991c847fc7e59984648f13d42f5b9382820ff1 --- /dev/null +++ b/starVLA/dataloader/infoflow_sampler.py @@ -0,0 +1,340 @@ +"""Deterministic hierarchical sampling primitives for InfoFlow experiments.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +import hashlib +import json +import os +from pathlib import Path +from typing import Any, Iterator, Sequence + +import numpy as np +from torch.utils.data import Sampler + + +def _identity_files(dataset_root: Path) -> list[Path]: + files = list((dataset_root / "data").rglob("*.parquet")) + files.extend((dataset_root / "videos").rglob("*.mp4")) + meta_root = dataset_root / "meta" + if meta_root.exists(): + for pattern in ("*.json", "*.jsonl", "*.parquet"): + files.extend(meta_root.rglob(pattern)) + return sorted(set(files), key=lambda path: path.relative_to(dataset_root).as_posix()) + + +def compute_dataset_content_fingerprint(dataset_root: str | Path) -> dict[str, Any]: + """Hash immutable LeRobot data, video, and task/episode metadata files once.""" + root = Path(dataset_root) + digest = hashlib.sha256() + entries = [] + total_bytes = 0 + for file_path in _identity_files(root): + relative_path = file_path.relative_to(root).as_posix() + file_size = file_path.stat().st_size + entries.append({"path": relative_path, "bytes": file_size}) + total_bytes += file_size + digest.update(relative_path.encode("utf-8")) + digest.update(str(file_size).encode("ascii")) + with open(file_path, "rb") as data_file: + for chunk in iter(lambda: data_file.read(4 * 1024 * 1024), b""): + digest.update(chunk) + return { + "sha256": digest.hexdigest(), + "file_count": len(entries), + "total_bytes": total_bytes, + "files": entries, + } + + +def validate_dataset_content_inventory( + dataset_root: str | Path, fingerprint: dict[str, Any] +) -> None: + """Cheap per-rank check; full content hashing belongs to fingerprint generation.""" + root = Path(dataset_root) + current_entries = [ + { + "path": path.relative_to(root).as_posix(), + "bytes": path.stat().st_size, + } + for path in _identity_files(root) + ] + if current_entries != fingerprint.get("files"): + raise RuntimeError( + f"Dataset file inventory differs from its content fingerprint: {root}" + ) + if len(current_entries) != fingerprint.get("file_count"): + raise RuntimeError("Dataset content fingerprint has an inconsistent file count") + if sum(entry["bytes"] for entry in current_entries) != fingerprint.get("total_bytes"): + raise RuntimeError("Dataset content fingerprint has an inconsistent byte count") + + +def _normalize_probabilities(values: Sequence[float], name: str) -> tuple[float, ...]: + probabilities = np.asarray(values, dtype=np.float64) + if probabilities.ndim != 1 or probabilities.size == 0: + raise ValueError(f"{name} must be a non-empty one-dimensional sequence") + if not np.all(np.isfinite(probabilities)) or np.any(probabilities <= 0): + raise ValueError(f"{name} must contain finite, strictly positive values") + probabilities /= probabilities.sum(dtype=np.float64) + return tuple(float(value) for value in probabilities) + + +def _stateless_seed(seed: int, epoch: int, global_draw_id: int) -> int: + payload = f"{int(seed)}:{int(epoch)}:{int(global_draw_id)}".encode("utf-8") + return int.from_bytes(hashlib.sha256(payload).digest()[:8], "little", signed=False) + + +@dataclass(frozen=True) +class EpisodeSupport: + trajectory_index: int + trajectory_id: int | str + anchors: tuple[int, ...] + + def __post_init__(self) -> None: + if not self.anchors: + raise ValueError("episode support must contain at least one anchor") + if len(set(self.anchors)) != len(self.anchors): + raise ValueError("episode anchors must be unique") + + +@dataclass(frozen=True) +class DatasetSupport: + name: str + probability: float + episodes: tuple[EpisodeSupport, ...] + episode_probabilities: tuple[float, ...] + + def __post_init__(self) -> None: + if not self.episodes: + raise ValueError(f"dataset {self.name!r} has no valid episodes") + if len(self.episodes) != len(self.episode_probabilities): + raise ValueError("episode support and probability lengths differ") + normalized = _normalize_probabilities(self.episode_probabilities, "episode_probabilities") + if not np.allclose(normalized, self.episode_probabilities, rtol=0, atol=1e-12): + raise ValueError("episode probabilities must already sum to one") + + +@dataclass(frozen=True) +class LogicalDraw: + epoch: int + global_draw_id: int + + def __post_init__(self) -> None: + if self.epoch < 0 or self.global_draw_id < 0: + raise ValueError("epoch and global_draw_id must be non-negative") + + +@dataclass(frozen=True) +class SampleTrace: + manifest_sha256: str + epoch: int + global_draw_id: int + dataset_index: int + dataset_name: str + trajectory_index: int + trajectory_id: int | str + anchor_index: int + base_index: int + baseline_probability: float + proposal_sha256: str | None = None + proposal_probability: float | None = None + importance_weight: float = 1.0 + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +class ValidAnchorManifest: + schema_version = 1 + + def __init__(self, datasets: Sequence[DatasetSupport], identity: dict[str, Any] | None = None): + if not datasets: + raise ValueError("manifest must contain at least one dataset") + dataset_probabilities = _normalize_probabilities( + [dataset.probability for dataset in datasets], "dataset_probabilities" + ) + if not np.allclose( + dataset_probabilities, + [dataset.probability for dataset in datasets], + rtol=0, + atol=1e-12, + ): + raise ValueError("dataset probabilities must already sum to one") + self.datasets = tuple(datasets) + self.identity = identity or {} + self._dataset_probabilities = np.asarray( + [dataset.probability for dataset in self.datasets], dtype=np.float64 + ) + self._episode_probabilities = tuple( + np.asarray(dataset.episode_probabilities, dtype=np.float64) + for dataset in self.datasets + ) + canonical = json.dumps(self.to_dict(include_hash=False), sort_keys=True, separators=(",", ":")) + self.sha256 = hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + @classmethod + def from_hierarchy( + cls, + dataset_names: Sequence[str], + dataset_weights: Sequence[float], + episodes: Sequence[Sequence[tuple[int | str, Sequence[int], float]]], + identity: dict[str, Any] | None = None, + ) -> "ValidAnchorManifest": + if len(dataset_names) != len(dataset_weights) or len(dataset_names) != len(episodes): + raise ValueError("dataset names, weights, and episode supports must have equal lengths") + dataset_probabilities = _normalize_probabilities(dataset_weights, "dataset_weights") + supports = [] + for dataset_index, (name, dataset_episodes) in enumerate(zip(dataset_names, episodes)): + if not dataset_episodes: + raise ValueError(f"dataset {name!r} has no valid episodes") + episode_probabilities = _normalize_probabilities( + [entry[2] for entry in dataset_episodes], f"episode_weights[{dataset_index}]" + ) + episode_supports = tuple( + EpisodeSupport( + trajectory_index=trajectory_index, + trajectory_id=trajectory_id, + anchors=tuple(int(anchor) for anchor in anchors), + ) + for trajectory_index, (trajectory_id, anchors, _) in enumerate(dataset_episodes) + ) + supports.append( + DatasetSupport( + name=str(name), + probability=dataset_probabilities[dataset_index], + episodes=episode_supports, + episode_probabilities=episode_probabilities, + ) + ) + return cls(supports, identity=identity) + + @classmethod + def from_dict(cls, payload: dict[str, Any]) -> "ValidAnchorManifest": + datasets = [] + for dataset_payload in payload["datasets"]: + episodes = tuple( + EpisodeSupport( + trajectory_index=int(episode["trajectory_index"]), + trajectory_id=episode["trajectory_id"], + anchors=tuple(int(anchor) for anchor in episode["anchors"]), + ) + for episode in dataset_payload["episodes"] + ) + datasets.append( + DatasetSupport( + name=str(dataset_payload["name"]), + probability=float(dataset_payload["probability"]), + episodes=episodes, + episode_probabilities=tuple( + float(value) for value in dataset_payload["episode_probabilities"] + ), + ) + ) + manifest = cls(datasets, identity=payload.get("identity", {})) + return manifest + + @classmethod + def load(cls, path: str | Path) -> "ValidAnchorManifest": + with open(path, "r", encoding="utf-8") as manifest_file: + return cls.from_dict(json.load(manifest_file)) + + def save(self, path: str | Path) -> None: + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(destination.suffix + ".tmp") + with open(temporary, "w", encoding="utf-8") as manifest_file: + json.dump(self.to_dict(), manifest_file, indent=2) + os.replace(temporary, destination) + + def sample(self, seed: int, draw: LogicalDraw) -> SampleTrace: + rng = np.random.default_rng(_stateless_seed(seed, draw.epoch, draw.global_draw_id)) + dataset_index = int(rng.choice(len(self.datasets), p=self._dataset_probabilities)) + dataset = self.datasets[dataset_index] + trajectory_index = int( + rng.choice(len(dataset.episodes), p=self._episode_probabilities[dataset_index]) + ) + episode = dataset.episodes[trajectory_index] + anchor_index = int(rng.integers(len(episode.anchors))) + baseline_probability = ( + dataset.probability + * dataset.episode_probabilities[trajectory_index] + / len(episode.anchors) + ) + return SampleTrace( + manifest_sha256=self.sha256, + epoch=int(draw.epoch), + global_draw_id=int(draw.global_draw_id), + dataset_index=dataset_index, + dataset_name=dataset.name, + trajectory_index=episode.trajectory_index, + trajectory_id=episode.trajectory_id, + anchor_index=anchor_index, + base_index=episode.anchors[anchor_index], + baseline_probability=float(baseline_probability), + ) + + def to_dict(self, include_hash: bool = True) -> dict[str, Any]: + payload = { + "schema_version": self.schema_version, + "identity": self.identity, + "datasets": [asdict(dataset) for dataset in self.datasets], + } + if include_hash: + payload["sha256"] = self.sha256 + return payload + + +class StatelessDrawSampler(Sampler[LogicalDraw]): + """Yield rank-disjoint logical draw ids; sampling itself stays inside the dataset.""" + + def __init__( + self, + num_local_draws: int, + rank: int = 0, + world_size: int = 1, + start_global_draw_id: int = 0, + ) -> None: + if num_local_draws < 0: + raise ValueError("num_local_draws must be non-negative") + if world_size <= 0 or rank < 0 or rank >= world_size: + raise ValueError("rank must satisfy 0 <= rank < world_size") + self.num_local_draws = int(num_local_draws) + self.rank = int(rank) + self.world_size = int(world_size) + self.start_global_draw_id = int(start_global_draw_id) + if self.start_global_draw_id < 0: + raise ValueError("start_global_draw_id must be non-negative") + self.epoch = 0 + self._resume_epoch = 0 + + def set_resume_position(self, epoch: int, start_global_draw_id: int) -> None: + """Resume at an inclusive global draw cursor within one epoch.""" + if epoch < 0 or start_global_draw_id < 0: + raise ValueError("resume epoch and draw id must be non-negative") + self.epoch = int(epoch) + self._resume_epoch = int(epoch) + self.start_global_draw_id = int(start_global_draw_id) + + def set_epoch(self, epoch: int) -> None: + if epoch < 0: + raise ValueError("epoch must be non-negative") + self.epoch = int(epoch) + if self.epoch != self._resume_epoch: + self.start_global_draw_id = 0 + + def _first_local_draw_id(self) -> int: + return max( + 0, + (self.start_global_draw_id - self.rank + self.world_size - 1) + // self.world_size, + ) + + def __iter__(self) -> Iterator[LogicalDraw]: + for local_draw_id in range(self._first_local_draw_id(), self.num_local_draws): + yield LogicalDraw( + epoch=self.epoch, + global_draw_id=local_draw_id * self.world_size + self.rank, + ) + + def __len__(self) -> int: + return max(0, self.num_local_draws - self._first_local_draw_id()) diff --git a/starVLA/dataloader/lerobot_datasets.py b/starVLA/dataloader/lerobot_datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..6e45cbe8c5a9145a2a56ed4a314c7e664d18547f --- /dev/null +++ b/starVLA/dataloader/lerobot_datasets.py @@ -0,0 +1,163 @@ +# Copyright 2025 NVIDIA Corp. and affiliates. All rights reserved. +# Modified by [Fangjing Wang/ SUST University] in [2025]. +# Modification: [return raw data and suport multi-dataset mixture]. +# Modified by [Jinhui YE/ HKUST University] in [2025]. +# Modification: [suport topdowm processing, suport param from config]. + +import logging +from pathlib import Path +from typing import Sequence +from omegaconf import OmegaConf + +from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotSingleDataset, LeRobotMixtureDataset +from starVLA.dataloader.gr00t_lerobot.registry import ( + ROBOT_TYPE_CONFIG_MAP, + DATASET_NAMED_MIXTURES, + EmbodimentTag, +) + +logger = logging.getLogger(__name__) + +def collate_fn(batch): + return batch + +def make_LeRobotSingleDataset( + data_root_dir: Path | str, + data_name: str, + robot_type: str, + delete_pause_frame: bool = False, + data_cfg: dict | None = None, +) -> LeRobotSingleDataset: + """ + Make a LeRobotSingleDataset object. + + :param data_root_dir: The root directory of the dataset. + :param data_name: The name of the dataset. + :param robot_type: The robot type config to use. + :param crop_obs_camera: Whether to crop the observation camera images. + :return: A LeRobotSingleDataset object. + """ + + data_config = ROBOT_TYPE_CONFIG_MAP[robot_type] + modality_config = data_config.modality_config() + transforms = data_config.transform() + dataset_path = data_root_dir / data_name + embodiment_tag = getattr(data_config, "embodiment_tag", None) + if embodiment_tag is None: + print(f"Warning: DataConfig for robot_type={robot_type!r} has no embodiment_tag, using {EmbodimentTag.NEW_EMBODIMENT} as default") + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + + video_backend = data_cfg.get("video_backend", "decord") if data_cfg else "torchvision_av" + + # Opt-in factory hook: a DataConfig may define ``make_dataset(dataset_name=..., **ds_kwargs)`` + # to swap in a custom dataset class (e.g. with per-task filtering / chunk stride). + # When absent, fall through to the default LeRobotSingleDataset construction below. + if hasattr(data_config, "make_dataset"): + return data_config.make_dataset( + dataset_path=dataset_path, + modality_configs=modality_config, + transforms=transforms, + embodiment_tag=embodiment_tag, + video_backend=video_backend, + delete_pause_frame=delete_pause_frame, + data_cfg=data_cfg, + dataset_name=data_name, + ) + + return LeRobotSingleDataset( + dataset_path=dataset_path, + modality_configs=modality_config, + transforms=transforms, + embodiment_tag=embodiment_tag, + video_backend=video_backend, # decord is more efficiency | torchvision_av for video.av1 + delete_pause_frame=delete_pause_frame, + data_cfg=data_cfg, + dataset_name=data_name, + ) + +def get_vla_dataset( + data_cfg: dict, + mode: str = "train", + balance_dataset_weights: bool = False, + balance_trajectory_weights: bool = False, + seed: int = 42, + **kwargs: dict, +) -> LeRobotMixtureDataset: + """ + Get a LeRobotMixtureDataset object. + """ + data_root_dir = data_cfg.data_root_dir + data_mix = data_cfg.data_mix + delete_pause_frame = data_cfg.get("delete_pause_frame", False) + mixture_spec = DATASET_NAMED_MIXTURES[data_mix] + logger.info(f"[dataloader] Using mixture '{data_mix}': {[(d, w, r) for d, w, r in mixture_spec]}") + included_datasets, filtered_mixture_spec = set(), [] + for d_name, d_weight, robot_type in mixture_spec: + dataset_key = (d_name, robot_type) + if dataset_key in included_datasets: + print(f"Skipping Duplicate Dataset: `{(d_name, d_weight, robot_type)}`") + continue + + included_datasets.add(dataset_key) + filtered_mixture_spec.append((d_name, d_weight, robot_type)) + + dataset_mixture = [] + for d_name, d_weight, robot_type in filtered_mixture_spec: + dataset_mixture.append((make_LeRobotSingleDataset(Path(data_root_dir), d_name, robot_type, delete_pause_frame=delete_pause_frame, data_cfg=data_cfg), d_weight)) + + return LeRobotMixtureDataset( + dataset_mixture, + mode=mode, + balance_dataset_weights=balance_dataset_weights, + balance_trajectory_weights=balance_trajectory_weights, + seed=seed, + data_cfg=data_cfg, + **kwargs, + ) + + + +if __name__ == "__main__": + import argparse + import os + + parser = argparse.ArgumentParser() + parser.add_argument("--config_yaml", type=str, default="./examples/LIBERO/train_files/bar/starvla_cotrain_libero.yaml", help="Path to YAML config") + parser.add_argument("--data_mix", type=str, default=None, help="Override data_mix from config") + parser.add_argument("--data_root_dir", type=str, default=None, help="Override data_root_dir from config") + args = parser.parse_args() + + if os.getenv("DEBUGPY_ENABLE", "0") == "1": + import debugpy + debugpy.listen(("0.0.0.0", 10092)) + print("Rank 0 waiting for debugger attach on port 10092...") + debugpy.wait_for_client() + + cfg = OmegaConf.load(args.config_yaml) + vla_dataset_cfg = cfg.datasets.vla_data + vla_dataset_cfg.data_root_dir = Path(vla_dataset_cfg.data_root_dir) + if args.data_mix is not None: + vla_dataset_cfg.data_mix = args.data_mix + if args.data_root_dir is not None: + vla_dataset_cfg.data_root_dir = Path(args.data_root_dir) + + dataset = get_vla_dataset(data_cfg=vla_dataset_cfg) + from torch.utils.data import DataLoader + train_dataloader = DataLoader( + dataset, + batch_size=2, + num_workers=1, # For Debug + collate_fn=collate_fn, + ) + + cfg.output_dir = "./results/debug" + output_dir = Path(cfg.output_dir) + dataset.save_dataset_statistics(output_dir / "dataset_statistics.json") + + from tqdm import tqdm + count = 0 + for batch in tqdm(train_dataloader, desc="Processing Batches"): + if count > 3: + break + count += 1 + pass \ No newline at end of file diff --git a/starVLA/dataloader/temporal_proposal.py b/starVLA/dataloader/temporal_proposal.py new file mode 100644 index 0000000000000000000000000000000000000000..4ed3fca3da099ebf45af25d29972d6000356c21f --- /dev/null +++ b/starVLA/dataloader/temporal_proposal.py @@ -0,0 +1,229 @@ +"""Offline temporal proposals that preserve StarVLA's dataset/episode marginals.""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +from typing import Any, Sequence + +import numpy as np + +from starVLA.dataloader.infoflow_sampler import ( + LogicalDraw, + SampleTrace, + ValidAnchorManifest, + _stateless_seed, +) + + +class TemporalProposal: + schema_version = 1 + valid_score_kinds = { + "synthetic", + "random", + "action_change", + "residual", + "likelihood", + } + + def __init__( + self, + manifest: ValidAnchorManifest, + alpha: float, + anchor_probabilities: Sequence[Sequence[Sequence[float]]], + score_kind: str, + scorer_sha256: str, + validation_certificate: dict[str, Any] | None = None, + legacy_implementation_sha256: str | None = None, + ) -> None: + if not 0.0 <= alpha < 1.0: + raise ValueError("alpha must satisfy 0 <= alpha < 1") + if score_kind not in self.valid_score_kinds: + raise ValueError(f"unsupported score_kind: {score_kind}") + if score_kind == "likelihood" and not validation_certificate: + raise ValueError("likelihood proposals require a validation certificate") + if not scorer_sha256: + raise ValueError("scorer_sha256 is required") + if len(anchor_probabilities) != len(manifest.datasets): + raise ValueError("proposal dataset count does not match the manifest") + + normalized_probabilities = [] + importance_weights = [] + normalized_masses = [] + for dataset_index, (dataset, dataset_probabilities) in enumerate( + zip(manifest.datasets, anchor_probabilities) + ): + if len(dataset_probabilities) != len(dataset.episodes): + raise ValueError( + f"proposal episode count differs for dataset {dataset_index}" + ) + normalized_dataset = [] + dataset_weights = [] + normalized_dataset_masses = [] + for episode_index, (episode, proposal_mass) in enumerate( + zip(dataset.episodes, dataset_probabilities) + ): + mass = np.asarray(proposal_mass, dtype=np.float64) + if mass.shape != (len(episode.anchors),): + raise ValueError( + f"proposal anchor count differs for dataset {dataset_index}, " + f"episode {episode_index}" + ) + if not np.all(np.isfinite(mass)) or np.any(mass < 0) or mass.sum() <= 0: + raise ValueError("proposal masses must be finite, non-negative, and non-zero") + residual_probability = mass / mass.sum(dtype=np.float64) + baseline_probability = np.full(len(mass), 1.0 / len(mass), dtype=np.float64) + safe_probability = (1.0 - alpha) * baseline_probability + alpha * residual_probability + weights = baseline_probability / safe_probability + if not np.isclose(safe_probability.sum(dtype=np.float64), 1.0, atol=1e-12): + raise ValueError("safe temporal proposal does not sum to one") + if np.max(weights) > 1.0 / (1.0 - alpha) + 1e-12: + raise ValueError("importance weight exceeds the analytical safe-mixture bound") + normalized_dataset.append(safe_probability) + dataset_weights.append(weights) + normalized_dataset_masses.append(mass.copy()) + normalized_probabilities.append(tuple(normalized_dataset)) + importance_weights.append(tuple(dataset_weights)) + normalized_masses.append(tuple(normalized_dataset_masses)) + + self.manifest_sha256 = manifest.sha256 + self._manifest_support = manifest.to_dict(include_hash=False)["datasets"] + # Schema-v1 proposals included a source-file hash in their content hash. + # Preserve that value only when loading an existing artifact; it is not + # compared with the current source and new proposals do not emit it. + self.legacy_implementation_sha256 = legacy_implementation_sha256 + self.alpha = float(alpha) + self.score_kind = score_kind + self.scorer_sha256 = scorer_sha256 + self.validation_certificate = validation_certificate + self.proposal_masses = tuple(normalized_masses) + self.anchor_probabilities = tuple(normalized_probabilities) + self.importance_weights = tuple(importance_weights) + canonical = json.dumps(self.to_dict(include_hash=False), sort_keys=True, separators=(",", ":")) + self.sha256 = hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + @classmethod + def from_masses( + cls, + manifest: ValidAnchorManifest, + alpha: float, + masses: Sequence[Sequence[Sequence[float]]], + score_kind: str = "synthetic", + scorer_sha256: str = "synthetic-fixture", + validation_certificate: dict[str, Any] | None = None, + ) -> "TemporalProposal": + return cls( + manifest=manifest, + alpha=alpha, + anchor_probabilities=masses, + score_kind=score_kind, + scorer_sha256=scorer_sha256, + validation_certificate=validation_certificate, + ) + + @classmethod + def from_dict( + cls, manifest: ValidAnchorManifest, payload: dict[str, Any] + ) -> "TemporalProposal": + if payload.get("schema_version") != cls.schema_version: + raise ValueError( + f"unsupported temporal proposal schema: {payload.get('schema_version')}" + ) + proposal = cls( + manifest=manifest, + alpha=float(payload["alpha"]), + anchor_probabilities=payload["proposal_masses"], + score_kind=str(payload["score_kind"]), + scorer_sha256=str(payload["scorer_sha256"]), + validation_certificate=payload.get("validation_certificate"), + legacy_implementation_sha256=payload.get("implementation_sha256"), + ) + return proposal + + @classmethod + def load( + cls, path: str | Path, manifest: ValidAnchorManifest + ) -> "TemporalProposal": + with open(path, "r", encoding="utf-8") as proposal_file: + return cls.from_dict(manifest, json.load(proposal_file)) + + def save(self, path: str | Path) -> None: + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(destination.suffix + ".tmp") + with open(temporary, "w", encoding="utf-8") as proposal_file: + json.dump(self.to_dict(), proposal_file, indent=2) + os.replace(temporary, destination) + + def sample( + self, manifest: ValidAnchorManifest, seed: int, draw: LogicalDraw + ) -> SampleTrace: + if manifest.to_dict(include_hash=False)["datasets"] != self._manifest_support: + raise ValueError("temporal proposal support differs from the active manifest") + rng = np.random.default_rng(_stateless_seed(seed, draw.epoch, draw.global_draw_id)) + dataset_index = int(rng.choice(len(manifest.datasets), p=manifest._dataset_probabilities)) + dataset = manifest.datasets[dataset_index] + trajectory_index = int( + rng.choice( + len(dataset.episodes), p=manifest._episode_probabilities[dataset_index] + ) + ) + episode = dataset.episodes[trajectory_index] + if self.alpha == 0.0: + # Preserve the exact baseline RNG path, not merely its distribution. + anchor_index = int(rng.integers(len(episode.anchors))) + else: + anchor_index = int( + rng.choice( + len(episode.anchors), + p=self.anchor_probabilities[dataset_index][trajectory_index], + ) + ) + baseline_probability = ( + dataset.probability + * dataset.episode_probabilities[trajectory_index] + / len(episode.anchors) + ) + proposal_probability = ( + dataset.probability + * dataset.episode_probabilities[trajectory_index] + * self.anchor_probabilities[dataset_index][trajectory_index][anchor_index] + ) + return SampleTrace( + manifest_sha256=manifest.sha256, + epoch=draw.epoch, + global_draw_id=draw.global_draw_id, + dataset_index=dataset_index, + dataset_name=dataset.name, + trajectory_index=episode.trajectory_index, + trajectory_id=episode.trajectory_id, + anchor_index=anchor_index, + base_index=episode.anchors[anchor_index], + baseline_probability=float(baseline_probability), + proposal_sha256=self.sha256, + proposal_probability=float(proposal_probability), + importance_weight=float( + self.importance_weights[dataset_index][trajectory_index][anchor_index] + ), + ) + + def to_dict(self, include_hash: bool = True) -> dict[str, Any]: + payload = { + "schema_version": self.schema_version, + "manifest_sha256": self.manifest_sha256, + "alpha": self.alpha, + "score_kind": self.score_kind, + "scorer_sha256": self.scorer_sha256, + "validation_certificate": self.validation_certificate, + "proposal_masses": [ + [probabilities.tolist() for probabilities in dataset] + for dataset in self.proposal_masses + ], + } + if self.legacy_implementation_sha256 is not None: + payload["implementation_sha256"] = self.legacy_implementation_sha256 + if include_hash: + payload["sha256"] = self.sha256 + return payload diff --git a/starVLA/model/framework/VLM4A/QwenGR00T.py b/starVLA/model/framework/VLM4A/QwenGR00T.py new file mode 100644 index 0000000000000000000000000000000000000000..f1ab3786d58307c28020332841b9cdf375db7305 --- /dev/null +++ b/starVLA/model/framework/VLM4A/QwenGR00T.py @@ -0,0 +1,332 @@ +# Copyright 2025 starVLA community. All rights reserved. +# Licensed under the MIT License, Version 1.0 (the "License"); +# Implemented by [Junqiu YU / Fudan University] in [2025]. +# Design and Merged by [Jinhui YE / HKUST University] in [2025]. +""" +Qwen-GR00T Framework +A lightweight implementation that Qwen-VL + Flow-matching head to directly predict continuous actions +Flow-matching header is copyright from GR00T N1.5, +""" + +import sys +from pathlib import Path + +# Add workspace root to Python path if not already there +_workspace_root = Path(__file__).parent.parent.parent.parent.parent +if str(_workspace_root) not in sys.path: + sys.path.insert(0, str(_workspace_root)) + +from dataclasses import dataclass, field +from typing import List, Optional, Tuple + +import numpy as np +import torch +from PIL import Image + +from deployment.model_server.tools.image_tools import to_pil_preserve +from starVLA.training.infoflow import aggregate_repeated_losses +from starVLA.training.trainer_utils import initialize_overwatch + +logger = initialize_overwatch(__name__) + +# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels) +IGNORE_INDEX = -100 + + +from starVLA.model.framework.base_framework import baseframework +from starVLA.model.framework.share_tools import merge_framework_config +from starVLA.model.modules.action_model.GR00T_ActionHeader import FlowmatchingActionHead, get_action_model +from starVLA.model.modules.vlm import get_vlm_model +from starVLA.model.tools import FRAMEWORK_REGISTRY +from starVLA.training.trainer_utils.trainer_tools import resize_images + + +# ────────────────────────────────────────────────────────────────────── +# Default Config for QwenGR00T +# - Documents every framework-level parameter with type + description +# - YAML values override these defaults; extra YAML keys are preserved +# ────────────────────────────────────────────────────────────────────── +@dataclass +class QwenGR00TDefaultConfig: + """QwenGR00T framework default parameters. + + All fields can be overridden by the corresponding key in the YAML + ``framework:`` section. Extra YAML keys not listed here are kept + as-is (Config-as-API flexibility). + """ + + # --- Registry identifier --- + name: str = "QwenGR00T" + + # === VLM backbone (Qwen2.5-VL / Qwen3-VL) === + qwenvl: dict = field( + default_factory=lambda: { + # Path to base VLM checkpoint (local or HF hub id) + "base_vlm": "./playground/Pretrained_models/Qwen3-VL-4B-Instruct", + # Attention implementation: "flash_attention_2" | "eager" | "sdpa" + "attn_implementation": "flash_attention_2", + # VLM hidden dimension (used for cross-attention alignment) + "vl_hidden_dim": 2048, + } + ) + + # # === DINO encoder (optional multi-view spatial tokens) === Dino is not used in this QwenGR00T version, we can add it later when we want to use it + # dino: dict = field(default_factory=lambda: { + # # DINO backbone variant: "dinov2_vits14" | "dinov2_vitb14" | ... + # "dino_backbone": "dinov2_vits14", + # }) + + # === Action head (Flow-matching / DiT diffusion) === + action_model: dict = field( + default_factory=lambda: { + # DiT model size: "DiT-B" | "DiT-L" | "DiT-XL" + "action_model_type": "DiT-B", + # Hidden dim for action model (auto-aligned at runtime) + "action_hidden_dim": 1024, + "hidden_size": 1024, + # Whether to add positional embeddings in the action head + "add_pos_embed": True, + "max_seq_len": 1024, + # Dimensionality of each action vector (e.g., 7 for 6-DoF + gripper) + "action_dim": 7, + # State dimension (proprioception input) + "state_dim": 7, + # Canonical chunk length (number of action steps the head predicts). + # Legacy YAMLs may use future_action_window_size = action_horizon - 1; + # apply_config_compat normalises both directions. + "action_horizon": 8, + # Repeat factor for flow-matching loss (more noise samples per batch) + "repeated_diffusion_steps": 8, + # Beta distribution params for noise schedule + "noise_beta_alpha": 1.5, + "noise_beta_beta": 1.0, + "noise_s": 0.999, + "num_timestep_buckets": 1000, + # Inference denoising steps + "num_inference_timesteps": 4, + # Number of vision tokens fed to action head + "num_target_vision_tokens": 32, + # === DiT Transformer sub-config === + "diffusion_model_cfg": { + # Cross-attention dim (aligned to VLM hidden_size at runtime) + "cross_attention_dim": 2048, + "dropout": 0.2, + "final_dropout": True, + "interleave_self_attention": True, + "norm_type": "ada_norm", + "num_layers": 16, + "output_dim": 1024, + "positional_embeddings": None, + }, + } + ) + + # # === Training precision flag === This is unnecessary, unused parameter + # reduce_in_full_precision: bool = True + + +@FRAMEWORK_REGISTRY.register("QwenGR00T") +class Qwen_GR00T(baseframework): + """ + Multimodal vision-language-action model (GR00T variant). + + Components: + - Qwen2.5-VL / Qwen3-VL backbone for fused language/vision token embeddings + - Flow-matching (DiT) diffusion head for continuous action sequence modeling + + Focus: Predict future continuous actions conditioned on images + instruction. + """ + + def __init__( + self, + config: Optional[dict] = None, + **kwargs, + ) -> None: + """ + Construct all submodules and cache key configuration values. + + Args: + config: Hierarchical configuration (OmegaConf/dict) containing framework + trainer sections. + **kwargs: Reserved for future overrides (unused). + """ + super().__init__() + # Merge framework defaults with YAML config (YAML wins on conflicts) + self.config = merge_framework_config(QwenGR00TDefaultConfig, config) + self.qwen_vl_interface = get_vlm_model(config=self.config) + # Florence conditions the DiT through its 1024-wide projection. + self.config.framework.action_model.diffusion_model_cfg.cross_attention_dim = ( + self.qwen_vl_interface.model.config.hidden_size + ) + + self.action_model: FlowmatchingActionHead = get_action_model(config=self.config) + # ZeRO-3 partitions all trainable parameters together, so the freshly + # initialized action head must match the BF16 Qwen backbone. + qwen_dtype = next(self.qwen_vl_interface.model.parameters()).dtype + self.action_model.to(dtype=qwen_dtype) + + # `action_horizon` is the single source of truth for chunk length. + # Legacy aliases (`future_action_window_size`, `past_action_window_size`) + # are normalised upstream by `share_tools.apply_config_compat`, so we + # only ever read `action_horizon` here. + self.action_horizon = int(self.config.framework.action_model.action_horizon) + + def forward( + self, + examples: List[dict] = None, + **kwargs, + ) -> Tuple: + """ """ + return_per_sample_loss = bool(kwargs.pop("return_per_sample_loss", False)) + flow_noise_seed = kwargs.pop("flow_noise_seed", None) + if kwargs: + raise TypeError(f"Unexpected forward arguments: {sorted(kwargs)}") + batch_images = [example["image"] for example in examples] # [B，[PLT]] + instructions = [example["lang"] for example in examples] # [B, str] + actions = [example["action"] for example in examples] # label [B， len, 7] + + state = [example["state"] for example in examples] if "state" in examples[0] else None # [B, 1, state_dim] + + # Step 1: QWenVL input format + qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions) + with torch.autocast("cuda", dtype=torch.bfloat16): + qwenvl_outputs = self.qwen_vl_interface( + **qwen_inputs, + output_attentions=False, + output_hidden_states=True, + return_dict=True, + ) + # last_hidden_state: [B, seq_len, H] + last_hidden = qwenvl_outputs.hidden_states[-1] # [B, L, H] + + # Step 4: Action Expert Forward and Loss + with torch.autocast("cuda", dtype=torch.float32): + actions = torch.tensor( + np.array(actions), device=last_hidden.device, dtype=last_hidden.dtype + ) # [B, T_full, action_dim] + actions_target = actions[:, -self.action_horizon :, :] # (B, action_horizon, action_dim) + + repeated_diffusion_steps = int( + self.config.trainer.repeated_diffusion_steps + ) + actions_target_repeated = actions_target.repeat(repeated_diffusion_steps, 1, 1) + last_hidden_repeated = last_hidden.repeat(repeated_diffusion_steps, 1, 1) + state_repeated = None + if state is not None: + state = torch.tensor(np.array(state), device=last_hidden.device, dtype=last_hidden.dtype) + state_repeated = state.repeat(repeated_diffusion_steps, 1, 1) + + action_loss = self.action_model( + last_hidden_repeated, actions_target_repeated, state_repeated, + reduction="none" if return_per_sample_loss else "mean", + flow_noise_seed=flow_noise_seed, + ) # (B, chunk_len, action_dim) + if return_per_sample_loss: + action_loss = aggregate_repeated_losses( + action_loss, + repeat_count=repeated_diffusion_steps, + batch_size=len(examples), + ) + + return {"action_loss": action_loss} + + @torch.inference_mode() + def predict_action( + self, + examples: List[dict], + **kwargs: str, + ) -> np.ndarray: + """ + Steps: + 1. Resize images to training resolution (if specified) + 2. Encode with QwenVL (hidden states retained) + 6. Return normalized action trajectory + Returns: + dict: + normalized_actions (np.ndarray): Shape [B, T, action_dim], diffusion-sampled normalized actions. + """ + if type(examples) is not list: + examples = [examples] + batch_images = [to_pil_preserve(example["image"]) for example in examples] # [B，[PLT]] + instructions = [example["lang"] for example in examples] # [B, str] + + state = [example["state"] for example in examples] if "state" in examples[0] else None # [B, 1, state_dim] + + train_obs_image_size = getattr(self.config.datasets.vla_data, "image_size", None) + if train_obs_image_size: + batch_images = resize_images(batch_images, target_size=train_obs_image_size) + + # Step 1: QWenVL input format + qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions) + with torch.autocast("cuda", dtype=torch.bfloat16): + qwenvl_outputs = self.qwen_vl_interface( + **qwen_inputs, + output_attentions=False, + output_hidden_states=True, + return_dict=True, + ) + + # last_hidden_state: [B, seq_len, H] + last_hidden = qwenvl_outputs.hidden_states[-1] # [B, L, H] + + state = ( + torch.from_numpy(np.array(state)).to(last_hidden.device, dtype=last_hidden.dtype) + if state is not None + else None + ) + + # Step 4: Action Expert Forward + with torch.autocast("cuda", dtype=torch.float32): + pred_actions = self.action_model.predict_action(last_hidden, state) + + normalized_actions = pred_actions.detach().cpu().numpy() + return {"normalized_actions": normalized_actions} + + +if __name__ == "__main__": + import argparse + import os + + from omegaconf import OmegaConf + + parser = argparse.ArgumentParser() + parser.add_argument( + "--config_yaml", + type=str, + default="examples/LIBERO/train_files/starvla_cotrain_libero.yaml", + help="Path to YAML config", + ) + args, clipargs = parser.parse_known_args() + + if os.getenv("DEBUGPY_ENABLE", "0") == "1": + import debugpy + + debugpy.listen(("0.0.0.0", 10092)) + print("Rank 0 waiting for debugger attach on port 10092...") + debugpy.wait_for_client() + + cfg = OmegaConf.load(args.config_yaml) + + model: Qwen_GR00T = Qwen_GR00T(cfg) + print(model) + + image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)) + sample = { + "action": np.random.uniform(-1, 1, size=(16, 7)).astype(np.float16), + "image": [image], + "lang": "This is a fake instruction for testing.", + } + sample2 = sample.copy() + sample2["lang"] = "Another fake instruction for testing." + + batch = [sample, sample2] + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = model.to(device) + forward_output = model(batch) + action_loss = forward_output["action_loss"] + print(f"Action Loss: {action_loss.item()}") + + predict_output = model.predict_action(examples=[sample]) + normalized_actions = predict_output["normalized_actions"] + print(f"Unnormalized Action: {normalized_actions}") + + print("Finished") diff --git a/starVLA/model/framework/VLM4A/__init__.py b/starVLA/model/framework/VLM4A/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..81330e2ca135c0a5d25e94d8060a267b50661678 --- /dev/null +++ b/starVLA/model/framework/VLM4A/__init__.py @@ -0,0 +1,8 @@ +"""VLM4A — VLM-for-Action frameworks. + +All VLA frameworks that use a Vision-Language Model (Qwen-VL, Florence, etc.) +as the perception backbone live here. + +The registry decorator on each class registers them in FRAMEWORK_REGISTRY, +so `build_framework(cfg)` can find them after auto-import. +""" diff --git a/starVLA/model/framework/__init__.py b/starVLA/model/framework/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cca1db73ef45e5732d73d74dc50a7da42d34096f --- /dev/null +++ b/starVLA/model/framework/__init__.py @@ -0,0 +1,8 @@ +"""Framework package internals. + +Public framework APIs are defined in `starVLA.model.framework.base_framework`. +""" + +from starVLA.model.tools import FRAMEWORK_REGISTRY + +__all__ = ["FRAMEWORK_REGISTRY"] diff --git a/starVLA/model/framework/base_framework.py b/starVLA/model/framework/base_framework.py new file mode 100644 index 0000000000000000000000000000000000000000..8c4c5fe9c08ebc3cb93801c9ac68731e2b9907a5 --- /dev/null +++ b/starVLA/model/framework/base_framework.py @@ -0,0 +1,289 @@ +""" +Base framework abstraction providing: +- Pretrained loading (config + normalization stats + weights) +- Action space utilities (dimension, stats, (un)normalization) +- Trainable module discovery helper +Note: No device placement or optimizer concerns handled here (delegated to trainer). +""" + +import importlib +import pkgutil +from pathlib import Path +from typing import Any, Dict, List + +import numpy as np +import torch +from transformers import PretrainedConfig, PreTrainedModel + +from starVLA.model.framework.share_tools import dict_to_namespace, read_mode_config +from starVLA.model.tools import FRAMEWORK_REGISTRY, FrameworkTools, auto_get_trainable_modules +from starVLA.training.trainer_utils import initialize_overwatch + +logger = initialize_overwatch(__name__) +_FRAMEWORKS_IMPORTED = False +_FRAMEWORK_IMPORT_ERRORS: dict[str, str] = {} +_OPTIONAL_EXTERNAL_MODULES = {"qwen_vl_utils"} + + +def _import_optional_framework(module_path: str) -> None: + """Import a framework plugin without making unrelated extras mandatory.""" + try: + importlib.import_module(module_path) + except ModuleNotFoundError as exc: + missing_root = (exc.name or "").split(".", 1)[0] + if missing_root not in _OPTIONAL_EXTERNAL_MODULES: + raise + _FRAMEWORK_IMPORT_ERRORS[module_path] = str(exc) + logger.warning("Skipping optional framework %s: %s", module_path, exc) + + +def _auto_import_framework_modules() -> None: + global _FRAMEWORKS_IMPORTED + if _FRAMEWORKS_IMPORTED: + return + + _SKIP = {"__init__", "base_framework", "share_tools"} + framework_dir = Path(__file__).resolve().parent + + # Scan top-level modules (backwards compat) + for _, module_name, is_pkg in pkgutil.iter_modules([str(framework_dir)]): + if module_name in _SKIP: + continue + if is_pkg: + # Scan sub-packages (VLM4A/, WM4A/, etc.) + sub_dir = framework_dir / module_name + for _, sub_name, _ in pkgutil.iter_modules([str(sub_dir)]): + if sub_name.startswith("_"): + continue + _import_optional_framework( + f"starVLA.model.framework.{module_name}.{sub_name}" + ) + else: + _import_optional_framework(f"starVLA.model.framework.{module_name}") + + _FRAMEWORKS_IMPORTED = True + + +def build_framework(cfg): # The single entry point for building different model frameworks + """ + Build a framework model from config. + Args: + cfg: Config object containing `cfg.framework.name`. + Returns: + nn.Module: Instantiated framework model. + """ + if not hasattr(cfg, "framework") or not hasattr(cfg.framework, "name"): + raise ValueError("Missing `cfg.framework.name`. The framework API now only accepts `framework.name`.") + + _auto_import_framework_modules() + + framework_id = cfg.framework.name + if framework_id not in FRAMEWORK_REGISTRY._registry: + available = sorted(FRAMEWORK_REGISTRY._registry.keys()) + skipped = "; ".join( + f"{module}: {error}" + for module, error in sorted(_FRAMEWORK_IMPORT_ERRORS.items()) + ) + raise NotImplementedError( + f"Framework `{framework_id}` is not implemented. Available frameworks: {available}. " + f"Skipped optional frameworks: {skipped or 'none'}" + ) + + model_class = FRAMEWORK_REGISTRY[framework_id] + return model_class(cfg) + + +# PreTrainedModel, AutoModel, PretrainedConfig, are so good, find sometime to study them +# TODO @JinhuiYE find sometime to merge yaml config with transformer config + + +class baseframework(PreTrainedModel): + """ + Lightweight base class for higher-level VLA model assemblies. + Subclasses are expected to: + - Accept a structured config + - Register components in __init__ + - Use provided helpers for action normalization handling + """ + + def __init__(self, hf_config=PretrainedConfig()) -> None: + """ + Initialize base nn.Module. Subclasses add components. + """ + + super().__init__(hf_config) + + # ------------------------------------------------------------------ + # Soft-constraint interface: subclasses should override these. + # Default implementations raise NotImplementedError so that IDE + # tooling (e.g. pylance, mypy) flags missing overrides, while + # still allowing PreTrainedModel instantiation (no ABC). + # ------------------------------------------------------------------ + + def forward(self, examples: List[dict], **kwargs) -> dict: + """Training forward pass. + + Args: + examples: List[dict], each dict requires at least: + - image: List[PIL.Image] + - lang: str + - action: np.ndarray shaped [T, action_dim] + + Returns: + dict: Must contain ``"action_loss"`` (torch.Tensor scalar). + May contain extra keys for logging (e.g. ``"kl_loss"``). + """ + raise NotImplementedError( + f"{type(self).__name__} must implement forward(examples) -> dict with 'action_loss' key." + ) + + def predict_action(self, examples: List[dict], **kwargs) -> dict: + """Inference: predict future actions from observations. + + Args: + examples: Same schema as *forward* (minus ``action`` which is optional). + **kwargs: Framework-specific inference options (e.g. ``use_ddim``). + + Returns: + dict: Must contain ``"normalized_actions"`` (np.ndarray [B, T, action_dim]). + """ + raise NotImplementedError( + f"{type(self).__name__} must implement predict_action(examples) -> dict with 'normalized_actions' key." + ) + + # ------------------------------------------------------------------ + # Unified loss interface for Trainer + # ------------------------------------------------------------------ + + def supports_training_tag(self, tag: str) -> bool: + """Return whether this framework can consume batches for *tag*.""" + if tag == "vla": + return type(self).forward is not baseframework.forward + if tag == "vlm": + return hasattr(self, "qwen_vl_interface") or type(self).forward_vlm is not baseframework.forward_vlm + return False + + def compute_loss(self, tag: str, batch, loss_scale: dict = None) -> Dict[str, torch.Tensor] | None: + """Unified forward entry-point: route to the right forward by *tag*. + + The trainer calls ``model.compute_loss(tag, batch)`` for every + ``(tag, batch)`` pair produced by :class:`DataLoaderManager`. + The model internally dispatches: + + - ``"vla"`` → ``self.forward(batch)`` + - ``"vlm"`` → ``self.forward_vlm(batch)`` + + Subclasses can override this to add more tags (e.g. ``"world"``). + + Args: + tag: dataset type tag (``"vla"``, ``"vlm"``, …) + batch: the batch produced by the corresponding DataLoader. + loss_scale: ``{"vla": 1.0, "vlm": 0.1}`` per-tag loss multiplier. + Defaults to 1.0 for unspecified tags. + + Returns: + dict[str, Tensor] | None: keyed losses (e.g. ``{"action_loss": ...}``). + Returns ``None`` when this framework does not support the + incoming dataloader tag so the trainer can ``continue``. + """ + if not self.supports_training_tag(tag): + return None + + scale = (loss_scale or {}).get(tag, 1.0) + + if tag == "vla": + out = self.forward(batch) + elif tag == "vlm": + out = self.forward_vlm(batch) + else: + return None + + # Apply loss scale and filter to Tensor values only + return {k: v * scale for k, v in out.items() if isinstance(v, torch.Tensor)} + + def forward_vlm(self, batch) -> Dict[str, torch.Tensor]: + """VLM forward pass (default implementation). + + Delegates to ``self.qwen_vl_interface(**batch)`` which is present on + every framework subclass that uses a Qwen VL backbone. + + Subclasses may override to add custom VLM logic. + + Args: + batch: dict produced by the VLM dataloader. + + Returns: + dict: Must contain ``"vlm_loss"`` (torch.Tensor scalar). + """ + if not hasattr(self, "qwen_vl_interface"): + raise NotImplementedError( + f"{type(self).__name__} has no `qwen_vl_interface`. " + "Override forward_vlm() to support VLM training." + ) + out = self.qwen_vl_interface(**batch) + return {"vlm_loss": out.loss} + + @classmethod + def from_pretrained( + cls, + pretrained_checkpoint: str, + **kwargs, + ) -> None: + """ + Restore a model instance from a saved checkpoint. + + Workflow: + 1. Resolve checkpoint path + 2. Load config + dataset normalization statistics + 3. Build model with loaded config + 4. Load state_dict strictly (reports missing/unexpected keys) + 5. Attach normalization stats for later un-normalization + + Args: + pretrained_checkpoint: Path to .pt file inside run/checkpoints directory. + **kwargs: Extra constructor overrides passed to subclass. + + Returns: + baseframework: Instantiated model (left on CPU; caller decides device). + + Raises: + RuntimeError: If state_dict key mismatch occurs under strict=True. + FileNotFoundError: If underlying files are missing (surfaced earlier). + """ + pretrained_checkpoint = Path(pretrained_checkpoint) + model_config, norm_stats = read_mode_config(pretrained_checkpoint) # read config and norm_stats + + config = dict_to_namespace(model_config) + model_config = config + model_config.trainer.pretrained_checkpoint = None + + FrameworkModel = build_framework(cfg=model_config) + # set for action un-norm + FrameworkModel.norm_stats = norm_stats + # Load from Checkpoint (Custom --> should load both *projector* and *llm* weights) + if pretrained_checkpoint.suffix == ".safetensors": + from safetensors.torch import load_file + + model_state_dict = load_file(str(pretrained_checkpoint)) + else: + model_state_dict = torch.load(pretrained_checkpoint, map_location="cpu") + # logger.info(f"Loading model weights from `{pretrained_checkpoint}`") + model_keys = set(FrameworkModel.state_dict().keys()) + checkpoint_keys = set(model_state_dict.keys()) + try: + FrameworkModel.load_state_dict(model_state_dict, strict=True) + except RuntimeError as e: + # must keep all keys matched + common_keys = model_keys.intersection(checkpoint_keys) + missing_keys = model_keys - common_keys + unexpected_keys = checkpoint_keys - common_keys + if missing_keys: + logger.warning(f"Missing keys in state_dict: {missing_keys}") + if unexpected_keys: + logger.warning(f"Unexpected keys in state_dict: {unexpected_keys}") + + raise e + + # **ensure model is on GPU** + FrameworkModel = FrameworkModel + return FrameworkModel diff --git a/starVLA/model/framework/share_tools.py b/starVLA/model/framework/share_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..8606daad182d80ba3d59dee136c8e982c07b1ee1 --- /dev/null +++ b/starVLA/model/framework/share_tools.py @@ -0,0 +1,559 @@ +""" +Shared configuration / utility helpers for framework components: +- NamespaceWithGet: lightweight namespace behaving like a dict +- OmegaConf conversion helpers +- Config merging decorator for model __init__ +- Checkpoint config/statistics loading +""" + +import functools +import inspect +import json +import os +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +from omegaconf import OmegaConf + +from starVLA.training.trainer_utils import initialize_overwatch + +# Initialize Overwatch =>> Wraps `logging.Logger` +overwatch = initialize_overwatch(__name__) + + +class NamespaceWithGet(SimpleNamespace): + def get(self, key, default=None): + """ + Return attribute value if present, else default (dict-like API). + + Args: + key: Attribute name. + default: Fallback if attribute missing. + + Returns: + Any: Stored value or default. + """ + return getattr(self, key, default) + + def items(self): + """ + Iterate (key, value) pairs like dict.items(). + + Returns: + Generator[Tuple[str, Any], None, None] + """ + return ((key, getattr(self, key)) for key in self.__dict__) + + def __iter__(self): + """ + Return iterator over attribute keys (enables dict unpacking **obj). + + Returns: + Iterator[str] + """ + return iter(self.__dict__) + + def to_dict(self): + """ + Recursively convert nested NamespaceWithGet objects into plain dicts. + + Returns: + dict: Fully materialized dictionary structure. + """ + return {key: value.to_dict() if isinstance(value, NamespaceWithGet) else value for key, value in self.items()} + + +def dict_to_namespace(d): + """ + Create an OmegaConf config from a plain dictionary. + + Args: + d: Input dictionary. + + Returns: + OmegaConf: DictConfig instance. + """ + return OmegaConf.create(d) + + +def _to_omegaconf(x: Any): + """ + Convert diverse input types into an OmegaConf object. + + Accepted types: + - None -> empty DictConfig + - str path -> load YAML/JSON via OmegaConf.load + - dict -> DictConfig + - DictConfig / ListConfig -> returned unchanged + - NamespaceWithGet / SimpleNamespace -> converted via vars()/to_dict() + + Args: + x: Input candidate. + + Returns: + OmegaConf: Normalized configuration node. + """ + if x is None: + return OmegaConf.create({}) + if isinstance(x, OmegaConf.__class__): # fallback, typically not hit + return x + try: + # OmegaConf node detection + from omegaconf import DictConfig, ListConfig + + if isinstance(x, (DictConfig, ListConfig)): + return x + except Exception: + pass + + if isinstance(x, str): + # treat as path + return OmegaConf.load(x) + if isinstance(x, dict): + return OmegaConf.create(x) + if isinstance(x, NamespaceWithGet) or isinstance(x, SimpleNamespace): + # convert to plain dict + try: + d = x.to_dict() if hasattr(x, "to_dict") else vars(x) + except Exception: + d = vars(x) + return OmegaConf.create(d) + # fallback: try to create + return OmegaConf.create(x) + + +def merge_pram_config(init): + """ + Decorator for __init__ to unify config handling. + + Behavior: + 1. Extract 'config' kwarg / arg (path | dict | OmegaConf | namespace) + 2. Convert to OmegaConf + 3. Merge with explicitly passed init parameters (explicit overrides file) + 4. Attach merged config to self.config + 5. Call original __init__ with merged config + + Args: + init: Original __init__ function. + + Returns: + Wrapped initializer. + """ + + @functools.wraps(init) + def wrapper(self, *args, **kwargs): + # Map positional args to parameter names (excluding self) + sig = inspect.signature(init) + param_names = [name for i, (name, p) in enumerate(sig.parameters.items()) if i > 0] + + init_kwargs = {} + for name, val in zip(param_names, args): + init_kwargs[name] = val + # override with explicit kwargs + init_kwargs.update(kwargs) + + # get provided config (if any) + provided_config = init_kwargs.get("config", None) + + loaded_cfg = _to_omegaconf(provided_config) + + # build params cfg from explicit init args (other than config) + params = {k: v for k, v in init_kwargs.items() if k != "config"} + params_cfg = OmegaConf.create(params) if params else OmegaConf.create({}) + + # merge: loaded_cfg <- params_cfg (params override file) + merged = OmegaConf.merge(loaded_cfg, params_cfg) + + # set on instance + try: + # prefer attaching OmegaConf directly + self.config = merged + except Exception: + # fallback to dict + self.config = OmegaConf.to_container(merged, resolve=True) + + # prepare kwargs for original init: ensure config is the merged OmegaConf + call_kwargs = dict(init_kwargs) + call_kwargs["config"] = merged + + # call original __init__ using keyword args only (safer) + return init(self, **call_kwargs) + + return wrapper + + +def merge_framework_config(default_config_cls, cfg): + """ + Merge a framework's default config (dataclass) with the incoming YAML config. + + Rules: + - default_config_cls provides documented defaults for `cfg.framework` + - YAML values (cfg.framework) override matching defaults + - Extra YAML keys not in defaults are preserved (Config-as-API flexibility) + - Missing YAML keys fall back to defaults (less YAML boilerplate) + + The merge only touches the `cfg.framework` sub-tree; datasets / trainer / etc. + are left untouched. + + Args: + default_config_cls: A dataclass **class** (not instance) whose fields() define + the default framework config with type hints and comments. + cfg: The full OmegaConf config (must contain cfg.framework). + + Returns: + cfg: The same config object with cfg.framework replaced by the merged result. + """ + import dataclasses + + from omegaconf import DictConfig, OmegaConf + + # 1. Instantiate defaults and convert to OmegaConf + defaults_instance = default_config_cls() + defaults_dict = dataclasses.asdict(defaults_instance) + defaults_omega = OmegaConf.create(defaults_dict) + + # 2. Extract the YAML framework section + if hasattr(cfg, "framework"): + # Unwrap AccessTrackedConfig if needed + yaml_fw = cfg.framework + if hasattr(yaml_fw, "_cfg"): + yaml_fw = yaml_fw._cfg + if not isinstance(yaml_fw, DictConfig): + yaml_fw = OmegaConf.create(yaml_fw if isinstance(yaml_fw, dict) else {}) + else: + yaml_fw = OmegaConf.create({}) + + # 3. Merge: defaults first, YAML overrides (YAML wins on conflicts) + merged_fw = OmegaConf.merge(defaults_omega, yaml_fw) + + # 4. Write back into the original cfg + # Handle both OmegaConf and AccessTrackedConfig transparently + if hasattr(cfg, "_cfg") and isinstance(cfg._cfg, DictConfig): + # AccessTrackedConfig caches child wrappers in _children dict. + # After replacing the underlying DictConfig node, the old child wrapper + # still points to the pre-merge node (stale data). We must invalidate + # the cache so the next attribute access creates a fresh wrapper around + # the merged node. + # + # However, the old child's _local_accessed set records which keys were + # already read (e.g. "name" from build_framework). Deleting the child + # would lose that tracking info, causing save_accessed_config to omit + # those keys from config.yaml. So we preserve and restore it. + cfg._cfg.framework = merged_fw + if hasattr(cfg, "_children") and "framework" in cfg._children: + old_accessed = cfg._children["framework"]._local_accessed.copy() + del cfg._children["framework"] # invalidate stale cache + new_child = cfg.framework # re-create child around merged_fw + new_child._local_accessed.update(old_accessed) # restore tracking + elif isinstance(cfg, DictConfig): + cfg.framework = merged_fw + else: + # Fallback — try direct attribute setting + try: + cfg.framework = merged_fw + except Exception: + overwatch.warning("Could not write merged framework config back to cfg.") + + return cfg + + +def populate_layerwise_dit_cfg(cfg, *, dit_hidden_dim: int, num_dit_layers: int): + """ + Populate ``framework.action_model.diffusion_model_cfg`` with the DiT shape + fields required by ``LayerwiseFlowmatchingActionHead``. + + Why this helper exists: + The action head is intentionally agnostic of the VLM backbone — it only + consumes ``diffusion_model_cfg``. Each framework (QwenPI, QwenPI_v3, + ...) is responsible for deciding the DiT shape (depth + hidden) from + whatever source it likes (LLM hidden, a compressed projector dim, ...) + and writing it here BEFORE calling ``get_action_model``. + + Fields written (override any stale YAML values): + - num_layers = num_dit_layers + - input_embedding_dim = dit_hidden_dim + - cross_attention_dim = dit_hidden_dim (encoder is pre-projected) + - num_attention_heads = dit_hidden_dim // attention_head_dim + (uses existing attention_head_dim if set, else 64) + + Args: + cfg: Full OmegaConf config. + dit_hidden_dim: DiT internal hidden dim. + num_dit_layers: Number of DiT cross-attention layers. + + Returns: + The (mutated) diffusion_model_cfg node. + """ + dit_cfg = cfg.framework.action_model.diffusion_model_cfg + head_dim = dit_cfg.get("attention_head_dim", None) or 64 + dit_cfg.attention_head_dim = head_dim + dit_cfg.num_layers = int(num_dit_layers) + dit_cfg.input_embedding_dim = int(dit_hidden_dim) + dit_cfg.cross_attention_dim = int(dit_hidden_dim) + dit_cfg.num_attention_heads = int(dit_hidden_dim) // int(head_dim) + return dit_cfg + + +def read_model_config(pretrained_checkpoint): + """ + Load global model configuration and dataset normalization statistics + associated with a saved checkpoint (.pt). + + Expected directory layout: + /checkpoints/.pt + /config.json + /dataset_statistics.json + + Args: + pretrained_checkpoint: Path to a .pt checkpoint file. + + Returns: + tuple: + global_cfg (dict): Loaded config.json contents. + norm_stats (dict): Dataset statistics for (de)normalization. + + Raises: + FileNotFoundError: If checkpoint or required JSON files are missing. + AssertionError: If file suffix or structure invalid. + """ + if os.path.isfile(pretrained_checkpoint): + overwatch.info(f"Loading from local checkpoint path `{(checkpoint_pt := Path(pretrained_checkpoint))}`") + + # [Validate] Checkpoint Path should look like + # `...//checkpoints/.pt|.safetensors` + assert checkpoint_pt.suffix in {".pt", ".safetensors"} + run_dir = checkpoint_pt.parents[1] + + # Get paths for `config.json`, `dataset_statistics.json` and pretrained checkpoint + config_json, dataset_statistics_json = run_dir / "config.json", run_dir / "dataset_statistics.json" + assert config_json.exists(), f"Missing `config.json` for `{run_dir = }`" + assert dataset_statistics_json.exists(), f"Missing `dataset_statistics.json` for `{run_dir = }`" + + # Otherwise =>> try looking for a match on `model_id_or_path` on the HF Hub (`model_id_or_path`) + # Load VLA Config (and corresponding base VLM `ModelConfig`) from `config.json` + with open(config_json, "r") as f: + global_cfg = json.load(f) + + # Normalise legacy / pre-v0.21 configs to current schema (idempotent; + # also ensures `past_action_window_size`, `action_horizon`, + # `future_action_window_size`, etc. are all present for downstream code). + try: + _oc = OmegaConf.create(global_cfg) + apply_config_compat(_oc) + global_cfg = OmegaConf.to_container(_oc, resolve=True) + except Exception as e: + overwatch.warning(f"apply_config_compat failed on `{config_json}`: {e}") + + # Load Dataset Statistics for Action Denormalization + with open(dataset_statistics_json, "r") as f: + norm_stats = json.load(f) + else: + overwatch.error(f"❌ Pretrained checkpoint `{pretrained_checkpoint}` does not exist.") + raise FileNotFoundError(f"Pretrained checkpoint `{pretrained_checkpoint}` does not exist.") + return global_cfg, norm_stats + + +def read_mode_config(pretrained_checkpoint): + """ + Same as read_model_config (legacy duplicate kept for backward compatibility). + + Args: + pretrained_checkpoint: Path to a .pt checkpoint file. + + Returns: + tuple: + vla_cfg (dict) + norm_stats (dict) + """ + if os.path.isfile(pretrained_checkpoint): + overwatch.info(f"Loading from local checkpoint path `{(checkpoint_pt := Path(pretrained_checkpoint))}`") + + # [Validate] Checkpoint Path should look like + # `...//checkpoints/.pt|.safetensors` + assert checkpoint_pt.suffix in {".pt", ".safetensors"} + run_dir = checkpoint_pt.parents[1] + + # Get paths for `config.json`, `dataset_statistics.json` and pretrained checkpoint + config_yaml, dataset_statistics_json = run_dir / "config.yaml", run_dir / "dataset_statistics.json" + assert config_yaml.exists(), f"Missing `config.yaml` for `{run_dir = }`" + assert dataset_statistics_json.exists(), f"Missing `dataset_statistics.json` for `{run_dir = }`" + + # Otherwise =>> try looking for a match on `model_id_or_path` on the HF Hub (`model_id_or_path`) + # Load VLA Config (and corresponding base VLM `ModelConfig`) from `config.json` + try: + ocfg = OmegaConf.load(str(config_yaml)) + # Normalise legacy / pre-v0.21 configs to current schema (idempotent). + apply_config_compat(ocfg) + global_cfg = OmegaConf.to_container(ocfg, resolve=True) + except Exception as e: + overwatch.error(f"❌ Failed to load YAML config `{config_yaml}`: {e}") + raise + + # Load Dataset Statistics for Action Denormalization + with open(dataset_statistics_json, "r") as f: + norm_stats = json.load(f) + else: + overwatch.error(f"❌ Pretrained checkpoint `{pretrained_checkpoint}` does not exist.") + raise FileNotFoundError(f"Pretrained checkpoint `{pretrained_checkpoint}` does not exist.") + return global_cfg, norm_stats + + +# ============================================================================= +# Config compatibility / "tightening" layer (introduced in version_id "0.21"). +# +# Goal: keep user-facing YAMLs short and unambiguous while preserving full +# back-compat for old checkpoints' config.yaml. See bar/config_收紧.md for +# the design rationale. +# +# This function is *idempotent* — calling it multiple times yields the same +# result. It does NOT touch framework class signatures; instead it normalises +# the OmegaConf tree so that downstream framework __init__ code (which still +# reads e.g. `future_action_window_size`) keeps working unchanged. +# ============================================================================= + +CONFIG_VERSION = "0.21" + + +def apply_config_compat(cfg, *, strict: bool = False): + """ + Normalise an arbitrary (old or new) starVLA training config into the + current `version_id == "0.21"` schema. + + Performed transformations (each applied only when needed): + + 1. `framework.action_model.action_horizon` ↔ `future_action_window_size` + - `action_horizon` is canonical (preferred user-facing name). + - `future_action_window_size = action_horizon - 1` is auto-filled so + framework code that still reads the old key keeps working. + - If both are present and inconsistent, a warning is emitted and + `action_horizon` wins. + + 2. `framework.action_model.diffusion_model_cfg.output_dim` + - Auto-filled from `framework.action_model.hidden_size` when missing. + + 3. `framework.action_model.diffusion_model_cfg.cross_attention_dim` + - Auto-filled from `framework.qwenvl.vl_hidden_dim` when missing. + Frameworks that further override this at runtime (e.g. QwenGR00T) + are unaffected. + + 4. `framework.action_model.action_hidden_dim` + - Auto-filled from `hidden_size` when missing. OFT-family frameworks + still overwrite this from VLM hidden_size at runtime. + + 5. `framework.action_model.past_action_window_size` + - Auto-filled to `0` when missing. All released starVLA frameworks + run with past=0; the field is therefore dropped from user YAMLs + and only re-materialised here for legacy code that still reads it. + + 6. `cfg.version_id` is stamped to `"0.21"`. + + Args: + cfg: An OmegaConf DictConfig (or anything _to_omegaconf can wrap). + strict: If True, raise on inconsistent values instead of warning. + + Returns: + The same `cfg` object (mutated in place) for chaining convenience. + """ + from omegaconf import OmegaConf + + if cfg is None: + return cfg + + src_version = OmegaConf.select(cfg, "version_id", default=None) + + # ---- 1. action_horizon ↔ future_action_window_size ---- + am_path = "framework.action_model" + am = OmegaConf.select(cfg, am_path, default=None) + if am is not None: + ah = OmegaConf.select(am, "action_horizon", default=None) + fw = OmegaConf.select(am, "future_action_window_size", default=None) + + if ah is None and fw is not None: + ah = int(fw) + 1 + OmegaConf.update(cfg, f"{am_path}.action_horizon", ah, force_add=True) + elif ah is not None and fw is None: + fw = int(ah) - 1 + OmegaConf.update(cfg, f"{am_path}.future_action_window_size", fw, force_add=True) + elif ah is not None and fw is not None and int(ah) != int(fw) + 1: + msg = ( + f"[apply_config_compat] inconsistent action_horizon={ah} vs " + f"future_action_window_size={fw}; expected action_horizon == future + 1. " + "Using action_horizon as canonical." + ) + if strict: + raise ValueError(msg) + overwatch.warning(msg) + OmegaConf.update(cfg, f"{am_path}.future_action_window_size", int(ah) - 1, force_add=True) + + # ---- 2 & 3. diffusion_model_cfg auto-fill ---- + dm_path = f"{am_path}.diffusion_model_cfg" + dm = OmegaConf.select(cfg, dm_path, default=None) + if dm is not None: + hidden_size = OmegaConf.select(am, "hidden_size", default=None) + if OmegaConf.select(dm, "output_dim", default=None) is None and hidden_size is not None: + OmegaConf.update(cfg, f"{dm_path}.output_dim", int(hidden_size), force_add=True) + + if OmegaConf.select(dm, "cross_attention_dim", default=None) is None: + vl_hidden = OmegaConf.select(cfg, "framework.qwenvl.vl_hidden_dim", default=None) + if vl_hidden is not None: + OmegaConf.update(cfg, f"{dm_path}.cross_attention_dim", int(vl_hidden), force_add=True) + # else: leave None — framework __init__ may auto-bind it + + # ---- 4. action_hidden_dim fallback ---- + if OmegaConf.select(am, "action_hidden_dim", default=None) is None: + hidden_size = OmegaConf.select(am, "hidden_size", default=None) + if hidden_size is not None: + OmegaConf.update(cfg, f"{am_path}.action_hidden_dim", int(hidden_size), force_add=True) + + # ---- 5. past_action_window_size default ---- + if OmegaConf.select(am, "past_action_window_size", default=None) is None: + OmegaConf.update(cfg, f"{am_path}.past_action_window_size", 0, force_add=True) + + # ---- 6. stamp version ---- + if src_version != CONFIG_VERSION: + try: + OmegaConf.update(cfg, "version_id", CONFIG_VERSION, force_add=True) + except Exception: + try: + cfg.version_id = CONFIG_VERSION + except Exception: + pass + overwatch.info(f"[apply_config_compat] normalised config from version_id={src_version!r} to {CONFIG_VERSION!r}") + + return cfg + + +# ────────────────────────────────────────────────────────────────────── +# Discretised proprioceptive state → instruction prefix (π₀.5 style) +# ────────────────────────────────────────────────────────────────────── +import numpy as _np +from typing import List as _List + + +def state2str_transform(state: "_np.ndarray", num_bins: int = 256) -> str: + """Quantise a state vector into ``num_bins`` uniform bins over [-1, 1] + and return space-separated bin indices. + + Example: [-0.5, 0.1, 0.8] -> "95 133 203" + """ + discretized_state = _np.digitize(state, bins=_np.linspace(-1, 1, num_bins + 1)[:-1]) - 1 + return " ".join(map(str, discretized_state)) + + +def add_discretized_state_to_instruction( + instructions: "_List[str]", + states: "_List[_np.ndarray]", + num_bins: int = 256, +) -> "_List[str]": + """Append discretised proprioceptive state tokens to each instruction. + + Format: `` [STATE] [ACTION]`` + Lets the VLM attend to the robot state purely through its existing + text-token pathway — no extra encoder required (π₀.5 style). + """ + updated_instructions = [] + for instr, state in zip(instructions, states): + state_str = state2str_transform(state[0], num_bins=num_bins) + updated_instructions.append(f"{instr} [STATE] {state_str} [ACTION]") + return updated_instructions diff --git a/starVLA/model/modules/action_model/GR00T_ActionHeader.py b/starVLA/model/modules/action_model/GR00T_ActionHeader.py new file mode 100644 index 0000000000000000000000000000000000000000..67e102de57894a2548ce4f8df9ab3eb09032df4e --- /dev/null +++ b/starVLA/model/modules/action_model/GR00T_ActionHeader.py @@ -0,0 +1,483 @@ +# Copyright 2025 NVIDIA Corp. and affiliates. All rights reserved. +# Modified by [Junqiu YU/ Fudan University] in [2025]. +# Modification: [rm and add some connect adapter to match with starVLA, e.g., "rm "]. +# Action repeat is inspired by CogACT + + +from dataclasses import dataclass, field + +import torch +import torch.nn.functional as F +from torch import nn +from torch.distributions import Beta +from transformers import PretrainedConfig +from transformers.feature_extraction_utils import BatchFeature + +from starVLA.model.modules.action_model.flow_matching_head.action_encoder import ( + SinusoidalPositionalEncoding, + swish, +) +from starVLA.model.modules.action_model.flow_matching_head.cross_attention_dit import DiT +from starVLA.training.infoflow import reduce_flow_matching_loss + +# TODO try to meger DiT Modules with follow_match_head, they are just the same arch, but diff loss, use diffusers package will be simple + + +class CategorySpecificLinear(nn.Module): + def __init__(self, num_categories, input_dim, hidden_dim): + super().__init__() + self.num_categories = num_categories + # For each category, we have separate weights and biases. + self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim)) + self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim)) + + def forward(self, x, cat_ids): + selected_W = self.W[cat_ids] + selected_b = self.b[cat_ids] + # import ipdb; ipdb.set_trace() + return torch.bmm(x, selected_W) + selected_b.unsqueeze(1) + + +class CategorySpecificMLP(nn.Module): + def __init__(self, num_categories, input_dim, hidden_dim, output_dim): + super().__init__() + self.num_categories = num_categories + self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim) + self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim) + + def forward(self, x, cat_ids): + hidden = F.relu(self.layer1(x, cat_ids)) + return self.layer2(hidden, cat_ids) + + +class MLP(nn.Module): + def __init__(self, input_dim, hidden_dim, output_dim): + super().__init__() + self.layer1 = nn.Linear(input_dim, hidden_dim) + self.layer2 = nn.Linear(hidden_dim, output_dim) + + def forward(self, x): + return self.layer2(F.relu(self.layer1(x))) + + +class ActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size): + super().__init__() + self.hidden_size = hidden_size + self.action_dim = action_dim + self.layer1 = nn.Linear(action_dim, hidden_size) + self.layer2 = nn.Linear(2 * hidden_size, hidden_size) + self.layer3 = nn.Linear(hidden_size, hidden_size) + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions, timesteps): + """ + actions: shape (B, T, action_dim) + timesteps: shape (B,) -- a single scalar per batch item + returns: shape (B, T, hidden_size) + """ + B, T, _ = actions.shape + + # 1) Expand each batch's single scalar time 'tau' across all T steps + # so that shape => (B, T) + # e.g. if timesteps is (B,), replicate across T + if timesteps.dim() == 1 and timesteps.shape[0] == B: + # shape (B,) => (B,T) + timesteps = timesteps.unsqueeze(1).expand(-1, T) + else: + raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.") + + # 2) Standard action MLP step for shape => (B, T, w) + a_emb = self.layer1(actions) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then layer2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.layer2(x)) + + # 5) Finally W3 => (B, T, w) + x = self.layer3(x) + return x + + +class MultiEmbodimentActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size, num_embodiments): + super().__init__() + self.hidden_size = hidden_size + self.num_embodiments = num_embodiments + + # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} + self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size) # (d -> w) + self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size) # (2w -> w) + self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size) # (w -> w) + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions, timesteps, cat_ids): + """ + actions: shape (B, T, action_dim) + timesteps: shape (B,) -- a single scalar per batch item + cat_ids: shape (B,) + returns: shape (B, T, hidden_size) + """ + B, T, _ = actions.shape + + # 1) Expand each batch's single scalar time 'tau' across all T steps + # so that shape => (B, T) + # e.g. if timesteps is (B,), replicate across T + if timesteps.dim() == 1 and timesteps.shape[0] == B: + # shape (B,) => (B,T) + timesteps = timesteps.unsqueeze(1).expand(-1, T) + else: + raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.") + + # 2) Standard action MLP step for shape => (B, T, w) + a_emb = self.W1(actions, cat_ids) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.W2(x, cat_ids)) + + # 5) Finally W3 => (B, T, w) + x = self.W3(x, cat_ids) + return x + + +@dataclass +class FlowmatchingActionHeadConfig(PretrainedConfig): + """NOTE: N1.5 uses XEmbFlowmatchingPolicyHeadConfig as action head""" + + add_pos_embed: bool = field(default=True, metadata={"help": "Whether to add positional embedding"}) + diffusion_model_cfg: dict = field(default=None, metadata={"help": "Diffusion model configuration."}) + input_embedding_dim: int = field(default=1536, metadata={"help": "Input embedding channel dimension."}) + + hidden_size: int = field(default=1024, metadata={"help": "Input embedding dimension."}) + max_seq_len: int = field(default=1024, metadata={"help": "Maxium Sequence Length"}) + action_dim: int = field(default=None, metadata={"help": "Action dimension."}) + action_horizon: int = field(default=None, metadata={"help": "Action horizon."}) + noise_beta_alpha: float = field(default=1.5, metadata={"help": ""}) + noise_beta_beta: float = field(default=1.0, metadata={"help": ""}) + noise_s: float = field(default=0.999, metadata={"help": "Flow matching noise Beta distribution s."}) + num_timestep_buckets: int = field(default=1000, metadata={"help": "Number of timestep discretization buckets."}) + num_inference_timesteps: int = field( + default=None, + metadata={"help": "Number of inference steps for noise diffusion."}, + ) + max_num_embodiments: int = field(default=32, metadata={"help": "Number of embodiments."}) + tune_projector: bool = field(default=True, metadata={"help": "Whether to tune the projector."}) + tune_diffusion_model: bool = field(default=True, metadata={"help": "Whether to tune the diffusion model."}) + load_pretrained_det_decode_layer_path: str = field( + default=None, metadata={"help": "Path to pretrained detection model."} + ) + detection_coeff: float = field(default=1.0, metadata={"help": "Detection coefficient."}) + + freeze_decode_layer: bool = field(default=False) + expand_batch: int = field(default=None) + use_vlln: bool = field(default=True) + + vl_self_attention_cfg: dict = field(default=None) + num_target_vision_tokens: int = field(default=32, metadata={"help": "Number of target vision tokens."}) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + for key, value in kwargs.items(): + setattr(self, key, value) + + +DiTConfig = { + "DiT-B": {"input_embedding_dim": 768, "attention_head_dim": 64, "num_attention_heads": 12}, + "DiT-L": {"input_embedding_dim": 1536, "attention_head_dim": 48, "num_attention_heads": 32}, +} + + +class FlowmatchingActionHead(nn.Module): + def __init__( + self, + full_config, + ): + super().__init__() + config = full_config.framework.action_model + self.full_config = full_config + + # ------------------------------------------------------------------ + # DiT architecture selection + # action_model_type: "DiT-B" | "DiT-L" + # DiT-B → input_embedding_dim=768, heads=12, head_dim=64 + # DiT-L → input_embedding_dim=1536, heads=32, head_dim=48 + # diffusion_model_cfg overrides/extends the base DiT shape. + # In particular, diffusion_model_cfg.cross_attention_dim MUST be + # set by the framework to match the VLM hidden size BEFORE calling + # get_action_model(), e.g.: + # cfg.framework.action_model.diffusion_model_cfg.cross_attention_dim + # = vlm.model.config.hidden_size + # ------------------------------------------------------------------ + action_model_type = config.action_model_type + action_model_cfg = DiTConfig[action_model_type] + self.input_embedding_dim = action_model_cfg["input_embedding_dim"] + + diffusion_model_cfg = config.diffusion_model_cfg + diffusion_model_cfg = {**action_model_cfg, **diffusion_model_cfg} + self.model = DiT(**diffusion_model_cfg) + + # ------------------------------------------------------------------ + # Action horizon (chunk length sent to the DiT) + # Single source of truth: `action_horizon` (e.g. 8). + # Legacy YAMLs that only provide `future_action_window_size` are + # normalised to `action_horizon` upstream by + # `share_tools.apply_config_compat`, so this code never touches + # the legacy alias. + # ------------------------------------------------------------------ + self.action_horizon = int(config.action_horizon) + + # ------------------------------------------------------------------ + # Action / state dimensions + # action_dim: DoF of the robot action (e.g. 7 for 6-DoF + gripper) + # state_dim: proprioception dimension; set to 0/None to disable + # the state_encoder branch entirely. + # ------------------------------------------------------------------ + self.action_dim = config.action_dim + + # ------------------------------------------------------------------ + # Inference denoising steps + # num_inference_timesteps: Euler steps during predict_action(). + # Typically 4–10; fewer = faster but less accurate. + # ------------------------------------------------------------------ + self.num_inference_timesteps = config.num_inference_timesteps + + # ------------------------------------------------------------------ + # hidden_size: intermediate MLP width for state_encoder / action_decoder. + # Decoupled from input_embedding_dim so you can use a smaller hidden + # for the MLP without changing the DiT latent size. + # ------------------------------------------------------------------ + self.hidden_size = config.hidden_size + + self.state_encoder = ( + MLP( + input_dim=config.state_dim, + hidden_dim=self.hidden_size, + output_dim=self.input_embedding_dim, + ) + if config.state_dim + else None + ) + + self.action_encoder = ActionEncoder( + action_dim=config.action_dim, + hidden_size=self.input_embedding_dim, + ) + self.action_decoder = MLP( + input_dim=self.model.config.output_dim, + hidden_dim=self.hidden_size, + output_dim=self.action_dim, + ) + + # ------------------------------------------------------------------ + # future_tokens: learnable query tokens prepended before the action + # sequence so the DiT has dedicated "planning" slots. + # num_target_vision_tokens controls how many such tokens are added. + # ------------------------------------------------------------------ + self.future_tokens = nn.Embedding(config.num_target_vision_tokens, self.input_embedding_dim) + nn.init.normal_(self.future_tokens.weight, mean=0.0, std=0.02) + + # ------------------------------------------------------------------ + # Positional embedding over the action sequence + # add_pos_embed: whether to add sinusoidal-style learned PE + # max_seq_len: max supported action sequence length + # ------------------------------------------------------------------ + if config.add_pos_embed: + self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim) + nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02) + + # ------------------------------------------------------------------ + # Flow-matching noise schedule (Beta distribution) + # noise_beta_alpha / noise_beta_beta: Beta(α, β) shape params. + # noise_s: upper-clip of the sampled value so t ∈ [0, noise_s]. + # num_timestep_buckets: discretise continuous t into N buckets for + # the timestep encoder inside DiT. + # ------------------------------------------------------------------ + self.beta_dist = Beta(config.noise_beta_alpha, config.noise_beta_beta) + self.num_timestep_buckets = config.num_timestep_buckets + self.config = config + + def sample_time(self, batch_size, device, dtype): + sample = self.beta_dist.sample([batch_size]).to(device, dtype=dtype).clamp(max=self.config.noise_s) + return (self.config.noise_s - sample) / self.config.noise_s + + def _sample_training_noise_and_time( + self, + actions: torch.Tensor, + flow_noise_seed: int | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + def sample() -> tuple[torch.Tensor, torch.Tensor]: + noise = torch.randn( + actions.shape, + device=actions.device, + dtype=actions.dtype, + ) + time = self.sample_time( + actions.shape[0], + device=actions.device, + dtype=actions.dtype, + ) + return noise, time + + if flow_noise_seed is None: + return sample() + devices = ( + [torch.cuda.current_device()] + if actions.is_cuda + else [] + ) + # Flow draws must not advance the dropout stream. + with torch.random.fork_rng(devices=devices): + torch.manual_seed(int(flow_noise_seed)) + if actions.is_cuda: + torch.cuda.manual_seed(int(flow_noise_seed)) + return sample() + + def prepare_input(self, batch: dict) -> BatchFeature: + return BatchFeature(data=batch) + + def predict_velocity( + self, + vl_embs: torch.Tensor, + actions: torch.Tensor, + time: torch.Tensor, + state: torch.Tensor = None, + encoder_attention_mask=None, + ) -> torch.Tensor: + """Evaluate the conditional action vector field at continuous time.""" + if time.shape != (actions.shape[0],): + raise ValueError( + f"time must have shape {(actions.shape[0],)}, got {tuple(time.shape)}" + ) + time_float = time.to(device=actions.device, dtype=torch.float32).clamp( + min=0.0, max=1.0 + ) + timesteps = (time_float * self.num_timestep_buckets).floor().long() + timesteps = timesteps.clamp(max=self.num_timestep_buckets) + action_features = self.action_encoder(actions, timesteps) + + state_features = self.state_encoder(state) if state is not None else None + if self.config.add_pos_embed: + pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=actions.device) + action_features = action_features + self.position_embedding(pos_ids).unsqueeze(0) + + future_tokens = self.future_tokens.weight.unsqueeze(0).expand(vl_embs.shape[0], -1, -1) + sa_embs = ( + torch.cat((state_features, future_tokens, action_features), dim=1) + if state_features is not None + else torch.cat((future_tokens, action_features), dim=1) + ) + model_output = self.model( + hidden_states=sa_embs, + encoder_hidden_states=vl_embs, + encoder_attention_mask=encoder_attention_mask, + timestep=timesteps, + return_all_hidden_states=False, + ) + prediction = self.action_decoder(model_output) + return prediction[:, -actions.shape[1] :] + + def forward( + self, + vl_embs: torch.Tensor, + actions: torch.Tensor, + state: torch.Tensor = None, + encoder_attention_mask=None, + reduction: str = "mean", + flow_noise_seed: int | None = None, + ): + """ + vl_embs: shape (B, seq_length, feature_dim) + actions: shape (B, action_horizon, action_dim) + """ + # Embed noised action trajectory. + noise, t = self._sample_training_noise_and_time( + actions, + flow_noise_seed, + ) + t = t[:, None, None] # shape (B,1,1) for broadcast + + noisy_trajectory = (1 - t) * noise + t * actions + velocity = actions - noise + + pred_actions = self.predict_velocity( + vl_embs, + noisy_trajectory, + t[:, 0, 0], + state=state, + encoder_attention_mask=encoder_attention_mask, + ) + + # Slice out only the action portion of pred and target. + elementwise_loss = (pred_actions - velocity) ** 2 + return reduce_flow_matching_loss(elementwise_loss, reduction=reduction) + + @torch.no_grad() + def predict_action( + self, + vl_embs: torch.Tensor, + state: torch.Tensor = None, + encoder_attention_mask=None, + ) -> torch.Tensor: + # Set initial actions as the sampled noise. + batch_size = vl_embs.shape[0] + device = vl_embs.device + actions = torch.randn( + size=(batch_size, self.action_horizon, self.action_dim), + dtype=vl_embs.dtype, + device=device, + ) + + num_steps = self.num_inference_timesteps + dt = 1.0 / num_steps + + # Run denoising steps. + for t in range(num_steps): + t_cont = t / float(num_steps) # e.g. goes 0, 1/N, 2/N, ... + times = torch.full( + (batch_size,), t_cont, device=device, dtype=torch.float32 + ) + pred_velocity = self.predict_velocity( + vl_embs, + actions, + times, + state=state, + encoder_attention_mask=encoder_attention_mask, + ) + + # Update actions using euler integration. + actions = actions + dt * pred_velocity + return actions + + @property + def device(self): + return next(iter(self.parameters())).device + + @property + def dtype(self): + return next(iter(self.parameters())).dtype + + +def get_action_model(config=None): + """ + Factory: build FlowmatchingActionHead from global framework config. + + Args: + config: Global config (expects config.framework.action_model namespace). + + Returns: + FlowmatchingActionHead: Initialized FlowMatchingActionHead. + """ + return FlowmatchingActionHead(full_config=config) + + +if __name__ == "__main__": + # TODO make each backbone.py can be debug independently + + pass diff --git a/starVLA/model/modules/action_model/__init__.py b/starVLA/model/modules/action_model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c16f8f21133c45b778d3ccda0f5e84580c5c235 --- /dev/null +++ b/starVLA/model/modules/action_model/__init__.py @@ -0,0 +1 @@ +"""Action models included in the Florence-Bridge OPTR release.""" diff --git a/starVLA/model/modules/action_model/flow_matching_head/__init__.py b/starVLA/model/modules/action_model/flow_matching_head/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3159bfe65645499015bd92609b99d476d69544e9 --- /dev/null +++ b/starVLA/model/modules/action_model/flow_matching_head/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/starVLA/model/modules/action_model/flow_matching_head/action_encoder.py b/starVLA/model/modules/action_model/flow_matching_head/action_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..e767827402098ca666115195f66cccf6890b4945 --- /dev/null +++ b/starVLA/model/modules/action_model/flow_matching_head/action_encoder.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn as nn + + +def swish(x): + return x * torch.sigmoid(x) + + +class SinusoidalPositionalEncoding(nn.Module): + """ + Produces a sinusoidal encoding of shape (B, T, w) + given timesteps of shape (B, T). + """ + + def __init__(self, embedding_dim): + super().__init__() + self.embedding_dim = embedding_dim + + def forward(self, timesteps): + # timesteps: shape (B, T) + # We'll compute sin/cos frequencies across dim T + timesteps = timesteps.float() # ensure float + + B, T = timesteps.shape + device = timesteps.device + + half_dim = self.embedding_dim // 2 + # typical log space frequencies for sinusoidal encoding + exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * ( + torch.log(torch.tensor(10000.0)) / half_dim + ) + # Expand timesteps to (B, T, 1) then multiply + freqs = timesteps.unsqueeze(-1) * exponent.exp() # (B, T, half_dim) + + sin = torch.sin(freqs) + cos = torch.cos(freqs) + enc = torch.cat([sin, cos], dim=-1) # (B, T, w) + + return enc + + +class ActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size): + super().__init__() + self.hidden_size = hidden_size + + # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} + self.W1 = nn.Linear(action_dim, hidden_size) # (d -> w) + self.W2 = nn.Linear(2 * hidden_size, hidden_size) # (2w -> w) + self.W3 = nn.Linear(hidden_size, hidden_size) # (w -> w) + + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions, timesteps): + """ + actions: shape (B, T, action_dim) + timesteps: shape (B,) -- a single scalar per batch item + returns: shape (B, T, hidden_size) + """ + B, T, _ = actions.shape + + # 1) Expand each batch's single scalar time 'tau' across all T steps + # so that shape => (B, T) + # e.g. if timesteps is (B,), replicate across T + if timesteps.dim() == 1 and timesteps.shape[0] == B: + # shape (B,) => (B,T) + timesteps = timesteps.unsqueeze(1).expand(-1, T) + else: + raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.") + + # 2) Standard action MLP step for shape => (B, T, w) + a_emb = self.W1(actions) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.W2(x)) + + # 5) Finally W3 => (B, T, w) + x = self.W3(x) + + return x diff --git a/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py b/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..4394c6444e08d27053ce5d139326bd4e7a3fc7e8 --- /dev/null +++ b/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py @@ -0,0 +1,377 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional + +import torch +import torch.nn.functional as F +from diffusers import ConfigMixin, ModelMixin +from diffusers.configuration_utils import register_to_config +from diffusers.models.attention import Attention, FeedForward +from diffusers.models.embeddings import ( + SinusoidalPositionalEmbedding, + TimestepEmbedding, + Timesteps, +) +from torch import nn + + +class TimestepEncoder(nn.Module): + def __init__(self, embedding_dim, compute_dtype=torch.float32): + super().__init__() + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward(self, timesteps): + dtype = next(self.parameters()).dtype + timesteps_proj = self.time_proj(timesteps).to(dtype) + timesteps_emb = self.timestep_embedder(timesteps_proj) # (N, D) + return timesteps_emb + + +class AdaLayerNorm(nn.Module): + def __init__( + self, + embedding_dim: int, + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-5, + chunk_dim: int = 0, + ): + super().__init__() + self.chunk_dim = chunk_dim + output_dim = embedding_dim * 2 + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, output_dim) + self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine) + + def forward( + self, + x: torch.Tensor, + temb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + temb = self.linear(self.silu(temb)) + scale, shift = temb.chunk(2, dim=1) + x = self.norm(x) * (1 + scale[:, None]) + shift[:, None] + return x + + +class BasicTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + attention_bias: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen' + norm_eps: float = 1e-5, + final_dropout: bool = False, + attention_type: str = "default", + positional_embeddings: Optional[str] = None, + num_positional_embeddings: Optional[int] = None, + ff_inner_dim: Optional[int] = None, + ff_bias: bool = True, + attention_out_bias: bool = True, + ): + super().__init__() + self.dim = dim + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + self.dropout = dropout + self.cross_attention_dim = cross_attention_dim + self.activation_fn = activation_fn + self.attention_bias = attention_bias + self.norm_elementwise_affine = norm_elementwise_affine + self.positional_embeddings = positional_embeddings + self.num_positional_embeddings = num_positional_embeddings + self.norm_type = norm_type + + if positional_embeddings and (num_positional_embeddings is None): + raise ValueError( + "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." + ) + + if positional_embeddings == "sinusoidal": + self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) + else: + self.pos_embed = None + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if norm_type == "ada_norm": + self.norm1 = AdaLayerNorm(dim) + else: + self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim, + upcast_attention=upcast_attention, + out_bias=attention_out_bias, + ) + + # 3. Feed-forward + self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) + self.ff = FeedForward( + dim, + dropout=dropout, + activation_fn=activation_fn, + final_dropout=final_dropout, + inner_dim=ff_inner_dim, + bias=ff_bias, + ) + if final_dropout: + self.final_dropout = nn.Dropout(dropout) + else: + self.final_dropout = None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + temb: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + + # 0. Self-Attention + if self.norm_type == "ada_norm": + norm_hidden_states = self.norm1(hidden_states, temb) + else: + norm_hidden_states = self.norm1(hidden_states) + + if self.pos_embed is not None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, # @JinhuiYE original attention_mask=attention_mask + ) + if self.final_dropout: + attn_output = self.final_dropout(attn_output) + + hidden_states = attn_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + # 4. Feed-forward + norm_hidden_states = self.norm3(hidden_states) + ff_output = self.ff(norm_hidden_states) + + hidden_states = ff_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + return hidden_states + + +class DiT(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + # register_to_config auto-registers constructor params into config, enabling access via self.config.xxx instead of self.xxx + @register_to_config # Registers passed params to config. TODO: replace with our singleton pattern, implement a mergeable @merge_param_config + def __init__( + self, + num_attention_heads: int = 8, + attention_head_dim: int = 64, + output_dim: int = 26, + num_layers: int = 12, + dropout: float = 0.1, + attention_bias: bool = True, + activation_fn: str = "gelu-approximate", + num_embeds_ada_norm: Optional[int] = 1000, + upcast_attention: bool = False, + norm_type: str = "ada_norm", + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-5, + max_num_positional_embeddings: int = 512, + compute_dtype=torch.float32, + final_dropout: bool = True, + positional_embeddings: Optional[str] = "sinusoidal", + interleave_self_attention=False, + cross_attention_dim: Optional[int] = None, + **kwargs, + ): + super().__init__() + self.attention_head_dim = attention_head_dim + self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim + self.gradient_checkpointing = False + + # Timestep encoder + # self.config.compute_dtype may not exist, handle it in advance + compute_dtype = getattr(self.config, "compute_dtype", torch.float32) + self.timestep_encoder = ( + TimestepEncoder( # TODO BUG: self.config.compute_dtype doesn't error during training but fails at eval + embedding_dim=self.inner_dim, compute_dtype=compute_dtype + ) + ) + + all_blocks = [] + for idx in range(self.config.num_layers): + + use_self_attn = idx % 2 == 1 and interleave_self_attention + curr_cross_attention_dim = cross_attention_dim if not use_self_attn else None + + all_blocks += [ + BasicTransformerBlock( + self.inner_dim, + self.config.num_attention_heads, + self.config.attention_head_dim, + dropout=self.config.dropout, + activation_fn=self.config.activation_fn, + attention_bias=self.config.attention_bias, + upcast_attention=self.config.upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=self.config.norm_elementwise_affine, + norm_eps=self.config.norm_eps, + positional_embeddings=positional_embeddings, + num_positional_embeddings=self.config.max_num_positional_embeddings, + final_dropout=final_dropout, + cross_attention_dim=curr_cross_attention_dim, + ) + ] + self.transformer_blocks = nn.ModuleList(all_blocks) + + # Output blocks + self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim) + self.proj_out_2 = nn.Linear(self.inner_dim, self.config.output_dim) + print( + "Total number of DiT parameters: ", + sum(p.numel() for p in self.parameters() if p.requires_grad), + ) + + def forward( + self, + hidden_states: torch.Tensor, # Shape: (B, T, D) + encoder_hidden_states: torch.Tensor, # Shape: (B, S, D) + timestep: Optional[torch.LongTensor] = None, + return_all_hidden_states: bool = False, + encoder_attention_mask=None, + ): + # Encode timesteps + temb = self.timestep_encoder(timestep) + + # Process through transformer blocks - single pass through the blocks + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + + all_hidden_states = [hidden_states] + + # Process through transformer blocks + for idx, block in enumerate(self.transformer_blocks): + if idx % 2 == 1 and self.config.interleave_self_attention: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + temb=temb, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + temb=temb, + ) + all_hidden_states.append(hidden_states) + + # Output processing + conditioning = temb + shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + if return_all_hidden_states: + return self.proj_out_2(hidden_states), all_hidden_states + else: + return self.proj_out_2(hidden_states) + + +class SelfAttentionTransformer(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + num_attention_heads: int = 8, + attention_head_dim: int = 64, + output_dim: int = 26, + num_layers: int = 12, + dropout: float = 0.1, + attention_bias: bool = True, + activation_fn: str = "gelu-approximate", + num_embeds_ada_norm: Optional[int] = 1000, + upcast_attention: bool = False, + max_num_positional_embeddings: int = 512, + compute_dtype=torch.float32, + final_dropout: bool = True, + positional_embeddings: Optional[str] = "sinusoidal", + interleave_self_attention=False, + ): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim + self.gradient_checkpointing = False + + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + self.inner_dim, + self.config.num_attention_heads, + self.config.attention_head_dim, + dropout=self.config.dropout, + activation_fn=self.config.activation_fn, + attention_bias=self.config.attention_bias, + upcast_attention=self.config.upcast_attention, + positional_embeddings=positional_embeddings, + num_positional_embeddings=self.config.max_num_positional_embeddings, + final_dropout=final_dropout, + ) + for _ in range(self.config.num_layers) + ] + ) + print( + "Total number of SelfAttentionTransformer parameters: ", + sum(p.numel() for p in self.parameters() if p.requires_grad), + ) + + def forward( + self, + hidden_states: torch.Tensor, # Shape: (B, T, D) + return_all_hidden_states: bool = False, + ): + # Process through transformer blocks - single pass through the blocks + hidden_states = hidden_states.contiguous() + all_hidden_states = [hidden_states] + + # Process through transformer blocks + for idx, block in enumerate(self.transformer_blocks): + hidden_states = block(hidden_states) + all_hidden_states.append(hidden_states) + + if return_all_hidden_states: + return hidden_states, all_hidden_states + else: + return hidden_states diff --git a/starVLA/model/modules/vlm/Florence2.py b/starVLA/model/modules/vlm/Florence2.py new file mode 100644 index 0000000000000000000000000000000000000000..a2e000b9e5eb285efe8fa5e8fcc1049f7ad4adc2 --- /dev/null +++ b/starVLA/model/modules/vlm/Florence2.py @@ -0,0 +1,320 @@ +# 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]. + +from typing import Optional + +import torch +from starVLA.training.trainer_utils import initialize_overwatch +from PIL import Image +from transformers import AutoModelForCausalLM, AutoProcessor +from transformers.modeling_outputs import CausalLMOutputWithPast + +logger = initialize_overwatch(__name__) + +# IGNORE_INDEX = -100 +# IMAGE_TOKEN_INDEX = 151655 +# VIDEO_TOKEN_INDEX = 151656 +# DEFAULT_IMAGE_TOKEN = "" +# DEFAULT_VIDEO_TOKEN = "