Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,246 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+wandb/
+
+# C extensions
+
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+*.swp
+*.swo
+gear_working_dir/
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+/groot/control/wbc_checkpoints/
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+.vscode/
+.cursor/
+.cursor*
+.claude/
+CLAUDE.md
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+# IDE
+.idea/
+.vscode/
+
+# log
+outputs/
+logs/
+*logs_rl*
+!external_dependencies/OpenHomie/HomieRL/legged_gym/logs/exported/policies/*.onnx
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+outputs/
+osmo/
+source/
+groot/rl/data/
+
+.DS_Store
+sim_dependencies/
+.venv/*
+
+
+# Mujoco
+MUJOCO_LOG.TXT
+playground/
+
+*.code-workspace
+batch_*.sh
+logs_*/
+runs/
+out/
+recordings/
+pyrightconfig.json
+*.npz
+*.nbize.py
+
+# Git worktrees
+/wt
+
+# Gear working directory
+gear_working_dir/
+
+episode_data/
+
+# third party packages
+third_parties/
+
+# redis related
+*.rdb
+download_model.py
+.dockerignore
+
+# Sysid files
+plots/
+sysid_data/
+isaac_data/
+
+# External dependencies
+external_dependencies/ws_lidar_slam_ros2/build
+external_dependencies/ws_lidar_slam_ros2/install
+external_dependencies/ws_lidar_slam_ros2/log
+external_dependencies/ws_slam/
+external_dependencies/Livox-SDK2/
+external_dependencies/mfm
+# external_dependencies/genmo
+external_dependencies/whole_body_tracking
+inputs/
+outputs/
+
+groot/dexmg/grootrobosuite/docs/images/
+external_dependencies/mfm/
+dump.rdb
+
+# Open loop eval plots
+offline_open_loop_eval_plots/
+
+# yam eval data
+groot/control/envs/yam/data/eval/*
+/models
+/data
+
+checkpoints/
+video_pred_output*/

COPYRIGHT

@@ -0,0 +1,4 @@
+Copyright (c) 2025 NVIDIA Corporation. All rights reserved.
+
+Licensed under the Apache License, Version 2.0.
+See LICENSE for the full license text.

LICENSE

@@ -0,0 +1,190 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to the Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by the Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding any notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   Copyright 2025 NVIDIA Corporation
+
+   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.

README.md

@@ -0,0 +1,274 @@
+# NVIDIA DreamZero: World Action Models Are Zero-Shot Policies
+A research project from [NVIDIA GEAR Lab](https://research.nvidia.com/labs/gear/).
+
+[![NVIDIA](https://img.shields.io/badge/NVIDIA-76B900?style=flat&logo=nvidia&logoColor=white)](https://www.nvidia.com) [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](LICENSE) [![arXiv](https://img.shields.io/badge/arXiv-2602.15922-b31b1b.svg)](https://arxiv.org/abs/2602.15922)
+
+[[Project Page](https://dreamzero0.github.io/)] [[Paper](https://arxiv.org/abs/2602.15922)]
+
+DreamZero is a World Action Model that jointly predicts actions and videos, achieving strong zero-shot performance on unseen tasks. This release package contains everything needed to load a pretrained DreamZero model and run distributed inference via a WebSocket server.
+
+## News
+
+- **02/27:** DreamZero is **#1 on both [MolmoSpaces]([https://huggingface.co/spaces/ai2-adapt/MolmoSpaces](https://molmospaces.allen.ai/leaderboard)) and [RoboArena]([https://robo-arena.github.io/](https://robo-arena.github.io/leaderboard))**! DreamZero-DROID is trained *from scratch* using only the DROID dataset — no pretraining on large-scale robot data, unlike competing VLAs. This demonstrates the strength of video-model backbones for generalist robot policies (VAMs/WAMs).
+- **02/27:** Released **DreamZero-AgiBot checkpoint** and **post-training code** for efficient few-shot adaptation. Post-train on just ~30 minutes of play data for your specific robot, and see the robot do basic language following and pick-and-place (see YAM experiments in our paper for more detail).
+- **02/20:** Released the **full training codebase, preprocessed dataset, and guide for new embodiments** to replicate the DreamZero-DROID checkpoint and train on your own robot. See [Adding a New Embodiment to DreamZero](docs/DATASET_TO_GEAR_AND_TRAIN.md) for a step-by-step walkthrough.
+
+## Features
+
+**Available Now**
+- Pretrained DreamZero-DROID model checkpoint [[Huggingface](https://huggingface.co/GEAR-Dreams/DreamZero-DROID)]
+- Pretrained DreamZero-AgiBot checkpoint (for post-training on new embodiments) [[Huggingface](https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot)]
+- Distributed WebSocket inference server (GB200, H100)
+- DiT caching for optimized inference (~0.6s on GB200, ~3s on H100)
+- DROID simulation evaluation support
+- [RoboArena](https://robo-arena.github.io/) integration (DROID real)
+- Video generation and saving (MP4)
+- LoRA and full fine-tuning training scripts
+- Training on new embodiments (AgiBot, YAM) — see [guide](docs/DATASET_TO_GEAR_AND_TRAIN.md)
+
+**Coming Soon**
+- [PolaRiS](https://polaris-evals.github.io/) simulation environment support
+- [Genie 3.0](https://arxiv.org/abs/2601.02078) sim environment support for DreamZero-AgiBot
+
+## Testing Out DreamZero in Simulation with API
+We provide an inference script that directly evaluates a hosted DreamZero-DROID policy on [`sim_evals`](https://github.com/arhanjain/sim-evals). To test out the policy, first request access to the API via this form [link](https://forms.gle/zCj5zjDvHsoeuMXU7). Then, follow these instructions to install [`sim_evals`](https://github.com/arhanjain/sim-evals) and launch evaluation.
+
+```bash
+# Clone repository
+git clone --recurse-submodules https://github.com/arhanjain/sim-evals.git
+cd sim-evals
+
+# Install uv
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+# Activate uv environment
+uv sync
+source .venv/bin/activate
+
+# [Optional] update pytorch versions
+pip install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu129
+
+# Download assets (may need to export HF_TOKEN=<YOUR_HUGGINGFACE_TOKEN> first)
+uvx hf download owhan/DROID-sim-environments --repo-type dataset --local-dir assets
+
+# Run eval script
+cd ..
+python eval_utils/run_sim_eval.py --host <API_HOST> --port <API_PORT> 
+```
+
+The outputs are saved in `runs` directory.
+
+
+## Quick Start
+
+### Prerequisites
+
+- **Python**: 3.11
+- **Hardware**: Multi-GPU setup (tested on GB200, H100)
+  - Minimum: 2 GPUs for distributed inference
+- **CUDA**: Compatible GPU with CUDA 12.9+
+
+### Installation
+
+1. **Create conda environment:**
+```bash
+conda create -n dreamzero python=3.11
+conda activate dreamzero
+```
+
+2. **Install dependencies (PyTorch 2.8+ with CUDA 12.9+):**
+```bash
+pip install -e . --extra-index-url https://download.pytorch.org/whl/cu129
+```
+
+3. **Install flash attention:**
+```bash
+MAX_JOBS=8 pip install --no-build-isolation flash-attn
+```
+
+4. **[GB200 ONLY, SKIP FOR H100] Install Transformer Engine:**
+```bash
+pip install --no-build-isolation transformer_engine[pytorch]
+```
+
+5. **[GB200 ONLY FOR TENSORRT, SKIP FOR H100] Install Tensorrt:**
+```bash
+pip install tensorrt==10.13.2.6 tensorrt_cu13==10.13.2.6 tensorrt_cu13_libs==10.13.2.6 tensorrt_cu13_bindings==10.13.2.6 --no-deps
+pip install transformer_engine==2.10.0 transformer_engine_cu12==2.10.0 transformer_engine_torch==2.10.0
+```
+
+## Downloading Pretrained Checkpoints
+
+### DreamZero-DROID (for inference)
+
+We release a 14B pretrained DROID checkpoint on [Huggingface](https://huggingface.co/GEAR-Dreams/DreamZero-DROID). To download the checkpoint, run
+
+```bash
+hf download GEAR-Dreams/DreamZero-DROID --repo-type model --local-dir <path/to/checkpoint>
+```
+
+### DreamZero-AgiBot (for fine-tuning on new embodiments)
+
+To fine-tune DreamZero on a new embodiment (e.g. YAM, AgiBot), download the pretrained [DreamZero-AgiBot](https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot) checkpoint (~45GB) to `./checkpoints/DreamZero-AgiBot`:
+
+```bash
+git clone https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot ./checkpoints/DreamZero-AgiBot
+```
+
+Or with the Hugging Face CLI:
+
+```bash
+hf download GEAR-Dreams/DreamZero-AgiBot --repo-type model --local-dir ./checkpoints/DreamZero-AgiBot
+```
+
+The YAM and AgiBot training scripts use `pretrained_model_path=./checkpoints/DreamZero-AgiBot` by default. See the [new embodiment guide](docs/DATASET_TO_GEAR_AND_TRAIN.md) for usage.
+
+## Running the Inference Server
+
+### Command Overview
+
+The inference server uses PyTorch distributed training utilities to parallelize the model across multiple GPUs:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.run --standalone --nproc_per_node=2 socket_test_optimized_AR.py --port 5000 --enable-dit-cache --model-path <path/to/checkpoint>
+```
+
+(Optional only for GB200) Tensorrt enables faster generation
+```bash
+export LOAD_TRT_ENGINE=<path/to/checkpoint>/tensorrt/wan/WanModel_nvfp4.trt 
+export DYNAMIC_CACHE_SCHEDULE=true 
+CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.run --standalone --nproc_per_node=2 /mnt/aws-lfs-02/shared/seonghyeony/dreamzero/socket_test_optimized_AR.py --port 8000 --enable-dit-cache --model-path <path/to/checkpoint>
+```
+To verify the server is working, run a test client. The first few inferences will take a few minutes to warm up. After warming up, inference takes ~0.6s on GB200 and ~3s on H100.
+
+```
+python test_client_AR.py --port 5000
+```
+
+### Command-line Arguments
+
+- `--port`: Port number for the WebSocket server (default: 8000)
+- `--model-path`: Path to the pretrained model checkpoint directory
+- `--enable-dit-cache`: Enable caching in DiT layers for faster inference (recommended)
+- `--max-chunk-size`: Override max_chunk_size for inference (optional)
+- `--timeout-seconds`: Server timeout in seconds (default: 50000)
+- `--index`: Index for output directory naming (default: 0)
+
+
+### Output
+
+The server saves:
+- **Videos**: Generated video predictions as MP4 files in `{model_path}/real_world_eval_gen_{date}_{index}/{checkpoint_name}/`
+- **Input observations**: Saved per message in `{output_dir}/inputs/{msg_index}_{timestamp}/`
+
+
+## Training
+
+> **Training on a new embodiment?** See [Adding a New Embodiment to DreamZero](docs/DATASET_TO_GEAR_AND_TRAIN.md) for a complete guide on converting your dataset, configuring modalities, and launching training. <em>Make sure to align the 3 camera view order to ensure positive transfer.</em>
+
+### Downloading Pretrained Base Model Weights
+
+DreamZero is built on top of [Wan2.1-I2V-14B-480P](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P) and uses the [umt5-xxl](https://huggingface.co/google/umt5-xxl) tokenizer. Download both before training:
+
+```bash
+pip install "huggingface_hub[cli]"
+
+# You may need to set your HuggingFace token:
+# export HF_TOKEN=<YOUR_HUGGINGFACE_TOKEN>
+
+# Download Wan2.1 model weights (~28GB)
+hf download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P
+
+# Download umt5-xxl tokenizer
+hf download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl
+```
+
+> **Note:** The training script will auto-download these if they are not found at the configured paths, but pre-downloading is recommended to avoid delays at launch.
+
+### DROID Dataset
+
+We release the preprocessed DROID dataset used to train DreamZero on HuggingFace: [GEAR-Dreams/DreamZero-DROID-Data](https://huggingface.co/datasets/GEAR-Dreams/DreamZero-DROID-Data).
+
+This dataset is derived from the [DROID 1.0.1](https://droid-dataset.github.io/) dataset with the following modifications:
+- Converted from RLDS/TFDS format to [LeRobot](https://github.com/huggingface/lerobot) v2.0 format
+- Idle frames removed using [Physical Intelligence's idle frame detector](https://github.com/Physical-Intelligence/openpi/blob/main/examples/droid/README_train.md#data-filtering) (`droid_sample_ranges_v1_0_1.json`)
+- Episodes without language annotations are filtered out
+- Successful episodes only (episodes with non-zero reward)
+- 3 camera views: `exterior_image_1_left`, `exterior_image_2_left`, `wrist_image_left`
+
+**To download the preprocessed dataset (~131GB):**
+
+```bash
+huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot
+```
+
+If you want to reproduce the dataset conversion from raw DROID 1.0.1 yourself (or modify the filtering), see [docs/DROID_CONVERSION.md](docs/DROID_CONVERSION.md).
+
+### Running Training
+
+```bash
+# Configure paths (override defaults as needed)
+export DROID_DATA_ROOT="./data/droid_lerobot"
+export OUTPUT_DIR="./checkpoints/dreamzero_droid"
+export NUM_GPUS=4
+
+# Point to your downloaded model weights (if not using default paths)
+export WAN_CKPT_DIR="./checkpoints/Wan2.1-I2V-14B-480P"
+export TOKENIZER_DIR="./checkpoints/umt5-xxl"
+
+# Launch training
+bash scripts/train/droid_training.sh
+```
+
+**Using Wan2.2-TI2V-5B backbone (5B params, lower VRAM):** To train with the smaller Wan2.2-TI2V-5B model instead of Wan2.1-I2V-14B, see [docs/WAN22_BACKBONE.md](docs/WAN22_BACKBONE.md) and run `bash scripts/train/droid_training_wan22.sh`.
+
+### Training Configuration
+
+The training script uses Hydra for configuration and DeepSpeed ZeRO Stage 2 for distributed training. Key defaults:
+
+| Parameter | Default | Description |
+|---|---|---|
+| `NUM_GPUS` | 4 | Number of GPUs |
+| `per_device_train_batch_size` | 1 | Batch size per GPU |
+| `learning_rate` | 1e-5 | Learning rate |
+| `max_steps` | 10 | Max training steps (increase for full training) |
+| `warmup_ratio` | 0.05 | Warmup ratio |
+| `weight_decay` | 1e-5 | Weight decay |
+| `image_resolution_width` | 320 | Image width |
+| `image_resolution_height` | 176 | Image height |
+| `num_frames` | 33 | Number of video frames |
+| `action_horizon` | 24 | Action prediction horizon |
+| `save_lora_only` | true | Only save LoRA weights |
+| `bf16` | true | Use bfloat16 precision |
+
+> **Note:** `max_steps=10` is set for a quick sanity check. For full training, increase this to your desired number of steps and configure `save_steps` / `save_strategy` accordingly.
+
+
+## Citation
+
+If you use DreamZero in your research, please cite:
+
+```bibtex
+@misc{ye2026worldactionmodelszeroshot,
+      title={World Action Models are Zero-shot Policies}, 
+      author={Seonghyeon Ye and Yunhao Ge and Kaiyuan Zheng and Shenyuan Gao and Sihyun Yu and George Kurian and Suneel Indupuru and You Liang Tan and Chuning Zhu and Jiannan Xiang and Ayaan Malik and Kyungmin Lee and William Liang and Nadun Ranawaka and Jiasheng Gu and Yinzhen Xu and Guanzhi Wang and Fengyuan Hu and Avnish Narayan and Johan Bjorck and Jing Wang and Gwanghyun Kim and Dantong Niu and Ruijie Zheng and Yuqi Xie and Jimmy Wu and Qi Wang and Ryan Julian and Danfei Xu and Yilun Du and Yevgen Chebotar and Scott Reed and Jan Kautz and Yuke Zhu and Linxi "Jim" Fan and Joel Jang},
+      year={2026},
+      eprint={2602.15922},
+      archivePrefix={arXiv},
+      primaryClass={cs.RO},
+      url={https://arxiv.org/abs/2602.15922}, 
+}
+```
+
+## License
+
+This project is licensed under the [Apache License 2.0](LICENSE).
+
+## Support
+
+For issues and questions:
+- Check the troubleshooting section above
+- Review server logs for detailed error messages
+- Verify your checkpoint is compatible with this release
+
+[![Star History Chart](https://api.star-history.com/svg?repos=dreamzero0/dreamzero&type=Date)](https://star-history.com/#dreamzero0/dreamzero&Date)

docker/Dockerfile

@@ -0,0 +1,71 @@
+# DreamZero 训练 Docker 镜像
+#
+# 构建:
+#   docker build -t dreamzero:latest -f docker/Dockerfile .
+#
+# 运行 (单节点 8 GPU):
+#   docker run --gpus all --shm-size=64g \
+#     -v /path/to/checkpoints:/workspace/dreamzero/checkpoints \
+#     -v /path/to/data:/workspace/dreamzero/data \
+#     -v /path/to/output:/workspace/dreamzero/output \
+#     dreamzero:latest \
+#     data=dreamzero/libero \
+#     train_architecture=full \
+#     max_steps=200000 \
+#     output_dir=./output/libero_full
+
+FROM nvidia/cuda:12.4.1-devel-ubuntu22.04
+
+LABEL description="DreamZero: Video Diffusion Model for Scalable Robot Learning"
+LABEL maintainer="DreamZero Team"
+
+# ============ 系统依赖 ============
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3.11 \
+    python3.11-dev \
+    python3.11-distutils \
+    python3-pip \
+    git \
+    ffmpeg \
+    libsm6 \
+    libxext6 \
+    libglib2.0-0 \
+    libxrender-dev \
+    libgl1-mesa-glx \
+    wget \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+# 设置 python3.11 为默认
+RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1
+
+# ============ 安装 PyTorch ============
+RUN pip install --no-cache-dir \
+    torch==2.4.0 \
+    torchvision==0.19.0 \
+    --index-url https://download.pytorch.org/whl/cu124
+
+# ============ 复制代码 ============
+WORKDIR /workspace/dreamzero
+COPY groot/ ./groot/
+COPY scripts/ ./scripts/
+COPY pyproject.toml requirements.txt ./
+
+# ============ 安装依赖 ============
+RUN pip install --no-cache-dir -e ".[train]"
+
+# ============ NCCL 配置（多节点优化） ============
+ENV NCCL_DEBUG=WARN
+ENV NCCL_IB_DISABLE=0
+ENV NCCL_SOCKET_IFNAME=^docker0,lo
+ENV NCCL_IB_TIMEOUT=22
+ENV NCCL_IB_RETRY_CNT=4
+
+# ============ 中国用户镜像 ============
+ENV HF_ENDPOINT=https://hf-mirror.com
+
+# ============ 默认入口 ============
+ENTRYPOINT ["torchrun", "--standalone", "--nproc_per_node", "8", \
+    "/workspace/dreamzero/groot/vla/experiment/experiment.py"]
+CMD ["--help"]

docs/DATASET_TO_GEAR_AND_TRAIN.md

@@ -0,0 +1,471 @@
+# Adding a New Embodiment to DreamZero
+
+How to take a LeRobot v2 dataset for a new robot, convert it to GEAR format, define its modality config, and train a DreamZero policy.
+
+Throughout this guide, replace `<EMBODIMENT>` with your robot's name (e.g. `myrobot`, `franka`, `aloha`).
+
+---
+
+## Overview
+
+```
+Step 1  Convert LeRobot v2 dataset → GEAR metadata
+Step 2  Register the embodiment tag
+Step 3  Add modality config + transforms to base YAML
+Step 4  Create a dataset YAML
+Step 5  Create a training script
+Step 6  Train
+```
+
+---
+
+## Step 1: Convert Dataset to GEAR Format
+
+The converter reads a LeRobot v2 dataset and generates the metadata files DreamZero needs. It does **not** modify your parquet files or videos — it only writes to `meta/`.
+
+### Expected input structure
+
+```
+your_dataset/
+├── data/
+│   └── chunk-000/
+│       ├── episode_000000.parquet
+│       └── ...
+├── videos/
+│   └── chunk-000/
+│       ├── observation.images.cam0/
+│       │   ├── episode_000000.mp4
+│       │   └── ...
+│       └── observation.images.cam1/
+│           └── ...
+└── meta/
+    └── info.json          # must contain: features, total_episodes, fps
+```
+
+### Run the converter
+
+```bash
+python scripts/data/convert_lerobot_to_gear.py \
+    --dataset-path /path/to/your_dataset \
+    --embodiment-tag <EMBODIMENT> \
+    --state-keys '{"joint_pos": [0, 6], "gripper_pos": [6, 7]}' \
+    --action-keys '{"joint_pos": [0, 6], "gripper_pos": [6, 7]}' \
+    --relative-action-keys joint_pos gripper_pos \
+    --task-key annotation.task
+```
+
+`--state-keys` and `--action-keys` tell the converter how to split a packed vector column into named sub-keys. The JSON maps sub-key name → `[start_index, end_index]`. Omit these flags to let the converter auto-detect.
+
+### Arguments
+
+| Argument | Default | Description |
+|---|---|---|
+| `--dataset-path` | *(required)* | Path to the LeRobot v2 dataset |
+| `--output-path` | *(in-place)* | Write to a different directory instead of in-place |
+| `--embodiment-tag` | `xdof` | Tag for `meta/embodiment.json`; must match the key you use in Step 3 |
+| `--state-keys` | *(auto)* | JSON: sub-key name → `[start, end]` index range |
+| `--action-keys` | *(auto)* | JSON: sub-key name → `[start, end]` index range |
+| `--relative-action-keys` | *(none)* | Sub-key names to compute relative action stats for |
+| `--task-key` | *(auto)* | Column name for language/task annotations |
+| `--fps` | *(from info.json)* | Override dataset FPS |
+| `--action-horizon` | `24` | Horizon for relative stats computation |
+| `--force` | `false` | Overwrite existing metadata files |
+
+### Generated files
+
+The converter creates these under `meta/`:
+
+| File | Contents |
+|---|---|
+| `modality.json` | Maps state, action, video, and annotation keys with index ranges and dtypes |
+| `embodiment.json` | `{"embodiment_tag": "<EMBODIMENT>"}` |
+| `stats.json` | Per-feature statistics (mean, std, min, max, q01, q99) |
+| `relative_stats_dreamzero.json` | Relative action statistics (action − reference state) |
+| `tasks.jsonl` | Unique task descriptions |
+| `episodes.jsonl` | Per-episode metadata (index, tasks, length) |
+
+---
+
+## Step 2: Register the Embodiment Tag
+
+1. Add to the enum in `groot/vla/data/schema/embodiment_tags.py`:
+
+```python
+class EmbodimentTag(str, Enum):
+    ...
+    MY_ROBOT = "<EMBODIMENT>"
+```
+
+2. Add to `VALID_EMBODIMENT_TAGS` in `scripts/data/convert_lerobot_to_gear.py` (if you want the converter to accept the tag without `--force`):
+
+```python
+VALID_EMBODIMENT_TAGS = [
+    ...,
+    "<EMBODIMENT>",
+]
+```
+
+---
+
+## Step 3: Add Modality Config and Transforms
+
+Edit `groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml`.
+
+### 3a. Understanding modality
+
+**Modality** connects your dataset columns to the training pipeline. The `modality.json` from Step 1 contains entries like:
+
+```json
+{
+  "state": {
+    "joint_pos": {"original_key": "observation.state", "start": 0, "end": 6},
+    "gripper_pos": {"original_key": "observation.state", "start": 6, "end": 7}
+  },
+  "action": {
+    "joint_pos": {"original_key": "action", "start": 0, "end": 6},
+    "gripper_pos": {"original_key": "action", "start": 6, "end": 7}
+  },
+  "video": {
+    "cam0": {"original_key": "observation.images.cam0"}
+  },
+  "annotation": {
+    "task": {"original_key": "annotation.task"}
+  }
+}
+```
+
+The YAML config must reference **exactly these key names** with type prefixes:
+
+| Modality | YAML key format | Example |
+|---|---|---|
+| State | `state.<name>` | `state.joint_pos` |
+| Action | `action.<name>` | `action.joint_pos` |
+| Video | `video.<name>` | `video.cam0` |
+| Language | `annotation.<name>` | `annotation.task` |
+
+If the YAML keys don't match `modality.json`, training will fail with missing-key errors.
+
+### 3b. Add `modality_config_<EMBODIMENT>`
+
+```yaml
+modality_config_<EMBODIMENT>:
+  video:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_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, 24]
+    eval_delta_indices: [0]
+    modality_keys:                        # one entry per camera, matching modality.json
+      - video.cam0
+      - video.cam1
+      - video.cam2
+  state:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:                        # matching modality.json state keys
+      - state.joint_pos
+      - state.gripper_pos
+  action:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_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]
+    modality_keys:                        # matching modality.json action keys
+      - action.joint_pos
+      - action.gripper_pos
+  language:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - annotation.task
+```
+
+**`delta_indices` explained:**
+
+- **Video** — frame offsets to sample (25 entries = 25 frames from the trajectory).
+- **State / Language** — `[0]` = current timestep only.
+- **Action** — future offsets (24 entries = 24-step action chunk).
+
+Adjust these to match your `num_frames` and `action_horizon` training settings.
+
+### 3c. Add `transform_<EMBODIMENT>`
+
+```yaml
+transform_<EMBODIMENT>:
+  _target_: groot.vla.data.transform.ComposedModalityTransform
+  transforms:
+    # Video
+    - <<: *totensor_cfg
+      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
+    - <<: *crop_cfg
+      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
+    - <<: *resize_cfg
+      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
+    - <<: *color_jitter_cfg
+      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
+    - <<: *to_numpy_cfg
+      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
+
+    # State
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_<EMBODIMENT>.state.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_<EMBODIMENT>.state.modality_keys}
+      normalization_modes:
+        state.joint_pos: q99              # every state key needs a normalization mode
+        state.gripper_pos: q99
+
+    # Action
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_<EMBODIMENT>.action.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_<EMBODIMENT>.action.modality_keys}
+      normalization_modes:
+        action.joint_pos: q99             # every action key needs a normalization mode
+        action.gripper_pos: q99
+
+    # Concat
+    - _target_: groot.vla.data.transform.ConcatTransform
+      video_concat_order: ${modality_config_<EMBODIMENT>.video.modality_keys}
+      state_concat_order: ${modality_config_<EMBODIMENT>.state.modality_keys}
+      action_concat_order: ${modality_config_<EMBODIMENT>.action.modality_keys}
+
+    # Model-specific (required, don't change)
+    - ${model_specific_transform}
+```
+
+Every state and action key **must** appear in `normalization_modes`. The strategy is typically `q99`.
+
+### 3d. Register in the global maps
+
+Add your embodiment to each of these four maps (at the bottom of the base YAML):
+
+```yaml
+modality_configs:
+  ...
+  <EMBODIMENT>: ${modality_config_<EMBODIMENT>}
+
+transforms:
+  ...
+  <EMBODIMENT>: ${transform_<EMBODIMENT>}
+
+metadata_versions:
+  ...
+  <EMBODIMENT>: '0221'
+
+fps:
+  ...
+  <EMBODIMENT>: 30          # set to your dataset's FPS
+```
+
+---
+
+## Step 4: Create a Dataset YAML
+
+Create `groot/vla/configs/data/dreamzero/<EMBODIMENT>_relative.yaml`:
+
+```yaml
+# @package _global_
+
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_
+
+max_state_dim: 64
+use_global_metadata: false
+relative_action: true
+relative_action_per_horizon: false
+relative_action_keys:
+  - joint_pos                    # sub-key names (without state./action. prefix)
+  - gripper_pos                  # that should use relative actions
+max_chunk_size: 5
+dataset_shard_sampling_rate: 0.1
+mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec
+single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID
+
+<EMBODIMENT>_data_root: ???     # set via CLI or env var
+
+train_dataset:
+  _target_: ${mixture_dataset_cls}
+  _convert_: object
+  mixture_spec:
+    - dataset_path:
+        <EMBODIMENT>:            # must match key in modality_configs/transforms
+          - ${<EMBODIMENT>_data_root}
+      dataset_weight: 1.0
+      distribute_weights: true
+
+  dataset_class: ${single_dataset_cls}
+  all_modality_configs: ${modality_configs}
+  all_transforms: ${transforms}
+  metadata_versions: ${metadata_versions}
+  fps: ${fps}
+  dataset_kwargs:
+    video_backend: decord
+    use_global_metadata: ${use_global_metadata}
+    max_chunk_size: ${max_chunk_size}
+    relative_action: ${relative_action}
+    relative_action_keys: ${relative_action_keys}
+    relative_action_per_horizon: ${relative_action_per_horizon}
+  mixture_kwargs:
+    training: true
+    balance_dataset_weights: false
+    seed: 42
+    shard_sampling_rate: ${dataset_shard_sampling_rate}
+```
+
+The critical things to get right:
+
+- `<EMBODIMENT>` in `mixture_spec.dataset_path` must match the key in `modality_configs` and `transforms`.
+- `relative_action_keys` lists the sub-key names (without `state.`/`action.` prefix) that exist in **both** state and action modalities.
+
+---
+
+## Step 5: Create a Training Script
+
+Create `scripts/train/<EMBODIMENT>_training.sh`:
+
+```bash
+#!/bin/bash
+export HYDRA_FULL_ERROR=1
+
+# ============ CONFIGURATION ============
+DATA_ROOT=${DATA_ROOT:?"Set DATA_ROOT to your GEAR-converted dataset"}
+OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero_<EMBODIMENT>_lora"}
+
+if [ -z "${NUM_GPUS:-}" ]; then
+  NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
+fi
+NUM_GPUS=${NUM_GPUS:-8}
+
+WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"}
+TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"}
+# =======================================
+
+# Auto-download weights if missing
+if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then
+    huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN_CKPT_DIR"
+fi
+if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then
+    huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR"
+fi
+
+if [ ! -d "$DATA_ROOT" ]; then
+    echo "ERROR: Dataset not found at $DATA_ROOT"
+    exit 1
+fi
+if [ ! -f "$DATA_ROOT/meta/embodiment.json" ]; then
+    echo "ERROR: meta/embodiment.json missing — run convert_lerobot_to_gear.py first"
+    exit 1
+fi
+
+torchrun --nproc_per_node $NUM_GPUS --standalone \
+    groot/vla/experiment/experiment.py \
+    report_to=wandb \
+    data=dreamzero/<EMBODIMENT>_relative \
+    wandb_project=dreamzero \
+    train_architecture=lora \
+    num_frames=33 \
+    action_horizon=24 \
+    num_views=3 \
+    model=dreamzero/vla \
+    model/dreamzero/action_head=wan_flow_matching_action_tf \
+    model/dreamzero/transform=dreamzero_cotrain \
+    num_frame_per_block=2 \
+    num_action_per_block=24 \
+    num_state_per_block=1 \
+    seed=42 \
+    training_args.learning_rate=1e-5 \
+    training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \
+    save_steps=10000 \
+    training_args.warmup_ratio=0.05 \
+    output_dir=$OUTPUT_DIR \
+    per_device_train_batch_size=4 \
+    max_steps=100000 \
+    weight_decay=1e-5 \
+    save_total_limit=10 \
+    upload_checkpoints=false \
+    bf16=true \
+    tf32=true \
+    eval_bf16=true \
+    dataloader_pin_memory=false \
+    dataloader_num_workers=1 \
+    image_resolution_width=320 \
+    image_resolution_height=176 \
+    save_lora_only=true \
+    max_chunk_size=4 \
+    frame_seqlen=880 \
+    save_strategy=steps \
+    <EMBODIMENT>_data_root=$DATA_ROOT \
+    dit_version=$WAN_CKPT_DIR \
+    text_encoder_pretrained_path=$WAN_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \
+    image_encoder_pretrained_path=$WAN_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \
+    vae_pretrained_path=$WAN_CKPT_DIR/Wan2.1_VAE.pth \
+    tokenizer_path=$TOKENIZER_DIR \
+    pretrained_model_path=./checkpoints/DreamZero-AgiBot \
+    ++action_head_cfg.config.skip_component_loading=true \
+    ++action_head_cfg.config.defer_lora_injection=true
+```
+
+### Key parameters to adjust per embodiment
+
+| Parameter | Default | When to change |
+|---|---|---|
+| `num_views` | `3` | Number of cameras your robot has |
+| `action_horizon` | `24` | Must match the number of action `delta_indices` |
+| `num_frames` | `33` | Must be `len(video delta_indices) + num_frame_per_block * (blocks - 1)` |
+| `image_resolution_width` | `320` | Match your camera resolution (or desired resize) |
+| `image_resolution_height` | `176` | Match your camera resolution (or desired resize) |
+| `max_steps` | `100000` | Scale with dataset size |
+| `per_device_train_batch_size` | `4` | Adjust for GPU memory |
+
+---
+
+## Step 6: Train
+
+### Download the pretrained checkpoint
+
+The training scripts load from a pretrained DreamZero checkpoint for LoRA fine-tuning. Download [DreamZero-AgiBot](https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot) (~45GB) to `./checkpoints/DreamZero-AgiBot`:
+
+```bash
+git clone https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot ./checkpoints/DreamZero-AgiBot
+```
+
+Or with the Hugging Face CLI:
+
+```bash
+hf download GEAR-Dreams/DreamZero-AgiBot --repo-type model --local-dir ./checkpoints/DreamZero-AgiBot
+```
+
+### Launch training
+
+```bash
+DATA_ROOT=/path/to/your_dataset bash scripts/train/<EMBODIMENT>_training.sh
+
+# With overrides:
+DATA_ROOT=/path/to/your_dataset OUTPUT_DIR=./checkpoints/run1 NUM_GPUS=4 \
+    bash scripts/train/<EMBODIMENT>_training.sh
+```
+
+---
+
+## Pre-Training Checklist
+
+- [ ] `meta/embodiment.json` exists and has the correct tag
+- [ ] `meta/modality.json` state/action/video/annotation keys are populated
+- [ ] `meta/stats.json` and `meta/relative_stats_dreamzero.json` exist
+- [ ] `meta/tasks.jsonl` and `meta/episodes.jsonl` exist
+- [ ] Embodiment tag in `embodiment.json` matches the key in `modality_configs` / `transforms` / `metadata_versions` / `fps`
+- [ ] YAML `modality_keys` match `modality.json` keys exactly (with `state.`/`action.`/`video.`/`annotation.` prefix)
+- [ ] Every state and action key appears in `normalization_modes` in the transform block
+- [ ] `relative_action_keys` are sub-key names that exist in both state and action
+- [ ] Wan2.1-I2V-14B-480P and umt5-xxl weights are available
+- [ ] DreamZero-AgiBot checkpoint is downloaded to `./checkpoints/DreamZero-AgiBot`
+
+---
+
+## Quick Reference: Existing Embodiments
+
+| Embodiment | Data Config | Layout |
+|---|---|---|
+| `oxe_droid` | `droid_relative.yaml` | 3 cameras, joint_position + gripper_position |
+| `agibot` | `agibot_relative.yaml` | 3 cameras, 6 state keys, 7 action keys |
+| `yam` | `yam_relative.yaml` | 3 cameras (top/left/right), bimanual left/right joint_pos + gripper_pos |
+
+Use these as concrete examples when building your own config.

docs/DROID_CONVERSION.md

@@ -0,0 +1,62 @@
+# Converting DROID from Scratch
+
+If you want to reproduce the DreamZero DROID dataset conversion yourself (or modify the filtering), follow the steps below. This requires the raw DROID 1.0.1 dataset in RLDS format and the idle filter ranges JSON.
+
+> **Most users should skip this** and simply download the preprocessed dataset:
+> ```bash
+> huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot
+> ```
+
+## Step 1: Install conversion dependencies
+
+```bash
+pip install tensorflow tensorflow-datasets polars av
+```
+
+## Step 2: Download the raw DROID 1.0.1 dataset
+
+This requires `gsutil` ([Google Cloud CLI](https://cloud.google.com/storage/docs/gsutil_install)). The full dataset is ~1.7TB.
+
+```bash
+gsutil -m cp -r gs://gresearch/robotics/droid/1.0.1 ./data/droid/1.0.1
+```
+
+> **Important:** Use version 1.0.1, not 1.0.0. Version 1.0.1 contains the complete set of language annotations (~75k episodes).
+
+## Step 3: Download the idle filter ranges
+
+This JSON file maps each episode to the frame ranges that should be kept (non-idle frames). It was originally computed by [Physical Intelligence](https://github.com/Physical-Intelligence/openpi) for training pi0-DROID models.
+
+```bash
+gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json ./data/keep_ranges.json
+```
+
+## Step 4: Run the conversion
+
+```bash
+python scripts/data/convert_droid.py \
+    ./data/droid/1.0.1 \
+    ./data/droid_lerobot \
+    --keep-ranges-path ./data/keep_ranges.json \
+    --filter-failed \
+    -n 16
+```
+
+For a quick test with a small subset:
+```bash
+python scripts/data/convert_droid.py \
+    ./data/droid/1.0.1 \
+    ./data/droid_lerobot_test \
+    --keep-ranges-path ./data/keep_ranges.json \
+    --filter-failed \
+    --first-n 5 \
+    -n 4
+```
+
+## Script reference
+
+See [`scripts/data/convert_droid.py`](scripts/data/convert_droid.py) for full usage:
+
+```
+python scripts/data/convert_droid.py --help
+```

docs/WAN22_BACKBONE.md

@@ -0,0 +1,167 @@
+# Training DreamZero with Wan2.2-TI2V-5B Backbone
+
+This guide explains how to train DreamZero on the DROID dataset using **Wan2.2-TI2V-5B** as the backbone instead of the default Wan2.1-I2V-14B.
+
+## Architecture Differences
+
+| Component | Wan2.1-I2V-14B | Wan2.2-TI2V-5B |
+|-----------|-----------------|----------------|
+| DiT dim | 5120 | 3072 |
+| DiT layers | 32 | 30 |
+| DiT heads | 16 | 24 |
+| FFN dim | 13824 | 14336 |
+| VAE latent channels | 16 | 48 |
+| VAE spatial stride | 8× | 16× |
+| Model type | i2v | ti2v |
+
+**FFN** = Feed-Forward Network: the two-layer MLP in each transformer block (Linear → GELU → Linear). FFN dim is the intermediate hidden size (e.g. 14336 for 5B).
+
+DreamZero uses a **CausalWanModel** wrapper that extends the base Wan architecture with **action/state registers** for robot policy learning. The same `CausalWanModel` class supports both Wan2.1 and Wan2.2 backbones via configuration—no new class is required. The config switches the architecture parameters (dim, in_dim, out_dim, etc.) and uses `WanVideoVAE38` for the 48-channel Wan2.2 VAE.
+
+**What action/state registers do:** The DiT sees a single sequence `[video_tokens | action_register]` where the action register is encoded action and state features (one chunk per block). All tokens share the same transformer (with causal masking and RoPE). The model learns to predict **video noise** (dynamics) and **action noise** (policy): the action-register slice is decoded by `action_decoder` to produce action noise predictions. So the model is conditioned on current state and (noisy) actions and learns to denoise both video and actions for closed-loop policy learning.
+
+## Causal masking, RoPE, and sequence layout
+
+### Causal masking
+
+In attention, **causal masking** means each position can only attend to **past and current** positions (no future). So token at index `i` can see keys at indices `j ≤ i`. That keeps the model autoregressive: it never uses future video frames or future actions when predicting the current step. In CausalWanModel the masking is **blockwise**: the first frame attends to itself; each later block of video frames can attend to the first frame plus previous (and optionally current) blocks. Action and state tokens have their own causal pattern so each action chunk only sees past video and past actions/state. This matches policy learning where you condition on observed history and predict the next action chunk.
+
+### RoPE (Rotary Position Embeddings)
+
+**RoPE** encodes position by rotating query and key vectors in a complex plane with position-dependent angles. Unlike adding a position vector, RoPE makes attention scores depend on the *relative* position of query and key, which generalizes better to longer sequences. In CausalWanModel:
+
+- **Video tokens** use **3D RoPE**: separate frequency components for frame index (time), height, and width of the patch grid. So each token knows its (t, h, w) in the video.
+- **Action and state tokens** use **1D RoPE**: a single position index along the sequence (frame/block index). So the model knows the temporal order of action chunks and state.
+
+Freqs are built in `_create_freqs()` from the patch grid size (F, H, W) and concatenated with separate 1D freqs for the action register.
+
+### Tokens, blocks, and chunks
+
+- **Token**: The smallest unit the transformer sees. After **patch_embedding** (stride 1×2×2 on the latent), one frame yields a 2D grid of tokens; the total per frame is **frame_seqlen** (e.g. 50 for 160×320). So one **token** = one patch (e.g. 1×2×2 in latent space).
+
+- **Block (image block)**: A group of consecutive **frames**, not tokens. **num_frame_per_block** (e.g. 2) frames form one “image block.” So with 33 frames you get multiple blocks. **num_image_blocks** = `(num_frames - 1) // num_frame_per_block`. Blocks are used for blockwise causal attention and to align video with action/state.
+
+- **Chunk**: In policy terms, an **action chunk** is the sequence of actions the policy outputs for one block (e.g. **num_action_per_block** = 24 actions per block). The **action register** in the DiT has one chunk per image block: for each block there are `num_action_per_block` action tokens and `num_state_per_block` state tokens. So the register length is `num_image_blocks * (num_action_per_block + num_state_per_block)`. “Chunk” and “block” are often used together: one video block corresponds to one action chunk (and one state token) in the register.
+
+Summary: **tokens** = patch-level units (50 per frame); **blocks** = groups of frames (e.g. 2 frames per block); **chunks** = per-block action (and state) outputs that are packed into the action register.
+
+## Inference: blocks, chunks, and closed-loop
+
+### How blocks and chunks are used when predicting actions
+
+At **inference**, the model predicts **one action chunk** per call, conditioned on **one block** of video (and current state):
+
+1. **Input**: A short video of the current block — e.g. `num_frame_per_block` frames (e.g. 2) — plus current **state** and (during the denoising loop) **noisy actions** for the chunk being predicted. The first time in a trajectory, the “context” is the first frame (and optionally a warm-up pass with no action to fill the KV cache).
+
+2. **DiT input**: The sequence is `[video_tokens for this block | action_register]`. The action register holds encoded **noisy** actions and **state** for this block only (one chunk). So the DiT sees: “this block of video + this chunk’s noisy actions and state.”
+
+3. **KV cache**: To keep inference causal and efficient, the model uses a **KV cache** over previous blocks. So for the *next* block, the cache already contains keys/values for earlier frames; the DiT only runs on the **new** block’s tokens plus the new action register. `current_start_frame` tells the DiT which block we’re on so RoPE and cache indexing are correct.
+
+4. **Output**: The DiT predicts **video noise** and **action noise**. The action noise is decoded by `action_decoder` into a prediction for the **current chunk**. The scheduler then updates the noisy action toward clean; after `num_inference_steps` denoising steps you get **one clean action chunk** (e.g. 24 actions).
+
+So: **one block of frames** (and state) in → **one action chunk** out. Blocks and chunks are aligned: one image block ↔ one action chunk in the register.
+
+### How DreamZero does closed-loop inference
+
+Closed-loop execution reuses the same block/chunk logic in a loop:
+
+1. **Observe**: Robot has current observation (e.g. image history + state). The policy is called with this observation (e.g. via `lazy_joint_video_action` or `lazy_joint_video_action_causal`).
+
+2. **Predict**: The action head runs the diffusion loop for the **current block**: it encodes the observed frames to latent, runs the DiT (with KV cache and `current_start_frame`) for each denoising step, and returns one denoised **action chunk** (e.g. 24 actions).
+
+3. **Execute**: The robot **executes** that chunk (e.g. 24 steps at 5 Hz → ~4.8 s). No new model call during execution.
+
+4. **Repeat**: After execution, new observation is available. The policy is called again with the new video (e.g. last N frames). If the task/language is unchanged, `current_start_frame` is incremented by `num_frame_per_block` and the KV cache is reused; the DiT only processes the **new** block and predicts the **next** action chunk. If the task or language changes (or the cache is full), the cache and `current_start_frame` are reset.
+
+So closed-loop = **repeated “one block in → one chunk out”** with KV cache across steps so the model never re-processes past frames.
+
+### What changes when you swap the backbone to 5B
+
+The **inference algorithm and API stay the same** for 14B vs 5B:
+
+- Same **block/chunk layout**: `num_frame_per_block`, `num_action_per_block`, `num_state_per_block` (and thus one block → one chunk) are defined by config and data; they do not depend on which backbone (14B vs 5B) you use.
+- Same **closed-loop flow**: `lazy_joint_video_action`, KV cache, `current_start_frame`, and the denoising loop are in the **action head** and are shared. The policy still calls the same methods (`get_action`, `lazy_joint_video_action`, etc.).
+- Same **backbone role**: The backbone only produces conditioning (e.g. text embeddings). The action head owns the DiT, VAE, and action/state encoders. So “swapping to 5B” means swapping the **action head config** (and checkpoints) to the Wan22 5B DiT + VAE38 + 160×320; the high-level inference path (backbone → action_head → one chunk) is unchanged.
+
+What **does** change with 5B:
+
+- **DiT size and layout**: 5B uses a smaller DiT (dim 3072, 30 layers, 24 heads), **frame_seqlen = 50** (for 160×320), and **no** first-frame latent concat (`concat_first_frame_latent=False`). First frame is conditioned via **CLIP** in the context, not as extra channel in the latent.
+- **VAE and resolution**: 5B uses **WanVideoVAE38** (48 channels, 16× spatial) and **160×320** video. So latent is 10×20; tokens per frame = 50.
+- **Conditioning**: 5B uses CLIP image embedding for the first frame in the context; 14B can concatenate the first-frame latent to the DiT input. The action head handles this inside the same `_forward_inference` / `_forward_blocks`; no change to the external inference API.
+
+So: **blocks and chunks** are used the same way for predicting actions at inference; **closed-loop** is the same loop of “observe → predict one chunk → execute → repeat” with KV cache; **swapping to 5B** keeps that flow and only changes the internal model (DiT/VAE) and resolution/conditioning.
+
+## Prerequisites
+
+1. **Wan2.2-TI2V-5B** weights:
+   ```bash
+   huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir ./checkpoints/Wan2.2-TI2V-5B
+   ```
+   Or clone from [Wan2.2 GitHub](https://github.com/Wan-Video/Wan2.2) and follow their download instructions.
+
+2. **Image encoder (CLIP)**: Wan2.2-TI2V-5B does not include the CLIP image encoder. Use the one from Wan2.1:
+   ```bash
+   huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P
+   ```
+   Only `models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth` is needed.
+
+3. **DROID dataset** in LeRobot format:
+   ```bash
+   huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot
+   ```
+
+## Quick Start
+
+```bash
+# Set paths (optional - defaults shown)
+export WAN22_CKPT_DIR=./checkpoints/Wan2.2-TI2V-5B
+export IMAGE_ENCODER_DIR=./checkpoints/Wan2.1-I2V-14B-480P  # for CLIP only
+export DROID_DATA_ROOT=./data/droid_lerobot
+
+# Run training
+bash scripts/train/droid_training_wan22.sh
+```
+
+## Configuration Details
+
+The Wan2.2 config (`wan_flow_matching_action_tf_wan22.yaml`) overrides:
+
+- **model/dreamzero/action_head**: `wan_flow_matching_action_tf_wan22`
+- **diffusion_model_cfg**: Wan2.2 architecture (dim=3072, in_dim=48, out_dim=48, etc.)
+- **vae_cfg**: `WanVideoVAE38` (48-channel Wan2.2 VAE)
+- **frame_seqlen**: 50 (patch output per frame)
+- **target_video_height / target_video_width**: 160 and 320 so latent spatial size is **even** (10×20 after VAE38 16×), avoiding a dynamics-loss crop. Previously 176×320 gave latent 11×20 (odd height); we use **160×320** (H×W) so both latent dimensions are even after the DiT’s stride-(1,2,2) patch embedding.
+
+For other resolutions, `frame_seqlen` must match patch output per frame; use H and W divisible by 32 for even latent:
+- 160×320 (H×W): latent 10×20 → 50
+- 176×320: latent 11×20 → 50 (odd H; loss uses crop)
+- 640×352: 220
+
+## Using with Custom Training Scripts
+
+To use Wan2.2 in your own training script, add:
+
+```bash
+model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \
+dit_version=$WAN22_CKPT_DIR \
+text_encoder_pretrained_path=$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \
+image_encoder_pretrained_path=$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \
+vae_pretrained_path=$WAN22_CKPT_DIR/Wan2.2_VAE.pth
+```
+(Do not pass `frame_seqlen`; the Wan22 config uses 50.)
+
+## File Layout
+
+```
+dreamzero/
+├── groot/vla/configs/model/dreamzero/action_head/
+│   ├── wan_flow_matching_action_tf.yaml      # Wan2.1 (default)
+│   └── wan_flow_matching_action_tf_wan22.yaml  # Wan2.2-TI2V-5B
+├── scripts/train/
+│   ├── droid_training.sh           # Wan2.1 backbone
+│   └── droid_training_wan22.sh     # Wan2.2 backbone
+└── docs/
+    └── WAN22_BACKBONE.md          # This file
+```
+
+The action head (`wan_flow_matching_action_tf.py`) automatically detects Wan2.2 vs Wan2.1 based on `in_dim` (48 vs 16) and `vae.z_dim` (48 vs 16), and loads the correct checkpoint files from the appropriate HuggingFace repos when local paths are not found.

docs/data-preparation.md

@@ -0,0 +1,194 @@
+# DreamZero 数据准备指南
+
+---
+
+## 1. 数据格式要求
+
+DreamZero 期望每个训练样本包含以下字段：
+
+| 字段 | 类型 | 形状 | 说明 |
+|------|------|------|------|
+| `video` | uint8 numpy | [T, V, H, W, 3] | T=帧数, V=视角数 |
+| `state` | float32 numpy | [state_horizon, max_state_dim] | 机器人状态（需 padding） |
+| `action` | float32 numpy | [action_horizon, max_action_dim] | 动作序列（需 padding） |
+| `language` | str | - | 任务描述 |
+
+### 参数约束
+
+```
+action_horizon / (lat_T - 1) = num_action_per_block / num_frame_per_block
+(lat_T - 1) / state_horizon = num_frame_per_block / num_state_per_block
+lat_T = num_frames // 4           # Wan2.2 VAE 4x 时间下采样
+```
+
+**标准参数**（已验证）:
+- `num_frames=12, action_horizon=12, state_horizon=1`
+- `num_frame_per_block=2, num_action_per_block=12, num_state_per_block=1`
+- `max_state_dim=44, max_action_dim=32`
+
+---
+
+## 2. State/Action Padding
+
+State 和 Action 统一 padding 到固定维度：
+
+```python
+import numpy as np
+
+MAX_STATE_DIM = 44
+MAX_ACTION_DIM = 32
+
+def pad_state(state: np.ndarray) -> np.ndarray:
+    """Pad state to MAX_STATE_DIM."""
+    d = state.shape[-1]
+    padded = np.zeros((MAX_STATE_DIM,), dtype=np.float32)
+    padded[:d] = state.astype(np.float32)
+    return padded
+
+def pad_action(action_chunk: np.ndarray) -> np.ndarray:
+    """Pad action to [horizon, MAX_ACTION_DIM]."""
+    d = action_chunk.shape[-1]
+    padded = np.zeros((*action_chunk.shape[:-1], MAX_ACTION_DIM), dtype=np.float32)
+    padded[..., :d] = action_chunk.astype(np.float32)
+    return padded
+```
+
+---
+
+## 3. 添加新数据集步骤
+
+### 3.1 创建 Dataset 类
+
+```python
+# groot/vla/data/dataset/my_dataset.py
+from pathlib import Path
+import numpy as np
+from groot.vla.model.dreamzero.transform.dreamzero_cotrain import DreamTransform
+
+class MyDataset:
+    def __init__(self, dataset_dir, num_frames=12, action_horizon=12,
+                 state_horizon=1, num_views=1, ...):
+        # 1. 扫描数据文件
+        # 2. 构建 episode 列表
+        # 3. 创建 DreamTransform 实例
+        
+        self.transform = DreamTransform(
+            default_instruction="Perform the task.",
+            max_state_dim=44, max_action_dim=32,
+            state_horizon=state_horizon,
+            action_horizon=action_horizon,
+            num_views=num_views,
+            embodiment_tag_mapping={"my_robot": 17},
+            tokenizer_path="/path/to/umt5-xxl",
+        )
+        # 必须设置 metadata
+        self.transform.set_metadata(self.merged_metadata["my_robot"])
+        self.transform.train()
+    
+    def __getitem__(self, idx):
+        # 返回格式:
+        return dict(self.transform({
+            "video": video,           # [T, V, H, W, C] uint8
+            "state": state,           # [T, D] float32
+            "action": action,         # [T, D] float32
+            "annotation.human.action.task_description": text,
+        }))
+```
+
+### 3.2 创建 Hydra 配置
+
+```yaml
+# groot/vla/configs/data/dreamzero/my_dataset.yaml
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_
+
+my_dataset_dir: /path/to/data
+
+num_frames: 12
+action_horizon: 12
+state_horizon: 1
+num_views: 1
+num_frame_per_block: 2
+num_action_per_block: 12
+num_state_per_block: 1
+max_state_dim: 44
+max_action_dim: 32
+max_chunk_size: 4
+image_resolution_height: 160
+image_resolution_width: 320
+frame_seqlen: 50
+
+train_dataset:
+  _target_: groot.vla.data.dataset.my_dataset.MyDataset
+  _convert_: object
+  dataset_dir: ${my_dataset_dir}
+  num_frames: ${num_frames}
+  action_horizon: ${action_horizon}
+  state_horizon: ${state_horizon}
+  num_views: ${num_views}
+  video_height: ${image_resolution_height}
+  video_width: ${image_resolution_width}
+```
+
+### 3.3 训练
+
+```bash
+torchrun ... experiment.py \
+  data=dreamzero/my_dataset \
+  num_frames=12 action_horizon=12 num_views=1
+```
+
+---
+
+## 4. Embodiment Tag
+
+每个数据集需要注册一个 embodiment tag（用于识别机器人种类）：
+
+```python
+# 内置 tag 映射 (部分)
+embodiment_tag_mapping = {
+    "oxe_droid": 17,   # DROID 机器人
+    "libero": 18,      # LIBERO
+    "panda": 19,       # Franka Panda
+    "manifeel": 20,    # ManiFeel
+    "robotwin": 21,    # RoboTwin
+}
+```
+
+如果不确定用哪个 tag，使用 `"oxe_droid": 17`。
+
+---
+
+## 5. 数据验证
+
+训练前验证数据格式：
+
+```python
+# 快速验证
+ds = MyDataset(dataset_dir="/path/to/data", max_episodes=3)
+sample = ds[0]
+print(sample.keys())
+for k, v in sample.items():
+    if hasattr(v, 'shape'):
+        print(f"  {k}: {v.shape}, {v.dtype}")
+
+# 验证 DreamTransform 输出
+transformed = dict(ds.transform({
+    "video": sample["video"],
+    "state": sample["state"],
+    "action": sample["action"],
+    "annotation.human.action.task_description": "test",
+}))
+```
+
+---
+
+## 6. 常见数据问题
+
+| 问题 | 原因 | 解决 |
+|------|------|------|
+| `lat_T = 0` | num_frames < 4 | 设置 num_frames >= 4 |
+| `reshape error` | num_frame_per_block 不对齐 | 确保 (lat_T-1) % num_frame_per_block == 0 |
+| `InterpolationKeyError` | config 缺少字段 | 检查 YAML 包含所有 `${...}` 引用 |
+| VAE dtype error | bf16 vs float32 不匹配 | VAE 加载时用 `dtype=self.dtype` |

docs/training-guide.md

@@ -0,0 +1,222 @@
+# DreamZero 训练指南
+
+---
+
+## 1. 超参调优
+
+### 全量 SFT 推荐参数
+
+| 参数 | 推荐值 | 说明 |
+|------|--------|------|
+| `learning_rate` | 1e-5 | 全量微调，比 LoRA 小 10x |
+| `weight_decay` | 1e-5 | 轻微权重衰减 |
+| `warmup_ratio` | 0.05 | 5% 步数线性预热 |
+| `lr_scheduler_type` | cosine | 余弦退火 |
+| `adam_beta1` | 0.95 | Adam beta1 |
+| `adam_beta2` | 0.999 | Adam beta2 |
+
+### LoRA 推荐参数
+
+| 参数 | 推荐值 |
+|------|--------|
+| `learning_rate` | 1e-4 |
+| LoRA rank | 64 |
+| LoRA alpha | 128 |
+| LoRA target modules | q_proj, k_proj, v_proj, o_proj |
+
+### 学习率调整策略
+
+- **从零开始训练**: `learning_rate=1e-5`, `warmup_ratio=0.05`
+- **从 checkpoint resume**: 学习率通常重置为初始值或减半
+- **Loss 震荡**: 减小学习率至 5e-6
+- **Loss 不降**: 检查数据预处理是否正确
+
+---
+
+## 2. Batch Size 和 Global Batch Size
+
+```bash
+# global_batch_size = per_device_train_batch_size * num_gpus * gradient_accumulation_steps
+# 32 GPU 训练: 推荐 global_batch_size=128
+global_batch_size=128
+per_device_train_batch_size=1
+# gradient_accumulation_steps = 128 / (1 * 32) = 4
+```
+
+---
+
+## 3. 显存优化决策树
+
+```
+训练时 OOM?
+├── 使用 8-bit Adam? → 未使用 → 加 optim=adamw_bnb_8bit
+├── 开 gradient_checkpointing? → 未开 → 加 model.gradient_checkpointing=true
+├── 可降低视频分辨率? → 是 → 降低 image_resolution_width/height
+├── 可减少帧数? → 是 → 减少 num_frames
+└── 以上都不行 → 使用 ZeRO-3 + CPU offload
+```
+
+### 各优化节约的显存
+
+| 优化手段 | 显存节约 | 计算开销 |
+|---------|---------|---------|
+| 8-bit Adam | ~50% 优化器显存 | 无（更慢的更新） |
+| Gradient Checkpointing | ~30% 激活显存 | ~15% 计算时间 |
+| ZeRO-3 (vs ZeRO-2) | ~20% 总显存 | ~5% 通信开销 |
+| CPU Offload | 额外 ~30% | ~20% 时间 |
+
+---
+
+## 4. Checkpoint 管理
+
+### 保存
+
+```bash
+save_steps=2000          # 每 2000 步保存
+save_total_limit=4       # 保留最近 4 个 + 最优
+save_strategy=steps      # 按步数保存
+```
+
+### Resume
+
+```bash
+# 从指定 checkpoint 恢复
+torchrun ... experiment.py \
+  ++training_args.resume_from_checkpoint=./output/libero_full/checkpoint-50000
+```
+
+**注意**: ZeRO-3 checkpoint 按 rank 分片保存，恢复时 GPU 数量必须一致。
+
+### ZeRO-3 Checkpoint 合并
+
+如果需要在不同 GPU 数之间迁移 checkpoint：
+
+```bash
+# 使用 DeepSpeed 的 zero_to_fp32.py
+python zero_to_fp32.py \
+  --checkpoint_dir ./checkpoint-50000 \
+  --output_file ./pytorch_model.bin
+```
+
+---
+
+## 5. 多节点训练注意事项
+
+### NCCL 配置
+
+```bash
+export NCCL_DEBUG=WARN           # 减少日志量
+export NCCL_IB_DISABLE=0         # 启用 InfiniBand
+export NCCL_IB_TIMEOUT=22        # IB 超时
+export NCCL_IB_RETRY_CNT=4       # IB 重试次数
+export NCCL_SOCKET_IFNAME=^docker0,lo  # 排除 docker/loopback
+```
+
+### 常见问题
+
+| 问题 | 原因 | 解决 |
+|------|------|------|
+| NCCL timeout | 网络配置错误 | 检查 NCCL_IB_*、NCCL_SOCKET_IFNAME |
+| 训练挂起 | DataLoader worker 死锁 | 设置 `dataloader_pin_memory=false` |
+| Loss 不一致 | 数据加载不统一 | 检查 `seed` 和 DataLoader shuffle |
+| OOM | 显存不足 | 参考显存优化决策树 |
+
+---
+
+## 6. DataLoader 优化
+
+| 场景 | num_workers | pin_memory | prefetch_factor |
+|------|-------------|------------|-----------------|
+| 单 GPU | 4 | false | 2 |
+| 单节点 8 GPU | 8 | true | 4 |
+| 多节点 | 8 | true | 4 |
+
+**注意**: `dataloader_pin_memory=true` 在 decord 视频后端可能导致问题，如遇训练挂起可设为 false。
+
+---
+
+## 7. 实验跟踪（Wandb）
+
+```bash
+# 配置 Wandb
+export WANDB_API_KEY=your_key
+export WANDB_ENTITY=your_entity
+
+# 启动训练
+torchrun ... experiment.py \
+  report_to=wandb \
+  wandb_project=dreamzero-sft \
+  ...
+```
+
+Wandb 记录内容：
+- `loss`: 总损失
+- `dynamics_loss_avg`: 视频预测损失
+- `action_loss_avg`: 动作预测损失
+- `learning_rate`: 学习率
+- `grad_norm`: 梯度范数
+- `training_step_time`: 每步时间
+- `model_forward_time`: 前向时间
+
+---
+
+## 8. 从 LoRA 切换到全量 SFT
+
+```bash
+# LoRA
+train_architecture=lora
+optim=adamw_torch
+learning_rate=1e-4
+
+# 全量 SFT
+train_architecture=full
+optim=adamw_bnb_8bit   # 必须！否则 5B 优化器状态 ~40GB
+learning_rate=1e-5
+```
+
+---
+
+## 9. 基准 Loss 参考
+
+| Benchmark | Step 1 | Step 5 | Step 100 | Step 1k | Step 10k | Step 100k |
+|-----------|--------|--------|----------|---------|----------|-----------|
+| LIBERO | ~2-3 | ~1.0-1.5 | ~0.8-1.0 | ~0.6-0.8 | ~0.5-0.7 | ~0.35-0.5 |
+| ManiFeel | ~1.5-2.5 | ~0.8-1.2 | ~0.6-0.8 | ~0.5-0.6 | ~0.35-0.5 | ~0.25-0.35 |
+| RoboTwin | ~3-5 | ~1.5-2.0 | ~1.0-1.5 | ~0.8-1.2 | ~0.6-0.9 | ~0.4-0.6 |
+
+Loss 持续高于参考范围的可能原因：
+- 学习率不合适
+- 数据预��理不一致（state/action 未正确 padding）
+- 视频帧率与 state 频率不对齐
+
+---
+
+## 10. 故障排除
+
+### "CUDA out of memory"
+
+```bash
+# 添加这些参数
+per_device_train_batch_size=1
+model.gradient_checkpointing=true
+optim=adamw_bnb_8bit
+training_args.deepspeed="groot/vla/configs/deepspeed/zero2_offload.json"
+```
+
+### "RuntimeError: Expected all tensors to be on the same device"
+
+可能是 VAE 或模型设备不匹配。检查 `experiment.py` 中的设备分配。
+
+### 训练突然变慢
+
+- 检查 CPU 内存（dataloader_workers 可能导致内存交换）
+- 检查磁盘 IO（数据是否在 SSD 上）
+- 减少 dataloader_num_workers
+
+### NCCL 初始化失败
+
+```bash
+export NCCL_DEBUG=INFO        # 查看详细 NCCL 日志
+export NCCL_IB_DISABLE=1      # 如果 IB 不可用
+export NCCL_SOCKET_IFNAME=eth0  # 指定网络接口
+```

eval_utils/policy_client.py

@@ -0,0 +1,93 @@
+"""Client for communicating with a policy server.
+
+Adapted from https://github.com/robo-arena/roboarena/
+
+"""
+
+import logging
+import time
+from typing import Dict, Tuple
+
+import websockets.sync.client
+from typing_extensions import override
+
+from openpi_client.base_policy import BasePolicy
+from openpi_client import msgpack_numpy
+
+# The websockets library by default sends a ping every 20 seconds and
+# expects a pong response within 20 seconds. However, the sever may not
+# send a pong response immediately if it is busy processing a request.
+# Increase the ping interval and timeout so that the client can wait
+# for a longer time before closing the connection.
+PING_INTERVAL_SECS = 60
+PING_TIMEOUT_SECS = 600
+
+class WebsocketClientPolicy(BasePolicy):
+    """Implements the Policy interface by communicating with a server over websocket.
+
+    See WebsocketPolicyServer for a corresponding server implementation.
+    """
+
+    def __init__(self, host: str = "0.0.0.0", port: int = 8000) -> None:
+        self._uri = f"ws://{host}:{port}"
+        self._packer = msgpack_numpy.Packer()
+        self._ws, self._server_metadata = self._wait_for_server()
+
+    def get_server_metadata(self) -> Dict:
+        return self._server_metadata
+
+    def _wait_for_server(self) -> Tuple[websockets.sync.client.ClientConnection, Dict]:
+        logging.info(f"Waiting for server at {self._uri}...")
+        try:
+            conn = websockets.sync.client.connect(
+                self._uri, 
+                compression=None, 
+                max_size=None,
+                ping_interval=PING_INTERVAL_SECS,
+                ping_timeout=PING_TIMEOUT_SECS,
+            )
+            metadata = msgpack_numpy.unpackb(conn.recv())
+            return conn, metadata
+        except:
+            logging.info("Connection to server with ws:// failed. Trying wss:// ...")
+            
+        self._uri = "wss://" + self._uri.split("//")[1]
+        conn = websockets.sync.client.connect(
+            self._uri, 
+            compression=None, 
+            max_size=None,
+            ping_interval=PING_INTERVAL_SECS,
+            ping_timeout=PING_TIMEOUT_SECS,
+        )
+        metadata = msgpack_numpy.unpackb(conn.recv())
+        return conn, metadata
+
+    @override
+    def infer(self, obs: Dict) -> Dict:  # noqa: UP006
+        # Notify server that we're calling the infer endpoint (as opposed to the reset endpoint)
+        obs["endpoint"] = "infer"
+
+        data = self._packer.pack(obs)
+        self._ws.send(data)
+        response = self._ws.recv()
+        if isinstance(response, str):
+            # we're expecting bytes; if the server sends a string, it's an error.
+            raise RuntimeError(f"Error in inference server:\n{response}")
+        return msgpack_numpy.unpackb(response)
+
+    @override
+    def reset(self, reset_info: Dict) -> None:
+        # Notify server that we're calling the reset endpoint (as opposed to the infer endpoint)
+        reset_info["endpoint"] = "reset"
+
+        data = self._packer.pack(reset_info)
+        self._ws.send(data)
+        response = self._ws.recv()
+        return response
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    client = WebsocketClientPolicy()
+    actions = client.infer({})
+    print(f"Actions received: {actions}")
+    client.reset({})

eval_utils/policy_server.py

@@ -0,0 +1,130 @@
+"""Server for serving a policy over websockets.
+
+Adapted from https://github.com/robo-arena/roboarena/
+
+"""
+
+
+import asyncio
+import dataclasses
+import logging
+import traceback
+
+from openpi_client.base_policy import BasePolicy
+from openpi_client import msgpack_numpy
+import websockets.asyncio.server
+import websockets.frames
+
+
+@dataclasses.dataclass
+class PolicyServerConfig:
+    # Resolution that images get resized to client-side, None means no resizing.
+    # It's beneficial to resize images to the desired resolution client-side for faster communication.
+    image_resolution: tuple[int, int] | None = (224, 224)
+    # Whether or not wrist camera image(s) should be sent.
+    needs_wrist_camera: bool = True
+    # Number of external cameras to send.
+    n_external_cameras: int = 1  # can be in [0, 1, 2]
+    # Whether or not stereo camera image(s) should be sent.
+    needs_stereo_camera: bool = False
+    # Whether or not the unique eval session id should be sent (e.g. for policies that want to keep track of history).
+    needs_session_id: bool = False
+    # Which action space to use.
+    action_space: str = "joint_position"  # can be in ["joint_position", "joint_velocity", "cartesian_position", "cartesian_velocity"]
+
+
+class WebsocketPolicyServer:
+    """
+    Serves a policy using the websocket protocol.
+
+    Interface:
+      Observation:
+        - observation/wrist_image_left: (H, W, 3) if needs_wrist_camera is True
+        - observation/wrist_image_right: (H, W, 3) if needs_wrist_camera is True and needs_stereo_camera is True
+        - observation/exterior_image_{i}_left: (H, W, 3) if n_external_cameras >= 1
+        - observation/exterior_image_{i}_right: (H, W, 3) if needs_stereo_camera is True
+        - session_id: (1,) if needs_session_id is True
+        - observation/joint_position: (7,)
+        - observation/cartesian_position: (6,)
+        - observation/gripper_position: (1,)
+        - prompt: str, the natural language task instruction for the policy
+    
+      Action:
+        - action: (N, 8,) or (N, 7,): either 7 movement actions (for joint action spaces) or 6 (for cartesian) plus one dimension for gripper position
+                           --> all N actions will get executed on the robot before the server is queried again
+
+    """
+
+    def __init__(
+        self,
+        policy: BasePolicy,
+        server_config: PolicyServerConfig,
+        host: str = "0.0.0.0",
+        port: int = 8000,
+    ) -> None:
+        self._policy = policy
+        self._server_config = server_config
+        self._host = host
+        self._port = port
+        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:
+            await server.serve_forever()
+
+    async def _handler(self, websocket: websockets.asyncio.server.ServerConnection):
+        logging.info(f"Connection from {websocket.remote_address} opened")
+        packer = msgpack_numpy.Packer()
+
+        # Send server config to client to configure what gets sent to server.
+        await websocket.send(packer.pack(dataclasses.asdict(self._server_config)))
+
+        while True:
+            try:
+                obs = msgpack_numpy.unpackb(await websocket.recv())
+                
+                endpoint = obs["endpoint"]
+                del obs["endpoint"]
+                if endpoint == "reset":
+                    self._policy.reset(obs)
+                    to_return = "reset successful"
+                else:
+                    action = self._policy.infer(obs)
+                    to_return = packer.pack(action)
+                await websocket.send(to_return)
+            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
+
+
+if __name__ == "__main__":
+    import numpy as np
+
+    class DummyPolicy(BasePolicy):
+        def infer(self, obs):
+            return np.zeros((1, 8), dtype=np.float32)
+        
+        def reset(self, reset_info):
+            pass
+    
+    logging.basicConfig(level=logging.INFO)
+    policy = DummyPolicy()
+    server = WebsocketPolicyServer(policy, PolicyServerConfig())
+    server.serve_forever()
+        

eval_utils/run_sim_eval.py

@@ -0,0 +1,219 @@
+"""
+Example script for running 10 rollouts of a DROID policy on the example environment.
+
+Usage:
+
+First, make sure you download the simulation assets and unpack them into the root directory of this package.
+
+Then, in a separate terminal, launch the policy server on localhost:8000 
+-- make sure to set XLA_PYTHON_CLIENT_MEM_FRACTION to avoid JAX hogging all the GPU memory.
+
+For example, to launch a pi0-FAST-DROID policy (with joint position control), 
+run the command below in a separate terminal from the openpi "karl/droid_policies" branch:
+
+XLA_PYTHON_CLIENT_MEM_FRACTION=0.5 uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_droid_jointpos --policy.dir=s3://openpi-assets-simeval/pi0_fast_droid_jointpos
+
+Finally, run the evaluation script:
+
+python run_eval.py --episodes 10 --headless
+"""
+
+import uuid
+
+import tyro
+import argparse
+import gymnasium as gym
+import torch
+import cv2
+import mediapy
+import numpy as np
+from datetime import datetime
+from pathlib import Path
+from PIL import Image
+from tqdm import tqdm
+
+from openpi_client import image_tools
+from sim_evals.inference.abstract_client import InferenceClient
+from policy_client import WebsocketClientPolicy
+
+
+class DreamZeroJointPosClient(InferenceClient):
+    def __init__(self, 
+                remote_host:str = "localhost", 
+                remote_port:int = 6000,
+                open_loop_horizon:int = 8,
+    ) -> None:
+        self.client = WebsocketClientPolicy(remote_host, remote_port)
+        self.open_loop_horizon = open_loop_horizon
+        self.actions_from_chunk_completed = 0
+        self.pred_action_chunk = None
+        self.session_id = str(uuid.uuid4())
+
+    def visualize(self, request: dict):
+        """
+        Return the camera views how the model sees it
+        """
+        curr_obs = self._extract_observation(request)
+        right_img = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224)
+        wrist_img = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224)
+        left_img = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224)
+        combined = np.concatenate([right_img, wrist_img, left_img], axis=1)
+        return combined
+
+    def reset(self):
+        self.actions_from_chunk_completed = 0
+        self.pred_action_chunk = None
+        self.session_id = str(uuid.uuid4())
+
+    def infer(self, obs: dict, instruction: str) -> dict:
+        """
+        Infer the next action from the policy in a server-client setup
+        """
+        curr_obs = self._extract_observation(obs)
+        if (
+            self.actions_from_chunk_completed == 0
+            or self.actions_from_chunk_completed >= self.open_loop_horizon
+        ):
+            self.actions_from_chunk_completed = 0
+            request_data = {
+                "observation/exterior_image_0_left": image_tools.resize_with_pad(curr_obs["right_image"], 180, 320),
+                "observation/exterior_image_1_left": image_tools.resize_with_pad(curr_obs["left_image"], 180, 320),
+                "observation/wrist_image_left": image_tools.resize_with_pad(curr_obs["wrist_image"], 180, 320),
+                "observation/joint_position": curr_obs["joint_position"].astype(np.float64),
+                "observation/cartesian_position": np.zeros((6,), dtype=np.float64),  # dummy cartesian position
+                "observation/gripper_position": curr_obs["gripper_position"].astype(np.float64),
+                "prompt": instruction,
+                "session_id": self.session_id,
+            }
+            for k, v in request_data.items():
+                print(f"{k}: {v.shape if not isinstance(v, str) else v}")
+            
+            result = self.client.infer(request_data)
+            actions = result["actions"] if isinstance(result, dict) else result
+            assert len(actions.shape) == 2, f"Expected 2D array, got shape {actions.shape}"
+            assert actions.shape[-1] == 8, f"Expected 8 action dimensions (7 joints + 1 gripper), got {actions.shape[-1]}"
+            self.pred_action_chunk = actions
+
+
+        action = self.pred_action_chunk[self.actions_from_chunk_completed]
+        self.actions_from_chunk_completed += 1
+
+        # binarize gripper action
+        if action[-1].item() > 0.5:
+            action = np.concatenate([action[:-1], np.ones((1,))])
+        else:
+            action = np.concatenate([action[:-1], np.zeros((1,))])
+
+        img1 = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224)
+        img2 = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224)
+        img3 = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224)
+        both = np.concatenate([img1, img2, img3], axis=1)
+
+        return {"action": action, "viz": both}
+
+    def _extract_observation(self, obs_dict, *, save_to_disk=False):
+        # Assign images
+        right_image = obs_dict["policy"]["external_cam"][0].clone().detach().cpu().numpy()
+        left_image = obs_dict["policy"]["external_cam_2"][0].clone().detach().cpu().numpy()
+        wrist_image = obs_dict["policy"]["wrist_cam"][0].clone().detach().cpu().numpy()
+
+        # Capture proprioceptive state
+        robot_state = obs_dict["policy"]
+        joint_position = robot_state["arm_joint_pos"].clone().detach().cpu().numpy()
+        gripper_position = robot_state["gripper_pos"].clone().detach().cpu().numpy()
+
+        if save_to_disk:
+            combined_image = np.concatenate([right_image, wrist_image], axis=1)
+            combined_image = Image.fromarray(combined_image)
+            combined_image.save("robot_camera_views.png")
+
+        return {
+            "right_image": right_image,
+            "left_image": left_image,
+            "wrist_image": wrist_image,
+            "joint_position": joint_position,
+            "gripper_position": gripper_position,
+        }
+
+
+
+
+def main(
+        episodes: int = 10,
+        scene: int = 1,
+        headless: bool = True,
+        host: str = "localhost",
+        port: int = 6000,
+        ):
+    # launch omniverse app with arguments (inside function to prevent overriding tyro)
+    from isaaclab.app import AppLauncher
+    parser = argparse.ArgumentParser(description="Tutorial on creating an empty stage.")
+    AppLauncher.add_app_launcher_args(parser)
+    args_cli, _ = parser.parse_known_args()
+    args_cli.enable_cameras = True
+    args_cli.headless = headless
+    app_launcher = AppLauncher(args_cli)
+    simulation_app = app_launcher.app
+
+    # All IsaacLab dependent modules should be imported after the app is launched
+    import sim_evals.environments # noqa: F401
+    from isaaclab_tasks.utils import parse_env_cfg
+
+
+    # Initialize the env
+    env_cfg = parse_env_cfg(
+        "DROID",
+        device=args_cli.device,
+        num_envs=1,
+        use_fabric=True,
+    )
+    instruction = None
+    match scene:
+        case 1:
+            instruction = "put the cube in the bowl"
+        case 2:
+            instruction = "pick up the can and put it in the mug"
+        case 3:
+            instruction = "put the banana in the bin"
+        case _:
+            raise ValueError(f"Scene {scene} not supported")
+        
+    env_cfg.set_scene(scene)
+    env = gym.make("DROID", cfg=env_cfg)
+
+    obs, _ = env.reset()
+    obs, _ = env.reset() # need second render cycle to get correctly loaded materials
+    client = DreamZeroJointPosClient(remote_host=host, remote_port=port)
+
+
+    video_dir = Path("runs") / datetime.now().strftime("%Y-%m-%d") / datetime.now().strftime("%H-%M-%S")
+    video_dir.mkdir(parents=True, exist_ok=True)
+    video = []
+    ep = 0
+    max_steps = env.env.max_episode_length
+    with torch.no_grad():
+        for ep in range(episodes):
+            for _ in tqdm(range(max_steps), desc=f"Episode {ep+1}/{episodes}"):
+                ret = client.infer(obs, instruction)
+                if not headless:
+                    cv2.imshow("Right Camera", cv2.cvtColor(ret["viz"], cv2.COLOR_RGB2BGR))
+                    cv2.waitKey(1)
+                video.append(ret["viz"])
+                action = torch.tensor(ret["action"])[None]
+                obs, _, term, trunc, _ = env.step(action)
+                if term or trunc:
+                    break
+
+            client.reset()
+            mediapy.write_video(
+                video_dir / f"episode_{ep}.mp4",
+                video,
+                fps=15,
+            )
+            video = []
+
+    env.close()
+    simulation_app.close()
+
+if __name__ == "__main__":
+    args = tyro.cli(main)

eval_utils/serve_dreamzero_wan22.py

@@ -0,0 +1,391 @@
+"""
+Serve the DreamZero 5B implementation (Wan2.2-TI2V-5B) over the websocket policy server.
+
+This is the 5B model: Wan2.2 diffusion backbone, 48-channel VAE38, frame_seqlen=50 (160×320
+latent 10×20). Inference is causal with KV caching: first request in a session uses 1 frame
+and warms the cache; subsequent requests use FRAMES_PER_CHUNK=4 frames and append to the cache.
+On session_id change (or explicit reset), buffers and action_head.current_start_frame are cleared.
+
+The checkpoint at model_path should be DreamZero with Wan22 5B (model/dreamzero/action_head=
+wan_flow_matching_action_tf_wan22, data droid_relative_wan22 → 160×320). GrootSimPolicy loads
+that checkpoint and runs inference; it is the correct policy class for DreamZero.
+
+Usage (single GPU):
+
+  torchrun --nproc_per_node=1 eval_utils/serve_dreamzero_wan22.py --model_path ./checkpoints/dreamzero_droid_wan22_smoke --port 8000
+
+  # Or single process:
+  python eval_utils/serve_dreamzero_wan22.py --model_path ./checkpoints/dreamzero_droid_wan22_smoke --port 8000
+
+Client: send observations per PolicyServerConfig (policy_server.py). Video is resized to the
+checkpoint's expected resolution (e.g. 180×320) so the eval transform accepts it; the 5B action
+head resizes to 160×320 internally. Override with --image_height/--image_width if needed.
+Response is an action chunk (N, 8). Use session_id for episode boundaries.
+"""
+
+import datetime
+import logging
+import os
+import sys
+
+import imageio
+
+logger = logging.getLogger(__name__)
+
+import cv2
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import init_device_mesh
+import tyro
+
+# Avoid FailOnRecompileLimitHit when serving: the flow scheduler's torch.compile'd
+# multistep_uni_p_bh_update recompiles under varying shapes/inputs (e.g. batch size,
+# step_index, order). Increase limits so the server doesn't hit the default cap.
+_dynamo = torch._dynamo.config
+if hasattr(_dynamo, "cache_size_limit"):
+    _dynamo.cache_size_limit = 1000
+if hasattr(_dynamo, "recompile_limit"):
+    _dynamo.recompile_limit = 800
+if hasattr(_dynamo, "accumulated_cache_size_limit"):
+    _dynamo.accumulated_cache_size_limit = 1000
+if hasattr(_dynamo, "accumulated_recompile_limit"):
+    _dynamo.accumulated_recompile_limit = 2000
+from pathlib import Path
+from tianshou.data import Batch
+
+# Add repo root for imports
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from openpi_client.base_policy import BasePolicy
+
+from eval_utils.policy_server import WebsocketPolicyServer, PolicyServerConfig
+from groot.vla.model.n1_5.sim_policy import GrootSimPolicy
+from groot.vla.data.schema import EmbodimentTag
+from groot.vla.data.transform import ComposedModalityTransform
+
+
+# DreamZero Wan 5B is trained with 160×320 (droid_relative_wan22). Fallback if we cannot read from policy.
+DEFAULT_IMAGE_HEIGHT = 160
+DEFAULT_IMAGE_WIDTH = 320
+FRAMES_PER_CHUNK = 4  # matches 5B num_frame_per_block for causal chunked inference
+
+
+def _get_expected_video_resolution(policy: GrootSimPolicy) -> tuple[int, int]:
+    """Get (height, width) the policy's eval_transform expects for video (from checkpoint
+    metadata). Resolution in metadata is (width, height); we return (height, width) for resize.
+    DreamZero Wan 5B (droid_relative_wan22) uses 160×320; other configs may use e.g. 180×320.
+    """
+    eval_transform = getattr(policy, "eval_transform", None)
+    if eval_transform is None:
+        return (DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH)
+    if not isinstance(eval_transform, ComposedModalityTransform):
+        return (DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH)
+    for t in eval_transform.transforms:
+        if hasattr(t, "original_resolutions") and getattr(t, "original_resolutions", None):
+            res = t.original_resolutions
+            if res:
+                # original_resolutions values are (width, height)
+                w, h = next(iter(res.values()))
+                return (int(h), int(w))
+    return (DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH)
+
+
+def _resize_frames_to_resolution(frames: np.ndarray, target_h: int, target_w: int) -> np.ndarray:
+    """Resize video frames to (target_h, target_w). Accepts (H,W,C) or (T,H,W,C)."""
+    if frames.ndim == 3:
+        if (frames.shape[0], frames.shape[1]) != (target_h, target_w):
+            frames = cv2.resize(frames, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
+        return frames
+    out = np.stack(
+        [cv2.resize(f, (target_w, target_h), interpolation=cv2.INTER_LINEAR) for f in frames],
+        axis=0,
+    )
+    return out
+
+
+def _maybe_init_distributed():
+    """Initialize process group for single-GPU or multi-GPU. Required by GrootSimPolicy."""
+    if dist.is_initialized():
+        return
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29500")
+    dist.init_process_group(backend="nccl", rank=0, world_size=1)
+    torch.cuda.set_device(0)
+
+
+# Modality key mappings: client observation keys -> model input keys per embodiment.
+# Client sends: observation/exterior_image_0_left, exterior_image_1_left, wrist_image_left.
+VIDEO_KEY_MAPPING = {
+    "oxe_droid": {
+        "observation/exterior_image_0_left": "video.exterior_image_1_left",
+        "observation/exterior_image_1_left": "video.exterior_image_2_left",
+        "observation/wrist_image_left": "video.wrist_image_left",
+    },
+}
+STATE_KEY_MAPPING = {
+    "oxe_droid": ("state.joint_position", "state.gripper_position"),
+}
+LANGUAGE_KEY_MAPPING = {
+    "oxe_droid": "annotation.language.action_text",
+}
+
+
+class DreamZeroWan225BPolicy(BasePolicy):
+    """
+    Wraps GrootSimPolicy for the DreamZero 5B implementation (Wan2.2-TI2V-5B).
+
+    Converts roboarena observation/action format to DROID/Batch. Video is resized to the
+    resolution expected by the policy's eval_transform (from checkpoint metadata) so
+    VideoToTensor validation passes. The 5B action head then resizes to 160×320 internally.
+    First call in a session uses 1 frame; later calls use 4 frames (FRAMES_PER_CHUNK).
+    Session reset clears frame buffers and action_head.current_start_frame.
+    """
+
+    def __init__(
+        self,
+        groot_policy: GrootSimPolicy,
+        image_height: int,
+        image_width: int,
+        embodiment_tag: str = "oxe_droid",
+        save_video_pred: bool = False,
+        video_output_dir: str = "./video_pred_output",
+    ):
+        super().__init__()
+        self._policy = groot_policy
+        self._image_height = image_height
+        self._image_width = image_width
+        self._embodiment_tag = (
+            embodiment_tag if embodiment_tag in VIDEO_KEY_MAPPING else "oxe_droid"
+        )
+        video_keys = list(VIDEO_KEY_MAPPING[self._embodiment_tag].values())
+        self._frame_buffers = {k: [] for k in video_keys}
+        self._is_first_call = True
+        self._current_session_id = None
+        self._save_video_pred = save_video_pred
+        self._video_output_dir = video_output_dir
+        self._video_pred_latents: list[torch.Tensor] = []
+        self._current_prompt: str = ""
+
+    def _convert_observation(self, obs: dict) -> dict:
+        """Convert roboarena observation format to model Batch format.
+        Incoming frames are resized to the policy's expected (height, width) so
+        eval_transform's VideoToTensor check passes.
+        """
+        image_key_mapping = VIDEO_KEY_MAPPING[self._embodiment_tag]
+        for roboarena_key, model_key in image_key_mapping.items():
+            if roboarena_key in obs:
+                data = obs[roboarena_key]
+                if isinstance(data, np.ndarray):
+                    data = _resize_frames_to_resolution(
+                        data, self._image_height, self._image_width
+                    )
+                    if data.ndim == 4:
+                        self._frame_buffers[model_key].extend(list(data))
+                    else:
+                        self._frame_buffers[model_key].append(data)
+
+        num_frames = 1 if self._is_first_call else FRAMES_PER_CHUNK
+        converted = {}
+        for model_key, buffer in self._frame_buffers.items():
+            if len(buffer) > 0:
+                if len(buffer) >= num_frames:
+                    frames_to_use = buffer[-num_frames:]
+                else:
+                    frames_to_use = buffer.copy()
+                    while len(frames_to_use) < num_frames:
+                        frames_to_use.insert(0, buffer[0])
+                video = np.stack(frames_to_use, axis=0)
+                converted[model_key] = video
+
+        state_joint_key, state_gripper_key = STATE_KEY_MAPPING[self._embodiment_tag]
+        if "observation/joint_position" in obs:
+            joint_pos = np.asarray(obs["observation/joint_position"])
+            if joint_pos.ndim == 1:
+                joint_pos = joint_pos.reshape(1, -1)
+            converted[state_joint_key] = joint_pos.astype(np.float64)
+        else:
+            converted[state_joint_key] = np.zeros((1, 7), dtype=np.float64)
+
+        if "observation/gripper_position" in obs:
+            gripper_pos = np.asarray(obs["observation/gripper_position"])
+            if gripper_pos.ndim == 1:
+                gripper_pos = gripper_pos.reshape(1, -1)
+            converted[state_gripper_key] = gripper_pos.astype(np.float64)
+        else:
+            converted[state_gripper_key] = np.zeros((1,1), dtype=np.float64)
+
+        text_prompt = obs.get("prompt", "")
+        logger.info("Text prompt: %s", text_prompt)
+        if text_prompt:
+            self._current_prompt = text_prompt
+        lang_key = LANGUAGE_KEY_MAPPING[self._embodiment_tag]
+        converted[lang_key] = text_prompt
+        return converted
+
+    def _convert_action(self, action_dict: dict) -> np.ndarray:
+        """Convert model action dict to (N, 8) array (7 joint + 1 gripper)."""
+        joint_action = None
+        gripper_action = None
+        for key, value in action_dict.items():
+            if ("joint_position" in key or "joint_pos" in key) and "gripper" not in key:
+                joint_action = value
+            elif "gripper_position" in key or "gripper" in key:
+                gripper_action = value
+        if joint_action is None:
+            return np.zeros((1, 8), dtype=np.float32)
+        if isinstance(joint_action, torch.Tensor):
+            joint_action = joint_action.cpu().numpy()
+        if joint_action.ndim == 1:
+            joint_action = joint_action.reshape(1, -1)
+        N = joint_action.shape[0]
+        if gripper_action is not None:
+            if isinstance(gripper_action, torch.Tensor):
+                gripper_action = gripper_action.cpu().numpy()
+            if gripper_action.ndim == 1:
+                gripper_action = gripper_action.reshape(-1, 1)
+            if gripper_action.shape[-1] > 1:
+                gripper_action = gripper_action[..., :1]
+        else:
+            gripper_action = np.zeros((N, 1), dtype=np.float32)
+        return np.concatenate([joint_action, gripper_action], axis=-1).astype(np.float32)
+
+    def infer(self, obs: dict) -> np.ndarray:
+        session_id = obs.get("session_id")
+        if session_id is not None and session_id != self._current_session_id:
+            if self._current_session_id is not None:
+                self.reset({})
+            self._current_session_id = session_id
+
+        converted_obs = self._convert_observation(obs)
+        batch = Batch(obs=converted_obs)
+        with torch.no_grad():
+            result_batch, video_pred = self._policy.lazy_joint_forward_causal(batch)
+        if self._save_video_pred and video_pred is not None:
+            self._video_pred_latents.append(video_pred.detach())
+        action_dict = {}
+        action_chunk_dict = result_batch.act
+        for k in dir(action_chunk_dict):
+            if k.startswith("action."):
+                action_dict[k] = getattr(action_chunk_dict, k)
+        action = self._convert_action(action_dict)
+        if self._is_first_call:
+            self._is_first_call = False
+        return action
+
+    def _save_predicted_video(self) -> None:
+        """Decode accumulated video prediction latents through the VAE and save as mp4."""
+        if not self._video_pred_latents:
+            return
+        try:
+            from einops import rearrange
+
+            action_head = self._policy.trained_model.action_head
+            latents = torch.cat(self._video_pred_latents, dim=2)
+            with torch.no_grad():
+                frames = action_head.vae.decode(
+                    latents,
+                    tiled=action_head.tiled,
+                    tile_size=(action_head.tile_size_height, action_head.tile_size_width),
+                    tile_stride=(action_head.tile_stride_height, action_head.tile_stride_width),
+                )
+            frames = rearrange(frames, "B C T H W -> B T H W C")[0]
+            frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8)
+
+            os.makedirs(self._video_output_dir, exist_ok=True)
+            timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S")
+            n_latent_frames = latents.shape[2]
+            existing = [f for f in os.listdir(self._video_output_dir) if f.endswith(".mp4")]
+            safe_prompt = self._current_prompt.replace(" ", "_")
+            safe_prompt = "".join(c for c in safe_prompt if c.isalnum() or c in "_-.")
+            if len(safe_prompt) > 80:
+                safe_prompt = safe_prompt[:80]
+            if not safe_prompt:
+                safe_prompt = "no_prompt"
+            output_path = os.path.join(
+                self._video_output_dir,
+                f"{len(existing):06}_{safe_prompt}_{timestamp}.mp4",
+            )
+            imageio.mimsave(output_path, list(frames), fps=5, codec="libx264")
+            logger.info("Saved video prediction (%d frames) to %s", len(frames), output_path)
+        except Exception as e:
+            logger.warning("Failed to save video prediction: %s", e)
+
+    def reset(self, reset_info: dict) -> None:
+        if self._save_video_pred:
+            self._save_predicted_video()
+        self._video_pred_latents.clear()
+        self._current_prompt = ""
+        for key in self._frame_buffers:
+            self._frame_buffers[key] = []
+        self._is_first_call = True
+        self._current_session_id = None
+        if hasattr(self._policy.trained_model, "action_head") and hasattr(
+            self._policy.trained_model.action_head, "current_start_frame"
+        ):
+            self._policy.trained_model.action_head.current_start_frame = 0
+
+
+def main(
+    model_path: str = "./checkpoints/dreamzero_droid_wan22_smoke",
+    embodiment_tag: str = "oxe_droid",
+    tokenizer_path: str | None = None,
+    port: int = 8000,
+    host: str = "0.0.0.0",
+    image_height: int | None = None,
+    image_width: int | None = None,
+    save_video_pred: bool = False,
+    video_output_dir: str = "./video_pred_output",
+) -> None:
+    logging.basicConfig(level=logging.INFO, force=True)
+
+    _maybe_init_distributed()
+    device_mesh = init_device_mesh("cuda", mesh_shape=(1,), mesh_dim_names=("ip",))
+
+    logger.info("Loading DreamZero Wan22 policy from %s (embodiment=%s)", model_path, embodiment_tag)
+    checkpoint_name = os.path.basename(model_path.rstrip("/"))
+    video_output_dir = os.path.join(video_output_dir, checkpoint_name)
+    policy = GrootSimPolicy(
+        embodiment_tag=EmbodimentTag(embodiment_tag),
+        model_path=model_path,
+        tokenizer_path_override=tokenizer_path,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        device_mesh=device_mesh,
+    )
+    if image_height is not None and image_width is not None:
+        h, w = image_height, image_width
+        logger.info("Using CLI video resolution: %dx%d", h, w)
+    else:
+        h, w = _get_expected_video_resolution(policy)
+        logger.info("Using checkpoint video resolution: %dx%d (HxW)", h, w)
+    wrapper = DreamZeroWan225BPolicy(
+        groot_policy=policy,
+        image_height=h,
+        image_width=w,
+        embodiment_tag=embodiment_tag,
+        save_video_pred=save_video_pred,
+        video_output_dir=video_output_dir,
+    )
+
+    server_config = PolicyServerConfig(
+        image_resolution=(h, w),
+        needs_wrist_camera=True,
+        n_external_cameras=2,
+        needs_stereo_camera=False,
+        needs_session_id=True,
+        action_space="joint_position",
+    )
+    logger.info("Starting WebsocketPolicyServer on %s:%d (DreamZero 5B, %dx%d)", host, port, h, w)
+    server = WebsocketPolicyServer(
+        policy=wrapper,
+        server_config=server_config,
+        host=host,
+        port=port,
+    )
+    server.serve_forever()
+
+
+if __name__ == "__main__":
+    tyro.cli(main)

groot/control/tensorrt_utils.py

@@ -0,0 +1,852 @@
+import torch
+import os
+import subprocess
+import tensorrt as trt
+import sys
+import atexit
+import ctypes
+import modelopt.torch.quantization as mtq
+from typing import Dict, List, Tuple
+import shutil
+
+import numpy as np
+import torch
+
+
+FP8_DEFAULT_CONFIG = {
+    "quant_cfg": {
+        "*weight_quantizer": {"num_bits": (4, 3), "axis": None},
+        "*input_quantizer": {"num_bits": (4, 3), "axis": None},
+        "*output_quantizer": {"enable": False},
+        "*[qkv]_bmm_quantizer": {"num_bits": (4, 3), "axis": None},
+        "*softmax_quantizer": {
+            "num_bits": (4, 3),
+            "axis": None,
+        },
+        "default": {"enable": False},
+    },
+    "algorithm": "max",
+}
+
+NVFP4_DEFAULT_CONFIG = {
+    "quant_cfg": {
+        "*weight_quantizer": {
+            "num_bits": (2, 1),
+            "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+            "axis": None,
+            "enable": True,
+        },
+        "*input_quantizer": {
+            "num_bits": (2, 1),
+            "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+            "axis": None,
+            "enable": True,
+        },
+        "*output_quantizer": {"enable": False},
+        "*[qkv]_bmm_quantizer": {"num_bits": (4, 3), "axis": None},
+        "*softmax_quantizer": {
+            "num_bits": (4, 3),
+            "axis": None,
+        },
+        "default": {"enable": False},
+    },
+    "algorithm": "max",
+}
+
+    
+
+def wan_quantize(
+    policy,
+    quantization_config,
+    model_type,
+    forward_loop,
+):
+    """Quantize the VLA model using ModelOpt - simplified to use calc_mse_for_single_trajectory."""
+
+    # Configure quantization - disable problematic layers
+    if "quant_cfg" in quantization_config:
+        quantization_config["quant_cfg"]["*patch_embedding*"] = {"enable": False}
+    #    if model_type == "14B" or model_type == "ar_14B":
+    #        # Workaround: until we understand the issue https://nvbugspro.nvidia.com/bug/5612316
+    #        quantization_config["quant_cfg"]["*.self_attn.o.*"] = {"enable": False}
+    #        quantization_config["quant_cfg"]["*.cross_attn.o.*"] = {"enable": False}
+
+    policy.trained_model.action_head.model = mtq.quantize(
+        policy.trained_model.action_head.model, quantization_config, forward_loop=forward_loop
+    )
+    mtq.print_quant_summary(policy.trained_model.action_head.model)
+
+    return
+
+
+def wan_trt_quantize_and_load_engine(
+    policy,
+    cfg,
+    onnx_path,
+    engine_path,
+    model_type,
+    forward_loop,
+):
+    if (
+        os.path.exists(os.path.dirname(engine_path))
+        and cfg.inference_mode == "trt_build"
+    ):
+        shutil.rmtree(os.path.dirname(engine_path))
+
+    quantization_config = None
+    if cfg.quantize_dtype == "fp8":
+        quantization_config = FP8_DEFAULT_CONFIG.copy()
+    elif cfg.quantize_dtype == "nvfp4":
+        quantization_config = NVFP4_DEFAULT_CONFIG.copy()
+    else:
+        print(f"Quantization type {cfg.quantize_dtype} not supported. Skipping quantization.")
+
+    if quantization_config is not None and cfg.inference_mode == "trt_build":
+        #policy.trained_model.action_head.model.to(torch.float16)
+        wan_quantize(
+            policy,
+            quantization_config,
+            model_type=model_type,
+            forward_loop=forward_loop,
+        )
+
+    if  cfg.inference_mode == "trt_build":
+        policy.trained_model.action_head.model.to(torch.float16)
+
+        print("Export model:", policy.trained_model.action_head.model)
+
+        test_inputs = create_wan_test_inputs(policy, device="cuda", model_type=model_type)
+        min_shape = None
+        max_shape = None
+        opt_shape = None
+
+        if model_type == "ar_14B":
+
+            policy.trained_model.action_head.model.forward = policy.trained_model.action_head.model._forward_inference_trt
+            dynamic_axes = {
+                "kv_cache_packed": {3: "kv_cache_len"},
+            }
+            min_shape = "kv_cache_packed:40x2x1x880x40x128"
+            max_shape = "kv_cache_packed:40x2x1x8800x40x128"
+            opt_shape = "kv_cache_packed:40x2x1x7920x40x128"
+        elif model_type == "ar_14B_droid":
+            policy.trained_model.action_head.model.forward = policy.trained_model.action_head.model._forward_inference_trt
+            dynamic_axes = {
+                "kv_cache_packed": {3: "kv_cache_len"},
+            }
+            min_shape = "kv_cache_packed:40x2x1x880x40x128"
+            max_shape = "kv_cache_packed:40x2x1x8800x40x128"
+            opt_shape = "kv_cache_packed:40x2x1x7920x40x128"
+        elif model_type == "ar_5B_n6":
+            policy.trained_model.action_head.model.forward = policy.trained_model.action_head.model._forward_inference_trt
+            dynamic_axes = {
+                "kv_cache_packed": {3: "kv_cache_len"},
+            }
+            min_shape = "kv_cache_packed:30x2x1x220x24x128"
+            max_shape = "kv_cache_packed:30x2x1x3080x24x128"
+            opt_shape = "kv_cache_packed:30x2x1x2860x24x128"
+        else:
+            dynamic_axes = None
+
+        if cfg.quantize_dtype == "nvfp4":
+            export_to_onnx_fp4(policy.trained_model.action_head.model, test_inputs, onnx_path, dynamic_axes=dynamic_axes)
+        else:
+            export_to_onnx(
+                policy.trained_model.action_head.model,
+                test_inputs,
+                onnx_path,
+                model_type=model_type,
+                quantization_mode=cfg.quantize_dtype,
+                dynamic_axes=dynamic_axes,
+            )
+
+        build_tensorrt_engine(onnx_path, engine_path, min_shape, max_shape, opt_shape)
+
+    trt_wan_model = load_tensorrt_engine(engine_path, model_type=model_type)
+    policy.trained_model.action_head.model = trt_wan_model
+
+def export_to_onnx_fp4(model, test_inputs, onnx_save_path, dynamic_axes=None):
+    from modelopt.torch._deploy.utils.torch_onnx import OnnxBytes
+    from modelopt.torch._deploy.utils.torch_onnx import get_onnx_bytes_and_metadata
+
+    print("exporting to onnx fp4")
+    try:
+        onnx_bytes, _ = get_onnx_bytes_and_metadata(model=model, dummy_input=test_inputs, dynamic_axes=dynamic_axes)
+        onnx_model = OnnxBytes.from_bytes(onnx_bytes)
+    except Exception as e:
+        print(f"Error exporting model to ONNX: {e}")
+        return
+    save_dir = os.path.dirname(os.path.abspath(onnx_save_path))
+    os.makedirs(save_dir, exist_ok=True)
+    for filename, file_bytes in onnx_model.onnx_model.items():
+        file_path = os.path.join(save_dir, filename)
+        with open(file_path, "wb") as f:
+            f.write(file_bytes)
+        print(f"exported onnx to {file_path}")
+
+
+def export_to_onnx(
+    pytorch_model,
+    test_inputs,
+    onnx_path="tensorrt/wan_model.onnx",
+    model_type="5B",
+    quantization_mode="fp8",
+    dynamic_axes=None,
+):
+    #
+    if model_type == "5B":
+        return export_to_onnx_5B(pytorch_model, test_inputs, onnx_path, dynamic_axes)
+    elif model_type == "14B":
+        return export_to_onnx_14B(pytorch_model, test_inputs, onnx_path, dynamic_axes)
+    elif model_type == "ar_14B" or model_type == "ar_14B_droid":
+        return export_to_onnx_ar_14B(pytorch_model, test_inputs, onnx_path, dynamic_axes)
+    else:
+        raise ValueError(f"Model type {model_type} not supported")
+
+
+def export_to_onnx_ar_14B(pytorch_model, test_inputs, onnx_path="tensorrt/wan_model.onnx", dynamic_axes=None):
+    """Export PyTorch model to ONNX"""
+    print("Exporting AR 14B model to ONNX...", onnx_path)
+
+    # Create directory if it doesn't exist
+    os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
+    pytorch_model.eval()
+    pytorch_model.to(torch.float16)
+
+    input_names = [
+        "x",
+        "timestep",
+        "context",
+        "kv_cache_packed",
+        "y",
+        "clip_feature",
+        "action",
+        "timestep_action",
+        "state",
+    ]
+    output_names = ["video_noise_pred", "action_noise_pred"]
+
+    try:
+        with torch.no_grad():
+            torch.onnx.export(
+                pytorch_model,
+                test_inputs,
+                onnx_path,
+                export_params=True,
+                opset_version=20,
+                do_constant_folding=True,
+                input_names=input_names,
+                output_names=output_names,
+                dynamic_axes=dynamic_axes,
+            )
+        print(f"  ONNX model exported to: {onnx_path}")
+        return onnx_path
+
+    except Exception as e:
+        import traceback
+        print(f"  ERROR: ONNX export failed. Exception type: {type(e)}")
+        print("Traceback:")
+        traceback.print_exc()
+        return None
+
+
+
+def export_to_onnx_5B(pytorch_model, test_inputs, onnx_path="tensorrt/wan_model.onnx"):
+    """Export PyTorch model to ONNX"""
+    print("Exporting model to ONNX...")
+
+    # Create directory if it doesn't exist
+    os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
+    pytorch_model.eval()
+    pytorch_model.to(torch.float16)
+
+    x, action, timestep, context, state, embodiment_id = test_inputs
+
+    # Define input names for better ONNX graph
+    input_names = ["x", "action", "timestep", "context", "state", "embodiment_id"]
+    output_names = ["video_noise_pred", "action_noise_pred"]
+
+    try:
+        with torch.no_grad():
+            torch.onnx.export(
+                pytorch_model,
+                (x, action, timestep, context, state, embodiment_id),
+                onnx_path,
+                export_params=True,
+                opset_version=20,
+                do_constant_folding=True,
+                input_names=input_names,
+                output_names=output_names,
+            )
+        print(f"  ONNX model exported to: {onnx_path}")
+        return onnx_path
+
+    except Exception as e:
+        import traceback
+        print(f"  ERROR: ONNX export failed. Exception type: {type(e)}")
+        print("Traceback:")
+        traceback.print_exc()
+        return None
+
+
+def export_to_onnx_14B(pytorch_model, test_inputs, onnx_path="tensorrt/wan_model.onnx"):
+    """Export PyTorch model to ONNX"""
+    print("Exporting model to ONNX...")
+
+    # Create directory if it doesn't exist
+    os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
+    pytorch_model.eval()
+    pytorch_model.to(torch.float16)
+
+    x, action, timestep, context, state, embodiment_id, clip_feature, y = test_inputs
+
+    # Define input names for better ONNX graph
+    input_names = [
+        "x",
+        "action",
+        "timestep",
+        "context",
+        "state",
+        "embodiment_id",
+        "clip_feature",
+        "y",
+    ]
+    output_names = ["video_noise_pred", "action_noise_pred"]
+
+    try:
+        with torch.no_grad():
+            torch.onnx.export(
+                pytorch_model,
+                (x, action, timestep, context, state, embodiment_id, clip_feature, y),
+                onnx_path,
+                export_params=True,
+                opset_version=20,
+                do_constant_folding=True,
+                input_names=input_names,
+                output_names=output_names,
+            )
+        print(f"  ONNX model exported to: {onnx_path}")
+        return onnx_path
+
+    except Exception as e:
+        import traceback
+        print(f"  ERROR: ONNX export failed. Exception type: {type(e)}")
+        print("Traceback:")
+        traceback.print_exc()
+        return None
+
+
+def build_tensorrt_engine(onnx_path, engine_path="tensorrt/wan_model.trt", min_shape=None, max_shape=None, opt_shape=None):
+    """Build TensorRT engine from ONNX using trtexec"""
+    print("Building TensorRT engine with trtexec...")
+
+    if not os.path.exists(onnx_path):
+        print(f"  ERROR: ONNX file not found: {onnx_path}")
+        return None
+
+    # Create directory if it doesn't exist
+    os.makedirs(os.path.dirname(engine_path), exist_ok=True)
+
+    # Build engine using trtexec (much faster than torch_tensorrt)
+    trtexec_bin = shutil.which("trtexec") or "/opt/tensorrt/bin/trtexec"
+    cmd = [
+        trtexec_bin,
+        f"--onnx={onnx_path}",
+        f"--saveEngine={engine_path}",
+        "--fp8",
+        "--fp16",
+        "--bf16",
+        "--separateProfileRun",
+        "--profilingVerbosity=detailed",
+        "--memPoolSize=workspace:65536",
+        "--dumpProfile",
+        "--dumpLayerInfo",
+        "--useCudaGraph",
+        "--verbose",
+    ]
+
+    if min_shape is not None:
+        cmd.append(f"--minShapes={min_shape}")
+    if max_shape is not None:
+        cmd.append(f"--maxShapes={max_shape}")
+    if opt_shape is not None:
+        cmd.append(f"--optShapes={opt_shape}")
+
+    # Create log file for trtexec output
+    log_file = engine_path.replace(".trt", "_build.log")
+
+    try:
+        print(f"  Running: {' '.join(cmd)}")
+        print(f"  Logging output to: {log_file}")
+
+        with open(log_file, "w") as f:
+            result = subprocess.run(cmd, stdout=f, stderr=subprocess.STDOUT, text=True, timeout=600)
+
+        if result.returncode == 0:
+            print(f"  TensorRT engine built successfully: {engine_path}")
+            print(f"  Build log saved to: {log_file}")
+            return engine_path
+        else:
+            print(f"  ERROR: trtexec failed with return code {result.returncode}")
+            print(f"  Check build log for details: {log_file}")
+            # Print last few lines of log file for immediate feedback
+            try:
+                with open(log_file, "r") as f:
+                    lines = f.readlines()
+                    if lines:
+                        print("  Last few lines from build log:")
+                        for line in lines[-10:]:  # Show last 10 lines
+                            print(f"    {line.rstrip()}")
+            except:
+                pass
+            return None
+
+    except subprocess.TimeoutExpired:
+        print("  ERROR: trtexec timed out after 5 minutes")
+        print(f"  Partial build log saved to: {log_file}")
+        return None
+    except Exception as e:
+        print(f"  ERROR: Failed to run trtexec: {e}")
+        return None
+
+
+def torch_type(trt_type):
+    mapping = {
+        trt.float32: torch.float32,  # Added missing FLOAT mapping
+        trt.float16: torch.float16,
+        trt.bfloat16: torch.bfloat16,
+        trt.int8: torch.int8,
+        trt.int32: torch.int32,
+        trt.bool: torch.bool,
+        trt.uint8: torch.uint8,
+        trt.int64: torch.int64,
+    }
+    if trt_type in mapping:
+        return mapping[trt_type]
+
+    raise TypeError(
+        f"Could not resolve TensorRT datatype to an equivalent torch datatype. {trt_type}"
+    )
+
+
+class Engine(object):
+    def __init__(self, file, plugins=[]):
+        super().__init__()
+
+        self.logger = trt.Logger(trt.Logger.ERROR)
+        trt.init_libnvinfer_plugins(self.logger, "")
+
+        self.plugins = [ctypes.CDLL(plugin, ctypes.RTLD_GLOBAL) for plugin in plugins]
+        self.file = file
+        self.load(file)
+
+        def destroy(self):
+            del self.execution_context
+            del self.handle
+
+        atexit.register(destroy, self)
+        self.print()
+
+    def print(self):
+
+        print("============= TRT Engine Detail =============")
+        print(f"Engine file: {self.file}")
+        print(f"Inputs: {len(self.in_meta)}")
+        for ib, item in enumerate(self.in_meta):
+            tensor_name, shape, dtype = item[:3]
+            print(f"   {ib}. {tensor_name}: {'x'.join(map(str, shape))} [{dtype}]")
+
+        print(f"Outputs: {len(self.out_meta)}")
+        for ib, item in enumerate(self.out_meta):
+            tensor_name, shape, dtype = item[:3]
+            print(f"   {ib}. {tensor_name}: {'x'.join(map(str, shape))} [{dtype}]")
+        print("=============================================")
+
+    def load(self, file):
+        runtime = trt.Runtime(self.logger)
+
+        with open(file, "rb") as f:
+            self.handle = runtime.deserialize_cuda_engine(f.read())
+            assert (
+                self.handle is not None
+            ), f"Failed to deserialize the cuda engine from file: {file}"
+
+        self.execution_context = self.handle.create_execution_context()
+        self.meta, self.in_meta, self.out_meta = [], [], []
+        for tensor_name in self.handle:
+            shape = self.handle.get_tensor_shape(tensor_name)
+            print(f"Tensor name: {tensor_name}, shape: {shape}")
+            dtype = torch_type(self.handle.get_tensor_dtype(tensor_name))
+            if self.handle.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT:
+                self.in_meta.append([tensor_name, shape, dtype])
+            else:
+                self.out_meta.append([tensor_name, shape, dtype])
+
+    def __call__(self, *args, **inputs):
+        return self.forward(*args, **inputs)
+
+    def set_runtime_tensor_shape(self, name, shape):
+        self.execution_context.set_input_shape(name, shape)
+
+    def forward(self, *args, **kwargs):
+        return_list = kwargs.pop("return_list", False)
+        reference_tensors = []
+        stream = torch.cuda.current_stream()
+        for iarg, x in enumerate(args):
+            name, shape, dtype = self.in_meta[iarg]
+            runtime_shape = self.execution_context.get_tensor_shape(name)
+            assert isinstance(x, torch.Tensor), f"Unsupported tensor type: {type(x)}"
+            assert runtime_shape == x.shape, f"Invalid input shape: {runtime_shape} != {x.shape}"
+            assert (
+                dtype == x.dtype
+            ), f"Invalid tensor dtype, excepted dtype is {dtype}, but got {x.dtype}"
+            assert x.is_cuda, f"Invalid tensor device, excepted device is cuda, but got {x.device}"
+            x = x.cuda().contiguous()
+            self.execution_context.set_tensor_address(name, x.data_ptr())
+            reference_tensors.append(x)
+
+        for name, shape, dtype in self.in_meta:
+            if name not in kwargs:
+                continue
+
+            runtime_shape = self.execution_context.get_tensor_shape(name)
+            x = kwargs[name]
+            assert isinstance(x, torch.Tensor), f"Unsupported tensor[{name}] type: {type(x)}"
+            assert (
+                runtime_shape == x.shape
+            ), f"Invalid input[{name}] shape: {x.shape}, but the expected shape is: {runtime_shape}"
+            assert (
+                dtype == x.dtype
+            ), f"Invalid tensor[{name}] dtype, expected dtype is {dtype}, but got {x.dtype}"
+            assert (
+                x.is_cuda
+            ), f"Invalid tensor[{name}] device, expected device is cuda, but got {x.device}"
+            x = x.cuda().contiguous()
+            self.execution_context.set_tensor_address(name, x.data_ptr())
+            reference_tensors.append(x)
+
+        for item in self.out_meta:
+            name = item[0]
+            runtime_shape = self.execution_context.get_tensor_shape(name)
+            output_tensor = torch.zeros(
+                *runtime_shape, dtype=item[2], device=reference_tensors[0].device
+            )
+            self.execution_context.set_tensor_address(name, output_tensor.data_ptr())
+            reference_tensors.append(output_tensor)
+
+        self.execution_context.execute_async_v3(stream.cuda_stream)
+        stream.synchronize()
+        assert len(reference_tensors) == len(self.in_meta) + len(
+            self.out_meta
+        ), f"Invalid input tensors. The expected I/O tensors are {len(self.in_meta) + len(self.out_meta)}, but got {len(reference_tensors)}"
+
+        if return_list:
+            return [
+                reference_tensors[len(self.in_meta) + i] for i, item in enumerate(self.out_meta)
+            ]
+        else:
+            return {
+                item[0]: reference_tensors[len(self.in_meta) + i]
+                for i, item in enumerate(self.out_meta)
+            }
+
+
+class WanTrtModel5B(torch.nn.Module):
+    def __init__(self, eng_path: str):
+        super().__init__()
+        self.engine = Engine(eng_path)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        action: torch.Tensor,
+        timestep: torch.Tensor,
+        context: torch.Tensor,
+        state: torch.Tensor,
+        embodiment_id: torch.Tensor,
+    ):
+
+        self.engine.set_runtime_tensor_shape("x", x.shape)
+        self.engine.set_runtime_tensor_shape("action", action.shape)
+        self.engine.set_runtime_tensor_shape("context", context.shape)
+        self.engine.set_runtime_tensor_shape("state", state.shape)
+
+        output = self.engine(
+            x=x.to(torch.float16),
+            action=action.to(torch.float16),
+            timestep=timestep.to(torch.float16),
+            context=context.to(torch.float16),
+            state=state.to(torch.float16),
+            embodiment_id=embodiment_id.to(torch.int32),
+        )
+        if "out.0" in output: # for nvfp4 model export through modelopt
+            return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous()
+        else:
+            return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous()
+
+
+class WanTrtModel14B(torch.nn.Module):
+    def __init__(self, eng_path: str):
+        super().__init__()
+        self.engine = Engine(eng_path)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        action: torch.Tensor,
+        timestep: torch.Tensor,
+        context: torch.Tensor,
+        state: torch.Tensor,
+        embodiment_id: torch.Tensor,
+        clip_feature: torch.Tensor,
+        y: torch.Tensor,
+    ):
+
+        self.engine.set_runtime_tensor_shape("x", x.shape)
+        self.engine.set_runtime_tensor_shape("action", action.shape)
+        self.engine.set_runtime_tensor_shape("context", context.shape)
+        self.engine.set_runtime_tensor_shape("state", state.shape)
+        self.engine.set_runtime_tensor_shape("clip_feature", clip_feature.shape)
+        self.engine.set_runtime_tensor_shape("y", y.shape)
+
+        output = self.engine(
+            x=x.to(torch.float16),
+            action=action.to(torch.float16),
+            timestep=timestep.to(torch.float16),
+            context=context.to(torch.float16),
+            state=state.to(torch.float16),
+            embodiment_id=embodiment_id.to(torch.int32),
+            clip_feature=clip_feature.to(torch.float16),
+            y=y.to(torch.float16),
+        )
+        if "out.0" in output: # for nvfp4 model export through modelopt
+            return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous()
+        else:
+            return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous()
+
+
+class WanTrtModelAr5B(torch.nn.Module):
+    """TRT wrapper for ar_5B_n6 model type - uses kv_cache but no clip_feature."""
+    def __init__(self, eng_path: str):
+        super().__init__()
+        self.engine = Engine(eng_path)
+
+    def forward(
+        self,
+        x,
+        timestep,
+        context,
+        kv_cache: list[torch.Tensor],
+        y=None,
+        action=None,
+        timestep_action=None,
+        state=None,
+    ):
+
+        kv_cache_packed = torch.stack(kv_cache, dim=0)
+
+        self.engine.set_runtime_tensor_shape("x", x.shape)
+        self.engine.set_runtime_tensor_shape("timestep", timestep.shape)
+        self.engine.set_runtime_tensor_shape("context", context.shape)
+        self.engine.set_runtime_tensor_shape("kv_cache_packed", kv_cache_packed.shape)
+        # self.engine.set_runtime_tensor_shape("y", y.shape)
+        self.engine.set_runtime_tensor_shape("action", action.shape)
+        self.engine.set_runtime_tensor_shape("timestep_action", timestep_action.shape)
+        self.engine.set_runtime_tensor_shape("state", state.shape)
+
+
+        output = self.engine(
+            x.to(torch.float16),
+            timestep.to(torch.float16),
+            context.to(torch.float16),
+            kv_cache_packed.to(torch.float16),
+            # y.to(torch.float16),
+            action.to(torch.float16),
+            timestep_action.to(torch.float16),
+            state.to(torch.float16),
+        )
+
+        if "out.0" in output: # for nvfp4 model export through modelopt
+            return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous()
+        else:
+            return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous()
+
+
+class WanTrtModelAr14B(torch.nn.Module):
+    def __init__(self, eng_path: str):
+        super().__init__()
+        self.engine = Engine(eng_path)
+
+    def forward(
+        self,
+        x,
+        timestep,
+        context,
+        kv_cache: list[torch.Tensor],
+        y=None,
+        clip_feature=None,
+        action=None,
+        timestep_action=None,
+        state=None,
+    ):
+
+        kv_cache_packed = torch.stack(kv_cache, dim=0)
+
+        self.engine.set_runtime_tensor_shape("x", x.shape)
+        self.engine.set_runtime_tensor_shape("timestep", timestep.shape)
+        self.engine.set_runtime_tensor_shape("context", context.shape)
+        self.engine.set_runtime_tensor_shape("kv_cache_packed", kv_cache_packed.shape)
+        self.engine.set_runtime_tensor_shape("y", y.shape)
+        self.engine.set_runtime_tensor_shape("clip_feature", clip_feature.shape)
+        self.engine.set_runtime_tensor_shape("action", action.shape)
+        self.engine.set_runtime_tensor_shape("timestep_action", timestep_action.shape)
+        self.engine.set_runtime_tensor_shape("state", state.shape)
+
+
+        output = self.engine(
+            x.to(torch.float16),
+            timestep.to(torch.float16),
+            context.to(torch.float16),
+            kv_cache_packed.to(torch.float16),
+            y.to(torch.float16),
+            clip_feature.to(torch.float16),
+            action.to(torch.float16),
+            timestep_action.to(torch.float16),
+            state.to(torch.float16),
+        )
+
+        if "out.0" in output: # for nvfp4 model export through modelopt
+            return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous()
+        else:
+            return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous()
+
+def load_tensorrt_engine(engine_path="tensorrt/wan_model.trt", model_type="5B"):
+    """Load TensorRT engine"""
+    if model_type == "5B":
+        trt_inference = WanTrtModel5B(engine_path)
+    elif model_type == "ar_5B_n6" or model_type == "ar_5B":
+        trt_inference = WanTrtModelAr5B(engine_path)
+    elif model_type == "14B":
+        trt_inference = WanTrtModel14B(engine_path)
+    elif model_type == "ar_14B" or model_type == "ar_14B_droid":
+        trt_inference = WanTrtModelAr14B(engine_path)
+    else:
+        raise ValueError(f"Model type {model_type} not supported")
+    return trt_inference
+
+
+def create_wan_test_inputs(policy, device="cuda", model_type="5B"):
+    # Get dtype from model parameters
+    dtype = torch.float16
+
+    # Use hardcoded dimensions from the original working version of the script
+    if model_type == "5B":
+        x = torch.randn(1, 48, 13, 22, 40, dtype=dtype, device=device)
+        action = torch.randn(1, 48, 32, dtype=dtype, device=device)
+        timestep = torch.randn(1, dtype=dtype, device=device)
+        context = torch.randn(1, 512, 4096, dtype=dtype, device=device)
+        state = torch.randn(1, 1, 64, dtype=dtype, device=device)
+        embodiment_id = torch.zeros(1, dtype=torch.int32, device=device)
+        timestep_action = torch.randn(1, 48, dtype=dtype, device=device)
+        seq_len = torch.tensor(440, dtype=torch.int32, device=device)
+        return x, action, timestep, context, state, embodiment_id, timestep_action, seq_len
+    elif model_type == "ar_5B_n6":
+        # ar_5B_n6 uses _forward_inference_trt which requires kv_cache_packed
+        # Shape from dynamic_axes: kv_cache_packed:30x2x1x220x24x128
+        # Note: 5B model doesn't use clip_feature (unlike 14B), but still needs y
+        x = torch.randn(1, 48, 2, 22, 40, dtype=dtype, device=device)
+        timestep = torch.randn(1, 2, dtype=dtype, device=device)
+        context = torch.randn(1, 512, 4096, dtype=dtype, device=device)
+        # y = torch.randn(1, 52, 2, 22, 40, dtype=dtype, device=device)  # y is required by _forward_inference_trt
+        action = torch.randn(1, 48, 32, dtype=dtype, device=device)
+        timestep_action = torch.randn(1, 48, dtype=dtype, device=device)
+        state = torch.randn(1, 1, 64, dtype=dtype, device=device)
+    
+        num_heads = 24
+        head_dim = 128
+        num_layers = 30
+        B = 1
+    
+        kv_cache = []
+        for _ in range(num_layers):
+            kv_cache.append(
+                torch.zeros([2, B, 13*220, num_heads, head_dim], dtype=dtype, device=device)
+            )
+    
+        kv_cache_packed = torch.stack(kv_cache, dim=0)
+        # Return order matches _forward_inference_trt signature: x, timestep, context, kv_cache_packed, y, action, timestep_action, state
+        return (x, timestep, context, kv_cache_packed, action, timestep_action, state)
+    elif model_type == "14B":
+        x = torch.randn(1, 16, 13, 44, 80, dtype=dtype, device=device)
+        action = torch.randn(1, 48, 32, dtype=dtype, device=device)
+        timestep = torch.randn(1, dtype=dtype, device=device)
+        context = torch.randn(1, 512, 4096, dtype=dtype, device=device)
+        state = torch.randn(1, 1, 64, dtype=dtype, device=device)
+        embodiment_id = torch.zeros(1, dtype=torch.int32, device=device)
+        clip_feature = torch.randn(1, 257, 1280, dtype=dtype, device=device)
+        y = torch.randn(1, 20, 13, 44, 80, dtype=dtype, device=device)
+        return x, action, timestep, context, state, embodiment_id, clip_feature, y
+    elif model_type == "ar_14B":
+        clip_feature = torch.randn(1, 257, 1280, dtype=dtype, device=device)
+        y = torch.randn(1, 20, 2, 44, 80, dtype=dtype, device=device)
+        timestep_action = torch.randn(1, 48, dtype=dtype, device=device)
+        x = torch.randn(1, 16, 2, 44, 80, dtype=dtype, device=device)
+        timestep = torch.randn(1, 2, dtype=dtype, device=device)
+        context = torch.randn(1, 512, 4096, dtype=dtype, device=device)
+        seq_len = torch.tensor(1760, dtype=torch.int32, device=device)
+        action = torch.randn(1, 48, 32, dtype=dtype, device=device)
+        state = torch.randn(1, 1, 64, dtype=dtype, device=device)
+        embodiment_id = torch.zeros(1, dtype=torch.int32, device=device)
+    
+        num_heads = 40 
+        head_dim = 5120 // num_heads
+        num_layers = 40 
+        B = 1
+    
+        kv_cache = []
+        for _ in range(num_layers):
+            kv_cache.append(
+                torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device)
+            )
+    
+        crossattn_k_cache = []
+        for _ in range(num_layers):
+            crossattn_k_cache.append(
+                torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device)
+            )
+        kv_cache_packed = torch.stack(kv_cache, dim=0)
+        crossattn_packed = torch.stack(crossattn_k_cache, dim=0)
+        return (x, timestep, context, kv_cache_packed, y, clip_feature, action, timestep_action, state)
+    elif model_type == "ar_14B_droid":
+        clip_feature = torch.randn(1, 257, 1280, dtype=dtype, device=device)
+        y = torch.randn(1, 20, 2, 44, 80, dtype=dtype, device=device)
+        timestep_action = torch.randn(1, 24, dtype=dtype, device=device)
+        x = torch.randn(1, 16, 2, 44, 80, dtype=dtype, device=device)
+        timestep = torch.randn(1, 2, dtype=dtype, device=device)
+        context = torch.randn(1, 512, 4096, dtype=dtype, device=device)
+        seq_len = torch.tensor(1760, dtype=torch.int32, device=device)
+        action = torch.randn(1, 24, 32, dtype=dtype, device=device)
+        state = torch.randn(1, 1, 64, dtype=dtype, device=device)
+        embodiment_id = torch.zeros(1, dtype=torch.int32, device=device)
+    
+        num_heads = 40 
+        head_dim = 5120 // num_heads
+        num_layers = 40 
+        B = 1
+    
+        kv_cache = []
+        for _ in range(num_layers):
+            kv_cache.append(
+                torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device)
+            )
+    
+        crossattn_k_cache = []
+        for _ in range(num_layers):
+            crossattn_k_cache.append(
+                torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device)
+            )
+        kv_cache_packed = torch.stack(kv_cache, dim=0)
+        crossattn_packed = torch.stack(crossattn_k_cache, dim=0)
+        return (x, timestep, context, kv_cache_packed, y, clip_feature, action, timestep_action, state)
+
+

groot/vla/common/__init__.py

@@ -0,0 +1 @@
+

groot/vla/common/utils/__init__.py

@@ -0,0 +1,3 @@
+from .data_structure import *  # noqa: F403
+from .io import *  # noqa: F403
+from .misc import *  # noqa: F403

groot/vla/common/utils/data_structure/__init__.py

@@ -0,0 +1,2 @@
+from .shape_utils import *  # noqa: F403
+from .tree_utils import *  # noqa: F403

groot/vla/common/utils/data_structure/shape_utils.py

@@ -0,0 +1,283 @@
+"""
+Shape inference methods
+"""
+
+from functools import partial
+import math
+from typing import List, Tuple, Union
+import warnings
+
+import numpy as np
+import torch
+
+# fmt: off
+__all__ = [
+    "shape_convnd",
+    "shape_conv1d", "shape_conv2d", "shape_conv3d",
+    "shape_transpose_convnd",
+    "shape_transpose_conv1d", "shape_transpose_conv2d", "shape_transpose_conv3d",
+    "shape_poolnd",
+    "shape_maxpool1d", "shape_maxpool2d", "shape_maxpool3d",
+    "shape_avgpool1d", "shape_avgpool2d", "shape_avgpool3d",
+    "shape_slice",
+    "check_shape"
+]
+# fmt: on
+
+
+def _get_shape(x):
+    "single object"
+    if isinstance(x, np.ndarray):
+        return tuple(x.shape)
+    else:
+        return tuple(x.size())
+
+
+def _expands(dim, *xs):
+    "repeat vars like kernel and stride to match dim"
+
+    def _expand(x):
+        if isinstance(x, int):
+            return (x,) * dim
+        else:
+            assert len(x) == dim
+            return x
+
+    return map(lambda x: _expand(x), xs)
+
+
+_HELPER_TENSOR = torch.zeros((1,))
+
+
+def shape_slice(input_shape, slice):
+    """
+    Credit to Adam Paszke for the trick. Shape inference without instantiating
+    an actual tensor.
+    The key is that `.expand()` does not actually allocate memory
+    Still needs to allocate a one-element HELPER_TENSOR.
+    """
+    shape = _HELPER_TENSOR.expand(*input_shape)[slice]
+    if hasattr(shape, "size"):
+        return tuple(shape.size())
+    return (1,)
+
+
+class ShapeSlice:
+    """
+    shape_slice inference with easy []-operator
+    """
+
+    def __init__(self, input_shape):
+        self.input_shape = input_shape
+
+    def __getitem__(self, slice):
+        return shape_slice(self.input_shape, slice)
+
+
+def check_shape(
+    value: Union[Tuple, List, torch.Tensor, np.ndarray],
+    expected: Union[Tuple, List, torch.Tensor, np.ndarray],
+    err_msg="",
+    mode="raise",
+):
+    """
+    Args:
+        value: np array or torch Tensor
+        expected:
+          - list[int], tuple[int]: if any value is None, will match any dim
+          - np array or torch Tensor: must have the same dimensions
+        mode:
+          - "raise": raise ValueError, shape mismatch
+          - "return": returns True if shape matches, otherwise False
+          - "warning": warnings.warn
+    """
+    assert mode in ["raise", "return", "warning"]
+    if torch.is_tensor(value):
+        actual_shape = value.size()
+    elif hasattr(value, "shape"):
+        actual_shape = value.shape
+    else:
+        assert isinstance(value, (list, tuple))
+        actual_shape = value
+        assert all(
+            isinstance(s, int) for s in actual_shape
+        ), f"actual shape: {actual_shape} is not a list of ints"
+
+    if torch.is_tensor(expected):
+        expected_shape = expected.size()
+    elif hasattr(expected, "shape"):
+        expected_shape = expected.shape
+    else:
+        assert isinstance(expected, (list, tuple))
+        expected_shape = expected
+
+    err_msg = f" for {err_msg}" if err_msg else ""
+
+    if len(actual_shape) != len(expected_shape):
+        err_msg = (
+            f"Dimension mismatch{err_msg}: actual shape {actual_shape} "
+            f"!= expected shape {expected_shape}."
+        )
+        if mode == "raise":
+            raise ValueError(err_msg)
+        elif mode == "warning":
+            warnings.warn(err_msg)
+        return False
+
+    for s_a, s_e in zip(actual_shape, expected_shape):
+        if s_e is not None and s_a != s_e:
+            err_msg = (
+                f"Shape mismatch{err_msg}: actual shape {actual_shape} "
+                f"!= expected shape {expected_shape}."
+            )
+            if mode == "raise":
+                raise ValueError(err_msg)
+            elif mode == "warning":
+                warnings.warn(err_msg)
+            return False
+    return True
+
+
+def shape_convnd(
+    dim,
+    input_shape,
+    out_channels,
+    kernel_size,
+    stride=1,
+    padding=0,
+    dilation=1,
+    has_batch=False,
+):
+    """
+    http://pytorch.org/docs/nn.html#conv1d
+    http://pytorch.org/docs/nn.html#conv2d
+    http://pytorch.org/docs/nn.html#conv3d
+
+    Args:
+        dim: supports 1D to 3D
+        input_shape:
+        - 1D: [channel, length]
+        - 2D: [channel, height, width]
+        - 3D: [channel, depth, height, width]
+        has_batch: whether the first dim is batch size or not
+    """
+    if has_batch:
+        assert (
+            len(input_shape) == dim + 2
+        ), "input shape with batch should be {}-dimensional".format(dim + 2)
+    else:
+        assert (
+            len(input_shape) == dim + 1
+        ), "input shape without batch should be {}-dimensional".format(dim + 1)
+    if stride is None:
+        # for pooling convention in PyTorch
+        stride = kernel_size
+    kernel_size, stride, padding, dilation = _expands(dim, kernel_size, stride, padding, dilation)
+    if has_batch:
+        batch = input_shape[0]
+        input_shape = input_shape[1:]
+    else:
+        batch = None
+    _, *img = input_shape
+    new_img_shape = [
+        math.floor(
+            (img[i] + 2 * padding[i] - dilation[i] * (kernel_size[i] - 1) - 1) // stride[i] + 1
+        )
+        for i in range(dim)
+    ]
+    return ((batch,) if has_batch else ()) + (out_channels, *new_img_shape)
+
+
+def shape_poolnd(
+    dim, input_shape, kernel_size, stride=None, padding=0, dilation=1, has_batch=False
+):
+    """
+    The only difference from infer_shape_convnd is that `stride` default is None
+    """
+    if has_batch:
+        out_channels = input_shape[1]
+    else:
+        out_channels = input_shape[0]
+    return shape_convnd(
+        dim,
+        input_shape,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        has_batch,
+    )
+
+
+def shape_transpose_convnd(
+    dim,
+    input_shape,
+    out_channels,
+    kernel_size,
+    stride=1,
+    padding=0,
+    output_padding=0,
+    dilation=1,
+    has_batch=False,
+):
+    """
+    http://pytorch.org/docs/nn.html#convtranspose1d
+    http://pytorch.org/docs/nn.html#convtranspose2d
+    http://pytorch.org/docs/nn.html#convtranspose3d
+
+    Args:
+        dim: supports 1D to 3D
+        input_shape:
+        - 1D: [channel, length]
+        - 2D: [channel, height, width]
+        - 3D: [channel, depth, height, width]
+        has_batch: whether the first dim is batch size or not
+    """
+    if has_batch:
+        assert (
+            len(input_shape) == dim + 2
+        ), "input shape with batch should be {}-dimensional".format(dim + 2)
+    else:
+        assert (
+            len(input_shape) == dim + 1
+        ), "input shape without batch should be {}-dimensional".format(dim + 1)
+    kernel_size, stride, padding, output_padding, dilation = _expands(
+        dim, kernel_size, stride, padding, output_padding, dilation
+    )
+    if has_batch:
+        batch = input_shape[0]
+        input_shape = input_shape[1:]
+    else:
+        batch = None
+    _, *img = input_shape
+    new_img_shape = [
+        (img[i] - 1) * stride[i] - 2 * padding[i] + kernel_size[i] + output_padding[i]
+        for i in range(dim)
+    ]
+    return ((batch,) if has_batch else ()) + (out_channels, *new_img_shape)
+
+
+shape_conv1d = partial(shape_convnd, 1)
+shape_conv2d = partial(shape_convnd, 2)
+shape_conv3d = partial(shape_convnd, 3)
+
+
+shape_transpose_conv1d = partial(shape_transpose_convnd, 1)
+shape_transpose_conv2d = partial(shape_transpose_convnd, 2)
+shape_transpose_conv3d = partial(shape_transpose_convnd, 3)
+
+
+shape_maxpool1d = partial(shape_poolnd, 1)
+shape_maxpool2d = partial(shape_poolnd, 2)
+shape_maxpool3d = partial(shape_poolnd, 3)
+
+
+"""
+http://pytorch.org/docs/nn.html#avgpool1d
+http://pytorch.org/docs/nn.html#avgpool2d
+http://pytorch.org/docs/nn.html#avgpool3d
+"""
+shape_avgpool1d = partial(shape_maxpool1d, dilation=1)
+shape_avgpool2d = partial(shape_maxpool2d, dilation=1)
+shape_avgpool3d = partial(shape_maxpool3d, dilation=1)

groot/vla/common/utils/data_structure/tree_utils.py

@@ -0,0 +1,219 @@
+"""
+Utils to handle nested data structures
+
+Install dm_tree first:
+https://tree.readthedocs.io/en/latest/api.html
+"""
+
+import collections
+from typing import Any, Iterable, List, Tuple, TypeVar
+
+import numpy as np
+
+try:
+    import tree
+
+except ImportError:
+    raise ImportError("Please install dm_tree first: `pip install dm_tree`")
+
+
+def is_sequence(obj):
+    """
+    Returns:
+      True if the sequence is a collections.Sequence and not a string.
+    """
+    return isinstance(obj, collections.abc.Sequence) and not isinstance(obj, str)
+
+
+def is_mapping(obj):
+    """
+    Returns:
+      True if the sequence is a collections.Mapping
+    """
+    return isinstance(obj, collections.abc.Mapping)
+
+
+def tree_value_at_path(obj, paths: Tuple):
+    try:
+        for p in paths:
+            obj = obj[p]
+        return obj
+    except Exception as e:
+        raise ValueError(f"{e}\n\n-- Incorrect nested path {paths} for object: {obj}.")
+
+
+def tree_assign_at_path(obj, paths: Tuple, value):
+    try:
+        for p in paths[:-1]:
+            obj = obj[p]
+        if len(paths) > 0:
+            obj[paths[-1]] = value
+    except Exception as e:
+        raise ValueError(f"{e}\n\n-- Incorrect nested path {paths} for object: {obj}.")
+
+
+def copy_non_leaf(obj):
+    """
+    Deepcopy the nested structure, but does NOT copy the leaf values like Tensors
+    """
+    return tree.map_structure(lambda x: x, obj)
+
+
+# =======================================================================
+# Copyright 2018 DeepMind Technologies Limited. All rights reserved.
+#
+# 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.
+
+# Tensor framework-agnostic utilities for manipulating nested structures.
+
+ElementType = TypeVar("ElementType")
+
+
+def fast_map_structure(func, *structure):
+    """Faster map_structure implementation which skips some error checking."""
+    flat_structure = (tree.flatten(s) for s in structure)
+    entries = zip(*flat_structure)
+    # Arbitrarily choose one of the structures of the original sequence (the last)
+    # to match the structure for the flattened sequence.
+    return tree.unflatten_as(structure[-1], [func(*x) for x in entries])
+
+
+def stack_sequence_fields(sequence: Iterable[ElementType]) -> ElementType:
+    """Stacks a list of identically nested objects.
+
+    This takes a sequence of identically nested objects and returns a single
+    nested object whose ith leaf is a stacked numpy array of the corresponding
+    ith leaf from each element of the sequence.
+
+    For example, if `sequence` is:
+
+    ```python
+    [{
+          'action': np.array([1.0]),
+          'observation': (np.array([0.0, 1.0, 2.0]),),
+          'reward': 1.0
+     }, {
+          'action': np.array([0.5]),
+          'observation': (np.array([1.0, 2.0, 3.0]),),
+          'reward': 0.0
+     }, {
+          'action': np.array([0.3]),1
+          'observation': (np.array([2.0, 3.0, 4.0]),),
+          'reward': 0.5
+     }]
+    ```
+
+    Then this function will return:
+
+    ```python
+    {
+        'action': np.array([....])         # array shape = [3 x 1]
+        'observation': (np.array([...]),)  # array shape = [3 x 3]
+        'reward': np.array([...])          # array shape = [3]
+    }
+    ```
+
+    Note that the 'observation' entry in the above example has two levels of
+    nesting, i.e it is a tuple of arrays.
+
+    Args:
+      sequence: a list of identically nested objects.
+
+    Returns:
+      A nested object with numpy.
+
+    Raises:
+      ValueError: If `sequence` is an empty sequence.
+    """
+    # Handle empty input sequences.
+    if not sequence:
+        raise ValueError("Input sequence must not be empty")
+
+    # Default to asarray when arrays don't have the same shape to be compatible
+    # with old behaviour.
+    try:
+        return fast_map_structure(lambda *values: np.stack(values), *sequence)
+    except ValueError:
+        return fast_map_structure(lambda *values: np.asarray(values), *sequence)
+
+
+def unstack_sequence_fields(struct: ElementType, batch_size: int) -> List[ElementType]:
+    """Converts a struct of batched arrays to a list of structs.
+
+    This is effectively the inverse of `stack_sequence_fields`.
+
+    Args:
+      struct: An (arbitrarily nested) structure of arrays.
+      batch_size: The length of the leading dimension of each array in the struct.
+        This is assumed to be static and known.
+
+    Returns:
+      A list of structs with the same structure as `struct`, where each leaf node
+       is an unbatched element of the original leaf node.
+    """
+
+    return [tree.map_structure(lambda s, i=i: s[i], struct) for i in range(batch_size)]
+
+
+def broadcast_structures(*args: Any) -> Any:
+    """Returns versions of the arguments that give them the same nested structure.
+
+    Any nested items in *args must have the same structure.
+
+    Any non-nested item will be replaced with a nested version that shares that
+    structure. The leaves will all be references to the same original non-nested
+    item.
+
+    If all *args are nested, or all *args are non-nested, this function will
+    return *args unchanged.
+
+    Example:
+    ```
+    a = ('a', 'b')
+    b = 'c'
+    tree_a, tree_b = broadcast_structure(a, b)
+    tree_a
+    > ('a', 'b')
+    tree_b
+    > ('c', 'c')
+    ```
+
+    Args:
+      *args: A Sequence of nested or non-nested items.
+
+    Returns:
+      `*args`, except with all items sharing the same nest structure.
+    """
+    if not args:
+        return
+
+    reference_tree = None
+    for arg in args:
+        if tree.is_nested(arg):
+            reference_tree = arg
+            break
+
+    if reference_tree is None:
+        reference_tree = args[0]
+
+    def mirror_structure(value, reference_tree):
+        if tree.is_nested(value):
+            # Use check_types=True so that the types of the trees we construct aren't
+            # dependent on our arbitrary choice of which nested arg to use as the
+            # reference_tree.
+            tree.assert_same_structure(value, reference_tree, check_types=True)
+            return value
+        else:
+            return tree.map_structure(lambda _: value, reference_tree)
+
+    return tuple(mirror_structure(arg, reference_tree) for arg in args)

groot/vla/common/utils/io/__init__.py

@@ -0,0 +1,6 @@
+from .config_utils import *  # noqa: F403
+from .file_utils import *  # noqa: F403
+from .hdf5_utils import *  # noqa: F403
+from .json_utils import *  # noqa: F403
+from .print_utils import *  # noqa: F403
+from .termcolor import *  # noqa: F403

groot/vla/common/utils/io/config_utils.py

@@ -0,0 +1,260 @@
+from copy import deepcopy
+import importlib.resources
+import os
+import sys
+
+import hydra
+from omegaconf import DictConfig, OmegaConf
+import tree
+
+from ..misc.functional_utils import call_once, is_mapping, is_sequence, meta_decorator
+from .print_utils import to_scientific_str
+
+_CLASS_REGISTRY = {}  # for instantiation
+
+
+def resource_file_path(pkg_name, fname) -> str:
+    with importlib.resources.path(pkg_name, fname) as p:
+        return str(p)
+
+
+def print_config(cfg: DictConfig):
+    print(cfg.pretty(resolve=True))
+
+
+def is_hydra_initialized():
+    return hydra.utils.HydraConfig.initialized()
+
+
+def hydra_config():
+    # https://github.com/facebookresearch/hydra/issues/377
+    # HydraConfig() is a singleton
+    if is_hydra_initialized():
+        return hydra.utils.HydraConfig().cfg.hydra
+    else:
+        return None
+
+
+def hydra_override_arg_list() -> list[str]:
+    """
+    Returns:
+        list ["lr=0.2", "batch=64", ...]
+    """
+    if is_hydra_initialized():
+        return hydra_config().overrides.task
+    else:
+        return []
+
+
+def hydra_override_name():
+    if is_hydra_initialized():
+        return hydra_config().job.override_dirname
+    else:
+        return ""
+
+
+def hydra_original_dir(*subpaths):
+    return os.path.join(hydra.utils.get_original_cwd(), *subpaths)
+
+
+@call_once(on_second_call="noop")
+def register_omegaconf_resolvers():
+    import numpy as np
+
+    OmegaConf.register_new_resolver("scientific", lambda v, i=0: to_scientific_str(v, i))
+    OmegaConf.register_new_resolver("_optional", lambda v: f"_{v}" if v else "")
+    OmegaConf.register_new_resolver("optional_", lambda v: f"{v}_" if v else "")
+    OmegaConf.register_new_resolver("_optional_", lambda v: f"_{v}_" if v else "")
+    OmegaConf.register_new_resolver("__optional", lambda v: f"__{v}" if v else "")
+    OmegaConf.register_new_resolver("optional__", lambda v: f"{v}__" if v else "")
+    OmegaConf.register_new_resolver("__optional__", lambda v: f"__{v}__" if v else "")
+    OmegaConf.register_new_resolver("iftrue", lambda cond, v_default: cond if cond else v_default)
+    OmegaConf.register_new_resolver("ifelse", lambda cond, v1, v2="": v1 if cond else v2)
+    OmegaConf.register_new_resolver(
+        "ifequal", lambda query, key, v1, v2: v1 if query == key else v2
+    )
+    OmegaConf.register_new_resolver("intbool", lambda cond: 1 if cond else 0)
+    OmegaConf.register_new_resolver("mult", lambda *x: np.prod(x).tolist())
+    OmegaConf.register_new_resolver("add", lambda *x: sum(x))
+    OmegaConf.register_new_resolver("div", lambda x, y: x / y)
+    OmegaConf.register_new_resolver("intdiv", lambda x, y: x // y)
+
+    # try each key until the key exists. Useful for multiple classes that have different
+    # names for the same key
+    def _try_key(cfg, *keys):
+        for k in keys:
+            if k in cfg:
+                return cfg[k]
+        raise KeyError(f"no key in {keys} is valid")
+
+    OmegaConf.register_new_resolver("trykey", _try_key)
+    # replace `resnet.gn.ws` -> `resnet_gn_ws`, because omegaconf doesn't support
+    # keys with dots. Useful for generating run name with dots
+    OmegaConf.register_new_resolver("underscore_to_dots", lambda s: s.replace("_", "."))
+
+    def _no_instantiate(cfg):
+        cfg = deepcopy(cfg)
+        cfg[_NO_INSTANTIATE] = True
+        return cfg
+
+    OmegaConf.register_new_resolver("no_instantiate", _no_instantiate)
+
+
+# ========================================================
+# ================== Instantiation tools  ================
+# ========================================================
+
+
+def register_callable(name, class_type):
+    if isinstance(class_type, str):
+        class_type, name = name, class_type
+    assert callable(class_type)
+    _CLASS_REGISTRY[name] = class_type
+
+
+@meta_decorator
+def register_class(cls, alias=None):
+    """
+    Decorator
+    """
+    assert callable(cls)
+    _CLASS_REGISTRY[cls.__name__] = cls
+    if alias:
+        assert is_sequence(alias)
+        for a in alias:
+            _CLASS_REGISTRY[str(a)] = cls
+    return cls
+
+
+def omegaconf_to_dict(cfg, resolve: bool = True, enum_to_str: bool = False):
+    """
+    Convert arbitrary nested omegaconf objects to primitive containers
+
+    WARNING: cannot use tree lib because it gets confused on DictConfig and ListConfig
+    """
+    kw = dict(resolve=resolve, enum_to_str=enum_to_str)
+    if OmegaConf.is_config(cfg):
+        return OmegaConf.to_container(cfg, **kw)
+    elif is_sequence(cfg):
+        return type(cfg)(omegaconf_to_dict(c, **kw) for c in cfg)
+    elif is_mapping(cfg):
+        return {k: omegaconf_to_dict(c, **kw) for k, c in cfg.items()}
+    else:
+        return cfg
+
+
+def omegaconf_save(cfg, *paths: str, resolve: bool = True):
+    """
+    Save omegaconf to yaml
+    """
+    from .file_utils import f_join
+
+    OmegaConf.save(cfg, f_join(*paths), resolve=resolve)
+
+
+def get_class(path):
+    """
+    First try to find the class in the registry first,
+    if it doesn't exist, use importlib to locate it
+    """
+    if path in _CLASS_REGISTRY:
+        return _CLASS_REGISTRY[path]
+    else:
+        assert "." in path, (
+            f"Because {path} is not found in class registry, " f"it must be a full module path"
+        )
+    try:
+        from importlib import import_module
+
+        module_path, _, class_name = path.rpartition(".")
+        mod = import_module(module_path)
+        try:
+            class_type = getattr(mod, class_name)
+        except AttributeError:
+            raise ImportError("Class {} is not in module {}".format(class_name, module_path))
+        return class_type
+    except ValueError as e:
+        print("Error initializing class " + path, file=sys.stderr)
+        raise e
+
+
+_DELETE_ARG = "__delete__"
+_NO_INSTANTIATE = "__no_instantiate__"  # return config as-is
+_OMEGA_MISSING = "???"
+
+
+def _get_instantiate_params(cfg, kwargs=None):
+    params = cfg
+    f_args, f_kwargs = (), {}
+    for k, value in params.items():
+        if k in ["cls", "class"]:
+            continue
+        elif k == "*args":
+            assert is_sequence(value), '"*args" value must be a sequence'
+            f_args = list(value)
+            continue
+        if value == _OMEGA_MISSING:
+            if kwargs and k in kwargs:
+                value = kwargs[k]
+            else:
+                raise ValueError(f'Missing required keyword arg "{k}" in cfg: {cfg}')
+        if value == _DELETE_ARG:
+            continue
+        else:
+            f_kwargs[k] = value
+    return f_args, f_kwargs
+
+
+def _instantiate_single(cfg):
+    if is_mapping(cfg) and ("cls" in cfg or "class" in cfg):
+        assert bool("cls" in cfg) != bool("class" in cfg), (
+            "to instantiate from config, "
+            'one and only one of "cls" or "class" key should be provided'
+        )
+        if _NO_INSTANTIATE in cfg:
+            no_instantiate = cfg.pop(_NO_INSTANTIATE)
+            if no_instantiate:
+                cfg = deepcopy(cfg)
+                return cfg
+            else:
+                return _instantiate_single(cfg)
+
+        cls = cfg.get("class", cfg.get("cls"))
+        args, kwargs = _get_instantiate_params(cfg)
+        try:
+            class_type = get_class(cls)
+            return class_type(*args, **kwargs)
+        except Exception as e:
+            raise RuntimeError(f"Error instantiating {cls}: {e}")
+    else:
+        return None
+
+
+def instantiate(_cfg_, **kwargs):
+    """
+    Any dict with "cls" or "class" key is considered instantiable.
+
+    Any key that has the special value "__delete__"
+    will not be passed to the constructor
+
+    **kwargs only apply to the top level object if it's a dict, otherwise raise error
+    """
+    assert OmegaConf.is_config(_cfg_) or isinstance(_cfg_, (list, tuple)) or is_mapping(_cfg_), (
+        '"cfg" must be a dict, list, tuple, or OmegaConf config to be instantiated. '
+        f"Current its type is {type(_cfg_)}"
+    )
+
+    _cfg_ = omegaconf_to_dict(_cfg_, resolve=True)
+
+    if kwargs:
+        if is_mapping(_cfg_):
+            _cfg_ = _cfg_.copy()
+            _cfg_.update(kwargs)
+            _cfg_ = {k: v for k, v in _cfg_.items() if v != _DELETE_ARG}
+        else:
+            raise RuntimeError(
+                f"**kwargs specified, but the top-level cfg is not a dict. "
+                f"It has type {type(_cfg_)}"
+            )
+
+    return tree.traverse(_instantiate_single, _cfg_, top_down=False)

groot/vla/common/utils/io/file_utils.py

@@ -0,0 +1,707 @@
+"""
+File system utils.
+"""
+
+import glob
+import os
+import pickle
+import shutil
+import sys
+from typing import Callable, Union
+
+from ..data_structure.tree_utils import is_sequence
+
+__all__ = [
+    "create_tar",
+    "dump_pickle",
+    "dump_text",
+    "dump_text_lines",
+    "extract_tar",
+    "f_add_ext",
+    "f_append_before_ext",
+    "f_copy",
+    "f_copytree",
+    "f_exists",
+    "f_expand",
+    "f_ext",
+    "f_glob",
+    "f_has_ext",
+    "f_join",
+    "f_listdir",
+    "f_mkdir",
+    "f_mkdir_in_path",
+    "f_move",
+    "f_not_empty",
+    "f_remove",
+    "f_size",
+    "f_split_path",
+    "f_time",
+    "get_dir",
+    "get_file_lock",
+    "get_package_root",
+    "get_parent_dir",
+    "get_script_dir",
+    "get_script_file_name",
+    "get_script_self_path",
+    "host_id",
+    "host_name",
+    "insert_before_ext",
+    "is_abs_path",
+    "is_dir",
+    "is_file",
+    "is_relative_path",
+    "last_part_in_path",
+    "load_pickle",
+    "load_text",
+    "load_text_lines",
+    "md5_checksum",
+    "move_with_backup",
+    "next_available_file_name",
+    "owner_name",
+    "pickle_dump",
+    "pickle_load",
+    "read_text",
+    "read_text_lines",
+    "text_dump",
+    "text_load",
+    "timestamp_file_name",
+    "utf_open",
+    "write_text",
+    "write_text_lines",
+]
+
+f_ext = os.path.splitext
+
+f_size = os.path.getsize
+
+is_file = os.path.isfile
+
+is_dir = os.path.isdir
+
+get_dir = os.path.dirname
+
+
+def owner_name(filepath):
+    """
+    Returns: file owner name, unix only
+    """
+    import pwd
+
+    return pwd.getpwuid(os.stat(filepath).st_uid).pw_name
+
+
+def host_name():
+    "Get host name, alias with ``socket.gethostname()``"
+    from socket import gethostname
+
+    return gethostname()
+
+
+def host_id():
+    """
+    Returns: first part of hostname up to '.'
+    """
+    return host_name().split(".")[0]
+
+
+def utf_open(fname, mode):
+    """
+    Wrapper for codecs.open
+    """
+    import codecs
+
+    return codecs.open(fname, mode=mode, encoding="utf-8")
+
+
+def f_not_empty(*fpaths):
+    """
+    Returns:
+        True if and only if the file exists and file size > 0
+          if fpath is a dir, if and only if dir exists and has at least 1 file
+    """
+    fpath = f_join(*fpaths)
+    if not os.path.exists(fpath):
+        return False
+
+    if os.path.isdir(fpath):
+        return len(os.listdir(fpath)) > 0
+    else:
+        return os.path.getsize(fpath) > 0
+
+
+def f_expand(fpath):
+    return os.path.expandvars(os.path.expanduser(fpath))
+
+
+def f_exists(*fpaths):
+    return os.path.exists(f_join(*fpaths))
+
+
+def f_join(*fpaths):
+    """
+    Join file paths and expand special symbols like `~` for home dir
+    """
+
+    def pack_varargs(args):
+        """
+        Pack *args or a single list arg as list
+
+        def f(*args):
+            arg_list = pack_varargs(args)
+            # arg_list is now packed as a list
+        """
+        assert isinstance(args, tuple), "please input the tuple `args` as in *args"
+        if len(args) == 1 and is_sequence(args[0]):
+            return args[0]
+        else:
+            return args
+
+    fpaths = pack_varargs(fpaths)
+    fpath = f_expand(os.path.join(*fpaths))
+    if isinstance(fpath, str):
+        fpath = fpath.strip()
+    return fpath
+
+
+def f_listdir(
+    *fpaths,
+    filter_ext=None,
+    filter=None,
+    sort=True,
+    full_path=False,
+    nonexist_ok=True,
+    recursive=False,
+):
+    """
+    Args:
+        full_path: True to return full paths to the dir contents
+        filter: function that takes in file name and returns True to include
+        nonexist_ok: True to return [] if the dir is non-existent, False to raise
+        sort: sort the file names by alphabetical
+        recursive: True to use os.walk to recursively list files. Note that `filter`
+            will be applied to the relative path string to the root dir.
+            e.g. filter will take "a/data1.txt" and "a/b/data3.txt" as input, instead of
+            just the base file names "data1.txt" and "data3.txt".
+            if False, will simply call os.listdir()
+    """
+    assert not (filter_ext and filter), "filter_ext and filter are mutually exclusive"
+    dir_path = f_join(*fpaths)
+    if not os.path.exists(dir_path) and nonexist_ok:
+        return []
+    if recursive:
+        files = [
+            os.path.join(os.path.relpath(root, dir_path), file)
+            for root, _, files in os.walk(dir_path)
+            for file in files
+        ]
+    else:
+        files = os.listdir(dir_path)
+    if filter is not None:
+        files = [f for f in files if filter(f)]
+    elif filter_ext is not None:
+        files = [f for f in files if f.endswith(filter_ext)]
+    if sort:
+        files.sort()
+    if full_path:
+        return [os.path.join(dir_path, f) for f in files]
+    else:
+        return files
+
+
+def f_mkdir(*fpaths):
+    """
+    Recursively creates all the subdirs
+    If exist, do nothing.
+    """
+    fpath = f_join(*fpaths)
+    os.makedirs(fpath, exist_ok=True)
+    return fpath
+
+
+def f_mkdir_in_path(*fpaths):
+    """
+    fpath is a file,
+    recursively creates all the parent dirs that lead to the file
+    If exist, do nothing.
+    """
+    os.makedirs(get_dir(f_join(*fpaths)), exist_ok=True)
+
+
+def last_part_in_path(fpath):
+    """
+    https://stackoverflow.com/questions/3925096/how-to-get-only-the-last-part-of-a-path-in-python
+    """
+    return os.path.basename(os.path.normpath(f_expand(fpath)))
+
+
+def is_abs_path(*fpath):
+    return os.path.isabs(f_join(*fpath))
+
+
+def is_relative_path(*fpath):
+    return not is_abs_path(f_join(*fpath))
+
+
+def f_time(*fpath):
+    "File modification time"
+    return str(os.path.getctime(f_join(*fpath)))
+
+
+def f_append_before_ext(fpath, suffix):
+    """
+    Append a suffix to file name and retain its extension
+    """
+    name, ext = f_ext(fpath)
+    return name + suffix + ext
+
+
+def f_add_ext(fpath, ext):
+    """
+    Append an extension if not already there
+    Args:
+      ext: will add a preceding `.` if doesn't exist
+    """
+    if not ext.startswith("."):
+        ext = "." + ext
+    if fpath.endswith(ext):
+        return fpath
+    else:
+        return fpath + ext
+
+
+def f_has_ext(fpath, ext):
+    "Test if file path is a text file"
+    _, actual_ext = f_ext(fpath)
+    return actual_ext == "." + ext.lstrip(".")
+
+
+def f_glob(*fpath):
+    return glob.glob(f_join(*fpath), recursive=True)
+
+
+def f_remove(*fpath, verbose=False, dry_run=False):
+    """
+    If exist, remove. Supports both dir and file. Supports glob wildcard.
+    """
+    import errno
+
+    assert isinstance(verbose, bool)
+    fpath = f_join(fpath)
+    if dry_run:
+        print("Dry run, delete:", fpath)
+        return
+    for f in glob.glob(fpath):
+        try:
+            shutil.rmtree(f)
+        except OSError as e:
+            if e.errno == errno.ENOTDIR:
+                try:
+                    os.remove(f)
+                except Exception as e:  # final resort safeguard
+                    pass
+    if verbose:
+        print(f'Deleted "{fpath}"')
+
+
+def f_copy(fsrc, fdst, ignore=None, include=None, exists_ok=True, verbose=False):
+    """
+    Supports both dir and file. Supports glob wildcard.
+    """
+    import errno
+
+    fsrc, fdst = f_expand(fsrc), f_expand(fdst)
+    for f in glob.glob(fsrc):
+        try:
+            f_copytree(f, fdst, ignore=ignore, include=include, exist_ok=exists_ok)
+        except OSError as e:
+            if e.errno == errno.ENOTDIR:
+                shutil.copy(f, fdst)
+            else:
+                raise
+    if verbose:
+        print(f'Copied "{fsrc}" to "{fdst}"')
+
+
+def _f_copytree(
+    src,
+    dst,
+    symlinks=False,
+    ignore=None,
+    exist_ok=True,
+    copy_function=shutil.copy2,
+    ignore_dangling_symlinks=False,
+):
+    """Copied from python standard lib shutil.copytree
+    except that we allow exist_ok
+    Use f_copytree as entry
+    """
+    names = os.listdir(src)
+    if ignore is not None:
+        ignored_names = ignore(src, names)
+    else:
+        ignored_names = set()
+
+    os.makedirs(dst, exist_ok=exist_ok)
+    errors = []
+    for name in names:
+        if name in ignored_names:
+            continue
+        srcname = os.path.join(src, name)
+        dstname = os.path.join(dst, name)
+        try:
+            if os.path.islink(srcname):
+                linkto = os.readlink(srcname)
+                if symlinks:
+                    # We can't just leave it to `copy_function` because legacy
+                    # code with a custom `copy_function` may rely on copytree
+                    # doing the right thing.
+                    os.symlink(linkto, dstname)
+                    shutil.copystat(srcname, dstname, follow_symlinks=not symlinks)
+                else:
+                    # ignore dangling symlink if the flag is on
+                    if not os.path.exists(linkto) and ignore_dangling_symlinks:
+                        continue
+                    # otherwise let the copy occurs. copy2 will raise an error
+                    if os.path.isdir(srcname):
+                        _f_copytree(srcname, dstname, symlinks, ignore, exist_ok, copy_function)
+                    else:
+                        copy_function(srcname, dstname)
+            elif os.path.isdir(srcname):
+                _f_copytree(srcname, dstname, symlinks, ignore, exist_ok, copy_function)
+            else:
+                # Will raise a SpecialFileError for unsupported file types
+                copy_function(srcname, dstname)
+        # catch the Error from the recursive copytree so that we can
+        # continue with other files
+        except shutil.Error as err:
+            errors.extend(err.args[0])
+        except OSError as why:
+            errors.append((srcname, dstname, str(why)))
+    try:
+        shutil.copystat(src, dst)
+    except OSError as why:
+        # Copying file access times may fail on Windows
+        if getattr(why, "winerror", None) is None:
+            errors.append((src, dst, str(why)))
+    if errors:
+        raise shutil.Error(errors)
+    return dst
+
+
+def _include_patterns(*patterns):
+    """Factory function that can be used with copytree() ignore parameter.
+
+    Arguments define a sequence of glob-style patterns
+    that are used to specify what files to NOT ignore.
+    Creates and returns a function that determines this for each directory
+    in the file hierarchy rooted at the source directory when used with
+    shutil.copytree().
+    """
+
+    def _ignore_patterns(path, names):
+        import fnmatch
+
+        keep = set(name for pattern in patterns for name in fnmatch.filter(names, pattern))
+        ignore = set(
+            name
+            for name in names
+            if name not in keep and not os.path.isdir(os.path.join(path, name))
+        )
+        return ignore
+
+    return _ignore_patterns
+
+
+def f_copytree(fsrc, fdst, symlinks=False, ignore=None, include=None, exist_ok=True):
+    fsrc, fdst = f_expand(fsrc), f_expand(fdst)
+    assert (ignore is None) or (include is None), "ignore= and include= are mutually exclusive"
+    if ignore:
+        ignore = shutil.ignore_patterns(*ignore)
+    elif include:
+        ignore = _include_patterns(*include)
+    _f_copytree(fsrc, fdst, ignore=ignore, symlinks=symlinks, exist_ok=exist_ok)
+
+
+def f_move(fsrc, fdst):
+    fsrc, fdst = f_expand(fsrc), f_expand(fdst)
+    for f in glob.glob(fsrc):
+        shutil.move(f, fdst)
+
+
+def f_split_path(fpath, normpath=True):
+    """
+    Splits path into a list of its component folders
+
+    Args:
+        normpath: call os.path.normpath to remove redundant '/' and
+            up-level references like ".."
+    """
+    if normpath:
+        fpath = os.path.normpath(fpath)
+    allparts = []
+    while 1:
+        parts = os.path.split(fpath)
+        if parts[0] == fpath:  # sentinel for absolute paths
+            allparts.insert(0, parts[0])
+            break
+        elif parts[1] == fpath:  # sentinel for relative paths
+            allparts.insert(0, parts[1])
+            break
+        else:
+            fpath = parts[0]
+            allparts.insert(0, parts[1])
+    return allparts
+
+
+def get_script_dir():
+    """
+    Returns: the dir of current script
+    """
+    return os.path.dirname(os.path.realpath(sys.argv[0]))
+
+
+def get_script_file_name():
+    """
+    Returns: the dir of current script
+    """
+    return os.path.basename(sys.argv[0])
+
+
+def get_script_self_path():
+    """
+    Returns: the dir of current script
+    """
+    return os.path.realpath(sys.argv[0])
+
+
+def get_parent_dir(location, abspath=False):
+    """
+    Args:
+      location: current directory or file
+
+    Returns:
+        parent directory absolute or relative path
+    """
+    _path = os.path.abspath if abspath else os.path.relpath
+    return _path(f_join(location, os.pardir))
+
+
+def md5_checksum(*fpath):
+    """
+    File md5 signature
+    """
+    import hashlib
+
+    hash_md5 = hashlib.md5()
+    with open(f_join(*fpath), "rb") as f:
+        for chunk in iter(lambda: f.read(65536), b""):
+            hash_md5.update(chunk)
+    return hash_md5.hexdigest()
+
+
+def create_tar(fsrc, output_tarball, include=None, ignore=None, compress_mode="gz"):
+    """
+    Args:
+        fsrc: source file or folder
+        output_tarball: output tar file name
+        compress_mode: ``gz``, ``bz2``, ``xz`` or ``''`` (empty for uncompressed write)
+        include: include pattern, will trigger copy to temp directory
+        ignore: ignore pattern, will trigger copy to temp directory
+    """
+    import tarfile
+    import tempfile
+
+    fsrc, output_tarball = f_expand(fsrc), f_expand(output_tarball)
+    assert compress_mode in ["gz", "bz2", "xz", ""]
+    src_base = os.path.basename(fsrc)
+
+    tempdir = None
+    if include or ignore:
+        tempdir = tempfile.mkdtemp()
+        tempdest = f_join(tempdir, src_base)
+        f_copy(fsrc, tempdest, include=include, ignore=ignore)
+        fsrc = tempdest
+
+    with tarfile.open(output_tarball, "w:" + compress_mode) as tar:
+        tar.add(fsrc, arcname=src_base)
+
+    if tempdir:
+        f_remove(tempdir)
+
+
+def extract_tar(source_tarball, output_dir=".", members=None):
+    """
+    Args:
+        source_tarball: extract members from archive
+        output_dir: default to current working dir
+        members: must be a subset of the list returned by getmembers()
+    """
+    import tarfile
+
+    source_tarball, output_dir = f_expand(source_tarball), f_expand(output_dir)
+    with tarfile.open(source_tarball, "r:*") as tar:
+        tar.extractall(output_dir, members=members)
+
+
+def move_with_backup(*fpath, suffix=".bak"):
+    """
+    Ensures that a path is not occupied. If there is a file, rename it by
+    adding @suffix. Resursively backs up everything.
+
+    Args:
+        fpath: file path to clear
+        suffix: Add to backed up files (default: {'.bak'})
+    """
+    fpath = str(f_join(*fpath))
+    if os.path.exists(fpath):
+        move_with_backup(fpath + suffix)
+        shutil.move(fpath, fpath + suffix)
+
+
+def insert_before_ext(name, insert):
+    """
+    log.txt -> log.ep50.txt
+    """
+    name, ext = os.path.splitext(name)
+    return name + insert + ext
+
+
+def timestamp_file_name(fname):
+    from datetime import datetime
+
+    timestr = datetime.now().strftime("_%H-%M-%S_%m-%d-%y")
+    return insert_before_ext(fname, timestr)
+
+
+def next_available_file_name(
+    *fpath,
+    suffix_template: Union[str, Callable[[int], str]] = "_v{i+1}",
+    before_ext: bool = True,
+):
+    """
+    Args:
+        suffix_template: a format string using "i" variable or
+            lambda int -> str
+        before_ext: True to insert suffix before the extension
+    """
+
+    def fstring(fmt_str, **kwargs):
+        """
+        Simulate python f-string but without `f`
+        """
+        import shlex
+
+        locals().update(kwargs)
+        return eval("f" + shlex.quote(fmt_str))
+
+    orig_file_path = f_join(*fpath)
+    i = 0
+    fpath = orig_file_path
+    while os.path.exists(fpath):
+        if isinstance(suffix_template, str):
+            suffix = fstring(suffix_template, i=i)
+        elif callable(suffix_template):
+            suffix = suffix_template(i)
+            assert isinstance(suffix, str)
+        else:
+            raise NotImplementedError(f"Unsupported suffix template {suffix_template}")
+        if before_ext:
+            fpath = insert_before_ext(orig_file_path, suffix)
+        else:
+            fpath = orig_file_path + suffix
+        i += 1
+    return fpath
+
+
+def get_file_lock(*fpath, timeout: int = 15, logging_level="critical"):
+    """
+    NFS-safe filesystem-backed lock. `pip install flufl.lock`
+    https://flufllock.readthedocs.io/en/stable/apiref.html
+
+    Args:
+        fpath: should be a path on NFS so that every process can see it
+        timeout: seconds
+    """
+    import logging
+
+    from flufl.lock import Lock
+
+    logging.getLogger("flufl.lock").setLevel(logging_level.upper())
+    return Lock(f_join(*fpath), lifetime=timeout)
+
+
+def load_pickle(*fpaths):
+    with open(f_join(*fpaths), "rb") as fp:
+        return pickle.load(fp)
+
+
+def dump_pickle(data, *fpaths):
+    with open(f_join(*fpaths), "wb") as fp:
+        pickle.dump(data, fp)
+
+
+def load_text(*fpaths, by_lines=False):
+    with open(f_join(*fpaths), "r") as fp:
+        if by_lines:
+            return fp.readlines()
+        else:
+            return fp.read()
+
+
+def load_text_lines(*fpaths):
+    return load_text(*fpaths, by_lines=True)
+
+
+def dump_text(s, *fpaths):
+    with open(f_join(*fpaths), "w") as fp:
+        fp.write(s)
+
+
+def dump_text_lines(lines: list[str], *fpaths, add_newline=True):
+    with open(f_join(*fpaths), "w") as fp:
+        for line in lines:
+            print(line, file=fp, end="\n" if add_newline else "")
+
+
+def get_package_root() -> str:
+    import importlib.util
+    import inspect
+
+    # Get the current frame
+    current_frame = inspect.currentframe()
+    if current_frame is None:
+        raise ImportError("Cannot determine the package name from __package__")
+
+    # Get the caller module
+    caller_module = inspect.getmodule(current_frame.f_back)
+    if caller_module is None:
+        raise ImportError("Cannot determine the package name from __package__")
+
+    # Get the package name
+    package_name = caller_module.__package__
+    if not package_name:
+        raise ImportError("Cannot determine the package name from __package__")
+
+    # Get the top-level package name
+    top_package_name = package_name.split(".")[0]
+
+    # Find the package specification
+    spec = importlib.util.find_spec(top_package_name)
+
+    if spec and spec.origin:
+        # Get the directory containing the package's __init__.py file
+        package_dir = os.path.dirname(spec.origin)
+        return package_dir
+    else:
+        raise ImportError(f"Cannot find the package {top_package_name}")
+
+
+# aliases to be consistent with other load_* and dump_*
+pickle_load = load_pickle
+pickle_dump = dump_pickle
+text_load = load_text
+read_text = load_text
+read_text_lines = load_text_lines
+write_text = dump_text
+write_text_lines = dump_text_lines
+text_dump = dump_text

groot/vla/common/utils/io/hdf5_utils.py

@@ -0,0 +1,84 @@
+import h5py
+import numpy as np
+from pydantic import BaseModel
+
+
+def hdf5_save(data: BaseModel | dict, group: h5py.Group) -> None:
+    """Recursively save Pydantic model or dict to HDF5 group."""
+    if isinstance(data, BaseModel):
+        # Convert to dict and exclude None values
+        data_dict = data.model_dump(mode="python", exclude_none=True)
+    else:
+        data_dict = data
+
+    for key, value in data_dict.items():
+        if isinstance(value, np.ndarray):
+            group.create_dataset(key, data=value)
+        elif isinstance(value, (BaseModel, dict)):
+            subgroup = group.create_group(key)
+            hdf5_save(value, subgroup)
+        else:
+            # For primitive types, convert to numpy array
+            try:
+                group.create_dataset(key, data=np.array(value))
+            except TypeError:
+                raise ValueError(f"Unsupported type: {type(value)} for key: {key}")
+
+
+def hdf5_load(group: h5py.Group) -> dict:
+    """Recursively load HDF5 group to Pydantic model or dict."""
+    data_dict = {}
+    for key, value in group.items():
+        if isinstance(value, h5py.Dataset):
+            data_dict[key] = value[()]
+        elif isinstance(value, h5py.Group):
+            data_dict[key] = hdf5_load(value)
+    return data_dict
+
+
+def hdf5_is_subset(this: h5py.Group, other: h5py.Group, verbose: bool = False) -> bool:
+    """Check if this HDF5 group is a subset of another HDF5 group."""
+    for key, value in this.items():
+        if key not in other:
+            if verbose:
+                print(f"Key {key} not in other")
+            return False
+        elif isinstance(value, h5py.Group):
+            if not isinstance(other[key], h5py.Group):
+                if verbose:
+                    print(f"Key {key} is not a group in other")
+                return False
+            if not hdf5_is_subset(value, other[key], verbose):
+                if verbose:
+                    print(f"Key {key} is not a subset of other")
+                return False
+        elif isinstance(value, h5py.Dataset):
+            if not isinstance(other[key], h5py.Dataset):
+                if verbose:
+                    print(f"Key {key} is not a dataset in other")
+                return False
+            if not np.array_equal(value, other[key]):
+                if verbose:
+                    print(f"Key {key} is not equal in other")
+                return False
+        elif isinstance(value, h5py.Datatype):
+            if not isinstance(other[key], h5py.Datatype):
+                if verbose:
+                    print(f"Key {key} is not a datatype in other")
+                return False
+            if value != other[key]:
+                if verbose:
+                    print(f"Key {key} is not equal in other")
+                return False
+        else:
+            # try to compare
+            if value != other[key]:
+                if verbose:
+                    print(f"Key {key} is not equal in other")
+                return False
+    return True
+
+
+def hdf5_is_equal(this: h5py.Group, other: h5py.Group, verbose: bool = False) -> bool:
+    """Check if this HDF5 group is equal to another HDF5 group."""
+    return hdf5_is_subset(this, other, verbose) and hdf5_is_subset(other, this, verbose)

groot/vla/common/utils/io/json_utils.py

@@ -0,0 +1,270 @@
+"""
+JSON, YAML, and python config file utilities
+"""
+
+from io import StringIO
+import json
+import os.path as path
+
+import yaml
+
+from ..misc.functional_utils import make_recursive_func
+from .file_utils import f_join
+
+__all__ = [
+    "json_load",
+    "json_loads",
+    "jsonl_load",
+    "yaml_load",
+    "yaml_loads",
+    "json_dump",
+    "json_dumps",
+    "jsonl_dump",
+    "yaml_dump",
+    "yaml_dumps",
+    "json_or_yaml_load",
+    "json_or_yaml_dump",
+    "Jsonl",
+    # ---------------- Aliases -----------------
+    "load_json",
+    "loads_json",
+    "load_jsonl",
+    "load_yaml",
+    "loads_yaml",
+    "dump_json",
+    "dumps_json",
+    "dump_jsonl",
+    "dump_yaml",
+    "dumps_yaml",
+    "load_json_or_yaml",
+    "dump_json_or_yaml",
+]
+
+from typing import Dict, List
+
+from typing_extensions import Literal
+
+
+def json_load(*file_path, **kwargs):
+    file_path = f_join(file_path)
+    with open(file_path, "r") as fp:
+        return json.load(fp, **kwargs)
+
+
+def json_loads(string, **kwargs):
+    return json.loads(string, **kwargs)
+
+
+def jsonl_load(*file_path, **kwargs):
+    file_path = f_join(file_path)
+    data = []
+    for line in open(file_path):
+        data.append(json.loads(line, **kwargs))
+    return data
+
+
+@make_recursive_func
+def any_to_primitive(x):
+    try:
+        import torch
+    except ImportError:
+        raise ImportError("torch is required for any_to_primitive")
+    import numpy as np
+
+    if isinstance(x, (np.ndarray, np.number, torch.Tensor)):
+        return x.tolist()
+    else:
+        return x
+
+
+def json_dump(data, *file_path, convert_to_primitive=False, **kwargs):
+    if convert_to_primitive:
+        data = any_to_primitive(data)
+    file_path = f_join(file_path)
+    with open(file_path, "w") as fp:
+        json.dump(data, fp, **kwargs)
+
+
+def json_dumps(data, convert_to_primitive=False, **kwargs):
+    """
+    Returns: string
+    """
+    if convert_to_primitive:
+        data = any_to_primitive(data)
+    return json.dumps(data, **kwargs)
+
+
+def jsonl_dump(data, *file_path):
+    from .file_utils import is_sequence
+
+    assert is_sequence(data)
+    data = any_to_primitive(data)
+    file_path = f_join(file_path)
+    with open(file_path, "w") as fp:
+        for line in data:
+            print(json.dumps(line), file=fp, flush=True)
+
+
+def yaml_load(*file_path, loader=yaml.safe_load, **kwargs):
+    file_path = f_join(file_path)
+    with open(file_path, "r") as fp:
+        return loader(fp, **kwargs)
+
+
+def yaml_loads(string, *, loader=yaml.safe_load, **kwargs):
+    return loader(string, **kwargs)
+
+
+def yaml_dump(data, *file_path, dumper=yaml.safe_dump, convert_to_primitive=False, **kwargs):
+    if convert_to_primitive:
+        data = any_to_primitive(data)
+    file_path = f_join(file_path)
+    indent = kwargs.pop("indent", 2)
+    default_flow_style = kwargs.pop("default_flow_style", False)
+    sort_keys = kwargs.pop("sort_keys", False)  # preserves original dict order
+    with open(file_path, "w") as fp:
+        dumper(
+            data,
+            stream=fp,
+            indent=indent,
+            default_flow_style=default_flow_style,
+            sort_keys=sort_keys,
+            **kwargs,
+        )
+
+
+def yaml_dumps(data, *, dumper=yaml.safe_dump, convert_to_primitive=False, **kwargs):
+    "Returns: string"
+    if convert_to_primitive:
+        data = any_to_primitive(data)
+    stream = StringIO()
+    indent = kwargs.pop("indent", 2)
+    default_flow_style = kwargs.pop("default_flow_style", False)
+    sort_keys = kwargs.pop("sort_keys", False)  # preserves original dict order
+    dumper(
+        data,
+        stream,
+        indent=indent,
+        default_flow_style=default_flow_style,
+        sort_keys=sort_keys,
+        **kwargs,
+    )
+    return stream.getvalue()
+
+
+# ==================== auto-recognize extension ====================
+def json_or_yaml_load(*file_path, **loader_kwargs):
+    """
+    Args:
+        file_path: JSON or YAML loader depends on the file extension
+
+    Raises:
+        IOError: if extension is not ".json", ".yml", or ".yaml"
+    """
+    file_path = str(f_join(file_path))
+    if file_path.endswith(".json"):
+        return json_load(file_path, **loader_kwargs)
+    elif file_path.endswith(".yml") or file_path.endswith(".yaml"):
+        return yaml_load(file_path, **loader_kwargs)
+    else:
+        raise IOError(
+            f'unknown file extension: "{file_path}", '
+            f'loader supports only ".json", ".yml", ".yaml"'
+        )
+
+
+def json_or_yaml_dump(data, *file_path, **dumper_kwargs):
+    """
+    Args:
+        file_path: JSON or YAML loader depends on the file extension
+
+    Raises:
+        IOError: if extension is not ".json", ".yml", or ".yaml"
+    """
+    file_path = str(f_join(file_path))
+    if file_path.endswith(".json"):
+        return json_dump(data, file_path, **dumper_kwargs)
+    elif file_path.endswith(".yml") or file_path.endswith(".yaml"):
+        return yaml_dump(data, file_path, **dumper_kwargs)
+    else:
+        raise IOError(
+            f'unknown file extension: "{file_path}", '
+            f'dumper supports only ".json", ".yml", ".yaml"'
+        )
+
+
+# ---------------- Aliases -----------------
+# add aliases where verb goes first, json_load -> load_json
+load_json = json_load
+load_yaml = yaml_load
+load_jsonl = jsonl_load
+loads_json = json_loads
+loads_yaml = yaml_loads
+dump_json = json_dump
+dump_jsonl = jsonl_dump
+dump_yaml = yaml_dump
+dumps_json = json_dumps
+dumps_yaml = yaml_dumps
+load_json_or_yaml = json_or_yaml_load
+dump_json_or_yaml = json_or_yaml_dump
+
+# ==================== Jsonl ====================
+
+
+class Jsonl:
+    """
+    Both reader and writer, as if everything's in-memory
+    """
+
+    def __init__(self, *file_path, mode: Literal["r", "w", "a"] = "a"):
+        """
+        Args:
+            mode:
+            - 'r': file must already exists
+            - 'w': overwrite the file regardless of whether it exists or not
+            - 'a': create a new file if doesn't exist, or append to an existing file
+        """
+        assert mode in "rwa"
+        self._file_path = str(f_join(file_path))
+        self._mode = mode
+        if mode == "r":
+            assert path.exists(self._file_path)
+            self._fp = None
+        else:
+            self._fp = open(self._file_path, mode)
+        if path.exists(self._file_path) and mode != "w":
+            self.data = jsonl_load(self._file_path)
+        else:
+            self.data = []
+
+    def append(self, data: Dict):
+        if self._mode == "r":
+            raise RuntimeError("Jsonl read mode cannot call append()")
+        self.data.append(data)
+        print(json_dumps(data), file=self._fp, flush=True)
+
+    def extend(self, data_list: List[Dict]):
+        for data in data_list:
+            self.append(data)
+
+    def close(self):
+        if self._fp is not None:
+            self._fp.close()
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+    def __len__(self):
+        return len(self.data)
+
+    def __iter__(self):
+        return iter(self.data)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
+
+    def __bool__(self):
+        return bool(self.data)

groot/vla/common/utils/io/print_utils.py

@@ -0,0 +1,362 @@
+from datetime import datetime
+import io
+import logging
+import os
+import pprint
+import shlex
+import string
+import sys
+import textwrap
+import time
+import traceback
+from typing import Callable, Union
+
+import numpy as np
+from typing_extensions import Literal
+
+from ..misc.functional_utils import meta_decorator
+from ..misc.misc_utils import match_patterns
+
+
+def to_readable_count_str(value: int, precision: int = 2) -> str:
+    assert value >= 0
+    labels = [" ", "K", "M", "B", "T"]
+    num_digits = int(np.floor(np.log10(value)) + 1 if value > 0 else 1)
+    num_groups = int(np.ceil(num_digits / 3))
+    num_groups = min(num_groups, len(labels))  # don't abbreviate beyond trillions
+    shift = -3 * (num_groups - 1)
+    value = value * (10**shift)
+    index = num_groups - 1
+    rem = value - int(value)
+    if precision > 0 and rem > 0.01:
+        fmt = f"{{:.{precision}f}}"
+        rem_str = fmt.format(rem).lstrip("0")
+    else:
+        rem_str = ""
+    return f"{int(value):,d}{rem_str} {labels[index]}"
+
+
+def to_scientific_str(value, precision: int = 1, capitalize: bool = False) -> str:
+    """
+    0.0015 -> "1.5e-3"
+    """
+    if value == 0:
+        return "0"
+    return f"{value:.{precision}e}".replace("e-0", "E-" if capitalize else "e-")
+
+
+def print_str(*args, **kwargs):
+    """
+    Same as print() signature but returns a string
+    """
+    sstream = io.StringIO()
+    kwargs.pop("file", None)
+    print(*args, **kwargs, file=sstream)
+    return sstream.getvalue()
+
+
+def fstring(fmt_str, **kwargs):
+    """
+    Simulate python f-string but without `f`
+    """
+    locals().update(kwargs)
+    return eval("f" + shlex.quote(fmt_str))
+
+
+def get_format_keys(fmt_str):
+    keys = []
+    for literal, field_name, fmt_spec, conversion in string.Formatter().parse(fmt_str):
+        if field_name:
+            keys.append(field_name)
+    return keys
+
+
+def get_timestamp(milli_precision: int = 3):
+    fmt = "%y-%m-%d %H:%M:%S"
+    if milli_precision > 0:
+        fmt += ".%f"
+    stamp = datetime.now().strftime(fmt)
+    if milli_precision > 0:
+        stamp = stamp[:-milli_precision]
+    return stamp
+
+
+def pretty_repr_str(obj, **kwargs):
+    """
+    Useful to produce __repr__()
+    """
+    if isinstance(obj, str):
+        cls_name = obj
+    else:
+        cls_name = obj.__class__.__name__
+    kw_strs = [k + "=" + pprint.pformat(v, indent=2, compact=True) for k, v in kwargs.items()]
+    new_line = len(cls_name) + sum(len(kw) for kw in kw_strs) > 84
+    if new_line:
+        kw = ",\n".join(kw_strs)
+        return f"{cls_name}(\n{textwrap.indent(kw, '  ')}\n)"
+    else:
+        kw = ", ".join(kw_strs)
+        return f"{cls_name}({kw})"
+
+
+def pprint_(*objs, **kwargs):
+    """
+    Use pprint to format the objects
+    """
+    print(
+        *[pprint.pformat(obj, indent=2) if not isinstance(obj, str) else obj for obj in objs],
+        **kwargs,
+    )
+
+
+def get_exception_info(to_str: bool = False):
+    """
+    Returns:
+        {'type': ExceptionType, 'value': ExceptionObject, 'trace': <traceback str>}
+    """
+    typ_, value, trace = sys.exc_info()
+    return {
+        "type": typ_.__name__ if to_str else typ_,
+        "value": str(value) if to_str else value,
+        "trace": "".join(traceback.format_exception(typ_, value, trace)),
+    }
+
+
+class DebugPrinter:
+    """
+    Debug print, usage: dprint = DebugPrint(enabled=True)
+    dprint(...)
+    """
+
+    def __init__(self, enabled, tensor_summary: Literal["shape", "shape+dtype", "none"] = "shape"):
+        """
+        Args:
+            tensor_summary:
+                - shape: only prints shape
+                - shape+dtype: also prints dtype and device
+                - none: print full tensor
+        """
+        self.enabled = enabled
+        assert tensor_summary in ["shape", "shape+dtype", "none"]
+        self.tensor_summary = tensor_summary
+
+    def __call__(self, *args, **kwargs):
+        if not self.enabled:
+            return
+        args = [self._process_arg(a) for a in args]
+        pprint_(*args, **kwargs)
+
+    def _process_arg(self, arg):
+        import numpy as np
+        import torch
+
+        if torch.is_tensor(arg):
+            if self.tensor_summary == "shape":
+                return str(list(arg.size()))
+            elif self.tensor_summary == "shape+dtype":
+                return f"{arg.dtype}{list(arg.size())}|{arg.device}"
+        elif isinstance(arg, np.ndarray):
+            if self.tensor_summary == "shape":
+                return str(list(arg.shape))
+            elif self.tensor_summary == "shape+dtype":
+                return f"{arg.dtype}{list(arg.shape)}"
+        return arg
+
+
+@meta_decorator
+def watch(func, seconds: int = 5, max_times: int = 0, keep_returns: bool = False):
+    """
+    Decorator: executes a function repeated with the args and
+        emulate `watch -n` capability
+
+    See `gpustat` repo: https://github.com/wookayin/gpustat/pull/41/files
+
+    Args:
+        max_times: watch for `max_times` and then exit. If 0, never exits
+        keep_returns: if True, will keep the return value from the function
+            and return as a list at the end
+    """
+    from blessings import Terminal
+
+    def _wrapped(*args, **kwargs):
+        term = Terminal()
+        N = 0
+        returns = []
+        with term.fullscreen():
+            while True:
+                try:
+                    with term.location(0, 0):
+                        ret = func(*args, **kwargs)
+                        print(term.clear_eos, end="")
+                        if keep_returns:
+                            returns.append(ret)
+                    N += 1
+                    if max_times > 0 and N >= max_times:
+                        break
+                    time.sleep(seconds)
+                except KeyboardInterrupt:
+                    break
+        return returns
+
+    return _wrapped
+
+
+class PrintRedirection(object):
+    """
+    Context manager: temporarily redirects stdout and stderr
+    """
+
+    def __init__(self, stdout=None, stderr=None):
+        """
+        Args:
+          stdout: if None, defaults to sys.stdout, unchanged
+          stderr: if None, defaults to sys.stderr, unchanged
+        """
+        if stdout is None:
+            stdout = sys.stdout
+        if stderr is None:
+            stderr = sys.stderr
+        self._stdout, self._stderr = stdout, stderr
+
+    def __enter__(self):
+        self._old_out, self._old_err = sys.stdout, sys.stderr
+        self._old_out.flush()
+        self._old_err.flush()
+        sys.stdout, sys.stderr = self._stdout, self._stderr
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.flush()
+        # restore the normal stdout and stderr
+        sys.stdout, sys.stderr = self._old_out, self._old_err
+
+    def flush(self):
+        "Manually flush the replaced stdout/stderr buffers."
+        self._stdout.flush()
+        self._stderr.flush()
+
+
+class PrintToFile(PrintRedirection):
+    """
+    Print to file and save/close the handle at the end.
+    """
+
+    def __init__(self, out_file=None, err_file=None):
+        """
+        Args:
+          out_file: file path
+          err_file: file path. If the same as out_file, print both stdout
+              and stderr to one file in order.
+        """
+        self.out_file, self.err_file = out_file, err_file
+        if out_file:
+            out_file = os.path.expanduser(out_file)
+            self.out_file = open(out_file, "w")
+        if err_file:
+            err_file = os.path.expanduser(out_file)
+            if err_file == out_file:  # redirect both stdout/err to one file
+                self.err_file = self.out_file
+            else:
+                self.err_file = open(os.path.expanduser(out_file), "w")
+        super().__init__(stdout=self.out_file, stderr=self.err_file)
+
+    def __exit__(self, *args):
+        super().__exit__(*args)
+        if self.out_file:
+            self.out_file.close()
+        if self.err_file:
+            self.err_file.close()
+
+
+def PrintSuppress(no_out=True, no_err=False):
+    """
+    Args:
+      no_out: stdout writes to sys.devnull
+      no_err: stderr writes to sys.devnull
+    """
+    out_file = os.devnull if no_out else None
+    err_file = os.devnull if no_err else None
+    return PrintToFile(out_file=out_file, err_file=err_file)
+
+
+class PrintString(PrintRedirection):
+    """
+    Redirect stdout and stderr to strings.
+    """
+
+    def __init__(self):
+        self.out_stream = io.StringIO()
+        self.err_stream = io.StringIO()
+        super().__init__(stdout=self.out_stream, stderr=self.err_stream)
+
+    def stdout(self):
+        "Returns: stdout as one string."
+        return self.out_stream.getvalue()
+
+    def stderr(self):
+        "Returns: stderr as one string."
+        return self.err_stream.getvalue()
+
+    def stdout_by_line(self):
+        "Returns: a list of stdout line by line, ignore trailing blanks"
+        return self.stdout().rstrip().split("\n")
+
+    def stderr_by_line(self):
+        "Returns: a list of stderr line by line, ignore trailing blanks"
+        return self.stderr().rstrip().split("\n")
+
+
+# ==================== Logging filters ====================
+class ExcludeLoggingFilter(logging.Filter):
+    """
+    Usage: logging.getLogger('name').addFilter(
+        ExcludeLoggingFilter(['info mess*age', 'Warning: *'])
+    )
+    Supports wildcard.
+    https://relaxdiego.com/2014/07/logging-in-python.html
+    """
+
+    def __init__(self, patterns):
+        super().__init__()
+        self._patterns = patterns
+
+    def filter(self, record):
+        if match_patterns(record.msg, include=self._patterns):
+            return False
+        else:
+            return True
+
+
+class ReplaceStringLoggingFilter(logging.Filter):
+    def __init__(self, patterns, replacer: Callable):
+        super().__init__()
+        self._patterns = patterns
+        assert callable(replacer)
+        self._replacer = replacer
+
+    def filter(self, record):
+        if match_patterns(record.msg, include=self._patterns):
+            record.msg = self._replacer(record.msg)
+
+
+def logging_exclude_pattern(
+    logger_name,
+    patterns: Union[str, list[str], Callable, list[Callable], None],
+):
+    """
+    Args:
+        patterns: see groot.vla.common.utils.misc_utils.match_patterns
+    """
+    logging.getLogger(logger_name).addFilter(ExcludeLoggingFilter(patterns))
+
+
+def logging_replace_string(
+    logger_name,
+    patterns: Union[str, list[str], Callable, list[Callable], None],
+    replacer: Callable,
+):
+    """
+    Args:
+        patterns: see groot.vla.common.utils.misc_utils.match_patterns
+    """
+    logging.getLogger(logger_name).addFilter(ReplaceStringLoggingFilter(patterns, replacer))

groot/vla/common/utils/io/termcolor.py

@@ -0,0 +1,186 @@
+# coding: utf-8
+# Copyright (c) 2008-2011 Volvox Development 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.
+#
+# 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.
+#
+# Original Author: Konstantin Lepa <konstantin.lepa@gmail.com>
+# Updated by Jim Fan
+
+"""ANSII Color formatting for output in terminal."""
+import io
+import os
+from typing import List, Optional, Union
+
+__ALL__ = ["color_text", "cprint"]
+
+STYLES = dict(
+    list(
+        zip(
+            ["bold", "dark", "", "underline", "blink", "", "reverse", "concealed"],
+            list(range(1, 9)),
+        )
+    )
+)
+del STYLES[""]
+
+
+HIGHLIGHTS = dict(
+    list(
+        zip(
+            ["grey", "red", "green", "yellow", "blue", "magenta", "cyan", "white"],
+            list(range(40, 48)),
+        )
+    )
+)
+
+
+COLORS = dict(
+    list(
+        zip(
+            ["grey", "red", "green", "yellow", "blue", "magenta", "cyan", "white"],
+            list(range(30, 38)),
+        )
+    )
+)
+
+
+def _strip_bg_prefix(color):
+    "on_red -> red"
+    if color.startswith("on_"):
+        return color[len("on_") :]
+    else:
+        return color
+
+
+RESET = "\033[0m"
+
+
+def color_text(
+    text,
+    color: Optional[str] = None,
+    bg_color: Optional[str] = None,
+    styles: Optional[Union[str, List[str]]] = None,
+):
+    """Colorize text.
+
+    Available text colors:
+        red, green, yellow, blue, magenta, cyan, white.
+
+    Available text highlights:
+        on_red, on_green, on_yellow, on_blue, on_magenta, on_cyan, on_white.
+
+    Available attributes:
+        bold, dark, underline, blink, reverse, concealed.
+
+    Example:
+        colored('Hello, World!', 'red', 'on_grey', ['blue', 'blink'])
+        colored('Hello, World!', 'green')
+    """
+    if os.getenv("ANSI_COLORS_DISABLED") is None:
+        fmt_str = "\033[%dm%s"
+        if color is not None:
+            text = fmt_str % (COLORS[color], text)
+
+        if bg_color is not None:
+            bg_color = _strip_bg_prefix(bg_color)
+            text = fmt_str % (HIGHLIGHTS[bg_color], text)
+
+        if styles is not None:
+            if isinstance(styles, str):
+                styles = [styles]
+            for style in styles:
+                text = fmt_str % (STYLES[style], text)
+
+        text += RESET
+    return text
+
+
+def cprint(
+    *args,
+    color: Optional[str] = None,
+    bg_color: Optional[str] = None,
+    styles: Optional[Union[str, List[str]]] = None,
+    **kwargs,
+):
+    """Print colorize text.
+
+    It accepts arguments of print function.
+    """
+    sstream = io.StringIO()
+    print(*args, sep=kwargs.pop("sep", None), end="", file=sstream)
+    text = sstream.getvalue()
+    print((color_text(text, color, bg_color, styles)), **kwargs)
+
+
+if __name__ == "__main__":
+    print("Current terminal type: %s" % os.getenv("TERM"))
+    print("Test basic colors:")
+    cprint("Grey color", color="grey")
+    cprint("Red color", color="red")
+    cprint("Green color", color="green")
+    cprint("Yellow color", color="yellow")
+    cprint("Blue color", color="blue")
+    cprint("Magenta color", color="magenta")
+    cprint("Cyan color", color="cyan")
+    cprint("White color", color="white")
+    print(("-" * 78))
+
+    print("Test highlights:")
+    cprint("On grey color", bg_color="on_grey")
+    cprint("On red color", bg_color="on_red")
+    cprint("On green color", bg_color="on_green")
+    cprint("On yellow color", bg_color="on_yellow")
+    cprint("On blue color", bg_color="on_blue")
+    cprint("On magenta color", bg_color="on_magenta")
+    cprint("On cyan color", bg_color="on_cyan")
+    cprint("On white color", color="grey", bg_color="on_white")
+    print("-" * 78)
+
+    print("Test attributes:")
+    cprint("Bold grey color", color="grey", styles="bold")
+    cprint("Dark red color", color="red", styles=["dark"])
+    cprint("Underline green color", color="green", styles=["underline"])
+    cprint("Blink yellow color", color="yellow", styles=["blink"])
+    cprint("Reversed blue color", color="blue", styles=["reverse"])
+    cprint("Concealed Magenta color", color="magenta", styles=["concealed"])
+    cprint(
+        "Bold underline reverse cyan color",
+        color="cyan",
+        styles=["bold", "underline", "reverse"],
+    )
+    cprint(
+        "Dark blink concealed white color",
+        color="white",
+        styles=["dark", "blink", "concealed"],
+    )
+    print(("-" * 78))
+
+    print("Test mixing:")
+    cprint(
+        "Underline red on grey color",
+        color="red",
+        bg_color="on_grey",
+        styles="underline",
+    )
+    cprint(
+        "Reversed green on red color",
+        color="green",
+        bg_color="on_red",
+        styles="reverse",
+    )

groot/vla/common/utils/misc/__init__.py

@@ -0,0 +1,5 @@
+from .functional_utils import *  # noqa: F403
+from .image_utils import *  # noqa: F403
+from .misc_utils import *  # noqa: F403
+from .torch_utils import *  # noqa: F403
+from .video_utils import *  # noqa: F403

groot/vla/common/utils/misc/array_tensor_utils.py

@@ -0,0 +1,372 @@
+"""
+Functions that work on nested structures of torch.Tensor or numpy array
+"""
+
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import tree
+
+from ..data_structure.tree_utils import (
+    copy_non_leaf,
+    is_sequence,
+    tree_assign_at_path,
+    tree_value_at_path,
+)
+from .functional_utils import make_recursive_func
+
+
+def is_array_tensor(obj):
+    return isinstance(obj, (np.ndarray, torch.Tensor))
+
+
+def is_numpy(obj):
+    return isinstance(obj, np.ndarray)
+
+
+def is_tensor(obj):
+    return torch.is_tensor(obj)
+
+
+def any_stack(xs: List, *, dim: int = 0):
+    """
+    Works for both torch Tensor and numpy array
+    """
+
+    def _any_stack_helper(*xs):
+        x = xs[0]
+        if isinstance(x, np.ndarray):
+            return np.stack(xs, axis=dim)
+        elif torch.is_tensor(x):
+            return torch.stack(xs, dim=dim)
+        elif isinstance(x, float):
+            # special treatment for float, defaults to float32
+            return np.array(xs, dtype=np.float32)
+        else:
+            return np.array(xs)
+
+    return tree.map_structure(_any_stack_helper, *xs)
+
+
+def any_concat(xs: List, *, dim: int = 0):
+    """
+    Works for both torch Tensor and numpy array
+    """
+
+    def _any_concat_helper(*xs):
+        x = xs[0]
+        if isinstance(x, np.ndarray):
+            return np.concatenate(xs, axis=dim)
+        elif torch.is_tensor(x):
+            return torch.cat(xs, dim=dim)
+        elif isinstance(x, float):
+            # special treatment for float, defaults to float32
+            return np.array(xs, dtype=np.float32)
+        else:
+            return np.array(xs)
+
+    return tree.map_structure(_any_concat_helper, *xs)
+
+
+def any_chunk(x, chunks: int, *, dim: int = 0, strict: bool = True) -> List[Any]:
+    """
+    Works for both torch Tensor and numpy array
+
+    Returns:
+        list of chunked nested structures
+    """
+    assert chunks >= 1
+
+    x_copies = [copy_non_leaf(x) for _ in range(chunks)]
+
+    def _any_chunk_helper(path, x):
+        if is_array_tensor(x):
+            if isinstance(x, np.ndarray):
+                chunked_values = np.split(x, chunks, axis=dim)
+            else:
+                chunked_values = torch.chunk(x, chunks, dim=dim)
+
+            if path:
+                for xc, chunked in zip(x_copies, chunked_values):
+                    tree_assign_at_path(xc, path, chunked)
+            else:  # top-level, no nested path
+                for i, chunked in enumerate(chunked_values):
+                    x_copies[i] = chunked
+        else:
+            if strict:
+                raise NotImplementedError(f"Cannot chunk type {type(x)}")
+            else:
+                return
+
+    tree.map_structure_with_path(_any_chunk_helper, x)
+    return x_copies
+
+
+def chunk_seq(arr, chunks: int, check_divide=True):
+    """
+    Args:
+        check_divide: True to force arr must divide n
+    """
+    k, m = divmod(len(arr), chunks)
+    if check_divide and m != 0:
+        raise ValueError(f"Array len {len(arr)} does not divide chunks {chunks}")
+    return (arr[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(chunks))
+
+
+@make_recursive_func
+def any_zeros_like(x: Union[Dict, np.ndarray, torch.Tensor, int, float, np.number]):
+    """Returns a zero-filled object of the same (d)type and shape as the input.
+
+    The difference between this and `np.zeros_like()` is that this works well
+    with `np.number`, `int`, `float`, and `jax.numpy.DeviceArray` objects without
+    converting them to `np.ndarray`s.
+
+    Args:
+      x: The object to replace with 0s.
+
+    Returns:
+      A zero-filed object of the same (d)type and shape as the input.
+    """
+    if isinstance(x, (int, float, np.number)):
+        return type(x)(0)
+    elif is_tensor(x):
+        return torch.zeros_like(x)
+    elif is_numpy(x):
+        return np.zeros_like(x)
+    else:
+        raise ValueError(
+            f"Input ({type(x)}) must be either a numpy array, a tensor, an int, or a float."
+        )
+
+
+@make_recursive_func
+def any_ones_like(x: Union[Dict, np.ndarray, torch.Tensor, int, float, np.number]):
+    """Returns a one-filled object of the same (d)type and shape as the input.
+    The difference between this and `np.ones_like()` is that this works well
+    with `np.number`, `int`, `float`, and `jax.numpy.DeviceArray` objects without
+    converting them to `np.ndarray`s.
+    Args:
+      x: The object to replace with 1s.
+    Returns:
+      A one-filed object of the same (d)type and shape as the input.
+    """
+    if isinstance(x, (int, float, np.number)):
+        return type(x)(1)
+    elif is_tensor(x):
+        return torch.ones_like(x)
+    elif is_numpy(x):
+        return np.ones_like(x)
+    else:
+        raise ValueError(
+            f"Input ({type(x)}) must be either a numpy array, a tensor, an int, or a float."
+        )
+
+
+@make_recursive_func
+def any_zero_(x: Union[Dict, np.ndarray, torch.Tensor]):
+    """
+    Apply in-place zero-out to a tensor, i.e. x.zero_()
+    """
+    if is_tensor(x):
+        x.zero_()
+    elif is_numpy(x):
+        x.fill(0)
+    else:
+        raise ValueError(f"Input ({type(x)}) must be either a numpy array or a tensor")
+
+
+@make_recursive_func
+def any_fill_(x: Union[Dict, np.ndarray, torch.Tensor], value):
+    """
+    Apply in-place zero-out to a tensor, i.e. x.zero_()
+    """
+    if is_tensor(x):
+        x.fill_(value)
+    elif is_numpy(x):
+        x.fill(value)
+    else:
+        raise ValueError(f"Input ({type(x)}) must be either a numpy array or a tensor")
+
+
+def get_batch_size(x, strict: bool = False) -> int:
+    """
+    Args:
+        x: can be any arbitrary nested structure of np array and torch tensor
+        strict: True to check all batch sizes are the same
+    """
+
+    def _get_batch_size(x):
+        if isinstance(x, np.ndarray):
+            return x.shape[0]
+        elif torch.is_tensor(x):
+            return x.size(0)
+        else:
+            return len(x)
+
+    xs = tree.flatten(x)
+
+    if strict:
+        batch_sizes = [_get_batch_size(x) for x in xs]
+        assert all(
+            b == batch_sizes[0] for b in batch_sizes
+        ), f"batch sizes must all be the same in nested structure: {batch_sizes}"
+        return batch_sizes[0]
+    else:
+        return _get_batch_size(xs[0])
+
+
+@make_recursive_func
+def add_batch_dim(x):
+    if is_numpy(x):
+        return np.expand_dims(x, axis=0)
+    elif is_tensor(x):
+        return x.unsqueeze(0)
+    else:
+        raise NotImplementedError(f"Unsupported data structure: {type(x)}")
+
+
+@make_recursive_func
+def remove_batch_dim(x):
+    if is_numpy(x):
+        return np.squeeze(x, axis=0)
+    elif is_tensor(x):
+        return x.squeeze(0)
+    else:
+        raise NotImplementedError(f"Unsupported data structure: {type(x)}")
+
+
+@make_recursive_func
+def any_to_primitive(x):
+    if isinstance(x, (np.ndarray, np.number, torch.Tensor)):
+        return x.tolist()
+    else:
+        return x
+
+
+@make_recursive_func
+def any_get_shape(x):
+    if is_numpy(x):
+        return tuple(x.shape)
+    elif is_tensor(x):
+        return tuple(x.size())
+    else:
+        raise NotImplementedError(f"Unsupported data structure: {type(x)}")
+
+
+@make_recursive_func
+def any_mean(x, dim: Optional[int] = None, keepdim: bool = False):
+    if is_numpy(x):
+        return np.mean(x, axis=dim, keepdims=keepdim)
+    elif is_tensor(x):
+        return torch.mean(x, dim=dim, keepdim=keepdim)
+    else:
+        raise NotImplementedError(f"Unsupported data structure: {type(x)}")
+
+
+@make_recursive_func
+def any_variance(x, dim: Optional[int] = None, keepdim: bool = False, unbiased: bool = False):
+    if is_numpy(x):
+        return np.var(x, axis=dim, keepdims=keepdim, ddof=1 if unbiased else 0)
+    elif is_tensor(x):
+        return torch.var(x, dim=dim, keepdim=keepdim, unbiased=unbiased)
+    else:
+        raise NotImplementedError(f"Unsupported data structure: {type(x)}")
+
+
+@make_recursive_func
+def any_describe_str(x, shape_only=False):
+    """
+    Describe type, shape, device, data type (of np array/tensor)
+    Very useful for debugging
+    """
+    t = type(x)
+    tname = type(x).__name__
+    if is_numpy(x):
+        shape = list(x.shape)
+        if x.size == 1:
+            if shape_only:
+                return f"np scalar: {x.item()} {shape}"
+            else:
+                return f"np scalar: {x.item()} {shape} {x.dtype}"
+        else:
+            if shape_only:
+                return f"np: {shape}"
+            else:
+                return f"np: {shape} {x.dtype}"
+    elif is_tensor(x):
+        shape = list(x.size())
+        if x.numel() == 1:
+            if shape_only:
+                return f"torch scalar: {x.item()} {shape}"
+            else:
+                return f"torch scalar: {x.item()} {shape} {x.dtype} {x.device}"
+        else:
+            if shape_only:
+                return f"torch: {shape}"
+            else:
+                return f"torch: {shape} {x.dtype} {x.device}"
+    elif is_sequence(x):
+        return f"{tname}[{len(x)}]"
+    elif isinstance(x, str):
+        return x
+    elif x is None:
+        return "None"
+    elif np.issubdtype(t, np.number) or np.issubdtype(t, np.bool_):
+        return f"{tname}: {x}"
+    else:
+        return f"{tname}"
+
+
+def any_describe(x, msg="", *, shape_only=False):
+    # from omlet.utils import yaml_dumps
+    from pprint import pprint
+
+    if isinstance(x, str) and msg != "":
+        x, msg = msg, x
+
+    if msg:
+        msg += ": "
+    print(msg, end="")
+    pprint(any_describe_str(x, shape_only=shape_only))
+
+
+@make_recursive_func
+def any_slice(x, slice):
+    """
+    Args:
+        slice: you can use np.s_[...] to return the slice object
+    """
+    if is_array_tensor(x):
+        return x[slice]
+    else:
+        return x
+
+
+def any_assign(x, assign_value, slice):
+    """
+    Recursive version of x[slice] = assign_value
+    If structures of x and assign_value do not match, we will respect `assign_value`
+    E.g. x = {'a': ..., 'b': ...}, assign_value = {'a': ...}, then 'b' will not change
+
+    Use np.s_[...] to get advanced slicing
+    """
+
+    def _any_assign_helper(path, v):
+        y = tree_value_at_path(x, path)
+        y[slice] = v
+
+    tree.map_structure_with_path(_any_assign_helper, assign_value)
+
+
+@make_recursive_func
+def any_transpose_first_two_axes(x):
+    """
+    util to convert between (L, B, ...) and (B, L, ...)
+    """
+    if is_numpy(x):
+        return np.swapaxes(x, 0, 1)
+    elif is_tensor(x):
+        return torch.swapaxes(x, 0, 1)
+    else:
+        raise ValueError(f"Input ({type(x)}) must be either a numpy array or a tensor.")

groot/vla/common/utils/misc/functional_utils.py

@@ -0,0 +1,635 @@
+"""
+Inspect, meta, etc.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import pprint
+import sys
+import types
+from typing import Any, Dict, Literal
+import warnings
+
+from ..data_structure.tree_utils import is_mapping, is_sequence
+
+
+def state_dict_class(keys: list[str]):
+    """
+    Just like pytorch nn.Module
+    Add the following methods to the class:
+        state_dict() -> dict of attribute keys
+        load_state_dict(sdict)  restore states
+    """
+
+    def _wrap_class(cls):
+        assert inspect.isclass(cls)
+
+        def state_dict(self):
+            return {k: getattr(self, k) for k in keys}
+
+        def load_state_dict(self, states: Dict[str, Any]):
+            if not set(keys).issubset(set(states.keys())):
+                raise ValueError(f"states does not have all the required keys: {keys}")
+            for k in keys:
+                setattr(self, k, states[k])
+
+        @property
+        def state_keys(self):
+            return keys
+
+        cls.state_dict = state_dict
+        cls.load_state_dict = load_state_dict
+        cls.state_keys = state_keys
+        return cls
+
+    return _wrap_class
+
+
+def implements_method(object, method: str):
+    """
+    Returns:
+        True if object implements a method
+    """
+    return hasattr(object, method) and callable(getattr(object, method))
+
+
+def assert_implements_method(object, method: str | list[str]):
+    if isinstance(method, str):
+        method = [method]
+    for m in method:
+        assert implements_method(object, m), (
+            f"object {object.__class__} does not " f"implement method {m}()"
+        )
+
+
+def meta_decorator(decor):
+    """
+    a decorator, allowing the wrapped decorator to be used as:
+    @decorator(*args, **kwargs)
+    def callable()
+      -- or --
+    @decorator  # without parenthesis, args and kwargs will use default
+    def callable()
+
+    Args:
+      decor: a decorator whose first argument is a callable (function or class
+        to be decorated), and the rest of the arguments can be omitted as default.
+        decor(f, ... the other arguments must have default values)
+
+    Warning:
+      decor can NOT be a function that receives a single, callable argument.
+      See stackoverflow: http://goo.gl/UEYbDB
+    """
+    import functools
+
+    def single_callable(args, kwargs):
+        return len(args) == 1 and len(kwargs) == 0 and callable(args[0])
+
+    @functools.wraps(decor)
+    def new_decor(*args, **kwargs):
+        if single_callable(args, kwargs):
+            # this is the double-decorated f.
+            # It should not run on a single callable.
+            return decor(args[0])
+        else:
+            # decorator arguments
+            return lambda real_f: decor(real_f, *args, **kwargs)
+
+    return new_decor
+
+
+@meta_decorator
+def make_recursive_func(fn, *, with_path=False):
+    """
+    Decorator that turns a function that works on a single array/tensor to working on
+    arbitrary nested structures.
+    """
+    import functools
+
+    import tree
+
+    @functools.wraps(fn)
+    def _wrapper(tensor_struct, *args, **kwargs):
+        if with_path:
+            return tree.map_structure_with_path(
+                lambda paths, x: fn(paths, x, *args, **kwargs), tensor_struct
+            )
+        else:
+            return tree.map_structure(lambda x: fn(x, *args, **kwargs), tensor_struct)
+
+    return _wrapper
+
+
+@meta_decorator
+def deprecated(func, msg="", action="warning", type=""):
+    """
+    Function/class decorator: designate deprecation.
+
+    Args:
+      msg: string message.
+      action: string mode
+      - 'warning': (default) prints `msg` to stderr
+      - 'noop': do nothing, just for source code annotation purposes
+      - 'raise': raise DeprecatedError(`msg`)
+    """
+    action = action.lower()
+    type = type.lower()
+    ALL_ACTIONS = ["warn", "warning", "noop", "raise"]
+    if action not in ALL_ACTIONS:
+        raise ValueError(f"Unknown action {action}. Choose from {ALL_ACTIONS}.")
+    ALL_TYPES = {
+        "": DeprecationWarning,
+        "pending": PendingDeprecationWarning,
+        "future": FutureWarning,
+    }
+    if type not in ALL_TYPES:
+        raise ValueError(f"Unknown type {type}. Choose from {ALL_TYPES.keys()}.")
+    if not msg:
+        msg = "This is a deprecated feature."
+
+    WarningExceptionCls = ALL_TYPES[type]
+
+    # only does the deprecation when being called
+    @functools.wraps(func)
+    def _deprecated(*args, **kwargs):
+        if action in ["warning", "warn"]:
+            warnings.warn(msg, WarningExceptionCls)
+        elif action == "raise":
+            raise WarningExceptionCls(msg)
+        return func(*args, **kwargs)
+
+    return _deprecated
+
+
+@meta_decorator
+def call_once(func, on_second_call: Literal["noop", "raise", "warn"] = "noop"):
+    """
+    Decorator to ensure that a function is only called once.
+
+    Args:
+      on_second_call (str): what happens when the function is called a second time.
+    """
+    assert on_second_call in [
+        "noop",
+        "raise",
+        "warn",
+    ], "mode must be one of 'noop', 'raise', 'warn'"
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if wrapper._called:
+            if on_second_call == "raise":
+                raise RuntimeError(f"{func.__name__} has already been called. Can only call once.")
+            elif on_second_call == "warn":
+                warnings.warn(f"{func.__name__} has already been called. Should only call once.")
+        else:
+            wrapper._called = True
+            return func(*args, **kwargs)
+
+    wrapper._called = False
+    return wrapper
+
+
+class NoopObject:
+    """
+    Object that does nothing when called any method
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.init_args = args
+        self.init_kwargs = kwargs
+
+    def __getattr__(self, name):
+        def _func(*args, **kwargs):
+            pass
+
+        return _func
+
+
+class NoopContext:
+    """
+    Placeholder context manager that does nothing.
+    We could have written simply as:
+
+    @contextmanager
+    def noop_context(*args, **kwargs):
+        yield
+
+    but the returned context manager cannot be called twice, i.e.
+    my_noop = NoopContext()
+    with my_noop:
+        do1()
+    with my_noop: # trigger generator error
+        do2()
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
+def make_registry_metaclass(class_name):
+    """
+    Usage:
+
+      TrainerRegistry = make_registry_metaclass('TrainerRegistry')
+
+      class BaseTrainer(metaclass=TrainerRegistry):
+          pass
+
+      class MyTrainer(BaseTrainer):
+          pass
+
+      TrainerRegistry['MyTrainer'] -> MyTrainer class  # syntax enabled by metaclass
+      TrainerRegistry.get_class('MyTrainer')  # same as above
+      TrainerRegistry.registry -> full dict of {name: trainer_class}
+
+    Templated definition:
+        class TrainerRegistry(type):
+            registry = {}
+
+            def __new__(cls, name, bases, attr):
+                new_cls = super().__new__(cls, name, bases, attr)
+                TrainerRegistry.registry[name] = new_cls
+                return new_cls
+
+            def get_class(cls, name):
+                if name not in cls.registry:
+                    raise KeyError(
+                        f"Trainer class {name} not found in registry: "
+                        f"{pprint.pformat(cls.registry)}"
+                    )
+                return cls.registry[name]"""
+
+    def new__(cls, name, bases, attr):
+        """
+        Change the attr dict to dynamically add methods and attributes
+        """
+        new_cls = type.__new__(cls, name, bases, attr)
+        cls.registry[name] = new_cls
+        return new_cls
+
+    def get_class(cls, name):
+        if name not in cls.registry:
+            existing_cls = list(cls.registry.keys())
+            raise KeyError(f"{class_name} class '{name}' not found in registry: {existing_cls}")
+        return cls.registry[name]
+
+    def instantiate(cls_, cls, **kwargs):
+        Cls = cls_.get_class(cls)
+        return Cls(**kwargs)
+
+    class _BracketOperator(type):
+        def __getitem__(cls, name):
+            return get_class(cls, name)
+
+    return types.new_class(
+        class_name,
+        bases=(type,),
+        kwds={"metaclass": _BracketOperator},
+        exec_body=lambda ns: ns.update(
+            {
+                "registry": {},
+                "__new__": new__,
+                "get_class": classmethod(get_class),
+                "instantiate": classmethod(instantiate),
+            }
+        ),
+    )
+
+
+class ClassRegistry:
+    """
+    May be a preferred way over make_registry_metaclass if your code does not support
+    metaclass well, e.g. pickle or Ray
+
+    Use in conjunction with `__init_subclass__` hook in your base class
+
+    class BaseClass:
+        registry = ClassRegistry()
+
+        def __init_subclass__(cls, **kwargs):
+            cls.registry.add(cls)
+            super().__init_subclass__(**kwargs)
+
+    print(registry)
+    """
+
+    def __init__(self, base_class_name: str = None):
+        self.registry = {}
+        self._base_class_name = base_class_name
+
+    def add(self, cls):
+        self.registry[cls.__name__] = cls
+
+    def get(self, name):
+        if name not in self.registry:
+            existing_cls = list(self.registry.keys())
+            base_name = self._base_class_name + " " if self._base_class_name else ""
+            raise KeyError(f"{base_name} subclass '{name}' not found in registry: {existing_cls}")
+        return self.registry[name]
+
+    def __str__(self):
+        return pprint.pformat(self.registry)
+
+    def __getitem__(self, name):
+        return self.get(name)
+
+    def instantiate(self, cls, **kwargs):
+        return self.get(cls)(**kwargs)
+
+
+# ========================================================
+# =================== Inspect utils ====================
+# ========================================================
+
+
+def func_parameters(func):
+    return inspect.signature(func).parameters
+
+
+def func_has_arg(func, arg_name):
+    return arg_name in func_parameters(func)
+
+
+def pack_varargs(args):
+    """
+    Pack *args or a single list arg as list
+
+    def f(*args):
+        arg_list = pack_varargs(args)
+        # arg_list is now packed as a list
+    """
+    assert isinstance(args, tuple), "please input the tuple `args` as in *args"
+    if len(args) == 1 and is_sequence(args[0]):
+        return args[0]
+    else:
+        return args
+
+
+def enable_list_arg(func):
+    """
+    Function decorator.
+    If a function only accepts varargs (*args),
+    make it support a single list arg as well
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        args = pack_varargs(args)
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def enable_varargs(func):
+    """
+    Function decorator.
+    If a function only accepts a list arg, make it support varargs as well
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        args = pack_varargs(args)
+        return func(args, **kwargs)
+
+    return wrapper
+
+
+def pack_kwargs(args, kwargs):
+    """
+    Pack **kwargs or a single dict arg as dict
+
+    def f(*args, **kwargs):
+        kwdict = pack_kwargs(args, kwargs)
+        # kwdict is now packed as a dict
+    """
+    if len(args) == 1 and is_mapping(args[0]):
+        assert not kwargs, "cannot have both **kwargs and a dict arg"
+        return args[0]  # single-dict
+    else:
+        assert not args, "cannot have positional args if **kwargs exist"
+        return kwargs
+
+
+def merge_kwargs(args, kwargs) -> Dict:
+    """
+    Merge all dicts in `args` and keywords in kwargs.
+
+    E.g. merge_kwargs({"a.b": 1, "a.c": 2}, foo=6, bar=8)
+       -> {"a.b": 1, "a.c": 2, "foo": 6, "bar": 8}
+    """
+    kw_all = {}
+    for arg in args:
+        assert is_mapping(arg), f"{arg} is not a dict."
+        kw_all.update(arg)
+    kw_all.update(kwargs)
+    return kw_all
+
+
+def enable_dict_arg(func):
+    """
+    Function decorator.
+    If a function only accepts varargs (*args),
+    make it support a single list arg as well
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        kwargs = pack_kwargs(args, kwargs)
+        return func(**kwargs)
+
+    return wrapper
+
+
+def enable_kwargs(func):
+    """
+    Function decorator.
+    If a function only accepts a dict arg, make it support kwargs as well
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        kwargs = pack_kwargs(args, kwargs)
+        return func(kwargs)
+
+    return wrapper
+
+
+def has_keys(D, keys: list):
+    assert is_mapping(D)
+    return all(key in D for key in keys)
+
+
+def assert_has_keys(D, keys: list):
+    assert is_mapping(D), "Input is not a dict"
+    for key in keys:
+        if key not in D:
+            raise KeyError(f'Required key "{key}" is missing in dict {D}')
+    return True
+
+
+def method_decorator(decorator):
+    """
+    Decorator of decorator: transform a decorator that only works on normal
+    functions to a decorator that works on class methods
+    From Django form: https://goo.gl/XLjxKK
+    """
+
+    @functools.wraps(decorator)
+    def wrapped_decorator(method):
+        @functools.wraps(method)
+        def wrapper(self, *args, **kwargs):
+            def bound_func(*args2, **kwargs2):
+                return method(self, *args2, **kwargs2)
+
+            return decorator(bound_func)(*args, **kwargs)
+
+        return wrapper
+
+    return wrapped_decorator
+
+
+def accepts_varargs(func):
+    """
+    If a function accepts *args
+    """
+    params = inspect.signature(func).parameters
+    return any(param.kind == inspect.Parameter.VAR_POSITIONAL for param in params.values())
+
+
+def accepts_kwargs(func):
+    """
+    If a function accepts **kwargs
+    """
+    params = inspect.signature(func).parameters
+    return any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values())
+
+
+def is_signature_compatible(func, *args, **kwargs):
+    sig = inspect.signature(func)
+    try:
+        sig.bind(*args, **kwargs)
+        return True
+    except TypeError:
+        return False
+
+
+def make_list(x):
+    """
+    Turns a singleton object to a list. If already a list, no change.
+    """
+    if is_sequence(x):
+        return x
+    else:
+        return [x]
+
+
+def make_tuple(elem, repeats):
+    """
+    E.g. expand a singleton x into (x, x, x)
+        useful for things like image_size or kernal, which can be a single int/float
+        or a tuple of fixed size
+    """
+    if is_sequence(elem):
+        assert len(elem) == repeats, f"length of input must be {repeats}: {elem}"
+        return elem
+    else:
+        return (elem,) * repeats
+
+
+def accumulate(iterable, fn=lambda x, y: x + y):
+    """
+    Return running totals
+    # _accumulate([1,2,3,4,5]) --> 1 3 6 10 15
+    # _accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
+    """
+    it = iter(iterable)
+    try:
+        total = next(it)
+    except StopIteration:
+        return
+    yield total
+    for element in it:
+        total = fn(total, element)
+        yield total
+
+
+class DecoratorContextManager:
+    """
+    Allow a context manager to be used as a decorator
+    From torch.auto_grad.grad_mode
+    """
+
+    def __call__(self, func):
+        if inspect.isgeneratorfunction(func):
+            return self._wrap_generator(func)
+
+        @functools.wraps(func)
+        def decorate_context(*args, **kwargs):
+            with self.__class__():
+                return func(*args, **kwargs)
+
+        return decorate_context
+
+    def _wrap_generator(self, func):
+        """Wrap each generator invocation with the context manager"""
+
+        @functools.wraps(func)
+        def generator_context(*args, **kwargs):
+            gen = func(*args, **kwargs)
+
+            # Generators are suspended and unsuspended at `yield`, hence we
+            # make sure the grad mode is properly set every time the execution
+            # flow returns into the wrapped generator and restored when it
+            # returns through our `yield` to our caller (see PR #49017).
+            cls = type(self)
+            try:
+                # Issuing `None` to a generator fires it up
+                with cls():
+                    response = gen.send(None)
+
+                while True:
+                    try:
+                        # Forward the response to our caller and get its next request
+                        request = yield response
+
+                    except GeneratorExit:
+                        # Inform the still active generator about its imminent closure
+                        with cls():
+                            gen.close()
+                        raise
+
+                    except BaseException:
+                        # Propagate the exception thrown at us by the caller
+                        with cls():
+                            response = gen.throw(*sys.exc_info())
+
+                    else:
+                        # Pass the last request to the generator and get its response
+                        with cls():
+                            response = gen.send(request)
+
+            # We let the exceptions raised above by the generator's `.throw` or
+            # `.send` methods bubble up to our caller, except for StopIteration
+            except StopIteration as e:
+                # The generator informed us that it is done: take whatever its
+                # returned value (if any) was and indicate that we're done too
+                # by returning it (see docs for python's return-statement).
+                return e.value
+
+        return generator_context
+
+    def __enter__(self) -> None:
+        raise NotImplementedError
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        raise NotImplementedError

groot/vla/common/utils/misc/image_utils.py

@@ -0,0 +1,225 @@
+"""
+Visualizations
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from typing import Literal
+import warnings
+
+import cv2
+import imageio
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from .array_tensor_utils import any_describe
+from .misc_utils import global_once
+from .torch_utils import torch_normalize
+
+
+def to_image(img, channel_order="auto"):
+    """
+    Returns:
+        numpy image of shape [H, W, C]
+        in "auto" mode, we assume C == 3
+    """
+    assert channel_order in ["hwc", "chw", "auto"]
+    if torch.is_tensor(img):
+        img = img.cpu().numpy()
+    assert isinstance(img, np.ndarray)
+    if img.ndim == 4:
+        assert img.shape[0] == 1
+        img = img[0]
+    assert img.ndim == 3
+    if channel_order == "auto":
+        # use C==3 to detect order
+        if img.shape[0] == 3:
+            channel_order = "chw"
+        else:
+            assert img.shape[-1] == 3, "image should either have [3,H,W] or [H,W,3]"
+            channel_order = "hwc"
+    img = img.astype(np.uint8)
+    if channel_order == "chw":
+        return np.transpose(img, (1, 2, 0))
+    else:
+        return img
+
+
+def imshow(img):
+    plt.imshow(to_image(img))
+
+
+def imsave(img, path):
+    imageio.imsave(os.path.expanduser(path), to_image(img))
+
+
+def imread(path, channel_order="chw", format="torch"):
+    assert channel_order in ["hwc", "chw"]
+    assert format in ["numpy", "torch"]
+    img = imageio.imread(path)
+    if channel_order == "chw":
+        img = np.transpose(img, (2, 0, 1))  # hwc -> chw
+    if format == "torch":
+        return torch.from_numpy(img)
+    else:
+        return img
+
+
+class Cv2Display:
+    def __init__(
+        self,
+        window_name="display",
+        image_size=None,
+        channel_order="auto",
+        bgr2rgb=True,
+        step_sleep=0,
+        enabled=True,
+    ):
+        """
+        Use cv2.imshow() to pop a window, requires virtual desktop GUI
+
+        Args:
+            channel_order: auto, hwc, or chw
+            image_size: None to use the original image size, otherwise resize
+            step_sleep: sleep for a few seconds
+        """
+        self._window_name = window_name
+        if isinstance(image_size, int):
+            image_size = (image_size, image_size)
+        else:
+            assert image_size is None or len(image_size) == 2
+        self._image_size = image_size
+        assert channel_order in ["auto", "chw", "hwc"]
+        self._channel_order = channel_order
+        self._bgr2rgb = bgr2rgb
+        self._step_sleep = step_sleep
+        self._enabled = enabled
+
+    def _resize(self, img):
+        if self._image_size is None:
+            return img
+        H, W = img.shape[:2]
+        Ht, Wt = self._image_size  # target
+        return cv2.resize(
+            img,
+            self._image_size,
+            interpolation=cv2.INTER_AREA if Ht < H else cv2.INTER_LINEAR,
+        )
+
+    def _reorder(self, img):
+        if self._channel_order == "chw":
+            return np.transpose(img, (1, 2, 0))
+        elif self._channel_order == "hwc":
+            return img
+        else:
+            if img.shape[0] in [1, 3]:  # chw
+                return np.transpose(img, (1, 2, 0))
+            else:
+                return img
+
+    def __call__(self, img):
+        if not self._enabled:
+            return
+        import torch
+
+        # prevent segfault in IsaacGym
+        display_var = os.environ.get("DISPLAY", None)
+        if not display_var:
+            os.environ["DISPLAY"] = ":0.0"
+
+        if torch.is_tensor(img):
+            img = img.detach().cpu().numpy()
+
+        img = self._resize(self._reorder(img))
+        if self._bgr2rgb:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        time.sleep(self._step_sleep)
+        cv2.imshow(self._window_name, img)
+        cv2.waitKey(1)
+
+        if display_var is not None:
+            os.environ["DISPLAY"] = display_var
+
+    def close(self):
+        if not self._enabled:
+            return
+        cv2.destroyWindow(self._window_name)
+
+
+# ---------------- Image tensor handling -----------------
+def sanity_check_image_tensor(
+    img: torch.Tensor, on_error: Literal["raise", "warn", "ignore"] = "raise"
+):
+    """
+    Check if the input image tensor is all integers, which is wrong for any NN input.
+    This is a common case if the user forgets to normalize the image first
+    """
+    assert on_error in [
+        "raise",
+        "warn",
+        "ignore",
+    ], 'on_error must be "raise", "warn", or "ignore"'
+    if not img.dtype.is_floating_point:
+        msg = f"Image tensor is not floating point format, but {img.dtype}!"
+        if on_error == "raise":
+            raise ValueError(msg)
+        elif on_error == "warn":
+            warnings.warn(msg)
+        else:
+            return False
+    # check if all values in the image are close to an integer
+    if (img - torch.round(img)).abs().max() < 1e-5:
+        msg = (
+            "Input image is all close to integers, "
+            "are you sure you have normalized it before passing it to a NN?"
+        )
+        if on_error == "raise":
+            raise ValueError(msg)
+        elif on_error == "warn":
+            warnings.warn(msg)
+        else:
+            return False
+    return True
+
+
+@torch.no_grad()
+def basic_image_tensor_preprocess(
+    img,
+    mean: tuple[float, float, float] = (0.5, 0.5, 0.5),
+    std: tuple[float, float, float] = (0.5, 0.5, 0.5),
+    shape: tuple[int, int] | None = None,
+):
+    """
+    Check for resize, and divide by 255
+    """
+    import kornia
+
+    assert torch.is_tensor(img)
+    assert img.dim() >= 4, any_describe(img)
+    original_shape = list(img.size())
+    img = img.float()
+    img = img.flatten(0, img.dim() - 4)
+    assert img.dim() == 4
+
+    input_size = img.size()[-2:]
+    if global_once("groot.vla.common.utils.image_utils.basic_image_preprocess:input_size"):
+        assert img.max() > 2, "img should be between [0, 255] before normalize"
+
+    if shape and input_size != shape:
+        if global_once("groot.vla.common.utils.image_utils.basic_image_preprocess:transform"):
+            warnings.warn(
+                f'{"Down" if shape < input_size else "Up"}sampling image'
+                f" from original resolution {input_size}x{input_size}"
+                f" to {shape}x{shape}"
+            )
+        img = kornia.geometry.transform.resize(img, shape).clamp(0.0, 255.0)
+
+    B, C, H, W = img.size()
+    assert C % 3 == 0, "channel must divide 3"
+    img = img.view(B * C // 3, 3, H, W)
+    img = torch_normalize(img / 255.0, mean=mean, std=std)
+    original_shape[-2:] = H, W
+    return img.view(original_shape)

groot/vla/common/utils/misc/misc_utils.py

@@ -0,0 +1,261 @@
+import codecs
+from collections import Counter
+import fnmatch
+import hashlib
+import os
+import pickle
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from typing_extensions import Literal
+
+
+def set_os_envs(envs: Optional[Dict[str, Any]] = None):
+    """
+    Special value __delete__ or None indicates that the ENV_VAR should be removed
+    """
+    if envs is None:
+        envs = {}
+    DEL = {None, "__delete__"}
+    for k, v in envs.items():
+        if v in DEL:
+            os.environ.pop(k, None)
+    os.environ.update({k: str(v) for k, v in envs.items() if v not in DEL})
+
+
+def argmax(L):
+    return max(zip(L, range(len(L))))[1]
+
+
+def _match_patterns_helper(element, patterns):
+    for p in patterns:
+        if callable(p) and p(element):
+            return True
+        if fnmatch.fnmatch(element, p):
+            return True
+    return False
+
+
+def match_patterns(
+    item: str,
+    include: Union[str, List[str], Callable, List[Callable], None] = None,
+    exclude: Union[str, List[str], Callable, List[Callable], None] = None,
+    *,
+    precedence: Literal["include", "exclude"] = "exclude",
+):
+    """
+    Args:
+        include: None to disable `include` filter and delegate to exclude
+        precedence: "include" or "exclude"
+    """
+    assert precedence in ["include", "exclude"]
+    if exclude is None:
+        exclude = []
+    if isinstance(exclude, (str, Callable)):
+        exclude = [exclude]
+    if isinstance(include, (str, Callable)):
+        include = [include]
+    if include is None:
+        # exclude is the sole veto vote
+        return not _match_patterns_helper(item, exclude)
+
+    if precedence == "include":
+        return _match_patterns_helper(item, include)
+    else:
+        if _match_patterns_helper(item, exclude):
+            return False
+        else:
+            return _match_patterns_helper(item, include)
+
+
+def filter_patterns(
+    items: List[str],
+    include: Union[str, List[str], Callable, List[Callable], None] = None,
+    exclude: Union[str, List[str], Callable, List[Callable], None] = None,
+    *,
+    precedence: Literal["include", "exclude"] = "exclude",
+    ordering: Literal["original", "include"] = "original",
+):
+    """
+    Args:
+        ordering: affects the order of items in the returned list. Does not affect the
+            content of the returned list.
+            - "original": keep the ordering of items in the input list
+            - "include": order items by the order of include patterns
+    """
+    assert ordering in ["original", "include"]
+    if include is None or isinstance(include, str) or ordering == "original":
+        return [
+            x
+            for x in items
+            if match_patterns(x, include=include, exclude=exclude, precedence=precedence)
+        ]
+    else:
+        items = items.copy()
+        ret = []
+        for inc in include:
+            for i, item in enumerate(items):
+                if item is None:
+                    continue
+                if match_patterns(item, include=inc, exclude=exclude, precedence=precedence):
+                    ret.append(item)
+                    items[i] = None
+        return ret
+
+
+def getitem_nested(cfg, key: str):
+    """
+    Recursively get key, if key has '.' in it
+    """
+    keys = key.split(".")
+    for k in keys:
+        assert k in cfg, f'{k} in key "{key}" does not exist in config'
+        cfg = cfg[k]
+    return cfg
+
+
+def setitem_nested(cfg, key: str, value):
+    """
+    Recursively get key, if key has '.' in it
+    """
+    keys = key.split(".")
+    for k in keys[:-1]:
+        assert k in cfg, f'{k} in key "{key}" does not exist in config'
+        cfg = cfg[k]
+    cfg[keys[-1]] = value
+
+
+def getattr_nested(obj, key: str):
+    """
+    Recursively get attribute
+    """
+    keys = key.split(".")
+    for k in keys:
+        assert hasattr(obj, k), f'{k} in attribute "{key}" does not exist'
+        obj = getattr(obj, k)
+    return obj
+
+
+def setattr_nested(obj, key: str, value):
+    """
+    Recursively set attribute
+    """
+    keys = key.split(".")
+    for k in keys[:-1]:
+        assert hasattr(obj, k), f'{k} in attribute "{key}" does not exist'
+        obj = getattr(obj, k)
+    setattr(obj, keys[-1], value)
+
+
+class PeriodicEvent:
+    """
+    triggers every period
+    """
+
+    def __init__(self, period: int, initial_value=0):
+        self._period = period
+        assert self._period >= 1
+        self._last_threshold = initial_value
+        self._last_value = initial_value
+        self._trigger_counts = 0
+
+    def __call__(self, new_value=None, increment=None):
+        assert bool(new_value is None) != bool(increment is None), (
+            "you must specify one and only one of new_value or increment, " "but not both"
+        )
+        d = self._period
+        if new_value is None:
+            new_value = self._last_value + increment
+        assert new_value >= self._last_value, (
+            f"value must be monotonically increasing. "
+            f"Current value {new_value} < last value {self._last_value}"
+        )
+        self._last_value = new_value
+        if new_value - self._last_threshold >= d:
+            self._last_threshold += (new_value - self._last_threshold) // d * d
+            self._trigger_counts += 1
+            return True
+        else:
+            return False
+
+    @property
+    def trigger_counts(self):
+        return self._trigger_counts
+
+    @property
+    def current_value(self):
+        return self._last_value
+
+
+class Once:
+    def __init__(self):
+        self._triggered = False
+
+    def __call__(self):
+        if not self._triggered:
+            self._triggered = True
+            return True
+        else:
+            return False
+
+    def __bool__(self):
+        raise RuntimeError("`Once` objects should be used by calling ()")
+
+
+_GLOBAL_ONCE_SET = set()
+_GLOBAL_NTIMES_COUNTER = Counter()
+
+
+def global_once(name):
+    """
+    Try this to automate the name:
+    https://gist.github.com/techtonik/2151727#gistcomment-2333747
+    """
+    if name in _GLOBAL_ONCE_SET:
+        return False
+    else:
+        _GLOBAL_ONCE_SET.add(name)
+        return True
+
+
+def global_n_times(name, n: int):
+    """
+    Triggers N times
+    """
+    assert n >= 1
+    if _GLOBAL_NTIMES_COUNTER[name] < n:
+        _GLOBAL_NTIMES_COUNTER[name] += 1
+        return True
+    else:
+        return False
+
+
+class Every:
+    def __init__(self, n: int, on_first: bool = False):
+        assert n > 0
+        self._i = 0 if on_first else 1
+        self._n = n
+
+    def __call__(self):
+        return self._i % self._n == 0
+
+    def __bool__(self):
+        raise RuntimeError("`Every` objects should be used by calling ()")
+
+
+def encode_base64(obj) -> str:
+    return codecs.encode(pickle.dumps(obj), "base64").decode()
+
+
+def decode_base64(s: str):
+    return pickle.loads(codecs.decode(s.encode(), "base64"))
+
+
+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

groot/vla/common/utils/misc/torch_utils.py

@@ -0,0 +1,748 @@
+from __future__ import annotations
+
+from copy import deepcopy
+import os
+import random
+import time
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import tree
+from typing_extensions import Literal
+
+from ..data_structure.tree_utils import tree_value_at_path
+from ..io.file_utils import f_join
+from ..io.print_utils import to_readable_count_str
+from .functional_utils import assert_implements_method, implements_method
+
+
+def weight_init(m):
+    """Custom weight init for Conv2D and Linear layers."""
+    if isinstance(m, nn.Linear):
+        nn.init.orthogonal_(m.weight.data)
+        if hasattr(m.bias, "data"):
+            m.bias.data.fill_(0.0)
+    elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+        gain = nn.init.calculate_gain("relu")
+        nn.init.orthogonal_(m.weight.data, gain)
+        if hasattr(m.bias, "data"):
+            m.bias.data.fill_(0.0)
+
+
+def get_seed(
+    seed: Union[int, str, None],
+    handle_invalid_seed: Literal["none", "system", "raise"] = "none",
+) -> Optional[int]:
+    """
+    Args:
+      seed:
+        "system": use scrambled int based on system time
+        None or int < 0: invalid seed values, see `handle_invalid_seed`
+        int >= 0: returns seed
+      handle_invalid_seed: None or int < 0
+        - "none": returns None
+        - "system": returns scrambled int based on system time
+        - "raise": raise Exception
+    """
+    handle_invalid_seed = handle_invalid_seed.lower()
+    assert handle_invalid_seed in ["none", "system", "raise"]
+    if isinstance(seed, str):
+        assert seed in ["system"]
+        invalid = False
+    else:
+        assert seed is None or isinstance(seed, int)
+        invalid = seed is None or seed < 0
+
+    if seed == "system" or invalid and handle_invalid_seed == "system":
+        # https://stackoverflow.com/questions/27276135/python-random-system-time-seed
+        t = int(time.time() * 100000)
+        return (
+            ((t & 0xFF000000) >> 24)
+            + ((t & 0x00FF0000) >> 8)
+            + ((t & 0x0000FF00) << 8)
+            + ((t & 0x000000FF) << 24)
+        )
+    elif invalid:
+        if handle_invalid_seed == "none":
+            return None
+        elif handle_invalid_seed == "raise":
+            raise ValueError(
+                f"Invalid random seed: {seed}, " f'must be a non-negative integer or "system"'
+            )
+        else:
+            raise NotImplementedError
+    else:
+        return seed
+
+
+def set_deterministic(flag: bool = True):
+    if not flag:
+        return
+
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+    os.environ["HOROVOD_FUSION_THRESHOLD"] = "0"
+    import torch.backends.cudnn as cudnn
+
+    cudnn.deterministic = True
+    cudnn.benchmark = False
+    if hasattr(torch, "use_deterministic_algorithms"):
+        # only available in PyTorch >= 1.9
+        torch.use_deterministic_algorithms(True)
+    elif hasattr(torch, "set_deterministic"):
+        # only available in PyTorch >= 1.7
+        torch.set_deterministic(True)
+
+
+def set_seed_everywhere(
+    seed: Optional[Union[int, str]],
+    deterministic=False,
+    set_tensorflow=False,
+    handle_invalid_seed: Literal["none", "system", "raise"] = "none",
+) -> Optional[int]:
+    """
+    References:
+    - https://github.com/NVIDIA/framework-determinism/blob/master/pytorch.md
+    - https://pytorch.org/docs/stable/notes/randomness.html
+    - CUBLAS env var:
+        https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
+
+    Args:
+        seed: see `get_seed()`
+        handle_invalid_seed: see `get_seed()`
+    """
+    set_deterministic(deterministic)
+
+    seed = get_seed(seed, handle_invalid_seed=handle_invalid_seed)
+    if seed is None:
+        return None
+
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    if set_tensorflow:
+        try:
+            import tensorflow as tf
+
+            tf.random.set_seed(seed)
+        except ImportError:
+            pass
+    return seed
+
+
+class eval_mode(object):
+    def __init__(self, *models):
+        self.models = models
+
+    def __enter__(self):
+        self.prev_states = []
+        for model in self.models:
+            self.prev_states.append(model.training)
+            model.train(False)
+
+    def __exit__(self, *args):
+        for model, state in zip(self.models, self.prev_states):
+            model.train(state)
+        return False
+
+
+def get_device(x, strict: bool = False) -> int:
+    """
+    Args:
+        x: can be any arbitrary nested structure of np array and torch tensor
+        strict: True to check all batch sizes are the same
+    """
+    xs = tree.flatten(x)
+
+    def _get_device(x):
+        if torch.is_tensor(x):
+            return x.device
+        elif isinstance(x, nn.Module):
+            return get_module_device(x)
+        else:
+            return None
+
+    if strict:
+        devices = [_get_device(x) for x in xs]
+        assert all(
+            b == devices[0] for b in devices
+        ), f"devices must all be the same in nested structure: {devices}"
+        return devices[0]
+    else:
+        return _get_device(xs[0])
+
+
+def load_torch(*fpath: str, map_location="cpu") -> dict:
+    """
+    Default maps to "cpu"
+    """
+    fpath = str(f_join(fpath))
+    try:
+        return torch.load(fpath, map_location=map_location)
+    except RuntimeError as e:
+        raise RuntimeError(f"{e}\n\n --- Error loading {fpath}")
+
+
+def save_torch(D, *fpath):
+    """
+    Supports both (D, fpath) and (fpath, D) arg order, as long as one of them is a str
+    """
+    if isinstance(D, str):
+        assert not isinstance(fpath, str), "Either torch_save(D, fpath) " "or torch_save(fpath, D)"
+        fpath, D = D, fpath
+    torch.save(D, str(f_join(fpath)))
+
+
+# Aliases for consistency with load_pickle, load_text, load_json/yaml, etc.
+torch_load = load_torch
+torch_save = save_torch
+dump_torch = save_torch
+
+
+def torch_compute_stats(x, precision: int = 2):
+    x = x.to(dtype=torch.float32)
+    return (
+        f"mean|std: {torch.mean(x):.{precision}f} +/- {torch.std(x):.{precision}f}, "
+        f"median: {torch.median(x):.{precision}f}, "
+        f"max: {torch.max(x):.{precision}f}, min: {torch.min(x):.{precision}f}"
+    )
+
+
+def tensor_hash(x: torch.Tensor, mode: str = "mean"):
+    if isinstance(x, np.ndarray):
+        x = torch.from_numpy(x)
+    x = x.float().abs()
+    if mode == "sum":
+        x = x.sum()
+    elif mode == "mean":
+        x = x.mean()
+    else:
+        raise NotImplementedError
+    return float(x)
+
+
+def torch_flatten_indices(indices: torch.Tensor, shape: Tuple[int]):
+    """
+    Convert M dim indices to 1D indices with the given shape
+
+    Args:
+        indices: BxM, batch_size x M-dimensional
+    """
+    offsets = np.array(shape)  # e.g. [3, 4, 5, 6]
+    offsets = np.append(offsets[1:], 1)  # [4, 5, 6, 1]
+    offsets = np.cumprod(offsets[::-1])[::-1]  # [4*5*6, 5*6, 6, 1]
+    offsets = torch.tensor(offsets.copy(), dtype=torch.long)
+    assert offsets.size() == (len(shape),)
+    return (indices * offsets.to(device=indices.device)).sum(dim=1)
+
+
+def torch_multi_index_select(x: torch.Tensor, indices: torch.Tensor):
+    """
+    Args:
+        x: N dim
+        indices: [B x M], M <= N, will select the first M-D from N-D
+
+    Returns:
+        (N - M + 1) dim
+    """
+    assert indices.ndim == 2
+    B, idx_dim = indices.size()
+    x_shape = x.size()
+    assert len(x_shape) >= idx_dim
+    remainder_dim = len(x_shape) - idx_dim
+    if remainder_dim == 0:
+        x = torch.flatten(x)
+    else:
+        x = x.view(-1, *x_shape[-remainder_dim:])  # flatten the first M dims
+    # convert indices to a 1D flattened array
+    indices = torch_flatten_indices(indices, x_shape[:idx_dim])
+    selected = x[indices]
+    return selected
+
+
+# ========== module operations =========
+def set_requires_grad(model, requires_grad):
+    if torch.is_tensor(model):
+        model.requires_grad = requires_grad
+    else:
+        for param in model.parameters():
+            param.requires_grad = requires_grad
+
+
+def freeze_params(model):
+    set_requires_grad(model, False)
+    if not torch.is_tensor(model):
+        model.eval()
+
+
+def unfreeze_params(model):
+    set_requires_grad(model, True)
+    if not torch.is_tensor(model):
+        model.train()
+
+
+def clip_grad_value(model, max_value):
+    with torch.no_grad():
+        nn.utils.clip_grad_value_(model.parameters(), max_value)
+
+
+def clip_grad_norm(model, max_norm, norm_type=2):
+    """
+    Returns:
+        Total norm of the parameters (viewed as a single vector).
+    """
+    with torch.no_grad():
+        return nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_norm, norm_type=norm_type)
+
+
+def implements_state_dict(object, requires_load_method: bool = False):
+    cond = implements_method(object, "state_dict")
+    if requires_load_method:
+        return cond and implements_method(object, "load_state_dict")
+    else:
+        return cond
+
+
+def unwrap_ddp_model(model):
+    if hasattr(model, "module") and len(list(model.children())) == 1:
+        model = model.module
+    return model
+
+
+class DDPMethodWrapper(nn.Module):
+    """
+    Wraps another module's method as forward(), because DDP only works on forward()
+    This module can be wrapped with DDP and directly called.
+    It will not save any extra parameters
+    """
+
+    def __init__(self, net: nn.Module, method_name: str):
+        super().__init__()
+        self.net = net
+        assert_implements_method(net, method_name)
+        self._method_name = method_name
+
+    def forward(self, *args, **kwargs):
+        return getattr(self.net, self._method_name)(*args, **kwargs)
+
+    def state_dict(self):
+        return {}
+
+
+def to_state_dict(objects, to_cpu: bool = False, copy: bool = False, unwrap_ddp: bool = False):
+    """
+    Anything that has state_dict() method, e.g. nn.Module, Optimizer, LRScheduler, etc.
+
+    Args:
+        to_cpu: True to copy to CPU. The original tensors will still be on GPU.
+        copy: takes effect if and only if to_cpu is False
+    """
+
+    def _transfer(x):
+        if torch.is_tensor(x):
+            x = x.detach()
+            if to_cpu:
+                return x.cpu()
+            elif copy:
+                return x.clone()
+        return x
+
+    def _to_state_dict(m):
+        if implements_state_dict(m):
+            if isinstance(m, nn.Module) and unwrap_ddp:
+                m = unwrap_ddp_model(m)
+            return tree.map_structure(_transfer, m.state_dict())
+        else:
+            return _transfer(m)
+
+    return tree.map_structure(_to_state_dict, objects)
+
+
+def load_state_dict(objects, states, strip_prefix=None, strict=False):
+    """
+    Args:
+        strict: objects and states must match exactly
+        strip_prefix: only match the keys that have the prefix, and strip it
+    """
+
+    def _load(paths, obj):
+        if not implements_method(obj, "load_state_dict"):
+            raise ValueError(f"Object {type(obj)} does not support load_state_dict() method")
+        try:
+            state = tree_value_at_path(states, paths)
+        except ValueError:  # paths do not exist in `states` structure
+            if strict:
+                raise
+            else:
+                return
+        if strip_prefix:
+            assert isinstance(strip_prefix, str)
+            state = {
+                k[len(strip_prefix) :]: v for k, v in state.items() if k.startswith(strip_prefix)
+            }
+        if isinstance(obj, nn.Module):
+            return obj.load_state_dict(state, strict=strict)
+        else:
+            return obj.load_state_dict(state)
+
+    return tree.map_structure_with_path(_load, objects)
+
+
+def count_parameters(model):
+    return sum(x.numel() for x in model.parameters())
+
+
+def readable_count_parameters(model, precision: int = 2):
+    return to_readable_count_str(count_parameters(model), precision=precision)
+
+
+def get_module_device(model):
+    """
+    Returns:
+        first model parameter's device
+    """
+    return next(model.parameters()).device
+
+
+def maybe_transfer_module(model, device):
+    """
+    Transfer a module to another device if and only if they are on different devices.
+    Assumes that the module's first parameter determines the module device, i.e.
+    no model parallelism.
+
+    Returns:
+        True if module is transferred to a different device, False otherwise
+    """
+    if device is None:
+        return False
+    device = torch.device(device)
+    if get_module_device(model) != device:
+        model.to(device=device)
+        return True
+    else:
+        return False
+
+
+def clone_model(model):
+    with torch.no_grad():
+        new_model = deepcopy(model).to(get_module_device(model))
+    # new_model.load_state_dict(model.state_dict())
+    return new_model
+
+
+def update_soft_params(net, target_net, tau):
+    for param, target_param in zip(net.parameters(), target_net.parameters()):
+        target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)
+
+
+def tie_weights(src, trg):
+    # TODO deprecate this
+    assert type(src) is type(trg)
+    trg.weight = src.weight
+    trg.bias = src.bias
+
+
+def torch_normalize(tensor: torch.Tensor, mean, std, inplace=False):
+    """
+    Adapted from https://pytorch.org/docs/stable/_modules/torchvision/transforms/functional.html#normalize
+
+    Normalize a tensor image with mean and standard deviation.
+
+    .. note::
+        This transform acts out of place by default, i.e., it does not mutates the input tensor.
+
+    See :class:`~torchvision.transforms.Normalize` for more details.
+
+    Args:
+        tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
+        mean (sequence): Sequence of means for each channel.
+        std (sequence): Sequence of standard deviations for each channel.
+        inplace(bool,optional): Bool to make this operation inplace.
+
+    Returns:
+        Tensor: Normalized Tensor image.
+    """
+    if not torch.is_tensor(tensor):
+        raise TypeError("tensor should be a torch tensor. Got {}.".format(type(tensor)))
+
+    if not inplace:
+        tensor = tensor.clone()
+
+    dtype = tensor.dtype
+    mean = torch.as_tensor(mean, dtype=dtype, device=tensor.device)
+    std = torch.as_tensor(std, dtype=dtype, device=tensor.device)
+    if (std == 0).any():
+        raise ValueError(
+            f"std evaluated to zero after conversion to {dtype}, leading to division by zero."
+        )
+    if mean.ndim == 1:
+        mean = mean[:, None, None]
+    if std.ndim == 1:
+        std = std[:, None, None]
+    tensor.sub_(mean).div_(std)
+    return tensor
+
+
+def contains_rnn(net: nn.Module) -> bool:
+    for m in net.modules():
+        if isinstance(m, nn.RNNBase):
+            return True
+    return False
+
+
+def multi_one_hot(x, num_classes: List[int], to_float=True):
+    """
+    Concatenates multiple one-hot matrices, useful for embedding MultiDiscrete action space
+
+    Args:
+        x: torch.long, [*N, D]
+        num_classes: list len == D, match the last dim of x
+
+    Returns:
+        [*N, sum(num_classes)]
+    """
+    from torch.nn.functional import one_hot
+
+    assert x.dtype == torch.long
+    assert x.dim() >= 2, x.size()
+    assert len(num_classes) == x.size(-1), f"{len(num_classes)} != {x.size(1)}"
+    result = torch.cat(
+        [one_hot(t, c) for t, c in zip(torch.unbind(x, dim=-1), num_classes)], dim=-1
+    )
+    if to_float:
+        return result.float()
+    else:
+        return result
+
+
+def _random_derangement(n):
+    while True:
+        v = [i for i in range(n)]
+        for j in range(n - 1, -1, -1):
+            p = random.randint(0, j)
+            if v[p] == j:
+                break
+            else:
+                v[j], v[p] = v[p], v[j]
+        else:
+            if v[0] != 0:
+                return tuple(v)
+
+
+def random_derangement(n, format: Literal["list", "numpy", "torch"] = "torch"):
+    """
+    Early refusal algorithm, described at
+    https://stackoverflow.com/questions/25200220/generate-a-random-derangement-of-a-list
+    Derangement is permuation without fixed point, useful for constructing negative
+    pairs in contrastive learning.
+    """
+    assert format in ["list", "numpy", "torch"]
+    D = _random_derangement(n)
+    if format == "list":
+        return D
+    elif format == "numpy":
+        return np.array(D, dtype=np.long)
+    elif format == "torch":
+        return torch.tensor(D, dtype=torch.long)
+    else:
+        raise NotImplementedError(f"Unknown format {format}")
+
+
+def classify_accuracy(
+    output,
+    target,
+    topk: Union[int, List[int], Tuple[int]] = 1,
+    mask=None,
+    reduction="mean",
+    scale_100=False,
+):
+    """
+    Computes the accuracy over the k top predictions for the specified values of k.
+    Accuracy is a float between 0.0 and 1.0
+
+    Args:
+        topk: if int, return a single acc. If tuple, return a tuple of accs
+        mask: shape [batch_size,], binary mask of whether to include this sample or not
+    """
+    if isinstance(topk, int):
+        topk = [topk]
+        is_int = True
+    else:
+        is_int = False
+
+    batch_size = target.size(0)
+    assert output.size(0) == batch_size
+    if mask is not None:
+        assert mask.dim() == 1
+        assert mask.size(0) == batch_size
+
+    assert reduction in ["sum", "mean", "none"]
+    if reduction != "mean":
+        assert not scale_100, f"reduce={reduction} does not support scale_100=True"
+
+    with torch.no_grad():
+        maxk = max(topk)
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+        if mask is not None:
+            correct = mask * correct
+
+        mult = 100.0 if scale_100 else 1.0
+        res = []
+        for k in topk:
+            correct_k = correct[:k].int().sum(dim=0)
+            if reduction == "mean":
+                if mask is not None:
+                    # fmt: off
+                    res.append(
+                        float(correct_k.float().sum().mul_(mult / mask.sum().item()).item())
+                    )
+                    # fmt: on
+                else:
+                    res.append(float(correct_k.float().sum().mul_(mult / batch_size).item()))
+            elif reduction == "sum":
+                res.append(int(correct_k.sum().item()))
+            elif reduction == "none":
+                res.append(correct_k)
+            else:
+                raise NotImplementedError(f"Unknown reduce={reduction}")
+
+    if is_int:
+        assert len(res) == 1, "INTERNAL"
+        return res[0]
+    else:
+        return res
+
+
+def sequential_split_dataset(dataset: torch.utils.data.Dataset, split_portions: list[float]):
+    """
+    Split a dataset into multiple datasets, each with a different portion of the
+    original dataset. Uses torch.utils.data.Subset.
+    """
+    from .functional_utils import accumulate
+
+    assert len(split_portions) > 0, "split_portions must be a non-empty list"
+    assert all(0.0 <= p <= 1.0 for p in split_portions), f"{split_portions=}"
+    assert abs(sum(split_portions) - 1.0) < 1e-6, f"{sum(split_portions)=} != 1.0"
+    L = len(dataset)
+    assert L > 0, "dataset must be non-empty"
+    # split the list with proportions
+    lengths = [int(p * L) for p in split_portions]
+    # make sure the last split fills the full dataset
+    lengths[-1] += L - sum(lengths)
+    indices = list(range(L))
+
+    return [
+        torch.utils.data.Subset(dataset, indices[offset - length : offset])
+        for offset, length in zip(accumulate(lengths), lengths)
+    ]
+
+
+class RunningMeanStd:
+    def __init__(self):
+        """
+        Calulates the running mean and std of a data stream
+        https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm
+        """
+        self._mean = None
+        self._var = None
+        self._count = 0
+
+    @property
+    def mean(self):
+        return self._mean
+
+    @property
+    def var(self):
+        return self._var
+
+    @property
+    def std(self):
+        if isinstance(self._var, np.ndarray):
+            return np.sqrt(self._var)
+        else:
+            return self._var.sqrt()
+
+    @property
+    def count(self):
+        return self._count
+
+    def update(self, values: np.ndarray | torch.Tensor) -> None:
+        from .array_tensor_utils import any_mean, any_variance, get_batch_size
+
+        batch_mean = any_mean(values, dim=0)
+        # our running var calculation currently only supports unbiased=False
+        batch_var = any_variance(values, dim=0, unbiased=False)
+        batch_count = get_batch_size(values)
+        self.update_from_moments(batch_mean, batch_var, batch_count)
+
+    def update_from_moments(
+        self,
+        batch_mean: np.ndarray | torch.Tensor,
+        batch_var: np.ndarray | torch.Tensor,
+        batch_count: int,
+    ) -> None:
+        from .array_tensor_utils import any_get_shape
+
+        is_tensor = torch.is_tensor(batch_mean)
+        _zeros = batch_mean.new_zeros if is_tensor else np.zeros
+        if self._mean is None:
+            self._mean = _zeros(any_get_shape(batch_mean))
+        if self._var is None:
+            self._var = _zeros(any_get_shape(batch_var)) + 1.0
+
+        delta = batch_mean - self._mean
+        tot_count = self._count + batch_count
+        assert tot_count > 0, "count must be > 0"
+
+        new_mean = self._mean + delta * batch_count / tot_count
+        m_a = self._var * self._count
+        m_b = batch_var * batch_count
+        m_2 = m_a + m_b + delta * delta * self._count * batch_count / tot_count
+        new_var = m_2 / tot_count
+
+        self._mean = new_mean
+        self._var = new_var
+        self._count = tot_count
+
+
+class AverageMeter:
+    """Computes and stores the average and current value"""
+
+    def __init__(self, name="", fmt="f"):
+        self._name = name
+        self._fmt = fmt
+        self.reset()
+
+    def reset(self):
+        self._sum = 0.0
+        self._count = 0.0
+
+    @torch.no_grad()
+    def update(self, value, n=1):
+        if torch.is_tensor(value):
+            value = value.detach()
+        self._sum += value * n
+        self._count += n
+
+    @torch.no_grad()
+    def compute(self):
+        return float(self._sum / self._count)
+
+    def __float__(self):
+        return self.compute()
+
+    def __str__(self):
+        if self._fmt:
+            s = f"{float(self):{self._fmt}}"
+        else:
+            s = str(float(self))
+        if self._name:
+            return f"{self._name}: {s}"
+        return s

groot/vla/common/utils/misc/video_utils.py

@@ -0,0 +1,487 @@
+import json
+import subprocess
+
+import av
+import cv2
+import numpy as np
+import torchvision
+
+# Import decord with graceful fallback
+try:
+    import decord
+
+    DECORD_AVAILABLE = True
+except ImportError:
+    DECORD_AVAILABLE = False
+
+try:
+    import torchcodec
+
+    TORCHCODEC_AVAILABLE = True
+except (ImportError, RuntimeError):
+    TORCHCODEC_AVAILABLE = False
+
+
+def _get_video_info_ffmpeg(video_path: str) -> dict:
+    """Get video metadata using ffprobe."""
+    cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-select_streams",
+        "v:0",
+        "-show_entries",
+        "stream=nb_frames,duration,r_frame_rate",
+        "-of",
+        "json",
+        video_path,
+    ]
+
+    try:
+        output = subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8")
+        probe_data = json.loads(output)
+        stream = probe_data["streams"][0]
+
+        # Parse frame rate (comes as fraction like "15/1")
+        if "/" in stream["r_frame_rate"]:
+            num, den = map(int, stream["r_frame_rate"].split("/"))
+            fps = num / den
+        else:
+            fps = float(stream["r_frame_rate"])
+
+        # Get frame count and duration
+        nb_frames = int(stream.get("nb_frames", 0))
+        duration = float(stream.get("duration", 0))
+
+        # If nb_frames is not available, estimate from duration and fps
+        if nb_frames == 0 and duration > 0:
+            nb_frames = int(duration * fps)
+
+        return {
+            "nb_frames": nb_frames,
+            "fps": fps,
+            "duration": duration,
+        }
+    except (subprocess.CalledProcessError, json.JSONDecodeError, KeyError) as e:
+        raise ValueError(f"Failed to get video info for {video_path}: {e}")
+
+
+def _extract_frames_ffmpeg(video_path: str, frame_indices: list[int]) -> np.ndarray:
+    """Extract specific frames using ffmpeg."""
+    frames = []
+
+    for idx in frame_indices:
+        # Use ffmpeg to extract a specific frame
+        cmd = [
+            "ffmpeg",
+            "-i",
+            video_path,
+            "-vf",
+            f"select=eq(n\\,{idx})",
+            "-vframes",
+            "1",
+            "-f",
+            "image2pipe",
+            "-pix_fmt",
+            "rgb24",
+            "-vcodec",
+            "rawvideo",
+            "-",
+        ]
+
+        try:
+            output = subprocess.check_output(cmd, stderr=subprocess.DEVNULL)
+
+            # Check if output is empty (frame doesn't exist)
+            if len(output) == 0:
+                raise subprocess.CalledProcessError(1, cmd)
+
+            # Get frame dimensions by probing first
+            if len(frames) == 0:
+                info_cmd = [
+                    "ffprobe",
+                    "-v",
+                    "error",
+                    "-select_streams",
+                    "v:0",
+                    "-show_entries",
+                    "stream=width,height",
+                    "-of",
+                    "json",
+                    video_path,
+                ]
+                info_output = subprocess.check_output(info_cmd).decode("utf-8")
+                info_data = json.loads(info_output)
+                width = info_data["streams"][0]["width"]
+                height = info_data["streams"][0]["height"]
+
+            # Decode raw RGB data
+            frame_data = np.frombuffer(output, dtype=np.uint8)
+            frame = frame_data.reshape((height, width, 3))
+            frames.append(frame)
+
+        except subprocess.CalledProcessError:
+            # Frame might not exist, create a black frame
+            if len(frames) > 0:
+                frames.append(np.zeros_like(frames[0]))
+            else:
+                # Default fallback frame
+                frames.append(np.zeros((480, 640, 3), dtype=np.uint8))
+
+    return np.array(frames)
+
+
+def _extract_frames_at_timestamps_ffmpeg(video_path: str, timestamps: list[float]) -> np.ndarray:
+    """Extract frames at specific timestamps using ffmpeg."""
+    frames = []
+
+    for timestamp in timestamps:
+        cmd = [
+            "ffmpeg",
+            "-ss",
+            str(timestamp),
+            "-i",
+            video_path,
+            "-vframes",
+            "1",
+            "-f",
+            "image2pipe",
+            "-pix_fmt",
+            "rgb24",
+            "-vcodec",
+            "rawvideo",
+            "-",
+        ]
+
+        try:
+            output = subprocess.check_output(cmd, stderr=subprocess.DEVNULL)
+
+            # Check if output is empty (timestamp doesn't exist)
+            if len(output) == 0:
+                raise subprocess.CalledProcessError(1, cmd)
+
+            # Get frame dimensions
+            if len(frames) == 0:
+                info_cmd = [
+                    "ffprobe",
+                    "-v",
+                    "error",
+                    "-select_streams",
+                    "v:0",
+                    "-show_entries",
+                    "stream=width,height",
+                    "-of",
+                    "json",
+                    video_path,
+                ]
+                info_output = subprocess.check_output(info_cmd).decode("utf-8")
+                info_data = json.loads(info_output)
+                width = info_data["streams"][0]["width"]
+                height = info_data["streams"][0]["height"]
+
+            # Decode raw RGB data
+            frame_data = np.frombuffer(output, dtype=np.uint8)
+            frame = frame_data.reshape((height, width, 3))
+            frames.append(frame)
+
+        except subprocess.CalledProcessError:
+            # Timestamp might be out of bounds, use last frame or black frame
+            if len(frames) > 0:
+                frames.append(frames[-1])
+            else:
+                frames.append(np.zeros((480, 640, 3), dtype=np.uint8))
+
+    return np.array(frames)
+
+
+def _extract_all_frames_ffmpeg(video_path: str) -> tuple[np.ndarray, np.ndarray]:
+    """Extract all frames and their timestamps using ffmpeg."""
+    # Get video info
+    info = _get_video_info_ffmpeg(video_path)
+    fps = info["fps"]
+
+    # Extract all frames
+    cmd = [
+        "ffmpeg",
+        "-i",
+        video_path,
+        "-f",
+        "image2pipe",
+        "-pix_fmt",
+        "rgb24",
+        "-vcodec",
+        "rawvideo",
+        "-",
+    ]
+
+    try:
+        output = subprocess.check_output(cmd, stderr=subprocess.DEVNULL)
+
+        # Get frame dimensions
+        info_cmd = [
+            "ffprobe",
+            "-v",
+            "error",
+            "-select_streams",
+            "v:0",
+            "-show_entries",
+            "stream=width,height",
+            "-of",
+            "json",
+            video_path,
+        ]
+        info_output = subprocess.check_output(info_cmd).decode("utf-8")
+        info_data = json.loads(info_output)
+        width = info_data["streams"][0]["width"]
+        height = info_data["streams"][0]["height"]
+
+        # Decode all frames
+        frame_data = np.frombuffer(output, dtype=np.uint8)
+        total_pixels = len(frame_data) // 3
+        actual_frames = total_pixels // (width * height)
+
+        frames = frame_data[: actual_frames * width * height * 3].reshape(
+            (actual_frames, height, width, 3)
+        )
+
+        # Generate timestamps
+        timestamps = np.arange(actual_frames) / fps
+
+        return frames, timestamps
+
+    except subprocess.CalledProcessError as e:
+        raise ValueError(f"Failed to extract frames from {video_path}: {e}")
+
+
+def get_frames_by_indices(
+    video_path: str,
+    indices: list[int] | np.ndarray,
+    video_backend: str = "ffmpeg",
+    video_backend_kwargs: dict = {},
+) -> np.ndarray:
+    if video_backend == "decord":
+        if not DECORD_AVAILABLE:
+            raise ImportError("decord is not available. Install it with: pip install decord")
+        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 == "ffmpeg":
+        return _extract_frames_ffmpeg(video_path, list(indices))
+    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
+    else:
+        raise NotImplementedError
+
+
+def get_frames_by_timestamps(
+    video_path: str,
+    timestamps: list[float] | np.ndarray,
+    video_backend: str = "ffmpeg",
+    video_backend_kwargs: dict = {},
+    fps: None | float = None,
+) -> 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 "ffmpeg".
+        fps (float, optional): FPS of the video. Defaults to 30.
+    Returns:
+        np.ndarray: Frames at the specified timestamps.
+    """
+    if video_backend == "decord":
+        if not DECORD_AVAILABLE:
+            raise ImportError("decord is not available. Install it with: pip install decord")
+        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
+        )
+
+        # https://docs.pytorch.org/torchcodec/stable/generated/torchcodec.decoders.VideoStreamMetadata.html#torchcodec.decoders.VideoStreamMetadata
+        # Temporary fix: use 30 fps as the fps of the video (agibot) 
+        # TODO: get fps as parameter
+        if fps is None: 
+            fps = decoder.metadata.average_fps
+        interval = 1 / fps
+        timestamps = np.array(timestamps).astype(np.float64)
+
+        if np.all(timestamps == 0):
+            timestamps = np.arange(len(timestamps)) / fps
+
+        # Get video duration range from first and last frames
+        # This is a robust way to get valid timestamp range without depending on specific metadata attributes
+        first_frame = decoder.get_frames_at(indices=[0])
+        last_frame = decoder.get_frames_at(indices=[len(decoder) - 1])
+        min_pts = float(first_frame.pts_seconds[0])
+        max_pts = float(last_frame.pts_seconds[0])
+        
+        # Clamp timestamps to valid range to avoid RuntimeError
+        timestamps = np.clip(timestamps, min_pts, max_pts)
+
+        # Correct float precision issues in timestamps
+        # E.g. for 5fps video: [1.0, 1.20000005, 1.39999998] -> [1.0, 1.2, 1.4]
+        # Without this, the torchcodec will read the delayed frame (e.g. 1.39999998 -> 1.2)
+        # Round to nearest frame interval to prevent torchcodec from reading wrong frames
+        # Allow max 1% error from expected interval
+        if fps is None: 
+            closest_timestamps = np.round(timestamps / interval) * interval
+            # Re-clamp after rounding to ensure still in valid range
+            closest_timestamps = np.clip(closest_timestamps, min_pts, max_pts)
+            timestamp_errors = np.abs(closest_timestamps - timestamps) / interval
+            invalid_mask = timestamp_errors >= 0.01
+            if np.any(invalid_mask):
+                invalid_indices = np.where(invalid_mask)[0]
+                invalid_timestamps = timestamps[invalid_indices]
+                raise ValueError(
+                    f"Try to read invalid timestamps {invalid_timestamps} from video {video_path} (FPS: {fps})"
+                )
+
+            timestamps = closest_timestamps
+
+        return decoder.get_frames_played_at(seconds=timestamps).data.numpy()
+    elif video_backend == "ffmpeg":
+        return _extract_frames_at_timestamps_ffmpeg(video_path, list(timestamps))
+    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 == "torchvision_av":
+        # set backend
+        torchvision.set_video_backend("pyav")
+
+        # set a video stream reader
+        reader = torchvision.io.VideoReader(video_path, "video")
+
+        # set the first and last requested timestamps
+        # Note: previous timestamps are usually loaded, since we need to access the previous key frame
+        first_ts = timestamps[0]
+        last_ts = timestamps[-1]
+
+        # access closest key frame of the first requested frame
+        # Note: closest key frame timestamp is usally smaller than `first_ts` (e.g. key frame can be the first frame of the video)
+        # for details on what `seek` is doing see: https://pyav.basswood-io.com/docs/stable/api/container.html?highlight=inputcontainer#av.container.InputContainer.seek
+        reader.seek(first_ts, keyframes_only=True)
+
+        # Decode frames sequentially, storing the ones we need in a dictionary
+        # to map timestamps to frame data. This allows for easy re-ordering later.
+        found_frames_map = {}
+        tolerance = 0.001  # 1ms tolerance for timestamp matching
+        
+        for frame in reader:
+            current_ts = frame["pts"]
+
+            # Use tolerance-based matching instead of exact match
+            for ts in timestamps:
+                if ts not in found_frames_map and abs(current_ts - ts) < tolerance:
+                    found_frames_map[ts] = frame["data"]
+                    break
+
+            if current_ts >= last_ts + tolerance or len(found_frames_map) == len(timestamps):
+                break
+
+        reader.container.close()
+        reader = None
+        
+        # Debug: print timestamp matching results
+        print(f"[video_utils] Requested {len(timestamps)} timestamps: {timestamps[:4]}{'...' if len(timestamps) > 4 else ''}")
+        print(f"[video_utils] Found {len(found_frames_map)} frames with tolerance={tolerance}s")
+        if len(found_frames_map) < len(timestamps):
+            missing = [ts for ts in timestamps if ts not in found_frames_map]
+            print(f"[video_utils] WARNING: Missing timestamps: {missing[:4]}{'...' if len(missing) > 4 else ''}")
+        
+        frames = np.array(list(found_frames_map.values()))
+        return frames.transpose(0, 2, 3, 1)
+
+    else:
+        raise NotImplementedError
+
+
+def get_all_frames(
+    video_path: str,
+    video_backend: str = "ffmpeg",
+    video_backend_kwargs: dict = {},
+) -> tuple[np.ndarray, np.ndarray]:
+    """Get all frames from a video.
+
+    Returns:
+        tuple[np.ndarray, np.ndarray]: Frames and timestamps.
+    """
+    if video_backend == "decord":
+        if not DECORD_AVAILABLE:
+            raise ImportError("decord is not available. Install it with: pip install decord")
+        vr = decord.VideoReader(video_path, **video_backend_kwargs)
+        frames = vr.get_batch(range(len(vr))).asnumpy()
+        return frames, vr.get_frame_timestamp(range(len(vr)))[:, 0]
+    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 == "ffmpeg":
+        return _extract_all_frames_ffmpeg(video_path)
+    elif video_backend == "pyav":
+        container = av.open(video_path)
+        stream = container.streams.video[0]
+        assert stream.time_base is not None
+        frames = []
+        timestamps = []
+        for frame in container.decode(video=0):
+            frames.append(frame.to_ndarray(format="rgb24"))
+            timestamps.append(frame.pts * stream.time_base)
+        container.close()
+        return np.stack(frames), np.array(timestamps)
+
+    else:
+        raise NotImplementedError

groot/vla/configs/conf.yaml

@@ -0,0 +1,135 @@
+defaults:
+  - _self_  # all below configs will override this conf.yaml
+  - model: dreamzero/vla
+  - data: dreamzero/droid_horizon_relative
+  - override hydra/hydra_logging: disabled  # disable hydra logging
+  - override hydra/job_logging: disabled  # disable hydra job logging
+
+
+# === Model Arguments ===
+model: ???
+
+# === Data Arguments ===
+train_dataset: ???
+
+# ======== Trainer ========
+trainer:
+  _target_: groot.vla.experiment.VLATrainer
+  _partial_: true
+  _recursive_: false
+  callbacks:
+  model: ???                                  # model
+  train_dataset: ???                          # train_dataset
+  compute_dtype: ???                          # dtype_from_string(model.config.model_dtype)
+  benchmark_time: false                       # whether or not to benchmark time for training
+  # Legacy per-step profiling (profiles every N steps)
+  enable_profiling: false                     # (legacy) enable per-step profiling in training_step
+  profiling_steps: 5                          # (legacy) profile every N steps
+  # ProfCallback: window-based profiling
+  enable_prof_callback: false                 # enable ProfCallback for window-based profiling
+  profile_start_step: 50                      # session step to start profiling
+  profile_warmup_steps: 1                     # warmup steps for profiler
+  profile_active_steps: 3                     # active profiling steps
+  profile_record_shapes: false                # record tensor shapes (adds overhead)
+  profile_with_stack: false                   # record Python stack traces
+  profile_memory: false                       # record memory allocation
+
+# === Training Arguments ===
+
+wandb_project: ???  # needs to be specified by user
+output_dir: ???  # need to be specified by user
+load_from_yaml:  # need to be specified by user, will override the current config
+gear_credentials: null
+upload_checkpoints: false
+upload_every: 1000
+upload_last_n_checkpoints: 5
+remove_unused_columns: false
+bf16: false
+tf32: false
+global_batch_size: null
+raise_error_if_global_batch_size_not_set: false
+per_device_train_batch_size: 256
+per_device_eval_batch_size: 64
+gradient_accumulation_steps: 1
+dataloader_num_workers: 10
+dataloader_pin_memory: true
+dataloader_persistent_workers: true
+optim: adamw_torch
+learning_rate: 1e-4
+adam_beta1: 0.95
+adam_beta2: 0.999
+adam_epsilon: 1e-8
+weight_decay: 1e-6
+lr_scheduler_type: cosine
+warmup_ratio: 0.05
+logging_steps: 10.0
+num_train_epochs: 1000
+max_steps: -1
+save_strategy: steps
+save_steps: 500
+eval_strategy: "no"  # there has to be a double quote; otherwise a bare `no` will be interpreted as False
+save_total_limit: 8
+report_to: wandb
+seed: 42
+do_eval: false
+gradient_checkpointing: false
+ddp_find_unused_parameters: false
+ddp_bucket_cap_mb: 100
+ray_num_workers: ???
+eval_bf16: true
+torch_compile_mode: null
+
+pretrained_model_path: null
+only_tune_projectors: false
+
+save_llm: false
+save_lora_only: false
+save_value_model: false
+save_q_model: false
+
+download_cache: false
+
+training_args:
+  _target_: transformers.TrainingArguments
+  output_dir: ${output_dir}
+  run_name: ???                                # training_args.output_dir.split("/")[-1]
+  remove_unused_columns: ${remove_unused_columns}
+  deepspeed: ""
+  gradient_checkpointing: ${gradient_checkpointing}
+  bf16: ${bf16}
+  tf32: ${tf32}
+  per_device_train_batch_size: ${per_device_train_batch_size}
+  per_device_eval_batch_size: ${per_device_eval_batch_size}
+  gradient_accumulation_steps: ${gradient_accumulation_steps}
+  dataloader_num_workers: ${dataloader_num_workers}
+  dataloader_pin_memory: ${dataloader_pin_memory}
+  dataloader_persistent_workers: ${dataloader_persistent_workers}
+  optim: ${optim}
+  adam_beta1: ${adam_beta1}
+  adam_beta2: ${adam_beta2}
+  adam_epsilon: ${adam_epsilon}
+  learning_rate: ${learning_rate}
+  weight_decay: ${weight_decay}
+  warmup_ratio: ${warmup_ratio}
+  lr_scheduler_type: ${lr_scheduler_type}
+  logging_steps: ${logging_steps}
+  num_train_epochs: ${num_train_epochs}
+  max_steps: ${max_steps}
+  save_strategy: ${save_strategy}
+  save_steps: ${save_steps}
+  save_total_limit: ${save_total_limit}
+  report_to: ${report_to}
+  seed: ${seed}
+  do_eval: ${do_eval}
+  ddp_find_unused_parameters: ${ddp_find_unused_parameters}
+  ddp_bucket_cap_mb: ${ddp_bucket_cap_mb}
+  torch_compile_mode: ${torch_compile_mode}
+
+# === Profiling Arguments ===
+profile_dir: null
+
+# === Disable Hydra Config ===
+hydra:
+  output_subdir: null
+  run:
+    dir: .

groot/vla/configs/data/dreamzero/agibot_relative.yaml

@@ -0,0 +1,54 @@
+# @package _global_
+
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_ # this file will override the base
+
+max_state_dim: 64
+use_global_metadata: false
+relative_action: true
+relative_action_per_horizon: false
+relative_action_keys:
+  - left_arm_joint_position
+  - right_arm_joint_position
+  - left_effector_position
+  - right_effector_position
+  - head_position
+  - waist_position
+max_chunk_size: 5
+# Use 10% of data in shards before moving to next shard
+dataset_shard_sampling_rate: 0.1
+mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec
+single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID
+
+# Set your AGIbot dataset path here or override via CLI:
+#   agibot_data_root=/path/to/your/agibot_dataset
+agibot_data_root: ???
+
+train_dataset:
+  _target_: ${mixture_dataset_cls}
+  _convert_: object
+  mixture_spec:
+    - dataset_path:
+        agibot:
+          - ${agibot_data_root}
+      dataset_weight: 1.0
+      distribute_weights: true
+
+  dataset_class: ${single_dataset_cls}
+  all_modality_configs: ${modality_configs}
+  all_transforms: ${transforms}
+  metadata_versions: ${metadata_versions}
+  fps: ${fps}
+  dataset_kwargs:
+    video_backend: decord
+    use_global_metadata: ${use_global_metadata}
+    max_chunk_size: ${max_chunk_size}
+    relative_action: ${relative_action}
+    relative_action_keys: ${relative_action_keys}
+    relative_action_per_horizon: ${relative_action_per_horizon}
+  mixture_kwargs:
+    training: true
+    balance_dataset_weights: false
+    seed: 42
+    shard_sampling_rate: ${dataset_shard_sampling_rate}

groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml

@@ -0,0 +1,439 @@
+# @package _global_
+
+# Assume `model_specific_transform` is defined in the model config file
+
+################################################################################
+# Normalization Statistics
+# By default, we compute the normalization statistics for the datasets actually
+# used in the mixture. If you want to use the global metadata, set this to true.
+################################################################################
+
+use_global_metadata: false
+
+################################################################################
+# Dimension Information
+################################################################################
+
+num_frames: 49
+action_horizon: 48
+state_horizon: 1
+
+# image_resolution_width: 832
+# image_resolution_height: 480
+
+image_resolution_width: 480
+image_resolution_height: 256
+
+image_resolution_width_single_frame: 256
+image_resolution_height_single_frame: 256
+
+################################################################################
+# Anchored Video Transforms
+################################################################################
+totensor_cfg: &totensor_cfg
+  _target_: groot.vla.data.transform.VideoToTensor
+  apply_to: ???
+
+crop_cfg: &crop_cfg
+  _target_: groot.vla.data.transform.VideoCrop
+  apply_to: ???
+  scale: 0.95
+  mode: random
+
+
+resize_cfg: &resize_cfg
+  _target_: groot.vla.data.transform.VideoResize
+  apply_to: ???
+  height: ${image_resolution_height}
+  width: ${image_resolution_width}
+  interpolation: linear
+
+resize_cfg_single_frame: &resize_cfg_single_frame
+  _target_: groot.vla.data.transform.VideoResize
+  apply_to: ???
+  height: ${image_resolution_height_single_frame}
+  width: ${image_resolution_width_single_frame}
+  interpolation: linear
+
+color_jitter_cfg: &color_jitter_cfg
+  _target_: groot.vla.data.transform.VideoColorJitter
+  apply_to: ???
+  brightness: 0.3
+  contrast: 0.4
+  saturation: 0.5
+  hue: 0.08
+
+random_grayscale_cfg: &random_grayscale_cfg
+  _target_: groot.vla.data.transform.VideoRandomGrayscale
+  apply_to: ???
+  p: 0.1
+
+random_posterize_cfg: &random_posterize_cfg
+  _target_: groot.vla.data.transform.VideoRandomPosterize
+  apply_to: ???
+  bits: 4
+  p: 0.1
+
+normalize_cfg: &normalize_cfg
+  _target_: groot.vla.data.transform.VideoNormalize
+  apply_to: ???
+  mean: [0.5, 0.5, 0.5]
+  std: [0.5, 0.5, 0.5]
+
+
+to_numpy_cfg: &to_numpy_cfg
+  _target_: groot.vla.data.transform.VideoToNumpy
+  apply_to: ???
+
+
+################################################################################
+# oxe_droid (OXE Droid)
+################################################################################
+
+# Modality Configs
+modality_config_oxe_droid:
+  video:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_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, 24]
+    eval_delta_indices: [0]
+    modality_keys:
+      - video.exterior_image_1_left
+      - video.exterior_image_2_left
+      - video.wrist_image_left
+  state:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - state.joint_position
+      - state.gripper_position
+  action:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_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]
+    modality_keys:
+      - action.joint_position
+      - action.gripper_position
+  language:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - annotation.language.language_instruction
+      - annotation.language.language_instruction_2
+      - annotation.language.language_instruction_3
+  lapa_action:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - lapa_action
+
+# Transforms
+transform_oxe_droid:
+  _target_: groot.vla.data.transform.ComposedModalityTransform
+  transforms:
+    # Video transforms
+    - <<: *totensor_cfg
+      apply_to: ${modality_config_oxe_droid.video.modality_keys}
+    - <<: *crop_cfg
+      apply_to: ${modality_config_oxe_droid.video.modality_keys}
+    - <<: *resize_cfg
+      apply_to: ${modality_config_oxe_droid.video.modality_keys}
+    - <<: *color_jitter_cfg
+      apply_to: ${modality_config_oxe_droid.video.modality_keys}
+    - <<: *to_numpy_cfg
+      apply_to: ${modality_config_oxe_droid.video.modality_keys}
+
+    # State transforms
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_oxe_droid.state.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_oxe_droid.state.modality_keys}
+      normalization_modes:
+        state.joint_position: q99
+        state.gripper_position: q99
+  
+    # Action transforms
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_oxe_droid.action.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_oxe_droid.action.modality_keys}
+      normalization_modes:
+        action.joint_position: q99
+        action.gripper_position: q99
+
+    # ConcatTransform
+    - _target_: groot.vla.data.transform.ConcatTransform
+      video_concat_order: ${modality_config_oxe_droid.video.modality_keys}
+      state_concat_order: ${modality_config_oxe_droid.state.modality_keys}
+      action_concat_order: ${modality_config_oxe_droid.action.modality_keys}
+
+    # Model-specific transform
+    - ${model_specific_transform}
+
+################################################################################
+# agibot (AGIbot: state 32, action 22, 3 views)
+################################################################################
+
+modality_config_agibot:
+  video:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_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, 24]
+    eval_delta_indices: [-3, -2, -1, 0]
+    modality_keys:
+      - video.top_head
+      - video.hand_left
+      - video.hand_right
+  state:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - state.left_arm_joint_position
+      - state.right_arm_joint_position
+      - state.left_effector_position
+      - state.right_effector_position
+      - state.head_position
+      - state.waist_position
+  action:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_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]
+    modality_keys:
+      - action.left_arm_joint_position
+      - action.right_arm_joint_position
+      - action.left_effector_position
+      - action.right_effector_position
+      - action.head_position
+      - action.waist_position
+      - action.robot_velocity
+  language:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - annotation.language.action_text
+
+transform_agibot:
+  _target_: groot.vla.data.transform.ComposedModalityTransform
+  transforms:
+    # Video transforms
+    - <<: *totensor_cfg
+      apply_to: ${modality_config_agibot.video.modality_keys}
+    - <<: *crop_cfg
+      apply_to: ${modality_config_agibot.video.modality_keys}
+    - <<: *resize_cfg
+      apply_to: ${modality_config_agibot.video.modality_keys}
+    - <<: *color_jitter_cfg
+      apply_to: ${modality_config_agibot.video.modality_keys}
+    - <<: *to_numpy_cfg
+      apply_to: ${modality_config_agibot.video.modality_keys}
+
+    # State transforms
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_agibot.state.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_agibot.state.modality_keys}
+      normalization_modes:
+        state.left_arm_joint_position: q99
+        state.right_arm_joint_position: q99
+        state.left_effector_position: q99
+        state.right_effector_position: q99
+        state.head_position: q99
+        state.waist_position: q99
+
+    # Action transforms
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_agibot.action.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_agibot.action.modality_keys}
+      normalization_modes:
+        action.left_arm_joint_position: q99
+        action.right_arm_joint_position: q99
+        action.left_effector_position: q99
+        action.right_effector_position: q99
+        action.head_position: q99
+        action.waist_position: q99
+        action.robot_velocity: q99
+
+    # ConcatTransform
+    - _target_: groot.vla.data.transform.ConcatTransform
+      video_concat_order: ${modality_config_agibot.video.modality_keys}
+      state_concat_order: ${modality_config_agibot.state.modality_keys}
+      action_concat_order: ${modality_config_agibot.action.modality_keys}
+
+    # Model-specific transform
+    - ${model_specific_transform}
+
+################################################################################
+# yam (YAM: joint+gripper only from Dataset/YAM_play_data/meta/modality.json;
+#       state 14 dims, action 14 dims, 3 views)
+################################################################################
+
+modality_config_yam:
+  video:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_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, 24]
+    eval_delta_indices: [0]
+    modality_keys:
+      - video.top_camera-images-rgb
+      - video.left_camera-images-rgb
+      - video.right_camera-images-rgb
+  state:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - state.left_joint_pos
+      - state.left_gripper_pos
+      - state.right_joint_pos
+      - state.right_gripper_pos
+  action:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_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]
+    modality_keys:
+      - action.left_joint_pos
+      - action.left_gripper_pos
+      - action.right_joint_pos
+      - action.right_gripper_pos
+  language:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - annotation.task
+
+transform_yam:
+  _target_: groot.vla.data.transform.ComposedModalityTransform
+  transforms:
+    # Video transforms
+    - <<: *totensor_cfg
+      apply_to: ${modality_config_yam.video.modality_keys}
+    - <<: *crop_cfg
+      apply_to: ${modality_config_yam.video.modality_keys}
+    - <<: *resize_cfg
+      apply_to: ${modality_config_yam.video.modality_keys}
+    - <<: *color_jitter_cfg
+      apply_to: ${modality_config_yam.video.modality_keys}
+    - <<: *to_numpy_cfg
+      apply_to: ${modality_config_yam.video.modality_keys}
+
+    # State transforms
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_yam.state.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_yam.state.modality_keys}
+      normalization_modes:
+        state.left_joint_pos: q99
+        state.left_gripper_pos: q99
+        state.right_joint_pos: q99
+        state.right_gripper_pos: q99
+
+    # Action transforms
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_yam.action.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_yam.action.modality_keys}
+      normalization_modes:
+        action.left_joint_pos: q99
+        action.left_gripper_pos: q99
+        action.right_joint_pos: q99
+        action.right_gripper_pos: q99
+
+    # ConcatTransform
+    - _target_: groot.vla.data.transform.ConcatTransform
+      video_concat_order: ${modality_config_yam.video.modality_keys}
+      state_concat_order: ${modality_config_yam.state.modality_keys}
+      action_concat_order: ${modality_config_yam.action.modality_keys}
+
+    # Model-specific transform
+    - ${model_specific_transform}
+
+
+# libero_sim (LIBERO: 1 camera, 8-dim state, 7-dim action)
+modality_config_libero_sim:
+  video:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0, 1, 2, 3, 4, 5, 6, 7]
+    modality_keys:
+      - video.image
+  state:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - state.joint_pos
+      - state.gripper_pos
+  action:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+    modality_keys:
+      - action.joint_pos
+  language:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - annotation.task
+
+transform_libero_sim:
+  _target_: groot.vla.data.transform.ComposedModalityTransform
+  transforms:
+    - <<: *totensor_cfg
+      apply_to: ${modality_config_libero_sim.video.modality_keys}
+    - <<: *crop_cfg
+      apply_to: ${modality_config_libero_sim.video.modality_keys}
+    - <<: *resize_cfg
+      apply_to: ${modality_config_libero_sim.video.modality_keys}
+    - <<: *color_jitter_cfg
+      apply_to: ${modality_config_libero_sim.video.modality_keys}
+    - <<: *to_numpy_cfg
+      apply_to: ${modality_config_libero_sim.video.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_libero_sim.state.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_libero_sim.state.modality_keys}
+      normalization_modes:
+        state.joint_pos: q99
+        state.gripper_pos: q99
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_libero_sim.action.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_libero_sim.action.modality_keys}
+      normalization_modes:
+        action.joint_pos: q99
+    - _target_: groot.vla.data.transform.ConcatTransform
+      video_concat_order: ${modality_config_libero_sim.video.modality_keys}
+      state_concat_order: ${modality_config_libero_sim.state.modality_keys}
+      action_concat_order: ${modality_config_libero_sim.action.modality_keys}
+    - ${model_specific_transform}
+
+
+
+################################################################################
+# Modality Configs
+################################################################################
+
+modality_configs:
+  oxe_droid: ${modality_config_oxe_droid}
+  agibot: ${modality_config_agibot}
+  yam: ${modality_config_yam}
+  libero_sim: ${modality_config_libero_sim}
+
+################################################################################
+# Transforms
+################################################################################
+
+transforms:
+  oxe_droid: ${transform_oxe_droid}
+  agibot: ${transform_agibot}
+  yam: ${transform_yam}
+  libero_sim: ${transform_libero_sim}
+
+################################################################################
+# Metadata Versions
+################################################################################
+
+metadata_versions:
+  oxe_droid: '0221'
+  agibot: '0221'
+  yam: '0221'
+  libero_sim: 0221
+
+################################################################################
+# FPS (per embodiment, null means use dataset default)
+################################################################################
+
+fps:
+  yam: 30
+  libero_sim: 10

groot/vla/configs/data/dreamzero/droid_relative.yaml

@@ -0,0 +1,49 @@
+# @package _global_
+
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_ # this file will override the base
+
+max_state_dim: 64
+use_global_metadata: false
+relative_action: true
+relative_action_per_horizon: false
+relative_action_keys:
+  - joint_position
+max_chunk_size: 5
+# Use 10% of data in shards before moving to next shard
+dataset_shard_sampling_rate: 0.1
+mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec
+single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID
+
+# Set your DROID dataset path here or override via CLI:
+#   droid_data_root=/path/to/your/droid_dataset
+droid_data_root: ???
+
+train_dataset:
+  _target_: ${mixture_dataset_cls}
+  _convert_: object
+  mixture_spec:
+    - dataset_path:
+        oxe_droid:
+          - ${droid_data_root}
+      dataset_weight: 1.0
+      distribute_weights: true
+
+  dataset_class: ${single_dataset_cls}
+  all_modality_configs: ${modality_configs}
+  all_transforms: ${transforms}
+  metadata_versions: ${metadata_versions}
+  fps: ${fps}
+  dataset_kwargs:
+    video_backend: decord
+    use_global_metadata: ${use_global_metadata}
+    max_chunk_size: ${max_chunk_size}
+    relative_action: ${relative_action}
+    relative_action_keys: ${relative_action_keys}
+    relative_action_per_horizon: ${relative_action_per_horizon}
+  mixture_kwargs:
+    training: true
+    balance_dataset_weights: false
+    seed: 42
+    shard_sampling_rate: ${dataset_shard_sampling_rate}

groot/vla/configs/data/dreamzero/droid_relative_wan22.yaml

@@ -0,0 +1,51 @@
+# @package _global_
+# DROID data config for Wan 5B (Wan2.2): 320x160 so latent is 20x10 (even H,W) with WanVideoVAE38 (16x), frame_seqlen=50.
+# Extends base and droid settings directly to avoid Hydra nesting (dreamzero/droid_relative would double-resolve defaults).
+
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_
+
+# Wan 5B: 160x320 (HxW) -> latent 10x20, (10//2)*(20//2)=50. Use H,W divisible by 32 so latent is even (no crop in loss).
+image_resolution_width: 320
+image_resolution_height: 160
+
+max_state_dim: 64
+use_global_metadata: false
+relative_action: true
+relative_action_per_horizon: false
+relative_action_keys:
+  - joint_position
+max_chunk_size: 5
+dataset_shard_sampling_rate: 0.1
+mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec
+single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID
+droid_data_root: ???
+
+train_dataset:
+  _target_: ${mixture_dataset_cls}
+  _convert_: object
+  mixture_spec:
+    - dataset_path:
+        oxe_droid:
+          - ${droid_data_root}
+      dataset_weight: 1.0
+      distribute_weights: true
+
+  dataset_class: ${single_dataset_cls}
+  all_modality_configs: ${modality_configs}
+  all_transforms: ${transforms}
+  metadata_versions: ${metadata_versions}
+  fps: ${fps}
+  dataset_kwargs:
+    video_backend: decord
+    use_global_metadata: ${use_global_metadata}
+    max_chunk_size: ${max_chunk_size}
+    relative_action: ${relative_action}
+    relative_action_keys: ${relative_action_keys}
+    relative_action_per_horizon: ${relative_action_per_horizon}
+  mixture_kwargs:
+    training: true
+    balance_dataset_weights: false
+    seed: 42
+    shard_sampling_rate: ${dataset_shard_sampling_rate}

groot/vla/configs/data/dreamzero/libero.yaml

@@ -0,0 +1,85 @@
+# @package _global_
+# LIBERO data config for DreamZero SFT (Wan2.2)
+# Override base config with: model/dreamzero/action_head=wan_flow_matching_action_tf_wan22
+
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_
+
+libero_dataset_dir: /root/autodl-tmp/data/libero
+
+# Wan2.2 resolution
+image_resolution_width: 320
+image_resolution_height: 160
+
+# VAE 4x temporal: 8 frames
+num_frames: 12
+action_horizon: 12
+state_horizon: 1
+num_views: 1
+
+# Action head alignment
+num_frame_per_block: 2
+num_action_per_block: 12
+num_state_per_block: 1
+max_state_dim: 44
+max_action_dim: 32
+max_chunk_size: 4
+frame_seqlen: 50
+
+# Wan2.2 pretrained paths
+dit_version: /root/autodl-tmp/checkpoints/Wan2.2-TI2V-5B
+vae_pretrained_path: /root/autodl-tmp/checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth
+text_encoder_pretrained_path: /root/autodl-tmp/checkpoints/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth
+image_encoder_pretrained_path: /root/autodl-tmp/checkpoints/DreamZero-DROID/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+
+train_dataset:
+  _target_: groot.vla.data.dataset.libero.LiberoDataset
+  _convert_: object
+  dataset_dir: ${libero_dataset_dir}
+  num_frames: ${num_frames}
+  action_horizon: ${action_horizon}
+  state_horizon: ${state_horizon}
+  num_views: ${num_views}
+  video_height: ${image_resolution_height}
+  video_width: ${image_resolution_width}
+  max_episodes: null
+
+transform_libero:
+  _target_: groot.vla.data.transform.ComposedModalityTransform
+  transforms:
+    - ${model_specific_transform}
+
+transforms:
+  oxe_droid: ${transform_libero}
+
+modality_config_libero:
+  video:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0, 1, 2, 3, 4, 5, 6, 7]
+    modality_keys:
+      - video
+  state:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - state
+  action:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+    modality_keys:
+      - action
+  language:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - language
+
+modality_configs:
+  oxe_droid: ${modality_config_libero}
+
+metadata_versions:
+  oxe_droid: '0221'
+
+fps:
+  oxe_droid: 10

groot/vla/configs/data/dreamzero/libero_sim_relative.yaml

@@ -0,0 +1,58 @@
+# @package _global_
+# LIBERO data config for DreamZero GEAR pipeline (Wan2.2, 1 camera)
+# Follows the official DATASET_TO_GEAR_AND_TRAIN.md pattern.
+
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_
+
+# Wan2.2 resolution: 160x320 -> latent 10x20
+image_resolution_width: 320
+image_resolution_height: 160
+
+# LIBERO: 12 frames, 12-step action chunk, 1 camera
+use_global_metadata: false
+relative_action: true
+relative_action_per_horizon: false
+relative_action_keys:
+  - joint_position
+max_chunk_size: 5
+dataset_shard_sampling_rate: 0.1
+mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec
+single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID
+
+# Set dataset path via CLI or env var:
+#   libero_sim_data_root=/root/autodl-tmp/data/libero_gear
+libero_gear_root: /root/autodl-tmp/data/libero_gear
+
+train_dataset:
+  _target_: ${mixture_dataset_cls}
+  _convert_: object
+  mixture_spec:
+    - dataset_path:
+        libero_sim:
+          - ${libero_gear_root}/libero_spatial
+          - ${libero_gear_root}/libero_object
+          - ${libero_gear_root}/libero_goal
+          - ${libero_gear_root}/libero_10
+          - ${libero_gear_root}/libero_90
+      dataset_weight: 1.0
+      distribute_weights: true
+
+  dataset_class: ${single_dataset_cls}
+  all_modality_configs: ${modality_configs}
+  all_transforms: ${transforms}
+  metadata_versions: ${metadata_versions}
+  fps: ${fps}
+  dataset_kwargs:
+    video_backend: decord
+    use_global_metadata: ${use_global_metadata}
+    max_chunk_size: ${max_chunk_size}
+    relative_action: ${relative_action}
+    relative_action_keys: ${relative_action_keys}
+    relative_action_per_horizon: ${relative_action_per_horizon}
+  mixture_kwargs:
+    training: true
+    balance_dataset_weights: false
+    seed: 42
+    shard_sampling_rate: ${dataset_shard_sampling_rate}

groot/vla/configs/data/dreamzero/libero_sim_relative.yaml.bak

@@ -0,0 +1,54 @@
+# @package _global_
+# LIBERO data config for DreamZero GEAR pipeline (Wan2.2, 1 camera)
+# Follows the official DATASET_TO_GEAR_AND_TRAIN.md pattern.
+
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_
+
+# Wan2.2 resolution: 160x320 -> latent 10x20
+image_resolution_width: 320
+image_resolution_height: 160
+
+# LIBERO: 12 frames, 12-step action chunk, 1 camera
+use_global_metadata: false
+relative_action: true
+relative_action_per_horizon: false
+relative_action_keys:
+  - joint_pos
+max_chunk_size: 5
+dataset_shard_sampling_rate: 0.1
+mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec
+single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID
+
+# Set dataset path via CLI or env var:
+#   libero_sim_data_root=/root/autodl-tmp/data/libero_gear
+libero_sim_data_root: ???
+
+train_dataset:
+  _target_: ${mixture_dataset_cls}
+  _convert_: object
+  mixture_spec:
+    - dataset_path:
+        libero_sim:
+          - ${libero_sim_data_root}
+      dataset_weight: 1.0
+      distribute_weights: true
+
+  dataset_class: ${single_dataset_cls}
+  all_modality_configs: ${modality_configs}
+  all_transforms: ${transforms}
+  metadata_versions: ${metadata_versions}
+  fps: ${fps}
+  dataset_kwargs:
+    video_backend: decord
+    use_global_metadata: ${use_global_metadata}
+    max_chunk_size: ${max_chunk_size}
+    relative_action: ${relative_action}
+    relative_action_keys: ${relative_action_keys}
+    relative_action_per_horizon: ${relative_action_per_horizon}
+  mixture_kwargs:
+    training: true
+    balance_dataset_weights: false
+    seed: 42
+    shard_sampling_rate: ${dataset_shard_sampling_rate}

groot/vla/configs/data/dreamzero/manifeel.yaml

@@ -0,0 +1,86 @@
+# @package _global_
+# ManiFeel data config for DreamZero SFT (Wan2.2)
+# Override base config with: model/dreamzero/action_head=wan_flow_matching_action_tf_wan22
+
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_
+
+manifeel_dataset_dir: /root/autodl-tmp/tmp/manifeel_extracted
+
+# Wan2.2 resolution: 160x320 -> latent 10x20 -> frame_seqlen=50
+image_resolution_width: 320
+image_resolution_height: 160
+
+# VAE 4x temporal: 12 frames -> 12//4=3 latent frames
+# With num_frame_per_block=2, state_horizon=1 -> assertion: (3-1)/1 = 2 = 2/1
+num_frames: 12
+action_horizon: 12
+state_horizon: 1
+num_views: 3
+
+# Action head alignment (Wan2.2)
+num_frame_per_block: 2
+num_action_per_block: 12
+num_state_per_block: 1
+max_state_dim: 44
+max_action_dim: 32
+max_chunk_size: 4
+frame_seqlen: 50
+
+# Wan2.2 pretrained paths
+dit_version: /root/autodl-tmp/checkpoints/Wan2.2-TI2V-5B
+vae_pretrained_path: /root/autodl-tmp/checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth
+text_encoder_pretrained_path: /root/autodl-tmp/checkpoints/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth
+image_encoder_pretrained_path: /root/autodl-tmp/checkpoints/DreamZero-DROID/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+
+train_dataset:
+  _target_: groot.vla.data.dataset.manifeel.ManiFeelDataset
+  _convert_: object
+  dataset_dir: ${manifeel_dataset_dir}
+  num_frames: ${num_frames}
+  action_horizon: ${action_horizon}
+  state_horizon: ${state_horizon}
+  num_views: ${num_views}
+  video_height: ${image_resolution_height}
+  video_width: ${image_resolution_width}
+  max_episodes: null
+
+transform_manifeel:
+  _target_: groot.vla.data.transform.ComposedModalityTransform
+  transforms:
+    - ${model_specific_transform}
+
+transforms:
+  oxe_droid: ${transform_manifeel}
+
+modality_config_manifeel:
+  video:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0, 1, 2, 3, 4, 5, 6, 7]
+    modality_keys:
+      - video
+  state:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - state
+  action:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+    modality_keys:
+      - action
+  language:
+    _target_: groot.vla.data.dataset.ModalityConfig
+    delta_indices: [0]
+    modality_keys:
+      - language
+
+modality_configs:
+  oxe_droid: ${modality_config_manifeel}
+
+metadata_versions:
+  oxe_droid: '0221'
+
+fps:
+  oxe_droid: 10

groot/vla/configs/data/dreamzero/robotwin.yaml

@@ -0,0 +1,35 @@
+# @package _global_
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_
+
+robotwin_dataset_dir: /root/autodl-tmp/data/robotwin_gear
+image_resolution_width: 320
+image_resolution_height: 160
+num_frames: 12
+action_horizon: 12
+state_horizon: 1
+num_views: 1
+num_frame_per_block: 2
+num_action_per_block: 12
+num_state_per_block: 1
+max_state_dim: 44
+max_action_dim: 32
+max_chunk_size: 4
+frame_seqlen: 50
+dit_version: /root/autodl-tmp/checkpoints/Wan2.2-TI2V-5B
+vae_pretrained_path: /root/autodl-tmp/checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth
+text_encoder_pretrained_path: /root/autodl-tmp/checkpoints/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth
+image_encoder_pretrained_path: /root/autodl-tmp/checkpoints/DreamZero-DROID/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
+
+train_dataset:
+  _target_: groot.vla.data.dataset.robotwin.RobotWinDataset
+  _convert_: object
+  dataset_dir: ${robotwin_dataset_dir}
+  num_frames: ${num_frames}
+  action_horizon: ${action_horizon}
+  state_horizon: ${state_horizon}
+  num_views: ${num_views}
+  video_height: ${image_resolution_height}
+  video_width: ${image_resolution_width}
+  max_episodes: null

groot/vla/configs/data/dreamzero/yam_relative.yaml

@@ -0,0 +1,52 @@
+# @package _global_
+
+defaults:
+  - dreamzero/base_48_wan_fine_aug_relative
+  - _self_ # this file will override the base
+
+max_state_dim: 64
+use_global_metadata: false
+relative_action: true
+relative_action_per_horizon: false
+relative_action_keys:
+  - left_joint_pos
+  - left_gripper_pos
+  - right_joint_pos
+  - right_gripper_pos
+max_chunk_size: 5
+# Use 10% of data in shards before moving to next shard
+dataset_shard_sampling_rate: 0.1
+mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec
+single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID
+
+# Set your YAM dataset path here or override via CLI:
+#   yam_data_root=/path/to/your/yam_dataset
+yam_data_root: ???
+
+train_dataset:
+  _target_: ${mixture_dataset_cls}
+  _convert_: object
+  mixture_spec:
+    - dataset_path:
+        yam:
+          - ${yam_data_root}
+      dataset_weight: 1.0
+      distribute_weights: true
+
+  dataset_class: ${single_dataset_cls}
+  all_modality_configs: ${modality_configs}
+  all_transforms: ${transforms}
+  metadata_versions: ${metadata_versions}
+  fps: ${fps}
+  dataset_kwargs:
+    video_backend: decord
+    use_global_metadata: ${use_global_metadata}
+    max_chunk_size: ${max_chunk_size}
+    relative_action: ${relative_action}
+    relative_action_keys: ${relative_action_keys}
+    relative_action_per_horizon: ${relative_action_per_horizon}
+  mixture_kwargs:
+    training: true
+    balance_dataset_weights: false
+    seed: 42
+    shard_sampling_rate: ${dataset_shard_sampling_rate}

groot/vla/configs/deepspeed/zero2.json

@@ -0,0 +1,28 @@
+{
+  "checkpoint": {
+        "load_universal": false
+  },
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "zero_allow_untested_optimizer": true,
+  "fp16": {
+    "enabled": "auto",
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "initial_scale_power": 16,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "bf16": {
+    "enabled": "auto"
+  },
+    "zero_optimization": {
+        "stage": 2,
+        "overlap_comm": false,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 1e8
+    }
+}

groot/vla/configs/deepspeed/zero2_offload.json

@@ -0,0 +1,32 @@
+{
+  "checkpoint": {
+        "load_universal": false
+  },
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "zero_allow_untested_optimizer": true,
+  "fp16": {
+    "enabled": "auto",
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "initial_scale_power": 16,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "bf16": {
+    "enabled": "auto"
+  },
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": false,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 1e8
+    }
+}