diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..c359d1decacfd23401f554bdb2dc657e14a48567 --- /dev/null +++ b/.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*/ \ No newline at end of file diff --git a/COPYRIGHT b/COPYRIGHT new file mode 100644 index 0000000000000000000000000000000000000000..dd908a2529902a5e576649f2a9064c34a8b048d4 --- /dev/null +++ b/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. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..0a3f705806bf29ff9db3252d5da4d371fde38c20 --- /dev/null +++ b/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. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..263d7670258132e10e6d40cb0205bf9e613260fe --- /dev/null +++ b/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= 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 --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 +``` + +### 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 +``` + +(Optional only for GB200) Tensorrt enables faster generation +```bash +export LOAD_TRT_ENGINE=/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 +``` +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. Make sure to align the 3 camera view order to ensure positive transfer. + +### 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= + +# 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) diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..0ddbbf886d54b10346e75579b119f7707f4bbd0f --- /dev/null +++ b/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"] diff --git a/docs/DATASET_TO_GEAR_AND_TRAIN.md b/docs/DATASET_TO_GEAR_AND_TRAIN.md new file mode 100644 index 0000000000000000000000000000000000000000..3b9c462ce8723d831a49548b166639c5ee91af2f --- /dev/null +++ b/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 `` 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 \ + --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": ""}` | +| `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 = "" +``` + +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 = [ + ..., + "", +] +``` + +--- + +## 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.` | `state.joint_pos` | +| Action | `action.` | `action.joint_pos` | +| Video | `video.` | `video.cam0` | +| Language | `annotation.` | `annotation.task` | + +If the YAML keys don't match `modality.json`, training will fail with missing-key errors. + +### 3b. Add `modality_config_` + +```yaml +modality_config_: + 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_` + +```yaml +transform_: + _target_: groot.vla.data.transform.ComposedModalityTransform + transforms: + # Video + - <<: *totensor_cfg + apply_to: ${modality_config_.video.modality_keys} + - <<: *crop_cfg + apply_to: ${modality_config_.video.modality_keys} + - <<: *resize_cfg + apply_to: ${modality_config_.video.modality_keys} + - <<: *color_jitter_cfg + apply_to: ${modality_config_.video.modality_keys} + - <<: *to_numpy_cfg + apply_to: ${modality_config_.video.modality_keys} + + # State + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_.state.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_.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_.action.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_.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_.video.modality_keys} + state_concat_order: ${modality_config_.state.modality_keys} + action_concat_order: ${modality_config_.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: + ... + : ${modality_config_} + +transforms: + ... + : ${transform_} + +metadata_versions: + ... + : '0221' + +fps: + ... + : 30 # set to your dataset's FPS +``` + +--- + +## Step 4: Create a Dataset YAML + +Create `groot/vla/configs/data/dreamzero/_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 + +_data_root: ??? # set via CLI or env var + +train_dataset: + _target_: ${mixture_dataset_cls} + _convert_: object + mixture_spec: + - dataset_path: + : # must match key in modality_configs/transforms + - ${_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: + +- `` 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/_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__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/_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 \ + _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/_training.sh + +# With overrides: +DATA_ROOT=/path/to/your_dataset OUTPUT_DIR=./checkpoints/run1 NUM_GPUS=4 \ + bash scripts/train/_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. diff --git a/docs/DROID_CONVERSION.md b/docs/DROID_CONVERSION.md new file mode 100644 index 0000000000000000000000000000000000000000..e2f316dbc44ceaa0085355d8d5933ab1d5307764 --- /dev/null +++ b/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 +``` \ No newline at end of file diff --git a/docs/WAN22_BACKBONE.md b/docs/WAN22_BACKBONE.md new file mode 100644 index 0000000000000000000000000000000000000000..7e9bdefa9c90414496d71c35d00297256dad46b1 --- /dev/null +++ b/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. diff --git a/docs/data-preparation.md b/docs/data-preparation.md new file mode 100644 index 0000000000000000000000000000000000000000..785bb2a11ca16f1ecd7362c903c79f5a16a694ee --- /dev/null +++ b/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` | diff --git a/docs/training-guide.md b/docs/training-guide.md new file mode 100644 index 0000000000000000000000000000000000000000..3b0fb4c668d47f0ed1191bb050de0af6b9adedc3 --- /dev/null +++ b/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 # 指定网络接口 +``` diff --git a/eval_utils/policy_client.py b/eval_utils/policy_client.py new file mode 100644 index 0000000000000000000000000000000000000000..5511ff3afdad382e02fd04fc72f018e3fcfbb80c --- /dev/null +++ b/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({}) \ No newline at end of file diff --git a/eval_utils/policy_server.py b/eval_utils/policy_server.py new file mode 100644 index 0000000000000000000000000000000000000000..9b34bb3bd66064f9d6e83707c3181d7bc3c15624 --- /dev/null +++ b/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() + \ No newline at end of file diff --git a/eval_utils/run_sim_eval.py b/eval_utils/run_sim_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..cb194e77113a09db0e8756cf27a5bc31e0e04e51 --- /dev/null +++ b/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) diff --git a/eval_utils/serve_dreamzero_wan22.py b/eval_utils/serve_dreamzero_wan22.py new file mode 100644 index 0000000000000000000000000000000000000000..725f652ebadecef6ed24b01a20a052e44cb19871 --- /dev/null +++ b/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) diff --git a/groot/__init__.py b/groot/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/control/__init__.py b/groot/control/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/control/tensorrt_utils.py b/groot/control/tensorrt_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..74e25eb85698d6e5736d3a8708fe4882aa39db87 --- /dev/null +++ b/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) + + diff --git a/groot/vla/__init__.py b/groot/vla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/common/__init__.py b/groot/vla/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/groot/vla/common/__init__.py @@ -0,0 +1 @@ + diff --git a/groot/vla/common/utils/__init__.py b/groot/vla/common/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..762cc7c5bd7c6231b10f1155ed1f3c83ae933ac0 --- /dev/null +++ b/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 diff --git a/groot/vla/common/utils/data_structure/__init__.py b/groot/vla/common/utils/data_structure/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c3e40684ae7be6d8a216204a4d89b6c5352a72f2 --- /dev/null +++ b/groot/vla/common/utils/data_structure/__init__.py @@ -0,0 +1,2 @@ +from .shape_utils import * # noqa: F403 +from .tree_utils import * # noqa: F403 diff --git a/groot/vla/common/utils/data_structure/shape_utils.py b/groot/vla/common/utils/data_structure/shape_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..343524c5db1247a6c27f8322b983a5058ddaa17a --- /dev/null +++ b/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) diff --git a/groot/vla/common/utils/data_structure/tree_utils.py b/groot/vla/common/utils/data_structure/tree_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9cce53b6c14b1e7cf4ccd93790e964f47b2c84d8 --- /dev/null +++ b/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) diff --git a/groot/vla/common/utils/io/__init__.py b/groot/vla/common/utils/io/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8dc853ff2849d027d12c92d09d2ebaa553cfd97e --- /dev/null +++ b/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 diff --git a/groot/vla/common/utils/io/config_utils.py b/groot/vla/common/utils/io/config_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d8d541efa9c55107c37c0cd3bcb32566b0ce2f8b --- /dev/null +++ b/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) diff --git a/groot/vla/common/utils/io/file_utils.py b/groot/vla/common/utils/io/file_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3e53246b39997b519b1f131a440c64d500c00db2 --- /dev/null +++ b/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 diff --git a/groot/vla/common/utils/io/hdf5_utils.py b/groot/vla/common/utils/io/hdf5_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..03775968a73e9c3a39d1371472e10eb8f323323a --- /dev/null +++ b/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) diff --git a/groot/vla/common/utils/io/json_utils.py b/groot/vla/common/utils/io/json_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..34eb226264b6ae98f017ea46bbdd0bbdb549604f --- /dev/null +++ b/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) diff --git a/groot/vla/common/utils/io/print_utils.py b/groot/vla/common/utils/io/print_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..161746c9c83179dbb4a61902478e767854efaa36 --- /dev/null +++ b/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': } + """ + 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)) diff --git a/groot/vla/common/utils/io/termcolor.py b/groot/vla/common/utils/io/termcolor.py new file mode 100644 index 0000000000000000000000000000000000000000..bf5a13d3aebb39b02cee7d6f709696cafca92197 --- /dev/null +++ b/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 +# 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", + ) diff --git a/groot/vla/common/utils/misc/__init__.py b/groot/vla/common/utils/misc/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a3ab939903e7c546f80c32175d84560e747557a1 --- /dev/null +++ b/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 diff --git a/groot/vla/common/utils/misc/array_tensor_utils.py b/groot/vla/common/utils/misc/array_tensor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6585ca8a2693c38390b6236eb7c2a337bf3f2926 --- /dev/null +++ b/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.") diff --git a/groot/vla/common/utils/misc/functional_utils.py b/groot/vla/common/utils/misc/functional_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9a5a509feac792751ccfd7956736f932bea33679 --- /dev/null +++ b/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 diff --git a/groot/vla/common/utils/misc/image_utils.py b/groot/vla/common/utils/misc/image_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..46ae18cf8b9b89e577da99df8dd0919d07318c71 --- /dev/null +++ b/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) diff --git a/groot/vla/common/utils/misc/misc_utils.py b/groot/vla/common/utils/misc/misc_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ed486fec52a63fe1c97039c813ceaf03f5db5315 --- /dev/null +++ b/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 diff --git a/groot/vla/common/utils/misc/torch_utils.py b/groot/vla/common/utils/misc/torch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e735532c4289262d632adfee93ca68581f248533 --- /dev/null +++ b/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 diff --git a/groot/vla/common/utils/misc/video_utils.py b/groot/vla/common/utils/misc/video_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..481f169cef3adf518df273e1175603c31cfcfdd0 --- /dev/null +++ b/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 diff --git a/groot/vla/configs/conf.yaml b/groot/vla/configs/conf.yaml new file mode 100644 index 0000000000000000000000000000000000000000..48da05bc28ed7560307933d94396c4d6eeac715c --- /dev/null +++ b/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: . diff --git a/groot/vla/configs/data/dreamzero/agibot_relative.yaml b/groot/vla/configs/data/dreamzero/agibot_relative.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1972af4e9ab846f94bc9304a238b731e4c2deb14 --- /dev/null +++ b/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} diff --git a/groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml b/groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8358d612d73d4ec346f826934e1146d8a6ad3a63 --- /dev/null +++ b/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 diff --git a/groot/vla/configs/data/dreamzero/droid_relative.yaml b/groot/vla/configs/data/dreamzero/droid_relative.yaml new file mode 100644 index 0000000000000000000000000000000000000000..197ff702b465dddd0ef54cd572e59921bae04302 --- /dev/null +++ b/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} diff --git a/groot/vla/configs/data/dreamzero/droid_relative_wan22.yaml b/groot/vla/configs/data/dreamzero/droid_relative_wan22.yaml new file mode 100644 index 0000000000000000000000000000000000000000..42537e2b7b66be9b55c4b936eaa931b5fb3f17d1 --- /dev/null +++ b/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} diff --git a/groot/vla/configs/data/dreamzero/libero.yaml b/groot/vla/configs/data/dreamzero/libero.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c4d5eee7bc3ee0f7ea37e613eff891a9e4bd1444 --- /dev/null +++ b/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 diff --git a/groot/vla/configs/data/dreamzero/libero_sim_relative.yaml b/groot/vla/configs/data/dreamzero/libero_sim_relative.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7662e5993444e68a451fe66a3c12b1153912ccbf --- /dev/null +++ b/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} diff --git a/groot/vla/configs/data/dreamzero/libero_sim_relative.yaml.bak b/groot/vla/configs/data/dreamzero/libero_sim_relative.yaml.bak new file mode 100644 index 0000000000000000000000000000000000000000..64bffc5fe80ce6b511e12390dd79a3f172c10aae --- /dev/null +++ b/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} diff --git a/groot/vla/configs/data/dreamzero/manifeel.yaml b/groot/vla/configs/data/dreamzero/manifeel.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e1514b7dfc45aff6a7f74f6e1513e2b4fc46fca3 --- /dev/null +++ b/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 diff --git a/groot/vla/configs/data/dreamzero/robotwin.yaml b/groot/vla/configs/data/dreamzero/robotwin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1d399ec4f67241a19fd36c3824cb0c331f74c81f --- /dev/null +++ b/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 diff --git a/groot/vla/configs/data/dreamzero/yam_relative.yaml b/groot/vla/configs/data/dreamzero/yam_relative.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2670479f2c2369f1440f5c8333adf2310391bbd2 --- /dev/null +++ b/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} diff --git a/groot/vla/configs/deepspeed/zero2.json b/groot/vla/configs/deepspeed/zero2.json new file mode 100644 index 0000000000000000000000000000000000000000..b1a5a0b41b082ff36c8a0cec197bb61c7d78c2b1 --- /dev/null +++ b/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 + } +} diff --git a/groot/vla/configs/deepspeed/zero2_offload.json b/groot/vla/configs/deepspeed/zero2_offload.json new file mode 100644 index 0000000000000000000000000000000000000000..29c7f2f22577fa6ef48a9958285a4bbf5dc4026d --- /dev/null +++ b/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 + } +} diff --git a/groot/vla/configs/deepspeed/zero3.json b/groot/vla/configs/deepspeed/zero3.json new file mode 100644 index 0000000000000000000000000000000000000000..cfd068a9f91485787b0895e029bc3f0b1577adf1 --- /dev/null +++ b/groot/vla/configs/deepspeed/zero3.json @@ -0,0 +1,33 @@ +{ + "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": 3, + "overlap_comm": true, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": 5e8, + "stage3_prefetch_bucket_size": 5e8, + "stage3_param_persistence_threshold": 1e6, + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "stage3_gather_16bit_weights_on_model_save": true + } +} diff --git a/groot/vla/configs/deepspeed/zero3_multinode.json b/groot/vla/configs/deepspeed/zero3_multinode.json new file mode 100644 index 0000000000000000000000000000000000000000..8bf65a89f0d4e4e619b788bcd65058cce59b7c8b --- /dev/null +++ b/groot/vla/configs/deepspeed/zero3_multinode.json @@ -0,0 +1,34 @@ +{ + "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, + "bf16": { + "enabled": "auto" + }, + "communication_data_type": "bf16", + "zero_optimization": { + "stage": 3, + "overlap_comm": true, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": 5e8, + "allgather_bucket_size": 5e8, + "stage3_prefetch_bucket_size": 5e8, + "stage3_param_persistence_threshold": 1e6, + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "stage3_gather_16bit_weights_on_model_save": true + }, + "activation_checkpointing": { + "partition_activations": true, + "cpu_checkpointing": true, + "number_of_training_steps": 200000, + "synchronize_checkpoint_boundary": false, + "profile": false + } +} diff --git a/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf.yaml b/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8cccdc6319ac822f6f93ed9af9746034e74215ea --- /dev/null +++ b/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf.yaml @@ -0,0 +1,108 @@ +# @package _global_ +add_pos_embed: true +hidden_size: 64 +attn_dropout: 0.2 +repa_layer: 8 +repa_coeff: 1.0 +load_pretrained_det_decode_layer_path: null +expand_batch: null +dit_version: null +text_encoder_pretrained_path: null +image_encoder_pretrained_path: null +vae_pretrained_path: null +train_architecture: "lora" +num_frame_per_block: 1 +num_action_per_block: 32 +num_state_per_block: 1 +frame_seqlen: 880 +action_head_cfg: + _target_: groot.vla.model.dreamzero.action_head.wan_flow_matching_action_tf.WANPolicyHead + _convert_: object # to convert to a dict so that it can save the config to json + config: + _target_: groot.vla.model.dreamzero.action_head.wan_flow_matching_action_tf.WANPolicyHeadConfig + _recursive_: false # avoid instantiating noise_scheduler_cfg and obs_encoder_cfg + tiled: false + tile_size_height: 34 + tile_size_width: 34 + tile_stride_height: 18 + tile_stride_width: 16 + lora_rank: 4 + lora_alpha: 4 + num_frames: ${num_frames} + num_frame_per_block: ${num_frame_per_block} + lora_target_modules: "q,k,v,o,ffn.0,ffn.2" + init_lora_weights: "kaiming" + train_architecture: ${train_architecture} + use_gradient_checkpointing: true + add_pos_embed: ${add_pos_embed} + model_dtype: float32 + max_state_dim: ${max_state_dim} + max_action_dim: ${max_action_dim} + action_loss_embodiment_ids: [26, 17, 32] + hidden_size: ${hidden_size} + input_embedding_dim: 1536 + backbone_embedding_dim: ${backbone_hidden_size} + repa_layer: ${repa_layer} + repa_coeff: ${repa_coeff} + load_pretrained_det_decode_layer_path: ${load_pretrained_det_decode_layer_path} + freeze_decode_layer: false + expand_batch: ${expand_batch} + use_vlln: true + vl_self_attention_cfg: + _target_: groot.vla.model.n1_5.modules.cross_attention_dit.SelfAttentionTransformer + positional_embeddings: null + num_layers: 4 + num_attention_heads: 24 + attention_head_dim: 64 + dropout: ${attn_dropout} + final_dropout: true + diffusion_model_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_dit_action_casual_chunk.CausalWanModel + _convert_: object + diffusion_model_pretrained_path: ${dit_version} + model_type: "i2v" + frame_seqlen: ${frame_seqlen} + dim: 5120 + in_dim: 36 + ffn_dim: 13824 + out_dim: 16 + freq_dim: 256 + eps: 1e-6 + num_heads: 40 + num_layers: 40 + max_chunk_size: ${max_chunk_size} + num_frame_per_block: ${num_frame_per_block} + num_action_per_block: ${num_action_per_block} + num_state_per_block: ${num_state_per_block} + text_encoder_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_text_encoder.WanTextEncoder + _convert_: object + text_encoder_pretrained_path: ${text_encoder_pretrained_path} + image_encoder_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_image_encoder.WanImageEncoder + _convert_: object + image_encoder_pretrained_path: ${image_encoder_pretrained_path} + vae_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_vae.WanVideoVAE + _convert_: object + vae_pretrained_path: ${vae_pretrained_path} + + action_dim: ${max_action_dim} + action_horizon: ${action_horizon} + num_inference_timesteps: 4 # not used during training + noise_beta_alpha: 1.5 + noise_beta_beta: 1.0 + noise_s: 0.999 + num_timestep_buckets: 1000 + + # ========== DECOUPLED NOISE SAMPLING CONFIG ========== + # When true: video uses Beta(alpha,beta) biased high, action uses independent uniform + # When false (default): original behavior - video and action share same timestep + decouple_video_action_noise: false + # Beta params (only used when decouple_video_action_noise=true) + # Beta(3,1): mean=0.75, Beta(5,1): mean=0.83 (higher alpha = more bias to high noise) + video_noise_beta_alpha: 3.0 + video_noise_beta_beta: 1.0 + + tune_projector: true + tune_diffusion_model: true diff --git a/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf_wan22.yaml b/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf_wan22.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8884737442053cfeb28387daafc65789c4a9503f --- /dev/null +++ b/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf_wan22.yaml @@ -0,0 +1,39 @@ +# @package _global_ +# Wan2.2-TI2V-5B backbone config for DreamZero DROID training +# Uses CausalWanModel with Wan2.2 architecture (dim=3072, 48-channel VAE) +# Override base config with: model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 + +defaults: + - wan_flow_matching_action_tf + +# Wan2.2-TI2V-5B specific overrides +# frame_seqlen must match DiT patch_embedding output per frame: (H//2)*(W//2) for stride (1,2,2). +# WanVideoVAE38 has 16x spatial downscale. For even latent H,W (no crop in dynamics loss) use video +# resolution with H,W divisible by 32: 160x320 -> latent 10x20 -> patch 5x10=50. Else 176x320 -> 11x20 (odd H). +# Same as 5B pretrained: latent only (in_dim=48), first-frame via CLIP; no [x;y] concat. +# Action/state invariant: need (noisy_frames-1)//num_frame_per_block >= 1, so min 3 frames per half. +frame_seqlen: 50 + +action_head_cfg: + config: + # 160x320 gives latent 10x20 (even); avoids dynamics-loss crop. Match data=dreamzero/droid_relative_wan22. + target_video_height: 160 + target_video_width: 320 + diffusion_model_cfg: + model_type: "ti2v" + dim: 3072 + in_dim: 48 + concat_first_frame_latent: false + ffn_dim: 14336 + out_dim: 48 + freq_dim: 256 + eps: 1e-6 + num_heads: 24 + num_layers: 30 + max_state_dim: ${max_state_dim} + vae_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_vae.WanVideoVAE38 + _convert_: object + z_dim: 48 + dim: 160 + vae_pretrained_path: ${vae_pretrained_path} diff --git a/groot/vla/configs/model/dreamzero/backbone/identity.yaml b/groot/vla/configs/model/dreamzero/backbone/identity.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2e4cf9ced4117a2dc622308691e39551ef0ff2c6 --- /dev/null +++ b/groot/vla/configs/model/dreamzero/backbone/identity.yaml @@ -0,0 +1,9 @@ +# @package _global_ + +backbone_hidden_size: 0 +backbone_cfg: + _target_: groot.vla.model.dreamzero.backbone.identity.IdentityBackbone + +action_head_cfg: + config: + backbone_features_projector_cfg: null # no projection because no backbone features diff --git a/groot/vla/configs/model/dreamzero/transform/base.yaml b/groot/vla/configs/model/dreamzero/transform/base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e0a17998ecf620ee4ba79fce9746b2da62437800 --- /dev/null +++ b/groot/vla/configs/model/dreamzero/transform/base.yaml @@ -0,0 +1,43 @@ +# @package _global_ + +embodiment_tag_to_projector_index: + real_gr1_arms_only: 0 + real_gr1_arms_only_annotated: 1 + real_gr1_arms_waist: 2 + real_gr1_arms_waist_annotated: 3 + dexmg_gr1_arms_only_inspire: 4 + dexmg_gr1_arms_only_fourier: 5 + dexmg_gr1_arms_waist_fourier: 6 + robocasa_single_arm: 7 + onex_eve_gripper: 8 + robocasa_gr1_arms_only_inspire_hands: 9 + robocasa_gr1_arms_only_fourier_hands: 10 + robocasa_gr1_fixed_lower_body_inspire_hands: 11 + robocasa_gr1_fixed_lower_body_fourier_hands: 12 + robocasa_panda_omron: 13 + # s: 14 + robocasa_bimanual_panda_parallel_gripper: 15 + robocasa_bimanual_panda_inspire_hand: 16 + oxe_droid: 17 + oxe_fractal: 18 + oxe_language_table: 19 + oxe_bridge: 20 + real_panda_single_arm: 21 + hot3d_hands_only: 23 + gr1_unified: 24 + robocasa_gr1_arms_waist_fourier_hands: 25 + agibot: 26 + lapa: 27 + oxe_mutex: 28 + oxe_roboset: 29 + oxe_plex: 30 + dream: 31 + yam: 32 + xdof: 22 + gr1_unified_segmentation: 14 + language_table_sim: 7 + gr1_isaac: 0 + sim_behavior_r1_pro: 31 + mecka_hands: 27 + real_r1_pro_sharpa: 28 + libero_sim: 33 diff --git a/groot/vla/configs/model/dreamzero/transform/dreamzero_cotrain.yaml b/groot/vla/configs/model/dreamzero/transform/dreamzero_cotrain.yaml new file mode 100644 index 0000000000000000000000000000000000000000..326d24cbcd4699dedfa7a79d44c390144b285e5d --- /dev/null +++ b/groot/vla/configs/model/dreamzero/transform/dreamzero_cotrain.yaml @@ -0,0 +1,44 @@ +# @package _global_ + +defaults: + - base + - _self_ # this file will override model/transform/base + + +max_length: 512 +num_views: 1 +tokenizer_path: "google/umt5-xxl" +################################################################################ +# Data Collator +################################################################################ + +data_collator: + _target_: groot.vla.model.dreamzero.transform.dreamzero_cotrain.DefaultDataCollator + tokenizer_path: ${tokenizer_path} + max_length : ${max_length} + num_views: ${num_views} + embodiment_tag_mapping: ${embodiment_tag_to_projector_index} + +################################################################################ +# Model-Specific Transform +# Needs to change this model-specific transform for different models +################################################################################ + +num_visual_tokens_per_frame: 16 +max_state_dim: 64 +max_action_dim: 32 +language_dropout_prob: 0.0 + + +model_specific_transform: + _target_: groot.vla.model.dreamzero.transform.dreamzero_cotrain.DreamTransform + default_instruction: "Perform the default behavior." + language_dropout_prob: ${language_dropout_prob} + always_use_default_instruction: false + max_state_dim: ${max_state_dim} + max_action_dim: ${max_action_dim} + max_length: ${max_length} + state_horizon: ${state_horizon} + action_horizon: ${action_horizon} + embodiment_tag_mapping: ${embodiment_tag_to_projector_index} + tokenizer_path: ${tokenizer_path} diff --git a/groot/vla/configs/model/dreamzero/vla.yaml b/groot/vla/configs/model/dreamzero/vla.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4ab9fac6999409208720e71754ba1a9a0207b39f --- /dev/null +++ b/groot/vla/configs/model/dreamzero/vla.yaml @@ -0,0 +1,20 @@ +# @package _global_ + +defaults: + - _self_ # all below configs will override this file + - /model/dreamzero/backbone: identity + - /model/dreamzero/action_head: wan_flow_matching_action_tf + - /model/dreamzero/transform: dreamzero_cotrain + +model: + _target_: groot.vla.model.dreamzero.base_vla.VLA + _convert_: object + config: + _target_: groot.vla.model.dreamzero.base_vla.VLAConfig + _recursive_: false # avoid instantiating backbone_cfg and action_head_cfg + model_dtype: float32 + hidden_size: ${backbone_hidden_size} # this is defined in the backbone's yaml + action_horizon: ${action_head_cfg.config.action_horizon} # this is defined in the action_head's yaml + action_dim: ${action_head_cfg.config.action_dim} # this is defined in the action_head's yaml + backbone_cfg: ${backbone_cfg} # this is defined in the backbone's yaml + action_head_cfg: ${action_head_cfg} # this is defined in the action_head's yaml diff --git a/groot/vla/data/__init__.py b/groot/vla/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/data/conversion/__init__.py b/groot/vla/data/conversion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/data/conversion/gr1/__init__.py b/groot/vla/data/conversion/gr1/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/data/conversion/gr1/constants.py b/groot/vla/data/conversion/gr1/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..55dbd88fe0f14222ef91b2d23ebce40838d477b1 --- /dev/null +++ b/groot/vla/data/conversion/gr1/constants.py @@ -0,0 +1,27 @@ +from groot.vla.data.schema import EmbodimentTag + +RAW_HDF5_FILENAME = "state_action.hdf5" +RAW_METADATA_FILENAME = "metadata.json" +RAW_VIDEO_FILENAME = "ego_view.mp4" +RAW_ANNOTATION_FILENAME = "annotation.json" + +TRAINABLE_HDF5_FILENAME = "dataset.hdf5" +TRAINABLE_METADATA_FILENAME = "metadata.json" + +RAW_DATA_CONTROL_FREQUENCY = 20 + +INITIAL_ACTIONS_FILENAME = "initial_actions.npz" + +EMBODIMENT_TAG_TO_ANNOTATED_VERSION = { + EmbodimentTag.REAL_GR1_ARMS_ONLY: EmbodimentTag.REAL_GR1_ARMS_ONLY_ANNOTATED, + EmbodimentTag.REAL_GR1_ARMS_WAIST: EmbodimentTag.REAL_GR1_ARMS_WAIST_ANNOTATED, + # Special for 5DC-S + EmbodimentTag.ROBOCASA_GR1_ARMS_ONLY_FOURIER_HANDS: EmbodimentTag.REAL_GR1_ARMS_ONLY_ANNOTATED, + EmbodimentTag.ROBOCASA_GR1_ARMS_WAIST_FOURIER_HANDS: EmbodimentTag.REAL_GR1_ARMS_WAIST_ANNOTATED, +} + +# For chopped data +EPISODE_LENGTH_FILENAME = "episode_length.json" + +# Processed +PROCESSED_VIDEO_FILENAME = "ego_view_pad_res224_freq20.mp4" diff --git a/groot/vla/data/conversion/gr1/get_initial_actions.py b/groot/vla/data/conversion/gr1/get_initial_actions.py new file mode 100644 index 0000000000000000000000000000000000000000..885aed3e4c4ddbdf94ba56e4da2e17bd75102c9d --- /dev/null +++ b/groot/vla/data/conversion/gr1/get_initial_actions.py @@ -0,0 +1,128 @@ +from pathlib import Path + +import h5py +import numpy as np +import pandas as pd + +import groot.vla.common.utils as U +from groot.vla.data.conversion.gr1.constants import ( + INITIAL_ACTIONS_FILENAME, + TRAINABLE_HDF5_FILENAME, +) +from groot.vla.data.dataset.macro import ( + LE_ROBOT_EPISODE_FILENAME, + LE_ROBOT_INFO_FILENAME, + LE_ROBOT_METADATA_DIR, + LE_ROBOT_MODALITY_FILENAME, +) + + +def get_initial_actions(data_dir: str | Path): + hdf5_file = h5py.File(Path(data_dir) / TRAINABLE_HDF5_FILENAME, "r") + initial_actions = [] + + """ + initial_actions: dict[str, dict[str, np.ndarray]] + 0: (the dataset dimension) + trajectory_name: + action_key: + action: np.ndarray + """ + initial_actions = {} + for demo_name in hdf5_file["data"].keys(): + demo_group = hdf5_file["data"][demo_name] + initial_actions[demo_name] = {} + action_keys = list(demo_group["action"].keys()) + for action_key in action_keys: + initial_actions[demo_name][action_key] = demo_group["action"][action_key][0] + return [initial_actions] + + +def get_initial_actions_from_lerobot(data_dir: str | Path): + data_dir = Path(data_dir) + + # 1. Get modality for slicing action + meta_modality_path = data_dir / LE_ROBOT_METADATA_DIR / LE_ROBOT_MODALITY_FILENAME + meta_modality = U.load_json(meta_modality_path) + action_keys = meta_modality["action"].keys() + + # 2. Get episode paths + # 2.1. Get data_path_pattern + meta_info_path = data_dir / LE_ROBOT_METADATA_DIR / LE_ROBOT_INFO_FILENAME + meta_info = U.load_json(meta_info_path) + data_path_pattern = meta_info["data_path"] + chunk_size = meta_info["chunks_size"] + + # 2.2. Get episode info + episode_metadata_path = data_dir / LE_ROBOT_METADATA_DIR / LE_ROBOT_EPISODE_FILENAME + episode_metadata = U.load_jsonl(episode_metadata_path) + + initial_actions = {} + for episode_info in episode_metadata: + episode_index = episode_info["episode_index"] + episode_chunk = episode_index // chunk_size + episode_path = data_dir / data_path_pattern.format( + episode_chunk=episode_chunk, episode_index=episode_index + ) + if not episode_path.exists(): + raise ValueError(f"Episode path {episode_path} does not exist") + + episode_data = pd.read_parquet(episode_path) + + initial_action_concat = episode_data["action"].iloc[0] + trajectory_id = episode_info["episode_index"] + initial_actions[trajectory_id] = {} + for action_key in action_keys: + start = meta_modality["action"][action_key]["start"] + end = meta_modality["action"][action_key]["end"] + initial_actions[trajectory_id][action_key] = initial_action_concat[start:end] + return [initial_actions] + + +def save_initial_actions( + initial_actions: dict[str, dict[str, np.ndarray]], initial_actions_path: str | Path +): + np.savez(str(initial_actions_path), initial_actions) + + +def load_initial_actions(initial_actions_path: str | Path): + """ + initial_actions: list[dict[str, dict[str, np.ndarray]]] + 0: (the first dataset) + trajectory_name: + action_key: + action: np.ndarray + 1: (the second dataset) + ... + """ + initial_actions_npz = np.load(str(initial_actions_path), allow_pickle=True) + initial_actions = [] + initial_actions_array = initial_actions_npz[ + "arr_0" + ] # This is the default key when np.savez saves a list + for dataset_initial_actions in initial_actions_array: + initial_actions_for_this_dataset = {} + for trajectory_name, action_dict in dataset_initial_actions.items(): + initial_actions_for_this_dataset[trajectory_name] = action_dict + initial_actions.append(initial_actions_for_this_dataset) + return initial_actions + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Generate initial_actions.npz for a LeRobot dataset") + parser.add_argument("data_dir", type=str, help="Path to LeRobot dataset directory") + args = parser.parse_args() + + initial_actions = get_initial_actions_from_lerobot(args.data_dir) + save_initial_actions( + initial_actions, + Path(args.data_dir) / LE_ROBOT_METADATA_DIR / INITIAL_ACTIONS_FILENAME, + ) + + # Verify + loaded_initial_actions = load_initial_actions( + Path(args.data_dir) / LE_ROBOT_METADATA_DIR / INITIAL_ACTIONS_FILENAME + ) + print(f"Saved initial actions for {len(loaded_initial_actions[0])} trajectories") diff --git a/groot/vla/data/dataset/__init__.py b/groot/vla/data/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db17fc385679fa2b4e40e5c5c7631e43d3ae3638 --- /dev/null +++ b/groot/vla/data/dataset/__init__.py @@ -0,0 +1,5 @@ +from .lerobot import ModalityConfig + +__all__ = [ + "ModalityConfig", +] diff --git a/groot/vla/data/dataset/lerobot.py b/groot/vla/data/dataset/lerobot.py new file mode 100644 index 0000000000000000000000000000000000000000..f8efc6ed0f424f32cd4273be064680c33059106c --- /dev/null +++ b/groot/vla/data/dataset/lerobot.py @@ -0,0 +1,2499 @@ +from collections import defaultdict +import copy +import glob +import hashlib +import importlib +import json +from pathlib import Path +import time +from typing import Sequence, TypeVar + +import numpy as np +import pandas as pd +from pydantic import BaseModel, Field, ValidationError, field_validator +from torch.utils.data import Dataset +from tqdm import tqdm +import yaml +import torch + +from groot.vla.common.utils import get_all_frames, get_frames_by_timestamps +from groot.vla.data.conversion.gr1.get_initial_actions import load_initial_actions +from groot.vla.data.schema import ( + DatasetMetadata, + DatasetStatisticalValues, + EmbodimentTag, + LeRobotModalityMetadata, + LeRobotStateActionMetadata, +) +from groot.vla.data.transform import ComposedModalityTransform + +T_LeRobotMixtureDataset = TypeVar("T_LeRobotMixtureDataset", bound="LeRobotMixtureDataset") + +LE_ROBOT_MODALITY_FILENAME = "meta/modality.json" +LE_ROBOT_EPISODE_FILENAME = "meta/episodes.jsonl" +LE_ROBOT_TASKS_FILENAME = "meta/tasks.jsonl" +LE_ROBOT_INFO_FILENAME = "meta/info.json" +LE_ROBOT_STATS_FILENAME = "meta/stats.json" +LE_ROBOT_DATA_FILENAME = "data/*/*.parquet" +LE_ROBOT_TASK_EMBEDDINGS_FILENAME = "meta/task_embeddings.pt" +LE_ROBOT_DETAILED_GLOBAL_INSTRUCTION_FILENAME = "meta/episodes_detail_global_instruction.jsonl" +INITIAL_ACTIONS_FILENAME = "meta/initial_actions.npz" +METADATA_DIR = Path(importlib.import_module("groot.vla.data").__file__).parent / "metadata" # type: ignore +STEP_FILTER_FILENAME = "meta/step_filter.jsonl" +LEROBOT_RELATIVE_STATS_FILE_NAME = "meta/relative_stats_dreamzero.json" +LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME = "meta/relative_horizon_stats_dreamzero.json" + +# Special language keys that load from metadata files instead of parquet columns +METADATA_LANG_KEYS = ["detailed_global_instruction_medium", "detailed_global_instruction_concise"] + + +def calculate_dataset_statistics( + parquet_paths: list[Path], features: list[str] | None = None +) -> dict[str, DatasetStatisticalValues]: + """Calculate the dataset statistics of all columns for a list of parquet files. + + Args: + parquet_paths (list[Path]): List of paths to parquet files to process. + features (list[str] | None): List of feature names to compute statistics for. + If None, computes statistics for all columns in the data. + + Returns: + dict[str, DatasetStatisticalValues]: Dictionary mapping feature names to their + statistical values (mean, std, min, max, q01, q99). + """ + # Dataset statistics + all_low_dim_data_list = [] + # Collect all the data + for parquet_path in tqdm( + sorted(list(parquet_paths)), + desc="Collecting all parquet files...", + ): + # Load the parquet file + parquet_data = pd.read_parquet(parquet_path) + parquet_data = parquet_data + all_low_dim_data_list.append(parquet_data) + all_low_dim_data = pd.concat(all_low_dim_data_list, axis=0) + # Compute dataset statistics + dataset_statistics = {} + if features is None: + features = list(all_low_dim_data.columns) + for le_modality in features: + print(f"Computing statistics for {le_modality}...") + np_data = np.vstack( + [np.asarray(x, dtype=np.float32) for x in all_low_dim_data[le_modality]] + ) + dataset_statistics[le_modality] = DatasetStatisticalValues( + mean=np.mean(np_data, axis=0).tolist(), + std=np.std(np_data, axis=0).tolist(), + min=np.min(np_data, axis=0).tolist(), + max=np.max(np_data, axis=0).tolist(), + q01=np.quantile(np_data, 0.01, axis=0).tolist(), + q99=np.quantile(np_data, 0.99, axis=0).tolist(), + ) + return dataset_statistics + + +class ModalityConfig(BaseModel): + """Configuration for a modality defining how data should be sampled and loaded. + + This class specifies which indices to sample relative to a base index and which + keys to load for a particular modality (e.g., video, state, action). + """ + + delta_indices: list[int] + """Delta indices to sample relative to the current index. The returned data will correspond to the original data at a sampled base index + delta indices.""" + eval_delta_indices: list[int] | None = None + """Delta indices to sample relative to the current index for evaluation. If None, uses the same indices as delta_indices.""" + modality_keys: list[str] + """The keys to load for the modality in the dataset.""" + + def model_post_init(self, *args, **kwargs): + """Initialize eval_delta_indices to delta_indices if not provided.""" + super().model_post_init(*args, **kwargs) + if self.eval_delta_indices is None: + self.eval_delta_indices = self.delta_indices + + +class LeRobotSingleDataset(Dataset): + """ + Base dataset class for LeRobot that supports sharding. + """ + + def __init__( + self, + dataset_path: Path | str, + modality_configs: dict[str, ModalityConfig], + embodiment_tag: str | EmbodimentTag, + use_global_metadata: bool = True, + metadata_version: str | None = None, + video_backend: str = "ffmpeg", + video_backend_kwargs: dict | None = None, + transforms: ComposedModalityTransform | None = None, + discard_bad_trajectories: bool = True, + fps: float = None, + max_chunk_size: int = None, + relative_action: bool = False, + relative_action_keys: list[str] | None = None, + relative_action_per_horizon: bool = False, + ): + """ + Initialize the dataset. + + Args: + dataset_path (Path | str): The path to the dataset. + modality_configs (dict[str, ModalityConfig]): The configuration for each modality. The keys are the modality names, and the values are the modality configurations. + See `ModalityConfig` for more details. + use_global_metadata (bool): Whether to use global metadata for normalization. + metadata_version (str): The version of the metadata, if `use_global_metadata` is True. + video_backend (str): Backend for video reading. + video_backend_kwargs (dict): Keyword arguments for the video backend when initializing the video reader. + transforms (ComposedModalityTransform): The transforms to apply to the dataset. + embodiment_tag (EmbodimentTag): Overload the embodiment tag for the dataset. e.g. define it as "new_embodiment" + relative_action (bool): Whether to use relative action stats for normalization. If True, will load or calculate + relative action stats from relative_stats_dreamzero.json. If the file doesn't exist, stats will be calculated. + relative_action_keys (list[str] | None): List of action keys to apply relative action to (e.g., ['joint_position']). + If None and relative_action is True, applies to all action keys except those containing 'gripper'. + relative_action_per_horizon (bool): Whether to use per-horizon relative action stats. If True, will load or calculate + separate stats for each action horizon index from relative_horizon_stats_dreamzero.json. + """ + # first check if the path directory exists + if not Path(dataset_path).exists(): + raise FileNotFoundError(f"Dataset path {dataset_path} does not exist") + + self.modality_configs = modality_configs + self.use_global_metadata = use_global_metadata + self.metadata_version = metadata_version + self.video_backend = video_backend + self.video_backend_kwargs = video_backend_kwargs if video_backend_kwargs is not None else {} + self.fps = fps + self.max_chunk_size = max_chunk_size + self.transforms = ( + transforms if transforms is not None else ComposedModalityTransform(transforms=[]) + ) + self.discard_bad_trajectories = discard_bad_trajectories + self.relative_action = relative_action + self.relative_action_per_horizon = relative_action_per_horizon + # Determine which action keys should use relative action + if relative_action_keys is not None: + self.relative_action_keys = relative_action_keys + else: + # Default: apply to all action keys except those containing 'gripper' + self.relative_action_keys = None # Will be set after modality_configs is available + self._relative_action_keys_input = relative_action_keys # Store original input + self._dataset_path = Path(dataset_path) + self._dataset_name = self._dataset_path.name + self.tag = EmbodimentTag(embodiment_tag) + # For dream and lapa, we use the global metadata since the lapa_actions and dream_actions are already normalized + if self.tag == EmbodimentTag.DREAM or self.tag == EmbodimentTag.LAPA: + self.use_global_metadata = True + self._lerobot_modality_meta = self._get_lerobot_modality_meta() + self._lerobot_info_meta = self._get_lerobot_info_meta() + # Notice: We also include discarded trajectories in stats for larger state coverage, for questions please ask @Fengyuan Hu @Yuqi Xie + self._lerobot_stats_meta = self._get_lerobot_stats_meta() + + # Initialize trajectory info and chunk size early (needed for relative stats calculation) + self._trajectory_ids, self._trajectory_lengths = self._get_trajectories() + self._data_path_pattern = self._get_data_path_pattern() + self._chunk_size = self._get_chunk_size() + + # Set default relative_action_keys if not provided + if self.relative_action and self._relative_action_keys_input is None: + # Default: apply to all action keys except those containing 'gripper' + action_keys = self.modality_configs.get("action", ModalityConfig(delta_indices=[0], modality_keys=[])).modality_keys + self.relative_action_keys = [ + k.replace("action.", "") for k in action_keys + if "gripper" not in k.lower() + ] + print(f"Relative action will be applied to keys: {self.relative_action_keys}") + # Load relative action stats if relative_action is enabled + print("relative_action", self.relative_action) + self._lerobot_relative_stats_meta = self._get_lerobot_relative_stats_meta() if self.relative_action else {} + # Load per-horizon relative action stats if relative_action_per_horizon is enabled + print("relative_action_per_horizon", self.relative_action_per_horizon) + self._lerobot_relative_horizon_stats_meta = self._get_lerobot_relative_horizon_stats_meta() if self.relative_action_per_horizon else {} + self._metadata = self._get_metadata() + self._step_filter = self._get_step_filter() + self._all_steps = self._get_all_steps() + self._modality_keys = self._get_modality_keys() + self._delta_indices = self._get_delta_indices() + self._max_delta_index = self._get_max_delta_index() + self._dataset_name = self._dataset_path.name + + # NOTE(YL): method to predict the task progress + if "action.task_progress" in self._modality_keys["action"]: + from groot.vla.data.schema import StateActionMetadata + + print("we will add task progress to the action modality") + self._modality_keys["action"].append("action.task_progress") + self._metadata.modalities.action["task_progress"] = StateActionMetadata( + absolute=True, rotation_type=None, shape=(1,), continuous=True + ) + # assume the task progress is uniformly distributed between 0 and 1 + self._metadata.statistics.action["task_progress"] = DatasetStatisticalValues( + max=[1.0], min=[0.0], mean=[0.5], std=[0.2887], q01=[0.01], q99=[0.99] + ) + + self.set_transforms_metadata(self.metadata) + self.set_epoch(0) + + print(f"Initialized dataset {self.dataset_name} with {embodiment_tag}") + + # LeRobot-specific config (some already initialized above for relative stats) + self._video_path_pattern = self._get_video_path_pattern() + self._tasks = self._get_tasks() + self._detailed_global_instructions = self._get_detailed_global_instructions() + self.curr_traj_data = None + self.curr_traj_id = None + + # Check if the dataset is valid + self._check_integrity() + + @property + def dataset_path(self) -> Path: + """The path to the dataset that contains the METADATA_FILENAME file.""" + return self._dataset_path + + @property + def metadata(self) -> DatasetMetadata: + """The metadata for the dataset, loaded from metadata.json in the dataset directory""" + return self._metadata + + @property + def trajectory_ids(self) -> np.ndarray: + """The trajectory IDs in the dataset, stored as a 1D numpy array of strings.""" + return self._trajectory_ids + + @property + def trajectory_lengths(self) -> np.ndarray: + """The trajectory lengths in the dataset, stored as a 1D numpy array of integers. + The order of the lengths is the same as the order of the trajectory IDs. + """ + return self._trajectory_lengths + + @property + def all_steps(self) -> list[tuple[int, int]]: + """The trajectory IDs and base indices for all steps in the dataset. + Example: + self.trajectory_ids: [0, 1, 2] + self.trajectory_lengths: [3, 2, 4] + return: [ + ("traj_0", 0), ("traj_0", 1), ("traj_0", 2), + ("traj_1", 0), ("traj_1", 1), + ("traj_2", 0), ("traj_2", 1), ("traj_2", 2), ("traj_2", 3) + ] + """ + return self._all_steps + + @property + def modality_keys(self) -> dict: + """The modality keys for the dataset. The keys are the modality names, and the values are the keys for each modality. + + Example: { + "video": ["video.image_side_0", "video.image_side_1"], + "state": ["state.eef_position", "state.eef_rotation"], + "action": ["action.eef_position", "action.eef_rotation"], + "language": ["language.human.task"], + "timestamp": ["timestamp"], + "reward": ["reward"], + } + """ + return self._modality_keys + + @property + def delta_indices(self) -> dict[str, np.ndarray]: + """The delta indices for the dataset. The keys are the modality.key, and the values are the delta indices for each modality.key.""" + return self._delta_indices + + def _get_max_delta_index(self) -> int: + """Calculate the maximum delta index across all modalities. + + Returns: + int: The maximum delta index value. + """ + max_delta_index = 0 + for delta_index in self.delta_indices.values(): + max_delta_index = max(max_delta_index, delta_index.max()) + return max_delta_index + + @property + def max_delta_index(self) -> int: + """The maximum delta index across all modalities.""" + return self._max_delta_index + + @property + def dataset_name(self) -> str: + """The name of the dataset.""" + return self._dataset_name + + @property + def lerobot_modality_meta(self) -> LeRobotModalityMetadata: + """The metadata for the LeRobot dataset.""" + return self._lerobot_modality_meta + + @property + def lerobot_info_meta(self) -> dict: + """The metadata for the LeRobot dataset.""" + return self._lerobot_info_meta + + @property + def lerobot_stats_meta(self) -> dict[str, DatasetStatisticalValues]: + """The metadata for the LeRobot dataset.""" + return self._lerobot_stats_meta + + @property + def lerobot_relative_stats_meta(self) -> dict[str, DatasetStatisticalValues]: + """The relative action stats metadata for the LeRobot dataset.""" + return self._lerobot_relative_stats_meta + + @property + def lerobot_relative_horizon_stats_meta(self) -> dict[str, dict[str, list]]: + """The per-horizon relative action stats metadata for the LeRobot dataset. + + Format: {action_key: {stat_name: [[h0_vals], [h1_vals], ...]}} + """ + return self._lerobot_relative_horizon_stats_meta + + @property + def step_filter(self) -> dict[int, np.ndarray]: + """The step filter for the dataset.""" + return self._step_filter + + @property + def data_path_pattern(self) -> str: + """The path pattern for the LeRobot dataset.""" + return self._data_path_pattern + + @property + def video_path_pattern(self) -> str: + """The path pattern for the LeRobot dataset.""" + return self._video_path_pattern + + @property + def chunk_size(self) -> int: + """The chunk size for the LeRobot dataset.""" + return self._chunk_size + + @property + def tasks(self) -> pd.DataFrame: + """The tasks for the dataset.""" + return self._tasks + + def _get_lerobot_modality_meta(self) -> LeRobotModalityMetadata: + """Get the metadata for the LeRobot dataset.""" + if self.use_global_metadata: + assert ( + self.metadata_version is not None + ), "metadata_version must be provided if use_global_metadata is True" + modality_meta_path = ( + METADATA_DIR + / self.tag.value + / self.metadata_version + / Path(LE_ROBOT_MODALITY_FILENAME).name + ) + assert ( + modality_meta_path.exists() + ), f"Please provide a {Path(LE_ROBOT_MODALITY_FILENAME).name} file in {METADATA_DIR / self.tag.value / self.metadata_version}" + with open(modality_meta_path, "r") as f: + modality_meta = LeRobotModalityMetadata.model_validate(json.load(f)) + return modality_meta + else: + modality_meta_path = self.dataset_path / LE_ROBOT_MODALITY_FILENAME + assert ( + modality_meta_path.exists() + ), f"Please provide a {LE_ROBOT_MODALITY_FILENAME} file in {self.dataset_path}" + with open(modality_meta_path, "r") as f: + modality_meta = LeRobotModalityMetadata.model_validate(json.load(f)) + return modality_meta + + def _get_lerobot_info_meta(self) -> dict: + """Get the metadata for the LeRobot dataset.""" + info_meta_path = self.dataset_path / LE_ROBOT_INFO_FILENAME + with open(info_meta_path, "r") as f: + info_meta = json.load(f) + return info_meta + + def _get_lerobot_stats_meta(self) -> dict[str, DatasetStatisticalValues]: + """Get the metadata for the LeRobot dataset.""" + if self.use_global_metadata: + assert ( + self.metadata_version is not None + ), "metadata_version must be provided if use_global_metadata is True" + stats_path = ( + METADATA_DIR + / self.tag.value + / self.metadata_version + / Path(LE_ROBOT_STATS_FILENAME).name + ) + else: + stats_path = self.dataset_path / LE_ROBOT_STATS_FILENAME + try: + with open(stats_path, "r") as f: + stats: dict = json.load(f) + for name in ["num_trajectories", "total_trajectory_length"]: + stats.pop(name, None) + for name, stat in stats.items(): + stats[name] = DatasetStatisticalValues.model_validate(stat) + return stats + except (FileNotFoundError, ValidationError) as e: + if self.use_global_metadata: + raise ValueError( + f"{e}: Please provide a {Path(LE_ROBOT_STATS_FILENAME).name} file in {stats_path}" + " and ensure the metadata format is correct." + ) + print(f"Failed to load dataset statistics: {e}") + print(f"Calculating dataset statistics for {self.dataset_name}") + # Get all parquet files in the dataset paths + parquet_files = list((self.dataset_path).glob(LE_ROBOT_DATA_FILENAME)) + lowdim_features = [] + le_features = self.lerobot_info_meta["features"] + for feature in le_features: + if "float" in le_features[feature]["dtype"]: + lowdim_features.append(feature) + + stats = calculate_dataset_statistics(parquet_files, lowdim_features) + stats_serialized = {k: v.model_dump(mode="json") for k, v in stats.items()} + with open(stats_path, "w") as f: + json.dump(stats_serialized, f, indent=4) + return stats + + def _get_lerobot_relative_stats_meta(self) -> dict[str, DatasetStatisticalValues]: + """Get the relative action stats metadata for the LeRobot dataset. + + Returns: + dict[str, DatasetStatisticalValues]: Dictionary mapping action keys to their relative stats. + """ + # Determine the path for relative stats file + if self.use_global_metadata: + assert ( + self.metadata_version is not None + ), "metadata_version must be provided if use_global_metadata is True" + stats_path = ( + METADATA_DIR + / self.tag.value + / self.metadata_version + / Path(LEROBOT_RELATIVE_STATS_FILE_NAME).name + ) + assert ( + stats_path.exists() + ), f"Please provide a {Path(LEROBOT_RELATIVE_STATS_FILE_NAME).name} file in {METADATA_DIR / self.tag.value / self.metadata_version}" + else: + stats_path = self.dataset_path / LEROBOT_RELATIVE_STATS_FILE_NAME + + # Try to load existing relative stats + if stats_path.exists(): + print(f"Loading relative action stats from {stats_path}") + with open(stats_path, "r") as f: + stats: dict = json.load(f) + for name, stat in stats.items(): + stats[name] = DatasetStatisticalValues.model_validate(stat) + return stats + + # Calculate relative stats if file doesn't exist + print(f"Relative stats file not found at {stats_path}") + print(f"Calculating relative action stats for {self.dataset_name}") + + # Get action keys from modality configs, filtered by relative_action_keys + all_action_keys = self.modality_configs.get("action", ModalityConfig(delta_indices=[0], modality_keys=[])).modality_keys + if not all_action_keys: + print("No action keys found in modality configs, skipping relative stats calculation") + return {} + + # Filter to only the keys that should use relative action + action_keys_to_process = [] + for key in all_action_keys: + subkey = key.replace("action.", "") + if self.relative_action_keys is None or subkey in self.relative_action_keys: + action_keys_to_process.append(subkey) + + if not action_keys_to_process: + print("No action keys to process for relative stats") + return {} + + print(f"Will calculate relative stats for: {action_keys_to_process}") + + stats = {} + for action_key in action_keys_to_process: + print(f"Calculating relative stats for action key: {action_key}") + try: + relative_stats = self._calculate_relative_stats_for_key(action_key) + stats[action_key] = relative_stats + except Exception as e: + print(f"Failed to calculate relative stats for {action_key}: {e}") + continue + + if stats: + # Save the calculated stats + stats_serialized = {k: v.model_dump(mode="json") for k, v in stats.items()} + # Only save to dataset path (not global metadata path) + save_path = self.dataset_path / LEROBOT_RELATIVE_STATS_FILE_NAME + print(f"Saving relative action stats to {save_path}") + with open(save_path, "w") as f: + json.dump(stats_serialized, f, indent=4) + + return stats + + def _get_lerobot_relative_horizon_stats_meta(self) -> dict[str, dict[str, list]]: + """Get the per-horizon relative action stats metadata for the LeRobot dataset. + + Similar to _get_lerobot_relative_stats_meta but calculates separate stats for each + action horizon index. Will load from file if exists, otherwise calculate and save. + + Returns: + dict[str, dict[str, list]]: Nested dictionary where: + - Outer key is the action key (e.g., 'joint_position') + - Inner key is the stat name (e.g., 'max', 'min', 'mean', 'std', 'q01', 'q99') + - Value is a list of stat values per horizon index + """ + # Determine the path for per-horizon relative stats file + if self.use_global_metadata: + assert ( + self.metadata_version is not None + ), "metadata_version must be provided if use_global_metadata is True" + stats_path = ( + METADATA_DIR + / self.tag.value + / self.metadata_version + / Path(LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME).name + ) + assert ( + stats_path.exists() + ), f"Please provide a {Path(LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME).name} file in {METADATA_DIR / self.tag.value / self.metadata_version}" + else: + stats_path = self.dataset_path / LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME + + # Try to load existing per-horizon relative stats + if stats_path.exists(): + print(f"Loading per-horizon relative action stats from {stats_path}") + with open(stats_path, "r") as f: + stats: dict = json.load(f) + return stats + + # Calculate per-horizon relative stats if file doesn't exist + print(f"Per-horizon relative stats file not found at {stats_path}") + print(f"Calculating per-horizon relative action stats for {self.dataset_name}") + + # Get action keys from modality configs, filtered by relative_action_keys + all_action_keys = self.modality_configs.get("action", ModalityConfig(delta_indices=[0], modality_keys=[])).modality_keys + if not all_action_keys: + print("No action keys found in modality configs, skipping per-horizon relative stats calculation") + return {} + + # Filter to only the keys that should use relative action + action_keys_to_process = [] + for key in all_action_keys: + subkey = key.replace("action.", "") + if self.relative_action_keys is None or subkey in self.relative_action_keys: + action_keys_to_process.append(subkey) + + if not action_keys_to_process: + print("No action keys to process for per-horizon relative stats") + return {} + + print(f"Will calculate per-horizon relative stats for: {action_keys_to_process}") + + stats = {} + for action_key in action_keys_to_process: + print(f"Calculating per-horizon relative stats for action key: {action_key}") + try: + relative_stats = self._calculate_relative_stats_for_key_per_horizon(action_key) + stats[action_key] = relative_stats + except Exception as e: + print(f"Failed to calculate per-horizon relative stats for {action_key}: {e}") + continue + + if stats: + # Only save to dataset path (not global metadata path) + save_path = self.dataset_path / LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME + print(f"Saving per-horizon relative action stats to {save_path}") + with open(save_path, "w") as f: + json.dump(stats, f, indent=4) + + return stats + + def _calculate_relative_stats_for_key(self, action_key: str) -> DatasetStatisticalValues: + """Calculate relative action statistics for a specific action key. + + Args: + action_key: The action key to calculate stats for (e.g., 'joint_position') + + Returns: + DatasetStatisticalValues: The calculated statistics for the relative action. + """ + # Get state and action metadata from lerobot modality config + state_key = action_key # Assume state key matches action key + + # Get the modality metadata to find original column names and indices + state_meta = self.lerobot_modality_meta.state.get(state_key) + action_meta = self.lerobot_modality_meta.action.get(action_key) + + if state_meta is None: + raise ValueError(f"State key '{state_key}' not found in modality metadata") + if action_meta is None: + raise ValueError(f"Action key '{action_key}' not found in modality metadata") + + # Get the original column names (e.g., 'observation.state', 'action') + state_original_key = state_meta.original_key + action_original_key = action_meta.original_key + + # Get the indices to slice from the concatenated vectors + state_start, state_end = state_meta.start, state_meta.end + action_start, action_end = action_meta.start, action_meta.end + + state_delta_indices = self.modality_configs.get("state", ModalityConfig(delta_indices=[0], modality_keys=[])).delta_indices + action_delta_indices = self.modality_configs["action"].delta_indices + + print(f"Calculating relative stats for {action_key}:") + print(f" State: column='{state_original_key}', indices=[{state_start}:{state_end}]") + print(f" Action: column='{action_original_key}', indices=[{action_start}:{action_end}]") + + # # Calculate relative actions for all trajectories + all_relative_actions = [] + + # for traj_id in tqdm(self.trajectory_ids, desc=f"Calculating relative stats for {action_key}"): + max_trajs_for_stats = 10000 + traj_ids_to_process = self.trajectory_ids + if len(traj_ids_to_process) > max_trajs_for_stats: + # Randomly sample 500 trajectories + rng = np.random.default_rng(seed=42) + sampled_indices = rng.choice(len(traj_ids_to_process), size=max_trajs_for_stats, replace=False) + traj_ids_to_process = traj_ids_to_process[sampled_indices] + print(f"Sampling {max_trajs_for_stats} trajectories out of {len(self.trajectory_ids)} for stats calculation") + + # Calculate relative actions for sampled trajectories + all_relative_actions = [] + + for traj_id in tqdm(traj_ids_to_process, desc=f"Calculating relative stats for {action_key}"): + try: + # Load trajectory data + traj_data = self._load_trajectory_data(traj_id) + if traj_data is None: + continue + + # Check if columns exist + if state_original_key not in traj_data.columns or action_original_key not in traj_data.columns: + print(f"Missing columns: state='{state_original_key}' exists={state_original_key in traj_data.columns}, " + f"action='{action_original_key}' exists={action_original_key in traj_data.columns}") + continue + + # Load full state and action arrays, then slice to get the specific component + full_state_data = np.stack(traj_data[state_original_key].values) + full_action_data = np.stack(traj_data[action_original_key].values) + + # Slice to get just the component we care about (e.g., joint_position) + state_data = full_state_data[:, state_start:state_end] + action_data = full_action_data[:, action_start:action_end] + + # Calculate usable length based on action delta indices + usable_length = len(traj_data) - max(action_delta_indices) + + for i in range(usable_length): + # Get reference state (last state before action chunk) + ref_state_idx = state_delta_indices[-1] + i + if ref_state_idx >= len(state_data): + continue + ref_state = state_data[ref_state_idx] + + # Get action chunk + action_indices = [idx + i for idx in action_delta_indices] + if max(action_indices) >= len(action_data): + continue + actions = action_data[action_indices] + + # print("actions shape", actions.shape, "ref_state shape", ref_state.shape) + + # Calculate relative actions (action - reference state) + relative_actions = actions - ref_state + all_relative_actions.extend(relative_actions) + + except Exception as e: + print(f"Error processing trajectory {traj_id}: {e}") + continue + + if not all_relative_actions: + raise ValueError(f"No relative actions calculated for {action_key}") + + all_relative_actions = np.array(all_relative_actions) + print(f"Collected {len(all_relative_actions)} relative action samples for {action_key}") + + return DatasetStatisticalValues( + max=np.max(all_relative_actions, axis=0).tolist(), + min=np.min(all_relative_actions, axis=0).tolist(), + mean=np.mean(all_relative_actions, axis=0).tolist(), + std=np.std(all_relative_actions, axis=0).tolist(), + q01=np.quantile(all_relative_actions, 0.01, axis=0).tolist(), + q99=np.quantile(all_relative_actions, 0.99, axis=0).tolist(), + ) + + def _calculate_relative_stats_for_key_per_horizon( + self, action_key: str + ) -> dict[str, list]: + """Calculate relative action statistics for each delta index (horizon step) separately. + + Unlike `_calculate_relative_stats_for_key` which pools all horizon steps together, + this method calculates separate statistics for each action horizon index. + + Args: + action_key: The action key to calculate stats for (e.g., 'joint_position') + + Returns: + dict[str, list]: Dictionary where keys are stat names (max, min, mean, std, q01, q99) + and values are lists of stat values per horizon index. + Format: {"max": [[h0_vals], [h1_vals], ...], "min": [...], ...} + """ + # Get state and action metadata from lerobot modality config + state_key = action_key # Assume state key matches action key + + # Get the modality metadata to find original column names and indices + state_meta = self.lerobot_modality_meta.state.get(state_key) + action_meta = self.lerobot_modality_meta.action.get(action_key) + + if state_meta is None: + raise ValueError(f"State key '{state_key}' not found in modality metadata") + if action_meta is None: + raise ValueError(f"Action key '{action_key}' not found in modality metadata") + + # Get the original column names (e.g., 'observation.state', 'action') + state_original_key = state_meta.original_key + action_original_key = action_meta.original_key + + # Get the indices to slice from the concatenated vectors + state_start, state_end = state_meta.start, state_meta.end + action_start, action_end = action_meta.start, action_meta.end + + state_delta_indices = self.modality_configs.get("state", ModalityConfig(delta_indices=[0], modality_keys=[])).delta_indices + action_delta_indices = self.modality_configs["action"].delta_indices + + print(f"Calculating per-horizon relative stats for {action_key}:") + print(f" State: column='{state_original_key}', indices=[{state_start}:{state_end}]") + print(f" Action: column='{action_original_key}', indices=[{action_start}:{action_end}]") + print(f" Action delta indices: {action_delta_indices}") + + # Initialize separate lists for each horizon index + all_relative_actions_per_horizon: dict[int, list] = { + delta_idx: [] for delta_idx in action_delta_indices + } + + max_trajs_for_stats = 10000 + traj_ids_to_process = self.trajectory_ids + if len(traj_ids_to_process) > max_trajs_for_stats: + # Randomly sample trajectories + rng = np.random.default_rng(seed=42) + sampled_indices = rng.choice(len(traj_ids_to_process), size=max_trajs_for_stats, replace=False) + traj_ids_to_process = traj_ids_to_process[sampled_indices] + print(f"Sampling {max_trajs_for_stats} trajectories out of {len(self.trajectory_ids)} for stats calculation") + + for traj_id in tqdm(traj_ids_to_process, desc=f"Calculating per-horizon relative stats for {action_key}"): + try: + # Load trajectory data + traj_data = self._load_trajectory_data(traj_id) + if traj_data is None: + continue + + # Check if columns exist + if state_original_key not in traj_data.columns or action_original_key not in traj_data.columns: + continue + + # Load full state and action arrays, then slice to get the specific component + full_state_data = np.stack(traj_data[state_original_key].values) + full_action_data = np.stack(traj_data[action_original_key].values) + + # Slice to get just the component we care about (e.g., joint_position) + state_data = full_state_data[:, state_start:state_end] + action_data = full_action_data[:, action_start:action_end] + + # Calculate usable length based on action delta indices + usable_length = len(traj_data) - max(action_delta_indices) + + for i in range(usable_length): + # Get reference state (last state before action chunk) + ref_state_idx = state_delta_indices[-1] + i + if ref_state_idx >= len(state_data): + continue + ref_state = state_data[ref_state_idx] + + # Get action for each horizon index separately + for delta_idx in action_delta_indices: + action_idx = delta_idx + i + if action_idx >= len(action_data): + continue + action = action_data[action_idx] + + # Calculate relative action (action - reference state) + relative_action = action - ref_state + all_relative_actions_per_horizon[delta_idx].append(relative_action) + + except Exception as e: + print(f"Error processing trajectory {traj_id}: {e}") + continue + + # Calculate stats for each horizon index and organize by stat name + stat_names = ["max", "min", "mean", "std", "q01", "q99"] + stats_by_name: dict[str, list] = {name: [] for name in stat_names} + + for delta_idx in action_delta_indices: + relative_actions = all_relative_actions_per_horizon[delta_idx] + if not relative_actions: + print(f"Warning: No relative actions calculated for {action_key} at horizon index {delta_idx}") + # Add empty/placeholder values + for name in stat_names: + stats_by_name[name].append([]) + continue + + relative_actions_array = np.array(relative_actions) + print(f"Collected {len(relative_actions_array)} relative action samples for {action_key} at horizon {delta_idx}") + + stats_by_name["max"].append(np.max(relative_actions_array, axis=0).tolist()) + stats_by_name["min"].append(np.min(relative_actions_array, axis=0).tolist()) + stats_by_name["mean"].append(np.mean(relative_actions_array, axis=0).tolist()) + stats_by_name["std"].append(np.std(relative_actions_array, axis=0).tolist()) + stats_by_name["q01"].append(np.quantile(relative_actions_array, 0.01, axis=0).tolist()) + stats_by_name["q99"].append(np.quantile(relative_actions_array, 0.99, axis=0).tolist()) + + return stats_by_name + + def get_relative_stats_per_horizon( + self, + action_keys: list[str] | None = None, + save_to_file: bool = True, + ) -> dict[str, dict[str, list]]: + """Get relative action stats calculated separately for each horizon index. + + This is useful when you want different normalization for different action horizon + steps, e.g., near-future actions vs far-future actions might have different distributions. + + Args: + action_keys: List of action keys to calculate stats for. If None, uses + relative_action_keys (all action keys except gripper by default). + save_to_file: Whether to save the calculated stats to a file. + + Returns: + dict[str, dict[str, list]]: Nested dictionary where: + - Outer key is the action key (e.g., 'joint_position') + - Inner key is the stat name (e.g., 'max', 'min', 'mean', 'std', 'q01', 'q99') + - Value is a list of stat values per horizon index + + Example output format: + { + "joint_position": { + "max": [[h0_vals], [h1_vals], ...], + "min": [[h0_vals], [h1_vals], ...], + ... + } + } + """ + # Determine which action keys to process + if action_keys is None: + all_action_keys = self.modality_configs.get( + "action", ModalityConfig(delta_indices=[0], modality_keys=[]) + ).modality_keys + if not all_action_keys: + print("No action keys found in modality configs") + return {} + # Default: apply to all action keys except those containing 'gripper' + action_keys = [ + k.replace("action.", "") for k in all_action_keys + if "gripper" not in k.lower() + ] + + if not action_keys: + print("No action keys to process for per-horizon relative stats") + return {} + + print(f"Calculating per-horizon relative stats for: {action_keys}") + + all_stats: dict[str, dict[str, list]] = {} + + for action_key in action_keys: + print(f"Processing action key: {action_key}") + try: + stats_per_horizon = self._calculate_relative_stats_for_key_per_horizon(action_key) + all_stats[action_key] = stats_per_horizon + except Exception as e: + print(f"Failed to calculate per-horizon relative stats for {action_key}: {e}") + continue + + if save_to_file and all_stats: + # Save to the designated file + save_path = self.dataset_path / LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME + print(f"Saving per-horizon relative action stats to {save_path}") + + with open(save_path, "w") as f: + json.dump(all_stats, f, indent=4) + + return all_stats + + def load_relative_stats_per_horizon( + self, + stats_path: Path | str | None = None, + ) -> dict[str, dict[str, list]]: + """Load pre-computed per-horizon relative stats from a file. + + Args: + stats_path: Path to the stats file. If None, uses the default path in the dataset. + + Returns: + dict[str, dict[str, list]]: Nested dictionary of stats per horizon. + """ + if stats_path is None: + stats_path = self.dataset_path / LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME + else: + stats_path = Path(stats_path) + + if not stats_path.exists(): + print(f"Per-horizon relative stats file not found at {stats_path}") + return {} + + print(f"Loading per-horizon relative stats from {stats_path}") + with open(stats_path, "r") as f: + all_stats = json.load(f) + + return all_stats + + def _load_trajectory_data(self, traj_id: int) -> pd.DataFrame | None: + """Load trajectory data from parquet file. + + Args: + traj_id: The trajectory ID to load. + + Returns: + pd.DataFrame or None if loading fails. + """ + try: + chunk_index = traj_id // self.chunk_size + parquet_path = self.dataset_path / f"data/chunk-{chunk_index:03d}/episode_{traj_id:06d}.parquet" + if not parquet_path.exists(): + # Try alternative pattern + parquet_files = list(self.dataset_path.glob(f"data/*/episode_{traj_id:06d}.parquet")) + if parquet_files: + parquet_path = parquet_files[0] + else: + return None + return pd.read_parquet(parquet_path) + except Exception: + return None + + def _get_step_filter(self) -> dict[int, np.ndarray]: + """Get the step filter for the dataset.""" + step_filter_path = self.dataset_path / STEP_FILTER_FILENAME + step_filter = {} + if step_filter_path.exists(): + with open(step_filter_path, "r") as f: + for line in f: + episode_step_filter = json.loads(line) + trajectory_id = episode_step_filter["episode_index"] + all_indices = np.arange(self.trajectory_lengths[trajectory_id].item()) + indices_to_filter = np.array(episode_step_filter["step_indices"]) + step_filter[trajectory_id] = np.setdiff1d(all_indices, indices_to_filter) + else: + for trajectory_id in self.trajectory_ids: + step_filter[trajectory_id] = np.arange( + self.trajectory_lengths[trajectory_id].item() + ) + return step_filter + + def _get_metadata(self) -> DatasetMetadata: + """Get the metadata for the dataset. + + Returns: + dict: The metadata for the dataset. + """ + + # 1. Modality metadata + # 1.1. State and action modalities + simplified_modality_meta: dict[str, dict] = {} + for modality in ["state", "action"]: + simplified_modality_meta[modality] = {} + le_state_action_meta: dict[str, LeRobotStateActionMetadata] = getattr( + self.lerobot_modality_meta, modality + ) + for subkey in le_state_action_meta: + state_action_dtype = np.dtype(le_state_action_meta[subkey].dtype) + if np.issubdtype(state_action_dtype, np.floating): + continuous = True + else: + continuous = False + simplified_modality_meta[modality][subkey] = { + "absolute": le_state_action_meta[subkey].absolute, + "rotation_type": le_state_action_meta[subkey].rotation_type, + "shape": [ + le_state_action_meta[subkey].end - le_state_action_meta[subkey].start + ], + "continuous": continuous, + } + + # 1.2. Video modalities + le_info_path = self.dataset_path / LE_ROBOT_INFO_FILENAME + assert ( + le_info_path.exists() + ), f"Please provide a {LE_ROBOT_INFO_FILENAME} file in {self.dataset_path}" + with open(le_info_path, "r") as f: + le_info = json.load(f) + simplified_modality_meta["video"] = {} + for new_key in self.lerobot_modality_meta.video: + original_key = self.lerobot_modality_meta.video[new_key].original_key + if original_key is None: + original_key = new_key + le_video_meta = le_info["features"][original_key] + height = le_video_meta["shape"][le_video_meta["names"].index("height")] + width = le_video_meta["shape"][le_video_meta["names"].index("width")] + # NOTE(FH): different lerobot dataset versions have different keys for the number of channels and fps + try: + channels = le_video_meta["shape"][le_video_meta["names"].index("channel")] + fps = le_video_meta["video_info"]["video.fps"] + except (ValueError, KeyError): + # channels = le_video_meta["shape"][le_video_meta["names"].index("channels")] + channels = le_video_meta["info"]["video.channels"] + fps = le_video_meta["info"]["video.fps"] + simplified_modality_meta["video"][new_key] = { + "resolution": [width, height], + "channels": channels, + "fps": fps, + } + + # 2. Dataset statistics + dataset_statistics = {} + le_statistics = {k: v.model_dump() for k, v in self.lerobot_stats_meta.items()} + # Prepare relative stats if available + relative_stats = {} + if self.relative_action and hasattr(self, '_lerobot_relative_stats_meta'): + relative_stats = {k: v.model_dump() for k, v in self._lerobot_relative_stats_meta.items()} + + # Prepare per-horizon relative stats if available + per_horizon_stats = {} + if self.relative_action_per_horizon and hasattr(self, '_lerobot_relative_horizon_stats_meta'): + per_horizon_stats = self._lerobot_relative_horizon_stats_meta + + for our_modality in ["state", "action"]: + dataset_statistics[our_modality] = {} + for subkey in simplified_modality_meta[our_modality]: + dataset_statistics[our_modality][subkey] = {} + state_action_meta = self.lerobot_modality_meta.get_key_meta( + f"{our_modality}.{subkey}" + ) + assert isinstance(state_action_meta, LeRobotStateActionMetadata) + + # Check if we should use per-horizon relative stats for this action key + should_use_per_horizon = ( + our_modality == "action" + and self.relative_action_per_horizon + and subkey in per_horizon_stats + and (self.relative_action_keys is None or subkey in self.relative_action_keys) + ) + + # Use relative stats for action modality if relative_action is enabled and stats are available + # Also check if this subkey is in the list of keys that should use relative action + should_use_relative = ( + our_modality == "action" + and self.relative_action + and subkey in relative_stats + and (self.relative_action_keys is None or subkey in self.relative_action_keys) + ) + + if should_use_per_horizon: + # Use per-horizon relative action stats (format: {stat_name: [[h0_vals], [h1_vals], ...]}) + for stat_name in per_horizon_stats[subkey]: + dataset_statistics[our_modality][subkey][stat_name] = per_horizon_stats[subkey][stat_name] + print(f"Using per-horizon relative stats for {subkey}") + elif should_use_relative: + # Use relative action stats directly + for stat_name in relative_stats[subkey]: + dataset_statistics[our_modality][subkey][stat_name] = relative_stats[subkey][stat_name] + print(f"Using relative stats for {subkey}: {dataset_statistics[our_modality][subkey]}") + else: + # Use original absolute stats + le_modality = state_action_meta.original_key + for stat_name in le_statistics[le_modality]: + indices = np.arange( + state_action_meta.start, + state_action_meta.end, + ) + stat = np.array(le_statistics[le_modality][stat_name]) + dataset_statistics[our_modality][subkey][stat_name] = stat[indices].tolist() + + # 3. Full dataset metadata + metadata = DatasetMetadata( + statistics=dataset_statistics, # type: ignore + modalities=simplified_modality_meta, # type: ignore + embodiment_tag=self.tag, + ) + + return metadata + + def _get_trajectories(self) -> tuple[np.ndarray, np.ndarray]: + """Get the trajectories in the dataset.""" + # Get trajectory lengths, IDs, and whitelist from dataset metadata + episode_path = self.dataset_path / LE_ROBOT_EPISODE_FILENAME + with open(episode_path, "r") as f: + episode_metadata = [json.loads(line) for line in f] + trajectory_ids = [] + trajectory_lengths = [] + for episode in episode_metadata: + trajectory_ids.append(episode["episode_index"]) + trajectory_lengths.append(episode["length"]) + return np.array(trajectory_ids), np.array(trajectory_lengths) + + def _get_all_steps(self) -> list[tuple[int, int]]: + """Get the trajectory IDs and base indices for all steps in the dataset. + + Returns: + list[tuple[int, int]]: A list of (trajectory_id, base_index) tuples. + + Example: + self.trajectory_ids: [0, 1, 2] + self.step_filter: { + 0: [0, 1, 2], + 1: [0, 1], + 2: [0, 2, 3] + } + return: [ + (0, 0), (0, 1), (0, 2), + (1, 0), (1, 1), + (2, 0), (2, 2), (2, 3) + ] + """ + all_steps: list[tuple[int, int]] = [] + # All steps is used in single dataset, so we need to discard bad trajectories + # Mixture dataset directly use trajectory_ids, so we handle it by changing the sampling weights + discarded_episode_indices = [] + if self.discard_bad_trajectories: + discarded_episode_indices = self._lerobot_info_meta.get("discarded_episode_indices", []) + + for trajectory_id in self.trajectory_ids: + if trajectory_id in discarded_episode_indices: + continue + for base_index in self.step_filter[trajectory_id]: + all_steps.append((trajectory_id, base_index)) + return all_steps + + def _get_modality_keys(self) -> dict: + """Get the modality keys for the dataset. + + Returns: + dict: Dictionary mapping modality names to their keys. + """ + modality_keys = defaultdict(list) + for modality, config in self.modality_configs.items(): + modality_keys[modality] = config.modality_keys + return modality_keys + + def _get_delta_indices(self) -> dict[str, np.ndarray]: + """Restructure the delta indices to use modality.key as keys instead of just the modalities.""" + delta_indices: dict[str, np.ndarray] = {} + for config in self.modality_configs.values(): + for key in config.modality_keys: + delta_indices[key] = np.array(config.delta_indices) + return delta_indices + + def _get_data_path_pattern(self) -> str: + """Get the data path pattern for the LeRobot dataset.""" + return self.lerobot_info_meta["data_path"] + + def _get_video_path_pattern(self) -> str: + """Get the video path pattern for the LeRobot dataset.""" + return self.lerobot_info_meta["video_path"] + + def _get_chunk_size(self) -> int: + """Get the chunk size for the LeRobot dataset.""" + return self.lerobot_info_meta["chunks_size"] + + def _get_tasks(self) -> pd.DataFrame: + """Get the tasks for the dataset.""" + tasks_path = self.dataset_path / LE_ROBOT_TASKS_FILENAME + with open(tasks_path, "r") as f: + tasks = [json.loads(line) for line in f] + df = pd.DataFrame(tasks) + return df.set_index("task_index") + + def _get_task_embeddings(self) -> dict: + """Get the task embeddings for the dataset.""" + task_embeddings_path = self.dataset_path / LE_ROBOT_TASK_EMBEDDINGS_FILENAME + return torch.load(task_embeddings_path) + + def _get_detailed_global_instructions(self) -> dict[int, dict]: + """Get the detailed global instructions for the dataset. + + Loads from episodes_detail_global_instruction.jsonl if it exists. + + Returns: + dict[int, dict]: Mapping from episode_index to detailed instruction dict. + """ + detailed_instruction_path = self.dataset_path / LE_ROBOT_DETAILED_GLOBAL_INSTRUCTION_FILENAME + if not detailed_instruction_path.exists(): + return {} + with open(detailed_instruction_path, "r") as f: + instructions_list = [json.loads(line) for line in f] + return {entry["episode_index"]: entry for entry in instructions_list} + + def _check_integrity(self): + """Use the config to check if the keys are valid and detect silent data corruption.""" + ERROR_MSG_HEADER = f"Error occurred in initializing dataset {self.dataset_name}:\n" + + for modality, modality_config in self.modality_configs.items(): + if modality in ["lapa_action", "dream_actions", "rl_info", "task_embedding"]: + continue + for key in modality_config.modality_keys: + + if key == "action.task_progress": + continue + # Skip metadata-based language keys (they don't need modality metadata) + if modality == "language" and key.startswith("annotation."): + lang_subkey = key.replace("annotation.", "") + if lang_subkey in METADATA_LANG_KEYS: + continue + # Check if the key is valid + try: + self.lerobot_modality_meta.get_key_meta(key) + except Exception as e: + raise ValueError( + ERROR_MSG_HEADER + f"Unable to find key {key} in modality metadata:\n{e}" + ) + + def set_transforms_metadata(self, metadata: DatasetMetadata): + """Set the metadata for the transforms. This is useful for transforms that need to know the metadata, such as the normalization values.""" + self.transforms.set_metadata(metadata) + # Also set per-horizon statistics if available + if self.relative_action_per_horizon and hasattr(self, '_lerobot_relative_horizon_stats_meta'): + if hasattr(self.transforms, 'set_per_horizon_statistics'): + self.transforms.set_per_horizon_statistics(self._lerobot_relative_horizon_stats_meta) + + def set_epoch(self, epoch: int): + """Set the epoch for the dataset. + + Args: + epoch (int): The epoch to set. + """ + self.epoch = epoch + + def __len__(self) -> int: + """Get the total number of data points in the dataset. + + Returns: + int: the total number of data points in the dataset. + """ + return len(self.all_steps) + + def __str__(self) -> str: + """Get the description of the dataset.""" + return f"{self.dataset_name} ({len(self)} steps)" + + def __getitem__(self, index: int) -> dict: + """Get the data for a single step in a trajectory. + + Args: + index (int): The index of the step to get. + + Returns: + dict: The data for the step. + """ + trajectory_id, base_index = self.all_steps[index] + indices = { + key: delta_indices + base_index for key, delta_indices in self.delta_indices.items() + } + return self.transforms(self.get_step_data(trajectory_id, indices)) + + def get_step_data(self, trajectory_id: int, indices: dict[str, np.ndarray]) -> dict: + """Get the RAW data for a single step in a trajectory. No transforms are applied. + + Args: + trajectory_id (int): The name of the trajectory. + indices (dict[str, np.ndarray]): The indices for each modality. + + Returns: + dict: The RAW data for the step. + + Example return: + { + "video": { + "video.image_side_0": [B, T, H, W, C], + "video.image_side_1": [B, T, H, W, C], + }, + "state": { + "state.eef_position": [B, T, state_dim], + "state.eef_rotation": [B, T, state_dim], + }, + "action": { + "action.eef_position": [B, T, action_dim], + "action.eef_rotation": [B, T, action_dim], + }, + } + """ + data = {} + # Get the data for all modalities + self.curr_traj_data = self.get_trajectory_data(trajectory_id) + for modality in self.modality_keys: + # Get the data corresponding to each key in the modality + for key in self.modality_keys[modality]: + # Only load the data if the key is in the indices + if key in indices: + data[key] = self.get_data_by_modality( + trajectory_id, modality, key, indices[key] + ) + return data + + def get_parquet_path(self, trajectory_id: int) -> Path: + """Get the parquet path for a trajectory.""" + chunk_index = self.get_episode_chunk(trajectory_id) + return self.dataset_path / self.data_path_pattern.format( + episode_chunk=chunk_index, episode_index=trajectory_id + ) + + def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame: + """Get the data for a trajectory.""" + if self.curr_traj_id == trajectory_id and self.curr_traj_data is not None: + return self.curr_traj_data + else: + parquet_path = self.get_parquet_path(trajectory_id) + assert parquet_path.exists(), f"Parquet file not found at {parquet_path}" + return pd.read_parquet(parquet_path) + + def get_trajectory_index(self, trajectory_id: int) -> int: + """Get the index of the trajectory in the dataset by the trajectory ID. + This is useful when you need to get the trajectory length or sampling weight corresponding to the trajectory ID. + + Args: + trajectory_id (str): The ID of the trajectory. + + Returns: + int: The index of the trajectory in the dataset. + """ + trajectory_indices = np.where(self.trajectory_ids == trajectory_id)[0] + if len(trajectory_indices) != 1: + raise ValueError( + f"Error finding trajectory index for {trajectory_id}, found {trajectory_indices=}" + ) + return trajectory_indices[0] + + def get_episode_chunk(self, ep_index: int) -> int: + """Get the chunk index for an episode index.""" + return ep_index // self.chunk_size + + def retrieve_data_and_pad( + self, + array: np.ndarray, + step_indices: np.ndarray, + max_length: int, + padding_strategy: str = "first_last", + ) -> np.ndarray: + """Retrieve the data from the dataset and pad it if necessary. + + Args: + array (np.ndarray): The array to retrieve the data from. + step_indices (np.ndarray): The step indices to retrieve the data for. + max_length (int): The maximum length of the trajectory. + padding_strategy (str): The padding strategy, either "first_last" or "zero". + "first_last" uses first/last step data for padding, "zero" uses zero padding. + + Returns: + np.ndarray: The retrieved and padded data. + """ + # Get the padding indices + front_padding_indices = step_indices < 0 + end_padding_indices = step_indices >= max_length + padding_positions = np.logical_or(front_padding_indices, end_padding_indices) + # Retrieve the data with the non-padding indices + # If there exists some padding, Given T step_indices, the shape of the retrieved data will be (T', ...) where T' < T + raw_data = array[step_indices[~padding_positions]] + assert isinstance(raw_data, np.ndarray), f"{type(raw_data)=}" + # This is the shape of the output, (T, ...) + if raw_data.ndim == 1: + expected_shape = (len(step_indices),) + else: + expected_shape = (len(step_indices), *array.shape[1:]) + + # Pad the data + output = np.zeros(expected_shape) + # Assign the non-padded data + output[~padding_positions] = raw_data + # If there exists some padding, pad the data + if padding_positions.any(): + if padding_strategy == "first_last": + # Use first / last step data to pad + front_padding_data = array[0] + end_padding_data = array[-1] + output[front_padding_indices] = front_padding_data + output[end_padding_indices] = end_padding_data + elif padding_strategy == "zero": + # Use zero padding + output[padding_positions] = 0 + else: + raise ValueError(f"Invalid padding strategy: {padding_strategy}") + return output + + def get_video_path(self, trajectory_id: int, key: str) -> Path: + """Get the video file path for a specific trajectory and video key. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The video key (without 'video.' prefix). + + Returns: + Path: Path to the video file. + """ + chunk_index = self.get_episode_chunk(trajectory_id) + original_key = self.lerobot_modality_meta.video[key].original_key + if original_key is None: + original_key = key + video_filename = self.video_path_pattern.format( + episode_chunk=chunk_index, episode_index=trajectory_id, video_key=original_key + ) + return self.dataset_path / video_filename + + def get_video( + self, + trajectory_id: int, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray: + """Get the video frames for a trajectory by a base index. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (str): The ID of the trajectory. + key (str): The key of the video. + base_index (int): The base index of the trajectory. + + Returns: + np.ndarray: The video frames for the trajectory and frame indices. Shape: (T, H, W, C) + """ + # print(f"{step_indices=}") + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Ensure the indices are within the valid range + # This is equivalent to padding the video with extra frames at the beginning and end + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1) + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + # Get the sub-key + key = key.replace("video.", "") + video_path = self.get_video_path(trajectory_id, key) + # Get the action/state timestamps for each frame in the video + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + assert "timestamp" in self.curr_traj_data.columns, f"No timestamp found in {trajectory_id=}" + timestamp: np.ndarray = self.curr_traj_data["timestamp"].to_numpy() + # Get the corresponding video timestamps from the step indices + video_timestamp = timestamp[step_indices] + + # try: + return get_frames_by_timestamps( + video_path.as_posix(), + video_timestamp, + video_backend=self.video_backend, + video_backend_kwargs=self.video_backend_kwargs, + ) + # except: + # self.video_backend = "torchvision_av" + # return get_frames_by_timestamps( + # video_path.as_posix(), + # video_timestamp, + # video_backend=self.video_backend, + # video_backend_kwargs=self.video_backend_kwargs, + # ) + + + def get_state_or_action( + self, + trajectory_id: int, + modality: str, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray: + """Get the state or action data for a trajectory by a base index. + If the step indices are out of range, pad with the data: + if the data is stored in absolute format, pad with the first or last step data; + otherwise, pad with zero. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data. + key (str): The key of the data. + base_index (int): The base index of the trajectory. + + Returns: + np.ndarray: The data for the trajectory and step indices. + """ + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + + # Note [YL]: this handles action.task_progress if specified + if key == "action.task_progress": + # Get frame_index array and apply proper bounds checking and padding + frame_index_array = self.curr_traj_data["frame_index"].to_numpy() + # Use retrieve_data_and_pad to handle out-of-bounds indices + frame_index = self.retrieve_data_and_pad( + array=frame_index_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last", # Use first/last for task progress + ) + # get the task progress by using "frame index / trajectory length" + progress = frame_index / max_length + progress = progress.reshape(-1, 1) + return progress + + assert key.startswith(modality + "."), f"{key} must start with {modality + '.'}, got {key}" + # Get the sub-key, e.g. state.joint_angles -> joint_angles + subkey = key.replace(modality + ".", "") + # Get the lerobot key + le_state_or_action_cfg = getattr(self.lerobot_modality_meta, modality) + le_key = le_state_or_action_cfg[subkey].original_key + if le_key is None: + le_key = subkey + # Get the data array, shape: (T, D) + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + assert le_key in self.curr_traj_data.columns, f"No {le_key} found in {trajectory_id=}" + data_array: np.ndarray = np.stack(self.curr_traj_data[le_key]) # type: ignore + if data_array.ndim == 1: + assert ( + data_array.shape[0] == max_length + ), f"Expected 1D array with length {max_length}, got {data_array.shape} array" + data_array = data_array.reshape(-1, 1) + assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" + le_indices = np.arange( + le_state_or_action_cfg[subkey].start, + le_state_or_action_cfg[subkey].end, + ) + data_array = data_array[:, le_indices] + # Get the state or action configuration + state_or_action_cfg = getattr(self.metadata.modalities, modality)[subkey] + + # Pad the data + return self.retrieve_data_and_pad( + array=data_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last" if state_or_action_cfg.absolute else "zero", + ) + + def get_lapa_action( + self, + trajectory_id: int, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray | None: + """Get LAPA action data for a trajectory by step indices. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The key of the LAPA action data. + step_indices (np.ndarray): The step indices to retrieve data for. + + Returns: + np.ndarray | None: The LAPA action data, or None if the key is not found. + """ + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + # Check key in the current trajectory data + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + if ( + key not in self.curr_traj_data.columns + ): # this ensures that we can still load data w/o lapa actions. will store values that are None. + return None + # assert key in self.curr_traj_data.columns, f"{key} not found in {trajectory_id=}" + # Get the data array, shape: (T, D) + data_array: np.ndarray = np.stack(self.curr_traj_data[key]) # type: ignore + assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" + # Pad the data + return self.retrieve_data_and_pad( + array=data_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last", + ) + + def get_dream_actions( + self, + trajectory_id: int, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray | None: + """Get DREAM action data for a trajectory by step indices. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The key of the DREAM action data. + step_indices (np.ndarray): The step indices to retrieve data for. + + Returns: + np.ndarray | None: The DREAM action data, or None if the key is not found. + """ + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + # Check key in the current trajectory data + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + if ( + key not in self.curr_traj_data.columns + ): # this ensures that we can still load data w/o lapa actions. will store values that are None. + return None + # assert key in self.curr_traj_data.columns, f"{key} not found in {trajectory_id=}" + # Get the data array, shape: (T, D) + data_array: np.ndarray = np.stack(self.curr_traj_data[key]) # type: ignore + assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" + # Pad the data + return self.retrieve_data_and_pad( + array=data_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last", + ) + + def get_language( + self, + trajectory_id: int, + key: str, + step_indices: np.ndarray, + ) -> list[str]: + """Get the language annotation data for a trajectory by step indices. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The key of the annotation. + step_indices (np.ndarray): The step indices to retrieve data for. + + Returns: + list[str]: The annotation data for the trajectory and step indices. + If no matching data is found, return empty strings. + """ + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + # Get the end times corresponding to the closest indices + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, max_length - 1) + # Get the annotations + assert key.startswith( + "annotation." + ), f"Language key must start with 'annotation.', got {key}" + subkey = key.replace("annotation.", "") + # print("subkey", subkey) + + # Check if this is a metadata-based language key (detailed_global_instruction_medium/concise) + if subkey in METADATA_LANG_KEYS: + # print("return metadata language") + return self._get_language_from_metadata(trajectory_id, subkey, len(step_indices)) + + # Otherwise, load from parquet columns (original behavior) + annotation_meta = self.lerobot_modality_meta.annotation + assert annotation_meta is not None, f"Annotation metadata is None for {subkey}" + assert ( + subkey in annotation_meta + ), f"Annotation key {subkey} not found in metadata, available annotation keys: {annotation_meta.keys()}" + subkey_meta = annotation_meta[subkey] + original_key = subkey_meta.original_key + if original_key is None: + original_key = key + if pd.api.types.is_numeric_dtype(self.curr_traj_data[original_key]): + # Stored as list of integers + task_indices: list[int] = self.curr_traj_data[original_key].iloc[step_indices].tolist() + return self.tasks.loc[task_indices]["task"].tolist() + else: + # Stored as list of strings + return self.curr_traj_data[original_key].iloc[step_indices].astype(str).tolist() + + def _get_language_from_metadata( + self, + trajectory_id: int, + lang_key: str, + nframes: int, + ) -> list[str]: + """Get language instruction from metadata files for special language keys. + + Supports: + - detailed_global_instruction_medium: Longer, detailed description + - detailed_global_instruction_concise: Short summary + + Args: + trajectory_id (int): The ID of the trajectory (episode_index). + lang_key (str): The language key (e.g., "detailed_global_instruction_medium"). + nframes (int): Number of frames to return the instruction for. + + Returns: + list[str]: The instruction repeated for each frame (empty string if not found). + """ + if trajectory_id in self._detailed_global_instructions: + instruction = self._detailed_global_instructions[trajectory_id].get(lang_key, "") + # print("instruction", instruction) + else: + instruction = "" + return [instruction] * nframes + + def get_rl_info( + self, + trajectory_id: int, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray: + """Get the reward data for a trajectory by step indices. + + If the step indices are out of range, pad with first/last step data. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The key of the reward data. + step_indices (np.ndarray): The step indices to retrieve data for. + + Returns: + np.ndarray: The reward data for the trajectory and step indices. + """ + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + data_array: np.ndarray = np.stack(self.curr_traj_data[key]) # type: ignore + + if key == "rl_info.next.reward": + padding_strategy = "zero" + else: + padding_strategy = "first_last" + + # Pad the data + return self.retrieve_data_and_pad( + array=data_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy=padding_strategy, + ) + + def get_data_by_modality( + self, + trajectory_id: int, + modality: str, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray | list[str] | None: + """Get the data corresponding to the modality for a trajectory by step indices. + + This method dispatches to the appropriate specialized method based on the modality. + For the language modality, empty strings are returned if no matching data is found. + + Args: + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data (video, state, action, language, etc.). + key (str): The key of the data. + step_indices (np.ndarray): The step indices of the trajectory. + + Returns: + np.ndarray | list[str] | None: The data for the specified modality. + """ + if modality == "video": + return self.get_video(trajectory_id, key, step_indices) + elif modality == "state" or modality == "action": + return self.get_state_or_action(trajectory_id, modality, key, step_indices) + elif modality == "language": + return self.get_language(trajectory_id, key, step_indices) + elif modality == "lapa_action": + return self.get_lapa_action(trajectory_id, key, step_indices) + elif modality == "dream_actions": + return self.get_dream_actions(trajectory_id, key, step_indices) + elif modality == "rl_info": + return self.get_rl_info(trajectory_id, key, step_indices) + else: + raise ValueError(f"Invalid modality: {modality}") + + def get_initial_actions(self): + """Load initial actions from the dataset if available. + + Returns: + list: List containing initial actions if the file exists, empty list otherwise. + """ + initial_actions_path = self.dataset_path / INITIAL_ACTIONS_FILENAME + if initial_actions_path.exists(): + initial_actions = load_initial_actions(initial_actions_path) + return initial_actions # a single-element list of dict[str, dict[str, np.ndarray]] + else: + return [] + + + +class CachedLeRobotSingleDataset(LeRobotSingleDataset): + """A cached version of LeRobotSingleDataset that preloads all video frames into memory. + + This class caches video frames for each trajectory and key to improve access speed + when video frames need to be accessed multiple times. Recommended for small datasets + or when memory usage is not a concern. + """ + + def __init__(self, *args, **kwargs): + """Initialize the cached dataset and preload all video frames. + + Args: + *args: Arguments passed to parent LeRobotSingleDataset. + **kwargs: Keyword arguments passed to parent LeRobotSingleDataset. + """ + # Initialize img_resize attribute first to ensure it exists + super().__init__(*args, **kwargs) + cached_frames: dict[str, np.ndarray] = {} + + for key in self.modality_keys["video"]: + all_frames = [] + key = key.replace("video.", "") + for trajectory_id, trajectory_length in tqdm( + zip(self.trajectory_ids, self.trajectory_lengths), + total=len(self.trajectory_ids), + desc=f"Caching {key} frames", + ): + video_path = self.get_video_path(trajectory_id, key) + frames, _ = get_all_frames( + video_path.as_posix(), + video_backend=self.video_backend, + video_backend_kwargs=self.video_backend_kwargs, + ) + assert frames.ndim == 4, f"Expected 4D array, got {frames.shape} array" + assert frames.shape[3] == 3, f"Expected 3 channels, got {frames.shape[3]} channels" + assert ( + frames.shape[0] == trajectory_length + ), f"Expected {trajectory_length} frames, got {frames.shape[0]} frames" + all_frames.append(frames) + cached_frames[key] = np.concatenate(all_frames, axis=0) + print(f"{key}: {cached_frames[key].shape}") + self.cached_frames = cached_frames + self.start_indices = np.cumsum(self.trajectory_lengths) - self.trajectory_lengths + + def get_video(self, trajectory_id: int, key: str, step_indices: np.ndarray) -> np.ndarray: + """Get video frames from the cached data. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The video key (with 'video.' prefix). + step_indices (np.ndarray): The step indices to retrieve frames for. + + Returns: + np.ndarray: The video frames with shape (T, H, W, C). + """ + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Ensure the indices are within the valid range + # This is equivalent to padding the video with extra frames at the beginning and end + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1) + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + # Get the sub-key + key = key.replace("video.", "") + # Calculate the absolute indices + absolute_indices = self.start_indices[trajectory_index] + step_indices + return self.cached_frames[key][absolute_indices] + + +def safe_hash(input_tuple): + """Generate a safe hash from an input tuple. + + Creates a deterministic hash using SHA256 and returns the lower 128 bits. + This is used for deterministic random seed generation. + + Args: + input_tuple: The tuple to hash. + + Returns: + int: A 128-bit hash value. + """ + # keep 128 bits of the hash + tuple_string = repr(input_tuple).encode("utf-8") + sha256 = hashlib.sha256() + sha256.update(tuple_string) + + seed = int(sha256.hexdigest(), 16) + + return seed & 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF + + +class MixtureSpecElement(BaseModel): + """Specification element for a dataset mixture defining paths and weights. + + This class validates dataset paths by embodiment tag and handles weight distribution + across multiple dataset paths if requested. + """ + + dataset_path: dict[str, list[Path] | Path] = Field(..., description="The path to the dataset.") + dataset_weight: float = Field(..., description="The weight of the dataset in the mixture.") + distribute_weights: bool = Field( + default=False, + description="Whether to distribute the weights of the dataset across all the paths. If True, the weights will be evenly distributed across all the paths.", + ) + + @field_validator("dataset_path", mode="after") + def validate_dataset_path_keys(cls, v: dict[str, list[Path] | Path]) -> dict[str, list[Path]]: + """Validate dataset paths and expand glob patterns. + + Args: + v (dict[str, list[Path] | Path]): Dictionary mapping embodiment tags to paths. + + Returns: + dict[str, list[Path]]: Validated and expanded paths. + + Raises: + ValueError: If an invalid embodiment tag is provided. + """ + all_globbed_paths: dict[str, list[Path]] = {} + for embodiment_tag, paths in v.items(): + try: + _ = EmbodimentTag(embodiment_tag) + except ValueError: + raise ValueError(f"Invalid embodiment tag: {embodiment_tag}") + if isinstance(paths, Path): + paths = [paths] + globbed_paths = [] + for path in paths: + globbed_paths.extend(glob.glob(str(path))) + all_globbed_paths[embodiment_tag] = globbed_paths + return all_globbed_paths + + +class LeRobotMixtureDataset(Dataset): + """ + A mixture of multiple datasets. This class samples a single dataset based on the dataset weights and then calls the `__getitem__` method of the sampled dataset. + It is recommended to modify the single dataset class instead of this class. + """ + + def __init__( + self, + data_mixture: Sequence[tuple[LeRobotSingleDataset, float]], + training: bool, + balance_dataset_weights: bool = True, + balance_trajectory_weights: bool = True, + seed: int = 42, + allow_padding_at_end: bool = False, + metadata_config: dict = { + "percentile_mixing_method": "min_max", + }, + ): + """ + Initialize the mixture dataset. + + Args: + data_mixture (list[tuple[LeRobotSingleDataset, float]]): Datasets and their corresponding weights. + training (bool): If True, __getitem__ will return different samples every epoch; if False, __getitem__ will return the same sample every epoch. + balance_dataset_weights (bool): If True, the weight of dataset will be multiplied by the total trajectory length of each dataset. + balance_trajectory_weights (bool): If True, sample trajectories within a dataset weighted by their length; otherwise, use equal weighting. + seed (int): Random seed for sampling. + allow_padding_at_end (bool): If True, allow padding at the end of the dataset. + """ + datasets: list[LeRobotSingleDataset] = [] + dataset_sampling_weights: list[float] = [] + for dataset, weight in data_mixture: + datasets.append(dataset) + dataset_sampling_weights.append(weight) + self.datasets = datasets + self.balance_dataset_weights = balance_dataset_weights + self.balance_trajectory_weights = balance_trajectory_weights + self.seed = seed + self.training = training + self.allow_padding_at_end = allow_padding_at_end + + # Set properties for sampling + + # 1. Dataset lengths + self._dataset_lengths = np.array([len(dataset) for dataset in self.datasets]) + + # 2. Dataset sampling weights + self._dataset_sampling_weights = np.array(dataset_sampling_weights) + if self.balance_dataset_weights: + self._dataset_sampling_weights *= self._dataset_lengths + self._dataset_sampling_weights /= self._dataset_sampling_weights.sum() + + # 3. Trajectory sampling weights + self._trajectory_sampling_weights: list[np.ndarray] = [] + for dataset in self.datasets: + trajectory_sampling_weights = np.ones(len(dataset.trajectory_ids)) + if self.balance_trajectory_weights: + trajectory_sampling_weights *= np.array( + [ + len(dataset.step_filter[trajectory_id]) + for trajectory_id in dataset.trajectory_ids + ] + ) + + if dataset.discard_bad_trajectories: + bad_trajectory_indices = dataset.lerobot_info_meta.get( + "discarded_episode_indices", [] + ) + trajectory_sampling_weights[bad_trajectory_indices] = 0.0 + + if trajectory_sampling_weights.sum() == 0: + raise ValueError(f"No valid trajectories found for dataset {dataset}") + + trajectory_sampling_weights /= trajectory_sampling_weights.sum() + self._trajectory_sampling_weights.append(trajectory_sampling_weights) + + # 4. Primary dataset indices + self._primary_dataset_indices = np.array(dataset_sampling_weights) == 1.0 + + # Set the epoch and sample the first epoch + self.set_epoch(0) + + # Create a merged metadata for the mixture dataset (we don't need this in the future as eval will directly use `get_metadata`) + self.update_metadata(metadata_config) + + # Set the transforms to training or evaluation mode + if self.training: + for dataset in self.datasets: + dataset.transforms.train() + else: + for dataset in self.datasets: + dataset.transforms.eval() + + @property + def dataset_lengths(self) -> np.ndarray: + """The lengths of each dataset.""" + return self._dataset_lengths + + @property + def dataset_sampling_weights(self) -> np.ndarray: + """The sampling weights for each dataset.""" + return self._dataset_sampling_weights + + @property + def trajectory_sampling_weights(self) -> list[np.ndarray]: + """The sampling weights for each trajectory in each dataset.""" + return self._trajectory_sampling_weights + + @property + def primary_dataset_indices(self) -> np.ndarray: + """The indices of the primary datasets.""" + return self._primary_dataset_indices + + def __str__(self) -> str: + """Return a string representation of the mixture dataset with weights.""" + dataset_descriptions = [] + for dataset, weight in zip(self.datasets, self.dataset_sampling_weights): + dataset_description = { + "Dataset": str(dataset), + "Sampling weight": float(weight), + } + dataset_descriptions.append(dataset_description) + return yaml.dump({"Mixture dataset": dataset_descriptions}) + + @classmethod + def from_mixture_spec( + cls: type[T_LeRobotMixtureDataset], + mixture_spec: Sequence[MixtureSpecElement | dict], + dataset_class: type[LeRobotSingleDataset] | str, + all_modality_configs: dict[str, dict[str, ModalityConfig]], + all_transforms: dict[str, ComposedModalityTransform], + metadata_versions: dict[str, str], + fps: float = None, + dataset_kwargs: dict | None = None, + mixture_kwargs: dict | None = None, + ) -> T_LeRobotMixtureDataset: + """Initialize the mixture dataset from a specification. + + Args: + mixture_spec (Sequence[MixtureSpecElement | dict]): The specification for the mixture dataset. + dataset_class (type[LeRobotSingleDataset] | str): The dataset class or its string path. + all_modality_configs (dict[str, dict[str, ModalityConfig]]): The modality configs for each embodiment. + all_transforms (dict[str, ComposedModalityTransform]): The transforms for each embodiment. + metadata_versions (dict[str, str]): The metadata versions for each embodiment. + dataset_kwargs (dict | None): Additional keyword arguments for the dataset classes. + mixture_kwargs (dict | None): Additional keyword arguments for the mixture dataset. + + Returns: + LeRobotMixtureDataset: The initialized mixture dataset. + """ + if isinstance(dataset_class, str): + module_name, class_name = dataset_class.rsplit(".", 1) + module = importlib.import_module(module_name) + dataset_class = getattr(module, class_name) + assert not isinstance(dataset_class, str), f"{dataset_class} is a string" + assert issubclass( + dataset_class, LeRobotSingleDataset + ), f"{dataset_class} is not a subclass of LeRobotSingleDataset" + data_mixture = [] + + for dataset_spec in tqdm( + mixture_spec, + total=len(mixture_spec), + desc="Initializing datasets", + ): + start_time = time.time() + if isinstance(dataset_spec, dict): + dataset_spec = MixtureSpecElement.model_validate(dataset_spec) + datasets = [] + for embodiment_tag, paths in dataset_spec.dataset_path.items(): + if isinstance(paths, Path): + paths = [paths] + for dataset_path in paths: + if '.sh' in dataset_path or '.json' in dataset_path: + continue + assert ( + embodiment_tag in all_modality_configs + ), f"{embodiment_tag} not in modality_configs: {all_modality_configs.keys()}" + assert ( + embodiment_tag in all_transforms + ), f"{embodiment_tag} not in transforms: {all_transforms.keys()}" + dataset = dataset_class( + dataset_path=dataset_path, + embodiment_tag=EmbodimentTag(embodiment_tag), + modality_configs=copy.copy(all_modality_configs[embodiment_tag]), + transforms=copy.copy(all_transforms[embodiment_tag]), + metadata_version=metadata_versions[embodiment_tag], + fps=fps[embodiment_tag] if embodiment_tag in fps else None, + **(dataset_kwargs if dataset_kwargs is not None else {}), + ) + datasets.append(dataset) + dataset_lengths = np.array([len(dataset) for dataset in datasets]) + dataset_relative_lengths = dataset_lengths / dataset_lengths.sum() + for dataset, relative_length in zip(datasets, dataset_relative_lengths): + if dataset_spec.distribute_weights: + weight = relative_length * dataset_spec.dataset_weight + else: + weight = dataset_spec.dataset_weight + data_mixture.append((dataset, weight)) + + print( + f"Time taken to initialize {len(datasets)} datasets: {time.time() - start_time:.2f} seconds" + ) + + return cls( + data_mixture=data_mixture, + **(mixture_kwargs if mixture_kwargs is not None else {}), + ) + + def set_epoch(self, epoch: int): + """Set the epoch for the dataset. + + Args: + epoch (int): The epoch to set. + """ + self.epoch = epoch + # self.sampled_steps = self.sample_epoch() + + def sample_step(self, index: int) -> tuple[LeRobotSingleDataset, int, int]: + """Sample a single step from the mixture dataset. + + Args: + index (int): The index to sample (used for deterministic sampling). + + Returns: + tuple[LeRobotSingleDataset, int, int]: A tuple of (dataset, trajectory_id, step_index). + """ + # return self.sampled_steps[index] + + # Set seed + if self.training: + seed = safe_hash((self.epoch, index, self.seed)) + rng = np.random.default_rng(seed) + + # Sample dataset + dataset_index = rng.choice(len(self.datasets), p=self.dataset_sampling_weights) + dataset = self.datasets[dataset_index] + + if self.allow_padding_at_end: + # Sample trajectory + trajectory_index = rng.choice( + len(dataset.trajectory_ids), p=self.trajectory_sampling_weights[dataset_index] + ) + trajectory_id = dataset.trajectory_ids[trajectory_index] + + allowed_length = dataset.trajectory_lengths[trajectory_index] + else: + # Avoid padding at the end of the trajectory + max_delta_index = dataset.max_delta_index + trajectory_length = 0 + trajectory_id = None + while trajectory_length < max_delta_index + 1: + # Sample trajectory + trajectory_index = rng.choice( + len(dataset.trajectory_ids), + p=self.trajectory_sampling_weights[dataset_index], + ) + trajectory_id = dataset.trajectory_ids[trajectory_index] + trajectory_length = dataset.trajectory_lengths[trajectory_index] + assert trajectory_id is not None + + # Sample step + assert ( + trajectory_length >= max_delta_index + 1 + ), f"{trajectory_length=}, {max_delta_index=}" + allowed_length = trajectory_length - max_delta_index + # Get the allowed indices from the step filter + allowed_indices = dataset.step_filter[trajectory_id] + # Remove indices that are too large + allowed_indices = allowed_indices[allowed_indices <= allowed_length] + step_index = rng.choice(allowed_indices) + return dataset, trajectory_id, step_index + else: + length_cumsum = np.cumsum(self.dataset_lengths) + dataset_index = np.searchsorted(length_cumsum, index) + dataset = self.datasets[dataset_index] + assert ( + len(dataset._lerobot_info_meta.get("discarded_episode_indices", [])) == 0 + ), f"Find discarded episode indices in evaluation dataset {dataset.dataset_path}" + trajectory_id, step_index = dataset.all_steps[index - length_cumsum[dataset_index]] + return dataset, trajectory_id, step_index + + def __getitem__(self, index: int) -> dict: + """Get the data for a single trajectory and start index. + + Args: + index (int): The index of the trajectory to get. + + Returns: + dict: The data for the trajectory and start index. + """ + dataset, trajectory_id, step_index = self.sample_step(index) + indices = { + key: delta_indices + step_index for key, delta_indices in dataset.delta_indices.items() + } + return dataset.transforms(dataset.get_step_data(trajectory_id, indices)) + + def __len__(self) -> int: + """Get the length of a single epoch in the mixture. + + Returns: + int: The length of a single epoch in the mixture. + """ + if self.training: + return int((self.dataset_lengths * self.dataset_sampling_weights).sum()) + else: + return int(self.dataset_lengths.sum()) + + @staticmethod + def compute_overall_statistics( + per_task_stats: list[dict[str, dict[str, list[float] | np.ndarray]]], + dataset_sampling_weights: list[float] | np.ndarray, + percentile_mixing_method: str = "weighted_average", + ) -> dict[str, dict[str, list[float]]]: + """ + Computes overall statistics from per-task statistics using dataset sample weights. + + Args: + per_task_stats: List of per-task statistics. + Example format of one element in the per-task statistics list: + { + "state.gripper": { + "min": [...], + "max": [...], + "mean": [...], + "std": [...], + "q01": [...], + "q99": [...], + }, + ... + } + dataset_sampling_weights: List of sample weights for each task. + percentile_mixing_method: The method to mix the percentiles, either "weighted_average" or "weighted_std". + + Returns: + A dict of overall statistics per modality. + """ + # Normalize the sample weights to sum to 1 + dataset_sampling_weights = np.array(dataset_sampling_weights) + normalized_weights = dataset_sampling_weights / dataset_sampling_weights.sum() + + # Initialize overall statistics dict + overall_stats: dict[str, dict[str, list[float]]] = {} + + # Get the list of modality keys + modality_keys = per_task_stats[0].keys() + + for modality in modality_keys: + # Check if stats are per-horizon (2D) by examining the first task's mean + first_mean = np.array(per_task_stats[0][modality]["mean"]) + is_per_horizon = first_mean.ndim == 2 # Shape (horizon_len, action_dim) + + if is_per_horizon: + # Handle per-horizon stats (2D arrays) + stats_shape = first_mean.shape # (horizon_len, action_dim) + + # Initialize accumulators for means and variances + weighted_means = np.zeros(stats_shape) + weighted_squares = np.zeros(stats_shape) + + # Collect min, max, q01, q99 from all tasks + min_list = [] + max_list = [] + q01_list = [] + q99_list = [] + + for task_idx, task_stats in enumerate(per_task_stats): + w_i = normalized_weights[task_idx] + stats = task_stats[modality] + means = np.array(stats["mean"]) + stds = np.array(stats["std"]) + + # Update weighted sums for mean and variance + weighted_means += w_i * means + weighted_squares += w_i * (stds**2 + means**2) + + # Collect min, max, q01, q99 + min_list.append(np.array(stats["min"])) + max_list.append(np.array(stats["max"])) + q01_list.append(np.array(stats["q01"])) + q99_list.append(np.array(stats["q99"])) + + # Compute overall mean + overall_mean = weighted_means.tolist() + + # Compute overall variance and std deviation + overall_variance = weighted_squares - weighted_means**2 + overall_std = np.sqrt(np.maximum(overall_variance, 0)).tolist() + + # Compute overall min and max per dimension + # Stack along new axis: (num_tasks, horizon_len, action_dim) + overall_min = np.min(np.stack(min_list, axis=0), axis=0).tolist() + overall_max = np.max(np.stack(max_list, axis=0), axis=0).tolist() + + # Compute overall q01 and q99 per dimension + q01_array = np.stack(q01_list, axis=0) # (num_tasks, horizon_len, action_dim) + q99_array = np.stack(q99_list, axis=0) + if percentile_mixing_method == "weighted_average": + # Weighted average along task axis + weighted_q01 = np.average(q01_array, axis=0, weights=normalized_weights).tolist() + weighted_q99 = np.average(q99_array, axis=0, weights=normalized_weights).tolist() + elif percentile_mixing_method == "min_max": + weighted_q01 = np.min(q01_array, axis=0).tolist() + weighted_q99 = np.max(q99_array, axis=0).tolist() + else: + raise ValueError(f"Invalid percentile mixing method: {percentile_mixing_method}") + else: + # Handle regular stats (1D arrays) + num_dims = len(first_mean) + + # Initialize accumulators for means and variances + weighted_means = np.zeros(num_dims) + weighted_squares = np.zeros(num_dims) + + # Collect min, max, q01, q99 from all tasks + min_list = [] + max_list = [] + q01_list = [] + q99_list = [] + + for task_idx, task_stats in enumerate(per_task_stats): + w_i = normalized_weights[task_idx] + stats = task_stats[modality] + means = np.array(stats["mean"]) + stds = np.array(stats["std"]) + + # Update weighted sums for mean and variance + weighted_means += w_i * means + weighted_squares += w_i * (stds**2 + means**2) + + # Collect min, max, q01, q99 + min_list.append(stats["min"]) + max_list.append(stats["max"]) + q01_list.append(stats["q01"]) + q99_list.append(stats["q99"]) + + # Compute overall mean + overall_mean = weighted_means.tolist() + + # Compute overall variance and std deviation + overall_variance = weighted_squares - weighted_means**2 + overall_std = np.sqrt(np.maximum(overall_variance, 0)).tolist() + + # Compute overall min and max per dimension + overall_min = np.min(np.array(min_list), axis=0).tolist() + overall_max = np.max(np.array(max_list), axis=0).tolist() + + # Compute overall q01 and q99 per dimension + # Use weighted average of per-task quantiles + q01_array = np.array(q01_list) + q99_array = np.array(q99_list) + if percentile_mixing_method == "weighted_average": + weighted_q01 = np.average(q01_array, axis=0, weights=normalized_weights).tolist() + weighted_q99 = np.average(q99_array, axis=0, weights=normalized_weights).tolist() + elif percentile_mixing_method == "min_max": + weighted_q01 = np.min(q01_array, axis=0).tolist() + weighted_q99 = np.max(q99_array, axis=0).tolist() + else: + raise ValueError(f"Invalid percentile mixing method: {percentile_mixing_method}") + + # Store the overall statistics for the modality + overall_stats[modality] = { + "min": overall_min, + "max": overall_max, + "mean": overall_mean, + "std": overall_std, + "q01": weighted_q01, + "q99": weighted_q99, + } + + return overall_stats + + @staticmethod + def merge_metadata( + metadatas: list[DatasetMetadata], + dataset_sampling_weights: list[float], + percentile_mixing_method: str, + ) -> DatasetMetadata: + """Merge multiple metadata into one.""" + # Convert to dicts + metadata_dicts = [metadata.model_dump(mode="json") for metadata in metadatas] + # Create a new metadata dict + merged_metadata = {} + + # Check all metadata have the same embodiment tag + assert all( + metadata.embodiment_tag == metadatas[0].embodiment_tag for metadata in metadatas + ), "All metadata must have the same embodiment tag" + merged_metadata["embodiment_tag"] = metadatas[0].embodiment_tag + + # Merge the dataset statistics + dataset_statistics = {} + dataset_statistics["state"] = LeRobotMixtureDataset.compute_overall_statistics( + per_task_stats=[m["statistics"]["state"] for m in metadata_dicts], + dataset_sampling_weights=dataset_sampling_weights, + percentile_mixing_method=percentile_mixing_method, + ) + dataset_statistics["action"] = LeRobotMixtureDataset.compute_overall_statistics( + per_task_stats=[m["statistics"]["action"] for m in metadata_dicts], + dataset_sampling_weights=dataset_sampling_weights, + percentile_mixing_method=percentile_mixing_method, + ) + merged_metadata["statistics"] = dataset_statistics + + # Merge the modality configs + modality_configs = defaultdict(set) + for metadata in metadata_dicts: + for modality, configs in metadata["modalities"].items(): + modality_configs[modality].add(json.dumps(configs)) + merged_metadata["modalities"] = {} + for modality, configs in modality_configs.items(): + # Check that all modality configs correspond to the same tag matches + assert ( + len(configs) == 1 + ), f"Multiple modality configs for modality {modality}: {list(configs)}" + merged_metadata["modalities"][modality] = json.loads(configs.pop()) + + return DatasetMetadata.model_validate(merged_metadata) + + def update_metadata(self, metadata_config: dict) -> None: + """Merge multiple metadatas into one and set the transforms with the merged metadata. + + Args: + metadata_config (dict): Configuration for the metadata. + "percentile_mixing_method": The method to mix the percentiles, either "weighted_average" or "min_max". + weighted_average: Use the weighted average of the percentiles using the weight used in sampling the datasets. + min_max: Use the min of the 1st percentile and max of the 99th percentile. + """ + + self.merged_metadata: dict[str, DatasetMetadata] = {} + # Group metadata by tag + all_metadatas: dict[str, list[DatasetMetadata]] = {} + for dataset in self.datasets: + if dataset.tag.value not in all_metadatas: + all_metadatas[dataset.tag.value] = [] + all_metadatas[dataset.tag.value].append(dataset.metadata) + for tag, metadatas in all_metadatas.items(): + self.merged_metadata[tag] = self.merge_metadata( + metadatas=metadatas, + dataset_sampling_weights=self.dataset_sampling_weights.tolist(), + percentile_mixing_method=metadata_config["percentile_mixing_method"], + ) + for dataset in self.datasets: + dataset.set_transforms_metadata(self.merged_metadata[dataset.tag.value]) + + def get_initial_actions(self): + initial_actions = [] + for dataset in self.datasets: + if hasattr(dataset, "get_initial_actions"): + initial_actions.extend(dataset.get_initial_actions()) + return initial_actions diff --git a/groot/vla/data/dataset/lerobot_sharded.py b/groot/vla/data/dataset/lerobot_sharded.py new file mode 100644 index 0000000000000000000000000000000000000000..26d68e112b01e3cb20e850735cbc089bef918629 --- /dev/null +++ b/groot/vla/data/dataset/lerobot_sharded.py @@ -0,0 +1,1542 @@ +from concurrent.futures import Future, ThreadPoolExecutor +import json +from pathlib import Path +import time + +import numpy as np +import pandas as pd +import torch +import torch.distributed as dist +from torch.utils.data import IterableDataset, get_worker_info +import yaml + +from groot.vla.common.utils import get_frames_by_timestamps + +from .lerobot import LE_ROBOT_EPISODE_FILENAME, LeRobotMixtureDataset, LeRobotSingleDataset + + +class ShardedLeRobotSingleDataset(LeRobotSingleDataset): + """ + A single dataset with shards. + """ + + def __init__( + self, + *args, + num_steps_per_shard: int = int(1e7), + **kwargs, + ): + self.args = args + self.kwargs = kwargs + super().__init__(*args, **kwargs) + self.num_steps_per_shard = num_steps_per_shard + self.all_video_paths = self.get_all_video_paths() + self.all_parquet_paths = self.get_all_parquet_paths() + self.sharded_trajectories, self.shard_lengths = self.generate_shards() + self.frames_to_load = self.get_all_frames_to_load() + + # Set shard caching properties + self.shard_start_indices: dict[int, int] | None = None + self.cached_shard: dict[str, np.ndarray] | None = None + self.cached_df: pd.DataFrame | None = None + self.frame_indices_map: dict[int, dict[str, np.ndarray]] | None = None + self._executor = ThreadPoolExecutor(max_workers=1) + self._cache_job: Future | None = None + + @property + def num_shards(self) -> int: + """The number of shards.""" + return len(self.sharded_trajectories) + + def get_all_video_paths(self) -> dict[int, dict[str, Path]]: + """Get the video paths for all trajectories and all views. + + Returns: + dict[int, dict[str, Path]]: The video paths for all trajectories. + """ + video_paths = {} + for trajectory_id in self.trajectory_ids: + if isinstance(trajectory_id, np.integer): + trajectory_id = trajectory_id.item() + assert isinstance( + trajectory_id, int + ), f"trajectory_id must be an integer, got {type(trajectory_id)}" + video_paths[trajectory_id] = {} + for key in self.modality_keys["video"]: + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + video_paths[trajectory_id][key] = self.get_video_path( + trajectory_id, key.replace("video.", "") + ) + return video_paths + + def get_all_parquet_paths(self) -> dict[int, Path]: + """Get the parquet paths for all trajectories. + + Returns: + dict[int, Path]: The parquet paths for all trajectories. + """ + return { + trajectory_id: self.get_parquet_path(trajectory_id) + for trajectory_id in self.trajectory_ids + } + + def generate_shards(self) -> tuple[list[list[int]], np.ndarray]: + """Generate shards of trajectories. We recommend num_steps_per_shard >> average trajectory length. + + Args: + num_steps_per_shard (int): The number of steps per shard. + + Returns: + list[list[str]]: The shards of trajectories. + """ + sharded_trajectories = [[]] + curr_num_steps = 0 + curr_shard_index = 0 + discarded_episode_indices = [] + trajectory_ids = self.trajectory_ids + if self.discard_bad_trajectories: + discarded_episode_indices = self._lerobot_info_meta.get("discarded_episode_indices", []) + trajectory_ids = [ + trajectory_id + for trajectory_id in trajectory_ids + if trajectory_id not in discarded_episode_indices + ] + + assert ( + len(trajectory_ids) > 0 + ), f"No valid trajectories found for dataset {self.dataset_path}" + total_steps = np.sum( + [len(self.step_filter[trajectory_id]) for trajectory_id in trajectory_ids] + ).astype(int) + num_shards = np.ceil(total_steps / self.num_steps_per_shard).astype(int) + cutoffs = np.linspace(0, total_steps, num_shards + 1)[1:] # Exclude the first cutoff (0) + shard_lengths = [] + last_num_steps = 0 + for trajectory_id in trajectory_ids: + sharded_trajectories[-1].append(trajectory_id) + curr_num_steps += len(self.step_filter[trajectory_id]) + if curr_num_steps > cutoffs[curr_shard_index]: + sharded_trajectories.append([]) + curr_shard_index += 1 + shard_lengths.append(curr_num_steps - last_num_steps) + last_num_steps = curr_num_steps + shard_lengths.append(curr_num_steps - last_num_steps) + assert ( + curr_num_steps == total_steps + ), "Total steps not equal to the sum of trajectory lengths" + assert ( + len(shard_lengths) == num_shards + ), "Number of shards not equal to the number of cutoffs" + assert ( + len(sharded_trajectories) == num_shards + ), "Number of shards not equal to the number of cutoffs" + print(f"Generated {len(sharded_trajectories)} shards for dataset {self.dataset_path}") + return sharded_trajectories, np.array(shard_lengths) + + def get_all_frames_to_load(self): + """Generate a map of video frame indices to trajectory indices.""" + all_frames_to_load = {} + for trajectory_id in self.trajectory_ids: + all_frames_to_load[trajectory_id] = {} + for key in self.modality_keys["video"]: + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + filtered_indices = self.step_filter[trajectory_id] + if len(filtered_indices) > 0: + frames_to_load = np.unique( + np.concatenate( + [i + np.array(self.delta_indices[key]) for i in filtered_indices] + ) + ) + # Cap within the length of the trajectory and >= 0 + frames_to_load = frames_to_load[ + (frames_to_load < self.trajectory_lengths[trajectory_id]) + & (frames_to_load >= 0) + ] + else: + frames_to_load = np.array([]) + all_frames_to_load[trajectory_id][key] = frames_to_load + return all_frames_to_load + + @staticmethod + def get_shard( + trajectory_ids: list[int] | np.ndarray, + modality_keys: dict, + video_paths: dict[int, dict[str, Path]], + parquet_paths: dict[int, Path], + frames_to_load: dict[int, dict[str, np.ndarray]], + video_backend: str = "pyav", + video_backend_kwargs: dict | None = None, + ) -> tuple[ + dict[str, np.ndarray], dict[int, int], pd.DataFrame, dict[int, dict[str, np.ndarray]] + ]: + print("Caching shard") + start_time = time.time() + assert "video" in modality_keys, "No video modality found. No need to use caching." + cached_frames = {} + trajectory_start_indices = {} + frame_indices_map = {} + curr_step_index = 0 + cached_df = None + curr_frame_index = {key: 0 for key in modality_keys["video"]} + for trajectory_id in trajectory_ids: + trajectory_start_indices[trajectory_id] = curr_step_index + parquet_path = parquet_paths[trajectory_id] + parquet_df = pd.read_parquet(parquet_path) + # Check timestamps are in sync + parquet_timestamps = parquet_df["timestamp"].to_numpy() + trajectory_length = len(parquet_timestamps) + if isinstance(trajectory_id, np.integer): + trajectory_id = trajectory_id.item() + assert isinstance( + trajectory_id, int + ), f"trajectory_id must be an integer, got {type(trajectory_id)}" + frame_indices_map[trajectory_id] = {} + for key in modality_keys["video"]: + # Only load the frames that are needed + this_frames_to_load = frames_to_load[trajectory_id][key] + if len(this_frames_to_load) == 0: + continue + load_timestamps = parquet_timestamps[this_frames_to_load] + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + # Store a mapping that frame_indices_map[trajectory_id][key][frame_index] = index_in_concat_video_frames + frame_indices_map[trajectory_id][key] = ( + np.ones(len(parquet_timestamps), dtype=np.int32) * -1 + ) + frame_indices_map[trajectory_id][key][this_frames_to_load] = np.arange( + curr_frame_index[key], + curr_frame_index[key] + len(this_frames_to_load), + dtype=np.int32, + ) + curr_frame_index[key] += len(this_frames_to_load) + if key not in cached_frames: + cached_frames[key] = [] + frames = get_frames_by_timestamps( + video_paths[trajectory_id][key].as_posix(), + timestamps=load_timestamps, + video_backend=video_backend, + video_backend_kwargs=video_backend_kwargs or {}, + ) + cached_frames[key].append(frames) + if cached_df is None: + cached_df = parquet_df + else: + cached_df = pd.concat([cached_df, parquet_df]) + curr_step_index += trajectory_length + + # Concatenate the frames + for key in cached_frames: + cached_frames[key] = np.concatenate(cached_frames[key], axis=0) + end_time = time.time() + print(f"Cached shard in {end_time - start_time:.2f} seconds") + assert cached_df is not None, "Cached dataframe is None" + # Add global "index" column if missing (some dataset formats omit it) + if "index" not in cached_df.columns: + cached_df = cached_df.reset_index(drop=True) + cached_df["index"] = cached_df.index + return cached_frames, trajectory_start_indices, cached_df, frame_indices_map + + def start_cache_shard(self, shard_index: int) -> None: + """Start caching a shard in a background thread.""" + self._cache_job = self._executor.submit( + self.get_shard, + self.sharded_trajectories[shard_index], + self.modality_keys, + self.all_video_paths, + self.all_parquet_paths, + self.frames_to_load, + self.video_backend, + self.video_backend_kwargs, + ) + + def finish_cache_shard(self): + """Get the cached shard.""" + assert self._cache_job is not None + self.cached_shard, self.shard_start_indices, self.cached_df, self.frame_indices_map = ( + self._cache_job.result() + ) + self._cache_job = None # Clear the future to allow memory to be freed + + def delete_cached_shard(self): + """Delete the cached shard.""" + del self.cached_shard + del self.shard_start_indices + del self.cached_df + + def get_trajectories_in_shard(self) -> list[int]: + """Get the trajectories in a shard.""" + assert self.shard_start_indices is not None + return list(self.shard_start_indices.keys()) + + def get_video(self, trajectory_id: int, key: str, step_indices: np.ndarray) -> np.ndarray: + """Get the video frames from cached shards for a trajectory by a base index. + + Args: + trajectory_id (str): The ID of the trajectory. + key (str): The key of the video. + base_index (int): The base index of the trajectory. + + Returns: + np.ndarray: The video frames for the trajectory and frame indices. Shape: (T, H, W, C) + """ + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Ensure the indices are within the valid range + # This is equivalent to padding the video with extra frames at the beginning and end + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1) + # Calculate the absolute indices + assert ( + self.shard_start_indices is not None + and self.cached_shard is not None + and trajectory_id in self.shard_start_indices + and self.frame_indices_map is not None + and trajectory_id in self.frame_indices_map + and key in self.frame_indices_map[trajectory_id] + ), "Shard not cached. Please call `cache_next_shard` and `use_next_shard` first." + indices_in_shard = self.frame_indices_map[trajectory_id][key][step_indices] + assert np.all( + indices_in_shard != -1 + ), f"Indices in shard are not loaded for {trajectory_id=}, {key=}, {step_indices=}" + return self.cached_shard[key][indices_in_shard] + + def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame: + """Get the trajectory data.""" + assert self.cached_df is not None, "Cached dataframe is None" + traj_data = self.cached_df.loc[self.cached_df["episode_index"] == trajectory_id] + trajectory_index = self.get_trajectory_index(trajectory_id) + trajectory_length = self.trajectory_lengths[trajectory_index] + assert ( + len(traj_data) == trajectory_length + ), f"Trajectory length mismatch: {len(traj_data)} != {trajectory_length} {self.args} {self.kwargs}" + indices = traj_data["index"].to_numpy() + if len(indices) > 0: + start_index = indices[0] + expected_indices = np.arange(start_index, start_index + len(indices)) + assert np.array_equal( + indices, expected_indices + ), f"[{self}] Index sequence mismatch in trajectory data, {trajectory_id=}" + return traj_data + + +class ShardedLeRobotSubLangSingleActionChunkDatasetDROID(LeRobotSingleDataset): + """ + A single dataset with shards. + """ + + def __init__( + self, + *args, + num_steps_per_shard: int = int(1e7), + **kwargs, + ): + self.args = args + self.kwargs = kwargs + super().__init__(*args, **kwargs) + self.num_steps_per_shard = num_steps_per_shard + self.all_video_paths = self.get_all_video_paths() + self.all_parquet_paths = self.get_all_parquet_paths() + self.sharded_trajectories, self.shard_lengths = self.generate_shards() + + # Set shard caching properties + self.shard_start_indices: dict[int, int] | None = None + self.cached_shard: dict[str, np.ndarray] | None = None + self.cached_df: pd.DataFrame | None = None + self._executor = ThreadPoolExecutor(max_workers=1) + self._cache_job: Future | None = None + # self._traj_cache: dict[int, pd.DataFrame] = {} + # # Precompute language key once to avoid repeated scans + # self.language_key: str | None = next( + # ( + # modality_key + # for modality_name in self.modality_keys + # for modality_key in self.modality_keys[modality_name] + # if modality_key.startswith("annotation.") + # ), + # None, + # ) + # # Track current sample's chunk count for alignment across modalities + # self._current_num_chunks: dict[int, int] = {} + + @property + def num_shards(self) -> int: + """The number of shards.""" + return len(self.sharded_trajectories) + + def get_all_video_paths(self) -> dict[int, dict[str, Path]]: + """Get the video paths for all trajectories and all views. + + Returns: + dict[int, dict[str, Path]]: The video paths for all trajectories. + """ + video_paths = {} + for trajectory_id in self.trajectory_ids: + if isinstance(trajectory_id, np.integer): + trajectory_id = trajectory_id.item() + assert isinstance( + trajectory_id, int + ), f"trajectory_id must be an integer, got {type(trajectory_id)}" + video_paths[trajectory_id] = {} + for key in self.modality_keys["video"]: + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + video_paths[trajectory_id][key] = self.get_video_path( + trajectory_id, key.replace("video.", "") + ) + return video_paths + + def get_all_parquet_paths(self) -> dict[int, Path]: + """Get the parquet paths for all trajectories. + + Returns: + dict[int, Path]: The parquet paths for all trajectories. + """ + return { + trajectory_id: self.get_parquet_path(trajectory_id) + for trajectory_id in self.trajectory_ids + } + + def generate_shards(self) -> tuple[list[list[int]], np.ndarray]: + """Generate shards of trajectories. We recommend num_steps_per_shard >> average trajectory length. + + Args: + num_steps_per_shard (int): The number of steps per shard. + + Returns: + list[list[str]]: The shards of trajectories. + """ + sharded_trajectories = [[]] + curr_num_steps = 0 + curr_shard_index = 0 + discarded_episode_indices = [] + trajectory_ids = self.trajectory_ids + if self.discard_bad_trajectories: + discarded_episode_indices = self._lerobot_info_meta.get("discarded_episode_indices", []) + trajectory_ids = [ + trajectory_id + for trajectory_id in trajectory_ids + if trajectory_id not in discarded_episode_indices + ] + + assert len(trajectory_ids) > 0, "No valid trajectories found for dataset" + total_steps = np.sum( + [len(self.step_filter[trajectory_id]) for trajectory_id in trajectory_ids] + ).astype(int) + num_shards = np.ceil(total_steps / self.num_steps_per_shard).astype(int) + cutoffs = np.linspace(0, total_steps, num_shards + 1)[1:] # Exclude the first cutoff (0) + shard_lengths = [] + last_num_steps = 0 + for trajectory_id in trajectory_ids: + sharded_trajectories[-1].append(trajectory_id) + curr_num_steps += len(self.step_filter[trajectory_id]) + if curr_num_steps > cutoffs[curr_shard_index]: + sharded_trajectories.append([]) + curr_shard_index += 1 + shard_lengths.append(curr_num_steps - last_num_steps) + last_num_steps = curr_num_steps + shard_lengths.append(curr_num_steps - last_num_steps) + assert ( + curr_num_steps == total_steps + ), "Total steps not equal to the sum of trajectory lengths" + assert ( + len(shard_lengths) == num_shards + ), "Number of shards not equal to the number of cutoffs" + assert ( + len(sharded_trajectories) == num_shards + ), "Number of shards not equal to the number of cutoffs" + print(f"Generated {len(sharded_trajectories)} shards for dataset {self.dataset_path}") + return sharded_trajectories, np.array(shard_lengths) + + @staticmethod + def get_shard( + trajectory_ids: list[int] | np.ndarray, + modality_keys: dict, + video_paths: dict[int, dict[str, Path]], + parquet_paths: dict[int, Path], + video_backend: str = "pyav", + video_backend_kwargs: dict | None = None, + fps: float = None, + ) -> tuple[dict[str, np.ndarray], dict[int, int], pd.DataFrame]: + # Optional logging to avoid stdout overhead during tight loops + # (controlled by instance-level verbose flag) + # Using a staticmethod, we cannot read self.verbose; defer to caller to control prints + print("Caching shard") + start_time = time.time() + assert "video" in modality_keys, "No video modality found. No need to use caching." + cached_frames = {} + trajectory_start_indices = {} + curr_step_index = 0 + cached_df = None + for trajectory_id in trajectory_ids: + trajectory_start_indices[trajectory_id] = curr_step_index + parquet_path = parquet_paths[trajectory_id] + parquet_df = pd.read_parquet(parquet_path) + # Check timestamps are in sync + parquet_timestamps = parquet_df["timestamp"].to_numpy() + trajectory_length = len(parquet_timestamps) + if isinstance(trajectory_id, np.integer): + trajectory_id = trajectory_id.item() + assert isinstance( + trajectory_id, int + ), f"trajectory_id must be an integer, got {type(trajectory_id)}" + for key in modality_keys["video"]: + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + if key not in cached_frames: + cached_frames[key] = [] + frames = get_frames_by_timestamps( + video_paths[trajectory_id][key].as_posix(), + timestamps=parquet_timestamps, + video_backend=video_backend, + video_backend_kwargs=video_backend_kwargs, + fps=fps, + ) + cached_frames[key].append(frames) + if cached_df is None: + cached_df = parquet_df + else: + cached_df = pd.concat([cached_df, parquet_df]) + curr_step_index += trajectory_length + + # Concatenate the frames + for key in cached_frames: + cached_frames[key] = np.concatenate(cached_frames[key], axis=0) + end_time = time.time() + print(f"Cached shard in {end_time - start_time:.2f} seconds") + assert cached_df is not None, "Cached dataframe is None" + # Add global "index" column if missing (some dataset formats omit it) + if "index" not in cached_df.columns: + cached_df = cached_df.reset_index(drop=True) + cached_df["index"] = cached_df.index + return cached_frames, trajectory_start_indices, cached_df + + def start_cache_shard(self, shard_index: int) -> None: + """Start caching a shard in a background thread.""" + self._cache_job = self._executor.submit( + self.get_shard, + self.sharded_trajectories[shard_index], + self.modality_keys, + self.all_video_paths, + self.all_parquet_paths, + self.video_backend, + self.video_backend_kwargs, + self.fps, + ) + + def finish_cache_shard(self): + """Get the cached shard.""" + assert self._cache_job is not None + self.cached_shard, self.shard_start_indices, self.cached_df = self._cache_job.result() + self._cache_job = None # Clear the future to allow memory to be freed + + def delete_cached_shard(self): + """Delete the cached shard.""" + del self.cached_shard + del self.shard_start_indices + del self.cached_df + # self._traj_cache.clear() + + def get_trajectories_in_shard(self) -> list[int]: + """Get the trajectories in a shard.""" + assert self.shard_start_indices is not None + return list(self.shard_start_indices.keys()) + + def get_step_data(self, trajectory_id: int, indices: dict[str, np.ndarray]) -> dict: + """Get the RAW data for a single step in a trajectory. No transforms are applied. + + Args: + trajectory_id (int): The name of the trajectory. + indices (dict[str, np.ndarray]): The indices for each modality. + + Returns: + dict: The RAW data for the step. + + Example return: + { + "video": { + "video.image_side_0": [B, T, H, W, C], + "video.image_side_1": [B, T, H, W, C], + }, + "state": { + "state.eef_position": [B, T, state_dim], + "state.eef_rotation": [B, T, state_dim], + }, + "action": { + "action.eef_position": [B, T, action_dim], + "action.eef_rotation": [B, T, action_dim], + }, + } + """ + data = {} + # Get the data for all modalities + self.curr_traj_data = self.get_trajectory_data(trajectory_id) + for modality in self.modality_keys: + # Get the data corresponding to each key in the modality + for key in self.modality_keys[modality]: + # Only load the data if the key is in the indices + if key in indices: + data[key] = self.get_data_by_modality( + trajectory_id, modality, key, indices[key] + ) + # Skip this sample if state or action data is empty + if data[key] is not None and hasattr(data[key], '__len__') and len(data[key]) == 0: + return None + return data + + def get_video(self, trajectory_id: int, key: str, step_indices: np.ndarray) -> np.ndarray: + """Get the video frames from cached shards for a trajectory by uniformly sampling from language-consistent ranges. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The key of the video. + step_indices (np.ndarray): The step indices to retrieve data for. + + Returns: + np.ndarray: The video frames for the trajectory and frame indices. Shape: (T, H, W, C) + """ + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + trajectory_length = self.trajectory_lengths[trajectory_index] + + # Get trajectory data to access language annotations (reuse if already loaded) + # traj_data = ( + # self.curr_traj_data + # if getattr(self, "curr_traj_data", None) is not None + # else self.get_trajectory_data(trajectory_id) + # ) + traj_data = self.get_trajectory_data(trajectory_id) + # print("trajectory id", trajectory_id, step_indices, trajectory_index) + + # Get language annotations for all steps in the trajectory + # language_key = self.language_key + for modality in self.modality_keys: + for modality_key in self.modality_keys[modality]: + if modality_key.startswith("annotation."): + subkey = modality_key.replace("annotation.", "") + annotation_meta = self.lerobot_modality_meta.annotation + subkey_meta = annotation_meta[subkey] + language_key = subkey_meta.original_key + break + assert language_key is not None, "Language key not found" + if language_key in traj_data.columns: + language_annotations = traj_data[language_key].values + else: + # Fallback to original behavior if language annotations are not available + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, trajectory_length - 1) + assert ( + self.shard_start_indices is not None + and self.cached_shard is not None + and trajectory_id in self.shard_start_indices + ), "Shard not cached. Please call `cache_next_shard` and `use_next_shard` first." + indices_in_shard = self.shard_start_indices[trajectory_id] + step_indices + return self.cached_shard[key][indices_in_shard] + + # Find language-consistent ranges and uniformly sample from them + sampled_indices = self._uniform_sample_from_language_ranges( + step_indices, language_annotations, trajectory_length + ) + + # Ensure the sampled indices are within the valid range + sampled_indices = np.maximum(sampled_indices, 0) + sampled_indices = np.minimum(sampled_indices, trajectory_length - 1) + # print("sampled indices", sampled_indices) + + # Calculate the absolute indices + assert ( + self.shard_start_indices is not None + and self.cached_shard is not None + and trajectory_id in self.shard_start_indices + ), "Shard not cached. Please call `cache_next_shard` and `use_next_shard` first." + indices_in_shard = self.shard_start_indices[trajectory_id] + sampled_indices + return self.cached_shard[key][indices_in_shard] + + def get_data_by_modality( + self, + trajectory_id: int, + modality: str, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray | list[str] | None: + """Get the data corresponding to the modality for a trajectory by step indices. + + This method dispatches to the appropriate specialized method based on the modality. + For the language modality, empty strings are returned if no matching data is found. + + Args: + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data (video, state, action, language, etc.). + key (str): The key of the data. + step_indices (np.ndarray): The step indices of the trajectory. + + Returns: + np.ndarray | list[str] | None: The data for the specified modality. + """ + if modality == "video": + return self.get_video(trajectory_id, key, step_indices) + elif modality == "state": + return self.get_state(trajectory_id, modality, key, step_indices) + elif modality == "action": + return self.get_action(trajectory_id, modality, key, step_indices) + elif modality == "language": + return self.get_language(trajectory_id, key, step_indices) + elif modality == "lapa_action": + return self.get_lapa_action(trajectory_id, key, step_indices) + elif modality == "dream_actions": + return self.get_dream_actions(trajectory_id, key, step_indices) + elif modality == "rl_info": + return self.get_rl_info(trajectory_id, key, step_indices) + else: + raise ValueError(f"Invalid modality: {modality}") + + + def get_state( + self, + trajectory_id: int, + modality: str, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray: + """Get the state data for a trajectory by a base index. + If the step indices are out of range, pad with the data: + if the data is stored in absolute format, pad with the first or last step data; + otherwise, pad with zero. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data. + key (str): The key of the data. + base_index (int): The base index of the trajectory. + + Returns: + np.ndarray: The data for the trajectory and step indices. + """ + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + + # Note [YL]: this handles action.task_progress if specified + if key == "action.task_progress": + # Get frame_index array and apply proper bounds checking and padding + frame_index_array = self.curr_traj_data["frame_index"].to_numpy() + # Use retrieve_data_and_pad to handle out-of-bounds indices + frame_index = self.retrieve_data_and_pad( + array=frame_index_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last", # Use first/last for task progress + ) + # get the task progress by using "frame index / trajectory length" + progress = frame_index / max_length + progress = progress.reshape(-1, 1) + return progress + + assert key.startswith(modality + "."), f"{key} must start with {modality + '.'}, got {key}" + # Get the sub-key, e.g. state.joint_angles -> joint_angles + subkey = key.replace(modality + ".", "") + # Get the lerobot key + le_state_or_action_cfg = getattr(self.lerobot_modality_meta, modality) + le_key = le_state_or_action_cfg[subkey].original_key + if le_key is None: + le_key = subkey + # Get the data array, shape: (T, D) + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + assert le_key in self.curr_traj_data.columns, f"No {le_key} found in {trajectory_id=}" + data_array: np.ndarray = np.stack(self.curr_traj_data[le_key]) # type: ignore + if data_array.ndim == 1: + assert ( + data_array.shape[0] == max_length + ), f"Expected 1D array with length {max_length}, got {data_array.shape} array" + data_array = data_array.reshape(-1, 1) + assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" + le_indices = np.arange( + le_state_or_action_cfg[subkey].start, + le_state_or_action_cfg[subkey].end, + ) + data_array = data_array[:, le_indices] + # Get the state or action configuration + state_or_action_cfg = getattr(self.metadata.modalities, modality)[subkey] + + # Build sampled indices for state aligned with language and video sampling + # For state, select only the anchor index per 30-frame chunk (stride 30): + # [..., first_idx-30, first_idx, first_idx+30, ...] + # Stop on language change at the step anchor, bounds, or when reaching 16 anchors (to match 16 chunks). + trajectory_index = self.get_trajectory_index(trajectory_id) + trajectory_length = self.trajectory_lengths[trajectory_index] + # traj_data = ( + # self.curr_traj_data + # if getattr(self, "curr_traj_data", None) is not None + # else self.get_trajectory_data(trajectory_id) + # ) + # language_key = self.language_key + traj_data = self.get_trajectory_data(trajectory_id) + language_key = None + for modality_name in self.modality_keys: + for modality_key in self.modality_keys[modality_name]: + if modality_key.startswith("annotation."): + subkey = modality_key.replace("annotation.", "") + annotation_meta = self.lerobot_modality_meta.annotation + subkey_meta = annotation_meta[subkey] + language_key = subkey_meta.original_key + break + if language_key is not None and language_key in traj_data.columns and len(step_indices) > 0: + language_annotations = traj_data[language_key].values + first_idx = max(0, min(int(step_indices[0]), trajectory_length - 1)) + target_language = language_annotations[first_idx] + + # Get the number of chunks from video sampling to ensure alignment + target_num_chunks = None + # if first_idx in self._current_num_chunks: + if hasattr(self, '_current_num_chunks') and first_idx in self._current_num_chunks: + target_num_chunks = self._current_num_chunks[first_idx] + # print(f"State: Using target_num_chunks from video: {target_num_chunks}") + + max_frames = self.max_chunk_size # 16 anchors to align with 16 chunks as video/action + sampled_list: list[int] = [] + + def add_anchor(anchor_index: int) -> None: + nonlocal sampled_list + if len(sampled_list) >= max_frames: + return + # If we have a target number of chunks, stop when we reach it + if target_num_chunks is not None and len(sampled_list) >= target_num_chunks: + return + # Require full 32-length window to exist for alignment with action/video + if 0 <= anchor_index and anchor_index + 24 < trajectory_length: + sampled_list.append(int(anchor_index)) + + # Always include first_idx anchor + add_anchor(first_idx) + + # Expand outward in 32-frame steps + step = 1 + back_done = False + fwd_done = False + while len(sampled_list) < max_frames and (not back_done or not fwd_done): + # Stop if we've reached the target number of chunks + if target_num_chunks is not None and len(sampled_list) >= target_num_chunks: + break + + if not back_done: + back_anchor = first_idx - 24 * step + if back_anchor < 0: + back_done = True + elif language_annotations[back_anchor] != target_language: + back_done = True + else: + add_anchor(back_anchor) + if len(sampled_list) >= max_frames: + break + if not fwd_done: + fwd_anchor = first_idx + 24 * step + if fwd_anchor >= trajectory_length: + fwd_done = True + elif language_annotations[fwd_anchor] != target_language: + fwd_done = True + else: + add_anchor(fwd_anchor) + step += 1 + + if len(sampled_list) > 0: + sampled_indices = np.array(sorted(set(sampled_list)), dtype=int) + else: + sampled_indices = np.array([], dtype=int) + else: + # Fallback: use provided indices with bounds + sampled_indices = np.maximum(step_indices, 0) + sampled_indices = np.minimum(sampled_indices, trajectory_length - 1) + + # print("sampled indices for state", sampled_indices) + + # Pad the data using the computed sampled indices + return self.retrieve_data_and_pad( + array=data_array, + step_indices=sampled_indices, + max_length=max_length, + padding_strategy="first_last" if state_or_action_cfg.absolute else "zero", + ) + + def get_action( + self, + trajectory_id: int, + modality: str, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray: + """Get the action data for a trajectory by a base index. + If the step indices are out of range, pad with the data: + if the data is stored in absolute format, pad with the first or last step data; + otherwise, pad with zero. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data. + key (str): The key of the data. + base_index (int): The base index of the trajectory. + + Returns: + np.ndarray: The data for the trajectory and step indices. + """ + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + + # Note [YL]: this handles action.task_progress if specified + if key == "action.task_progress": + # Get frame_index array and apply proper bounds checking and padding + frame_index_array = self.curr_traj_data["frame_index"].to_numpy() + # Use retrieve_data_and_pad to handle out-of-bounds indices + frame_index = self.retrieve_data_and_pad( + array=frame_index_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last", # Use first/last for task progress + ) + # get the task progress by using "frame index / trajectory length" + progress = frame_index / max_length + progress = progress.reshape(-1, 1) + return progress + + assert key.startswith(modality + "."), f"{key} must start with {modality + '.'}, got {key}" + # Get the sub-key, e.g. state.joint_angles -> joint_angles + subkey = key.replace(modality + ".", "") + # Get the lerobot key + le_state_or_action_cfg = getattr(self.lerobot_modality_meta, modality) + le_key = le_state_or_action_cfg[subkey].original_key + if le_key is None: + le_key = subkey + # Get the data array, shape: (T, D) + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + assert le_key in self.curr_traj_data.columns, f"No {le_key} found in {trajectory_id=}" + data_array: np.ndarray = np.stack(self.curr_traj_data[le_key]) # type: ignore + if data_array.ndim == 1: + assert ( + data_array.shape[0] == max_length + ), f"Expected 1D array with length {max_length}, got {data_array.shape} array" + data_array = data_array.reshape(-1, 1) + assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" + le_indices = np.arange( + le_state_or_action_cfg[subkey].start, + le_state_or_action_cfg[subkey].end, + ) + data_array = data_array[:, le_indices] + # Get the state or action configuration + state_or_action_cfg = getattr(self.metadata.modalities, modality)[subkey] + + + # Build sampled indices for action aligned with language and video sampling + # Action runs at 30fps, so for each ±30-frame step around first_idx, + # collect a 30-length chunk with stride 1: [anchor ... anchor+29]. + # Stop on language change at the step anchor, bounds, or when reaching 480 frames (16 chunks * 30). + trajectory_index = self.get_trajectory_index(trajectory_id) + trajectory_length = self.trajectory_lengths[trajectory_index] + # traj_data = ( + # self.curr_traj_data + # if getattr(self, "curr_traj_data", None) is not None + # else self.get_trajectory_data(trajectory_id) + # ) + # language_key = self.language_key + traj_data = self.get_trajectory_data(trajectory_id) + language_key = None + for modality_name in self.modality_keys: + for modality_key in self.modality_keys[modality_name]: + if modality_key.startswith("annotation."): + subkey = modality_key.replace("annotation.", "") + annotation_meta = self.lerobot_modality_meta.annotation + subkey_meta = annotation_meta[subkey] + language_key = subkey_meta.original_key + break + if language_key is not None and language_key in traj_data.columns and len(step_indices) > 0: + language_annotations = traj_data[language_key].values + first_idx = max(0, min(int(step_indices[0]), trajectory_length - 1)) + target_language = language_annotations[first_idx] + + # Get the number of chunks from video sampling to ensure alignment + target_num_chunks = None + # if first_idx in self._current_num_chunks: + if hasattr(self, '_current_num_chunks') and first_idx in self._current_num_chunks: + target_num_chunks = self._current_num_chunks[first_idx] + # print(f"Using target_num_chunks from video: {target_num_chunks}") + + max_frames = 24 * self.max_chunk_size + per_step_offsets = list(range(24)) # 0..23 + sampled_list: list[int] = [] + + def add_step_set(anchor_index: int) -> None: + nonlocal sampled_list + # Ensure the whole 32-length chunk fits within bounds + if anchor_index < 0 or anchor_index + 24 >= trajectory_length: + return + # Ensure we don't overrun the max_frames cap with a partial chunk + if len(sampled_list) + 24 > max_frames: + return + # If we have a target number of chunks, stop when we reach it + if target_num_chunks is not None and len(sampled_list) // 24 >= target_num_chunks: + return + for offset in per_step_offsets: + idx = anchor_index + offset + sampled_list.append(int(idx)) + + # Always include first_idx chunk + add_step_set(first_idx) + + # Expand outward in 32-frame steps + step = 1 + back_done = False + fwd_done = False + while len(sampled_list) < max_frames and (not back_done or not fwd_done): + # Stop if we've reached the target number of chunks + if target_num_chunks is not None and len(sampled_list) // 24 >= target_num_chunks: + break + + if not back_done: + back_anchor = first_idx - 24 * step + if back_anchor < 0: + back_done = True + elif language_annotations[back_anchor] != target_language: + back_done = True + else: + add_step_set(back_anchor) + if len(sampled_list) >= max_frames: + break + if not fwd_done: + fwd_anchor = first_idx + 24 * step + if fwd_anchor >= trajectory_length: + fwd_done = True + elif language_annotations[fwd_anchor] != target_language: + fwd_done = True + else: + add_step_set(fwd_anchor) + step += 1 + + if len(sampled_list) > 0: + unique_sorted = np.array(sorted(set(sampled_list)), dtype=int) + # Enforce divisibility by 30 and the 480 cap + capped_size = min(unique_sorted.size, max_frames) + divisible_size = (capped_size // 24) * 24 + sampled_indices = unique_sorted[:divisible_size] + else: + sampled_indices = np.array([], dtype=int) + else: + # Fallback: use provided indices with bounds + sampled_indices = np.maximum(step_indices, 0) + sampled_indices = np.minimum(sampled_indices, trajectory_length - 1) + + # print("sampled indices for action", first_idx, sampled_indices, trajectory_length) + + # Pad the data using the computed sampled indices + action_data = self.retrieve_data_and_pad( + array=data_array, + step_indices=sampled_indices, + max_length=max_length, + padding_strategy="first_last" if state_or_action_cfg.absolute else "zero", + ) + # print("action data before convert", key) + # Calculate relative action on the fly if relative_action is enabled + # Only apply to keys that are in relative_action_keys + subkey = key.replace("action.", "") + should_convert_to_relative = ( + (self.relative_action or self.relative_action_per_horizon) + and len(sampled_indices) > 0 + and (self.relative_action_keys is None or subkey in self.relative_action_keys) + ) + if should_convert_to_relative: + # print("action data before convert", action_data[0], action_data[-1], key) + action_data = self._convert_to_relative_action( + action_data=action_data, + action_key=key, + sampled_indices=sampled_indices, + trajectory_id=trajectory_id, + chunk_size=24, + ) + # print("action data after convert", action_data[0], action_data[-1], key) + + return action_data + + def _convert_to_relative_action( + self, + action_data: np.ndarray, + action_key: str, + sampled_indices: np.ndarray, + trajectory_id: int, + chunk_size: int = 24, + ) -> np.ndarray: + """Convert absolute action to relative action by subtracting reference state. + + Args: + action_data: Absolute action data, shape (T, D) + action_key: The action key (e.g., 'action.left_arm_joints') + sampled_indices: The sampled indices for the action + trajectory_id: The trajectory ID + chunk_size: Size of each action chunk (default 24) + + Returns: + np.ndarray: Relative action data, shape (T, D) + """ + # Get corresponding state key (assume state key matches action key) + state_key = action_key.replace("action.", "state.") + subkey = action_key.replace("action.", "") + + # Get state data from trajectory + traj_data = self.get_trajectory_data(trajectory_id) + le_state_cfg = getattr(self.lerobot_modality_meta, "state", None) + + if le_state_cfg is None or subkey not in le_state_cfg: + # If no corresponding state key, return original action data + return action_data + + le_state_key = le_state_cfg[subkey].original_key + if le_state_key is None: + le_state_key = subkey + + if le_state_key not in traj_data.columns: + # If state column doesn't exist, return original action data + return action_data + + # Get state data array + state_array: np.ndarray = np.stack(traj_data[le_state_key]) + if state_array.ndim == 1: + state_array = state_array.reshape(-1, 1) + + # Apply same indices as action + le_indices = np.arange( + le_state_cfg[subkey].start, + le_state_cfg[subkey].end, + ) + state_array = state_array[:, le_indices] + + # Calculate relative action for each chunk + relative_action_data = action_data.copy() + num_chunks = len(sampled_indices) // chunk_size + + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * chunk_size + chunk_end = chunk_start + chunk_size + + # Get anchor index (first index of the chunk) + anchor_idx = sampled_indices[chunk_start] + + # Get reference state at anchor index + if anchor_idx < len(state_array): + reference_state = state_array[anchor_idx] + + # Subtract reference state from all actions in this chunk + relative_action_data[chunk_start:chunk_end] = ( + action_data[chunk_start:chunk_end] - reference_state + ) + + return relative_action_data + + def _uniform_sample_from_language_ranges( + self, + step_indices: np.ndarray, + language_annotations: np.ndarray, + trajectory_length: int + ) -> np.ndarray: + """Uniformly sample from language-consistent ranges based on the first index's language. + + Args: + step_indices (np.ndarray): Original step indices to sample. + language_annotations (np.ndarray): Language annotations for each step in the trajectory. + trajectory_length (int): Total length of the trajectory. + + Returns: + np.ndarray: New indices sampled uniformly from the language-consistent range of the first index. + """ + if len(step_indices) == 0: + return np.array([]) + + # Use only the first index to determine the target language + first_idx = max(0, min(step_indices[0], trajectory_length - 1)) + target_language = language_annotations[first_idx] + + # Build sampled indices by moving in ±32-frame steps from first_idx + # and adding 4 frames at 8-frame strides for each step, while: + # - staying within trajectory bounds, + # - keeping language consistent with target_language at the anchor step, + # - and limiting the total collected frames to 81. + max_frames = 8 * self.max_chunk_size + 1 + per_step_offsets = [0, 3, 6, 9, 12, 15, 18, 21] + sampled_list: list[int] = [] + + def add_step_set(anchor_index: int) -> None: + # Only add a complete 4-frame set if it fully fits and capacity allows + # Require full 32-frame window to exist for alignment with action/state + nonlocal sampled_list + if anchor_index < 0 or anchor_index + 23 >= trajectory_length: + return + if len(sampled_list) + len(per_step_offsets) > max_frames: + return + for offset in per_step_offsets: + idx = anchor_index + offset + sampled_list.append(int(idx)) + + # Always include the set at the first_idx + add_step_set(first_idx) + + # Expand outward in both directions in 32-frame steps + step = 1 + back_done = False + fwd_done = False + while len(sampled_list) < max_frames and (not back_done or not fwd_done): + # Backward step + if not back_done: + back_anchor = first_idx - 24 * step + if back_anchor < 0: + back_done = True + elif language_annotations[back_anchor] != target_language: + back_done = True + else: + add_step_set(back_anchor) + # Forward step + if len(sampled_list) >= max_frames: + break + if not fwd_done: + fwd_anchor = first_idx + 24 * step + if fwd_anchor >= trajectory_length: + fwd_done = True + elif language_annotations[fwd_anchor] != target_language: + fwd_done = True + else: + add_step_set(fwd_anchor) + step += 1 + + # De-duplicate and sort ascending for stable ordering + if len(sampled_list) == 0: + return np.array([]) + unique_sorted = np.array(sorted(set(sampled_list)), dtype=int) + # Ensure we return at most 81 frames + if unique_sorted.size > max_frames: + unique_sorted = unique_sorted[:max_frames] + + # Convert to 4n+1 format by adding one more frame at the end with 8-frame stride + if unique_sorted.size > 0: + # Get the last index and add one more frame with 8-frame stride + last_idx = unique_sorted[-1] + additional_idx = last_idx + 3 + + # Only add if it doesn't exceed trajectory bounds and max_frames + if additional_idx < trajectory_length and unique_sorted.size < max_frames: + unique_sorted = np.append(unique_sorted, additional_idx) + else: + # Trim to 8n+1 format. Require at least 9 frames so (noisy_frames-1)//num_frame_per_block >= 1 + # for action/state model invariant (CausalWanModel); otherwise return empty so sample is skipped. + if unique_sorted.size <= 8: + return np.array([]) + unique_sorted = unique_sorted[:-7] + + # ensure that unique_sorted has 4n+1 frames + assert unique_sorted.size % 8 == 1, f"unique_sorted size {unique_sorted.size} is not 4n+1" + + # Store the number of chunks for alignment with action/state + num_video_chunks = (unique_sorted.size - 1) // 8 + if not hasattr(self, '_current_num_chunks'): + self._current_num_chunks = {} + # Use first_idx as a key to track the current sample's chunk count + self._current_num_chunks[first_idx] = num_video_chunks + + # print("unique_sorted size", unique_sorted.size, "num_video_chunks", num_video_chunks) + return unique_sorted + + + def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame: + """Get the trajectory data.""" + assert self.cached_df is not None, "Cached dataframe is None" + + # Quick verification + if self.cached_df.empty: + raise ValueError("cached_df is completely empty!") + + # # Fast path: return cached slice if available + # if trajectory_id in self._traj_cache: + # return self._traj_cache[trajectory_id] + + available_episodes = self.cached_df["episode_index"].unique() + if trajectory_id not in available_episodes: + raise ValueError( + f"trajectory_id {trajectory_id} not found in cached_df. " + f"Available episodes: {sorted(available_episodes)}" + ) + + traj_data = self.cached_df.loc[self.cached_df["episode_index"] == trajectory_id] + trajectory_index = self.get_trajectory_index(trajectory_id) + trajectory_length = self.trajectory_lengths[trajectory_index] + assert ( + len(traj_data) == trajectory_length + ), f"Trajectory length mismatch: {len(traj_data)} != {trajectory_length} {self.args} {self.kwargs}" + indices = traj_data["index"].to_numpy() + if len(indices) > 0: + start_index = indices[0] + expected_indices = np.arange(start_index, start_index + len(indices)) + assert np.array_equal( + indices, expected_indices + ), f"[{self}] Index sequence mismatch in trajectory data, {trajectory_id=}" + # Store in cache to avoid repeated filtering on subsequent calls within a batch + # self._traj_cache[trajectory_id] = traj_data + return traj_data + + + +class ShardedLeRobotMixtureDataset(LeRobotMixtureDataset, IterableDataset): + """ + A mixture of multiple datasets. This class samples a single dataset based on the dataset weights and then calls the `__getitem__` method of the sampled dataset. + It is recommended to modify the single dataset class instead of this class. + """ + + def __init__( + self, + data_mixture: list[tuple[LeRobotSingleDataset, float]], + training: bool, + balance_dataset_weights: bool = True, + balance_trajectory_weights: bool = True, + seed: int = 42, + shard_sampling_rate: float = 0.5, + num_shards_to_sample: int = 2**20, + allow_padding_at_end: bool = False, + ): + """ + Initialize the mixture dataset. + + Args: + data_mixture (list[tuple[ShardedLeRobotSingleDataset, float]]): Datasets and their corresponding weights. + mode (str): If "train", __iter__ will yield different samples every epoch; if "val" or "test", __iter__ will yield the same sample every epoch. + balance_dataset_weights (bool): If True, the weight of dataset will be multiplied by the total trajectory length of each dataset. + balance_trajectory_weights (bool): If True, sample trajectories within a dataset weighted by their length; otherwise, use equal weighting. + seed (int): Random seed for sampling. + shard_sampling_rate (float): How much data per shard to sample, in a 0-1 scale. + num_shards_to_sample (int): The number of shards to sample. + """ + super().__init__( + data_mixture=data_mixture, + training=training, + balance_dataset_weights=balance_dataset_weights, + balance_trajectory_weights=balance_trajectory_weights, + seed=seed, + allow_padding_at_end=allow_padding_at_end, + ) + # Add type hint + self.datasets: list[ShardedLeRobotSingleDataset] = self.datasets + # Set properties + self.shard_sampling_rate = shard_sampling_rate + self.num_shards_to_sample = num_shards_to_sample + + # Calculate shard sampling weights + all_shard_sampling_weights = [] + all_shards = [] + for dataset_id, (dataset, weight) in enumerate( + zip(self.datasets, self._dataset_sampling_weights) + ): + shard_sampling_weights = dataset.shard_lengths / dataset.shard_lengths.sum() + all_shard_sampling_weights.append(shard_sampling_weights * weight) + all_shards.extend( + [(dataset_id, shard_idx) for shard_idx in range(shard_sampling_weights.shape[0])] + ) + all_shard_sampling_weights = np.concatenate(all_shard_sampling_weights) + all_shard_sampling_weights /= all_shard_sampling_weights.sum() + self._shard_sampling_weights = all_shard_sampling_weights + self._all_shards = all_shards + + # Generate shards sample schedule for all ranks and workers + self._shards_sample_schedule = self.generate_shards_sample_schedule() + + # Check shard sampling rate + assert 0 <= shard_sampling_rate <= 1, "Shard sampling rate must be between 0 and 1" + + # Set properties for distributed training + if dist.is_initialized(): + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + else: + self.rank = 0 + self.world_size = 1 + self.worker_id = None + self.num_workers = None + + @property + def dataset_sampling_weights(self) -> np.ndarray: + """The dataset sampling weights.""" + return self._dataset_sampling_weights + + @property + def shard_sampling_weights(self) -> list[np.ndarray]: + """The weights of each shard.""" + return self._shard_sampling_weights + + @property + def all_shards(self) -> list[tuple[int, int]]: + """The shards to sample.""" + return self._all_shards + + @property + def shards_sample_schedule(self) -> list[tuple[int, int]]: + """The shards sample schedule. + + Returns: + list[tuple[int, int]]: The shards to sample, in (dataset_index, shard_index). + """ + assert self._shards_sample_schedule is not None, "Shards sample schedule not set." + return self._shards_sample_schedule + + @property + def trajectory_sampling_weights(self): + """The trajectory sampling weights.""" + raise ValueError("ShardedRobotMixtureDataset does not support trajectory sampling weights.") + + @property + def primary_dataset_indices(self): + """The primary dataset indices.""" + raise ValueError("ShardedRobotMixtureDataset does not support primary dataset indices.") + + def reset_seed(self, seed: int): + self.seed = seed + self._shards_sample_schedule = self.generate_shards_sample_schedule() + + def generate_shards_sample_schedule(self): + if self.training: + rng = np.random.default_rng(self.seed) + sampled_shard_ids = rng.choice( + len(self.all_shards), size=self.num_shards_to_sample, p=self.shard_sampling_weights + ) + shards_sample_schedule = [self.all_shards[i] for i in sampled_shard_ids] + rng.shuffle(shards_sample_schedule) + else: + shards_sample_schedule = [ + self.all_shards[i % len(self.all_shards)] for i in range(self.num_shards_to_sample) + ] + return shards_sample_schedule + + def filter_shards_sample_schedule(self): + """Filter the shards sample schedule for each worker. + + Returns: + list[tuple[int, int]]: The shards to sample, in (dataset_index, shard_index). + """ + # Filter shards for each worker + filtered_schedule = [] + worker_info = get_worker_info() + # If we have multiple workers, further split shards among them + if worker_info is not None: + worker_id = worker_info.id + num_workers = worker_info.num_workers + else: + worker_id = 0 + num_workers = 1 + + if self.worker_id is None: + assert self.num_workers is None + self.worker_id = worker_id + self.num_workers = num_workers + else: + assert ( + self.worker_id == worker_id and self.num_workers == num_workers + ), "Worker ID or number of workers has been changed since it was set. This is not allowed." + + for i, shard in enumerate(self.shards_sample_schedule): + if i % (self.world_size * num_workers) == self.rank * num_workers + worker_id: + filtered_schedule.append(shard) + # print(f"Filtered shards for rank {self.rank}, worker {worker_id}: {filtered_schedule}") + return filtered_schedule + + def __str__(self) -> str: + dataset_descriptions = [] + for dataset, weight in zip(self.datasets, self.dataset_sampling_weights): + shard_lengths = dataset.shard_lengths + assert len(shard_lengths.shape) == 1, "Shard lengths must be a 1D array" + num_shards = shard_lengths.shape[0] + max_shard_length = int(shard_lengths.max()) + min_shard_length = int(shard_lengths.min()) + dataset_description = { + "Dataset": str(dataset), + "Sampling weight": float(weight), + "Num shards": num_shards, + "Max shard length": max_shard_length, + "Min shard length": min_shard_length, + } + dataset_descriptions.append(dataset_description) + return yaml.dump( + { + "Mixture dataset": dataset_descriptions, + "Rank": self.rank, + "World size": self.world_size, + } + ) + + def __iter__(self): + """Iterate over the dataset.""" + + # Not supported: balance_trajectory_weights=False + if not self.balance_trajectory_weights: + raise NotImplementedError( + "balance_trajectory_weights=False is not supported. Please use balance_dataset_weights=True instead." + ) + + self._shards_sample_schedule = self.filter_shards_sample_schedule() + self.curr_shard_index = -1 + self.cache_next_shard() + rng = np.random.default_rng(self.seed) + for i, (dataset_index, shard_index) in enumerate(self.shards_sample_schedule): + self.curr_shard_index += 1 + assert ( + i == self.curr_shard_index + ), f"Shard index mismatch: {i} != {self.curr_shard_index}" + dataset = self.datasets[dataset_index] + wait_start = time.time() + dataset.finish_cache_shard() + wait_end = time.time() + print( + f"Rank {self.rank}, Worker {self.worker_id}: Wait for shard {shard_index} in dataset {dataset_index} in {wait_end - wait_start:.2f} seconds" + ) + # Start caching the next shard immediately + self.cache_next_shard() + all_steps: list[tuple[int, int]] = [] + for trajectory_id in dataset.get_trajectories_in_shard(): + trajectory_index = dataset.get_trajectory_index(trajectory_id) + if self.allow_padding_at_end: + allowed_length = dataset.trajectory_lengths[trajectory_index] + else: + max_delta_index = dataset.max_delta_index + trajectory_length = dataset.trajectory_lengths[trajectory_index] + allowed_length = trajectory_length - max_delta_index + # Get the allowed indices from the step filter + allowed_indices = dataset.step_filter[trajectory_id] + # Remove indices that are too large + allowed_indices = allowed_indices[allowed_indices <= allowed_length] + for i in allowed_indices: + all_steps.append((trajectory_id, i)) + if self.training: + rng.shuffle(all_steps) + sampled_steps = all_steps[: int(dataset.num_steps_per_shard * self.shard_sampling_rate)] + for trajectory_id, step_index in sampled_steps: + # print( + # f"Loading step data from rank {self.rank}, worker {self.worker_id}: {dataset_index} {trajectory_id}, {step_index}" + # ) + indices = { + key: delta_indices + step_index + for key, delta_indices in dataset.delta_indices.items() + } + step_data = dataset.get_step_data(trajectory_id, indices) + # Skip samples where state or action would be empty + if step_data is not None: + yield dataset.transforms(step_data) + + # Delete the cached shard and shard start indices to free up memory + dataset.delete_cached_shard() + + def cache_next_shard(self): + """Cache the next shard in a background thread.""" + next_dataset_idx, next_shard_idx = self.shards_sample_schedule[self.curr_shard_index + 1] + self.datasets[next_dataset_idx].start_cache_shard(next_shard_idx) + + def __getitem__(self, index: int) -> dict: + raise NotImplementedError( + "__getitem__ is not supported for CachedRobotMixtureDataset. Please use __iter__ instead." + ) + + def __len__(self) -> int: + """The length of the dataset.""" + total_length = 0 + for dataset_idx, _ in self.shards_sample_schedule: + dataset = self.datasets[dataset_idx] + total_length += int(dataset.num_steps_per_shard * self.shard_sampling_rate) + return total_length diff --git a/groot/vla/data/dataset/libero.py b/groot/vla/data/dataset/libero.py new file mode 100644 index 0000000000000000000000000000000000000000..cd1ec782d74fb8308dc2d45c98d911cf8d12f413 --- /dev/null +++ b/groot/vla/data/dataset/libero.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""DreamZero LIBERO Dataset - reads LeRobot parquet, outputs DreamTransform format.""" +import os, random, json, logging +from pathlib import Path +import numpy as np +import torch +import torch.utils.data as data +import pyarrow.parquet as pq +import cv2 + +from groot.vla.model.dreamzero.transform.dreamzero_cotrain import DreamTransform +from groot.vla.data.schema.lerobot import ( + DatasetMetadata, DatasetStatistics, DatasetStatisticalValues, + DatasetModalities, VideoMetadata, StateActionMetadata, +) +from groot.vla.data.schema.embodiment_tags import EmbodimentTag + +logger = logging.getLogger(__name__) + +class LiberoDataset(data.Dataset): + """LIBERO dataset for DreamZero. Reads LeRobot parquet, applies DreamTransform.""" + + def __init__(self, dataset_dir, num_frames=8, action_horizon=12, + state_horizon=1, num_views=1, video_height=160, video_width=320, + max_episodes=None, seed=42, **kwargs): + self.dataset_dir = Path(dataset_dir) + self.num_frames = num_frames + self.action_horizon = action_horizon + self.state_horizon = state_horizon + self.num_views = num_views + self.video_height = video_height + self.video_width = video_width + self.seed = seed + self.max_episodes = max_episodes + self.rng = random.Random(seed) + + # Scan parquet files + data_dir = self.dataset_dir / "data" / "chunk-000" + parquet_files = sorted(data_dir.rglob("file-*.parquet")) + if not parquet_files: + raise FileNotFoundError("No parquet files in {}".format(data_dir)) + + # Build episode list from all parquet files + self.episodes = [] # (file_path, start_row, n_frames, task_index) + for pf in parquet_files: + table = pq.read_table(pf, columns=["episode_index", "task_index", "index"]) + # Find episode boundaries + ep_indices = table.column("episode_index").to_pylist() + task_indices = table.column("task_index").to_pylist() + + prev_ep = -1 + ep_start = 0 + for row_idx in range(len(ep_indices)): + if ep_indices[row_idx] != prev_ep: + if prev_ep != -1: + # Episode ended at row_idx-1 + n = row_idx - ep_start + self.episodes.append((str(pf), ep_start, n, task_indices[ep_start])) + ep_start = row_idx + prev_ep = ep_indices[row_idx] + # Last episode + n = len(ep_indices) - ep_start + if n > 0: + self.episodes.append((str(pf), ep_start, n, task_indices[ep_start])) + + if self.max_episodes: + self.episodes = self.episodes[:self.max_episodes] + + # Load tasks (jsonl or parquet) + self.tasks = {} + tasks_jsonl = self.dataset_dir / "meta" / "tasks.jsonl" + tasks_parquet = self.dataset_dir / "meta" / "tasks.parquet" + if tasks_jsonl.exists(): + with open(tasks_jsonl) as f: + for line in f: + task = json.loads(line) + self.tasks[task["task_index"]] = task["task"] + elif tasks_parquet.exists(): + import pandas as pd + tdf = pd.read_parquet(str(tasks_parquet)) + for idx, row in tdf.iterrows(): + self.tasks[int(row["task_index"])] = str(idx) + + self._build_metadata() + + self.transform = DreamTransform( + default_instruction="Perform the default behavior.", + language_dropout_prob=0.0, always_use_default_instruction=False, + max_state_dim=44, max_action_dim=32, max_length=512, + state_horizon=self.state_horizon, action_horizon=self.action_horizon, + num_views=self.num_views, + embodiment_tag_mapping={"oxe_droid": 17}, + tokenizer_path=kwargs.get("tokenizer_path", "/root/autodl-tmp/checkpoints/umt5-xxl"), + ) + self.transform.set_metadata(self.merged_metadata["oxe_droid"]) + self.transform.train() + + n_ep = len(self.episodes) + logger.info("LiberoDataset: {} episodes, {} tasks".format(n_ep, len(self.tasks))) + + def _build_metadata(self): + ds = DatasetStatisticalValues( + max=np.ones(1), min=np.zeros(1), mean=np.zeros(1), + std=np.ones(1), q01=np.zeros(1), q99=np.ones(1)) + self.merged_metadata = { + "oxe_droid": DatasetMetadata( + statistics=DatasetStatistics( + state={"joint_position": ds}, action={"joint_position": ds}), + modalities=DatasetModalities( + video={"image": VideoMetadata( + resolution=(self.video_width, self.video_height), channels=3, fps=10)}, + state={"joint_position": StateActionMetadata( + absolute=True, shape=(8,), continuous=True)}, + action={"joint_position": StateActionMetadata( + absolute=True, shape=(7,), continuous=True)}, + ), + embodiment_tag=EmbodimentTag.OXE_DROID, + ) + } + + def reset_seed(self, new_seed): + self.seed = new_seed + self.rng = random.Random(new_seed) + + def __len__(self): + return max(len(self.episodes) * 10, 1) + + def __getitem__(self, idx): + ep_path, ep_start, n_frames, task_idx = self.rng.choice(self.episodes) + + max_start = max(0, n_frames - self.num_frames - self.action_horizon) + t = self.rng.randint(0, max_start) if max_start > 0 else 0 + abs_start = ep_start + t + + # Read episode data from parquet + pf = pq.ParquetFile(ep_path) + # Only read the rows we need (from abs_start to abs_start+num_frames+action_horizon) + n_read = min(self.num_frames + self.action_horizon, n_frames - t) + table = pf.read_row_group(0, columns=[ + "observation.images.image", "observation.state", "action" + ]) + + # Read video frames + images_col = table.column("observation.images.image") + frames = [] + for i in range(self.num_frames): + fi = min(abs_start + i, ep_start + n_frames - 1) + raw_val = images_col[fi].as_py() + if isinstance(raw_val, dict): + img_bytes = raw_val["bytes"] + elif isinstance(raw_val, bytes): + img_bytes = raw_val + else: + raise TypeError("Unexpected image type: {}".format(type(raw_val))) + img = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + frames.append(img) + + video = np.stack(frames, axis=0) # [T, H, W, C] + video = video[:, np.newaxis, :, :, :] # [T, V=1, H, W, C] + + if video.shape[2] != self.video_height or video.shape[3] != self.video_width: + resized = np.zeros((self.num_frames, 1, self.video_height, self.video_width, 3), dtype=np.uint8) + for ti in range(self.num_frames): + resized[ti, 0] = cv2.resize(video[ti, 0], (self.video_width, self.video_height)) + video = resized + + # Read state + state_col = table.column("observation.state") + state_val = np.array(state_col[abs_start].as_py(), dtype=np.float32) + state = np.tile(state_val, (self.state_horizon, 1)) + + # Read actions + action_col = table.column("action") + actions = np.zeros((self.action_horizon, 7), dtype=np.float32) + for i in range(self.action_horizon): + fi = min(abs_start + i, ep_start + n_frames - 1) + actions[i] = np.array(action_col[fi].as_py(), dtype=np.float32) + + # Task description + task_name = self.tasks.get(task_idx, "Manipulate the object on the table") + + raw = { + "video": video, + "state": state, + "action": actions, + "annotation.human.action.task_description": task_name, + } + return dict(self.transform(raw)) diff --git a/groot/vla/data/dataset/macro.py b/groot/vla/data/dataset/macro.py new file mode 100644 index 0000000000000000000000000000000000000000..94077443e7660765e4a56a60ba88280ebcd7e830 --- /dev/null +++ b/groot/vla/data/dataset/macro.py @@ -0,0 +1,16 @@ +# Deprecated +METADATA_FILENAME = "metadata.json" +DATA_FILENAME = "dataset.hdf5" +FULL_SET_NAME = "FullSet" + +# LeRobot +LE_ROBOT_METADATA_FILENAME = "metadata.json" +LE_ROBOT_MODALITY_FILENAME = "modality.json" +LE_ROBOT_FEATURES_FILENAME = "features.json" +LE_ROBOT_STATISTICS_FILENAME = "stats.json" +LE_ROBOT_EMBODIMENT_FILENAME = "embodiment.json" +# LeRobot extras +LE_ROBOT_EPISODE_FILENAME = "episodes.jsonl" +LE_ROBOT_TASKS_FILENAME = "tasks.jsonl" +LE_ROBOT_INFO_FILENAME = "info.json" +LE_ROBOT_METADATA_DIR = "meta" diff --git a/groot/vla/data/dataset/manifeel.py b/groot/vla/data/dataset/manifeel.py new file mode 100644 index 0000000000000000000000000000000000000000..7b781493bd0452b5986c8d20fd19f6c9780d4a34 --- /dev/null +++ b/groot/vla/data/dataset/manifeel.py @@ -0,0 +1,173 @@ +import os, random, logging +from pathlib import Path +import numpy as np +import torch +import torch.utils.data as data +import zarr +import cv2 + +from groot.vla.model.dreamzero.transform.dreamzero_cotrain import DreamTransform +from groot.vla.data.schema.lerobot import ( + DatasetMetadata, DatasetStatistics, DatasetStatisticalValues, + DatasetModalities, VideoMetadata, StateActionMetadata, +) +from groot.vla.data.schema.embodiment_tags import EmbodimentTag + +logger = logging.getLogger(__name__) +# Per-dimension q01/q99 statistics for ManiFeel actions and states +# Computed across all 8 tasks (64200 frames) +ACTION_Q01 = np.array([-0.637143, -0.707143, -1.440000, -0.231429, -0.245357, -6.835714, 0.000000], dtype=np.float32) +ACTION_Q99 = np.array([0.827143, 0.805714, 2.537143, 0.170357, 0.188571, 6.717857, 0.040000], dtype=np.float32) +STATE_Q01 = np.array([0.435339, -0.076816, 0.002771, 0.650748, -0.696577, -0.078420, -0.031633], dtype=np.float32) +STATE_Q99 = np.array([0.564063, 0.087766, 0.193609, 0.999999, 0.753267, 0.077396, 0.025305], dtype=np.float32) + + +def _normalize_q99(x, q01, q99, eps=1e-8): + """Normalize to [-1, 1] using q01/q99.""" + return np.clip(2.0 * (x - q01) / (q99 - q01 + eps) - 1.0, -1.0, 1.0) + + + + +class ManiFeelDataset(data.Dataset): + """ManiFeel dataset for DreamZero. Reads Zarr, applies DreamTransform.""" + + def __init__(self, dataset_dir, num_frames=33, action_horizon=24, + state_horizon=1, num_views=3, video_height=160, video_width=320, + max_episodes=None, seed=42, **kwargs): + self.dataset_dir = Path(dataset_dir) + self.num_frames = num_frames + self.action_horizon = action_horizon + self.state_horizon = state_horizon + self.num_views = num_views + self.video_height = video_height + self.video_width = video_width + self.seed = seed + self.max_episodes = max_episodes + self.rng = random.Random(seed) + + self.zarr_dirs = sorted([d for d in self.dataset_dir.iterdir() + if d.is_dir() and (d / ".zgroup").exists()]) + if not self.zarr_dirs: + raise FileNotFoundError("No zarr dirs in " + str(dataset_dir)) + + self.episodes = [] + for zarr_dir in self.zarr_dirs: + root = zarr.open(str(zarr_dir), mode="r") + ends = root["meta/episode_ends"][:] + prev = 0 + for end in ends: + if end > prev: + self.episodes.append((str(zarr_dir), prev, int(end - prev))) + prev = end + if self.max_episodes: + self.episodes = self.episodes[:self.max_episodes] + + self.task_name = self._infer_task_name() + self._build_metadata() + + self.transform = DreamTransform( + default_instruction="Perform the default behavior.", + language_dropout_prob=0.0, always_use_default_instruction=False, + max_state_dim=44, max_action_dim=32, max_length=512, + state_horizon=self.state_horizon, action_horizon=self.action_horizon, + num_views=self.num_views, + embodiment_tag_mapping={"oxe_droid": 17}, + tokenizer_path=kwargs.get("tokenizer_path", "/root/autodl-tmp/checkpoints/umt5-xxl"), + ) + self.transform.set_metadata(self.merged_metadata["oxe_droid"]) + self.transform.train() + + n_ep = len(self.episodes) + logger.info(f"ManiFeelDataset: {n_ep} episodes, task={self.task_name}") + + def _infer_task_name(self): + for zarr_dir in self.zarr_dirs: + return zarr_dir.name.split("_")[0].capitalize() + return "Manipulate" + + def _build_metadata(self): + ds = DatasetStatisticalValues( + max=np.ones(1), min=np.zeros(1), mean=np.zeros(1), + std=np.ones(1), q01=np.zeros(1), q99=np.ones(1)) + self.merged_metadata = { + "oxe_droid": DatasetMetadata( + statistics=DatasetStatistics( + state={"joint_position": ds}, action={"joint_position": ds}), + modalities=DatasetModalities( + video={"exterior_image_1_left": VideoMetadata( + resolution=(self.video_width, self.video_height), channels=3, fps=10)}, + state={"joint_position": StateActionMetadata( + absolute=True, shape=(7,), continuous=True)}, + action={"joint_position": StateActionMetadata( + absolute=True, shape=(6,), continuous=True)}, + ), + embodiment_tag=EmbodimentTag.OXE_DROID, + ) + } + + def reset_seed(self, new_seed): + self.seed = new_seed + self.rng = random.Random(new_seed) + + def __len__(self): + return max(len(self.episodes) * 10, 1) + + def __getitem__(self, idx): + ep = self.rng.choice(self.episodes) + zarr_path, ep_start, ep_len = ep + root = zarr.open(zarr_path, mode="r") + + max_start = max(0, ep_len - self.num_frames - self.action_horizon) + t = self.rng.randint(0, max_start) if max_start > 0 else 0 + abs_start = ep_start + t + actual_frames = min(self.num_frames, ep_len - t) + + frames = [] + for cam in ["front", "side", "wrist"]: + cam_data = root["data/" + cam] + cam_frames = [] + for i in range(actual_frames): + fi = abs_start + i + if fi >= cam_data.shape[0]: + fi = cam_data.shape[0] - 1 + frame = cam_data[fi] + frame = (frame * 255).clip(0, 255).astype(np.uint8) + cam_frames.append(frame) + frames.append(np.stack(cam_frames)) + + video = np.stack(frames, axis=0) + v, t_f, h, w, c = video.shape + if h != self.video_height or w != self.video_width: + resized = np.zeros((v, t_f, self.video_height, self.video_width, c), dtype=np.uint8) + for vi in range(v): + for ti in range(t_f): + resized[vi, ti] = cv2.resize(video[vi, ti], + (self.video_width, self.video_height)) + video = resized + video = video.transpose(1, 0, 2, 3, 4) + + state = np.tile(np.array(root["data/state"][abs_start], dtype=np.float32), + (self.state_horizon, 1)) + action_data = root["data/action"] + actions = np.zeros((self.action_horizon, action_data.shape[1]), dtype=np.float32) + for i in range(self.action_horizon): + fi = min(abs_start + i, action_data.shape[0] - 1) + actions[i] = np.array(action_data[fi], dtype=np.float32) + + # Normalize state and action to [-1, 1] using q01/q99 + state = _normalize_q99(state, STATE_Q01, STATE_Q99) + # Pad 6-dim actions to 7-dim (dim 6 = gripper, 0 for 6-d tasks) + if actions.shape[1] < 7: + padded = np.zeros((actions.shape[0], 7), dtype=np.float32) + padded[:, :actions.shape[1]] = actions + actions = padded + actions = _normalize_q99(actions, ACTION_Q01, ACTION_Q99) + + raw = { + "video": video, + "state": state, + "action": actions, + "annotation.human.action.task_description": self.task_name, + } + return dict(self.transform(raw)) diff --git a/groot/vla/data/dataset/metadata.py b/groot/vla/data/dataset/metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..bc7cefef8796b7dcfea8580f0d20259a0db932e0 --- /dev/null +++ b/groot/vla/data/dataset/metadata.py @@ -0,0 +1,261 @@ +# import importlib +# from pathlib import Path + +# import numpy as np +# import pandas as pd +# from tqdm import tqdm + +# import groot.vla.common.utils as U +# from groot.vla.data.schema import ( +# EmbodimentTag, +# LeRobotModalityMetadata, +# LeRobotStateActionMetadata, +# DatasetMetadata, +# ) + +# from .macro import ( +# FULL_SET_NAME, +# LE_ROBOT_EMBODIMENT_FILENAME, +# LE_ROBOT_FEATURES_FILENAME, +# LE_ROBOT_METADATA_FILENAME, +# LE_ROBOT_MODALITY_FILENAME, +# LE_ROBOT_STATISTICS_FILENAME, +# ) +# from .registry import EMBODIMENT_TAGS_TO_DATASET_PATHS + +# METADATA_DIR = Path(importlib.import_module("groot.vla.data").__file__).parent / "metadata" # type: ignore + + +# def calculate_dataset_statistics( +# parquet_paths: list[Path], features: list[str] | None = None +# ) -> dict: +# """Calculate the dataset statistics of all columns for a list of parquet files.""" +# # Dataset statistics +# all_low_dim_data_list = [] +# # Collect all the data +# for parquet_path in tqdm( +# sorted(list(parquet_paths)), +# desc="Collecting all parquet files...", +# ): +# # Load the parquet file +# parquet_data = pd.read_parquet(parquet_path) +# parquet_data = parquet_data +# all_low_dim_data_list.append(parquet_data) +# all_low_dim_data = pd.concat(all_low_dim_data_list, axis=0) +# # Compute dataset statistics +# num_steps = len(all_low_dim_data.index) +# dataset_statistics: dict = { +# "num_trajectories": len(all_low_dim_data_list), +# "total_trajectory_length": num_steps, +# } +# if features is None: +# features = list(all_low_dim_data.columns) +# for le_modality in features: +# print(f"Computing statistics for {le_modality}...") +# np_data = np.vstack( +# [np.asarray(x, dtype=np.float32) for x in all_low_dim_data[le_modality]] # type: ignore +# ) +# dataset_statistics[le_modality] = { +# "mean": np.mean(np_data, axis=0).tolist(), +# "std": np.std(np_data, axis=0).tolist(), +# "min": np.min(np_data, axis=0).tolist(), +# "max": np.max(np_data, axis=0).tolist(), +# "q01": np.quantile(np_data, 0.01, axis=0).tolist(), +# "q99": np.quantile(np_data, 0.99, axis=0).tolist(), +# } +# return dataset_statistics + + +# def get_metadata( +# embodiment_tag: EmbodimentTag, +# metadata_version: str, +# regenerate_stats: bool = False, +# regenerate_metadata: bool = False, +# ) -> DatasetMetadata: +# """Get the metadata corresponding to the given embodiment tag and metadata version.""" +# metadata_dir = METADATA_DIR / embodiment_tag.value / metadata_version +# metadata_path = metadata_dir / LE_ROBOT_METADATA_FILENAME +# if metadata_path.exists() and not regenerate_metadata: +# metadata = DatasetMetadata.model_validate_json(metadata_path.read_text()) +# return metadata + + +# def get_metadata( +# embodiment_tag: EmbodimentTag, +# metadata_version: str, +# regenerate_stats: bool = False, +# regenerate_metadata: bool = False, +# ) -> TrainableDatasetMetadata_V1_2: +# """Get the metadata corresponding to the given embodiment tag and metadata version. + +# Args: +# embodiment_tag: The embodiment tag to load the metadata for. +# metadata_version: The version of the metadata to load. +# generate_metadata: Whether to generate the metadata if it does not exist. +# """ +# metadata_dir = METADATA_DIR / embodiment_tag.value / metadata_version +# metadata_path = metadata_dir / LE_ROBOT_METADATA_FILENAME +# if metadata_path.exists() and not regenerate_metadata: +# metadata = TrainableDatasetMetadata_V1_2.model_validate_json(metadata_path.read_text()) +# return metadata + +# assert ( +# embodiment_tag in EMBODIMENT_TAGS_TO_DATASET_PATHS +# ), f"Embodiment tag {embodiment_tag} not found in dataset registry. Available tags: {EMBODIMENT_TAGS_TO_DATASET_PATHS.keys()}" + +# dataset_paths = EMBODIMENT_TAGS_TO_DATASET_PATHS[embodiment_tag] +# # Load supporting metadata +# le_modality_meta_path = metadata_dir / LE_ROBOT_MODALITY_FILENAME +# le_features_path = metadata_dir / LE_ROBOT_FEATURES_FILENAME +# embodiment_meta_path = metadata_dir / LE_ROBOT_EMBODIMENT_FILENAME +# le_modality_meta = LeRobotModalityMetadata.model_validate_json( +# le_modality_meta_path.read_text() +# ) +# le_features = U.load_json(le_features_path) +# embodiment_meta = U.load_json(embodiment_meta_path) +# # Load stats +# if regenerate_stats: +# le_statistics = None +# else: +# le_statistics_path = metadata_dir / LE_ROBOT_STATISTICS_FILENAME +# le_statistics = U.load_json(le_statistics_path) +# # Generate metadata +# metadata, le_statistics = generate_metadata( +# embodiment_tag=embodiment_tag, +# dataset_paths=dataset_paths, +# le_modality_meta=le_modality_meta, +# le_features=le_features, +# embodiment_meta=embodiment_meta, +# le_statistics=le_statistics, +# ) + +# # Save metadata +# print(f"Generated metadata at {metadata_path}") +# metadata_path.write_text(metadata.model_dump_json(indent=4)) +# # Save stats +# if regenerate_stats: +# le_statistics_path = metadata_dir / LE_ROBOT_STATISTICS_FILENAME +# U.dump_json(le_statistics, le_statistics_path, indent=4) + +# return metadata + + +# def generate_metadata( +# embodiment_tag: EmbodimentTag, +# dataset_paths: list[Path], +# le_modality_meta: LeRobotModalityMetadata, +# le_features: dict, +# embodiment_meta: dict, +# le_statistics: dict | None = None, +# ): +# dataset_name = f"{embodiment_tag.value}:{FULL_SET_NAME}" + +# # Generate our custom modality metadata +# our_modality_meta: dict[str, dict] = {} +# for modality in ["state", "action"]: +# our_modality_meta[modality] = {} +# le_state_action_meta: dict[str, LeRobotStateActionMetadata] = getattr( +# le_modality_meta, modality +# ) +# for subkey in le_state_action_meta: +# state_action_dtype = np.dtype(le_state_action_meta[subkey].dtype) +# if np.issubdtype(state_action_dtype, np.floating): +# continuous = True +# else: +# continuous = False +# our_modality_meta[modality][subkey] = { +# "absolute": le_state_action_meta[subkey].absolute, +# "rotation_type": le_state_action_meta[subkey].rotation_type, +# "shape": [le_state_action_meta[subkey].end - le_state_action_meta[subkey].start], +# "continuous": continuous, +# } + +# # Add video modalities +# our_modality_meta["video"] = {} +# for new_key in le_modality_meta.video: +# original_key = le_modality_meta.video[new_key].original_key +# le_video_meta = le_features[original_key] +# height = le_video_meta["shape"][le_video_meta["names"].index("height")] +# width = le_video_meta["shape"][le_video_meta["names"].index("width")] +# channels = le_video_meta["shape"][le_video_meta["names"].index("channel")] +# if "info" in le_video_meta: +# fps = le_video_meta["info"]["video.fps"] +# elif "video_info" in le_video_meta: +# fps = le_video_meta["video_info"]["video.fps"] +# else: +# raise ValueError( +# f"Video modality {new_key} does not contain video_info or info: {le_video_meta.keys()}" +# ) +# our_modality_meta["video"][new_key] = { +# "resolution": [width, height], +# "channels": channels, +# "fps": fps, +# } + +# # Add annotation metadata +# our_modality_meta["annotation"] = {} +# if le_modality_meta.annotation is not None: +# for annotation_key in le_modality_meta.annotation: +# key_split = annotation_key.split(".") +# annotation_source = key_split[0] +# annotation_type = ".".join(key_split[1:]) +# if annotation_source not in our_modality_meta["annotation"]: +# our_modality_meta["annotation"][annotation_source] = [] +# our_modality_meta["annotation"][annotation_source].append(annotation_type) + +# lowdim_features = [] +# for feature in le_features: +# if "float" in le_features[feature]["dtype"]: +# lowdim_features.append(feature) + +# # Dataset statistics +# if le_statistics is None: +# print(f"Calculating dataset statistics for {dataset_name}") +# # Get all parquet files in the dataset paths +# parquet_files = [] +# for dataset_path in dataset_paths: +# parquet_files.extend(list(dataset_path.glob("data/*/*.parquet"))) +# le_statistics = calculate_dataset_statistics(parquet_files, lowdim_features) +# for le_modality in le_statistics: +# if not isinstance(le_statistics[le_modality], dict): +# continue +# for stat in le_statistics[le_modality]: +# le_statistics[le_modality][stat] = np.asarray(le_statistics[le_modality][stat]) + +# # Split statistics keys to our format +# dataset_statistics = { +# "num_trajectories": le_statistics["num_trajectories"], +# "total_trajectory_length": le_statistics["total_trajectory_length"], +# } +# for our_modality in ["state", "action"]: +# dataset_statistics[our_modality] = {} +# for subkey in our_modality_meta[our_modality]: +# dataset_statistics[our_modality][subkey] = {} +# state_action_meta = le_modality_meta.get_key_meta(f"{our_modality}.{subkey}") +# assert isinstance(state_action_meta, LeRobotStateActionMetadata) +# le_modality = state_action_meta.original_key +# for stat in le_statistics[le_modality]: +# indices = np.arange( +# state_action_meta.start, +# state_action_meta.end, +# ) +# dataset_statistics[our_modality][subkey][stat] = le_statistics[le_modality][stat][ +# indices +# ].tolist() + +# # Full dataset metadata +# metadata = TrainableDatasetMetadata_V1_2( +# dataset_name=dataset_name, +# dataset_statistics=dataset_statistics, # type: ignore +# modalities=our_modality_meta, # type: ignore +# embodiment=embodiment_meta, # type: ignore +# ) + +# # Convert stats from numpy to list +# for le_modality in le_statistics: +# if not isinstance(le_statistics[le_modality], dict): +# continue +# for stat in le_statistics[le_modality]: +# le_statistics[le_modality][stat] = le_statistics[le_modality][stat].tolist() + +# return metadata, le_statistics diff --git a/groot/vla/data/dataset/registry.py b/groot/vla/data/dataset/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..47bf672174db5bae46088bdef27fa408628ecdcc --- /dev/null +++ b/groot/vla/data/dataset/registry.py @@ -0,0 +1,14 @@ +from pathlib import Path + +from groot.vla.data.schema import EmbodimentTag + +# Dataset registry: maps embodiment tags to local dataset paths. +# Update these paths to point to your local dataset directories. +# See docs/DATASET_TO_GEAR_AND_TRAIN.md for dataset conversion instructions. +EMBODIMENT_TAGS_TO_DATASET_PATHS: dict[EmbodimentTag, list[Path]] = {} + +DATASET_PATHS_TO_EMBODIMENT_TAGS = { + path: dataset_tag + for dataset_tag, dataset_paths in EMBODIMENT_TAGS_TO_DATASET_PATHS.items() + for path in dataset_paths +} diff --git a/groot/vla/data/dataset/robotwin.py b/groot/vla/data/dataset/robotwin.py new file mode 100644 index 0000000000000000000000000000000000000000..6c6c08bb64fb4cfe1315fa1b0803e804c4573824 --- /dev/null +++ b/groot/vla/data/dataset/robotwin.py @@ -0,0 +1,239 @@ +#!/usr/bin/env python3 +"""DreamZero RoboTwin Dataset - reads qpos+videos+metas, outputs DreamTransform format.""" +import os, random, json, logging +from pathlib import Path +import numpy as np +import torch +import torch.utils.data as data +import cv2 + +from groot.vla.model.dreamzero.transform.dreamzero_cotrain import DreamTransform +from groot.vla.data.schema.lerobot import ( + DatasetMetadata, DatasetStatistics, DatasetStatisticalValues, + DatasetModalities, VideoMetadata, StateActionMetadata, +) +from groot.vla.data.schema.embodiment_tags import EmbodimentTag + +logger = logging.getLogger(__name__) + +# Per-dimension q01/q99 for RoboTwin velocity actions (14-dim) +# Computed across 8349 episodes +ACTION_Q01 = np.array([-0.002055, -0.003258, -0.002083, -0.002576, -0.002143, -0.004001, -0.006708, 0.0, -0.006258, -0.010985, -0.009753, -0.008991, -0.006935, -0.011197], dtype=np.float32) +ACTION_Q99 = np.array([0.004110, 0.004289, 0.002321, 0.006814, 0.002750, 0.002597, 0.006232, 0.0, 0.006063, 0.014186, 0.010926, 0.010506, 0.006669, 0.010019], dtype=np.float32) +STATE_Q01 = np.array([-0.2]*7 + [0.0]*7, dtype=np.float32) +STATE_Q99 = np.array([0.2]*7 + [0.0]*7, dtype=np.float32) + +def _normalize_q99(x, q01, q99, eps=1e-8): + return np.clip(2.0 * (x - q01) / (q99 - q01 + eps) - 1.0, -1.0, 1.0) + + +class RobotWinDataset(data.Dataset): + """RoboTwin dataset for DreamZero. Reads converted qpos+videos+metas format.""" + + def __init__(self, dataset_dir, num_frames=12, action_horizon=12, + state_horizon=1, num_views=1, video_height=160, video_width=320, + max_episodes=None, seed=42, **kwargs): + self.dataset_dir = Path(dataset_dir) + self.num_frames = num_frames + self.action_horizon = action_horizon + self.state_horizon = state_horizon + self.num_views = num_views + self.video_height = video_height + self.video_width = video_width + self.seed = seed + self.max_episodes = max_episodes + self.rng = random.Random(seed) + + # Scan task directories + self.tasks = sorted([d for d in self.dataset_dir.iterdir() if d.is_dir()]) + if not self.tasks: + raise FileNotFoundError("No task dirs in " + str(dataset_dir)) + + # Build episode list: (task_dir, ep_idx, n_video_frames, n_state_frames) + self.episodes = [] + for task_dir in self.tasks: + video_dir = task_dir / "videos" + qpos_dir = task_dir / "qpos" + meta_dir = task_dir / "metas" + if not video_dir.exists() or not qpos_dir.exists(): + continue + + video_files = sorted(video_dir.glob("episode*.mp4")) + qpos_files = sorted(qpos_dir.glob("episode*.pt")) + + for vf in video_files: + ep_name = vf.stem # episode0, episode1, etc. + qf = qpos_dir / (ep_name + ".pt") + if not qf.exists(): + continue + # Get frame count via cv2 + cap = cv2.VideoCapture(str(vf)) + n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + cap.release() + if n_frames < self.num_frames + self.action_horizon: + continue + # Load qpos to get state count + qpos_data = torch.load(str(qf)) + n_state = qpos_data.shape[0] + if n_state < self.num_frames + self.action_horizon: + continue + self.episodes.append((str(task_dir), str(vf), str(qf), n_frames, n_state, ep_name)) + + if not self.episodes: + raise FileNotFoundError("No valid episodes found (need >= {} frames)".format( + self.num_frames + self.action_horizon)) + + if self.max_episodes: + self.rng.shuffle(self.episodes) + self.episodes = self.episodes[:self.max_episodes] + + self._build_metadata() + + self.transform = DreamTransform( + default_instruction="Perform the default behavior.", + language_dropout_prob=0.0, always_use_default_instruction=False, + max_state_dim=44, max_action_dim=32, max_length=512, + state_horizon=self.state_horizon, action_horizon=self.action_horizon, + num_views=self.num_views, + embodiment_tag_mapping={"oxe_droid": 17}, + tokenizer_path=kwargs.get("tokenizer_path", "/root/autodl-tmp/checkpoints/umt5-xxl"), + ) + self.transform.set_metadata(self.merged_metadata["oxe_droid"]) + self.transform.train() + + n_ep = len(self.episodes) + logger.info("RobotWinDataset: {} episodes across {} tasks".format(n_ep, len(self.tasks))) + + def _build_metadata(self): + ds = DatasetStatisticalValues( + max=np.ones(1), min=np.zeros(1), mean=np.zeros(1), + std=np.ones(1), q01=np.zeros(1), q99=np.ones(1)) + self.merged_metadata = { + "oxe_droid": DatasetMetadata( + statistics=DatasetStatistics( + state={"joint_position": ds}, action={"joint_position": ds}), + modalities=DatasetModalities( + video={"image": VideoMetadata( + resolution=(self.video_width, self.video_height), channels=3, fps=30)}, + state={"joint_position": StateActionMetadata( + absolute=True, shape=(14,), continuous=True)}, + action={"joint_position": StateActionMetadata( + absolute=True, shape=(14,), continuous=True)}, + ), + embodiment_tag=EmbodimentTag.OXE_DROID, + ) + } + + def reset_seed(self, new_seed): + self.seed = new_seed + self.rng = random.Random(new_seed) + + def __len__(self): + return max(len(self.episodes) * 10, 1) + + def _load_task_desc(self, task_dir, ep_name): + """Load task description from metas directory.""" + meta_dir = task_dir / "metas" + if not meta_dir.exists(): + return "Manipulate the object on the table" + meta_files = sorted(meta_dir.glob("*.txt")) + if not meta_files: + return "Manipulate the object on the table" + # Extract episode index + ep_idx = int(ep_name.replace("episode", "")) + # Try to find the matching meta file + if ep_idx < len(meta_files): + with open(meta_files[ep_idx]) as f: + return f.read().strip() + return "Manipulate the object on the table" + + def _compute_actions(self, qpos_data, t, video_fps=30): + """Compute actions from state differences, matching video frame rate.""" + # qpos_data: [T_state, 14], sampled at 10Hz (from original 100Hz decimated) + # Video at 30fps + # We need 1 action per video frame + # If qpos has more timesteps than video frames, subsample qpos + # If qpos has fewer timesteps, interpolate + + # For simplicity: each video frame takes the action at the aligned state index + # action[i] = qpos[t + i + 1] - qpos[t + i] (velocity) + # But we need to map from video frame index to state index + + # Qpos is typically at ~10Hz, video at ~30Hz + # So ratio is state_per_video_frame = qpos_len / video_fps_per_segment + # For simplicity with num_frames=12 and action_horizon=12: + # just use state differences from qpos + + actions = np.zeros((self.action_horizon, 14), dtype=np.float32) + for i in range(self.action_horizon): + fi = min(t + i + 1, qpos_data.shape[0] - 1) + si = min(t + i, qpos_data.shape[0] - 1) + actions[i] = qpos_data[fi] - qpos_data[si] + return actions + + def __getitem__(self, idx): + task_dir, video_path, qpos_path, n_video_frames, n_qpos_frames, ep_name = \ + self.rng.choice(self.episodes) + + # Random start, ensuring enough frames for video + actions + max_start = max(0, n_video_frames - self.num_frames - self.action_horizon) + t = self.rng.randint(0, max_start) if max_start > 0 else 0 + + # Read video frames + cap = cv2.VideoCapture(video_path) + frames = [] + for i in range(self.num_frames): + fi = min(t + i, n_video_frames - 1) + cap.set(cv2.CAP_PROP_POS_FRAMES, fi) + ret, frame = cap.read() + if not ret: + frame = frames[-1] if frames else np.zeros((240, 320, 3), dtype=np.uint8) + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frames.append(frame) + cap.release() + + video = np.stack(frames, axis=0) # [T, H, W, C] + video = video[:, np.newaxis, :, :, :] # [T, V=1, H, W, C] + + # Resize if needed + if video.shape[2] != self.video_height or video.shape[3] != self.video_width: + resized = np.zeros((self.num_frames, 1, self.video_height, self.video_width, 3), dtype=np.uint8) + for ti in range(self.num_frames): + resized[ti, 0] = cv2.resize(video[ti, 0], (self.video_width, self.video_height)) + video = resized + + # Load qpos and compute state/action + qpos_data = torch.load(qpos_path).numpy() # [T_qpos, 14] + + # Map video frame index to qpos index + # qpos sampling rate = n_qpos_frames / video_duration_in_frames (approx) + video_to_qpos_ratio = n_qpos_frames / max(n_video_frames, 1) + qpos_start = int(t * video_to_qpos_ratio) + qpos_start = min(qpos_start, n_qpos_frames - 13) # ensure enough room + + # State: use the state at the aligned timestep + state_val = qpos_data[min(qpos_start, n_qpos_frames - 1)] + state = np.tile(state_val, (self.state_horizon, 1)) + + # Actions: velocity from consecutive qpos steps + actions = np.zeros((self.action_horizon, 14), dtype=np.float32) + for i in range(self.action_horizon): + fi = min(qpos_start + i + 1, n_qpos_frames - 1) + si = min(qpos_start + i, n_qpos_frames - 1) + actions[i] = qpos_data[fi] - qpos_data[si] + + # Task description + task_name = self._load_task_desc(Path(task_dir), ep_name) + + + # Normalize state and action to [-1, 1] + state = _normalize_q99(state, STATE_Q01, STATE_Q99) + actions = _normalize_q99(actions, ACTION_Q01, ACTION_Q99) + + raw = { + "video": video, + "state": state, + "action": actions, + "annotation.human.action.task_description": task_name, + } + return dict(self.transform(raw)) diff --git a/groot/vla/data/hf_datasets/__init__.py b/groot/vla/data/hf_datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/data/hf_datasets/dreamzero_libero.py b/groot/vla/data/hf_datasets/dreamzero_libero.py new file mode 100644 index 0000000000000000000000000000000000000000..c045f734d8c264631bb6c66b9bcaca2b0a1a04a3 --- /dev/null +++ b/groot/vla/data/hf_datasets/dreamzero_libero.py @@ -0,0 +1,79 @@ +"""LIBERO dataset loading script for HuggingFace datasets. + +Usage: + import datasets + ds = datasets.load_dataset("poet70/libero-benchmark", split="train", streaming=True) +""" + +import datasets +from datasets import Video, Value, Sequence + + +logger = datasets.logging.get_logger(__name__) + +FEATURES = datasets.Features({ + "episode_index": Value("int64"), + "frame_index": Value("int64"), + "state": Sequence(Value("float32"), length=44), + "action": Sequence(Sequence(Value("float32")), length=12), + "action_mask": Sequence(Value("bool"), length=12), + "text": Value("string"), + "video": Video(), +}) + + +class DreamZeroLIBERO(datasets.GeneratorBasedBuilder): + """LIBERO benchmark dataset for DreamZero.""" + + VERSION = datasets.Version("1.0.0") + BUILDER_CONFIGS = [ + datasets.BuilderConfig( + name="default", + version=VERSION, + description="LIBERO dataset for DreamZero SFT", + ), + ] + + def _info(self): + return datasets.DatasetInfo( + description="LIBERO benchmark dataset for DreamZero", + features=FEATURES, + homepage="https://huggingface.co/datasets/poet70/libero-benchmark", + license="cc-by-4.0", + ) + + def _split_generators(self, dl_manager): + # If the dataset is loaded from a local path, use it directly + # If from HF Hub, dl_manager will download + data_dir = self.config.data_dir + if data_dir is None: + # Default when loaded from HF Hub + data_dir = dl_manager.download_and_extract( + "https://huggingface.co/datasets/dreamzero/libero/resolve/main/data/train-00000.parquet" + ) + + return [ + datasets.SplitGenerator( + name=datasets.Split.TRAIN, + gen_kwargs={"data_dir": data_dir}, + ), + ] + + def _generate_examples(self, data_dir): + import pandas as pd + + data_dir_path = str(data_dir) + parquet_files = [f for f in sorted(os.listdir(data_dir_path)) if f.endswith(".parquet")] + + for parquet_file in parquet_files: + df = pd.read_parquet(os.path.join(data_dir_path, parquet_file)) + for idx, row in df.iterrows(): + yield idx, { + "episode_index": row["episode_index"], + "frame_index": row["frame_index"], + "state": row["state"], + "action": row["action"], + "action_mask": row["action_mask"], + "text": row["text"], + "video": row["video_path"], + } diff --git a/groot/vla/data/hf_datasets/dreamzero_manifeel.py b/groot/vla/data/hf_datasets/dreamzero_manifeel.py new file mode 100644 index 0000000000000000000000000000000000000000..c97d9b141a363d3cf4059973be5ae9cca682e188 --- /dev/null +++ b/groot/vla/data/hf_datasets/dreamzero_manifeel.py @@ -0,0 +1,78 @@ +"""ManiFeel dataset loading script for HuggingFace datasets. + +Usage: + import datasets + ds = datasets.load_dataset("dreamzero/manifeel", split="train", streaming=True) +""" + +import os +import datasets +from datasets import Video, Value, Sequence + + +FEATURES = datasets.Features({ + "task": Value("string"), + "episode_index": Value("int64"), + "frame_index": Value("int64"), + "state": Sequence(Value("float32"), length=44), + "action": Sequence(Sequence(Value("float32")), length=12), + "action_mask": Sequence(Value("bool"), length=12), + "text": Value("string"), + "video_0": Video(), + "video_1": Video(), + "video_2": Video(), +}) + + +class DreamZeroManiFeel(datasets.GeneratorBasedBuilder): + """ManiFeel benchmark dataset for DreamZero (3 views: front, wrist, side).""" + + VERSION = datasets.Version("1.0.0") + BUILDER_CONFIGS = [ + datasets.BuilderConfig( + name="default", + version=VERSION, + description="ManiFeel dataset for DreamZero SFT (3 views)", + ), + ] + + def _info(self): + return datasets.DatasetInfo( + description="ManiFeel benchmark dataset for DreamZero (3 views)", + features=FEATURES, + homepage="https://huggingface.co/datasets/poet70/manifeel-benchmark", + license="cc-by-4.0", + ) + + def _split_generators(self, dl_manager): + data_dir = self.config.data_dir + return [ + datasets.SplitGenerator( + name=datasets.Split.TRAIN, + gen_kwargs={"data_dir": data_dir}, + ), + ] + + def _generate_examples(self, data_dir): + import pandas as pd + + data_dir_path = str(data_dir) + parquet_files = sorted([ + f for f in os.listdir(data_dir_path) if f.endswith(".parquet") + ]) + + for parquet_file in parquet_files: + df = pd.read_parquet(os.path.join(data_dir_path, parquet_file)) + for idx, row in df.iterrows(): + yield idx, { + "task": row["task"], + "episode_index": row["episode_index"], + "frame_index": row["frame_index"], + "state": row["state"], + "action": row["action"], + "action_mask": row["action_mask"], + "text": row["text"], + "video_0": row["video_path_0"], + "video_1": row["video_path_1"], + "video_2": row["video_path_2"], + } diff --git a/groot/vla/data/hf_datasets/dreamzero_robotwin.py b/groot/vla/data/hf_datasets/dreamzero_robotwin.py new file mode 100644 index 0000000000000000000000000000000000000000..7e2be5f7fabc78ea6e77c1a14e82bd5c8baa56c5 --- /dev/null +++ b/groot/vla/data/hf_datasets/dreamzero_robotwin.py @@ -0,0 +1,74 @@ +"""RoboTwin dataset loading script for HuggingFace datasets. + +Usage: + import datasets + ds = datasets.load_dataset("dreamzero/robotwin", split="train", streaming=True) +""" + +import os +import datasets +from datasets import Video, Value, Sequence + + +FEATURES = datasets.Features({ + "task": Value("string"), + "episode_index": Value("int64"), + "frame_index": Value("int64"), + "state": Sequence(Value("float32"), length=44), + "action": Sequence(Sequence(Value("float32")), length=12), + "action_mask": Sequence(Value("bool"), length=12), + "text": Value("string"), + "video": Video(), +}) + + +class DreamZeroRoboTwin(datasets.GeneratorBasedBuilder): + """RoboTwin benchmark dataset for DreamZero.""" + + VERSION = datasets.Version("1.0.0") + BUILDER_CONFIGS = [ + datasets.BuilderConfig( + name="default", + version=VERSION, + description="RoboTwin dataset for DreamZero SFT", + ), + ] + + def _info(self): + return datasets.DatasetInfo( + description="RoboTwin benchmark dataset for DreamZero", + features=FEATURES, + homepage="https://huggingface.co/datasets/poet70/robotwin-benchmark", + license="cc-by-4.0", + ) + + def _split_generators(self, dl_manager): + data_dir = self.config.data_dir + return [ + datasets.SplitGenerator( + name=datasets.Split.TRAIN, + gen_kwargs={"data_dir": data_dir}, + ), + ] + + def _generate_examples(self, data_dir): + import pandas as pd + + data_dir_path = str(data_dir) + parquet_files = sorted([ + f for f in os.listdir(data_dir_path) if f.endswith(".parquet") + ]) + + for parquet_file in parquet_files: + df = pd.read_parquet(os.path.join(data_dir_path, parquet_file)) + for idx, row in df.iterrows(): + yield idx, { + "task": row["task"], + "episode_index": row["episode_index"], + "frame_index": row["frame_index"], + "state": row["state"], + "action": row["action"], + "action_mask": row["action_mask"], + "text": row["text"], + "video": row["video_path"], + } diff --git a/groot/vla/data/schema/__init__.py b/groot/vla/data/schema/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..786e3cb513a829954ce7709786f37a5393e722a2 --- /dev/null +++ b/groot/vla/data/schema/__init__.py @@ -0,0 +1,15 @@ +from .embodiment_tags import EmbodimentTag +from .lerobot import ( + DatasetMetadata, + DatasetModalities, + DatasetStatisticalValues, + DatasetStatistics, + LeRobotActionMetadata, + LeRobotModalityField, + LeRobotModalityMetadata, + LeRobotStateActionMetadata, + LeRobotStateMetadata, + RotationType, + StateActionMetadata, + VideoMetadata, +) diff --git a/groot/vla/data/schema/embodiment_tags.py b/groot/vla/data/schema/embodiment_tags.py new file mode 100644 index 0000000000000000000000000000000000000000..e7eaec04ec1e2a03dfc29e4863f3f15eae9e05e7 --- /dev/null +++ b/groot/vla/data/schema/embodiment_tags.py @@ -0,0 +1,348 @@ +from enum import Enum + + +class EmbodimentTag(Enum): + REAL_GR1_ARMS_ONLY = "real_gr1_arms_only" + """ + The real GR1 robot embodiment with arms only. + """ + + REAL_GR1_ARMS_ONLY_ANNOTATED = "real_gr1_arms_only_annotated" + """ + The real GR1 robot embodiment with arms only with annotations. + """ + + REAL_GR1_ARMS_WAIST = "real_gr1_arms_waist" + """ + The real GR1 robot embodiment with arms and waist. + """ + + REAL_GR1_ARMS_WAIST_ANNOTATED = "real_gr1_arms_waist_annotated" + """ + The real GR1 robot embodiment with arms and waist with annotations. + """ + + DEXMG_GR1_ARMS_ONLY_INSPIRE = "dexmg_gr1_arms_only_inspire" + """ + The DEXMG GR1 dataset with arms only and inspire hand. + """ + + DEXMG_GR1_ARMS_ONLY_FOURIER = "dexmg_gr1_arms_only_fourier" + """ + The DEXMG GR1 dataset with arms only and Fourier hand. + """ + + DEXMG_GR1_ARMS_WAIST_FOURIER = "dexmg_gr1_arms_waist_fourier" + """ + The DEXMG GR1 dataset with arms and waist and Fourier hand. + """ + + ROBOCASA_SINGLE_ARM = "robocasa_single_arm" + """ + The RoboCasa dataset with single arm. + """ + + ONE_X_EVE_GRIPPER = "onex_eve_gripper" + """ + The OneX Eve Robot with gripper. + """ + + ROBOCASA_GR1_ARMS_ONLY_INSPIRE_HANDS = "robocasa_gr1_arms_only_inspire_hands" + """ + The RoboCasa GR1 dataset with arms only and inspire hands. + """ + + ROBOCASA_GR1_ARMS_ONLY_FOURIER_HANDS = "robocasa_gr1_arms_only_fourier_hands" + """ + The RoboCasa GR1 dataset with arms only and Fourier hands. + """ + + ROBOCASA_GR1_FIXED_LOWER_BODY_INSPIRE_HANDS = "robocasa_gr1_fixed_lower_body_inspire_hands" + """ + The RoboCasa GR1 dataset with fixed lower body and inspire hands. + """ + + ROBOCASA_GR1_FIXED_LOWER_BODY_FOURIER_HANDS = "robocasa_gr1_fixed_lower_body_fourier_hands" + """ + The RoboCasa GR1 dataset with fixed lower body and Fourier hands. + """ + + ROBOCASA_GR1_ARMS_WAIST_FOURIER_HANDS = "robocasa_gr1_arms_waist_fourier_hands" + """ + The RoboCasa GR1 dataset with arms and waist and Fourier hands. + """ + + ROBOCASA_PANDA_OMRON = "robocasa_panda_omron" + """ + The RoboCasa dataset with panda omron. + """ + + ROBOCASA_SINGLE_ARM_PANDA_OMRON = "robocasa_single_arm_panda_omron" + """ + The RoboCasa dataset with single arm panda omron. + """ + + ROBOCASA_BIMANUAL_PANDA_PARALLEL_GRIPPER = "robocasa_bimanual_panda_parallel_gripper" + """ + The dexmg bimanual panda dataset with parallel grippers. + """ + + ROBOCASA_BIMANUAL_PANDA_INSPIRE_HAND = "robocasa_bimanual_panda_inspire_hand" + """ + The DEXMG bimanual panda dataset with inspire hands. + """ + + OXE_DROID = "oxe_droid" + """ + The Open X-Embodiment droid dataset. + """ + + OXE_FRACTAL = "oxe_fractal" + """ + The Open X-Embodiment fractal (RT-1) dataset. + """ + + OXE_LANGUAGE_TABLE = "oxe_language_table" + """ + The Open X-Embodiment language table dataset. + """ + + OXE_BRIDGE = "oxe_bridge" + """ + The Open X-Embodiment bridge dataset. + """ + + OXE_MUTEX = "oxe_mutex" + """ + The Open X-Embodiment mutex dataset. + """ + + OXE_ROBOSET = "oxe_roboset" + """ + The Open X-Embodiment Roboset dataset. + """ + + OXE_PLEX = "oxe_plex" + """ + The Open X-Embodiment Plex RoboSuite dataset. + """ + + REAL_PANDA_SINGLE_ARM = "real_panda_single_arm" + """ + The real single arm panda robot. + """ + + HOT3D_HANDS_ONLY = "hot3d_hands_only" + """ + The HOT3D dataset with hands only. + """ + + GR1_UNIFIED = "gr1_unified" + """ + The GR1 unified dataset. + """ + + GR1_ISAAC = "gr1_isaac" + """ + The GR1 Isaac dataset (Shiwei). + """ + + LAPA = "lapa" + """ + The datasets with LAPA actions. + """ + AGIBOT = "agibot" + + YAM = "yam" + + DREAM = "dream" + """ + The datasets with DREAM actions. + """ + + UNKNOWN = "unknown" + + GR1_UNIFIED_SEGMENTATION = "gr1_unified_segmentation" + """ + The GR1 unified dataset with segmentation. + """ + + LANGUAGE_TABLE_SIM = "language_table_sim" + """ + Simulated Language Table. + """ + + SIMPLER_ENV_GOOGLE = "simpler_env_google" + """ + SimplerEnv Google. + """ + + SIMPLER_ENV_WIDOWX = "simpler_env_widowx" + """ + SimplerEnv Widowx. + """ + + LIBERO_SIM = "libero_sim" + """ + The Libero Sim dataset. + """ + + DROID_SIM = "droid_sim" + """ + The Droid dataset in sim. + """ + + UNITREE_G1_UPPER_BODY = "unitree_g1_upper_body" + """ + The Unitree G1 dataset. + """ + + UNITREE_G1_FULL_BODY = "unitree_g1_full_body" + """ + The Unitree G1 dataset with full body. + """ + + UNITREE_G1_FULL_BODY_IN_SIM = "unitree_g1_full_body_in_sim" + """ + The Unitree G1 dataset in sim, data is collected with 50Hz. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT = "unitree_g1_full_body_with_height" + """ + The Unitree G1 dataset with full body and height command. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT_AND_EEF = "unitree_g1_full_body_with_height_and_eef" + """ + The Unitree G1 dataset with full body and height command and eef command. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT_EEF_NAV_CMD = "unitree_g1_full_body_with_height_eef_nav_cmd" + """ + The Unitree G1 dataset with full body and height command and eef command and navigate command. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT_NAV_CMD = "unitree_g1_full_body_with_height_nav_cmd" + """ + The Unitree G1 dataset with full body and height command and navigate command. + """ + + UNITREE_G1_FULL_BODY_WITH_WAIST_HEIGHT_NAV_CMD = ( + "unitree_g1_full_body_with_waist_height_nav_cmd" + ) + """ + The Unitree G1 dataset with full body and waist and height command and navigate command. + """ + + UNITREE_G1_FULL_BODY_WITH_WAIST_HEIGHT_NAV_CMD_AND_TASK_PROG = ( + "unitree_g1_full_body_with_waist_height_nav_cmd_and_task_prog" + ) + """ + The Unitree G1 dataset with full body and waist and height command and navigate command and task progress. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT_NAV_CMD_IN_SIM = ( + "unitree_g1_full_body_with_height_nav_cmd_in_sim" + ) + """ + The Unitree G1 dataset with full body and height command and navigate command in sim, data is collected with 50Hz. + """ + + G1_FIX_LOWER_RIGHT_HAND = "g1_fix_lower_right_hand" + """ + The G1 robot with fixed lower body and right hand. + """ + + GR1_UNIFIED_OFFLINE_RL = "gr1_unified_offline_rl" + """ + The GR1 robot with reward information for oflfine RL. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT_NAV_CMD_OAK_STEREO = ( + "unitree_g1_full_body_with_height_nav_cmd_oak_stereo" + ) + """ + The Unitree G1 dataset with full body and height command and navigate command and oak stereo. + """ + + XDOF = "xdof" + """ + The XDOF robot. + """ + + XDOF_H16 = "xdof_h16" + """ + The XDOF robot with action horizon 16. + """ + + XDOF_OSS_DATA = "xdof_oss_data" + """ + The XDOF data with conversions via jimmywu's processing pipeline. + """ + + SIM_BEHAVIOR_R1_PRO = "sim_behavior_r1_pro" + """ + The sim BEHAVIOR Galaxea R1 Pro robot with grippers. + """ + + # ------------- Deprecated G1 Embodiments ------------- + # below are deprecated G1 embodiments. Why deprecated? + # 1. The neck wasn't locked. So the the camera angles differe slightly between episodes and robots + # 2. Video modality was called `rs_view`. It is now called `ego_view` + # TODO: In future if we want to unify these embodiment tags, we can do so by renaming + # video.rs_view -> video.ego_view. After that we can remove the below deprecated embodiment tags. + DEPRECATED_UNITREE_G1_UPPER_BODY = "deprecated_unitree_g1_upper_body" + """ + The Unitree G1 dataset. + """ + + DEPRECATED_UNITREE_G1_FULL_BODY = "deprecated_unitree_g1_full_body" + """ + The Unitree G1 dataset with full body. + """ + + DEPRECATED_UNITREE_G1_FULL_BODY_WITH_HEIGHT = "deprecated_unitree_g1_full_body_with_height" + """ + The Unitree G1 dataset with full body and height command. + """ + + DEPRECATED_UNITREE_G1_FULL_BODY_WITH_HEIGHT_AND_EEF = ( + "deprecated_unitree_g1_full_body_with_height_and_eef" + ) + """ + The Unitree G1 dataset with full body and height command and eef command. + """ + + DEPRECATED_UNITREE_G1_FULL_BODY_WITH_HEIGHT_EEF_NAV_CMD = ( + "deprecated_unitree_g1_full_body_with_height_eef_nav_cmd" + ) + """ + The Unitree G1 dataset with full body and height command and eef command and navigate command. + """ + + DEPRECATED_UNITREE_G1_FULL_BODY_WITH_HEIGHT_NAV_CMD = ( + "deprecated_unitree_g1_full_body_with_height_nav_cmd" + ) + """ + The Unitree G1 dataset with full body and height command and navigate command. + """ + + GR1_UNIFIED_512 = "gr1_unified_512" + """ + The GR1 unified dataset with 512*320 resolution without cropping. + """ + + R1_PRO = "r1_pro" + """ + Real Galaxea R1 Pro. + """ + + R1_PRO_SINGLE_VIEW = "r1_pro_single-view" + """ + Real Galaxea R1 Pro with single camera view. + """ + + MECKA_HANDS = "mecka_hands" + """ + The Mecka robot with hands. + """ diff --git a/groot/vla/data/schema/lerobot.py b/groot/vla/data/schema/lerobot.py new file mode 100644 index 0000000000000000000000000000000000000000..4139f5697e79a00edf5e845bdd9720677fbed553 --- /dev/null +++ b/groot/vla/data/schema/lerobot.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from enum import Enum +from typing import Optional + +import numpy as np +from pydantic import BaseModel, Field, field_serializer, field_validator, model_validator + +from .embodiment_tags import EmbodimentTag + +# Common schema + + +class RotationType(Enum): + """Type of rotation representation""" + + AXIS_ANGLE = "axis_angle" + QUATERNION = "quaternion" + ROTATION_6D = "rotation_6d" + MATRIX = "matrix" + EULER_ANGLES_RPY = "euler_angles_rpy" + EULER_ANGLES_RYP = "euler_angles_ryp" + EULER_ANGLES_PRY = "euler_angles_pry" + EULER_ANGLES_PYR = "euler_angles_pyr" + EULER_ANGLES_YRP = "euler_angles_yrp" + EULER_ANGLES_YPR = "euler_angles_ypr" + + +# LeRobot schema + + +class LeRobotModalityField(BaseModel): + """Metadata for a LeRobot modality field.""" + + original_key: Optional[str] = Field( + default=None, + description="The original key of the modality in the LeRobot dataset", + ) + + +class LeRobotStateActionMetadata(LeRobotModalityField): + """Metadata for a LeRobot modality.""" + + start: int = Field( + ..., + description="The start index of the modality in the concatenated state/action vector", + ) + end: int = Field( + ..., + description="The end index of the modality in the concatenated state/action vector", + ) + rotation_type: Optional[RotationType] = Field( + default=None, description="The type of rotation for the modality" + ) + absolute: bool = Field(default=True, description="Whether the modality is absolute") + dtype: str = Field( + default="float64", + description="The data type of the modality. Defaults to float64.", + ) + range: Optional[tuple[float, float]] = Field( + default=None, + description="The range of the modality, if applicable. Defaults to None.", + ) + original_key: str = Field( + ..., + description="The original key of the modality in the LeRobot dataset.", + ) + + +class LeRobotStateMetadata(LeRobotStateActionMetadata): + """Metadata for a LeRobot state modality.""" + + original_key: str = Field( + default="observation.state", # LeRobot convention for states + description="The original key of the state modality in the LeRobot dataset", + ) + + +class LeRobotActionMetadata(LeRobotStateActionMetadata): + """Metadata for a LeRobot action modality.""" + + original_key: str = Field( + default="action", # LeRobot convention for actions + description="The original key of the action modality in the LeRobot dataset", + ) + + +class LeRobotModalityMetadata(BaseModel): + """Metadata for a LeRobot modality.""" + + state: dict[str, LeRobotStateMetadata] = Field( + ..., + description="The metadata for the state modality. The keys are the names of each split of the state vector.", + ) + action: dict[str, LeRobotActionMetadata] = Field( + ..., + description="The metadata for the action modality. The keys are the names of each split of the action vector.", + ) + video: dict[str, LeRobotModalityField] = Field( + ..., + description="The metadata for the video modality. The keys are the new names of each video modality.", + ) + annotation: Optional[dict[str, LeRobotModalityField]] = Field( + default=None, + description="The metadata for the annotation modality. The keys are the new names of each annotation modality.", + ) + + @model_validator(mode="after") + def check_original_keys(self): + # Generate default original keys for state + for key in self.state: + if self.state[key].original_key is None: + self.state[key].original_key = "observation.state" + + # Generate default original keys for action + for key in self.action: + if self.action[key].original_key is None: + self.action[key].original_key = "action" + + # Generate default original keys for video + for key in self.video: + if self.video[key].original_key is None: + self.video[key].original_key = "observation.images." + key + + # Generate default original keys for annotation + if self.annotation is not None: + for key in self.annotation: + if self.annotation[key].original_key is None: + self.annotation[key].original_key = "annotation." + key + + return self + + def get_key_meta(self, key: str) -> LeRobotModalityField: + """Get the metadata for a key in the LeRobot modality metadata. + + Args: + key (str): The key to get the metadata for. + + Returns: + LeRobotModalityField: The metadata for the key. + + Example: + lerobot_modality_meta = LeRobotModalityMetadata.model_validate(U.load_json(modality_meta_path)) + lerobot_modality_meta.get_key_meta("state.joint_shoulder_y") + lerobot_modality_meta.get_key_meta("video.main_camera") + lerobot_modality_meta.get_key_meta("annotation.human.action.task_description") + """ + split_key = key.split(".") + modality = split_key[0] + subkey = ".".join(split_key[1:]) + if modality == "state": + if subkey not in self.state: + raise ValueError( + f"Key: {key}, state key {subkey} not found in metadata, available state keys: {self.state.keys()}" + ) + return self.state[subkey] + elif modality == "action": + if subkey not in self.action: + raise ValueError( + f"Key: {key}, action key {subkey} not found in metadata, available action keys: {self.action.keys()}" + ) + return self.action[subkey] + elif modality == "video": + if subkey not in self.video: + raise ValueError( + f"Key: {key}, video key {subkey} not found in metadata, available video keys: {self.video.keys()}" + ) + return self.video[subkey] + elif modality == "annotation": + assert ( + self.annotation is not None + ), "Trying to get annotation metadata for a dataset with no annotations" + if subkey not in self.annotation: + raise ValueError( + f"Key: {key}, annotation key {subkey} not found in metadata, available annotation keys: {self.annotation.keys()}" + ) + return self.annotation[subkey] + else: + raise ValueError(f"Key: {key}, unexpected modality: {modality}") + + def get_original_key(self, key: str) -> str: + """Get the original key for a key in the LeRobot modality metadata. + + Args: + key (str): The key to get the original key for. + + Returns: + str: The original key for the key. + + Example: + lerobot_modality_meta = LeRobotModalityMetadata.model_validate(U.load_json(modality_meta_path)) + lerobot_modality_meta.get_original_key("state.joint_shoulder_y") + lerobot_modality_meta.get_original_key("video.main_camera") + lerobot_modality_meta.get_original_key("annotation.human.action.task_description") + """ + original_key = self.get_key_meta(key).original_key + assert original_key is not None, f"Key: {key}, original key is not set" + return original_key + + +# Dataset schema (parsed from LeRobot schema and simplified) + + +class DatasetStatisticalValues(BaseModel): + max: np.ndarray = Field(..., description="Maximum values") + min: np.ndarray = Field(..., description="Minimum values") + mean: np.ndarray = Field(..., description="Mean values") + std: np.ndarray = Field(..., description="Standard deviation") + q01: np.ndarray = Field(..., description="1st percentile values") + q99: np.ndarray = Field(..., description="99th percentile values") + + model_config = {"arbitrary_types_allowed": True} + + @field_serializer("*", when_used="json") + def serialize_ndarray(self, v: np.ndarray) -> list: + return v.tolist() # type: ignore + + @field_validator("*", mode="before") + def validate_ndarray(cls, v) -> np.ndarray: + return np.array(v) + + +class DatasetStatistics(BaseModel): + state: dict[str, DatasetStatisticalValues] = Field(..., description="Statistics of the state") + action: dict[str, DatasetStatisticalValues] = Field(..., description="Statistics of the action") + + +class VideoMetadata(BaseModel): + """Metadata of the video modality""" + + resolution: tuple[int, int] = Field(..., description="Resolution of the video") + channels: int = Field(..., description="Number of channels in the video", gt=0) + fps: float = Field(..., description="Frames per second", gt=0) + + +class StateActionMetadata(BaseModel): + absolute: bool = Field(..., description="Whether the state or action is absolute") + rotation_type: Optional[RotationType] = Field(None, description="Type of rotation, if any") + shape: tuple[int, ...] = Field(..., description="Shape of the state or action") + continuous: bool = Field(..., description="Whether the state or action is continuous") + + +class DatasetModalities(BaseModel): + video: dict[str, VideoMetadata] = Field(..., description="Metadata of the video") + state: dict[str, StateActionMetadata] = Field(..., description="Metadata of the state") + action: dict[str, StateActionMetadata] = Field(..., description="Metadata of the action") + + +class DatasetMetadata(BaseModel): + """Metadata of the trainable dataset + + Changes: + - Update to use the new RawCommitHashMetadataMetadata_V1_2 + """ + + statistics: DatasetStatistics = Field(..., description="Statistics of the dataset") + modalities: DatasetModalities = Field(..., description="Metadata of the modalities") + embodiment_tag: EmbodimentTag = Field(..., description="Embodiment tag of the dataset") diff --git a/groot/vla/data/transform/__init__.py b/groot/vla/data/transform/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bfa7eeeac139423679a12220cbde4db6161f0b38 --- /dev/null +++ b/groot/vla/data/transform/__init__.py @@ -0,0 +1,29 @@ +from .base import ( + ComposedModalityTransform, + InvertibleModalityTransform, + ModalityTransform, +) +from .concat import ConcatTransform +from .language import LanguageRemovePrefix, LanguageTransform +from .state_action import ( + PerHorizonActionTransform, + StateActionDropout, + StateActionPerturbation, + StateActionSinCosTransform, + StateActionToTensor, + StateActionTransform, +) +from .video import ( + VideoColorJitter, + VideoCrop, + VideoGrayscale, + VideoHorizontalFlip, + VideoRandomGrayscale, + VideoRandomPosterize, + VideoRandomRotation, + VideoResize, + VideoToNumpy, + VideoToTensor, + VideoTransform, + VideoNormalize, +) diff --git a/groot/vla/data/transform/base.py b/groot/vla/data/transform/base.py new file mode 100644 index 0000000000000000000000000000000000000000..228860049653dd3d1b0cb57150d8e561b7d4c935 --- /dev/null +++ b/groot/vla/data/transform/base.py @@ -0,0 +1,156 @@ +from abc import ABC, abstractmethod +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr + +from groot.vla.data.schema import DatasetMetadata + + +class ModalityTransform(BaseModel, ABC): + """ + Abstract class for transforming data modalities, e.g. video frame augmentation or action normalization. + """ + + apply_to: list[str] = Field(..., description="The keys to apply the transform to.") + training: bool = Field( + default=True, description="Whether to apply the transform in training mode." + ) + _dataset_metadata: DatasetMetadata | None = PrivateAttr(default=None) + + model_config = ConfigDict(arbitrary_types_allowed=True) + + @property + def dataset_metadata(self) -> DatasetMetadata: + assert ( + self._dataset_metadata is not None + ), "Dataset metadata is not set. Please call set_metadata() before calling apply()." + return self._dataset_metadata + + @dataset_metadata.setter + def dataset_metadata(self, value: DatasetMetadata): + self._dataset_metadata = value + + def set_metadata(self, dataset_metadata: DatasetMetadata): + """ + Set the dataset metadata. This is useful for transforms that need to know the dataset metadata, e.g. to normalize actions. + Subclasses can override this method if they need to do something more complex. + """ + self.dataset_metadata = dataset_metadata + + def __call__(self, data: dict[str, Any]) -> dict[str, Any]: + """Apply the transformation to the data corresponding to target_keys and return the processed data. + + Args: + data (dict[str, Any]): The data to transform. + example: data = { + "video.image_side_0": np.ndarray, + "action.eef_position": np.ndarray, + ... + } + + Returns: + dict[str, Any]: The transformed data. + example: transformed_data = { + "video.image_side_0": np.ndarray, + "action.eef_position": torch.Tensor, # Normalized and converted to tensor + ... + } + """ + return self.apply(data) + + @abstractmethod + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + """Apply the transformation to the data corresponding to keys matching the `apply_to` regular expression and return the processed data.""" + pass + + def train(self): + self.training = True + + def eval(self): + self.training = False + + +class InvertibleModalityTransform(ModalityTransform): + @abstractmethod + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + """Reverse the transformation to the data corresponding to keys matching the `apply_to` regular expression and return the processed data.""" + pass + + +class IdentityModalityTransform(ModalityTransform): + """Identity transform.""" + + apply_to: list[str] = Field( + default_factory=list, description="Will be ignored for identity transforms." + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + return data + + +class ComposedModalityTransform(ModalityTransform): + """Compose multiple modality transforms.""" + + transforms: list[ModalityTransform] = Field(..., description="The transforms to compose.") + apply_to: list[str] = Field( + default_factory=list, description="Will be ignored for composed transforms." + ) + training: bool = Field( + default=True, description="Whether to apply the transform in training mode." + ) + + model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True) + + def set_metadata(self, dataset_metadata: DatasetMetadata): + for transform in self.transforms: + transform.set_metadata(dataset_metadata) + # this is used to pass the list of transforms to concat transform + # concat transform needs needs to know what transforms were applied + # because it needs to compute the correct dimension of features + # post transform (during unapply). + # this attribute can also be used by other transforms to know what + # transforms were applied before it in the pipeline. + if hasattr(transform, "set_transform_pipeline"): + getattr(transform, "set_transform_pipeline")(self.transforms) + + def set_per_horizon_statistics(self, per_horizon_stats: dict[str, dict[str, list]]): + """Set per-horizon statistics for transforms that support it (e.g., PerHorizonActionTransform). + + Args: + per_horizon_stats: Dict from dataset.lerobot_relative_horizon_stats_meta + Format: {action_key: {stat_name: [[h0_vals], [h1_vals], ...]}} + """ + for transform in self.transforms: + if hasattr(transform, "set_per_horizon_statistics"): + transform.set_per_horizon_statistics(per_horizon_stats) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for i, transform in enumerate(self.transforms): + try: + data = transform(data) + except Exception as e: + raise ValueError(f"Error applying transform {i} to data: {e}") from e + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + for i, transform in enumerate(reversed(self.transforms)): + if isinstance(transform, InvertibleModalityTransform): + try: + data = transform.unapply(data) + except Exception as e: + step = len(self.transforms) - i - 1 + raise ValueError(f"Error unapplying transform {step} to data: {e}") from e + return data + + def train(self): + for transform in self.transforms: + transform.train() + self.training = True + + def eval(self): + for transform in self.transforms: + transform.eval() + self.training = False diff --git a/groot/vla/data/transform/concat.py b/groot/vla/data/transform/concat.py new file mode 100644 index 0000000000000000000000000000000000000000..e5474baffd56679fa2098e46abe4300fcff341a7 --- /dev/null +++ b/groot/vla/data/transform/concat.py @@ -0,0 +1,259 @@ +from typing import Any, Dict, List, Optional + +import numpy as np +from pydantic import Field, PrivateAttr +import torch + +from groot.vla.data.schema import DatasetMetadata, StateActionMetadata +from groot.vla.data.transform.base import InvertibleModalityTransform, ModalityTransform + + +class ConcatTransform(InvertibleModalityTransform): + """ + Concatenate the keys according to specified order. + """ + + # -- We inherit from ModalityTransform, so we keep apply_to as well -- + apply_to: list[str] = Field( + default_factory=list, description="Not used in this transform, kept for compatibility." + ) + + video_concat_order: list[str] = Field( + ..., + description="Concatenation order for each video modality. " + "Format: ['video.ego_view_pad_res224_freq20', ...]", + ) + + state_concat_order: Optional[list[str]] = Field( + default=None, + description="Concatenation order for each state modality. " + "Format: ['state.position', 'state.velocity', ...].", + ) + + action_concat_order: Optional[list[str]] = Field( + default=None, + description="Concatenation order for each action modality. " + "Format: ['action.position', 'action.velocity', ...].", + ) + + action_dims: dict[str, int] = Field( + default_factory=dict, + description="The dimensions of the action keys.", + ) + state_dims: dict[str, int] = Field( + default_factory=dict, + description="The dimensions of the state keys.", + ) + + action_dims_post_transform: dict[str, int] = Field( + default_factory=dict, + description="The new dimensions of the action keys after transform is applied.", + ) + state_dims_post_transform: dict[str, int] = Field( + default_factory=dict, + description="The new dimensions of the state keys after transform is applied.", + ) + # Store the transform pipeline to examine for dimension changes + _transform_pipeline: List[ModalityTransform] = PrivateAttr(default_factory=list) + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = { + "apply_to", + "video_concat_order", + "state_concat_order", + "action_concat_order", + } + else: + include = kwargs.pop("include", None) + + return super().model_dump(*args, include=include, **kwargs) + + def set_transform_pipeline(self, transforms: List[ModalityTransform]): + """Set the transform pipeline so this transform can examine it for dimension changes.""" + self._transform_pipeline = transforms + + def _get_target_rotations_from_pipeline(self) -> Dict[str, str]: + """Extract target_rotations from StateActionTransform instances in the pipeline.""" + target_rotations = {} + for transform in self._transform_pipeline: + if hasattr(transform, "target_rotations"): + transform_target_rotations = getattr(transform, "target_rotations", {}) + if transform_target_rotations: + target_rotations.update(transform_target_rotations) + return target_rotations + + def apply(self, data: Dict[str, Any]) -> Dict[str, Any]: + grouped_keys = {} + for key in data.keys(): + try: + modality, _ = key.split(".") + except: # noqa: E722 + ### Handle language annotation special case + if "annotation" in key: + modality = "language" + else: + modality = "others" + if modality not in grouped_keys: + grouped_keys[modality] = [] + grouped_keys[modality].append(key) + + if "video" in grouped_keys: + # Check if keys in video_concat_order, state_concat_order, action_concat_order are + # ineed contained in the data. If not, then the keys are misspecified + video_keys = grouped_keys["video"] + assert self.video_concat_order is not None, f"{self.video_concat_order=}, {video_keys=}" + assert all( + item in video_keys for item in self.video_concat_order + ), f"keys in video_concat_order are misspecified, \n{video_keys=}, \n{self.video_concat_order=}" + + # Process each video view + unsqueezed_videos = [] + for video_key in self.video_concat_order: + video_data = data.pop(video_key) + unsqueezed_video = np.expand_dims( + video_data, axis=-4 + ) # [..., H, W, C] -> [..., 1, H, W, C] + unsqueezed_videos.append(unsqueezed_video) + # Concatenate along the new axis + unsqueezed_video = np.concatenate(unsqueezed_videos, axis=-4) # [..., V, H, W, C] + + # Video + data["video"] = unsqueezed_video + + # "state" + if "state" in grouped_keys: + state_keys = grouped_keys["state"] + assert self.state_concat_order is not None, f"{self.state_concat_order=}" + assert all( + item in state_keys for item in self.state_concat_order + ), f"keys in state_concat_order are misspecified, \n{state_keys=}, \n{self.state_concat_order=}" + # Check the state dims + for key in self.state_concat_order: + target_shapes = [self.state_dims[key]] + if self.is_rotation_key(key): + target_shapes.extend( + [3, 4, 6] + ) # 3 -> axis_angle, 4 -> quaternion, 6 -> rotation_6d + target_shapes.append(self.state_dims[key] * 2) # Allow for sin-cos transform + assert ( + data[key].shape[-1] in target_shapes + ), f"State dim mismatch for {key=}, {data[key].shape[-1]=}, {target_shapes=}" + # Concatenate the state keys + # We'll have StateActionToTensor before this transform, so here we use torch.cat + data["state"] = torch.cat( + [data.pop(key) for key in self.state_concat_order], dim=-1 + ) # [T, D_state] + + if "action" in grouped_keys: + action_keys = grouped_keys["action"] + assert self.action_concat_order is not None, f"{self.action_concat_order=}" + # Check if all keys in concat_order are present + assert set(self.action_concat_order) == set( + action_keys + ), f"{set(self.action_concat_order)=}, {set(action_keys)=}" + # Record the action dims + for key in self.action_concat_order: + target_shapes = [self.action_dims[key]] + if self.is_rotation_key(key): + target_shapes.extend( + [3, 4, 6] + ) # 3 -> axis_angle, 4 -> quaternion, 6 -> rotation_6d + assert ( + data[key].shape[-1] in target_shapes + ), f"Action dim mismatch for {key=}, {data[key].shape[-1]=}, {target_shapes=}" + # Concatenate the action keys + # We'll have StateActionToTensor before this transform, so here we use torch.cat + data["action"] = torch.cat( + [data.pop(key) for key in self.action_concat_order], dim=-1 + ) # [T, D_action] + + return data + + def unapply(self, data: dict) -> dict: + start_dim = 0 + assert "action" in data, f"{data.keys()=}" + # For those dataset without actions (LAPA), we'll never run unapply + assert self.action_concat_order is not None, f"{self.action_concat_order=}" + action_tensor = data.pop("action") + for key in self.action_concat_order: + if key not in self.action_dims: + raise ValueError(f"Action dim {key} not found in action_dims.") + end_dim = start_dim + self.get_state_action_dims_post_transform(key) + data[key] = action_tensor[..., start_dim:end_dim] + start_dim = end_dim + if "state" in data: + assert self.state_concat_order is not None, f"{self.state_concat_order=}" + start_dim = 0 + state_tensor = data.pop("state") + for key in self.state_concat_order: + end_dim = start_dim + self.get_state_action_dims_post_transform(key) + data[key] = state_tensor[..., start_dim:end_dim] + start_dim = end_dim + return data + + def __call__(self, data: dict) -> dict: + return self.apply(data) + + def get_modality_metadata(self, key: str) -> StateActionMetadata: + modality, subkey = key.split(".") + assert self.dataset_metadata is not None, "Metadata not set" + modality_config = getattr(self.dataset_metadata.modalities, modality) + assert subkey in modality_config, f"{subkey=} not found in {modality_config=}" + assert isinstance( + modality_config[subkey], StateActionMetadata + ), f"Expected {StateActionMetadata} for {subkey=}, got {type(modality_config[subkey])=}" + return modality_config[subkey] + + def get_state_action_dims(self, key: str) -> int: + """Get the dimension of a state or action key from the dataset metadata.""" + modality_config = self.get_modality_metadata(key) + shape = modality_config.shape + assert len(shape) == 1, f"{shape=}" + return shape[0] + + def get_state_action_dims_post_transform(self, key: str) -> int: + """ + This function is used to get the dims of the state/action keys after transform is applied. + It is different from the `get_state_action_dims` function, because this function accounts for + the case where we apply transforms and the # of dims is change eg. after applying axis_angle transform on + quaternion, the dims change from 4D to 3D. + """ + modality_config = self.get_modality_metadata(key) + shape = modality_config.shape + assert len(shape) == 1, f"{shape=}" + + if self.is_rotation_key(key): + target_rotations = self._get_target_rotations_from_pipeline() + if key in target_rotations: + target_rotation = target_rotations[key] + if target_rotation == "axis_angle": + return 3 + elif target_rotation == "quaternion": + return 4 + elif target_rotation == "rotation_6d": + return 6 + elif target_rotation == "euler_angles": + return 3 + else: + raise ValueError(f"Unknown target rotation type: {target_rotation}") + else: + # No target rotation specified, return original dimension + return shape[0] + else: + return shape[0] + + def is_rotation_key(self, key: str) -> bool: + modality_config = self.get_modality_metadata(key) + return modality_config.rotation_type is not None + + def set_metadata(self, dataset_metadata: DatasetMetadata): + """Set the metadata and compute the dimensions of the state and action keys.""" + super().set_metadata(dataset_metadata) + # Pre-compute the dimensions of the state and action keys + if self.action_concat_order is not None: + for key in self.action_concat_order: + self.action_dims[key] = self.get_state_action_dims(key) + if self.state_concat_order is not None: + for key in self.state_concat_order: + self.state_dims[key] = self.get_state_action_dims(key) diff --git a/groot/vla/data/transform/language.py b/groot/vla/data/transform/language.py new file mode 100644 index 0000000000000000000000000000000000000000..6fbf4d5c6679f6ffacadc0220c3a94418fd23591 --- /dev/null +++ b/groot/vla/data/transform/language.py @@ -0,0 +1,86 @@ +from typing import Any, Union + +import numpy as np +from pydantic import Field, field_validator +from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast + +from groot.vla.data.transform.base import InvertibleModalityTransform, ModalityTransform + +T_Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast] + + +class LanguageTransform(InvertibleModalityTransform): + """Transform for language modalities. + + Attributes: + apply_to (list[str]): The keys in the modality to load and transform. + tokenizer (T_Tokenizer): The tokenizer to use. Can be either PreTrainedTokenizer, PreTrainedTokenizerFast, or path to the HuggingFace tokenizer + """ + + apply_to: list[str] = Field(..., description="The keys in the modality to load and transform.") + tokenizer: T_Tokenizer = Field(..., description="The tokenizer to use.") + + @field_validator("tokenizer") + def validate_tokenizer(cls, v: T_Tokenizer | str) -> T_Tokenizer: + if isinstance(v, str): + return AutoTokenizer.from_pretrained(v) + return v + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + """Tokenize the data, pad the tokens to the longest sequence, and concatenate the tokens. + + Args: + data (dict[str, Any]): The complete data dictionary. + + Returns: + dict[str, Any]: The processed data dictionary with the keys in `apply_to` replaced with the tokenized data. + """ + for key in self.apply_to: + data[key] = self.tokenizer( + data[key], return_tensors="pt", padding=True, truncation=True + ).input_ids + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + """Untokenize the data. + + Args: + data (dict[str, Any]): The processed data dictionary with the keys in `apply_to` replaced with the tokenized data. + + Returns: + dict[str, Any]: The untokenized data. + """ + for key in self.apply_to: + data[key] = self.tokenizer.decode(data[key], skip_special_tokens=True) + return data + + +class LanguageRemovePrefix(ModalityTransform): + apply_to: list[str] = Field( + ..., description="The keys in the modality to remove the prefix from." + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + """Remove the prefix from the language. + Expects: + - data[key] is a list of strings, shape (T,) + - OR data[key] is a list of lists of strings, shape (B, T) + + Args: + data (dict[str, Any]): The processed data dictionary with the keys in `apply_to` replaced with the tokenized data. + + Returns: + dict[str, Any]: The data with the prefix removed for key in `apply_to`. + """ + for key in self.apply_to: + value = data[key] + # Handle both batched (list of lists) and non-batched (list) language data. + if isinstance(value[0], np.ndarray): + # Batched case: list of lists of strings, shape (B, T) + data[key] = np.array( + [[lang.split(": ")[-1] for lang in sublist] for sublist in value] + ) + else: + # Non-batched case: list of strings, shape (T,) + data[key] = np.array([lang.split(": ")[-1] for lang in value]) + return data diff --git a/groot/vla/data/transform/state_action.py b/groot/vla/data/transform/state_action.py new file mode 100644 index 0000000000000000000000000000000000000000..2be1cc810f1dca0e2c9625d7775bc7b62b0c4816 --- /dev/null +++ b/groot/vla/data/transform/state_action.py @@ -0,0 +1,893 @@ +import functools +import random +from typing import Any, ClassVar + +import numpy as np +from pydantic import Field, PrivateAttr, field_validator, model_validator +import torch + +from groot.vla.data.schema import DatasetMetadata, RotationType, StateActionMetadata +from groot.vla.data.transform.base import InvertibleModalityTransform, ModalityTransform + + +class RotationTransform: + """Adapted from https://github.com/real-stanford/diffusion_policy/blob/548a52bbb105518058e27bf34dcf90bf6f73681a/diffusion_policy/model/common/rotation_transformer.py""" + + valid_reps = ["axis_angle", "euler_angles", "quaternion", "rotation_6d", "matrix"] + + def __init__(self, from_rep="axis_angle", to_rep="rotation_6d"): + """ + Valid representations + + Always use matrix as intermediate representation. + """ + if from_rep.startswith("euler_angles"): + from_convention = from_rep.split("_")[-1] + from_rep = "euler_angles" + from_convention = from_convention.replace("r", "X").replace("p", "Y").replace("y", "Z") + else: + from_convention = None + if to_rep.startswith("euler_angles"): + to_convention = to_rep.split("_")[-1] + to_rep = "euler_angles" + to_convention = to_convention.replace("r", "X").replace("p", "Y").replace("y", "Z") + else: + to_convention = None + assert from_rep != to_rep, f"from_rep and to_rep cannot be the same: {from_rep}" + assert from_rep in self.valid_reps, f"Invalid from_rep: {from_rep}" + assert to_rep in self.valid_reps, f"Invalid to_rep: {to_rep}" + + forward_funcs = list() + inverse_funcs = list() + + if from_rep != "matrix": + import pytorch3d.transforms as pt # noqa: F401 + + funcs = [getattr(pt, f"{from_rep}_to_matrix"), getattr(pt, f"matrix_to_{from_rep}")] + if from_convention is not None: + funcs = [functools.partial(func, convention=from_convention) for func in funcs] + forward_funcs.append(funcs[0]) + inverse_funcs.append(funcs[1]) + + if to_rep != "matrix": + funcs = [getattr(pt, f"matrix_to_{to_rep}"), getattr(pt, f"{to_rep}_to_matrix")] + if to_convention is not None: + funcs = [functools.partial(func, convention=to_convention) for func in funcs] + forward_funcs.append(funcs[0]) + inverse_funcs.append(funcs[1]) + + inverse_funcs = inverse_funcs[::-1] + + self.forward_funcs = forward_funcs + self.inverse_funcs = inverse_funcs + + @staticmethod + def _apply_funcs(x: torch.Tensor, funcs: list) -> torch.Tensor: + assert isinstance(x, torch.Tensor) + for func in funcs: + x = func(x) + return x + + def forward(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + return self._apply_funcs(x, self.forward_funcs) + + def inverse(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + return self._apply_funcs(x, self.inverse_funcs) + + +class Normalizer: + valid_modes = ["q99", "mean_std", "min_max", "binary", "scale"] + + def __init__(self, mode: str, statistics: dict): + self.mode = mode + self.statistics = statistics + for key, value in self.statistics.items(): + self.statistics[key] = torch.tensor(value) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + + # Normalize the tensor + if self.mode == "q99": + + # Range of q99 is [-1, 1] + q01 = self.statistics["q01"].to(x.dtype) + q99 = self.statistics["q99"].to(x.dtype) + + # In the case of q01 == q99, the normalization will be undefined + # So we set the normalized values to the original values + mask = q01 != q99 + normalized = torch.zeros_like(x) + + # Normalize the values where q01 != q99 + # Formula: 2 * (x - q01) / (q99 - q01) - 1 + normalized[..., mask] = (x[..., mask] - q01[..., mask]) / ( + q99[..., mask] - q01[..., mask] + ) + normalized[..., mask] = 2 * normalized[..., mask] - 1 + + # Set the normalized values to the original values where q01 == q99 + normalized[..., ~mask] = x[..., ~mask].to(x.dtype) + + # Clip the normalized values to be between -1 and 1 + normalized = torch.clamp(normalized, -1, 1) + + elif self.mode == "mean_std": + # Range of mean_std is not fixed, but can be positive or negative + mean = self.statistics["mean"].to(x.dtype) + std = self.statistics["std"].to(x.dtype) + + # In the case of std == 0, the normalization will be undefined + # So we set the normalized values to the original values + mask = std != 0 + normalized = torch.zeros_like(x) + + # Normalize the values where std != 0 + # Formula: (x - mean) / std + normalized[..., mask] = (x[..., mask] - mean[..., mask]) / std[..., mask] + + # Set the normalized values to the original values where std == 0 + normalized[..., ~mask] = x[..., ~mask].to(x.dtype) + + elif self.mode == "min_max": + # Range of min_max is [-1, 1] + min = self.statistics["min"].to(x.dtype) + max = self.statistics["max"].to(x.dtype) + + # In the case of min == max, the normalization will be undefined + # So we set the normalized values to the original values + mask = min != max + normalized = torch.zeros_like(x) + + # Normalize the values where min != max + # Formula: 2 * (x - min) / (max - min) - 1 + normalized[..., mask] = (x[..., mask] - min[..., mask]) / ( + max[..., mask] - min[..., mask] + ) + normalized[..., mask] = 2 * normalized[..., mask] - 1 + + # Set the normalized values to the original values where min == max + # normalized[..., ~mask] = x[..., ~mask].to(x.dtype) + # Set the normalized values to 0 where min == max + normalized[..., ~mask] = 0 + + if not (torch.all(normalized <= 1 + 1e-5) and torch.all(normalized >= -1 - 1e-5)): + print( + f"WARNING: normalization out of range: {x.shape=}, max: {torch.max(normalized)}, min: {torch.min(normalized)}" + ) + + # Clip the normalized values to be between -1 and 1 + normalized = torch.clamp(normalized, -1, 1) + + elif self.mode == "scale": + # Range of scale is [0, 1] + min = self.statistics["min"].to(x.dtype) + max = self.statistics["max"].to(x.dtype) + abs_max = torch.max(torch.abs(min), torch.abs(max)) + mask = abs_max != 0 + normalized = torch.zeros_like(x) + normalized[..., mask] = x[..., mask] / abs_max[..., mask] + normalized[..., ~mask] = 0 + + elif self.mode == "binary": + # Range of binary is [0, 1] + normalized = (x > 0.5).to(x.dtype) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + return normalized + + def inverse(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + if self.mode == "q99": + q01 = self.statistics["q01"].to(dtype=x.dtype, device=x.device) + q99 = self.statistics["q99"].to(dtype=x.dtype, device=x.device) + return (x + 1) / 2 * (q99 - q01) + q01 + elif self.mode == "mean_std": + mean = self.statistics["mean"].to(dtype=x.dtype, device=x.device) + std = self.statistics["std"].to(dtype=x.dtype, device=x.device) + return x * std + mean + elif self.mode == "min_max": + min = self.statistics["min"].to(dtype=x.dtype, device=x.device) + max = self.statistics["max"].to(dtype=x.dtype, device=x.device) + return (x + 1) / 2 * (max - min) + min + elif self.mode == "binary": + return (x > 0.5).to(dtype=x.dtype, device=x.device) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + +class PerHorizonNormalizer: + """Normalizer that applies different statistics for each horizon index. + + Unlike the standard Normalizer which uses single stats for all timesteps, + this normalizer uses per-horizon stats where each action horizon step has + its own normalization parameters. + + Args: + mode: Normalization mode ('q99', 'mean_std', 'min_max', 'binary', 'scale') + statistics: Dict with stat names as keys and 2D tensors of shape (horizon_len, action_dim) as values. + Format: {"q01": [[h0_vals], [h1_vals], ...], "q99": [[h0_vals], [h1_vals], ...], ...} + """ + valid_modes = ["q99", "mean_std", "min_max", "binary", "scale"] + + def __init__(self, mode: str, statistics: dict): + self.mode = mode + self.statistics = {} + for key, value in statistics.items(): + # Convert to tensor: shape (horizon_len, action_dim) + self.statistics[key] = torch.tensor(value) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Normalize tensor with per-horizon statistics. + + Args: + x: Input tensor of shape (horizon_len, action_dim) or (num_chunks * horizon_len, action_dim) + or (..., horizon_len, action_dim) + + Returns: + Normalized tensor of same shape + """ + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + + # Get the stats horizon length + stats_horizon_len = self.statistics["q01"].shape[0] + action_dim = x.shape[-1] + original_shape = x.shape + + # Handle case where x is flattened (num_chunks * horizon_len, action_dim) + # We need to reshape to (..., horizon_len, action_dim) for proper broadcasting + reshaped = False + if len(x.shape) == 2: + total_len = x.shape[0] + if total_len > stats_horizon_len and total_len % stats_horizon_len == 0: + # Flattened batch case: reshape to (num_chunks, horizon_len, action_dim) + num_chunks = total_len // stats_horizon_len + x = x.view(num_chunks, stats_horizon_len, action_dim) + reshaped = True + elif total_len == stats_horizon_len: + # Single chunk, add batch dimension for consistent processing + x = x.unsqueeze(0) + reshaped = True + + # Now x should have shape (..., horizon_len, action_dim) + horizon_len = x.shape[-2] + + if self.mode == "q99": + q01 = self.statistics["q01"].to(dtype=x.dtype, device=x.device)[:horizon_len] # (horizon_len, action_dim) + q99 = self.statistics["q99"].to(dtype=x.dtype, device=x.device)[:horizon_len] + + # Compute normalization with safe division + range_vals = q99 - q01 + range_vals = torch.where(range_vals == 0, torch.ones_like(range_vals), range_vals) + normalized = 2 * (x - q01) / range_vals - 1 + + # For zero-range values, keep original + mask = (q01 == q99).unsqueeze(0).expand_as(x) if len(x.shape) > 2 else (q01 == q99) + normalized = torch.where(mask, x, normalized) + normalized = torch.clamp(normalized, -1, 1) + + elif self.mode == "mean_std": + mean = self.statistics["mean"].to(dtype=x.dtype, device=x.device)[:horizon_len] + std = self.statistics["std"].to(dtype=x.dtype, device=x.device)[:horizon_len] + + std_safe = torch.where(std == 0, torch.ones_like(std), std) + normalized = (x - mean) / std_safe + + mask = (std == 0).unsqueeze(0).expand_as(x) if len(x.shape) > 2 else (std == 0) + normalized = torch.where(mask, x, normalized) + + elif self.mode == "min_max": + min_val = self.statistics["min"].to(dtype=x.dtype, device=x.device)[:horizon_len] + max_val = self.statistics["max"].to(dtype=x.dtype, device=x.device)[:horizon_len] + + range_vals = max_val - min_val + range_vals = torch.where(range_vals == 0, torch.ones_like(range_vals), range_vals) + normalized = 2 * (x - min_val) / range_vals - 1 + + mask = (min_val == max_val).unsqueeze(0).expand_as(x) if len(x.shape) > 2 else (min_val == max_val) + normalized = torch.where(mask, torch.zeros_like(x), normalized) + normalized = torch.clamp(normalized, -1, 1) + + elif self.mode == "scale": + min_val = self.statistics["min"].to(dtype=x.dtype, device=x.device)[:horizon_len] + max_val = self.statistics["max"].to(dtype=x.dtype, device=x.device)[:horizon_len] + abs_max = torch.max(torch.abs(min_val), torch.abs(max_val)) + + abs_max_safe = torch.where(abs_max == 0, torch.ones_like(abs_max), abs_max) + normalized = x / abs_max_safe + + mask = (abs_max == 0).unsqueeze(0).expand_as(x) if len(x.shape) > 2 else (abs_max == 0) + normalized = torch.where(mask, torch.zeros_like(x), normalized) + + elif self.mode == "binary": + normalized = (x > 0.5).to(x.dtype) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + # Reshape back to original shape + if reshaped: + normalized = normalized.view(original_shape) + + return normalized + + def inverse(self, x: torch.Tensor) -> torch.Tensor: + """Denormalize tensor with per-horizon statistics. + + Args: + x: Normalized tensor of shape (horizon_len, action_dim) or (num_chunks * horizon_len, action_dim) + or (..., horizon_len, action_dim) + + Returns: + Denormalized tensor of same shape + """ + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + + # Get the stats horizon length + stats_horizon_len = self.statistics["q01"].shape[0] + action_dim = x.shape[-1] + original_shape = x.shape + + # Handle case where x is flattened (num_chunks * horizon_len, action_dim) + reshaped = False + if len(x.shape) == 2: + total_len = x.shape[0] + if total_len > stats_horizon_len and total_len % stats_horizon_len == 0: + num_chunks = total_len // stats_horizon_len + x = x.view(num_chunks, stats_horizon_len, action_dim) + reshaped = True + elif total_len == stats_horizon_len: + x = x.unsqueeze(0) + reshaped = True + + horizon_len = x.shape[-2] + + if self.mode == "q99": + q01 = self.statistics["q01"].to(dtype=x.dtype, device=x.device)[:horizon_len] + q99 = self.statistics["q99"].to(dtype=x.dtype, device=x.device)[:horizon_len] + denormalized = (x + 1) / 2 * (q99 - q01) + q01 + elif self.mode == "mean_std": + mean = self.statistics["mean"].to(dtype=x.dtype, device=x.device)[:horizon_len] + std = self.statistics["std"].to(dtype=x.dtype, device=x.device)[:horizon_len] + denormalized = x * std + mean + elif self.mode == "min_max": + min_val = self.statistics["min"].to(dtype=x.dtype, device=x.device)[:horizon_len] + max_val = self.statistics["max"].to(dtype=x.dtype, device=x.device)[:horizon_len] + denormalized = (x + 1) / 2 * (max_val - min_val) + min_val + elif self.mode == "binary": + denormalized = (x > 0.5).to(dtype=x.dtype, device=x.device) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + # Reshape back to original shape + if reshaped: + denormalized = denormalized.view(original_shape) + + return denormalized + + +class PerHorizonActionTransform(InvertibleModalityTransform): + """Transform for action normalization with per-horizon statistics. + + This transform applies different normalization parameters for each action + horizon index, allowing for more precise normalization when the distribution + of relative actions varies across the prediction horizon. + + Args: + apply_to: List of action keys to transform (e.g., ['action.joint_position']) + normalization_modes: Dict mapping action keys to normalization modes + per_horizon_statistics: Dict mapping action keys to per-horizon stats. + Format: {action_key: {stat_name: [[h0_vals], [h1_vals], ...]}} + """ + + apply_to: list[str] = Field(..., description="The action keys to transform.") + normalization_modes: dict[str, str] = Field( + default_factory=dict, description="The normalization modes for each action key." + ) + per_horizon_statistics: dict[str, dict[str, list]] = Field( + default_factory=dict, description="Per-horizon statistics for each action key." + ) + + _normalizers: dict[str, PerHorizonNormalizer] = PrivateAttr(default_factory=dict) + _input_dtypes: dict[str, np.dtype | torch.dtype] = PrivateAttr(default_factory=dict) + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = {"apply_to", "normalization_modes"} + else: + include = kwargs.pop("include", None) + return super().model_dump(*args, include=include, **kwargs) + + def set_per_horizon_statistics(self, per_horizon_stats: dict[str, dict[str, list]]): + """Set per-horizon statistics and initialize normalizers. + + Args: + per_horizon_stats: Dict from dataset.lerobot_relative_horizon_stats_meta + Format: {action_key: {stat_name: [[h0_vals], [h1_vals], ...]}} + """ + self.per_horizon_statistics = per_horizon_stats + self._initialize_normalizers() + + def _initialize_normalizers(self): + """Initialize per-horizon normalizers from statistics.""" + for key in self.apply_to: + if key not in self.normalization_modes: + continue + + # Extract action subkey (e.g., 'action.joint_position' -> 'joint_position') + subkey = key.replace("action.", "") + + if subkey not in self.per_horizon_statistics: + print(f"Warning: No per-horizon statistics found for {subkey}, skipping") + continue + + statistics = self.per_horizon_statistics[subkey] + self._normalizers[key] = PerHorizonNormalizer( + mode=self.normalization_modes[key], + statistics=statistics + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + """Apply per-horizon normalization to action data. + + Args: + data: Dict with action keys mapping to tensors of shape (horizon_len, action_dim) + + Returns: + Dict with normalized action tensors + """ + for key in self.apply_to: + if key not in data: + continue + + if key not in self._input_dtypes: + input_dtype = data[key].dtype + assert isinstance( + input_dtype, torch.dtype + ), f"Unexpected input dtype: {input_dtype}. Expected type: {torch.dtype}" + self._input_dtypes[key] = input_dtype + + action = data[key] + if key in self._normalizers: + action = self._normalizers[key].forward(action) + data[key] = action + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + """Apply inverse per-horizon normalization (denormalization). + + Args: + data: Dict with normalized action tensors + + Returns: + Dict with denormalized action tensors + """ + for key in self.apply_to: + if key not in data: + continue + + action = data[key] + assert isinstance( + action, torch.Tensor + ), f"Unexpected action type: {type(action)}. Expected type: {torch.Tensor}" + + if key in self._normalizers: + action = self._normalizers[key].inverse(action) + + if key in self._input_dtypes: + original_dtype = self._input_dtypes[key] + if isinstance(original_dtype, np.dtype): + action = action.numpy().astype(original_dtype) + elif isinstance(original_dtype, torch.dtype): + action = action.to(original_dtype) + + data[key] = action + return data + + +class StateActionToTensor(InvertibleModalityTransform): + """ + Transforms states and actions to tensors. + """ + + input_dtypes: dict[str, np.dtype] = Field( + default_factory=dict, description="The input dtypes for each state key." + ) + output_dtypes: dict[str, torch.dtype] = Field( + default_factory=dict, description="The output dtypes for each state key." + ) + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = {"apply_to"} + else: + include = kwargs.pop("include", None) + + return super().model_dump(*args, include=include, **kwargs) + + @field_validator("input_dtypes", "output_dtypes", mode="before") + def validate_dtypes(cls, v): + for key, dtype in v.items(): + if isinstance(dtype, str): + if dtype.startswith("torch."): + dtype_split = dtype.split(".")[-1] + v[key] = getattr(torch, dtype_split) + elif dtype.startswith("np.") or dtype.startswith("numpy."): + dtype_split = dtype.split(".")[-1] + v[key] = np.dtype(dtype_split) + else: + raise ValueError(f"Invalid dtype: {dtype}") + return v + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + continue + value = data[key] + assert isinstance( + value, np.ndarray + ), f"Unexpected input type: {type(value)}. Expected type: {np.ndarray}" + data[key] = torch.from_numpy(value) + if key in self.output_dtypes: + data[key] = data[key].to(self.output_dtypes[key]) + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + continue + value = data[key] + assert isinstance( + value, torch.Tensor + ), f"Unexpected input type: {type(value)}. Expected type: {torch.Tensor}" + data[key] = value.numpy() + if key in self.input_dtypes: + data[key] = data[key].astype(self.input_dtypes[key]) + return data + + +class StateActionTransform(InvertibleModalityTransform): + """ + Class for state or action transform. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + normalization_modes (dict[str, str]): The normalization modes for each state key. + If a state key in apply_to is not present in the dictionary, it will not be normalized. + target_rotations (dict[str, str]): The target representations for each state key. + If a state key in apply_to is not present in the dictionary, it will not be rotated. + """ + + # Configurable attributes + apply_to: list[str] = Field(..., description="The keys in the modality to load and transform.") + normalization_modes: dict[str, str] = Field( + default_factory=dict, description="The normalization modes for each state key." + ) + target_rotations: dict[str, str] = Field( + default_factory=dict, description="The target representations for each state key." + ) + normalization_statistics: dict[str, dict] = Field( + default_factory=dict, description="The statistics for each state key." + ) + modality_metadata: dict[str, StateActionMetadata] = Field( + default_factory=dict, description="The modality metadata for each state key." + ) + + # Model variables + _rotation_transformers: dict[str, RotationTransform] = PrivateAttr(default_factory=dict) + _normalizers: dict[str, Normalizer] = PrivateAttr(default_factory=dict) + _input_dtypes: dict[str, np.dtype | torch.dtype] = PrivateAttr(default_factory=dict) + + # Model constants + _DEFAULT_MIN_MAX_STATISTICS: ClassVar[dict] = { + "rotation_6d": { + "min": [-1, -1, -1, -1, -1, -1], + "max": [1, 1, 1, 1, 1, 1], + }, + "euler_angles": { + "min": [-np.pi, -np.pi, -np.pi], + "max": [np.pi, np.pi, np.pi], + }, + "quaternion": { + "min": [-1, -1, -1, -1], + "max": [1, 1, 1, 1], + }, + "axis_angle": { + "min": [-np.pi, -np.pi, -np.pi], + "max": [np.pi, np.pi, np.pi], + }, + } + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = {"apply_to", "normalization_modes", "target_rotations"} + else: + include = kwargs.pop("include", None) + + return super().model_dump(*args, include=include, **kwargs) + + @field_validator("modality_metadata", mode="before") + def validate_modality_metadata(cls, v): + for modality_key, config in v.items(): + if isinstance(config, dict): + config = StateActionMetadata.model_validate(config) + else: + assert isinstance( + config, StateActionMetadata + ), f"Invalid source rotation config: {config}" + v[modality_key] = config + return v + + @model_validator(mode="after") + def validate_normalization_statistics(self): + for modality_key, normalization_statistics in self.normalization_statistics.items(): + if modality_key in self.normalization_modes: + normalization_mode = self.normalization_modes[modality_key] + if normalization_mode == "min_max": + assert ( + "min" in normalization_statistics and "max" in normalization_statistics + ), f"Min and max statistics are required for min_max normalization, but got {normalization_statistics}" + assert len(normalization_statistics["min"]) == len( + normalization_statistics["max"] + ), f"Min and max statistics must have the same length, but got {normalization_statistics['min']} and {normalization_statistics['max']}" + elif normalization_mode == "mean_std": + assert ( + "mean" in normalization_statistics and "std" in normalization_statistics + ), f"Mean and std statistics are required for mean_std normalization, but got {normalization_statistics}" + assert len(normalization_statistics["mean"]) == len( + normalization_statistics["std"] + ), f"Mean and std statistics must have the same length, but got {normalization_statistics['mean']} and {normalization_statistics['std']}" + elif normalization_mode == "q99": + assert ( + "q01" in normalization_statistics and "q99" in normalization_statistics + ), f"q01 and q99 statistics are required for q99 normalization, but got {normalization_statistics}" + assert len(normalization_statistics["q01"]) == len( + normalization_statistics["q99"] + ), f"q01 and q99 statistics must have the same length, but got {normalization_statistics['q01']} and {normalization_statistics['q99']}" + elif normalization_mode == "binary": + assert ( + len(normalization_statistics) == 1 + ), f"Binary normalization should only have one value, but got {normalization_statistics}" + assert normalization_statistics[0] in [ + 0, + 1, + ], f"Binary normalization should only have 0 or 1, but got {normalization_statistics[0]}" + else: + raise ValueError(f"Invalid normalization mode: {normalization_mode}") + return self + + def set_metadata(self, dataset_metadata: DatasetMetadata): + dataset_statistics = dataset_metadata.statistics + modality_metadata = dataset_metadata.modalities + + # Check that all state keys specified in apply_to have their modality_metadata + for key in self.apply_to: + split_key = key.split(".") + assert len(split_key) == 2, "State keys should have two parts: 'modality.key'" + if key not in self.modality_metadata: + modality, state_key = split_key + assert hasattr(modality_metadata, modality), f"{modality} config not found" + assert state_key in getattr( + modality_metadata, modality + ), f"{state_key} config not found in {modality}" + self.modality_metadata[key] = getattr(modality_metadata, modality)[state_key] + + # Check that all state keys specified in normalization_modes have their statistics in state_statistics + for key in self.normalization_modes: + split_key = key.split(".") + assert len(split_key) == 2, "State keys should have two parts: 'modality.key'" + modality, state_key = split_key + assert hasattr(dataset_statistics, modality), f"{modality} statistics not found" + assert state_key in getattr( + dataset_statistics, modality + ), f"{state_key} statistics not found" + assert ( + len(getattr(modality_metadata, modality)[state_key].shape) == 1 + ), f"{getattr(modality_metadata, modality)[state_key].shape=}" + self.normalization_statistics[key] = getattr(dataset_statistics, modality)[ + state_key + ].model_dump() + + # Initialize the rotation transformers + for key in self.target_rotations: + # Get the original representation of the state + from_rep = self.modality_metadata[key].rotation_type + assert from_rep is not None, f"Source rotation type not found for {key}" + + # Get the target representation of the state, will raise an error if the target representation is not valid + to_rep = RotationType(self.target_rotations[key]) + + # If the original representation is not the same as the target representation, initialize the rotation transformer + if from_rep != to_rep: + self._rotation_transformers[key] = RotationTransform( + from_rep=from_rep.value, to_rep=to_rep.value + ) + + # Initialize the normalizers + for key in self.normalization_modes: + modality, state_key = key.split(".") + # If the state has a nontrivial rotation, we need to handle it more carefully + # For absolute rotations, we need to convert them to the target representation and normalize them using min_max mode, + # since we can infer the bounds by the representation + # For relative rotations, we cannot normalize them as we don't know the bounds + if key in self._rotation_transformers: + # Case 1: Absolute rotation + if self.modality_metadata[key].absolute: + # Check that the normalization mode is valid + assert ( + self.normalization_modes[key] == "min_max" + ), "Absolute rotations that are converted to other formats must be normalized using `min_max` mode" + rotation_type = RotationType(self.target_rotations[key]).value + # If the target representation is euler angles, we need to parse the convention + if rotation_type.startswith("euler_angles"): + rotation_type = "euler_angles" + # Get the statistics for the target representation + statistics = self._DEFAULT_MIN_MAX_STATISTICS[rotation_type] + # Case 2: Relative rotation + else: + raise ValueError( + f"Cannot normalize relative rotations: {key} that's converted to {self.target_rotations[key]}" + ) + # If the state is not continuous, we should not use normalization modes other than binary + elif ( + not self.modality_metadata[key].continuous + and self.normalization_modes[key] != "binary" + ): + raise ValueError( + f"{key} is not continuous, so it should be normalized using `binary` mode" + ) + # Initialize the normalizer + else: + statistics = self.normalization_statistics[key] + self._normalizers[key] = Normalizer( + mode=self.normalization_modes[key], statistics=statistics + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + # We allow some keys to be missing in the data, and only process the keys that are present + continue + if key not in self._input_dtypes: + input_dtype = data[key].dtype + assert isinstance( + input_dtype, torch.dtype + ), f"Unexpected input dtype: {input_dtype}. Expected type: {torch.dtype}" + self._input_dtypes[key] = input_dtype + else: + assert ( + data[key].dtype == self._input_dtypes[key] + ), f"All states corresponding to the same key must be of the same dtype, input dtype: {data[key].dtype}, expected dtype: {self._input_dtypes[key]}" + # Rotate the state + state = data[key] + if key in self._rotation_transformers: + state = self._rotation_transformers[key].forward(state) + # Normalize the state + if key in self._normalizers: + state = self._normalizers[key].forward(state) + data[key] = state + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + continue + state = data[key] + assert isinstance( + state, torch.Tensor + ), f"Unexpected state type: {type(state)}. Expected type: {torch.Tensor}" + # Unnormalize the state + if key in self._normalizers: + state = self._normalizers[key].inverse(state) + # Change the state back to its original representation + if key in self._rotation_transformers: + state = self._rotation_transformers[key].inverse(state) + assert isinstance( + state, torch.Tensor + ), f"State should be tensor after unapplying transformations, but got {type(state)}" + # Only convert back to the original dtype if it's known, i.e. `apply` was called before + # If not, we don't know the original dtype, so we don't convert + if key in self._input_dtypes: + original_dtype = self._input_dtypes[key] + if isinstance(original_dtype, np.dtype): + state = state.numpy().astype(original_dtype) + elif isinstance(original_dtype, torch.dtype): + state = state.to(original_dtype) + else: + raise ValueError(f"Invalid input dtype: {original_dtype}") + data[key] = state + return data + + +class StateActionPerturbation(ModalityTransform): + """ + Class for state or action perturbation. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + std (float): Standard deviation of the noise to be added to the state or action. + """ + + # Configurable attributes + std: float = Field( + ..., description="Standard deviation of the noise to be added to the state or action." + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if not self.training: + # Don't perturb the data in eval mode + return data + if self.std < 0: + # If the std is negative, we don't add any noise + return data + for key in self.apply_to: + state = data[key] + assert isinstance(state, torch.Tensor) + transformed_data_min = torch.min(state) + transformed_data_max = torch.max(state) + noise = torch.randn_like(state) * self.std + state += noise + # Clip to the original range + state = torch.clamp(state, transformed_data_min, transformed_data_max) + data[key] = state + return data + + +class StateActionDropout(ModalityTransform): + """ + Class for state or action dropout. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + dropout_prob (float): Probability of dropping out a state or action. + """ + + # Configurable attributes + dropout_prob: float = Field(..., description="Probability of dropping out a state or action.") + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if not self.training: + # Don't drop out the data in eval mode + return data + if self.dropout_prob < 0: + # If the dropout probability is negative, we don't drop out any states + return data + if self.dropout_prob > 1e-9 and random.random() < self.dropout_prob: + for key in self.apply_to: + state = data[key] + assert isinstance(state, torch.Tensor) + state = torch.zeros_like(state) + data[key] = state + return data + + +class StateActionSinCosTransform(ModalityTransform): + """ + Class for state or action sin-cos transform. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + """ + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + state = data[key] + assert isinstance(state, torch.Tensor) + sin_state = torch.sin(state) + cos_state = torch.cos(state) + data[key] = torch.cat([sin_state, cos_state], dim=-1) + return data diff --git a/groot/vla/data/transform/video.py b/groot/vla/data/transform/video.py new file mode 100644 index 0000000000000000000000000000000000000000..e5a76df21d1d9b14d25f68de4a63da06cc98869a --- /dev/null +++ b/groot/vla/data/transform/video.py @@ -0,0 +1,1026 @@ +from typing import Any, Callable, ClassVar, Optional, Literal + +import albumentations as A +import cv2 +from einops import rearrange +import functools +import numpy as np +from pydantic import Field, PrivateAttr, field_validator +import torch +import torchvision.transforms.v2 as T + +from groot.vla.data.schema import DatasetMetadata +from groot.vla.data.transform.base import ModalityTransform + + +class VideoTransform(ModalityTransform): + # Configurable attributes + backend: str = Field( + default="torchvision", description="The backend to use for the transformations" + ) + + # Model variables + _train_transform: Callable | None = PrivateAttr(default=None) + _eval_transform: Callable | None = PrivateAttr(default=None) + _original_resolutions: dict[str, tuple[int, int]] = PrivateAttr(default_factory=dict) + + # Model constants + _INTERPOLATION_MAP: ClassVar[dict[str, dict[str, Any]]] = PrivateAttr( + { + "nearest": { + "albumentations": cv2.INTER_NEAREST, + "torchvision": T.InterpolationMode.NEAREST, + }, + "linear": { + "albumentations": cv2.INTER_LINEAR, + "torchvision": T.InterpolationMode.BILINEAR, + }, + "cubic": { + "albumentations": cv2.INTER_CUBIC, + "torchvision": T.InterpolationMode.BICUBIC, + }, + "area": { + "albumentations": cv2.INTER_AREA, + "torchvision": None, # Torchvision does not support this interpolation mode + }, + "lanczos4": { + "albumentations": cv2.INTER_LANCZOS4, # Lanczos with a 4x4 filter + "torchvision": T.InterpolationMode.LANCZOS, # Torchvision does not specify filter size, might be different from 4x4 + }, + "linear_exact": { + "albumentations": cv2.INTER_LINEAR_EXACT, + "torchvision": None, # Torchvision does not support this interpolation mode + }, + "nearest_exact": { + "albumentations": cv2.INTER_NEAREST_EXACT, + "torchvision": T.InterpolationMode.NEAREST_EXACT, + }, + "max": { + "albumentations": cv2.INTER_MAX, + "torchvision": None, + }, + } + ) + + @property + def train_transform(self) -> Callable: + assert ( + self._train_transform is not None + ), "Transform is not set. Please call set_metadata() before calling apply()." + return self._train_transform + + @train_transform.setter + def train_transform(self, value: Callable): + self._train_transform = value + + @property + def eval_transform(self) -> Callable | None: + return self._eval_transform + + @eval_transform.setter + def eval_transform(self, value: Callable | None): + self._eval_transform = value + + @property + def original_resolutions(self) -> dict[str, tuple[int, int]]: + assert ( + self._original_resolutions is not None + ), "Original resolutions are not set. Please call set_metadata() before calling apply()." + return self._original_resolutions + + @original_resolutions.setter + def original_resolutions(self, value: dict[str, tuple[int, int]]): + self._original_resolutions = value + + def check_input(self, data: dict[str, Any]): + if self.backend == "torchvision": + for key in self.apply_to: + assert isinstance(data[key], torch.Tensor), f"Video {key} is not a torch tensor" + assert data[key].ndim in [ + 4, + 5, + ], f"Expected video {key} to have 4 or 5 dimensions (T, C, H, W or T, B, C, H, W), got {data[key].ndim}" + elif self.backend == "albumentations": + for key in self.apply_to: + assert isinstance(data[key], np.ndarray), f"Video {key} is not a numpy array" + assert data[key].ndim in [ + 4, + 5, + ], f"Expected video {key} to have 4 or 5 dimensions (T, C, H, W or T, B, C, H, W), got {data[key].ndim}" + else: + raise ValueError(f"Backend {self.backend} not supported") + + def set_metadata(self, dataset_metadata: DatasetMetadata): + super().set_metadata(dataset_metadata) + self.original_resolutions = {} + for key in self.apply_to: + split_keys = key.split(".") + assert len(split_keys) == 2, f"Invalid key: {key}. Expected format: modality.key" + sub_key = split_keys[1] + if sub_key in dataset_metadata.modalities.video: + self.original_resolutions[key] = dataset_metadata.modalities.video[ + sub_key + ].resolution + else: + raise ValueError( + f"Video key {sub_key} not found in dataset metadata. Available keys: {dataset_metadata.modalities.video.keys()}" + ) + train_transform = self.get_transform(mode="train") + eval_transform = self.get_transform(mode="eval") + if self.backend == "albumentations": + self.train_transform = A.ReplayCompose(transforms=[train_transform]) # type: ignore + if eval_transform is not None: + self.eval_transform = A.ReplayCompose(transforms=[eval_transform]) # type: ignore + else: + assert train_transform is not None, "Train transform must be set" + self.train_transform = train_transform + self.eval_transform = eval_transform + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if self.training: + transform = self.train_transform + else: + transform = self.eval_transform + if transform is None: + return data + assert ( + transform is not None + ), "Transform is not set. Please call set_metadata() before calling apply()." + try: + self.check_input(data) + except AssertionError as e: + raise ValueError( + f"Input data does not match the expected format for {self.__class__.__name__}: {e}" + ) from e + + # Concatenate views + views = [data[key] for key in self.apply_to] + num_views = len(views) + is_batched = views[0].ndim == 5 + bs = views[0].shape[0] if is_batched else 1 + if isinstance(views[0], torch.Tensor): + views = torch.cat(views, 0) + elif isinstance(views[0], np.ndarray): + views = np.concatenate(views, 0) + else: + raise ValueError(f"Unsupported view type: {type(views[0])}") + if is_batched: + views = rearrange(views, "(v b) t c h w -> (v b t) c h w", v=num_views, b=bs) + # Apply the transform + if self.backend == "torchvision": + views = transform(views) + elif self.backend == "albumentations": + assert isinstance(transform, A.ReplayCompose), "Transform must be a ReplayCompose" + first_frame = views[0] + transformed = transform(image=first_frame) + replay_data = transformed["replay"] + transformed_first_frame = transformed["image"] + + if len(views) > 1: + # Apply the same transformations to the rest of the frames + transformed_frames = [ + transform.replay(replay_data, image=frame)["image"] for frame in views[1:] + ] + # Add the first frame back + transformed_frames = [transformed_first_frame] + transformed_frames + else: + # If there is only one frame, just make a list with one frame + transformed_frames = [transformed_first_frame] + + # Delete the replay data to save memory + del replay_data + views = np.stack(transformed_frames, 0) + + else: + raise ValueError(f"Backend {self.backend} not supported") + # Split views + if is_batched: + views = rearrange(views, "(v b t) c h w -> v b t c h w", v=num_views, b=bs) + else: + views = rearrange(views, "(v t) c h w -> v t c h w", v=num_views) + for key, view in zip(self.apply_to, views): + data[key] = view + return data + + @classmethod + def _validate_interpolation(cls, interpolation: str): + if interpolation not in cls._INTERPOLATION_MAP: + raise ValueError(f"Interpolation mode {interpolation} not supported") + + def _get_interpolation(self, interpolation: str, backend: str = "torchvision"): + """ + Get the interpolation mode for the given backend. + + Args: + interpolation (str): The interpolation mode. + backend (str): The backend to use. + + Returns: + Any: The interpolation mode for the given backend. + """ + return self._INTERPOLATION_MAP[interpolation][backend] + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + raise NotImplementedError( + "set_transform is not implemented for VideoTransform. Please implement this function to set the transforms." + ) + + +class VideoCrop(VideoTransform): + height: int | None = Field(default=None, description="The height of the input image") + width: int | None = Field(default=None, description="The width of the input image") + scale: float = Field( + ..., + description="The scale of the crop. The crop size is (width * scale, height * scale)", + ) + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the transform for the given mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: If mode is "train", return a random crop transform. If mode is "eval", return a center crop transform. + """ + # 1. Check the input resolution + assert ( + len(set(self.original_resolutions.values())) == 1 + ), f"All video keys must have the same resolution, got: {self.original_resolutions}" + if self.height is None: + assert self.width is None, "Height and width must be either both provided or both None" + self.width, self.height = self.original_resolutions[self.apply_to[0]] + else: + assert ( + self.width is not None + ), "Height and width must be either both provided or both None" + # 2. Create the transform + size = (int(self.height * self.scale), int(self.width * self.scale)) + if self.backend == "torchvision": + if mode == "train": + return T.RandomCrop(size) + elif mode == "eval": + return T.CenterCrop(size) + else: + raise ValueError(f"Crop mode {mode} not supported") + elif self.backend == "albumentations": + if mode == "train": + return A.RandomCrop(height=size[0], width=size[1], p=1) + elif mode == "eval": + return A.CenterCrop(height=size[0], width=size[1], p=1) + else: + raise ValueError(f"Crop mode {mode} not supported") + else: + raise ValueError(f"Backend {self.backend} not supported") + + def check_input(self, data: dict[str, Any]): + super().check_input(data) + # Check the input resolution + for key in self.apply_to: + if self.backend == "torchvision": + height, width = data[key].shape[-2:] + elif self.backend == "albumentations": + height, width = data[key].shape[-3:-1] + else: + raise ValueError(f"Backend {self.backend} not supported") + assert ( + height == self.height and width == self.width + ), f"Video {key} has invalid shape {height, width}, expected {self.height, self.width}" + + +class VideoRandomErasing(VideoTransform): + """Adds random rectangles overlaying the video. + + This discourages overfitting to the background. + """ + + probability: float = Field(default=0.2, description="Probability of applying the transform") + scale: tuple[float, float] = Field(default=(0.02, 0.33), description="Scale of the rectangle") + ratio: tuple[float, float] = Field( + default=(0.3, 3.3), description="Aspect ratio of the rectangle" + ) + value: Literal["random"] | tuple[float, float, float] = Field( + default="random", description="Color to fill the erased region with" + ) + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the transform for the given mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: If mode is "train", return a transform that adds random rectangles. If mode is "eval", return a no-op. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomErasing( + p=self.probability, scale=self.scale, ratio=self.ratio, value=self.value + ) + elif self.backend == "albumentations": + return A.Erasing( + p=self.probability, scale=self.scale, ratio=self.ratio, value=self.value + ) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoResize(VideoTransform): + height: int = Field(..., description="The height of the resize") + width: int = Field(..., description="The width of the resize") + interpolation: str = Field(default="linear", description="The interpolation mode") + antialias: bool = Field(default=True, description="Whether to apply antialiasing") + + @field_validator("interpolation") + def validate_interpolation(cls, v): + cls._validate_interpolation(v) + return v + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the resize transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The resize transform. + """ + interpolation = self._get_interpolation(self.interpolation, self.backend) + if interpolation is None: + raise ValueError( + f"Interpolation mode {self.interpolation} not supported for torchvision" + ) + if self.backend == "torchvision": + size = (self.height, self.width) + return T.Resize(size, interpolation=interpolation, antialias=self.antialias) + elif self.backend == "albumentations": + return A.Resize( + height=self.height, + width=self.width, + interpolation=interpolation, + p=1, + ) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoRandomRotation(VideoTransform): + degrees: float | tuple[float, float] = Field( + ..., description="The degrees of the random rotation" + ) + interpolation: str = Field("linear", description="The interpolation mode") + + @field_validator("interpolation") + def validate_interpolation(cls, v): + cls._validate_interpolation(v) + return v + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the random rotation transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: The random rotation transform. None for eval mode. + """ + if mode == "eval": + return None + interpolation = self._get_interpolation(self.interpolation, self.backend) + if interpolation is None: + raise ValueError( + f"Interpolation mode {self.interpolation} not supported for torchvision" + ) + if self.backend == "torchvision": + return T.RandomRotation(self.degrees, interpolation=interpolation) # type: ignore + elif self.backend == "albumentations": + return A.Rotate(limit=self.degrees, interpolation=interpolation, p=1) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoHorizontalFlip(VideoTransform): + p: float = Field(..., description="The probability of the horizontal flip") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the horizontal flip transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a horizontal flip transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomHorizontalFlip(self.p) + elif self.backend == "albumentations": + return A.HorizontalFlip(p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoGrayscale(VideoTransform): + p: float = Field(..., description="The probability of the grayscale transformation") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the grayscale transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a grayscale transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomGrayscale(self.p) + elif self.backend == "albumentations": + return A.ToGray(p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoColorJitter(VideoTransform): + brightness: float | tuple[float, float] = Field( + ..., description="The brightness of the color jitter" + ) + contrast: float | tuple[float, float] = Field( + ..., description="The contrast of the color jitter" + ) + saturation: float | tuple[float, float] = Field( + ..., description="The saturation of the color jitter" + ) + hue: float | tuple[float, float] = Field(..., description="The hue of the color jitter") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the color jitter transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a color jitter transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.ColorJitter( + brightness=self.brightness, + contrast=self.contrast, + saturation=self.saturation, + hue=self.hue, + ) + elif self.backend == "albumentations": + return A.ColorJitter( + brightness=self.brightness, + contrast=self.contrast, + saturation=self.saturation, + hue=self.hue, + p=1, + ) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoRandomGrayscale(VideoTransform): + p: float = Field(..., description="The probability of the grayscale transformation") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the grayscale transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a grayscale transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomGrayscale(self.p) + elif self.backend == "albumentations": + return A.ToGray(p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoRandomPosterize(VideoTransform): + bits: int = Field(..., description="The number of bits to posterize the image") + p: float = Field(..., description="The probability of the posterize transformation") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the posterize transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a posterize transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomPosterize(bits=self.bits, p=self.p) + elif self.backend == "albumentations": + return A.Posterize(num_bits=self.bits, p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoToTensor(VideoTransform): + + output_on_cuda: bool = Field( + default=False, + description="Output the tensor on CUDA if True.", + ) + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the to tensor transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The to tensor transform. + """ + if self.backend == "torchvision": + return functools.partial( + self.__class__.to_tensor, + output_on_cuda=self.output_on_cuda, + ) + else: + raise ValueError(f"Backend {self.backend} not supported") + + def check_input(self, data: dict): + """Check if the input data has the correct shape. + Expected video shape: [T, H, W, C], dtype np.uint8 + """ + for key in self.apply_to: + assert key in data, f"Key {key} not found in data. Available keys: {data.keys()}" + assert data[key].ndim in [ + 4, + 5, + ], f"Video {key} must have 4 or 5 dimensions, got {data[key].ndim}" + assert ( + data[key].dtype == np.uint8 + ), f"Video {key} must have dtype uint8, got {data[key].dtype}" + input_resolution = data[key].shape[-3:-1][::-1] + if key in self.original_resolutions: + expected_resolution = self.original_resolutions[key] + else: + expected_resolution = input_resolution + assert ( + input_resolution == expected_resolution + ), f"Video {key} has invalid resolution {input_resolution}, expected {expected_resolution}. Full shape: {data[key].shape}" + + @staticmethod + def to_tensor(frames: np.ndarray, output_on_cuda: bool) -> torch.Tensor: + """Convert numpy array to tensor efficiently. + + Args: + frames: numpy array of shape [T, H, W, C] in uint8 format + output_on_cuda: whether to output the tensor on CUDA if True. + Returns: + tensor of shape [T, C, H, W] in range [0, 1] + """ + frames = torch.from_numpy(frames) + if output_on_cuda: + frames = frames.cuda() + frames = frames.to(torch.float32) / 255.0 + return frames.permute(0, 3, 1, 2) # [T, C, H, W] + + +class VideoToNumpy(VideoTransform): + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the to numpy transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The to numpy transform. + """ + if self.backend == "torchvision": + return self.__class__.to_numpy + else: + raise ValueError(f"Backend {self.backend} not supported") + + @staticmethod + def to_numpy(frames: torch.Tensor) -> np.ndarray: + """Convert tensor back to numpy array efficiently. + + Args: + frames: tensor of shape [T, C, H, W] in range [0, 1] + Returns: + numpy array of shape [T, H, W, C] in uint8 format + """ + frames = (frames.permute(0, 2, 3, 1) * 255).to(torch.uint8) + return frames.cpu().numpy() + + +class VideoMergeTimeBatch(ModalityTransform): + """ + Merge the batch and time dimensions of the video. + """ + + apply_to: list[str] = Field(..., description="The keys of the modalities to merge") + + def apply(self, data: dict) -> dict: + warnings.warn( + "VideoMergeTimeBatch is deprecated. Use ComposedModalityTransform instead.", + DeprecationWarning, + ) + for key in self.apply_to: + data[key] = rearrange(data[key], "b t ... -> (b t) ...") + return data + + +class VideoSplitTimeBatch(ModalityTransform): + """ + Split the batch and time dimensions of the video. + """ + + apply_to: list[str] = Field(..., description="The keys of the modalities to split") + time_dim: int = Field(..., description="The dimension of the time dimension") + + def apply(self, data: dict) -> dict: + warnings.warn( + "VideoSplitTimeBatch is deprecated. Use ComposedModalityTransform instead.", + DeprecationWarning, + ) + for key in self.apply_to: + data[key] = rearrange(data[key], "(b t) ... -> b t ...", t=self.time_dim) + return data + + +class VideoFocusRect(ModalityTransform): + """ + Given a rectangle area in the video, apply focus effects on the target + rectangle, by applying blur and noise to the surrounding region. + + Mainly useful for EgoView + """ + + # Region coordinates in normalized space [0,1] + xtl: float = Field(2 / 12, description="Top-left x coordinate (normalized)", ge=0.0, le=1.0) + ytl: float = Field(3 / 8, description="Top-left y coordinate (normalized)", ge=0.0, le=1.0) + xbr: float = Field(10 / 12, description="Bottom-left x coordinate (normalized)", ge=0.0, le=1.0) + ybr: float = Field(1.0, description="Bottom-left y coordinate (normalized)", ge=0.0, le=1.0) + + # Content region parameters (in pixel coordinates, None means auto-detect) + content_y_min: Optional[int] = Field( + None, description="Top coordinate of content region (pixels)" + ) + content_y_max: Optional[int] = Field( + None, description="Bottom coordinate of content region (pixels)" + ) + content_x_min: Optional[int] = Field( + None, description="Left coordinate of content region (pixels)" + ) + content_x_max: Optional[int] = Field( + None, description="Right coordinate of content region (pixels)" + ) + + # Jitter amount for coordinates (in normalized space) + jitter: float = Field(0.05, description="Amount of random jitter to apply to coordinates") + + # Effect parameters + blur_kernel: int = Field(95, description="Gaussian blur kernel size") + noise_std: float = Field(0.3, description="Standard deviation of Gaussian noise") + blend_size: float = Field(0.1, description="Size of blending region as fraction of image size") + + # Effect probabilities during training + p_blur: float = Field(0.2, description="Probability of applying blur") + p_noise: float = Field(0.2, description="Probability of applying noise") + + def detect_padding(self, image: np.ndarray) -> tuple[slice, slice]: + """ + Detect padding in the image by finding non-black regions. + Returns slices for the content region (y_slice, x_slice). + """ + H, W = image.shape[:2] + + # If all content region parameters are provided, use them + if all( + param is not None + for param in [ + self.content_y_min, + self.content_y_max, + self.content_x_min, + self.content_x_max, + ] + ): + y_min = max(0, self.content_y_min) + y_max = min(H, self.content_y_max) + x_min = max(0, self.content_x_min) + x_max = min(W, self.content_x_max) + return slice(y_min, y_max), slice(x_min, x_max) + + if image.ndim == 3: + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + else: + gray = image + + # Find non-black regions + rows = np.any(gray > 0.01, axis=1) + cols = np.any(gray > 0.01, axis=0) + + # find first and last non-black pixel indices + rmin, rmax = np.where(rows)[0][[0, -1]] + cmin, cmax = np.where(cols)[0][[0, -1]] + + return slice(rmin, rmax + 1), slice(cmin, cmax + 1) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if not self.training: + # Do nothing in eval mode + return data + + for key in self.apply_to: + video = data[key] + + # Handle numpy array case + assert isinstance( + video, np.ndarray + ), f"Expected numpy array or torch tensor for {key}, got {type(video)}" + assert video.ndim in { + 3, + 4, + }, f"Expected [H, W, C] or [T, H, W, C] array for {key}, got shape {video.shape}" + transformed = self._transform_video(video) + + data[key] = transformed + + return data + + def _transform_video(self, video: np.ndarray) -> np.ndarray: + """ + Apply the focus rectangle transformation to a video. + + Args: + video (np.ndarray): Video tensor of shape [T, H, W, C] + + Returns: + np.ndarray: Transformed video + """ + # Handle both single frame and video inputs + is_single_frame = video.ndim == 3 + if is_single_frame: + video = video[np.newaxis] + + T, H, W, C = video.shape + + assert ( + self.p_blur + self.p_noise <= 1.0 + ), "Sum of blur and noise probabilities must be <= 1.0" + r = np.random.random() + apply_blur = r < self.p_blur + apply_noise = self.p_blur <= r < self.p_blur + self.p_noise + alpha = random.uniform(0.0, 1.0) # Noise blending factor + + # Apply jitter once to rectangle + xtl = self.xtl + np.random.uniform(-self.jitter, self.jitter) + ytl = self.ytl + np.random.uniform(-self.jitter, self.jitter) + xbr = self.xbr + np.random.uniform(-self.jitter, self.jitter) + ybr = self.ybr + np.random.uniform(-self.jitter, self.jitter) + xtl, ytl, xbr, ybr = [np.clip(x, 0.0, 1.0) for x in [xtl, ytl, xbr, ybr]] + + # Detect padding from first frame (assume consistent across frames) + y_slice, x_slice = self.detect_padding(video[0]) + content_h = y_slice.stop - y_slice.start + content_w = x_slice.stop - x_slice.start + + # Convert normalized coordinates relative to pixel space + x1 = int(xtl * content_w) + x_slice.start + y1 = int(ytl * content_h) + y_slice.start + x2 = int(xbr * content_w) + x_slice.start + y2 = int(ybr * content_h) + y_slice.start + + # Create mask for the inner rectangle + mask = np.zeros((H, W), dtype=np.float32) + pts = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]], dtype=np.int32) + cv2.fillPoly(mask, [pts], color=1.0) + + # Create a smooth blend mask around the target rectangle using distance transform + content_mask = np.zeros((H, W), dtype=np.uint8) + content_mask[y_slice, x_slice] = 1 + dist = cv2.distanceTransform(1 - (mask > 0).astype(np.uint8) * content_mask, cv2.DIST_L2, 3) + blend_radius = int(min(content_h, content_w) * self.blend_size) + blend_mask = np.clip(1.0 - dist / blend_radius, 0, 1) + blend_mask *= content_mask + blend_mask = blend_mask[..., np.newaxis] + + # Process all frames with same transformations + result = np.zeros_like(video) + for t in range(T): + frame = video[t] + modified = frame.copy() + content = modified[y_slice, x_slice] + + if apply_blur: + content = cv2.GaussianBlur(content, (self.blur_kernel, self.blur_kernel), 0) + + if apply_noise: + background = np.random.randint(0, 256, content.shape, dtype=np.uint8) / 255.0 + content = alpha * content + (1 - alpha) * background + content = np.clip(content, 0, 1) + + modified[y_slice, x_slice] = content + result[t] = frame * blend_mask + modified * (1 - blend_mask) + + return result[0] if is_single_frame else result + + +class VideoNormalize(VideoTransform): + mean: list[float] = Field(..., description="Mean for normalization") + std: list[float] = Field(..., description="Standard deviation for normalization") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the normalization transform. Same for train and eval mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The normalization transform. + """ + print("Using VideoNormalize transform") + if self.backend == "torchvision": + return T.Normalize(mean=self.mean, std=self.std) + elif self.backend == "albumentations": + return A.Normalize(mean=self.mean, std=self.std, max_pixel_value=1.0, p=1.0) + else: + raise ValueError(f"Backend {self.backend} not supported") + + def check_input(self, data: dict): + for key in self.apply_to: + assert key in data, f"Key {key} not found in data" + assert isinstance(data[key], torch.Tensor), f"Video {key} is not a torch tensor" + assert data[key].ndim in [4, 5], f"Video {key} must have 4 or 5 dimensions, got {data[key].ndim}" + assert data[key].dtype == torch.float32, f"Video {key} must be float32, got {data[key].dtype}" + assert data[key].min() >= 0.0 and data[key].max() <= 1.0, ( + f"Video {key} must be in [0,1] range before normalization" + ) + +# class VideoTransformLegacy(ModalityTransform): +# def __init__( +# self, +# modality_keys: list[str], +# backend: str = "torchvision", +# crop_cfg: CropConfig | None = None, +# resize_cfg: ResizeConfig | None = None, +# random_rotation_cfg: RandomRotationConfig | None = None, +# horizontal_flip_cfg: HorizontalFlipConfig | None = None, +# grayscale_cfg: GrayscaleConfig | None = None, +# color_jitter_cfg: ColorJitterConfig | None = None, +# strong_vision_aug: bool = False, +# ): +# """ +# Initialize the video transform. +# With the default settings, the input will be (T, H, W, C) where T is the number of frames. +# The output will be (K, T, C, H, W) where K is the number of video keys. + +# Args: +# modality_keys (list[str]): The keys of the modalities to load and transform. +# backend (str): The backend to use for the transformations. The default is "torchvision". +# crop_cfg (CropConfig | None): Configuration for the crop transformation. See CropConfig for more details. +# resize_cfg (ResizeConfig | None): Configuration for the resize transformation. See ResizeConfig for more details. +# random_rotation_cfg (RandomRotationConfig | None): Configuration for the random rotation transformation. See RandomRotationConfig for more details. +# horizontal_flip_cfg (HorizontalFlipConfig | None): Configuration for the horizontal flip transformation. See HorizontalFlipConfig for more details. +# grayscale_cfg (GrayscaleConfig | None): Configuration for the grayscale transformation. See GrayscaleConfig for more details. +# color_jitter_cfg (ColorJitterConfig | None): Configuration for the color jitter transformation. See ColorJitterConfig for more details. +# strong_vision_aug (bool): Whether to apply strong vision augmentation. The default is False. +# """ +# super().__init__(modality_keys) +# self.backend = backend +# self.crop_cfg = crop_cfg +# self.resize_cfg = resize_cfg +# self.random_rotation_cfg = random_rotation_cfg +# self.horizontal_flip_cfg = horizontal_flip_cfg +# self.grayscale_cfg = grayscale_cfg +# self.color_jitter_cfg = color_jitter_cfg +# self.strong_vision_aug = strong_vision_aug +# self.transforms = None + +# def set_metadata(self, dataset_metadata: TrainableDatasetMetadata_V1_1): +# super().set_metadata(dataset_metadata) +# # Get the original height and width +# video_metadata = dataset_metadata.modalities.video +# resolutions = {} +# for key in self.modality_keys: +# split_keys = key.split(".") +# assert len(split_keys) == 2, f"Invalid key: {key}. Expected format: modality.key" +# sub_key = split_keys[1] +# resolutions[key] = video_metadata[sub_key].resolution +# assert ( +# len(set(resolutions.values())) == 1 +# ), f"All video keys must have the same resolution, got: {resolutions}" +# width, height = resolutions[self.modality_keys[0]] + +# transforms = [] +# if self.crop_cfg is not None: +# self.crop_cfg.set_original_height_width(height, width) +# transforms.append(self.crop_cfg.get_transform(self.backend)) +# if self.resize_cfg is not None: +# transforms.append(self.resize_cfg.get_transform(self.backend)) +# if self.random_rotation_cfg is not None: +# transforms.append(self.random_rotation_cfg.get_transform(self.backend)) +# if self.horizontal_flip_cfg is not None: +# transforms.append(self.horizontal_flip_cfg.get_transform(self.backend)) +# if self.grayscale_cfg is not None: +# transforms.append(self.grayscale_cfg.get_transform(self.backend)) +# if self.color_jitter_cfg is not None: +# transforms.append(self.color_jitter_cfg.get_transform(self.backend)) + +# if self.backend == "torchvision": +# if len(transforms) == 0: +# transforms.append(T.Identity()) +# self.transforms = T.Compose(transforms) +# else: +# raise ValueError(f"Backend {self.backend} not supported") + +# if self.strong_vision_aug: +# import kornia.augmentation as K +# from kornia.augmentation import ImageSequential + +# assert ( +# self.backend == "torchvision" +# ), "Temporarily only support torchvision backend for strong augmentation" +# self.strong_transform = ImageSequential( +# K.RandomErasing(p=0.5, scale=(0.005, 0.01), ratio=(0.3, 1.3)), +# K.RandomSaltAndPepperNoise(p=0.5, amount=0.05, salt_vs_pepper=0.5), +# K.RandomCutMixV2(p=0.5, num_mix=1, cut_size=(0.98, 1.0)), +# random_apply=1, +# keepdim=True, +# same_on_batch=True, +# ) + +# def __call__(self, data: dict) -> dict[str, torch.Tensor | np.ndarray | Image.Image]: +# # Batch frames along the first dimension +# frames = [data[key] for key in self.modality_keys] # view x [T, H, W, C] +# n_view, n_frames = len(frames), len(frames[0]) +# frames = np.concatenate(frames, 0) # [view*T, H, W, C] + +# if self.backend == "torchvision": +# transformed_frames = self.transform_torchvision(frames) +# else: +# raise ValueError(f"Backend {self.backend} not supported") + +# # De-batch the frames +# transformed_frames = np.array(transformed_frames) # [view*T, H, W, C] +# H, W, C = transformed_frames.shape[-3:] +# transformed_frames = { +# key: x +# for key, x in zip( +# self.modality_keys, transformed_frames.reshape(n_view, n_frames, H, W, C) +# ) +# } + +# return transformed_frames + +# def check_input(self, data: dict): +# for key in self.modality_keys: +# assert key in data, f"Key {key} not found in data" +# video = data[key] +# assert isinstance(video, np.ndarray), f"Video {key} is not a numpy array" +# assert video.ndim == 4, f"Video {key} must have 4 dimensions, got {video.ndim}" +# assert video.dtype == np.uint8, f"Video {key} must have dtype uint8, got {video.dtype}" +# shape = video.shape +# expected_resolution = self.dataset_metadata.modalities.video[key].resolution +# assert ( +# shape[1:3] == expected_resolution +# ), f"Video {key} has invalid shape {shape}, expected {expected_resolution}" + +# def transform_torchvision( +# self, frames: np.ndarray +# ) -> list[torch.Tensor | np.ndarray | Image.Image]: +# """ +# frames: [view * T, H, W, C], np.uint8 +# """ +# if self.transforms is None: +# raise ValueError( +# "Transform is not set. Please call set_metadata() before calling __call__()" +# ) +# # Convert to batched tensor, using ToTensor() is too slow +# frames_tensor = torch.from_numpy(frames).to(torch.float32) / 255.0 +# frames_tensor = frames_tensor.permute(0, 3, 1, 2) # [view * T, C, H, W] +# transformed_frames = self.transforms(frames_tensor) +# if self.strong_vision_aug: +# transformed_frames = self.strong_transform(transformed_frames) + +# to_pil = T.ToPILImage() +# transformed_frames = [to_pil(frame) for frame in transformed_frames] +# return transformed_frames # type: ignore + + +# class IdentityTransform(ModalityTransform): +# def __call__(self, data: dict) -> dict: # type: ignore +# warnings.warn("IdentityTransform is used, further transformations is required.") +# output = {} +# for key in self.modality_keys: +# output[key] = data[key] +# return output diff --git a/groot/vla/experiment/__init__.py b/groot/vla/experiment/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2cf08b47119645e476dc0b3cfcb7ed54fa4aaec1 --- /dev/null +++ b/groot/vla/experiment/__init__.py @@ -0,0 +1 @@ +from groot.vla.experiment.experiment import VLATrainer \ No newline at end of file diff --git a/groot/vla/experiment/base.py b/groot/vla/experiment/base.py new file mode 100644 index 0000000000000000000000000000000000000000..7c285f115370edb04bc42cc72bf4baad89363b19 --- /dev/null +++ b/groot/vla/experiment/base.py @@ -0,0 +1,875 @@ +# Copyright 2024 NVIDIA CORPORATION & AFFILIATES +# +# 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. +# +# SPDX-License-Identifier: Apache-2.0 +# This file is modified from https://github.com/haotian-liu/LLaVA/ + +from abc import ABC +import contextlib +import json +import logging +import os +from pathlib import Path +import shutil +import threading +import time +from typing import Optional +import warnings + +from hydra.utils import instantiate +import numpy as np +from omegaconf import DictConfig, OmegaConf, open_dict +import torch +from torch.profiler import ProfilerActivity, profile +from torch.utils.data import DataLoader, Dataset, Sampler +import transformers +from transformers import TrainerCallback, set_seed +from transformers.trainer import ( + # ALL_LAYERNORM_LAYERS, # ShardedDDPOption, # Removed deprecated import + TRAINER_STATE_NAME, + TrainerState, + get_last_checkpoint, + get_parameter_names, + is_sagemaker_mp_enabled, +) + +import groot.vla.common.utils as U +from groot.vla.data.dataset.lerobot_sharded import ShardedLeRobotMixtureDataset +from groot.vla.data.schema import EmbodimentTag +from groot.vla.data.transform import ComposedModalityTransform +from groot.vla.experiment.utils import ( + compute_grad_accum_to_match_global_bs, + dtype_from_string, + get_checkpoint_path, + mprint, + safe_save_model_for_hf_trainer, +) +from groot.vla.utils.timer import ContextTimer + +# Fix resume: https://github.com/huggingface/transformers/pull/34632/files +np_core = np.core +allowlist = [np_core.multiarray._reconstruct, np.ndarray, np.dtype] +# numpy >1.25 defines numpy.dtypes.UInt32DType, but below works for +# all versions of numpy +allowlist += [type(np.dtype(np.uint32))] +torch.serialization.add_safe_globals(allowlist) + +# Define LayerNorm classes locally to replace deprecated ALL_LAYERNORM_LAYERS +LAYERNORM_LAYERS = [ + torch.nn.LayerNorm, + torch.nn.GroupNorm, + torch.nn.InstanceNorm1d, + torch.nn.InstanceNorm2d, + torch.nn.InstanceNorm3d, + torch.nn.LocalResponseNorm, + torch.nn.BatchNorm1d, + torch.nn.BatchNorm2d, + torch.nn.BatchNorm3d, + torch.nn.SyncBatchNorm, +] + + +class LossLoggerCallback(TrainerCallback): + """Callback that writes per-step loss metrics to a JSONL file for offline analysis.""" + + def __init__(self, output_path: str): + self.output_path = output_path + + def on_log(self, args, state, control, logs=None, **kwargs): + if not state.is_world_process_zero or logs is None: + return + entry = {"step": state.global_step} + for key in ("loss", "dynamics_loss_avg", "action_loss_avg", "learning_rate"): + if key in logs: + entry[key] = logs[key] + if len(entry) > 1: # more than just "step" + with open(self.output_path, "a") as f: + f.write(json.dumps(entry) + "\n") + + +class CheckpointFormatCallback(TrainerCallback): + """This callback format checkpoint to make them standalone. For now, it copies all config + files to /checkpoint-{step}/experiment_cfg/: + - conf.yaml + - initial_actions.npz + - metadata.json + """ + + def __init__( + self, run_name: str, exp_cfg_dir: Path | None = None, processor_dir: Path | None = None + ): + """ + Args: + run_name: Name of the experiment run + exp_cfg_dir: Path to the directory containing all experiment metadata + """ + self.exp_cfg_dir = exp_cfg_dir + self.processor_dir = processor_dir + + def on_save(self, args, state, control, **kwargs): + """Called after the trainer saves a checkpoint.""" + if state.is_world_process_zero: + checkpoint_dir = Path(args.output_dir) / f"checkpoint-{state.global_step}" + + # Copy experiment config directory if provided + if self.exp_cfg_dir is not None: + exp_cfg_dst = checkpoint_dir / self.exp_cfg_dir.name + if self.exp_cfg_dir.exists(): + print( + f"Copying experiment config directory {self.exp_cfg_dir} to {exp_cfg_dst}" + ) + shutil.copytree(self.exp_cfg_dir, exp_cfg_dst, dirs_exist_ok=True) + + # Copy processor directory if provided + if self.processor_dir is not None: + if self.processor_dir.exists(): + print(f"Copying processor directory {self.processor_dir} to {checkpoint_dir}") + shutil.copytree(self.processor_dir, checkpoint_dir, dirs_exist_ok=True) + + # Copy wandb_config.json if provided + wandb_config_src = Path(args.output_dir) / "wandb_config.json" + wandb_config_dst = checkpoint_dir / "wandb_config.json" + if wandb_config_src.exists(): + print(f"Copying wandb_config.json from {wandb_config_src} to {wandb_config_dst}") + shutil.copy2(wandb_config_src, wandb_config_dst) + + +class ProfCallback(transformers.TrainerCallback): + """Callback to manage PyTorch profiler during training. + + Dynamically starts/stops the profiler within a specified session step window. + After profiling completes, triggers optional S3 upload and removes itself. + + Args: + profile_dir: Directory to save profile traces + upload_callback: Optional callback to trigger S3 upload after profiling + profile_start_step: Session step to start profiling (default: 50) + profile_end_step: Session step to stop profiling + warmup_steps: Number of warmup steps for profiler schedule (default: 1) + active_steps: Number of active profiling steps (default: 5) + trainer: Trainer instance (required for self-removal after profiling) + record_shapes: Record tensor shapes in profiler (default: False) + with_stack: Record Python stack traces (default: True) + profile_memory: Record memory allocation/deallocation (default: False) + """ + + def __init__( + self, + profile_dir, + upload_callback=None, + profile_start_step=50, + profile_end_step=55, + warmup_steps=1, + active_steps=5, + trainer=None, + record_shapes=False, + with_stack=True, + profile_memory=False, + ): + self.profile_dir = profile_dir + self.upload_callback = upload_callback + self.profile_start_step = profile_start_step + self.profile_end_step = profile_end_step + self.warmup_steps = warmup_steps + self.active_steps = active_steps + self.trainer = trainer + self.record_shapes = record_shapes + self.with_stack = with_stack + self.profile_memory = profile_memory + self.upload_triggered = False + self.starting_global_step = None + self.session_step = 0 + self.prof = None + self.profiling_active = False + self.profiling_complete = False + self.removed_from_trainer = False + + def on_step_begin(self, args, state, control, **kwargs): + # Remove callback after upload triggered to eliminate all overhead + if self.profiling_complete and self.upload_triggered and not self.removed_from_trainer: + if self.trainer is not None and hasattr(self.trainer, "callback_handler"): + try: + self.trainer.callback_handler.callbacks.remove(self) + self.removed_from_trainer = True + logging.info( + f"Removed ProfCallback from trainer at global step {state.global_step}" + ) + except (ValueError, AttributeError) as e: + logging.warning(f"Failed to remove ProfCallback: {e}") + return + + # Early return if profiling already complete + if self.profiling_complete: + return + + # Record starting global step on first call + if self.starting_global_step is None: + self.starting_global_step = state.global_step + + # Calculate session step + self.session_step = state.global_step - self.starting_global_step + + # Start profiler when we reach the profiling window + if self.session_step == self.profile_start_step and self.prof is None: + logging.info( + f"Starting profiler at global step {state.global_step} (session step {self.session_step})" + ) + self.prof = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + schedule=torch.profiler.schedule( + skip_first=0, + wait=0, + warmup=self.warmup_steps, + active=self.active_steps, + repeat=1, + ), + profile_memory=self.profile_memory, + with_stack=self.with_stack, + record_shapes=self.record_shapes, + on_trace_ready=torch.profiler.tensorboard_trace_handler(str(self.profile_dir)), + ) + self.prof.__enter__() + self.profiling_active = True + + def on_step_end(self, args, state, control, **kwargs): + # Early return if profiling already complete + if self.profiling_complete: + return + + # Recalculate session_step to ensure accuracy + if self.starting_global_step is not None: + self.session_step = state.global_step - self.starting_global_step + + # Step profiler if active + if self.profiling_active and self.prof is not None: + self.prof.step() + + # Stop profiler when we reach the end of profiling window + if self.session_step == self.profile_end_step and self.prof is not None: + self.prof.__exit__(None, None, None) + self.profiling_active = False + + # Explicitly release profiler resources to minimize CUPTI overhead + # Combined with TEARDOWN_CUPTI=1 env var for full cleanup + del self.prof + self.prof = None + + # Force CUDA synchronization to ensure profiler cleanup completes + if torch.cuda.is_available(): + torch.cuda.synchronize() + + self.profiling_complete = True + logging.info( + f"Profiler stopped and resources released at global step {state.global_step} " + f"(session step {self.session_step})" + ) + + # Trigger upload if callback provided + if self.upload_callback: + logging.info(f"Triggering upload at global step {state.global_step}...") + self.upload_callback() + + # Mark as ready for callback removal + self.upload_triggered = True + + +class BaseSampler(Sampler): + """Sampler for dataset, which enables `set_epoch` for Dataset. + `set_epoch` will be called by huggingface Trainer at the end of each epoch. + `shuffle` is also supported for training set shuffling + """ + + def __init__(self, data_source: Dataset, shuffle: bool = False, seed: int = 0): + self.data_source = data_source + self.shuffle = shuffle + self.seed = seed + self.epoch = 0 + + def __iter__(self): + if self.shuffle: + g = torch.Generator() + g.manual_seed(self.seed + self.epoch) + # must not add rank here, or randomization will be different for each rank + return iter(torch.randperm(len(self.data_source), generator=g).tolist()) + return iter(range(len(self.data_source))) + + def set_epoch(self, epoch): + self.epoch = epoch + if hasattr(self.data_source, "set_epoch"): + # this is important for dataset + self.data_source.set_epoch(epoch) + + def __len__(self): + return len(self.data_source) + + +class BaseTrainer(transformers.Trainer): + + def __init__(self, **kwargs): + # Increase the cache size limit for torch._dynamo to + # accommodate videos with different numbers of frames. + torch._dynamo.config.cache_size_limit = 1000 + + self.compute_dtype = kwargs.pop("compute_dtype") + self.output_dir = kwargs.pop("output_dir") + self.timer = ContextTimer(self) + + self.world_size = int(os.environ.get("WORLD_SIZE", "1")) + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + self.global_rank = int(os.environ.get("RANK", "0")) + self.node_rank = int(os.environ.get("NODE_RANK", "0")) + + # Get distributed info + self.current_step = 0 + + # Profiling (legacy per-step profiling) + self.enable_profiling = kwargs.pop("enable_profiling", False) + self.profiling_steps = kwargs.pop("profiling_steps", 5) + # Pop new ProfCallback config options (handled in create_trainer, not here) + kwargs.pop("enable_prof_callback", None) + kwargs.pop("profile_start_step", None) + kwargs.pop("profile_warmup_steps", None) + kwargs.pop("profile_active_steps", None) + kwargs.pop("profile_record_shapes", None) + kwargs.pop("profile_with_stack", None) + kwargs.pop("profile_memory", None) + kwargs.pop("msc_profile_url", None) + kwargs.pop("profile_delete_after_upload", None) + if self.enable_profiling: + # Setup profiling directories + self.profile_dir = Path(self.output_dir) / "profiling" + self.memory_profile_dir = self.profile_dir / "memory" + self.torch_profile_dir = self.profile_dir / "torch" + + self.memory_profile_dir.mkdir(exist_ok=True, parents=True) + self.torch_profile_dir.mkdir(exist_ok=True, parents=True) + + # Start recording the memory history. + torch.cuda.memory._record_memory_history(max_entries=100000) + + super().__init__(**kwargs) + + self.loss_queues = {} + self.loss_queue_size = 10 + + def _get_train_sampler(self): + return BaseSampler(self.train_dataset, shuffle=True, seed=self.args.seed) + + def _get_eval_sampler(self, eval_dataset): + return BaseSampler(eval_dataset, shuffle=False) + + def training_step(self, model, inputs, num_items_in_batch=None): + enable_profile = self.enable_profiling and self.current_step % self.profiling_steps == 0 + if enable_profile: + profile_context = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=True, + with_stack=True, + ) + else: + profile_context = contextlib.nullcontext() + + start_time = time.time() + + with self.timer.with_label("training_step"), profile_context as prof: + output = super().training_step(model, inputs) + + time_taken = time.time() - start_time + print( + f"Rank {self.global_rank} time taken for training_step {self.current_step}: {time_taken:.2f} seconds" + ) + + if enable_profile: + trace_path = f"{self.torch_profile_dir}/trace_rank_{self.global_rank}_step_{self.current_step}.json.gz" + print(f"Rank {self.global_rank} exporting torch profile to {trace_path}") + prof.export_chrome_trace(trace_path) + + snapshot_path = f"{self.memory_profile_dir}/memory_snapshot_rank_{self.global_rank}_step_{self.current_step}.pickle" + print(f"Rank {self.global_rank} dumping memory snapshot to {snapshot_path}") + torch.cuda.memory._dump_snapshot(snapshot_path) + + self.current_step += 1 + return output + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + with self.timer.with_label("model_forward"): + outputs = model(inputs) + ### For additional losses, track and log their moving averages + for key, value in outputs.items(): + if key.endswith("_loss") and key != "loss": + # Initialize queue if not exists + if key not in self.loss_queues: + self.loss_queues[key] = [] + + # Add current loss value to queue + current_value = value.item() if torch.is_tensor(value) else value + self.loss_queues[key].append(current_value) + + # Keep only last N values + if len(self.loss_queues[key]) > self.loss_queue_size: + self.loss_queues[key].pop(0) + + # Log average every 10 steps + if self.current_step % self.loss_queue_size == 0: + avg_loss = sum(self.loss_queues[key]) / len(self.loss_queues[key]) + self.log({f"{key}_avg": avg_loss}) + + loss = outputs["loss"] + + return (loss, outputs) if return_outputs else loss + + def create_optimizer(self): + """ + Setup the optimizer. + + We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the + Trainer's init through `optimizers`, or subclass and override this method in a subclass. + """ + if is_sagemaker_mp_enabled(): + return super().create_optimizer() + + opt_model = self.model + + if self.optimizer is None: + decay_parameters = get_parameter_names(opt_model, LAYERNORM_LAYERS) + decay_parameters = [name for name in decay_parameters if "bias" not in name] + optimizer_grouped_parameters = [ + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n in decay_parameters and p.requires_grad) + ], + "weight_decay": self.args.weight_decay, + }, + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n not in decay_parameters and p.requires_grad) + ], + "weight_decay": 0.0, + }, + ] + + optimizer_cls, optimizer_kwargs = transformers.Trainer.get_optimizer_cls_and_kwargs( + self.args + ) + self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs) + + # DeepSpeed CPU Adam (ZeRO offload) expects 'bias_correction' in each param group. + # HuggingFace Trainer's AdamW does not set it, causing KeyError in cpu_adam.step(). + if getattr(self.args, "deepspeed", None): + for group in self.optimizer.param_groups: + group.setdefault("bias_correction", True) + + return self.optimizer + + def save_model(self, output_dir: Optional[str], _internal_call: bool): + + ## save tuned model separately + if self.is_deepspeed_enabled: + state_dict = self.accelerator.get_state_dict(self.deepspeed) + else: + state_dict = self.model.state_dict() + + if self.base_cfg.save_lora_only: + # Save only the trainable parameters + train_key = [k for k, v in self.model.named_parameters() if v.requires_grad] + lora_state_dict = {k: v for k, v in self.model.state_dict().items() if k in train_key} + state_dict = lora_state_dict + + if self.args.should_save: + ret = self.model.save_pretrained(output_dir, state_dict=state_dict) + + # can separately save the VLM model for downstream evalualtion + if self.base_cfg.save_llm: + llm_output_dir = os.path.join(output_dir, "llm") + self.model.backbone.model.save_pretrained(llm_output_dir) + + if self.base_cfg.save_value_model: + assert hasattr( + self.model.action_head, "value_model" + ), f"Value model not found in action head: {type(self.model.action_head)}" + value_model_output_dir = os.path.join(output_dir, "value_model") + self.model.action_head.value_model.save_pretrained(value_model_output_dir) + + return ret + + def train( + self, + resume_from_checkpoint=None, + trial=None, + ignore_keys_for_eval=None, + **kwargs, + ): + """Correctly set self.state from checkpoint so get_train_dataloader can read from it.""" + if resume_from_checkpoint is False: + resume_from_checkpoint = None + + if isinstance(resume_from_checkpoint, bool) and resume_from_checkpoint: + resume_from_checkpoint = get_last_checkpoint(self.args.output_dir) + if resume_from_checkpoint is None: + raise ValueError( + f"No valid checkpoint found in output directory ({self.args.output_dir})" + ) + + if resume_from_checkpoint is not None: + # In case of repeating the find_executable_batch_size, set `self._train_batch_size` properly + self.state = TrainerState.load_from_json( + os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME) + ) + return super().train(resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs) + + def get_train_dataloader(self) -> DataLoader: + """ + Returns the training [`~torch.utils.data.DataLoader`]. + + Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed + training if necessary) otherwise. + + Subclass and override this method if you want to inject some custom behavior. + """ + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + if not isinstance(train_dataset, (ShardedLeRobotMixtureDataset)): + return super().get_train_dataloader() + + # During resume, don't skip the data + self.args.ignore_data_skip = True + curr_global_step = self.state.global_step + print(f"Current global step: {curr_global_step}") + if curr_global_step > 0: + new_seed = train_dataset.seed + curr_global_step + train_dataset.reset_seed(new_seed) + print( + f"Resetting seed to {new_seed}. Please note that this will make the experiment non-reproducible." + ) + + print("Creating custom train dataloader") + # Handle the case where the dataset is an IterableDataset + data_collator = self.data_collator + data_collator = self._get_collator_with_removed_columns( + data_collator, description="training" + ) + + dataloader_params = { + "batch_size": self._train_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + } + # persistent_workers is only valid when num_workers > 0 (PyTorch raises otherwise) + if self.args.dataloader_num_workers > 0: + dataloader_params["persistent_workers"] = self.args.dataloader_persistent_workers + + return DataLoader(train_dataset, **dataloader_params) + + +class BaseExperiment(ABC): + def __init__(self, cfg: DictConfig): + # assert cfg.save_steps == 500, "save_steps must be 500 for standarized evaluation" + assert cfg.max_steps > 0, "max_steps must be > 0 for standarized evaluation" + # patched: reduced to 2 for disk space + assert cfg.save_total_limit >= 2, "save_total_limit must be >= 2" + + if cfg.load_from_yaml is not None: + # Override the default config with the loaded config. + loaded_cfg = OmegaConf.load(cfg.load_from_yaml) + cfg = loaded_cfg # overwrite + + # Check if evaluation transforms are valid. + assert cfg.transforms is not None, "Evaluation transforms are not provided." + for tag, transform_cfg in cfg.transforms.items(): + try: + # Check if the tag is a valid EmbodimentTag + _ = EmbodimentTag(tag) + # Check if the transform is a valid ComposedModalityTransform + transform = instantiate(transform_cfg) + assert isinstance(transform, ComposedModalityTransform), f"{transform=}" + except Exception as e: + raise ValueError(f"Evaluation transform {tag} is invalid: {e}") + + # Instantiate the training arguments. + cfg.training_args.output_dir = cfg.training_args.output_dir.rstrip("/") + cfg.training_args.run_name = cfg.training_args.output_dir.split("/")[-1] + print(f"Run name: {cfg.training_args.run_name}") + training_args = instantiate(cfg.training_args) + set_seed(training_args.seed) + + # Set the environment variables for wandb. + if "WANDB_PROJECT" not in os.environ: + os.environ["WANDB_PROJECT"] = cfg.wandb_project + if "WANDB_RUN_ID" not in os.environ: + runtime_id = os.environ.get("RUNTIME_ID", None) + """If a RUNTIME_ID is available in the environment, we use it as the wandb id, + which will allow to display the evaluation results and the training results + in the same wandb run. Otherwise, we create a new run.""" + if runtime_id: + os.environ["WANDB_RUN_ID"] = runtime_id + os.environ["WANDB_DIR"] = training_args.output_dir + + # Create the experiment config directory. + output_dir = Path(training_args.output_dir) + exp_cfg_dir = output_dir / "experiment_cfg" + exp_cfg_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(cfg, exp_cfg_dir / "conf.yaml", resolve=True) + + wandb_config_file = output_dir / "wandb_config.json" + with open(wandb_config_file, "w") as f: + json.dump( + { + "project": os.environ.get("WANDB_PROJECT", ""), + "run_id": os.environ.get("WANDB_RUN_ID", ""), + }, + f, + ) + + # Check if we are resuming training. + resume_path, continue_training = get_checkpoint_path(training_args.output_dir) + if not continue_training: + print(f"Models is ready under {training_args.output_dir}. Skip training.") + exit(0) + if resume_path: + print(f"Resuming training from {resume_path}") + resume_from_checkpoint = True + else: + # First time training. + resume_from_checkpoint = False + + # Instantiate the model. + model = self.create_model(cfg, training_args) + + if hasattr(model.action_head, "max_steps"): + model.action_head.max_steps = cfg.max_steps + + # Make sure model_dtype and training_args dtype are compatible. + compute_dtype = dtype_from_string(model.config.model_dtype) + + # Create the train dataset. + # Dump the metadata; necessary for policy to normalize the input and unnormalize the output + train_dataset = self.create_train_dataset(cfg, model) + print("Using dataset:") + print(train_dataset) + assert ( + train_dataset.merged_metadata is not None + ), "You must set metadata_config.merge=true in order to save the metadata." + + metadata_save_path = exp_cfg_dir / "metadata.json" + U.json_dump( + {k: v.model_dump(mode="json") for k, v in train_dataset.merged_metadata.items()}, + metadata_save_path, + indent=4, + ) + print("Successfully dumped metadata") + + val_dataset = self.create_val_dataset(cfg, model) + data_collator = self.create_data_collator(cfg, model) + trainer = self.create_trainer( + cfg=cfg, + exp_cfg_dir=exp_cfg_dir, + model=model, + training_args=training_args, + train_dataset=train_dataset, + val_dataset=val_dataset, + data_collator=data_collator, + compute_dtype=compute_dtype, + ) + self.cfg = cfg + self.exp_cfg_dir = exp_cfg_dir + self.training_args = training_args + self.resume_from_checkpoint = resume_from_checkpoint + self.train_dataset = train_dataset + self.trainer = trainer + + def create_model(self, cfg, training_args): + model = instantiate(cfg.model) + + if cfg.pretrained_model_path is not None: + mprint(f"Loading pretrained weights from: {cfg.pretrained_model_path}") + import json, gc + from safetensors.torch import load_file + + ckpt_dir = cfg.pretrained_model_path + safetensors_index_path = os.path.join(ckpt_dir, "model.safetensors.index.json") + safetensors_path = os.path.join(ckpt_dir, "model.safetensors") + + if os.path.exists(safetensors_index_path): + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + for shard_file in sorted(set(index["weight_map"].values())): + shard_path = os.path.join(ckpt_dir, shard_file) + mprint(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + model.load_state_dict(shard_state_dict, strict=False) + del shard_state_dict + gc.collect() + elif os.path.exists(safetensors_path): + state_dict = load_file(safetensors_path) + model.load_state_dict(state_dict, strict=False) + else: + raise FileNotFoundError( + f"No weights found at '{ckpt_dir}'. " + "Expected 'model.safetensors' or 'model.safetensors.index.json'." + ) + + if (hasattr(model, 'action_head') + and hasattr(model.action_head, 'inject_lora_after_loading') + and model.action_head.config.defer_lora_injection): + model.action_head.inject_lora_after_loading() + + mprint("Successfully loaded pretrained weights") + + model.config.resume_path = model.config._name_or_path = training_args.output_dir + mprint(f"{model}\n") + return model + + def create_train_dataset(self, cfg, model): + assert torch.distributed.is_initialized() + train_dataset = instantiate(cfg.train_dataset) + return train_dataset + + def create_val_dataset(self, cfg, model): + return None + + def create_data_collator(self, cfg, model): + return instantiate(cfg.data_collator) + + def create_trainer( + self, + cfg, + exp_cfg_dir, + model, + training_args, + train_dataset, + val_dataset, + data_collator, + compute_dtype, + ): + # Set the gradient accumulation steps. + if cfg.global_batch_size is not None: + global_bs = cfg.global_batch_size + bs = training_args.per_device_train_batch_size + grad_acc = compute_grad_accum_to_match_global_bs(global_bs, bs) + training_args.gradient_accumulation_steps = grad_acc + print( + f"Set global batch size to {global_bs}, set gradient accumulation steps to {grad_acc}" + ) + elif cfg.raise_error_if_global_batch_size_not_set: + raise ValueError( + "global_batch_size is not set. To ensure the scripts can be reproduced regardless of the number of nodes used, please set this." + ) + else: + warnings.warn( + "global_batch_size is not set. This is fine for debugging, but please set this for real experiments." + ) + + # Instantiate the partial trainer. + trainer_partial = instantiate( + cfg.trainer, + model=model, + output_dir=training_args.output_dir, + train_dataset=train_dataset, + eval_dataset=val_dataset, + compute_dtype=compute_dtype, + ) + + # Fully instantiate the trainer with dataclasses instances. + trainer = trainer_partial(data_collator=data_collator, args=training_args) + trainer.base_cfg = cfg + train_dl_len = len(trainer.get_train_dataloader()) + eval_dl_len = ( + len(trainer.get_eval_dataloader()) if val_dataset is not None else "no eval dataloader" + ) + + # Save the total training steps in the config. + with open_dict(cfg): + cfg.total_training_steps = train_dl_len * cfg.training_args.num_train_epochs + + # Save config. + OmegaConf.save(cfg, exp_cfg_dir / "conf.yaml", resolve=True) + + run_name = cfg.training_args.get("run_name", None) + ckpt_format_callback = CheckpointFormatCallback(run_name=run_name, exp_cfg_dir=exp_cfg_dir) + trainer.add_callback(ckpt_format_callback) + + loss_log_path = str(Path(training_args.output_dir) / "loss_log.jsonl") + trainer.add_callback(LossLoggerCallback(output_path=loss_log_path)) + + + # Add profiling callback (local profiling only, no S3 upload) + # Local: {output_dir}/profiling/rank_{id}/*.pt.trace.json + if cfg.trainer.get("enable_prof_callback", False): + output_dir = Path(training_args.output_dir) + global_rank = int(os.environ.get("RANK", "0")) + + # Get profiling configuration from trainer config + profile_start_step = cfg.trainer.get("profile_start_step", 50) + profile_warmup_steps = cfg.trainer.get("profile_warmup_steps", 1) + profile_active_steps = cfg.trainer.get("profile_active_steps", 5) + profile_record_shapes = cfg.trainer.get("profile_record_shapes", False) + profile_with_stack = cfg.trainer.get( + "profile_with_stack", False + ) # Default False to match omni (stack traces add significant file size) + profile_memory = cfg.trainer.get("profile_memory", False) + + # Calculate end step + profile_end_step = profile_start_step + profile_warmup_steps + profile_active_steps - 1 + + # Setup profile directory with rank subdirectory: {output_dir}/profiling/rank_{id}/ + profile_dir = output_dir / "profiling" / f"rank_{global_rank}" + profile_dir.mkdir(parents=True, exist_ok=True) + + mprint( + f"Profiling enabled: steps {profile_start_step}-{profile_end_step}, " + f"saving to {profile_dir}" + ) + + # Add ProfCallback + trainer.add_callback( + ProfCallback( + profile_dir=profile_dir, + upload_callback=None, + profile_start_step=profile_start_step, + profile_end_step=profile_end_step, + warmup_steps=profile_warmup_steps, + active_steps=profile_active_steps, + trainer=trainer, + record_shapes=profile_record_shapes, + with_stack=profile_with_stack, + profile_memory=profile_memory, + ) + ) + + mprint( + f"train dataloader length: {train_dl_len}\n" + f"eval dataloader length: {eval_dl_len}\n" + f"train dataset length: {len(trainer.train_dataset)}\n" + f"GPU memory before training: {torch.cuda.memory_allocated() / 1024 / 1024 / 1024} GB", + flush=True, + ) + return trainer + + def train(self): + # Start training. + self.trainer.train(resume_from_checkpoint=self.resume_from_checkpoint) + self.trainer.save_state() + safe_save_model_for_hf_trainer( + trainer=self.trainer, output_dir=self.training_args.output_dir + ) diff --git a/groot/vla/experiment/experiment.py b/groot/vla/experiment/experiment.py new file mode 100644 index 0000000000000000000000000000000000000000..e02d36c52999ffadd68ec0c00333f794b7db25d2 --- /dev/null +++ b/groot/vla/experiment/experiment.py @@ -0,0 +1,135 @@ +import logging +import os +import time +from typing import Optional + +import hydra +import numpy as np +from omegaconf import DictConfig +import torch + +from groot.vla.experiment.base import BaseExperiment, BaseTrainer +from groot.vla.utils.action_args_override_utils import apply_action_overrides + +logger = logging.getLogger(__name__) + + +INITIAL_ACTIONS_FILENAME = "initial_actions.npz" + + +class ForceRestart(ValueError): + pass + + +class VLATrainer(BaseTrainer): + def __init__(self, **kwargs): + self.benchmark_time = kwargs.pop("benchmark_time", False) + self.step_timer = None + self.num_trials = kwargs.pop("num_trials", 10) + self.curr_trial = 0 + self.all_times = [] + self.start_time = time.time() + self.restart_max_seconds = kwargs.pop("restart_max_seconds", 0) + import torch.distributed as dist + + self.rank = dist.get_rank() + + self.micro_global_step = 0 + + super().__init__(**kwargs) + + def training_step(self, model, inputs, *args, **kwargs): + self.micro_global_step += 1 + + if hasattr(self.model.action_head, "global_step"): + self.model.action_head.global_step = self.state.global_step + + if self.benchmark_time: + if self.state.global_step % 100 == 0: + if self.step_timer is not None: + elapsed_time = time.time() - self.step_timer + self.all_times.append(elapsed_time) + self.curr_trial += 1 + self.step_timer = time.time() + if self.curr_trial >= self.num_trials: + exit(0) + if self.state.global_step % self.state.save_steps == 1: + if self.restart_max_seconds > 0: + cur_time = time.time() + if (cur_time - self.start_time) > self.restart_max_seconds: + raise ForceRestart(f"Exceeded time limit {self.restart_max_seconds} seconds") + loss_dict = super().training_step(model, inputs, *args, **kwargs) + return loss_dict + + +class VLATrainerInferenceBenchmark(VLATrainer): + + def compute_loss(self, model, inputs, return_outputs=False): + + warmup_steps = 100 + measure_steps = 100 + + model.eval() + + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + with torch.inference_mode(): + for i in range(warmup_steps): + action = model.module.get_action(inputs) + action.keys() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start_event.record() + + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + with torch.inference_mode(): + for i in range(measure_steps): + action = model.module.get_action(inputs) + action.keys() + + end_event.record() + torch.cuda.synchronize() + elapsed_time = start_event.elapsed_time(end_event) + + time_per_step = elapsed_time / measure_steps + exit() + + +class VLAExperiment(BaseExperiment): + + def __init__(self, cfg: DictConfig): + super().__init__(cfg) + # Dump the initial actions + if hasattr(self.train_dataset, "get_initial_actions"): + # We only dump the initial actions for the real robot dataset + # Sim dataset doesn't have this function + """ + initial_actions: list[dict[str, dict[str, np.ndarray]]] + 0: (the first dataset) + trajectory_name: + action_key: + action: np.ndarray + 1: (the second dataset) + ... + """ + initial_actions = self.train_dataset.get_initial_actions() + if len(initial_actions) > 0: + initial_actions_path = self.exp_cfg_dir / INITIAL_ACTIONS_FILENAME + np.savez(str(initial_actions_path), initial_actions) + print("Successfully dumped initial actions") + else: + print("No initial actions to dump") + + +@hydra.main(config_path="../configs", config_name="conf", version_base=None) +def main(cfg): + # Automatically update action dim and action horizon keys if specified in the config + cfg = apply_action_overrides(cfg) + + experiment = VLAExperiment(cfg) + experiment.train() + + +if __name__ == "__main__": + main() diff --git a/groot/vla/experiment/utils.py b/groot/vla/experiment/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d54993b2096aa6cc38ada6151da19b83d26ae2e5 --- /dev/null +++ b/groot/vla/experiment/utils.py @@ -0,0 +1,165 @@ +""" +Originally trinity.train.utils +""" + +from dataclasses import dataclass +import os +import pathlib +from pathlib import Path +import re +import shutil + +import torch +import torch.nn as nn +from transformers import PretrainedConfig, Trainer + + +def dtype_from_string(dtype_str): + if dtype_str == "bfloat16": + return torch.bfloat16 + elif dtype_str == "float16": + return torch.float16 + elif dtype_str == "float32": + return torch.float32 + else: + raise ValueError(f"Unsupported dtype_str {dtype_str}") + + +def rprint(*args, **kwargs): + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + if world_size > 1: + return print(f"[dist-{rank}-of-{world_size}]", *args, **kwargs) + else: + return print(*args, **kwargs) + + +def mprint(*args, **kwargs): + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + if world_size > 1: + if rank == 0: + return print(f"[dist-{rank}-of-{world_size}]", *args, **kwargs) + else: + return + else: + return print(*args, **kwargs) + + +def is_local(model_name_or_path: str) -> bool: + return os.path.isdir(model_name_or_path) + + +def get_checkpoint_path(output_dir: str, checkpoint_prefix: str = "checkpoint") -> str | None: + output_dir = os.path.abspath(output_dir) + pathlib_dir = pathlib.Path(output_dir) + + if list(pathlib_dir.glob("config.json")): + # training has been finished + return output_dir, False + else: + try: + ordering_and_checkpoint_path = [] + glob_checkpoints = [ + str(x) + for x in pathlib.Path(output_dir).glob(f"{checkpoint_prefix}-*") + if os.path.isdir(x) + ] + for path in glob_checkpoints: + regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path) + if regex_match is not None and regex_match.groups() is not None: + ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path)) + checkpoints_sorted = sorted(ordering_and_checkpoint_path) + return checkpoints_sorted[-1][1], True + except IndexError: + return None, True + + +def prepare_config_for_training( + config: PretrainedConfig, model_args: dataclass, training_args: dataclass, data_args: dataclass +) -> None: + ## set default dtype + # config.model_dtype = "bfloat16" if training_args.bf16 else "float16" + + ## set tuning modules + config.tune_language_model = training_args.tune_language_model + config.tune_vision_tower = training_args.tune_vision_tower + config.tune_mm_projector = training_args.tune_mm_projector + + +def safe_save_model_for_hf_trainer(trainer: Trainer, output_dir: str): + """Collects the state dict and dump to disk.""" + if trainer.deepspeed: + torch.cuda.synchronize() + trainer.save_model(output_dir, _internal_call=True) + return + + state_dict = trainer.model.state_dict() + if trainer.args.should_save: + cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()} + del state_dict + trainer._save(output_dir, state_dict=cpu_state_dict) # noqa + + +def compute_grad_accum_to_match_global_bs(global_bs: int, bs: int): + num_devices = torch.distributed.get_world_size() + per_step_bs = bs * num_devices + assert global_bs % per_step_bs == 0, f"{global_bs=}, {per_step_bs=}" + num_grad_accum = global_bs // per_step_bs + return num_grad_accum + + +def get_training_param_info(model): + module_states = dict() + for module_name, module in model.named_children(): + key = f"{module_name}({module.__class__.__name__})" + if all([p.requires_grad for p in module.parameters()]): + module_states[key] = "true" + elif all([not p.requires_grad for p in module.parameters()]): + module_states[key] = "false" + else: + module_states[key] = get_training_param_info(module) + + return module_states + + +def get_param_count_tree(model: nn.Module): + """ + Calculate parameters for the model, structure them as a nested dictionary, + and save the result as a formatted JSON file. + """ + + def format_param_count(count: int) -> str: + """Format the count as a string in millions, e.g. 11M or 5.5M.""" + count_in_millions = count / 1e6 + # If the value is an integer, display without decimal places + if count_in_millions.is_integer(): + return f"{int(count_in_millions)}M" + else: + return f"{count_in_millions:.2f}M" + + def module_to_dict(module: nn.Module, module_name: str) -> dict: + """ + Recursively convert a module and its children into a nested dictionary. + The key is formatted as "module_name (ClassName, param_count)". + """ + + total_count = sum(p.numel() for p in module.parameters()) + formatted_total = format_param_count(total_count) + key = f"{module_name} ({module.__class__.__name__}, {formatted_total})" + + # Get immediate children modules + children = list(module.named_children()) + if children: + nested = {} + for child_name, child_module in children: + # Recursively convert child modules + nested.update(module_to_dict(child_module, child_name)) + return {key: nested} + else: + # Leaf module: return empty dict as value + return {key: {}} + + # Start from the top-level module (you can change the name "model" as needed) + nested_dict = module_to_dict(model, "model") + return nested_dict diff --git a/groot/vla/model/__init__.py b/groot/vla/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/__init__.py b/groot/vla/model/dreamzero/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/action_head/__init__.py b/groot/vla/model/dreamzero/action_head/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py b/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py new file mode 100644 index 0000000000000000000000000000000000000000..92b72a1bf7c55b0079496d2f9fbd4d2e94eb3e0c --- /dev/null +++ b/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py @@ -0,0 +1,1403 @@ +from dataclasses import dataclass, field +import logging +import time +from typing import TypeAlias, cast +import os + +from accelerate import load_checkpoint_and_dispatch + +from einops import rearrange +from hydra.utils import instantiate +from peft import LoraConfig, get_peft_model +import torch +from torch import nn +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +from safetensors.torch import load_file +import json +from huggingface_hub import hf_hub_download + + +logger = logging.getLogger(__name__) + +WAN_HF_REPO_ID = "Wan-AI/Wan2.1-I2V-14B-480P" +WAN22_HF_REPO_ID = "Wan-AI/Wan2.2-TI2V-5B" + + +def hf_download(filename: str, repo_id: str = WAN_HF_REPO_ID) -> str: + """Download a file from the specified HuggingFace repo to HF cache.""" + path = hf_hub_download(repo_id=repo_id, filename=filename) + return path + + +def ensure_file(path: str | None, hf_filename: str, repo_id: str = WAN_HF_REPO_ID) -> str: + """Return a valid local path: use `path` if it exists, otherwise download from HuggingFace.""" + if path is not None and os.path.exists(path): + return path + return hf_download(hf_filename, repo_id) + +from torch.distributions import Beta +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +from torchvision.transforms import v2 +from transformers import PretrainedConfig +from transformers.feature_extraction_utils import BatchFeature + +from groot.vla.model.n1_5.action_head.base_action_head import ActionHead +from groot.vla.model.dreamzero.modules.flow_match_scheduler import FlowMatchScheduler +from groot.vla.model.dreamzero.modules.vram_management import enable_vram_management, AutoWrappedModule, AutoWrappedLinear +from groot.vla.model.dreamzero.modules.wan_video_text_encoder import T5RelativeEmbedding, T5LayerNorm +from groot.vla.model.dreamzero.modules.flow_unipc_multistep_scheduler import FlowUniPCMultistepScheduler + + +KVCacheType: TypeAlias = torch.Tensor + +@dataclass +class WANPolicyHeadConfig(PretrainedConfig): + add_pos_embed: bool = field( + default=True, metadata={"help": "Whether to add positional embedding"} + ) + model_dtype: str = field(default="float32", metadata={"help": "Model data type."}) + diffusion_model_cfg: dict = field( + default=None, metadata={"help": "Diffusion model configuration."} + ) + input_embedding_dim: int = field( + default=1536, metadata={"help": "Input embedding channel dimension."} + ) + backbone_embedding_dim: int = field( + default=1536, metadata={"help": "Backbone embedding channel dimension."} + ) + tiled: bool = field(default=True, metadata={"help": "Whether to use tiled input."}) + tile_size_height: int = field(default=34, metadata={"help": "Tile size height."}) + tile_size_width: int = field(default=34, metadata={"help": "Tile size width."}) + tile_stride_height: int = field(default=18, metadata={"help": "Tile stride height."}) + tile_stride_width: int = field(default=16, metadata={"help": "Tile stride width."}) + num_frame_per_block: int = field(default=1, metadata={"help": "Number of frames per block."}) + # Target video (H, W) for Wan22 resize. When set, videos are resized to this before VAE so latent + # spatial size matches. Use height/width divisible by 32 for WanVideoVAE38 (16x) so latent H,W are even. + target_video_height: int | None = field(default=None, metadata={"help": "Target video height for resize (e.g. 160 for even latent with VAE38)."}) + target_video_width: int | None = field(default=None, metadata={"help": "Target video width for resize (e.g. 320)."}) + + lora_rank: int = field(default=4, metadata={"help": "LoRA rank."}) + lora_alpha: int = field(default=4, metadata={"help": "LoRA alpha."}) + lora_target_modules: str = field(default="q,k,v,o,ffn.0,ffn.2") + init_lora_weights: str = field(default="kaiming", metadata={"help": "LoRA initialization method."}) + train_architecture: str= field(default="lora", metadata={"help": "Train architecture."}) + skip_component_loading: bool = field(default=False, metadata={"help": "Skip loading individual component weights (used when loading from full pretrained model)."}) + + use_gradient_checkpointing: bool = field(default=True, metadata={"help": "Whether to use gradient checkpointing."}) + qformer_cfg: dict = field(default=None, metadata={"help": "Qformer configuration."}) + hidden_size: int = field(default=1024, metadata={"help": "Input embedding dimension."}) + max_seq_len: int = field(default=1024, metadata={"help": "Maxium Sequence Length"}) + action_dim: int = field(default=None, metadata={"help": "Action dimension."}) + action_horizon: int = field(default=None, metadata={"help": "Action horizon."}) + noise_beta_alpha: float = field(default=1.5, metadata={"help": ""}) + noise_beta_beta: float = field(default=1.0, metadata={"help": ""}) + noise_s: float = field( + default=0.999, metadata={"help": "Flow matching noise Beta distribution s."} + ) + # High noise emphasis for BASE (coupled) training - applies Beta distribution to BOTH video and action together + use_high_noise_emphasis: bool = field( + default=False, metadata={"help": "Use Beta distribution for noise sampling (biases BOTH video and action towards high noise levels together)."} + ) + high_noise_beta_alpha: float = field( + default=3.0, metadata={"help": "Beta alpha for high noise emphasis. Beta(3,1): mean=0.75, Beta(5,1): mean=0.83. Higher = more high noise bias."} + ) + # Decoupled noise sampling config for training-inference alignment + # When enabled: video uses Beta(alpha,beta) biased towards high noise, action uses independent uniform + decouple_video_action_noise: bool = field( + default=False, metadata={"help": "Decouple video/action noise: video uses Beta distribution (high noise bias), action uses independent uniform."} + ) + video_noise_beta_alpha: float = field( + default=3.0, metadata={"help": "Beta alpha for video noise. Beta(3,1): mean=0.75, Beta(5,1): mean=0.83. Higher alpha = more bias to high noise."} + ) + video_noise_beta_beta: float = field( + default=1.0, metadata={"help": "Beta beta for video noise. Keep at 1.0."} + ) + # Decoupled inference config - allows video to stay noisy while action fully denoises + decouple_inference_noise: bool = field( + default=False, metadata={"help": "Use decoupled noise schedules during inference (video stays noisy, action fully denoises)."} + ) + video_inference_final_noise: float = field( + default=0.8, metadata={"help": "Final noise level for video during decoupled inference (0.0-1.0). E.g., 0.8 means video ends at 80% noise."} + ) + num_timestep_buckets: int = field( + default=1000, metadata={"help": "Number of timestep discretization buckets."} + ) + num_inference_timesteps: int = field( + default=None, + metadata={"help": "Number of inference steps for noise diffusion."}, + ) + max_num_embodiments: int = field(default=32, metadata={"help": "Number of embodiments."}) + tune_projector: bool = field(default=True, metadata={"help": "Whether to tune the projector."}) + tune_diffusion_model: bool = field( + default=True, metadata={"help": "Whether to tune the diffusion model."} + ) + load_pretrained_det_decode_layer_path: str = field( + default=None, metadata={"help": "Path to pretrained detection model."} + ) + detection_coeff: float = field(default=1.0, metadata={"help": "Detection coefficient."}) + + freeze_decode_layer: bool = field(default=False) + expand_batch: int = field(default=None) + use_vlln: bool = field(default=True) + defer_lora_injection: bool = field(default=False, metadata={"help": "Defer LoRA injection until after loading pretrained weights."}) + + vl_self_attention_cfg: dict = field(default=None) + text_encoder_cfg: dict = field(default=None) + image_encoder_cfg: dict = field(default=None) + vae_cfg: dict = field(default=None) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + for key, value in kwargs.items(): + setattr(self, key, value) + + +class WANPolicyHead(ActionHead): + config_class = WANPolicyHeadConfig + supports_gradient_checkpointing = True + + def __init__( + self, + config: WANPolicyHeadConfig, + ): + super().__init__() + self.tiled = config.tiled + self.tile_size_height = config.tile_size_height + self.tile_size_width = config.tile_size_width + self.tile_stride_height = config.tile_stride_height + self.tile_stride_width = config.tile_stride_width + self.num_frame_per_block = config.num_frame_per_block + self.hidden_size = config.hidden_size + self.num_frames = config.num_frames + self.text_encoder = instantiate(config.text_encoder_cfg) + self.image_encoder = instantiate(config.image_encoder_cfg) + self.vae = instantiate(config.vae_cfg) + self.scheduler = FlowMatchScheduler(shift=5, sigma_min=0.0, extra_one_step=True) + self.model_names = ['text_encoder'] + + self.num_inference_steps = 16 + self.seed = 1140 + self.cfg_scale = 5.0 + self.denoising_strength = 1.0 + self.sigma_shift = 5.0 + self.kv_cache1: KVCacheType | None = None + self.kv_cache_neg: KVCacheType | None = None + self.crossattn_cache: KVCacheType | None = None + self.crossattn_cache_neg: KVCacheType | None = None + + self.global_step = 0 + self.max_steps = 0 + self.lora_rank = config.lora_rank + self.lora_alpha = config.lora_alpha + self.lora_target_modules = config.lora_target_modules + self.init_lora_weights = config.init_lora_weights + self.train_architecture = config.train_architecture + self.clip_feas = None + self.ys = None + self.current_start_frame = 0 + self.language = None + + self.ip_rank = 0 + self.ip_size = 1 + self.ip_group = None + + self._device = "cuda" + self.dynamic_cache_schedule = os.getenv("DYNAMIC_CACHE_SCHEDULE", "False").lower() == "true" + + + num_dit_steps = 8 + if os.getenv("NUM_DIT_STEPS") is not None: + num_dit_steps = int(os.getenv("NUM_DIT_STEPS")) + if num_dit_steps == 5: + self.dit_step_mask = [True, True, True, False, False, False, False, True, False, False, False, False, True, False, False, False] + elif num_dit_steps == 6: + self.dit_step_mask = [True, True, False, False, False, True, False, False, False, False, True, False, False, False, True, True] + elif num_dit_steps == 7: + self.dit_step_mask = [True, True, True, False, False, False, True, False, False, False, True, False, False, False, True, True] + elif num_dit_steps == 8: + self.dit_step_mask = [True, True, True, False, False, False, True, False, False, False, True, False, False, True, True, True] + else: + self.dit_step_mask = [True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True] + assert self.dit_step_mask[0] == True, "first step must be True" + + self.normalize_video = v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + + + self.use_gradient_checkpointing = config.use_gradient_checkpointing + if self.training: + self.scheduler.set_timesteps(1000, training=True) + + + self.input_embedding_dim = config.input_embedding_dim + + self.cpu_offload = False + + self.model = instantiate(config.diffusion_model_cfg) + self.action_dim = config.action_dim + self.action_horizon = config.action_horizon + self.num_inference_timesteps = config.num_inference_timesteps + + text_enc_path = ensure_file( + self.text_encoder.text_encoder_pretrained_path, + "models_t5_umt5-xxl-enc-bf16.pth", + ) + self.text_encoder.load_state_dict(torch.load(text_enc_path, map_location='cpu')) + + img_enc_path = ensure_file( + self.image_encoder.image_encoder_pretrained_path, + "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", + ) + self.image_encoder.model.load_state_dict(torch.load(img_enc_path, map_location='cpu'), strict=False) + + # Wan2.2 (WanVideoVAE38, z_dim=48) uses Wan2.2_VAE.pth; Wan2.1 uses Wan2.1_VAE.pth + vae_hf_filename = "Wan2.2_VAE.pth" if getattr(self.vae, "z_dim", 16) == 48 else "Wan2.1_VAE.pth" + vae_repo_id = WAN22_HF_REPO_ID if getattr(self.vae, "z_dim", 16) == 48 else WAN_HF_REPO_ID + vae_path = ensure_file( + self.vae.vae_pretrained_path, + vae_hf_filename, + repo_id=vae_repo_id, + ) + self.vae.model.load_state_dict(torch.load(vae_path, map_location='cpu')) + + if not config.skip_component_loading: + dit_dir = self.model.diffusion_model_pretrained_path + # Wan2.2 (in_dim=48) uses Wan2.2-TI2V-5B repo; Wan2.1 uses Wan2.1-I2V-14B-480P + dit_repo_id = WAN22_HF_REPO_ID if getattr(self.model, "in_dim", 16) == 48 else WAN_HF_REPO_ID + if dit_dir is None or not os.path.isdir(dit_dir): + index_path = hf_hub_download(repo_id=dit_repo_id, filename="diffusion_pytorch_model.safetensors.index.json") + dit_dir = os.path.dirname(index_path) + with open(index_path, 'r') as f: + index = json.load(f) + for shard_file in set(index["weight_map"].values()): + hf_hub_download(repo_id=dit_repo_id, filename=shard_file) + + if dit_dir is not None: + safetensors_path = os.path.join(dit_dir, "diffusion_pytorch_model.safetensors") + safetensors_index_path = os.path.join(dit_dir, "diffusion_pytorch_model.safetensors.index.json") + state_dict = {} + + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(dit_dir, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict = load_file(safetensors_path) + + else: + raise ValueError(f"No safetensors file found at {safetensors_path} or {safetensors_index_path}") + + missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False) + + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + + print("Successfully loaded pretrained weights") + else: + print("Skipping individual component loading (loading from full pretrained model)") + self.beta_dist = Beta(config.noise_beta_alpha, config.noise_beta_beta) + # Video noise Beta distribution (biased towards high noise levels when enabled) + self.video_beta_dist = Beta(config.video_noise_beta_alpha, config.video_noise_beta_beta) + # High noise emphasis Beta distribution for coupled training (applies to both video and action) + self.high_noise_beta_dist = Beta(config.high_noise_beta_alpha, 1.0) + # self.num_timestep_buckets = config.num_timestep_buckets + self.config = config + self._noise_logged = False + self.defer_lora_injection = config.defer_lora_injection + print("defer_lora_injection@@", self.defer_lora_injection) + self.set_trainable_parameters(config.tune_projector, config.tune_diffusion_model) + + def set_trainable_parameters(self, tune_projector: bool, tune_diffusion_model: bool): + self.tune_projector = tune_projector + self.tune_diffusion_model = tune_diffusion_model + for p in self.parameters(): + p.requires_grad = True + if not tune_diffusion_model: + self.model.requires_grad_(False) + print(f"Tune action head projector: {self.tune_projector}") + print(f"Tune action head diffusion model: {self.tune_diffusion_model}") + # Check if any parameters are still trainable. If not, print a warning. + if not tune_projector and not tune_diffusion_model: + for name, p in self.named_parameters(): + if p.requires_grad: + print(f"Action head trainable parameter: {name}") + if not any(p.requires_grad for p in self.parameters()): + print("Warning: No action head trainable parameters found.") + + if self.train_architecture == "lora" and not self.defer_lora_injection: + print("Adding LoRA to model") + for p in self.parameters(): + p.requires_grad = False + self.model = self.add_lora_to_model( + self.model, + lora_rank=self.lora_rank, + lora_alpha=self.lora_alpha, + lora_target_modules=self.lora_target_modules, + init_lora_weights=self.init_lora_weights, + ) + self.model.state_encoder.requires_grad_(True) + self.model.action_encoder.requires_grad_(True) + self.model.action_decoder.requires_grad_(True) + elif self.train_architecture == "lora" and self.defer_lora_injection: + print("Deferring LoRA injection until after pretrained weights are loaded") + else: + self.print_trainable_params() + + self.text_encoder.requires_grad_(False) + self.image_encoder.requires_grad_(False) + self.vae.requires_grad_(False) + if not self.defer_lora_injection: + self.print_trainable_params() + + + def print_trainable_params(self): + """Print trainable parameters of the diffusion model.""" + trainable_params = [] + total_params = 0 + trainable_total = 0 + + for name, param in self.model.named_parameters(): + total_params += param.numel() + if param.requires_grad: + trainable_params.append(name) + trainable_total += param.numel() + + print(f"Total parameters in diffusion model: {total_params:,}") + print(f"Trainable parameters in diffusion model: {trainable_total:,}") + # print(trainable_params) + + + def inject_lora_after_loading(self): + """ + Inject LoRA adapters after pretrained weights have been loaded. + This should be called when defer_lora_injection=True. + """ + if self.train_architecture == "lora": + print("Injecting LoRA after loading pretrained weights") + for p in self.parameters(): + p.requires_grad = False + self.model = self.add_lora_to_model( + self.model, + lora_rank=self.lora_rank, + lora_alpha=self.lora_alpha, + lora_target_modules=self.lora_target_modules, + init_lora_weights=self.init_lora_weights, + ) + self.model.state_encoder.requires_grad_(True) + self.model.action_encoder.requires_grad_(True) + self.model.action_decoder.requires_grad_(True) + # self.model.registers.requires_grad_(True) + # self.model.time_modality_projection.requires_grad_(True) + + self.text_encoder.requires_grad_(False) + self.image_encoder.requires_grad_(False) + self.vae.requires_grad_(False) + self.print_trainable_params() + else: + print("LoRA injection not needed (train_architecture != 'lora')") + + def set_frozen_modules_to_eval_mode(self): + """ + Huggingface will call model.train() at each training_step. To ensure + the expected behaviors for modules like dropout, batchnorm, etc., we + need to call model.eval() for the frozen modules. + """ + if self.training: + if not self.tune_diffusion_model: + self.model.eval() + self.text_encoder.eval() + self.image_encoder.eval() + self.vae.eval() + + + def enable_vram_management(self, num_persistent_param_in_dit=None): + dtype = next(iter(self.text_encoder.parameters())).dtype + enable_vram_management( + self.text_encoder, + module_map = { + torch.nn.Linear: AutoWrappedLinear, + torch.nn.Embedding: AutoWrappedModule, + T5RelativeEmbedding: AutoWrappedModule, + T5LayerNorm: AutoWrappedModule, + }, + module_config = dict( + offload_dtype=dtype, + offload_device="cpu", + onload_dtype=dtype, + onload_device="cpu", + computation_dtype=self.dtype, + computation_device='cuda', + ), + ) + + self.cpu_offload = True + + def load_models_to_device(self, loadmodel_names=[]): + # only load models to device if cpu_offload is enabled + if not self.cpu_offload: + return + # offload the unneeded models to cpu + for model_name in self.model_names: + if model_name not in loadmodel_names: + model = getattr(self, model_name) + if model is not None: + if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: + print("offloadd") + for module in model.modules(): + if hasattr(module, "offload"): + # print("offload", module) + module.offload() + else: + print("tocpu") + model.cpu() + # load the needed models to device + for model_name in loadmodel_names: + model = getattr(self, model_name) + if model is not None: + if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: + print("onload") + for module in model.modules(): + if hasattr(module, "onload"): + # print("onload", module) + module.onload() + else: + print("togpu") + model.to(self._device) + # fresh the cuda cache + torch.cuda.empty_cache() + + def _create_kv_caches( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + frame_seqlen: int, + ) -> tuple[KVCacheType, KVCacheType]: + """ + Initialize a Per-GPU KV cache for the Wan model. + Use the model's num_heads and head_dim (5B has 24 heads, 14B has 40). + """ + num_heads = self.model.num_heads + head_dim = self.model.dim // num_heads + kv_cache1: KVCacheType = [] + kv_cache_neg: KVCacheType = [] + for _ in range(self.model.num_layers): + kv_cache1.append( + torch.zeros([2, batch_size, 0, num_heads, head_dim], dtype=dtype, device=device), + ) + kv_cache_neg.append( + torch.zeros([2, batch_size, 0, num_heads, head_dim], dtype=dtype, device=device), + ) + + return kv_cache1, kv_cache_neg + + def _create_crossattn_caches( + self, batch_size: int, dtype: torch.dtype, device: torch.device, + ) -> tuple[KVCacheType, KVCacheType]: + """ + Initialize a Per-GPU cross-attention cache for the Wan model. + Use the model's num_heads and head_dim (5B has 24 heads, 14B has 40). + """ + num_heads = self.model.num_heads + head_dim = self.model.dim // num_heads + crossattn_cache: KVCacheType = [] + crossattn_cache_neg: KVCacheType = [] + + for _ in range(self.model.num_layers): + crossattn_cache.append( + torch.zeros([2, batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + ) + crossattn_cache_neg.append( + torch.zeros([2, batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + ) + + return crossattn_cache, crossattn_cache_neg + + def sample_time(self, batch_size, device, dtype): + sample = self.beta_dist.sample([batch_size]).to(device, dtype=dtype) + return (self.config.noise_s - sample) / self.config.noise_s + + def prepare_input(self, batch: dict) -> BatchFeature: + return BatchFeature(data=batch) + + def preprocess_image(self, image): + image = (image * (2 / 255) - 1).permute(0, 1, 4, 2, 3) + return image + + def encode_prompt(self, input_ids, attention_mask): + seq_lens = attention_mask.gt(0).sum(dim=1).long() + prompt_emb = self.text_encoder(input_ids, attention_mask) + prompt_emb = prompt_emb.clone().to(dtype=torch.bfloat16) + for i, v in enumerate(seq_lens): + prompt_emb[:, v:] = 0 + return prompt_emb + + def _ensure_vae_on_device(self, ref_tensor): + """Lazily move the VAE to the correct device/dtype on first use.""" + if not getattr(self, '_vae_device_ready', False): + self.vae.to(device=ref_tensor.device, dtype=torch.bfloat16) + self.vae.eval() + self._vae_device_ready = True + + def encode_video(self, input_video, tiled=True, tile_size=(34, 34), tile_stride=(18, 16)): + self._ensure_vae_on_device(input_video) + with torch.amp.autocast(dtype=torch.bfloat16, device_type=torch.device(self._device).type): + with torch.no_grad(): + latents = self.vae.encode(input_video, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + return latents + + def encode_image(self, image, num_frames, height, width): + with torch.amp.autocast(dtype=torch.bfloat16, device_type=torch.device(self._device).type): + batch_size = image.shape[0] + clip_context = self.image_encoder.encode_image(image) + image_input = image.transpose(1, 2) + image_zeros = torch.zeros(batch_size, 3, num_frames-1, height, width, dtype=torch.bfloat16, device=self._device) + self._ensure_vae_on_device(image_input) + with torch.no_grad(): + y = self.vae.encode(torch.concat([image_input, image_zeros], dim=2)) + # Build mask to match VAE output shape (VAE may use different spatial downsampling, e.g. WanVideoVAE38 uses patch_size=2 -> height/16) + # y shape is B * 16 * (1+(T-1)/4) * H_latent * W_latent + num_t = y.shape[2] + h_latent, w_latent = y.shape[3], y.shape[4] + msk = torch.zeros(batch_size, 4, num_t, h_latent, w_latent, dtype=y.dtype, device=self._device) + msk[:, :, 0:1, :, :] = 1 + new_image = y[:, :, 0:1] + # concat: B * (4+16) * (1+(T-1)/4) * H_latent * W_latent + y = torch.concat([msk, y], dim=1) + return clip_context, y, new_image + + def prepare_extra_input(self, latents=None): + return {} + + def add_lora_to_model(self, model, lora_rank=4, lora_alpha=4, lora_target_modules="q,k,v,o,ffn.0,ffn.2", init_lora_weights="kaiming") -> nn.Module: + # Add LoRA to UNet + self.lora_alpha = lora_alpha + if init_lora_weights == "kaiming": + init_lora_weights = True + + lora_config = LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + init_lora_weights=init_lora_weights, + target_modules=lora_target_modules.split(","), + ) + model = get_peft_model(model, lora_config) + for param in model.parameters(): + param.data = param.to(torch.float32) + return model + + def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: + # Set frozen modules to eval + self.set_frozen_modules_to_eval_mode() + + data = action_input + # Get embodiment ID. + embodiment_id = action_input.embodiment_id + # print("embodiment_id", embodiment_id) + has_real_action = action_input.has_real_action + action_mask = action_input.action_mask + + state_features = action_input.state + + actions = action_input.action + # assert the values of action is in between -1 and 1 + if actions.numel() > 0: + assert actions.min() >= -1.0 and actions.max() <= 1.0, "actions must be in [-1,1] range" + videos = data["images"] + + videos = rearrange(videos, "b t h w c -> b c t h w") + print("videos", videos.shape) + + + if videos.dtype == torch.uint8: + videos = videos.float() / 255.0 + b, c, t, h, w = videos.shape + videos = videos.permute(0, 2, 1, 3, 4) # [b, t, c, h, w] + videos = videos.reshape(b * t, c, h, w) + videos = self.normalize_video(videos) + videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) # back to [b, c, t, h, w] + assert videos.min() >= -1.0 and videos.max() <= 1.0, "videos must be in [-1,1] range" + + videos = videos.to(dtype=self.dtype) + + # shape of B * max_length * dim + prompt_embs = self.encode_prompt(data["text"], data["text_attention_mask"]) + + # Wan 5B: resize to target resolution so latent tokens/frame matches DiT. Use config target when set + # (e.g. 160x320 so latent is 10x20 with VAE38 16x → even H,W, no crop in dynamics loss); else 176x320. + target_h = getattr(self.config, "target_video_height", None) + target_w = getattr(self.config, "target_video_width", None) + if target_h is None or target_w is None: + if getattr(self.model, "frame_seqlen", None) in (50, 55): + target_h, target_w = 176, 320 + else: + target_h, target_w = None, None + if target_h is not None and target_w is not None: + _, _, _, h, w = videos.shape + if (h, w) != (target_h, target_w): + b, c, t, _, _ = videos.shape + videos = torch.nn.functional.interpolate( + videos.reshape(b * t, c, h, w), + size=(target_h, target_w), + mode="bilinear", + align_corners=False, + ).reshape(b, c, t, target_h, target_w) + + latents = self.encode_video(videos, self.tiled, (self.tile_size_height, self.tile_size_width), (self.tile_stride_height, self.tile_stride_width)) + + # print("latents shape", latents.shape, self.dtype) + _, _, num_frames, height, width = videos.shape + image = videos[:, :, :1].transpose(1, 2) + + clip_feas, ys, _ = self.encode_image(image, num_frames, height, width) + + latents = latents.to(self._device) + clip_feas = clip_feas.to(self._device) + ys = ys.to(self._device) + prompt_embs = prompt_embs.to(self._device) + + # Loss + noise = torch.randn_like(latents) + + # specific to autoregressive + noise = noise.transpose(1, 2) + latents = latents.transpose(1, 2) + + # ============ VIDEO TIMESTEP SAMPLING ============ + if self.config.decouple_video_action_noise: + # Decoupled mode: sample video from Beta distribution biased towards HIGH noise + video_noise_ratio = self.video_beta_dist.sample([noise.shape[0], noise.shape[1]]) + timestep_id = ((1.0 - video_noise_ratio) * self.scheduler.num_train_timesteps).long() + timestep_id = torch.clamp(timestep_id, 0, self.scheduler.num_train_timesteps - 1) + noise_mode = "DECOUPLED" + elif self.config.use_high_noise_emphasis: + # High noise emphasis mode (coupled): BOTH video and action use Beta distribution + noise_ratio = self.high_noise_beta_dist.sample([noise.shape[0], noise.shape[1]]) + timestep_id = ((1.0 - noise_ratio) * self.scheduler.num_train_timesteps).long() + timestep_id = torch.clamp(timestep_id, 0, self.scheduler.num_train_timesteps - 1) + noise_mode = "HIGH_NOISE_EMPHASIS" + else: + # Original: uniform sampling over full range + timestep_id = torch.randint(0, self.scheduler.num_train_timesteps, (noise.shape[0], noise.shape[1])) + noise_mode = "STANDARD" + + timestep_id_block = timestep_id[:, 1:].reshape( + timestep_id.shape[0], -1, self.num_frame_per_block) + timestep_id_block[:, :, 1:] = timestep_id_block[:, :, 0:1] + + if actions.numel() > 0: + noise_action = torch.randn_like(actions) + assert actions.shape[1] / (noise.shape[1]-1) == (self.model.num_action_per_block // self.num_frame_per_block), f"actions.shape, {actions.shape}, noise.shape, {noise.shape}, video.shape, {videos.shape}, latents.shape, {latents.shape}" + assert (noise.shape[1]-1) / state_features.shape[1] == (self.num_frame_per_block // self.model.num_state_per_block), f"state_features.shape, {state_features.shape}, noise.shape, {noise.shape}, video.shape, {videos.shape}, latents.shape, {latents.shape}" + + # ============ ACTION TIMESTEP SAMPLING ============ + if self.config.decouple_video_action_noise: + # Decoupled: sample action timestep independently with full range + timestep_action_id = torch.randint( + 0, + self.scheduler.num_train_timesteps, + (actions.shape[0], actions.shape[1]) + ) + action_mode = "INDEPENDENT" + else: + # Original coupled: action timestep derived from video timestep + timestep_action_id = timestep_id_block.repeat(1, 1, actions.shape[1]//(noise.shape[1]-1)) + timestep_action_id = timestep_action_id.reshape(timestep_action_id.shape[0], -1) + action_mode = "COUPLED" + + # Log noise mode once + if not self._noise_logged: + video_mean = timestep_id.float().mean().item() + action_mean = timestep_action_id.float().mean().item() + if noise_mode == "DECOUPLED": + print(f"[NOISE] Mode={noise_mode} | Video: Beta({self.config.video_noise_beta_alpha},1) mean_t={video_mean:.0f} | Action: {action_mode} Uniform mean_t={action_mean:.0f}") + elif noise_mode == "HIGH_NOISE_EMPHASIS": + print(f"[NOISE] Mode={noise_mode} | Video+Action: Beta({self.config.high_noise_beta_alpha},1) mean_t={video_mean:.0f} | Action: {action_mode}") + else: + print(f"[NOISE] Mode={noise_mode} | Video+Action: Uniform mean_t={video_mean:.0f} | Action: {action_mode}") + self._noise_logged = True + else: + noise_action = None + timestep_action_id = None + + timestep_id_block = timestep_id_block.reshape(timestep_id_block.shape[0], -1) + timestep_id = torch.concat([timestep_id[:, :1], timestep_id_block], dim=1) + _, num_frames, num_channels, height, width = noise.shape + # DiT patch_embedding uses stride (1,2,2), so sequence length is num_frames * (H//2) * (W//2) + tokens_per_frame = (height // 2) * (width // 2) + seq_len = num_frames * tokens_per_frame + + timestep = self.scheduler.timesteps[timestep_id].to(self._device) + noisy_latents = self.scheduler.add_noise(latents.flatten(0, 1), noise.flatten(0, 1), timestep.flatten(0, 1)).unflatten(0, (noise.shape[0], noise.shape[1])) + training_target = self.scheduler.training_target(latents, noise, timestep).transpose(1, 2) + + if actions.numel() > 0: + timestep_action = self.scheduler.timesteps[timestep_action_id].to(self._device) + noisy_actions = self.scheduler.add_noise( + actions.flatten(0, 1), + noise_action.flatten(0, 1), + timestep_action.flatten(0, 1), + ).unflatten(0, (noise_action.shape[0], noise_action.shape[1])) + training_target_action = self.scheduler.training_target(actions, noise_action, timestep_action) + else: + timestep_action = None + noisy_actions = None + training_target_action = None + + # Compute loss + with torch.amp.autocast(dtype=torch.bfloat16, device_type=torch.device(self._device).type): + if actions.numel() > 0: + video_noise_pred, action_noise_pred = self.model( + noisy_latents.transpose(1, 2), timestep=timestep, clip_feature=clip_feas, y=ys, context=prompt_embs, seq_len=seq_len, + state=state_features, embodiment_id=embodiment_id, + action=noisy_actions, timestep_action=timestep_action, + clean_x=latents.transpose(1, 2), + ) + else: + video_noise_pred, action_noise_pred = self.model( + noisy_latents.transpose(1, 2), timestep=timestep, timestep_action=timestep_action, + clip_feature=clip_feas, y=ys, context=prompt_embs, seq_len=seq_len, + state=state_features, embodiment_id=embodiment_id, + clean_x=latents.transpose(1, 2), + ) + + # Per-sample dynamics loss + # DiT patch_embedding uses stride (1,2,2), so output spatial size can be smaller than + # latent when H or W is odd (e.g. latent 11x20 -> model output 10x20). Crop target to match. + if training_target.shape != video_noise_pred.shape: + training_target = training_target[ + ..., : video_noise_pred.shape[3], : video_noise_pred.shape[4] + ] + dynamics_loss_per_sample = torch.nn.functional.mse_loss( + video_noise_pred.float(), training_target.float(), reduction='none' + ).mean(dim=(1,3,4)) # shape: [B, ...] + + weight_dynamics = dynamics_loss_per_sample * self.scheduler.training_weight(timestep.flatten(0, 1)).unflatten(0, (noise.shape[0], noise.shape[1])).to(self._device) + weighted_dynamics_loss = weight_dynamics.mean() + + if actions.numel() > 0: + action_loss_per_sample = torch.nn.functional.mse_loss( + action_noise_pred.float(), training_target_action.float(), reduction='none' + ) * action_mask # shape: [B, ...] + action_loss_per_sample = has_real_action[:, None].float() * action_loss_per_sample # apply has_real_action + weight_action = action_loss_per_sample.mean(dim=2) * self.scheduler.training_weight( + timestep_action.flatten(0, 1), + ).unflatten(0, (noise_action.shape[0], noise_action.shape[1])).to(self._device) + weighted_action_loss = weight_action.mean() + loss = weighted_dynamics_loss + weighted_action_loss + else: + weighted_action_loss = torch.tensor(0.0, device=self._device) + loss = weighted_dynamics_loss + # loss = dynamics_loss_per_sample.mean() + + # Record log + output_dict = { + "loss": loss, + "dynamics_loss": weighted_dynamics_loss, + "action_loss": weighted_action_loss, + } + + return BatchFeature(data=output_dict) + + def generate_noise(self, shape, seed=None, device="cpu", dtype=torch.float16): + generator = None if seed is None else torch.Generator(device).manual_seed(seed) + noise = torch.randn(shape, generator=generator, device=device, dtype=dtype) + return noise + + def _get_caches( + self, kv_caches_input: list[KVCacheType], + ) -> list[KVCacheType]: + if self.ip_size > 1: + assert self.cfg_scale != 1.0, "cfg_scale must be != 1.0 when ip_size > 1" + assert len(kv_caches_input) == 2 + if self.ip_rank == 0: + kv_caches = [kv_caches_input[0]] + else: + kv_caches = [kv_caches_input[1]] + else: + assert len(kv_caches_input) <= 2 + kv_caches = [kv_caches_input[0]] + if self.cfg_scale != 1.0: + kv_caches.append(kv_caches_input[1]) + return kv_caches + + def _prepare_text_inputs(self, data: BatchFeature) -> list[tuple[torch.Tensor, torch.Tensor]]: + + if self.ip_size > 1: + assert self.cfg_scale != 1.0, "cfg_scale must be != 1.0 when ip_size > 1" + if self.ip_rank == 0: + text_inputs = [(data["text"], data["text_attention_mask"])] + else: + text_inputs = [(data["text_negative"], data["text_attention_mask_negative"])] + else: + text_inputs = [(data["text"], data["text_attention_mask"])] + if self.cfg_scale != 1.0: + text_inputs.append((data["text_negative"], data["text_attention_mask_negative"])) + return text_inputs + + + def _run_diffusion_steps( + self, + noisy_input: torch.Tensor, + timestep: torch.Tensor, + action: torch.Tensor, + timestep_action: torch.Tensor, + state: torch.Tensor, + embodiment_id: torch.Tensor, + context: torch.Tensor, + seq_len: int, + y: torch.Tensor, + clip_feature: torch.Tensor, + kv_caches: list[KVCacheType], + crossattn_caches: list[KVCacheType], + kv_cache_metadata: dict[str, bool | int], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + predictions = [] + for index, prompt_emb in enumerate(context): + kv_cache = kv_caches[index] + crossattn_cache = crossattn_caches[index] + if not kv_cache_metadata["update_kv_cache"] and self.trt_engine is not None: + obs_noise_pred, action_noise_pred = self.trt_engine( + noisy_input, + timestep, + action=action, + timestep_action=timestep_action, + state=state, + context=prompt_emb, + y=y, + clip_feature=clip_feature, + kv_cache=kv_cache, + ) + else: + obs_noise_pred, action_noise_pred, updated_kv_caches = self.model( + noisy_input, + timestep, + action=action, + timestep_action=timestep_action, + state=state, + embodiment_id=embodiment_id, + context=prompt_emb, + seq_len=seq_len, + y=y, + clip_feature=clip_feature, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start_frame=kv_cache_metadata["start_frame"], + ) + if kv_cache_metadata["update_kv_cache"]: + for block_index, updated_kv_cache in enumerate(updated_kv_caches): + kv_cache[block_index] = updated_kv_cache.clone() + obs_noise_pred = obs_noise_pred.clone() + if action_noise_pred is not None: + action_noise_pred = action_noise_pred.clone() + else: + action_noise_pred = torch.tensor(0.0, device=obs_noise_pred.device) # dummy action noise prediction + predictions.append((obs_noise_pred, action_noise_pred)) + return self._exchange_predictions(predictions) + + def _exchange_predictions( + self, + predictions: list[tuple[torch.Tensor, torch.Tensor]], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + if self.ip_size == 1: + return predictions + + assert len(predictions) == 1 + my_predictions = list(predictions[0]) + + other_predictions = [torch.empty_like(pred) for pred in my_predictions] + + send_ops = [ + dist.P2POp(op=dist.isend, tensor=pred, group_peer=(self.ip_rank + 1) % self.ip_size, group=self.ip_group) + for pred in my_predictions + ] + recv_ops = [ + dist.P2POp(op=dist.irecv, tensor=other_pred, group_peer=(self.ip_rank + 1) % self.ip_size, group=self.ip_group) + for other_pred in other_predictions + ] + ops = send_ops + recv_ops + + reqs = dist.batch_isend_irecv(ops) + for req in reqs: + req.wait() + + output_predictions: list[tuple[torch.Tensor, torch.Tensor] | None] = [None for _ in range(self.ip_size)] + output_predictions[self.ip_rank] = tuple(my_predictions) + output_predictions[(self.ip_rank + 1) % self.ip_size] = tuple(other_predictions) + assert all(isinstance(pred, tuple) for pred in output_predictions) + return cast(list[tuple[torch.Tensor, torch.Tensor]], output_predictions) + + def should_run_model(self, index, current_timestep, prev_predictions): + + if not self.dynamic_cache_schedule: + return self.dit_step_mask[index] + + # Always run first 2 steps to establish history + if len(prev_predictions) < 2: + return True + + if self.skip_countdown > 1: + self.skip_countdown -= 1 + return False + elif self.skip_countdown == 1: + self.skip_countdown = 0 + return True + + v_last = prev_predictions[-1][1].flatten(1).float() + v_prev = prev_predictions[-2][1].flatten(1).float() + sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean() + + thresholds = [0.95, 0.93] + countdowns = [4, 2] + + for threshold, countdown in zip(thresholds, countdowns): + if sim > threshold: + self.skip_countdown = countdown + return False + + return True + + def lazy_joint_video_action(self, backbone_output: BatchFeature, action_input: BatchFeature, latent_video: torch.Tensor | None = None) -> BatchFeature: + start_time = time.perf_counter() + + # Tracking time taken on GPU for various operations. + start_text_encoder_event = torch.cuda.Event(enable_timing=True) + end_text_encoder_event = torch.cuda.Event(enable_timing=True) + start_image_encoder_event = torch.cuda.Event(enable_timing=True) + end_image_encoder_event = torch.cuda.Event(enable_timing=True) + start_vae_event = torch.cuda.Event(enable_timing=True) + end_vae_event = torch.cuda.Event(enable_timing=True) + start_kv_event = torch.cuda.Event(enable_timing=True) + end_kv_event = torch.cuda.Event(enable_timing=True) + start_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in range(self.num_inference_steps)] + end_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in range(self.num_inference_steps)] + + self.set_frozen_modules_to_eval_mode() + data = action_input + + videos = data["images"] + + embodiment_id = action_input.embodiment_id + state_features = action_input.state + + videos = rearrange(videos, "b t h w c -> b c t h w") + + if videos.dtype == torch.uint8: + videos = videos.float() / 255.0 + videos = videos.to(dtype=self.dtype) + b, c, t, h, w = videos.shape + videos = videos.permute(0, 2, 1, 3, 4) # [b, t, c, h, w] + videos = videos.reshape(b * t, c, h, w) + videos = self.normalize_video(videos) + videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) # back to [b, c, t, h, w] + assert videos.min() >= -1.0 and videos.max() <= 1.0, "videos must be in [-1,1] range" + videos = videos.to(dtype=self.dtype) + + state_features = state_features.to(dtype=torch.bfloat16) + videos = videos.to(dtype=torch.bfloat16) + + # Wan 5B: same as training — resize to target resolution so latent matches DiT + target_h = getattr(self.config, "target_video_height", None) + target_w = getattr(self.config, "target_video_width", None) + if target_h is None or target_w is None: + if getattr(self.model, "frame_seqlen", None) in (50, 55): + target_h, target_w = 176, 320 + else: + target_h, target_w = None, None + if target_h is not None and target_w is not None: + _, _, _, h, w = videos.shape + if (h, w) != (target_h, target_w): + b, c, t, _, _ = videos.shape + videos = torch.nn.functional.interpolate( + videos.reshape(b * t, c, h, w), + size=(target_h, target_w), + mode="bilinear", + align_corners=False, + ).reshape(b, c, t, target_h, target_w) + + if self.language is None: + print("language is None, reset current_start_frame to 0") + self.language = data["text"] + self.current_start_frame = 0 + elif not torch.equal(self.language, data["text"]): + print("language changed, reset current_start_frame to 0") + self.current_start_frame = 0 + self.language = data["text"] + elif videos.shape[2] == 1: + print("videos.shape[2] == 1, reset current_start_frame to 0") + self.current_start_frame = 0 + elif self.current_start_frame >= self.model.local_attn_size: + print("current_start_frame >= local_attn_size, reset current_start_frame to 0") + self.current_start_frame = 0 + + if self.ip_rank == 0: + print("videos shape", videos.shape, self.num_frames) + + start_text_encoder_event.record() + + text_inputs = self._prepare_text_inputs(data) + prompt_embs = [self.encode_prompt(text, attention_mask) for text, attention_mask in text_inputs] + + end_text_encoder_event.record() + + start_image_encoder_event.record() + + _, _, num_frames, height, width = videos.shape + if videos.shape[2] == 4 or videos.shape[2] == 9: + # special case for real-world eval where language is updated + image = videos[:, :, -1:].transpose(1, 2) + else: + image = videos[:, :, :1].transpose(1, 2) + + if self.current_start_frame == 0: + clip_feas, ys, image = self.encode_image(image, self.num_frames, height, width) + self.clip_feas = clip_feas.to(dtype=image.dtype) + self.ys = ys.to(dtype=image.dtype) + + assert self.clip_feas is not None and self.ys is not None, "clip_feas and ys must be set" + + end_image_encoder_event.record() + + start_vae_event.record() + + if latent_video is not None and self.current_start_frame != 0: + image = latent_video + if self.ip_rank == 0: + print("image shape@@", image.shape) + elif self.current_start_frame != 0: + # this is for real world execution + if (videos.shape[2] - 1) // 4 == self.num_frame_per_block: + print("no further action") + elif videos.shape[2] // 4 != self.num_frame_per_block: + # Repeating videos along dim 2. + repeat_factor = self.num_frame_per_block // (videos.shape[2] // 4) + videos = torch.repeat_interleave(videos, repeat_factor, dim=2) + + first_frame = videos[:, :, 0:1] # Extract first frame + videos = torch.cat([first_frame, videos], dim=2) + else: + first_frame = videos[:, :, 0:1] # Extract first frame + videos = torch.cat([first_frame, videos], dim=2) + + image = self.vae.encode( + videos, + tiled=self.tiled, + tile_size=(self.tile_size_height, self.tile_size_width), + tile_stride=(self.tile_stride_height, self.tile_stride_width), + ) + + end_vae_event.record() + + noise_obs = self.generate_noise((image.shape[0], image.shape[1], self.num_frame_per_block, image.shape[3], image.shape[4]), seed=self.seed, device='cuda', dtype=torch.bfloat16) + noise_action = self.generate_noise((image.shape[0], self.action_horizon, self.model.action_dim), seed=self.seed, device='cuda', dtype=torch.bfloat16) + batch_size, num_channels, num_frames, height, width = noise_obs.shape + ######### Generate video ######### + # DiT patch_embedding uses stride (1,2,2), so tokens per frame = (H//2)*(W//2) + tokens_per_frame = (height // 2) * (width // 2) + frame_seqlen = tokens_per_frame + seq_len = num_frames * frame_seqlen + + image = image.transpose(1, 2) + noise_obs = noise_obs.transpose(1, 2) + + if self.current_start_frame == 0: + # Reinitialize KV cache and crossattn cache for each new sequence. + self.kv_cache1, self.kv_cache_neg = self._create_kv_caches( + batch_size=batch_size, + dtype=noise_obs.dtype, + device=noise_obs.device, + frame_seqlen=frame_seqlen, + ) + self.crossattn_cache, self.crossattn_cache_neg = self._create_crossattn_caches( + batch_size=batch_size, + dtype=noise_obs.dtype, + device=noise_obs.device, + ) + + assert self.kv_cache1 is not None + assert self.kv_cache_neg is not None + assert self.crossattn_cache is not None + assert self.crossattn_cache_neg is not None + kv_caches = self._get_caches( + [self.kv_cache1, self.kv_cache_neg], + ) + crossattn_caches = self._get_caches( + [self.crossattn_cache, self.crossattn_cache_neg], + ) + + start_kv_event.record() + + if self.current_start_frame == 0: + timestep = torch.ones([batch_size, 1], device=noise_obs.device, dtype=torch.int64) * 0 + self._run_diffusion_steps( + noisy_input=image.transpose(1, 2), + timestep=timestep * 0, + action=None, + timestep_action=None, + state=None, + embodiment_id=None, + context=prompt_embs, + seq_len=frame_seqlen, + y=self.ys[:, :, 0:1], + clip_feature=self.clip_feas, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + kv_cache_metadata=dict( + start_frame=0, + update_kv_cache=True, + ), + ) + self.current_start_frame += 1 + + timestep = torch.ones([batch_size, self.num_frame_per_block], device=noise_obs.device, dtype=torch.int64) * 0 + + if self.current_start_frame != 1: + current_ref_latents = image[:, -self.num_frame_per_block:] + if self.current_start_frame <= self.ys.shape[2]: + y = self.ys[:, :, self.current_start_frame - self.num_frame_per_block : self.current_start_frame] + else: + y = self.ys[:, :, -self.num_frame_per_block:] + self._run_diffusion_steps( + noisy_input=current_ref_latents.transpose(1, 2), + timestep=timestep * 0, + action=None, + timestep_action=None, + state=None, + embodiment_id=None, + context=prompt_embs, + seq_len=seq_len, + y=y, + clip_feature=self.clip_feas, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + kv_cache_metadata=dict( + start_frame=self.current_start_frame - self.num_frame_per_block, + update_kv_cache=True, + ), + ) + + end_kv_event.record() + + noisy_input = noise_obs + noisy_input_action = noise_action + + # Step 3.1: Spatial denoising loop + + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=self.scheduler.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler_action = FlowUniPCMultistepScheduler( + num_train_timesteps=self.scheduler.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler.set_timesteps( + self.num_inference_steps, device=noise_obs.device, shift=self.sigma_shift) + sample_scheduler_action.set_timesteps( + self.num_inference_steps, device=noise_obs.device, shift=self.sigma_shift) + + # Decoupled inference: video sigmas end at video_final_noise instead of 0 + # This rescales the schedule so video still takes all denoising steps, + # but ends at a higher noise level (e.g., 1.0 → 0.9 → 0.8 instead of 1.0 → 0.5 → 0.0) + if self.config.decouple_inference_noise: + video_final_noise = self.config.video_inference_final_noise + # Rescale video sigmas: map [sigma_max, 0] -> [sigma_max, video_final_noise] + sigma_max = sample_scheduler.sigmas[0].item() + sample_scheduler.sigmas = sample_scheduler.sigmas * (sigma_max - video_final_noise) / sigma_max + video_final_noise + sample_scheduler.timesteps = (sample_scheduler.sigmas[:-1] * 1000).to(torch.int64) + if self.ip_rank == 0: + print(f"Decoupled inference: video sigmas {sigma_max:.3f} -> {sample_scheduler.sigmas[-1].item():.3f}") + + start_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in sample_scheduler.timesteps] + end_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in sample_scheduler.timesteps] + prev_predictions = [] + self.skip_countdown = 0 + dit_compute_steps = 0 + for index, current_timestep in enumerate(sample_scheduler.timesteps): + start_diffusion_events[index].record() + + # Get timesteps from respective schedulers + action_timestep = sample_scheduler_action.timesteps[index] + video_timestep = sample_scheduler.timesteps[index] # Already rescaled if decoupled + + # set current timestep + timestep = torch.ones( + [batch_size, self.num_frame_per_block], + device=noise_obs.device, + dtype=torch.int64, + ) * video_timestep + timestep_action = torch.ones( + [batch_size, self.action_horizon], + device=noise_obs.device, + dtype=torch.int64, + ) * action_timestep + + # check if we need to run the DIT step + should_run_model = self.should_run_model(index, current_timestep, prev_predictions) + if should_run_model: + dit_compute_steps += 1 + if self.current_start_frame + self.num_frame_per_block <= self.ys.shape[2]: + y = self.ys[:, :, self.current_start_frame : self.current_start_frame + self.num_frame_per_block] + else: + y = self.ys[:, :, -self.num_frame_per_block:] + predictions = self._run_diffusion_steps( + noisy_input=noisy_input.transpose(1, 2), + timestep=timestep, + action=noisy_input_action, + timestep_action=timestep_action, + state=state_features, + embodiment_id=embodiment_id, + context=prompt_embs, + seq_len=seq_len, + y=y, + clip_feature=self.clip_feas, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + kv_cache_metadata=dict( + start_frame=self.current_start_frame, + update_kv_cache=False, + ), + ) + flow_pred_cond, flow_pred_cond_action = predictions[0] + flow_pred_uncond, flow_pred_uncond_action = predictions[1] + + flow_pred = flow_pred_uncond + self.cfg_scale * (flow_pred_cond - flow_pred_uncond) + prev_predictions.append((current_timestep, flow_pred, flow_pred_cond_action)) + max_cache_size = 2 + if len(prev_predictions) > max_cache_size: + prev_predictions.pop(0) + + else: + assert len(prev_predictions) > 0, "prev_predictions must be set when skipping" + _, flow_pred, flow_pred_cond_action = prev_predictions[-1] + + end_diffusion_events[index].record() + + # Video: denoising step (uses rescaled schedule if decoupled) + noisy_input = sample_scheduler.step( + model_output=flow_pred.transpose(1, 2), + timestep=video_timestep, + sample=noisy_input, + step_index=index, + return_dict=False, + )[0] + + # Action: always fully denoises with standard schedule (1000->0) + noisy_input_action = sample_scheduler_action.step( + model_output=flow_pred_cond_action, + timestep=action_timestep, + sample=noisy_input_action, + step_index=index, + return_dict=False, + )[0] + + latents = noisy_input + latents_action = noisy_input_action + output = latents + + if self.current_start_frame == 1: + output = torch.cat([image, output], dim=1) + self.current_start_frame += self.num_frame_per_block + + # Do torch.cuda.synchronize() to ensure all operations are completed before timing. + # This isn't expected to affect inference performance since it's at the end of an inference step. + torch.cuda.synchronize() + + total_time = time.perf_counter() - start_time + text_encoder_time = start_text_encoder_event.elapsed_time(end_text_encoder_event) / 1000 + image_encoder_time = start_image_encoder_event.elapsed_time(end_image_encoder_event) / 1000 + vae_time = start_vae_event.elapsed_time(end_vae_event) / 1000 + kv_creation_time = start_kv_event.elapsed_time(end_kv_event) / 1000 + diffusion_times = [s.elapsed_time(e) for s, e in zip(start_diffusion_events, end_diffusion_events)] + diffusion_time = sum(diffusion_times) / 1000 + scheduler_time = total_time - kv_creation_time - diffusion_time - text_encoder_time - image_encoder_time - vae_time + + if self.ip_rank == 0: + print(f"Time taken: Total {total_time:.2f} seconds, " + f"Text Encoder {text_encoder_time:.2f} seconds, " + f"Image Encoder {image_encoder_time:.2f} seconds, " + f"VAE {vae_time:.2f} seconds, " + f"KV Cache Creation {kv_creation_time:.2f} seconds, " + f"Diffusion {diffusion_time:.2f} seconds, " + f"DIT Compute Steps {dit_compute_steps} steps, " + f"Scheduler {scheduler_time:.2f} seconds") + + return BatchFeature(data={"action_pred": latents_action, "video_pred": output.transpose(1, 2)}) + + def cache_predict_order1(self, current_timestep, timestep_1, f1, timestep_2, f2): + h_curr = current_timestep - timestep_1 + h_past = timestep_1 - timestep_2 + + v_prime = (f1 - f2) / h_past + + # Prediction + damping_factor = 0.25 + flow_pred = f1 + (v_prime * h_curr) * damping_factor + return flow_pred + + def post_initialize(self): + # Move models to the cuda device and set the dtype to bfloat16. + print("Moving models to the cuda device and setting the dtype to bfloat16.") + self.model.to(device=self._device, dtype=torch.bfloat16) + self.text_encoder.to(device=self._device, dtype=torch.bfloat16) + self.image_encoder.to(device=self._device, dtype=torch.bfloat16) + self.vae.to(device=self._device, dtype=torch.bfloat16) + import os + ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" + LOAD_TRT_ENGINE = os.getenv("LOAD_TRT_ENGINE", None) + + # Torch compile the modules. Skip _forward_blocks: Dynamo with fullgraph can fail on + # shape variation (e.g. x [1,50,C] vs e [1,200,C]); the block aligns e to x at runtime. + if not ENABLE_TENSORRT: + print("Torch compiling the TextEncoder, ImageEncoder, and VAE modules (Wan _forward_blocks not compiled).") + + self.text_encoder.forward = torch.compile( + mode="reduce-overhead", fullgraph=True, dynamic=False, + )(self.text_encoder.forward) + + self.image_encoder.model.visual.forward = torch.compile( + mode="reduce-overhead", fullgraph=True, dynamic=False, + )(self.image_encoder.model.visual.forward) + + self.vae.model.encode = torch.compile( + mode="reduce-overhead", fullgraph=True, dynamic=False, + )(self.vae.model.encode) + + self.trt_engine = None + if LOAD_TRT_ENGINE is not None: + print(f"Loading TRT engine from {LOAD_TRT_ENGINE}") + import groot.control.tensorrt_utils as trt_utils + model_path = LOAD_TRT_ENGINE + self.trt_engine = trt_utils.load_tensorrt_engine(model_path, model_type="ar_14B") + + def parallelize(self, device_mesh: DeviceMesh) -> None: + ip_mesh = device_mesh["ip"] + self.ip_rank = ip_mesh.get_local_rank() + self.ip_size = ip_mesh.size() + self.ip_group = ip_mesh.get_group() + + assert self.ip_size == 1 or self.ip_size == 2, "ip_size must be 1 or 2" + assert self.ip_rank >= 0 and self.ip_rank < self.ip_size, "ip_rank must be in [0, ip_size)" + + @property + def device(self): + return next(iter(self.parameters())).device + + @property + def dtype(self): + return next(iter(self.parameters())).dtype diff --git a/groot/vla/model/dreamzero/backbone/__init__.py b/groot/vla/model/dreamzero/backbone/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/backbone/base_backbone.py b/groot/vla/model/dreamzero/backbone/base_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..b768bcdef5fa0e6b08ee4634ec2e144e549edbb7 --- /dev/null +++ b/groot/vla/model/dreamzero/backbone/base_backbone.py @@ -0,0 +1,16 @@ +from abc import ABC, abstractmethod + +from torch import nn +from transformers.feature_extraction_utils import BatchFeature + + +class Backbone(ABC, nn.Module): + def __init__(self): + super(Backbone, self).__init__() + + @abstractmethod + def forward(self, backbone_input: BatchFeature) -> BatchFeature: + pass + + def prepare_input(self, batch: dict) -> BatchFeature: + pass diff --git a/groot/vla/model/dreamzero/backbone/identity.py b/groot/vla/model/dreamzero/backbone/identity.py new file mode 100644 index 0000000000000000000000000000000000000000..d63ad24139b9b796f0244025618f472e2fe0235a --- /dev/null +++ b/groot/vla/model/dreamzero/backbone/identity.py @@ -0,0 +1,49 @@ +import torch +from transformers.feature_extraction_utils import BatchFeature + +from groot.vla.model.dreamzero.backbone.base_backbone import Backbone + + +class IdentityBackbone(Backbone): + """ + This class allows pretraining the action head without depending on any backbone. + That's why it's called "identity" — it preserves the action head to be a standalone trainable model. + """ + + def set_trainable_parameters(self, **kwargs): + return + + def forward(self, backbone_input: BatchFeature) -> BatchFeature: + backbone_input_first_value = next(iter(backbone_input.values())) + B = backbone_input_first_value.shape[0] + + backbone_features = torch.empty( + B, 1, 0, dtype=torch.float32, device=backbone_input_first_value.device + ) + output_dict = { + "backbone_features": backbone_features, + } + + return BatchFeature(data=output_dict) + + def prepare_input(self, batch: dict) -> BatchFeature: + """ + Args: + batch: dict + Must contain at least one key-value pair to inform the batch size. + Expects the first dimension to be the batch size. See `forward`. + """ + if "action" in batch: + return BatchFeature(data={"action": batch["action"]}) + else: + # at inference time, we have to use either state or video + if "state" in batch: + return BatchFeature(data={"state": batch["state"]}) + elif "video" in batch: + # For video, it's tricky because it's a numpy array, which isn't compatible with BatchFeature's `to` method + # So instead, we make it a tensor and return it + video = batch["video"] + video_tensor = torch.from_numpy(video) + return BatchFeature(data={"video": video_tensor}) + else: + return BatchFeature(data=batch) diff --git a/groot/vla/model/dreamzero/base_vla.py b/groot/vla/model/dreamzero/base_vla.py new file mode 100644 index 0000000000000000000000000000000000000000..274b8640f44b2ac06b7fe8f60f43b4ce8dd51a28 --- /dev/null +++ b/groot/vla/model/dreamzero/base_vla.py @@ -0,0 +1,614 @@ +from dataclasses import dataclass, field +from typing import Tuple + +from hydra.utils import instantiate +import numpy as np +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh + +from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel +from transformers.feature_extraction_utils import BatchFeature +import tree + +BACKBONE_FEATURE_KEY = "backbone_features" +ACTION_KEY = "action_pred" +LOSS_KEY = "loss" +ERROR_MSG = "Error: unexpected input/output" +N_COLOR_CHANNELS = 3 + + +@dataclass +class VLAConfig(PretrainedConfig): + model_type = "vla" + backbone_cfg: PretrainedConfig = field( + default=None, metadata={"help": "Backbone configuration."} + ) + + action_head_cfg: PretrainedConfig = field( + default=None, metadata={"help": "Action head configuration."} + ) + + action_horizon: int = field(default=None, metadata={"help": "Action horizon."}) + + action_dim: int = field(default=None, metadata={"help": "Action dimension."}) + compute_dtype: str = field(default="float32", metadata={"help": "Compute dtype."}) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + for key, value in kwargs.items(): + setattr(self, key, value) + + +class VLA(PreTrainedModel): + supports_gradient_checkpointing = True + config_class = VLAConfig + """ + we expect the backbone output to have a key 'backbone_features' with shape (batch_size, n, hidden_size) + here n is variable and can be e.g. time, 1 or user specified + we expect the action head output to have a key 'action_pred' with shape (batch_size, time, action_dim) during inference time + we expect these to have type BatchFeature, and they can of course have many other user specified keys too + see discussion at https://nvidia.slack.com/archives/C07T1V7L886/p1732550624654139 + """ + + def __init__( + self, + config: VLAConfig, + ): + assert isinstance(config.backbone_cfg, dict) + assert isinstance(config.action_head_cfg, dict) + super().__init__(config) + self.backbone = instantiate(config.backbone_cfg) + self.action_head = instantiate(config.action_head_cfg) + self.action_horizon = config.action_horizon + self.action_dim = config.action_dim + self.compute_dtype = config.compute_dtype + + self.rank = dist.get_rank() if dist.is_initialized() else 0 + + def validate_inputs(self, inputs): + detected_error = False + error_msg = ERROR_MSG + if "action" in inputs: + action = inputs["action"] + type_ok = isinstance(action, torch.Tensor) + shape_ok = ( + len(action.shape) == 3 + and action.shape[1] % self.action_horizon == 0 + and action.shape[2] == self.action_dim + ) + if not type_ok: + error_msg += f"\n{action.dtype=}" + detected_error = True + if not shape_ok: + error_msg += f"\n{action.shape=}" + detected_error = True + + if "video" in inputs: + video = inputs["video"] + type_ok = isinstance(video, np.ndarray) + dtype_ok = video.dtype == np.uint8 + shape_ok = len(video.shape) == 6 and video.shape[3] == N_COLOR_CHANNELS + if not type_ok: + error_msg += f"\n{type(video)=}" + detected_error = True + if not dtype_ok: + error_msg += f"\n{video.dtype=}" + detected_error = True + if not shape_ok: + error_msg += f"\n{video.shape=}" + detected_error = True + + if detected_error: + raise ValueError(error_msg) + + def validate_data(self, action_head_outputs, backbone_outputs, is_training): + + fail_backbone = ( + not isinstance(backbone_outputs, BatchFeature) + or BACKBONE_FEATURE_KEY not in backbone_outputs + ) + + if fail_backbone: + error_msg = ERROR_MSG + error_msg += f"\n{isinstance(backbone_outputs, BatchFeature)=}" + error_msg += f"\n{BACKBONE_FEATURE_KEY in backbone_outputs=}" + error_msg += f"\n{backbone_outputs[BACKBONE_FEATURE_KEY].shape=}" + raise ValueError(error_msg) + + fail_action_head = (not isinstance(action_head_outputs, BatchFeature)) or not ( + ( + LOSS_KEY in action_head_outputs and is_training + ) # there might not be an action prediction during training + or ( + ACTION_KEY in action_head_outputs + and action_head_outputs[ACTION_KEY].shape[1] == self.action_horizon + and action_head_outputs[ACTION_KEY].shape[2] == self.action_dim + ) + ) + + if fail_action_head: + error_msg = ERROR_MSG + error_msg += f"\n{isinstance(action_head_outputs, BatchFeature)=}" + error_msg += f"\n{LOSS_KEY in action_head_outputs=}" + error_msg += f"\n{action_head_outputs[ACTION_KEY].shape=}" + error_msg += f"\n{self.action_horizon=}" + error_msg += f"\n{self.action_dim=}" + raise ValueError(error_msg) + + def forward( + self, + inputs: dict, + ) -> BatchFeature: + + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head(backbone_outputs, action_inputs) + + return action_head_outputs + + def get_action( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.get_action(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def joint_video_action( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.joint_video_action(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.lazy_joint_video_action(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action_causal( + self, + inputs: dict, + latent_video: torch.Tensor | None = None, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.lazy_joint_video_action(backbone_outputs, action_inputs, latent_video=latent_video) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action_causal_gt_cond( + self, + inputs: dict, + latent_video: torch.Tensor | None = None, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + + action_head_outputs = self.action_head.lazy_joint_video_action_causal_gt_cond(backbone_outputs, action_inputs, latent_video=latent_video) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action_efficient( + self, + inputs: dict, + prompt_embs: torch.Tensor | None = None, + prompt_emb_nega: torch.Tensor | None = None, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.lazy_joint_video_action_efficient(backbone_outputs, action_inputs, prompt_embs=prompt_embs, prompt_emb_nega=prompt_emb_nega) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def gt_video_action_pred( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.gt_video_action_pred(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def get_language( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + # Because the behavior of backbones remains the same for training and inference, we can use `forward` for backbones. + backbone_outputs = self.backbone.generate(backbone_inputs) + return backbone_outputs + + def get_video( + self, + inputs: dict, + ) -> BatchFeature: + _, video_inputs = self.prepare_input(inputs) + video_outputs = self.action_head.get_video(video_inputs) + return video_outputs + + def prepare_input(self, inputs) -> Tuple[BatchFeature, BatchFeature]: + self.validate_inputs(inputs) + backbone_inputs = self.backbone.prepare_input(inputs) + action_inputs = self.action_head.prepare_input(inputs) + + def to_device_with_maybe_dtype(x): + # Only cast to self.compute_dtype if the tensor is floating + if torch.is_floating_point(x): + return x.to(self.device, dtype=self.action_head.dtype) + else: + # Keep original dtype + return x.to(self.device) + + backbone_inputs = tree.map_structure(to_device_with_maybe_dtype, backbone_inputs) + action_inputs = tree.map_structure(to_device_with_maybe_dtype, action_inputs) + return backbone_inputs, action_inputs + + + @classmethod + def from_pretrained_for_tuning( + cls, + pretrained_model_name_or_path: str, + config: VLAConfig = None, # This config will now be USED + device_map: str = "auto", + dtype: torch.dtype = torch.bfloat16, + offload_state_dict: bool = True, + lora_weights_path: str | None = None, + ): + if config is None: + raise ValueError( + "A `config` object must be provided to build the model structure." + ) + + import os + import json + import gc + from safetensors.torch import load_file + + model = cls(config) + + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + if os.path.exists(safetensors_index_path): + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + missing_keys_accum = set() + unexpected_keys_accum = set() + shard_files = sorted(set(index["weight_map"].values())) + for shard_file in shard_files: + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + missing_keys, unexpected_keys = model.load_state_dict(shard_state_dict, strict=False) + if missing_keys: + missing_keys_accum.update(missing_keys) + if unexpected_keys: + unexpected_keys_accum.update(unexpected_keys) + # Free shard immediately + del shard_state_dict + gc.collect() + if missing_keys_accum: + print(f"Missing keys when loading sharded pretrained weights: {sorted(missing_keys_accum)} ... total={len(missing_keys_accum)}") + if unexpected_keys_accum: + print(f"Unexpected keys when loading sharded pretrained weights: {sorted(unexpected_keys_accum)} ... total={len(unexpected_keys_accum)}") + if not missing_keys_accum and not unexpected_keys_accum: + print("Successfully loaded pretrained base weights (sharded)") + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict = load_file(safetensors_path) + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + if not missing_keys and not unexpected_keys: + print("Successfully loaded pretrained base weights") + else: + raise FileNotFoundError( + f"No weights found at '{pretrained_model_name_or_path}'. " + "Expected 'model.safetensors' or 'model.safetensors.index.json'." + ) + + if lora_weights_path is not None: + print(f"Loading LoRA weights from: {lora_weights_path}") + model.load_lora_weight(lora_weights_path) + else: + if hasattr(model, 'action_head') and hasattr(model.action_head, 'inject_lora_after_loading') and model.action_head.config.defer_lora_injection: + print("Injecting LoRA adapters into action_head after loading pretrained weights") + model.action_head.inject_lora_after_loading() + + print(f"{cls}\n") + return model + + @classmethod + def load_lora( + cls, + pretrained_model_name_or_path: str + ): + from safetensors.torch import load_file + import os + import json + print("loading lora@@@@@") + + # Check for different checkpoint formats + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + state_dict = {} + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict.update(load_file(safetensors_path)) + + # Load config + print("loading config@@") + config_path = os.path.join(pretrained_model_name_or_path, "config.json") + with open(config_path, "r") as f: + config_dict = json.load(f) + config = VLAConfig(**config_dict) + print("loading model") + + # Disable defer_lora_injection so LoRA layers are created during init, + # matching the PEFT key hierarchy (base_model.model.*) in the checkpoint. + ah_cfg = config.action_head_cfg + inner = ah_cfg.get('config', ah_cfg) if isinstance(ah_cfg.get('config'), dict) else ah_cfg + if 'defer_lora_injection' in inner: + inner['defer_lora_injection'] = False + print("defer_lora_injection disabled for load_lora") + # Enable component loading so DiT base weights are loaded from pretrained + if 'skip_component_loading' in inner: + inner['skip_component_loading'] = False + print("skip_component_loading disabled for load_lora") + + # Instantiate model (LoRA layers now exist from init) + model = cls(config) + + # Remove .base_layer from keys if present + has_base_layer = any(".base_layer." in key for key in state_dict.keys()) + if has_base_layer: + print("Removing '.base_layer' from state dict keys") + state_dict = {k.replace(".base_layer.", "."): v for k, v in state_dict.items()} + + # Load weights + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + + print("Successfully loaded pretrained weights") + + print(f"{cls}\n") + return model + + def load_lora_weight(self, pretrained_model_name_or_path: str): + """Load only LoRA weights from a pretrained model without loading config.""" + from safetensors.torch import load_file + import os + import json + + print(f"Loading LoRA weights from {pretrained_model_name_or_path}") + + # Check for different checkpoint formats + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + state_dict = {} + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict.update(load_file(safetensors_path)) + else: + raise FileNotFoundError(f"No valid checkpoint found at {pretrained_model_name_or_path}") + + print("Loading LoRA weights into existing model") + + def rewrite_lora_state_dict_keys(state_dict, pattern, repl): + new_state_dict = {} + for k, v in state_dict.items(): + new_k = k.replace(pattern, repl) + new_state_dict[new_k] = v + return new_state_dict + + has_target_pattern = any("action_head.model.base_model.model" in key for key in state_dict.keys()) + + if not has_target_pattern: + print("Rewriting LoRA state dict keys from 'action_head.model' to 'action_head.model.base_model.model'") + state_dict = rewrite_lora_state_dict_keys( + state_dict, + pattern="action_head.model", + repl="action_head.model.base_model.model", + ) + else: + print("State dict already has 'action_head.model.base_model.model' pattern, skipping key rewrite") + + # Load only the weights into the existing model + missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=False) + + print("Successfully loaded LoRA state dict") + + if missing_keys: + print(f"Missing keys when loading LoRA weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading LoRA weights: {unexpected_keys}") + + print("Successfully loaded LoRA weights") + + @classmethod + def from_config_with_lora_weights( + cls, + config: VLAConfig, + pretrained_model_path: str, + ): + """Create VLA model from config and then load LoRA weights from pretrained model.""" + print(f"Creating VLA model from config and loading LoRA weights from {pretrained_model_path}") + + # 1. Create model from config (similar to vla.yaml) + model = cls(config) + print("Model created from config") + + # 2. Load LoRA weights into the created model + model.load_lora_weight(pretrained_model_path) + + return model + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str, + config: VLAConfig = None + ): + del config + + from safetensors.torch import load_file + import os + import json + print("loading pretrained@@@@@") + # Check for different checkpoint formats + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + state_dict = {} + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict.update(load_file(safetensors_path)) + + # Load config + print("loading config@@") + config_path = os.path.join(pretrained_model_name_or_path, "config.json") + with open(config_path, "r") as f: + config_dict = json.load(f) + config = VLAConfig(**config_dict) + print("loading model") + print("config.action_head_cfg", config.action_head_cfg) + # Always disable defer_lora_injection + # config.action_head_cfg is a dict, and defer_lora_injection is nested in config.action_head_cfg['config'] + if 'config' in config.action_head_cfg and isinstance(config.action_head_cfg['config'], dict): + if 'defer_lora_injection' in config.action_head_cfg['config']: + config.action_head_cfg['config']['defer_lora_injection'] = False + print("config.action_head_cfg['config']['defer_lora_injection'] disabled (set to False)") + elif 'defer_lora_injection' in config.action_head_cfg: + config.action_head_cfg['defer_lora_injection'] = False + print("config.action_head_cfg['defer_lora_injection'] disabled (set to False)") + + # Instantiate model + model = cls(config) + print("model", model) + # Remove .base_layer from keys (e.g., 'action_head.model.base_model.model.blocks.19.self_attn.v.base_layer.bias' -> 'action_head.model.base_model.model.blocks.19.self_attn.v.bias') + has_base_layer = any(".base_layer." in key for key in state_dict.keys()) + if has_base_layer: + print("Removing '.base_layer' from state dict keys") + new_state_dict = {} + for k, v in state_dict.items(): + new_k = k.replace(".base_layer.", ".") + new_state_dict[new_k] = v + state_dict = new_state_dict + + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + + print("Successfully loaded pretrained weights") + + print(f"{cls}\n") + return model + + def post_initialize(self): + self.action_head.post_initialize() + + def parallelize(self, device_mesh: DeviceMesh): + self.action_head.parallelize(device_mesh=device_mesh) + + +class CotrainVLA(VLA): + + def forward( + self, + inputs: dict, + ) -> BatchFeature: + if "cotrain" in inputs and inputs["cotrain"]: + return self.backbone.cotrain(inputs) + return super().forward(inputs) + + +def create_vla_with_pretrained_action_head(pretrained_vla_path: str, config: VLAConfig): + # 1. Instantiate a new VLAModel + vla = VLA(config) + + # 2. Load the pretrained VLAModel + pretrained_vla = VLA.from_pretrained(pretrained_vla_path) + + # 3. Replace the action head in the new VLAModel with the pretrained action head + vla.action_head = pretrained_vla.action_head + + # 4. Replace the action head config in the new VLAModel with the pretrained action head config + vla.config.action_head_cfg = pretrained_vla.config.action_head_cfg + + # 5. Return the new VLAModel + return vla + + +# register +AutoConfig.register("vla", VLAConfig) +AutoModel.register(VLAConfig, VLA) diff --git a/groot/vla/model/dreamzero/modules/__init__.py b/groot/vla/model/dreamzero/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/modules/attention.py b/groot/vla/model/dreamzero/modules/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..0ebc62a64a02b84cbecbc682e5612b0e290a87d4 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/attention.py @@ -0,0 +1,258 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import torch +import os + +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +import warnings + + +__all__ = [ + 'flash_attention', + 'attention', +] + + +def _gpu_supports_flash_attention(): + """FlashAttention requires Ampere (compute capability 8.0) or newer.""" + if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): + return False + try: + if not torch.cuda.is_available(): + return False + cap = torch.cuda.get_device_capability() + return cap[0] >= 8 + except Exception: + return False + + +def _sdpa_attention_fallback( + q, k, v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + dtype=torch.bfloat16, +): + """PyTorch SDPA fallback for GPUs that don't support FlashAttention (e.g. pre-Ampere).""" + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention on this GPU. ' + 'It can have a slight impact on quality.' + ) + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + if q_scale is not None: + q = q * q_scale + if softmax_scale is not None: + q = q * softmax_scale + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=None, is_causal=causal, dropout_p=dropout_p + ) + return out.transpose(1, 2).contiguous() + + +def flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + version=None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == 'cuda' and q.size(-1) <= 256 + + # Use PyTorch SDPA on pre-Ampere GPUs (FlashAttention requires Ampere or newer) + if not _gpu_supports_flash_attention(): + return _sdpa_attention_fallback( + q, k, v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + dtype=dtype, + ) + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor( + [lq] * b, dtype=torch.int32).to( + device=q.device, non_blocking=True) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor( + [lk] * b, dtype=torch.int32).to( + device=k.device, non_blocking=True) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + + if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn( + 'Flash attention 3 is not available, use flash attention 2 instead.' + ) + + # Check for TensorRT at runtime, not import time + if os.getenv("ENABLE_TENSORRT", "False").lower() == "true": + # use torch.nn.functional.scaled_dot_product_attention for tensorrt export + + # The input is (s, n, d), but sdpa needs (b, n, s, d). + # We add a batch dimension and transpose. + q = q.unsqueeze(0).transpose(1, 2) + k = k.unsqueeze(0).transpose(1, 2) + v = v.unsqueeze(0).transpose(1, 2) + + # Fix for ONNX export: repeat k and v to match q's batch size in cross-attention + if q.shape[0] != k.shape[0] and k.shape[0] == 1: + k = k.repeat(q.shape[0], 1, 1, 1) + v = v.repeat(q.shape[0], 1, 1, 1) + + attn_mask = None + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p) + + # Transpose back to (b, s, n, d) format. + out = out.transpose(1, 2).contiguous() + return out + + elif (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p, window_size are not supported in FA3 now. + x = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + seqused_q=None, + seqused_k=None, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic)[0].unflatten(0, (b, lq)) + else: + assert FLASH_ATTN_2_AVAILABLE + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic).unflatten(0, (b, lq)) + + # output + return x.type(out_dtype) + + +def attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + fa_version=None, +): + if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE: + return flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=fa_version, + ) + else: + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.' + ) + attn_mask = None + + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p) + + out = out.transpose(1, 2).contiguous() + return out diff --git a/groot/vla/model/dreamzero/modules/cudnn_attention.py b/groot/vla/model/dreamzero/modules/cudnn_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..c996baa33c4f8212652234a8f6d9ba7645d3c234 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/cudnn_attention.py @@ -0,0 +1,610 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Simplified wrapper around TransformerEngine's C++ pytorch backend. +This supports torch.compile(fullgraph=True). +Lowers to cudnn ultimately. +Only bf16 / fp16 is supported. +Only BSHD layout is supported. +Currently, tensors are made contiguous -- packed th2d, th3d not supported yet. +""" + +import math +from typing import Any, List, Optional, Tuple, Union + +import torch +import transformer_engine + +_TE_VER = tuple(int(x) for x in transformer_engine.__version__.split(".")[:2]) + + +try: + # transformer_engine >= 2.8.0 + import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils +except ImportError: + # transformer_engine < 2.8.0 + import transformer_engine.pytorch.dot_product_attention.utils as dpa_utils + +import transformer_engine_torch as tex +from transformer_engine.pytorch.constants import ( + TE_DType, +) +from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + AttnBiasType, + AttnMaskType, + FusedAttnBackend, + QKVLayout, +) + +if _TE_VER >= (2, 8): + from transformer_engine.pytorch.cpp_extensions.fused_attn import SoftmaxType + +from transformer_engine.pytorch.utils import get_cudnn_version + +__all__ = ["DotProductAttention"] + + +class DotProductAttention(torch.nn.Module): + def __init__( + self, + num_attention_heads: int, + kv_channels: Union[int, Tuple[int, int]], + num_gqa_groups: Optional[int] = None, + attention_dropout: float = 0.0, + qkv_format: str = "bshd", + attn_mask_type: str = "no_mask", + window_size: Optional[Tuple[int, int]] = None, + ) -> None: + super().__init__() + assert qkv_format == "bshd", "Only bshd layout is supported." + + self.qkv_format = qkv_format + self.attn_mask_type = attn_mask_type + + self.softmax_scale = 1.0 / math.sqrt(kv_channels if isinstance(kv_channels, int) else kv_channels[0]) + self.attention_dropout = attention_dropout + self.softmax_type = "vanilla" + self.window_size = dpa_utils.check_set_window_size(attn_mask_type) + + self.fused_attention = FusedAttention( + self.softmax_scale, + deterministic=False, + attention_dropout=self.attention_dropout, + softmax_type=self.softmax_type, + ) + + def forward( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + ) -> torch.Tensor: + """ + Dot Product Attention Layer. + """ + + qkv_layout = "bshd_bshd_bshd" + batch_size = query_layer.shape[0] + device = query_layer.device + + def _get_cu_seqlens(max_seqlen: int) -> torch.Tensor: + return torch.arange( + 0, + (batch_size + 1) * max_seqlen, + step=max_seqlen, + dtype=torch.int32, + device=device, + ) + + max_seqlen_q = query_layer.shape[1] + max_seqlen_kv = key_layer.shape[1] + cu_seqlens_q = _get_cu_seqlens(max_seqlen_q) + cu_seqlens_kv = _get_cu_seqlens(max_seqlen_kv) + + return self.fused_attention( + query_layer, + key_layer, + value_layer, + qkv_layout=qkv_layout, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=cu_seqlens_q, + cu_seqlens_kv_padded=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=self.attn_mask_type, + window_size=self.window_size, + fused_attention_backend=FusedAttnBackend["F16_arbitrary_seqlen"], + ) + + +def prepare_for_saving(*tensors) -> Tuple[list[torch.Tensor | None], list[torch.Tensor | None]]: + """Prepare tensors for saving. Needed because save_for_backward accepts only + torch.Tensor/torch.nn.Parameter types, while we want to be able to save + the internal TensorBase types too.""" + + tensor_list, tensor_objects_list = [], [] + for tensor in tensors: + if tensor is None or isinstance(tensor, torch.Tensor): + tensor_list.append(tensor) + tensor_objects_list.append(None) + else: + t, t_obj = tensor.prepare_for_saving() + tensor_list.extend(t) + tensor_objects_list.append(t_obj) + return tensor_list, tensor_objects_list + + +def restore_from_saved(tensors, saved_tensors) -> Tuple[Any, ...]: + """Recombine the tensor data and metadata during backward pass.""" + tensor_objects = [] + for tensor in tensors: + if tensor is None or isinstance(tensor, torch.Tensor): + tensor_objects.append(saved_tensors[0]) + saved_tensors = saved_tensors[1:] + else: + saved_tensors = tensor.restore_from_saved(saved_tensors) + tensor_objects.append(tensor) + + return tuple(tensor_objects) + + +class FusedAttention(torch.nn.Module): + def __init__( + self, + softmax_scale: float, + attention_dropout: float = 0.0, + deterministic: bool = False, + softmax_type: str = "vanilla", + ) -> None: + super().__init__() + self.softmax_scale = softmax_scale + self.attention_dropout = attention_dropout + self.deterministic = deterministic + self.softmax_type = softmax_type + + def forward( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + qkv_layout: str, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + cu_seqlens_q_padded: torch.Tensor, + cu_seqlens_kv_padded: torch.Tensor, + max_seqlen_q: int, + max_seqlen_kv: int, + attn_mask_type: str = "causal", + attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, + window_size: Optional[Tuple[int, int]] = None, + fused_attention_backend: tex.NVTE_Fused_Attn_Backend = tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend, + core_attention_bias_type: str = "no_bias", + core_attention_bias: Optional[torch.Tensor] = None, + fast_zero_fill: bool = True, + quantizers=None, + pad_between_seqs: bool = False, + softmax_offset: torch.Tensor | None = None, + softmax_scale: float = None, + attention_dropout: float = 0.0, + deterministic: bool = False, + softmax_type: str = "vanilla", + ) -> torch.Tensor: + """fused attention fprop""" + + cu_seqlens_q_padded = cu_seqlens_q + cu_seqlens_kv_padded = cu_seqlens_kv + + out_nominal_dtype = query_layer.dtype + output_tensors = fused_attn( + self.training, + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + query_layer, + key_layer, + value_layer, + out_nominal_dtype, + window_size, + core_attention_bias, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + None, # page_table_k, + None, # page_table_v, + self.softmax_scale, + self.attention_dropout if self.training else 0.0, + fast_zero_fill, + qkv_layout, + core_attention_bias_type, + attn_mask_type, + self.softmax_type, + self.deterministic, + softmax_offset, + ) + return output_tensors[0] + + +BACKEND_F16arb_ELTS_PER_THREADS = 16 + + +@torch.library.custom_op("groot::fused_attn", mutates_args=()) +def fused_attn( + is_training: bool, + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + fake_dtype: torch.dtype, + window_size: List[int], + attn_bias: torch.Tensor = None, + cu_seqlens_q_padded: torch.Tensor = None, + cu_seqlens_kv_padded: torch.Tensor = None, + page_table_k: torch.Tensor = None, + page_table_v: torch.Tensor = None, + attn_scale: Optional[float] = None, + dropout: float = 0.0, + fast_zero_fill: bool = True, + qkv_layout: str = "sbh3d", + attn_bias_type: str = "no_bias", + attn_mask_type: str = "padding", + softmax_type: str = "vanilla", + deterministic: bool = False, + softmax_offset: torch.Tensor = None, +) -> List[torch.Tensor]: + assert deterministic is not None + + rng_elts_per_thread = BACKEND_F16arb_ELTS_PER_THREADS + s_quantizer = None + o_quantizer = None + rng_gen = None + + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + + args = ( + max_seqlen_q, + max_seqlen_kv, + is_training, + attn_scale, + dropout, + fast_zero_fill, + QKVLayout[qkv_layout], + AttnBiasType[attn_bias_type], + AttnMaskType[attn_mask_type], + ) + + if _TE_VER >= (2, 8): + args += (SoftmaxType[softmax_type],) + + args += ( + tuple(window_size), + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + fake_dtype, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + page_table_k, + page_table_v, + s_quantizer, + o_quantizer, + attn_bias, + ) + + if _TE_VER >= (2, 8): + args += (softmax_offset,) + + args += ( + rng_gen, + rng_elts_per_thread, + ) + + if _TE_VER >= (2, 9): + # return_max_logit + args += (False,) + + if _TE_VER >= (2, 10): + # is_cuda_graph + args += (False,) + + output_tensors = tex.fused_attn_fwd(*args) + return output_tensors + + +@fused_attn.register_fake +def _( + is_training: bool, + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + fake_dtype: torch.dtype, + window_size: List[int], + attn_bias: torch.Tensor = None, + cu_seqlens_q_padded: torch.Tensor = None, + cu_seqlens_kv_padded: torch.Tensor = None, + page_table_k: torch.Tensor = None, + page_table_v: torch.Tensor = None, + attn_scale: Optional[float] = None, + dropout: float = 0.0, + fast_zero_fill: bool = True, + qkv_layout: str = "sbh3d", + attn_bias_type: str = "no_bias", + attn_mask_type: str = "padding", + softmax_type: str = "vanilla", + deterministic: bool = False, + softmax_offset: torch.Tensor = None, +) -> List[torch.Tensor]: + return [ + q.new_empty(tuple(q.shape[:-1]) + (v.shape[-1],)), + q.new_empty( + q.shape[0], q.shape[2], q.shape[1], 1, dtype=torch.float32 + ), # these are the softmax outputs from cudnn; will always be float32 + q.new_empty((2,), dtype=torch.int64), + ] + + +def fused_attn_bwd_setup_context(ctx, inputs, output) -> None: + ( + _, + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + _, + window_size, + _, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + _, + _, + attn_scale, + dropout, + fast_zero_fill, + qkv_layout, + attn_bias_type, + attn_mask_type, + softmax_type, + deterministic, + _, + ) = inputs + + out = output[0] + aux_ctx_tensors = output[1:] + qkvo_tensors = (q, k, v, out) + + # assume fwd and bwd always use the same high precision, i.e. torch.float16 or torch.bfloat16 + # used when some tensors are base tensors and loose the "dtype" attribute + ctx.nominal_dtype = q.dtype + + tensors_to_save, tensor_objects = prepare_for_saving( + *qkvo_tensors, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + *aux_ctx_tensors, + ) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + + ctx.QKV_quantizer = None + ctx.O_quantizer = None + ctx.dQKV_quantizer = None + ctx.dO_quantizer = None + ctx.dP_quantizer = None + ctx.S_quantizer = None + + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_kv = max_seqlen_kv + ctx.attn_scale = attn_scale + ctx.dropout_p = dropout + ctx.fast_zero_fill = fast_zero_fill + ctx.qkv_layout = qkv_layout + ctx.attn_bias_type = attn_bias_type + ctx.attn_mask_type = attn_mask_type + ctx.softmax_type = softmax_type + ctx.window_size = window_size + ctx.fused_attention_backend = FusedAttnBackend["F16_arbitrary_seqlen"] + ctx.deterministic = deterministic + + +@torch.library.custom_op("groot::fused_attn_bwd_op", mutates_args=()) +def fused_attn_bwd_op( + max_seqlen_q: int, + max_seqlen_kv: int, + attn_scale: float, + dropout: float, + fast_zero_fill: bool, + qkv_layout: str, + attn_bias_type: str, + attn_mask_type: str, + softmax_type: str, + window_size: List[int], + deterministic: bool, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + d_out: torch.Tensor, + dqkv_nominal_dtype: torch.dtype, + dqkv_te_dtype: torch.dtype, + aux_ctx_tensors: List[torch.Tensor], + cu_seqlens_q_padded: torch.Tensor, + cu_seqlens_kv_padded: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + args = ( + max_seqlen_q, + max_seqlen_kv, + attn_scale, + dropout, + fast_zero_fill, + QKVLayout[qkv_layout], + AttnBiasType[attn_bias_type], + AttnMaskType[attn_mask_type], + ) + + if _TE_VER >= (2, 8): + args += (SoftmaxType[softmax_type],) + + args += ( + window_size, + deterministic, + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + out, + d_out, + dqkv_nominal_dtype, + TE_DType[dqkv_te_dtype], + aux_ctx_tensors, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + None, # s_quantizer, + None, # dp_quantizer, + None, # dqkv_quantizer, + ) + + if _TE_VER >= (2, 10): + # is_cuda_graph + args += (False,) + + dq, dk, dv, *rest = tex.fused_attn_bwd(*args) + return dq, dk, dv + + +@fused_attn_bwd_op.register_fake +def _( + max_seqlen_q: int, + max_seqlen_kv: int, + attn_scale: float, + dropout: float, + fast_zero_fill: bool, + qkv_layout: str, + attn_bias_type: str, + attn_mask_type: str, + softmax_type: str, + window_size: List[int], + deterministic: bool, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + d_out: torch.Tensor, + dqkv_nominal_dtype: torch.dtype, + dqkv_te_dtype: torch.dtype, + aux_ctx_tensors: List[torch.Tensor], + cu_seqlens_q_padded: torch.Tensor, + cu_seqlens_kv_padded: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + + +def fused_attn_bwd_impl(ctx, grad): + d_out, _, _ = grad + d_out = d_out.contiguous() + + ( + q, + k, + v, + out, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + *other_tensors, + ) = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) + + aux_ctx_tensors = other_tensors + + if not aux_ctx_tensors[0].is_contiguous(): + aux_ctx_tensors[0] = aux_ctx_tensors[0].contiguous() + + with torch.cuda.nvtx.range("FusedAttnFunc.backward"): + assert ctx.fused_attention_backend != FusedAttnBackend["No_Backend"], ( + "Fused attention does not support this input combination." + ) + + # get nominal data type of dq, dk, dv + # FP16/BF16 attention: torch.float16 or torch.bfloat16 + dqkv_nominal_dtype = ctx.nominal_dtype + + # q, k, v, out, d_out, dq, dk, dv: torch.Tensor; torch.float16 or torch.bfloat16 + dq, dk, dv = fused_attn_bwd_op( + ctx.max_seqlen_q, + ctx.max_seqlen_kv, + ctx.attn_scale, + ctx.dropout_p, + ctx.fast_zero_fill, + ctx.qkv_layout, + ctx.attn_bias_type, + ctx.attn_mask_type, + ctx.softmax_type, + ctx.window_size, + ctx.deterministic, + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + out, + d_out, + dqkv_nominal_dtype, + d_out.dtype, + aux_ctx_tensors, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + ) + + output = ( + None, # is_training + None, # max_seqlen_q + None, # max_seqlen_kv + None, # cu_seqlens_q + None, # cu_seqlens_kv + dq, + dk, + dv, + None, # fake_dtype + None, # window_size + None, # d_bias, # attn_bias + None, # cu_seqlens_q_padded + None, # cu_seqlens_kv_padded + None, # page_table_k + None, # page_table_v + None, # attn_scale + None, # dropout + None, # fast_zero_fill + None, # qkv_layout + None, # attn_bias_type + None, # attn_mask_type + None, # softmax_type + None, # deterministic + None, # d_softmax_offset, # softmax_offset + ) + return output + + +fused_attn.register_autograd(fused_attn_bwd_impl, setup_context=fused_attn_bwd_setup_context) diff --git a/groot/vla/model/dreamzero/modules/flow_match_scheduler.py b/groot/vla/model/dreamzero/modules/flow_match_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..03b53b3b26237ef82492d174c78b971a22280a0f --- /dev/null +++ b/groot/vla/model/dreamzero/modules/flow_match_scheduler.py @@ -0,0 +1,92 @@ +import torch + + + +class FlowMatchScheduler(): + + def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003/1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.set_timesteps(num_inference_steps) + + + def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, shift=None): + if shift is not None: + self.shift = shift + sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + self.training = True + else: + self.training = False + + + def step(self, model_output, timestep, sample, to_final=False, **kwargs): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + if to_final or timestep_id + 1 >= len(self.timesteps): + sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 + else: + sigma_ = self.sigmas[timestep_id + 1] + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample + + + def return_to_timestep(self, timestep, sample, sample_stablized): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + model_output = (sample - sample_stablized) / sigma + return model_output + + + # def add_noise(self, original_samples, noise, timestep): + # if isinstance(timestep, torch.Tensor): + # timestep = timestep.cpu() + # timestep_id = torch.argmin((self.timesteps - timestep).abs()) + # sigma = self.sigmas[timestep_id] + # sample = (1 - sigma) * original_samples + sigma * noise + # return sample + + def add_noise(self, original_samples, noise, timestep): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim = 0) + sigma = self.sigmas[timestep_id].to(device=original_samples.device, dtype=original_samples.dtype) + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + sample = (1 - sigma) * original_samples + sigma * noise + return sample + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + + def training_weight(self, timestep): + # timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs()) + timestep_id = torch.argmin((self.timesteps.unsqueeze(1) - timestep.unsqueeze(0).to(self.timesteps.device)).abs(), dim = 0) + weights = self.linear_timesteps_weights[timestep_id] + return weights \ No newline at end of file diff --git a/groot/vla/model/dreamzero/modules/flow_unipc_multistep_scheduler.py b/groot/vla/model/dreamzero/modules/flow_unipc_multistep_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..df7ea56e278813a242cbb1cec873edc79649f8a5 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/flow_unipc_multistep_scheduler.py @@ -0,0 +1,680 @@ +# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py +# Convert unipc for flow matching +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +import math +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import ( + KarrasDiffusionSchedulers, + SchedulerMixin, + SchedulerOutput, +) + + +class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin): + """ + `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + solver_order (`int`, default `2`): + The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1` + due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for + unconditional sampling. + prediction_type (`str`, defaults to "flow_prediction"): + Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts + the flow of the diffusion process. + thresholding (`bool`, defaults to `False`): + Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such + as Stable Diffusion. + dynamic_thresholding_ratio (`float`, defaults to 0.995): + The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. + sample_max_value (`float`, defaults to 1.0): + The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`. + predict_x0 (`bool`, defaults to `True`): + Whether to use the updating algorithm on the predicted x0. + solver_type (`str`, default `bh2`): + Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2` + otherwise. + lower_order_final (`bool`, default `True`): + Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can + stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10. + disable_corrector (`list`, default `[]`): + Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)` + and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is + usually disabled during the first few steps. + solver_p (`SchedulerMixin`, default `None`): + Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`. + use_karras_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, + the sigmas are determined according to a sequence of noise levels {σi}. + use_exponential_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + steps_offset (`int`, defaults to 0): + An offset added to the inference steps, as required by some model families. + final_sigmas_type (`str`, defaults to `"zero"`): + The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final + sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. + """ + + _compatibles = [e.name for e in KarrasDiffusionSchedulers] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + solver_order: int = 2, + prediction_type: str = "flow_prediction", + shift: Optional[float] = 1.0, + use_dynamic_shifting=False, + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: float = 1.0, + predict_x0: bool = True, + solver_type: str = "bh2", + lower_order_final: bool = True, + disable_corrector: List[int] = [], + solver_p: SchedulerMixin = None, + timestep_spacing: str = "linspace", + steps_offset: int = 0, + final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" + ): + + if solver_type not in ["bh1", "bh2"]: + if solver_type in ["midpoint", "heun", "logrho"]: + self.register_to_config(solver_type="bh2") + else: + raise NotImplementedError( + f"{solver_type} is not implemented for {self.__class__}") + + self.predict_x0 = predict_x0 + # setable values + self.num_inference_steps = None + alphas = np.linspace(1, 1 / num_train_timesteps, + num_train_timesteps)[::-1].copy() + sigmas = 1.0 - alphas + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device='cuda') + + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + assert shift is not None, "shift must be not None when use_dynamic_shifting is False" + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.sigmas = sigmas + self.timesteps = sigmas * num_train_timesteps + + self.model_outputs = [None] * solver_order + self.timestep_list = [None] * solver_order + self.lower_order_nums = 0 + self.disable_corrector = disable_corrector + self.solver_p = solver_p + self.last_sample = None + + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device = None, + sigmas: np.ndarray[float] | None = None, + mu: float | None = None, + shift: float | None = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + Args: + num_inference_steps (`int`): + Total number of the spacing of the time steps. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + self.num_inference_steps = num_inference_steps + + if self.config.use_dynamic_shifting and mu is None: + raise ValueError( + " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`" + ) + + if sigmas is None: + sigmas = np.linspace( + self.sigma_max, + self.sigma_min, + num_inference_steps + 1, + ).copy()[:-1] + + if self.config.use_dynamic_shifting: + assert mu is not None + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + if shift is None: + shift = self.config.shift + assert isinstance(shift, float) + assert sigmas is not None + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + if self.config.final_sigmas_type == "sigma_min": + sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0])**0.5 + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + + timesteps = sigmas * self.config.num_train_timesteps + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) + + self.sigmas = torch.from_numpy(sigmas).to(device=device) + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) + + self.model_outputs = [None] * self.config.solver_order + self.lower_order_nums = 0 + self.last_sample = None + if self.solver_p: + self.solver_p.set_timesteps(self.num_inference_steps, device=device) + + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample + def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: + """ + "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the + prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by + s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing + pixels from saturation at each step. We find that dynamic thresholding results in significantly better + photorealism as well as better image-text alignment, especially when using very large guidance weights." + + https://arxiv.org/abs/2205.11487 + """ + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + + if dtype not in (torch.float32, torch.float64): + sample = sample.float( + ) # upcast for quantile calculation, and clamp not implemented for cpu half + + # Flatten sample for doing quantile calculation along each image + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + + abs_sample = sample.abs() # "a certain percentile absolute pixel value" + + s = torch.quantile( + abs_sample, self.config.dynamic_thresholding_ratio, dim=1) + s = torch.clamp( + s, min=1, max=self.config.sample_max_value + ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s = s.unsqueeze( + 1) # (batch_size, 1) because clamp will broadcast along dim=0 + sample = torch.clamp( + sample, -s, s + ) / s # "we threshold xt0 to the range [-s, s] and then divide by s" + + sample = sample.reshape(batch_size, channels, *remaining_dims) + sample = sample.to(dtype) + + return sample + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + @staticmethod + def _sigma_to_alpha_sigma_t(sigma): + return 1 - sigma, sigma + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma) + + def convert_model_output( + self, + model_output: torch.Tensor, + sample: torch.Tensor, + step_index: int, + ) -> torch.Tensor: + r""" + Convert the model output to the corresponding type the UniPC algorithm needs. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + step_index (`int`): + The current timestep index. + + Returns: + `torch.Tensor`: + The converted model output. + """ + if self.predict_x0: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[step_index] + x0_pred = sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + + return x0_pred + else: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[step_index] + epsilon = sample - (1 - sigma_t) * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + sigma_t = self.sigmas[step_index] + x0_pred = sample - sigma_t * model_output + x0_pred = self._threshold_sample(x0_pred) + epsilon = model_output + x0_pred + + return epsilon + + @torch.compile(mode="reduce-overhead", fullgraph=True, dynamic=False) + def multistep_uni_p_bh_update( + self, + model_output: torch.Tensor, + sample: torch.Tensor, + order: int, + step_index: int, + ) -> torch.Tensor: + """ + One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model at the current timestep. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + order (`int`): + The order of UniP at this timestep (corresponds to the *p* in UniPC-p). + step_index (`int`): + The current timestep index. + + Returns: + `torch.Tensor`: + The sample tensor at the previous timestep. + """ + model_output_list = self.model_outputs + + s0 = self.timestep_list[-1] + m0 = model_output_list[-1] + x = sample + + if self.solver_p: + x_t = self.solver_p.step(model_output, s0, x).prev_sample + return x_t + + sigma_t, sigma_s0 = self.sigmas[step_index + 1], self.sigmas[step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + + rks = [] + D1s = [] + for i in range(1, order): + si = step_index - i + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(torch.ones((), dtype=self.sigmas.dtype, device=self.sigmas.device)) + rks = torch.stack(rks, dim=0) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R, dim=0) + b = torch.stack(b, dim=0) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) # (B, K) + # for order 2, we use a simplified version + if order == 2: + rhos_p = torch.full((1,), 0.5, dtype=x.dtype, device=self.sigmas.device) + else: + rhos_p = torch.linalg.solve_ex(R[:-1, :-1], b[:-1])[0].to(x.dtype) + else: + D1s = None + rhos_p = None + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + else: + pred_res = 0 + x_t = x_t_ - alpha_t * B_h * pred_res + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + else: + pred_res = 0 + x_t = x_t_ - sigma_t * B_h * pred_res + + x_t = x_t.to(x.dtype) + return x_t + + @torch.compile(mode="reduce-overhead", fullgraph=True, dynamic=False) + def multistep_uni_c_bh_update( + self, + this_model_output: torch.Tensor, + last_sample: torch.Tensor, + this_sample: torch.Tensor, + order: int, + step_index: int, + ) -> torch.Tensor: + """ + One step for the UniC (B(h) version). + + Args: + this_model_output (`torch.Tensor`): + The model outputs at `x_t`. + this_timestep (`int`): + The current timestep `t`. + last_sample (`torch.Tensor`): + The generated sample before the last predictor `x_{t-1}`. + this_sample (`torch.Tensor`): + The generated sample after the last predictor `x_{t}`. + order (`int`): + The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`. + step_index (`int`): + The current timestep index. + + Returns: + `torch.Tensor`: + The corrected sample tensor at the current timestep. + """ + model_output_list = self.model_outputs + + m0 = model_output_list[-1] + x = last_sample + x_t = this_sample + model_t = this_model_output + + sigma_t, sigma_s0 = self.sigmas[step_index], self.sigmas[step_index - 1] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + + rks = [] + D1s = [] + for i in range(1, order): + si = step_index - (i + 1) + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(torch.ones((), dtype=self.sigmas.dtype, device=self.sigmas.device)) + rks = torch.stack(rks, dim=0) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R, dim=0) + b = torch.stack(b, dim=0) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + else: + D1s = None + + # for order 1, we use a simplified version + if order == 1: + rhos_c = torch.full((1,), 0.5, dtype=x.dtype, device=self.sigmas.device) + else: + rhos_c = torch.linalg.solve_ex(R, b)[0].to(x.dtype) + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) + x_t = x_t.to(x.dtype) + return x_t + + def index_for_timestep(self, timestep, schedule_timesteps): + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def step( + self, + model_output: torch.Tensor, + timestep: torch.Tensor, + sample: torch.Tensor, + step_index: int, + return_dict: bool = True, + ) -> SchedulerOutput | tuple: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with + the multistep UniPC. + + Args: + model_output (`torch.Tensor`): + The direct output from learned diffusion model. + timestep (`torch.Tensor`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + step_index (`int`): + The current timestep index. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. + + Returns: + [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + use_corrector = ( + step_index > 0 and + step_index - 1 not in self.disable_corrector and + self.last_sample is not None + ) + + model_output_convert = self.convert_model_output( + model_output=model_output, + sample=sample, + step_index=step_index, + ) + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + step_index=step_index, + ) + # We must clone the outputs of a CUDA graph'd computation. + sample = sample.clone() + + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep + + if self.config.lower_order_final: + this_order = min( + self.config.solver_order, + len(self.timesteps) - step_index, + ) + else: + this_order = self.config.solver_order + + # Warmup for multistep. + self.this_order = min(this_order, self.lower_order_nums + 1) + assert self.this_order > 0 + + self.last_sample = sample + # Pass the original non-converted model output, in case solver-p is used. + prev_sample = self.multistep_uni_p_bh_update( + model_output=model_output, + sample=sample, + order=self.this_order, + step_index=step_index, + ) + # We must clone the outputs of a CUDA graph'd computation. + prev_sample = prev_sample.clone() + + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + if not return_dict: + return (prev_sample,) + + return SchedulerOutput(prev_sample=prev_sample) + + def scale_model_input(self, sample: torch.Tensor, *args, + **kwargs) -> torch.Tensor: + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + + Args: + sample (`torch.Tensor`): + The input sample. + + Returns: + `torch.Tensor`: + A scaled input sample. + """ + return sample + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.IntTensor, + ) -> torch.Tensor: + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to( + device=original_samples.device, dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): + # mps does not support float64 + schedule_timesteps = self.timesteps.to( + original_samples.device, dtype=torch.float32) + timesteps = timesteps.to( + original_samples.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(original_samples.device) + timesteps = timesteps.to(original_samples.device) + + # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index + step_indices = [ + self.index_for_timestep(t, schedule_timesteps) + for t in timesteps + ] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + noisy_samples = alpha_t * original_samples + sigma_t * noise + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps diff --git a/groot/vla/model/dreamzero/modules/utils.py b/groot/vla/model/dreamzero/modules/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0d58e4e11dc844822d0588dac7f432e7e31b8daf --- /dev/null +++ b/groot/vla/model/dreamzero/modules/utils.py @@ -0,0 +1,182 @@ +import torch, os +from safetensors import safe_open +from contextlib import contextmanager +import hashlib + +@contextmanager +def init_weights_on_device(device = torch.device("meta"), include_buffers :bool = False): + + old_register_parameter = torch.nn.Module.register_parameter + if include_buffers: + old_register_buffer = torch.nn.Module.register_buffer + + def register_empty_parameter(module, name, param): + old_register_parameter(module, name, param) + if param is not None: + param_cls = type(module._parameters[name]) + kwargs = module._parameters[name].__dict__ + kwargs["requires_grad"] = param.requires_grad + module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) + + def register_empty_buffer(module, name, buffer, persistent=True): + old_register_buffer(module, name, buffer, persistent=persistent) + if buffer is not None: + module._buffers[name] = module._buffers[name].to(device) + + def patch_tensor_constructor(fn): + def wrapper(*args, **kwargs): + kwargs["device"] = device + return fn(*args, **kwargs) + + return wrapper + + if include_buffers: + tensor_constructors_to_patch = { + torch_function_name: getattr(torch, torch_function_name) + for torch_function_name in ["empty", "zeros", "ones", "full"] + } + else: + tensor_constructors_to_patch = {} + + try: + torch.nn.Module.register_parameter = register_empty_parameter + if include_buffers: + torch.nn.Module.register_buffer = register_empty_buffer + for torch_function_name in tensor_constructors_to_patch.keys(): + setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name))) + yield + finally: + torch.nn.Module.register_parameter = old_register_parameter + if include_buffers: + torch.nn.Module.register_buffer = old_register_buffer + for torch_function_name, old_torch_function in tensor_constructors_to_patch.items(): + setattr(torch, torch_function_name, old_torch_function) + +def load_state_dict_from_folder(file_path, torch_dtype=None): + state_dict = {} + for file_name in os.listdir(file_path): + if "." in file_name and file_name.split(".")[-1] in [ + "safetensors", "bin", "ckpt", "pth", "pt" + ]: + state_dict.update(load_state_dict(os.path.join(file_path, file_name), torch_dtype=torch_dtype)) + return state_dict + + +def load_state_dict(file_path, torch_dtype=None, device="cpu"): + if file_path.endswith(".safetensors"): + return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device) + else: + return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device) + + +def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"): + state_dict = {} + with safe_open(file_path, framework="pt", device=device) as f: + for k in f.keys(): + state_dict[k] = f.get_tensor(k) + if torch_dtype is not None: + state_dict[k] = state_dict[k].to(torch_dtype) + return state_dict + + +def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"): + state_dict = torch.load(file_path, map_location=device, weights_only=True) + if torch_dtype is not None: + for i in state_dict: + if isinstance(state_dict[i], torch.Tensor): + state_dict[i] = state_dict[i].to(torch_dtype) + return state_dict + + +def search_for_embeddings(state_dict): + embeddings = [] + for k in state_dict: + if isinstance(state_dict[k], torch.Tensor): + embeddings.append(state_dict[k]) + elif isinstance(state_dict[k], dict): + embeddings += search_for_embeddings(state_dict[k]) + return embeddings + + +def search_parameter(param, state_dict): + for name, param_ in state_dict.items(): + if param.numel() == param_.numel(): + if param.shape == param_.shape: + if torch.dist(param, param_) < 1e-3: + return name + else: + if torch.dist(param.flatten(), param_.flatten()) < 1e-3: + return name + return None + + +def build_rename_dict(source_state_dict, target_state_dict, split_qkv=False): + matched_keys = set() + with torch.no_grad(): + for name in source_state_dict: + rename = search_parameter(source_state_dict[name], target_state_dict) + if rename is not None: + print(f'"{name}": "{rename}",') + matched_keys.add(rename) + elif split_qkv and len(source_state_dict[name].shape)>=1 and source_state_dict[name].shape[0]%3==0: + length = source_state_dict[name].shape[0] // 3 + rename = [] + for i in range(3): + rename.append(search_parameter(source_state_dict[name][i*length: i*length+length], target_state_dict)) + if None not in rename: + print(f'"{name}": {rename},') + for rename_ in rename: + matched_keys.add(rename_) + for name in target_state_dict: + if name not in matched_keys: + print("Cannot find", name, target_state_dict[name].shape) + + +def search_for_files(folder, extensions): + files = [] + if os.path.isdir(folder): + for file in sorted(os.listdir(folder)): + files += search_for_files(os.path.join(folder, file), extensions) + elif os.path.isfile(folder): + for extension in extensions: + if folder.endswith(extension): + files.append(folder) + break + return files + + +def convert_state_dict_keys_to_single_str(state_dict, with_shape=True): + keys = [] + for key, value in state_dict.items(): + if isinstance(key, str): + if isinstance(value, torch.Tensor): + if with_shape: + shape = "_".join(map(str, list(value.shape))) + keys.append(key + ":" + shape) + keys.append(key) + elif isinstance(value, dict): + keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape)) + keys.sort() + keys_str = ",".join(keys) + return keys_str + + +def split_state_dict_with_prefix(state_dict): + keys = sorted([key for key in state_dict if isinstance(key, str)]) + prefix_dict = {} + for key in keys: + prefix = key if "." not in key else key.split(".")[0] + if prefix not in prefix_dict: + prefix_dict[prefix] = [] + prefix_dict[prefix].append(key) + state_dicts = [] + for prefix, keys in prefix_dict.items(): + sub_state_dict = {key: state_dict[key] for key in keys} + state_dicts.append(sub_state_dict) + return state_dicts + + +def hash_state_dict_keys(state_dict, with_shape=True): + keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape) + keys_str = keys_str.encode(encoding="UTF-8") + return hashlib.md5(keys_str).hexdigest() \ No newline at end of file diff --git a/groot/vla/model/dreamzero/modules/vram_management.py b/groot/vla/model/dreamzero/modules/vram_management.py new file mode 100644 index 0000000000000000000000000000000000000000..381d564caed84403b8b1062b53fed41b110668f2 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/vram_management.py @@ -0,0 +1,212 @@ +import torch, copy +from groot.vla.model.dreamzero.modules.utils import init_weights_on_device + + +def cast_to(weight, dtype, device): + r = torch.empty_like(weight, dtype=dtype, device=device) + r.copy_(weight) + return r + + +class AutoTorchModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def check_free_vram(self): + gpu_mem_state = torch.cuda.mem_get_info(self.computation_device) + used_memory = (gpu_mem_state[1] - gpu_mem_state[0]) / (1024 ** 3) + return used_memory < self.vram_limit + + def offload(self): + if self.state != 0: + self.to(dtype=self.offload_dtype, device=self.offload_device) + self.state = 0 + + def onload(self): + if self.state != 1: + self.to(dtype=self.onload_dtype, device=self.onload_device) + self.state = 1 + + def keep(self): + if self.state != 2: + self.to(dtype=self.computation_dtype, device=self.computation_device) + self.state = 2 + + +class AutoWrappedModule(AutoTorchModule): + def __init__(self, module: torch.nn.Module, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, **kwargs): + super().__init__() + self.module = module.to(dtype=offload_dtype, device=offload_device) + self.offload_dtype = offload_dtype + self.offload_device = offload_device + self.onload_dtype = onload_dtype + self.onload_device = onload_device + self.computation_dtype = computation_dtype + self.computation_device = computation_device + self.vram_limit = vram_limit + self.state = 0 + + def forward(self, *args, **kwargs): + if self.state == 2: + module = self.module + else: + if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: + module = self.module + elif self.vram_limit is not None and self.check_free_vram(): + self.keep() + module = self.module + else: + module = copy.deepcopy(self.module).to(dtype=self.computation_dtype, device=self.computation_device) + return module(*args, **kwargs) + + +class WanAutoCastLayerNorm(torch.nn.LayerNorm, AutoTorchModule): + def __init__(self, module: torch.nn.LayerNorm, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, **kwargs): + with init_weights_on_device(device=torch.device("meta")): + super().__init__(module.normalized_shape, eps=module.eps, elementwise_affine=module.elementwise_affine, bias=module.bias is not None, dtype=offload_dtype, device=offload_device) + self.weight = module.weight + self.bias = module.bias + self.offload_dtype = offload_dtype + self.offload_device = offload_device + self.onload_dtype = onload_dtype + self.onload_device = onload_device + self.computation_dtype = computation_dtype + self.computation_device = computation_device + self.vram_limit = vram_limit + self.state = 0 + + def forward(self, x, *args, **kwargs): + if self.state == 2: + weight, bias = self.weight, self.bias + else: + if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: + weight, bias = self.weight, self.bias + elif self.vram_limit is not None and self.check_free_vram(): + self.keep() + weight, bias = self.weight, self.bias + else: + weight = None if self.weight is None else cast_to(self.weight, self.computation_dtype, self.computation_device) + bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device) + with torch.amp.autocast(device_type=x.device.type): + x = torch.nn.functional.layer_norm(x.float(), self.normalized_shape, weight, bias, self.eps).type_as(x) + return x + + +class AutoWrappedLinear(torch.nn.Linear, AutoTorchModule): + def __init__(self, module: torch.nn.Linear, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, name="", **kwargs): + with init_weights_on_device(device=torch.device("meta")): + super().__init__(in_features=module.in_features, out_features=module.out_features, bias=module.bias is not None, dtype=offload_dtype, device=offload_device) + self.weight = module.weight + self.bias = module.bias + self.offload_dtype = offload_dtype + self.offload_device = offload_device + self.onload_dtype = onload_dtype + self.onload_device = onload_device + self.computation_dtype = computation_dtype + self.computation_device = computation_device + self.vram_limit = vram_limit + self.state = 0 + self.name = name + self.lora_A_weights = [] + self.lora_B_weights = [] + self.lora_merger = None + self.enable_fp8 = computation_dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz] + + def fp8_linear( + self, + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + device = input.device + origin_dtype = input.dtype + origin_shape = input.shape + input = input.reshape(-1, origin_shape[-1]) + + x_max = torch.max(torch.abs(input), dim=-1, keepdim=True).values + fp8_max = 448.0 + # For float8_e4m3fnuz, the maximum representable value is half of that of e4m3fn. + # To avoid overflow and ensure numerical compatibility during FP8 computation, + # we scale down the input by 2.0 in advance. + # This scaling will be compensated later during the final result scaling. + if self.computation_dtype == torch.float8_e4m3fnuz: + fp8_max = fp8_max / 2.0 + scale_a = torch.clamp(x_max / fp8_max, min=1.0).float().to(device=device) + scale_b = torch.ones((weight.shape[0], 1)).to(device=device) + input = input / (scale_a + 1e-8) + input = input.to(self.computation_dtype) + weight = weight.to(self.computation_dtype) + bias = bias.to(torch.bfloat16) + + result = torch._scaled_mm( + input, + weight.T, + scale_a=scale_a, + scale_b=scale_b.T, + bias=bias, + out_dtype=origin_dtype, + ) + new_shape = origin_shape[:-1] + result.shape[-1:] + result = result.reshape(new_shape) + return result + + def forward(self, x, *args, **kwargs): + # VRAM management + if self.state == 2: + weight, bias = self.weight, self.bias + else: + if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: + weight, bias = self.weight, self.bias + elif self.vram_limit is not None and self.check_free_vram(): + self.keep() + weight, bias = self.weight, self.bias + else: + weight = cast_to(self.weight, self.computation_dtype, self.computation_device) + bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device) + + # Linear forward + if self.enable_fp8: + out = self.fp8_linear(x, weight, bias) + else: + out = torch.nn.functional.linear(x, weight, bias) + + # LoRA + if len(self.lora_A_weights) == 0: + # No LoRA + return out + elif self.lora_merger is None: + # Native LoRA inference + for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights): + out = out + x @ lora_A.T @ lora_B.T + else: + # LoRA fusion + lora_output = [] + for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights): + lora_output.append(x @ lora_A.T @ lora_B.T) + lora_output = torch.stack(lora_output) + out = self.lora_merger(out, lora_output) + return out + + +def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, total_num_param=0, vram_limit=None, name_prefix=""): + for name, module in model.named_children(): + layer_name = name if name_prefix == "" else name_prefix + "." + name + for source_module, target_module in module_map.items(): + if isinstance(module, source_module): + num_param = sum(p.numel() for p in module.parameters()) + if max_num_param is not None and total_num_param + num_param > max_num_param: + module_config_ = overflow_module_config + else: + module_config_ = module_config + module_ = target_module(module, **module_config_, vram_limit=vram_limit, name=layer_name) + setattr(model, name, module_) + total_num_param += num_param + break + else: + total_num_param = enable_vram_management_recursively(module, module_map, module_config, max_num_param, overflow_module_config, total_num_param, vram_limit=vram_limit, name_prefix=layer_name) + return total_num_param + + +def enable_vram_management(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, vram_limit=None): + enable_vram_management_recursively(model, module_map, module_config, max_num_param, overflow_module_config, total_num_param=0, vram_limit=vram_limit) + model.vram_management_enabled = True \ No newline at end of file diff --git a/groot/vla/model/dreamzero/modules/wan2_1_attention.py b/groot/vla/model/dreamzero/modules/wan2_1_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..5ffb826bf16fa222a9fdf1a35ae27aeae0dbbf03 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan2_1_attention.py @@ -0,0 +1,348 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +import contextlib +import torch +from torch.profiler import profile, ProfilerActivity +import time +from typing import Optional +import os + +try: + import flash_attn_interface + + def is_hopper_gpu(): + if not torch.cuda.is_available(): + return False + device_name = torch.cuda.get_device_name(0).lower() + return "h100" in device_name or "hopper" in device_name + FLASH_ATTN_3_AVAILABLE = is_hopper_gpu() +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +try: + import transformer_engine + from groot.vla.model.dreamzero.modules.cudnn_attention import DotProductAttention + TRANSFORMER_ENGINE_AVAILABLE = True +except ModuleNotFoundError: + TRANSFORMER_ENGINE_AVAILABLE = False + +import warnings + + +def _gpu_supports_flash_attention(): + """FlashAttention requires Ampere (compute capability 8.0) or newer.""" + if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): + return False + try: + if not torch.cuda.is_available(): + return False + cap = torch.cuda.get_device_capability() + return cap[0] >= 8 + except Exception: + return False + + +def _sdpa_attention_fallback( + q, k, v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + dtype=torch.bfloat16, +): + """PyTorch SDPA fallback for GPUs that don't support FlashAttention (e.g. pre-Ampere).""" + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention on this GPU. ' + 'It can have a slight impact on quality.' + ) + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + if q_scale is not None: + q = q * q_scale + if softmax_scale is not None: + q = q * softmax_scale + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=None, is_causal=causal, dropout_p=dropout_p + ) + return out.transpose(1, 2).contiguous() + + +def flash_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: Optional[torch.Tensor] = None, + k_lens: Optional[torch.Tensor] = None, + dropout_p: float = 0., + softmax_scale: Optional[float] = None, + q_scale: Optional[float] = None, + causal: bool = False, + window_size: Optional[tuple[int, int]] = None, + deterministic: bool = False, + dtype: torch.dtype = torch.bfloat16, + version: Optional[int] = None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + version: int. 2 for flash attention 2, 3 for flash attention 3. + + Returns: + x: [B, Lq, Nq, C2]. + """ + if window_size is None: + window_size = (-1, -1) + if version is None: + version = 3 + + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == 'cuda' and q.size(-1) <= 256 + + # Use PyTorch SDPA on pre-Ampere GPUs (FlashAttention requires Ampere or newer) + if not _gpu_supports_flash_attention(): + return _sdpa_attention_fallback( + q, k, v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + dtype=dtype, + ) + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor([lq] * b, dtype=torch.int32, device=q.device) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor([lk] * b, dtype=torch.int32, device=k.device) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + + if version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn( + 'Flash attention 3 is not available, use flash attention 2 instead.' + ) + zeros = torch.zeros([1], dtype=torch.int32, device=q.device) + cu_seqlens_q = torch.cat([zeros, q_lens]).cumsum(0).to(torch.int32) + cu_seqlens_k = torch.cat([zeros, k_lens]).cumsum(0).to(torch.int32) + + # apply attention + if version == 3 and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p, window_size are not supported in FA3 now. + x = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic)[0].unflatten(0, (b, lq)) + elif FLASH_ATTN_2_AVAILABLE: + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic).unflatten(0, (b, lq)) + else: + raise ValueError(f"Invalid version: {version}") + + # output + return x.type(out_dtype) + + +class AttentionModule(torch.nn.Module): + def __init__( + self, + num_heads: int, + head_dim: int, + dropout_p: float = 0., + softmax_scale: Optional[float] = None, + q_scale: Optional[float] = None, + causal: bool = False, + window_size: Optional[tuple[int, int]] = None, + deterministic: bool = False, + dtype: torch.dtype = torch.bfloat16, + backend: Optional[str] = None, + ): + super().__init__() + if backend is None: + backend = "torch" + + if os.getenv("ATTENTION_BACKEND") is not None: + backend = os.getenv("ATTENTION_BACKEND") + else: + backend = "FA2" + + # Check for TensorRT at runtime, not import time + if os.getenv("ENABLE_TENSORRT", "False").lower() == "true": + backend = "torch" + + # Fall back to FA backend if TE is specified but not available + if backend == "TE" and not TRANSFORMER_ENGINE_AVAILABLE: + print("Warning: Transformer Engine is not available. Falling back to FA2 backend.") + backend = "FA2" + + assert backend in ["torch", "FA2", "FA3", "TE", "torch_onnx"] + self.backend = backend + + if backend == "torch": + def _torch_impl(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + out_dtype = q.dtype + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, + is_causal=causal, + dropout_p=dropout_p, + scale=softmax_scale, + ) + + out = out.transpose(1, 2).contiguous() + return out.to(out_dtype) + self.attn_func = _torch_impl + + elif backend == "torch_onnx": + def _torch_onnx_impl(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + out_dtype = q.dtype + # use torch.nn.functional.scaled_dot_product_attention for tensorrt export + + # The input is (s, n, d), but sdpa needs (b, h, s, d). + # We add a batch dimension and transpose. + q = q.unsqueeze(0).transpose(1, 2).to(dtype) + k = k.unsqueeze(0).transpose(1, 2).to(dtype) + v = v.unsqueeze(0).transpose(1, 2).to(dtype) + + # Fix for ONNX export: repeat k and v to match q's batch size in cross-attention + if q.shape[0] != k.shape[0] and k.shape[0] == 1: + k = k.repeat(q.shape[0], 1, 1, 1) + v = v.repeat(q.shape[0], 1, 1, 1) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, + is_causal=causal, + dropout_p=dropout_p, + scale=softmax_scale, + ) + + # Transpose back to (b, s, n, d) format. + out = out.transpose(1, 2).contiguous() + return out.to(out_dtype) + self.attn_func = _torch_onnx_impl + + elif backend == "TE" and TRANSFORMER_ENGINE_AVAILABLE: + self.attn_backend = DotProductAttention( + num_attention_heads=num_heads, + kv_channels=head_dim, + qkv_format="bshd", + attn_mask_type="causal" if causal else "no_mask", + window_size=window_size, + attention_dropout=dropout_p, + ) + + def _te_impl(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + out_dtype = q.dtype + return self.attn_backend( + query_layer=q.to(dtype), + key_layer=k.to(dtype), + value_layer=v.to(dtype), + ).to(out_dtype) + self.attn_func = _te_impl + + elif backend == "FA2" or backend == "FA3": + def _flash_attn_impl( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + q_lens: Optional[torch.Tensor], k_lens: Optional[torch.Tensor], + ) -> torch.Tensor: + return flash_attention( + q=q, k=k, v=v, + q_lens=q_lens, k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=3 if backend == "FA3" else 2, + ) + self.attn_func = _flash_attn_impl + + else: + raise ValueError(f"Invalid backend: {backend}") + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: Optional[torch.Tensor] = None, + k_lens: Optional[torch.Tensor] = None, + ): + if ( + self.backend == "torch" or + self.backend == "torch_onnx" or + (self.backend == "TE" and TRANSFORMER_ENGINE_AVAILABLE) + ): + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.' + ) + return self.attn_func(q, k, v) # type: ignore[call-arg] + else: + return self.attn_func(q, k, v, q_lens, k_lens) # type: ignore[call-arg] diff --git a/groot/vla/model/dreamzero/modules/wan2_1_submodule.py b/groot/vla/model/dreamzero/modules/wan2_1_submodule.py new file mode 100644 index 0000000000000000000000000000000000000000..b7bbe1187dddb4adf9ec2e9b2eb715b58c20e1e0 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan2_1_submodule.py @@ -0,0 +1,904 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import math + +import torch +import torch.nn as nn +import os +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from einops import repeat +from groot.vla.model.dreamzero.modules.attention import flash_attention + +__all__ = ['WanModel'] + +ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" + +def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor: + # preprocess + assert dim % 2 == 0 + half = dim // 2 + position = position.type(torch.float64) + + # calculation + sinusoid = torch.outer( + position, torch.pow(10000, -torch.arange(half, dtype=position.dtype, device=position.device).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +def rope_params(max_seq_len, dim, theta=10000): + if ENABLE_TENSORRT: + return rope_params_no_polar(max_seq_len, dim, theta) + else: + return rope_params_polar(max_seq_len, dim, theta) + + +# @amp.autocast(enabled=False) +def rope_params_polar(max_seq_len: int, dim: int, theta: float = 10000) -> torch.Tensor: + assert dim % 2 == 0 + freqs = torch.outer( + torch.arange(max_seq_len), + 1.0 / torch.pow(theta, + torch.arange(0, dim, 2).to(torch.float64).div(dim))) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + +def rope_params_no_polar(max_seq_len: int, dim: int, theta: float = 10000) -> torch.Tensor: + assert dim % 2 == 0 + inv_freq = 1.0 / torch.pow( + theta, + torch.arange(0, dim, 2).to(torch.float32) / dim + ) + t = torch.arange(max_seq_len, dtype=inv_freq.dtype) + freqs = torch.outer(t, inv_freq) + emb = torch.stack((freqs.cos(), freqs.sin()), dim=-1).flatten(-2) + return emb + +def rope_apply(x, grid_sizes, freqs): + if ENABLE_TENSORRT: + return rope_apply_no_polar(x, freqs) + else: + return rope_apply_polar(x, freqs) + +# @amp.autocast(enabled=False) +def rope_apply_polar(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + B, seq_len, n, _ = x.shape + + # precompute multipliers + x = torch.view_as_complex( + x.to(torch.float64).reshape(B, seq_len, n, -1, 2) + ) + + # apply rotary embedding + freqs = freqs.unsqueeze(0) + x = torch.view_as_real(x * freqs).flatten(3) + return x + +def rope_apply_no_polar(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + B, seq_len, n, D = x.shape + + # Reshape freqs to be broadcastable: (1, seq_len, 1, D) + freqs = freqs.unsqueeze(0).unsqueeze(2) + + x0, x1 = x.chunk(2, dim=-1) + freqs_cos, freqs_sin = freqs.chunk(2, dim=-1) + + rotated_x0 = x0 * freqs_cos - x1 * freqs_sin + rotated_x1 = x1 * freqs_cos + x0 * freqs_sin + x_rotated = torch.cat((rotated_x0, rotated_x1), dim=-1) + return x_rotated + + +def rope_action_apply(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block=32, num_state_per_block=1): + if ENABLE_TENSORRT: + return rope_action_apply_no_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block) + else: + return rope_action_apply_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block) + + +def rope_action_apply_no_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int, + num_action_per_block: int = 32, + num_state_per_block: int = 1, +) -> torch.Tensor: + B, seq_len, n, D = x.shape + + if action_register_length is not None: + chunk_size = action_register_length // (num_action_per_block + num_state_per_block) + freqs_1d_action = freqs_action[:chunk_size * num_action_per_block] + freqs_1d_state = freqs_state[:chunk_size * num_state_per_block] + freqs = torch.cat([freqs, freqs_1d_action, freqs_1d_state], dim=0) + + # Reshape freqs to be broadcastable: (1, seq_len, 1, D) + freqs = freqs.unsqueeze(0).unsqueeze(2) + + x0, x1 = x.chunk(2, dim=-1) + freqs_cos, freqs_sin = freqs.chunk(2, dim=-1) + + rotated_x0 = x0 * freqs_cos - x1 * freqs_sin + rotated_x1 = x1 * freqs_cos + x0 * freqs_sin + x_rotated = torch.cat((rotated_x0, rotated_x1), dim=-1) + + return x_rotated + + +# @amp.autocast(enabled=False) +def rope_action_apply_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + num_action_per_block: int | None = None, + num_state_per_block: int | None = None, +) -> torch.Tensor: + B, seq_len, n, _ = x.shape + + # precompute multipliers + x = torch.view_as_complex( + x.to(torch.float64).reshape(B, seq_len, n, -1, 2) + ) + + if action_register_length is not None: + assert num_action_per_block is not None + assert num_state_per_block is not None + + chunk_size = action_register_length // (num_action_per_block + num_state_per_block) + + freqs_1d_action = freqs_action[:chunk_size * num_action_per_block].view(chunk_size * num_action_per_block, 1, -1) + freqs_1d_state = freqs_state[:chunk_size * num_state_per_block].view(chunk_size * num_state_per_block, 1, -1) + freqs = torch.cat([freqs, freqs_1d_action, freqs_1d_state], dim=0) + + # apply rotary embedding + freqs = freqs.unsqueeze(0) + x = torch.view_as_real(x * freqs).flatten(3) + return x + + +class WanRMSNorm(nn.Module): + + def __init__(self, dim, eps=1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class WanLayerNorm(nn.LayerNorm): + + def __init__(self, dim, eps=1e-6, elementwise_affine=False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + +class WanSelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, seq_lens, freqs): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + x = flash_attention( + q=rope_apply(q, freqs), + k=rope_apply(k, freqs), + v=v, + k_lens=seq_lens, + window_size=self.window_size) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanT2VCrossAttention(WanSelfAttention): + + def forward(self, x, context, context_lens, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding. + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.q(x)).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + + # compute attention + x = flash_attention(q, k, v, k_lens=context_lens) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanGanCrossAttention(WanSelfAttention): + + def forward(self, x, context, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding. + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + qq = self.norm_q(self.q(context)).view(b, 1, -1, d) + + kk = self.norm_k(self.k(x)).view(b, -1, n, d) + vv = self.v(x).view(b, -1, n, d) + + # compute attention + x = flash_attention(qq, kk, vv) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanI2VCrossAttention(WanSelfAttention): + + def __init__(self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6): + super().__init__(dim, num_heads, window_size, qk_norm, eps) + + self.k_img = nn.Linear(dim, dim) + self.v_img = nn.Linear(dim, dim) + # self.alpha = nn.Parameter(torch.zeros((1, ))) + self.norm_k_img = WanRMSNorm( + dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, context, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + """ + context_img = context[:, :257] + context = context[:, 257:] + b, n, d = x.size(0), self.num_heads, self.head_dim + + q = self.norm_q(self.q(x)).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + # compute query, key, value + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + x = flash_attention(q, k, v, k_lens=None) + + k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d) + v_img = self.v_img(context_img).view(b, -1, n, d) + img_x = flash_attention(q, k_img, v_img, k_lens=None) + + # output + x = x.flatten(2) + img_x = img_x.flatten(2) + x = x + img_x + x = self.o(x) + return x + + +WAN_CROSSATTENTION_CLASSES = { + 't2v_cross_attn': WanT2VCrossAttention, + 'i2v_cross_attn': WanI2VCrossAttention, +} + + +class WanAttentionBlock(nn.Module): + + def __init__(self, + cross_attn_type, + dim, + ffn_dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, + eps) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, + num_heads, + (-1, -1), + qk_norm, + eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + e, + seq_lens, + grid_sizes, + freqs, + context, + context_lens, + ): + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, 6, C] + seq_lens(Tensor): Shape [B], length of each sequence in batch + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + # assert e.dtype == torch.float32 + # with amp.autocast(dtype=torch.float32): + e = (self.modulation + e).chunk(6, dim=1) + # assert e[0].dtype == torch.float32 + + # self-attention + y = self.self_attn( + self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes, + freqs, + ) + # with amp.autocast(dtype=torch.float32): + x = x + y * e[2] + + # cross-attention & ffn function + def cross_attn_ffn(x, context, context_lens, e): + x = x + self.cross_attn(self.norm3(x), context, context_lens) + y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3]) + # with amp.autocast(dtype=torch.float32): + x = x + y * e[5] + return x + + x = cross_attn_ffn(x, context, context_lens, e) + return x + + +class GanAttentionBlock(nn.Module): + + def __init__(self, + dim=1536, + ffn_dim=8192, + num_heads=12, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + # self.norm1 = WanLayerNorm(dim, eps) + # self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, + # eps) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + self.cross_attn = WanGanCrossAttention(dim, num_heads, + (-1, -1), + qk_norm, + eps) + + # modulation + # self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + context, + # seq_lens, + # grid_sizes, + # freqs, + # context, + # context_lens, + ): + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, 6, C] + seq_lens(Tensor): Shape [B], length of each sequence in batch + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + # assert e.dtype == torch.float32 + # with amp.autocast(dtype=torch.float32): + # e = (self.modulation + e).chunk(6, dim=1) + # assert e[0].dtype == torch.float32 + + # # self-attention + # y = self.self_attn( + # self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes, + # freqs) + # # with amp.autocast(dtype=torch.float32): + # x = x + y * e[2] + + # cross-attention & ffn function + def cross_attn_ffn(x, context): + token = context + self.cross_attn(self.norm3(x), context) + y = self.ffn(self.norm2(token)) + token # * (1 + e[4]) + e[3]) + # with amp.autocast(dtype=torch.float32): + # x = x + y * e[5] + return y + + x = cross_attn_ffn(x, context) + return x + + +class Head(nn.Module): + + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, C] + """ + # assert e.dtype == torch.float32 + # with amp.autocast(dtype=torch.float32): + e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1) + x = (self.head(self.norm(x) * (1 + e[1]) + e[0])) + return x + + +class MLPProj(torch.nn.Module): + + def __init__(self, in_dim, out_dim): + super().__init__() + + self.proj = torch.nn.Sequential( + torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim), + torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim), + torch.nn.LayerNorm(out_dim)) + + def forward(self, image_embeds): + clip_extra_context_tokens = self.proj(image_embeds) + return clip_extra_context_tokens + + +class RegisterTokens(nn.Module): + def __init__(self, num_registers: int, dim: int): + super().__init__() + self.register_tokens = nn.Parameter(torch.randn(num_registers, dim) * 0.02) + self.rms_norm = WanRMSNorm(dim, eps=1e-6) + + def forward(self): + return self.rms_norm(self.register_tokens) + + def reset_parameters(self): + nn.init.normal_(self.register_tokens, std=0.02) + + +class WanModel(ModelMixin, ConfigMixin): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + ignore_for_config = [ + 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size' + ] + _no_split_modules = ['WanAttentionBlock'] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__(self, + model_type='t2v', + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6): + r""" + Initialize the diffusion model backbone. + + Args: + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + window_size (`tuple`, *optional*, defaults to (-1, -1)): + Window size for local attention (-1 indicates global attention) + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + assert model_type in ['t2v', 'i2v'] + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + self.local_attn_size = 21 + + # embeddings + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), + nn.Linear(dim, dim)) + + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn' + self.blocks = nn.ModuleList([ + WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, + window_size, qk_norm, cross_attn_norm, eps) + for _ in range(num_layers) + ]) + + # head + self.head = Head(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + self.freqs = [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)) + ] + + if model_type == 'i2v': + self.img_emb = MLPProj(1280, dim) + + # initialize weights + self.init_weights() + + self.gradient_checkpointing = False + + def _set_gradient_checkpointing(self, module, value=False): + self.gradient_checkpointing = value + + def forward( + self, + x, + t, + context, + seq_len, + classify_mode=False, + concat_time_embeddings=False, + register_tokens=None, + cls_pred_branch=None, + gan_ca_blocks=None, + clip_fea=None, + y=None, + ): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + clip_fea (Tensor, *optional*): + CLIP image features for image-to-video mode + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_fea is not None and y is not None + # params + device = self.patch_embedding.weight.device + if any(freqs_i.device != device for freqs_i in self.freqs): + self.freqs = [freqs_i.to(device) for freqs_i in self.freqs] + + if y is not None: + x = [torch.cat([u, v.to(dtype=u.dtype)], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack( + [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + assert seq_lens.max() <= seq_len + x = torch.cat([ + torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], + dim=1) for u in x + ]) + + # time embeddings + # with amp.autocast(dtype=torch.float32): + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, t).type_as(x)) + e0 = self.time_projection(e).unflatten(1, (6, self.dim)) + # assert e.dtype == torch.float32 and e0.dtype == torch.float32 + + # context + context_lens = None + context = self.text_embedding( + torch.stack([ + torch.cat( + [u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) + for u in context + ])) + + if clip_fea is not None: + context_clip = self.img_emb(clip_fea) # bs x 257 x dim + context = torch.concat([context_clip, context], dim=1) + + # arguments + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens) + + def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + return custom_forward + + # TODO: Tune the number of blocks for feature extraction + final_x = None + if classify_mode: + assert register_tokens is not None + assert gan_ca_blocks is not None + assert cls_pred_branch is not None + + final_x = [] + registers = repeat(register_tokens(), "n d -> b n d", b=x.shape[0]) + # x = torch.cat([registers, x], dim=1) + + gan_idx = 0 + for ii, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, **kwargs, + use_reentrant=False, + ) + else: + x = block(x, **kwargs) + + if classify_mode and ii in [13, 21, 29]: + gan_token = registers[:, gan_idx: gan_idx + 1] + final_x.append(gan_ca_blocks[gan_idx](x, gan_token)) + gan_idx += 1 + + if classify_mode: + final_x = torch.cat(final_x, dim=1) + if concat_time_embeddings: + final_x = cls_pred_branch(torch.cat([final_x, 10 * e[:, None, :]], dim=1).view(final_x.shape[0], -1)) + else: + final_x = cls_pred_branch(final_x.view(final_x.shape[0], -1)) + + # head + x = self.head(x, e) + + # unpatchify + x = self.unpatchify(x, grid_sizes) + + if classify_mode: + return torch.stack(x), final_x + + return torch.stack(x) + + def unpatchify(self, x, grid_sizes, c=None): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + List[Tensor]: + Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] + """ + + c = self.out_dim if c is None else c + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum('fhwpqrc->cfphqwr', u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/groot/vla/model/dreamzero/modules/wan_video_camera_controller.py b/groot/vla/model/dreamzero/modules/wan_video_camera_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..8cbe658699e52c8cf3854d94ed4d0c356bb96989 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_camera_controller.py @@ -0,0 +1,202 @@ +import torch +import torch.nn as nn +import numpy as np +from einops import rearrange +import os +from typing_extensions import Literal + +class SimpleAdapter(nn.Module): + def __init__(self, in_dim, out_dim, kernel_size, stride, num_residual_blocks=1): + super(SimpleAdapter, self).__init__() + + # Pixel Unshuffle: reduce spatial dimensions by a factor of 8 + self.pixel_unshuffle = nn.PixelUnshuffle(downscale_factor=8) + + # Convolution: reduce spatial dimensions by a factor + # of 2 (without overlap) + self.conv = nn.Conv2d(in_dim * 64, out_dim, kernel_size=kernel_size, stride=stride, padding=0) + + # Residual blocks for feature extraction + self.residual_blocks = nn.Sequential( + *[ResidualBlock(out_dim) for _ in range(num_residual_blocks)] + ) + + def forward(self, x): + # Reshape to merge the frame dimension into batch + bs, c, f, h, w = x.size() + x = x.permute(0, 2, 1, 3, 4).contiguous().view(bs * f, c, h, w) + + # Pixel Unshuffle operation + x_unshuffled = self.pixel_unshuffle(x) + + # Convolution operation + x_conv = self.conv(x_unshuffled) + + # Feature extraction with residual blocks + out = self.residual_blocks(x_conv) + + # Reshape to restore original bf dimension + out = out.view(bs, f, out.size(1), out.size(2), out.size(3)) + + # Permute dimensions to reorder (if needed), e.g., swap channels and feature frames + out = out.permute(0, 2, 1, 3, 4) + + return out + + def process_camera_coordinates( + self, + direction: Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown"], + length: int, + height: int, + width: int, + speed: float = 1/54, + origin=(0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) + ): + if origin is None: + origin = (0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) + coordinates = generate_camera_coordinates(direction, length, speed, origin) + plucker_embedding = process_pose_file(coordinates, width, height) + return plucker_embedding + + + +class ResidualBlock(nn.Module): + def __init__(self, dim): + super(ResidualBlock, self).__init__() + self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, padding=1) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(dim, dim, kernel_size=3, padding=1) + + def forward(self, x): + residual = x + out = self.relu(self.conv1(x)) + out = self.conv2(out) + out += residual + return out + +class Camera(object): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + def __init__(self, entry): + fx, fy, cx, cy = entry[1:5] + self.fx = fx + self.fy = fy + self.cx = cx + self.cy = cy + w2c_mat = np.array(entry[7:]).reshape(3, 4) + w2c_mat_4x4 = np.eye(4) + w2c_mat_4x4[:3, :] = w2c_mat + self.w2c_mat = w2c_mat_4x4 + self.c2w_mat = np.linalg.inv(w2c_mat_4x4) + +def get_relative_pose(cam_params): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params] + abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params] + cam_to_origin = 0 + target_cam_c2w = np.array([ + [1, 0, 0, 0], + [0, 1, 0, -cam_to_origin], + [0, 0, 1, 0], + [0, 0, 0, 1] + ]) + abs2rel = target_cam_c2w @ abs_w2cs[0] + ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]] + ret_poses = np.array(ret_poses, dtype=np.float32) + return ret_poses + +def custom_meshgrid(*args): + # torch>=2.0.0 only + return torch.meshgrid(*args, indexing='ij') + + +def ray_condition(K, c2w, H, W, device): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + # c2w: B, V, 4, 4 + # K: B, V, 4 + + B = K.shape[0] + + j, i = custom_meshgrid( + torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype), + torch.linspace(0, W - 1, W, device=device, dtype=c2w.dtype), + ) + i = i.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5 # [B, HxW] + j = j.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5 # [B, HxW] + + fx, fy, cx, cy = K.chunk(4, dim=-1) # B,V, 1 + + zs = torch.ones_like(i) # [B, HxW] + xs = (i - cx) / fx * zs + ys = (j - cy) / fy * zs + zs = zs.expand_as(ys) + + directions = torch.stack((xs, ys, zs), dim=-1) # B, V, HW, 3 + directions = directions / directions.norm(dim=-1, keepdim=True) # B, V, HW, 3 + + rays_d = directions @ c2w[..., :3, :3].transpose(-1, -2) # B, V, 3, HW + rays_o = c2w[..., :3, 3] # B, V, 3 + rays_o = rays_o[:, :, None].expand_as(rays_d) # B, V, 3, HW + # c2w @ dirctions + rays_dxo = torch.linalg.cross(rays_o, rays_d) + plucker = torch.cat([rays_dxo, rays_d], dim=-1) + plucker = plucker.reshape(B, c2w.shape[1], H, W, 6) # B, V, H, W, 6 + # plucker = plucker.permute(0, 1, 4, 2, 3) + return plucker + + +def process_pose_file(cam_params, width=672, height=384, original_pose_width=1280, original_pose_height=720, device='cpu', return_poses=False): + if return_poses: + return cam_params + else: + cam_params = [Camera(cam_param) for cam_param in cam_params] + + sample_wh_ratio = width / height + pose_wh_ratio = original_pose_width / original_pose_height # Assuming placeholder ratios, change as needed + + if pose_wh_ratio > sample_wh_ratio: + resized_ori_w = height * pose_wh_ratio + for cam_param in cam_params: + cam_param.fx = resized_ori_w * cam_param.fx / width + else: + resized_ori_h = width / pose_wh_ratio + for cam_param in cam_params: + cam_param.fy = resized_ori_h * cam_param.fy / height + + intrinsic = np.asarray([[cam_param.fx * width, + cam_param.fy * height, + cam_param.cx * width, + cam_param.cy * height] + for cam_param in cam_params], dtype=np.float32) + + K = torch.as_tensor(intrinsic)[None] # [1, 1, 4] + c2ws = get_relative_pose(cam_params) # Assuming this function is defined elsewhere + c2ws = torch.as_tensor(c2ws)[None] # [1, n_frame, 4, 4] + plucker_embedding = ray_condition(K, c2ws, height, width, device=device)[0].permute(0, 3, 1, 2).contiguous() # V, 6, H, W + plucker_embedding = plucker_embedding[None] + plucker_embedding = rearrange(plucker_embedding, "b f c h w -> b f h w c")[0] + return plucker_embedding + + + +def generate_camera_coordinates( + direction: Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown"], + length: int, + speed: float = 1/54, + origin=(0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) +): + coordinates = [list(origin)] + while len(coordinates) < length: + coor = coordinates[-1].copy() + if "Left" in direction: + coor[9] += speed + if "Right" in direction: + coor[9] -= speed + if "Up" in direction: + coor[13] += speed + if "Down" in direction: + coor[13] -= speed + coordinates.append(coor) + return coordinates \ No newline at end of file diff --git a/groot/vla/model/dreamzero/modules/wan_video_dit.py b/groot/vla/model/dreamzero/modules/wan_video_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..305d3203cba05577ee0b2881de716999c8dc9e88 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_dit.py @@ -0,0 +1,818 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +import os +from typing import Tuple, Optional +from einops import rearrange +from groot.vla.model.dreamzero.modules.utils import hash_state_dict_keys +from groot.vla.model.dreamzero.modules.wan_video_camera_controller import SimpleAdapter +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +try: + from sageattention import sageattn + SAGE_ATTN_AVAILABLE = True +except ModuleNotFoundError: + SAGE_ATTN_AVAILABLE = False + + +def _gpu_supports_flash_attention(): + """FlashAttention requires Ampere (compute capability 8.0) or newer.""" + if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): + return False + try: + if not torch.cuda.is_available(): + return False + cap = torch.cuda.get_device_capability() + return cap[0] >= 8 + except Exception: + return False + + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin + +ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" +if ENABLE_TENSORRT: + # disable torch compile and transformer engine and flash attention for onnx/tensorrt export + FLASH_ATTN_COMPATIBILITY_MODE = True + DISABLE_TORCH_COMPILE = True +else: + DISABLE_TORCH_COMPILE = False + FLASH_ATTN_COMPATIBILITY_MODE = False +def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, compatibility_mode=False): + # Use PyTorch SDPA on pre-Ampere GPUs or when compatibility_mode (FlashAttention requires Ampere or newer) + if compatibility_mode or not _gpu_supports_flash_attention(): + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = F.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + elif FLASH_ATTN_3_AVAILABLE: + q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) + x = flash_attn_interface.flash_attn_func(q, k, v) + if isinstance(x,tuple): + x = x[0] + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + elif FLASH_ATTN_2_AVAILABLE: + q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) + x = flash_attn.flash_attn_func(q, k, v) + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + elif SAGE_ATTN_AVAILABLE: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = sageattn(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + else: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = F.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + return x + + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): + return (x * (1 + scale) + shift) + + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def sinusoidal_embedding_1d(dim, position): + sinusoid = torch.outer( + position.type(torch.float64), + torch.pow(10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2)) + ) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x.to(position.dtype) + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def rope_apply_no_polar_op(x, freqs_cis, num_heads): + freqs_cos, freqs_sin = freqs_cis + + x_reshaped = rearrange(x, "b s (h d) -> b s h d", h=num_heads) + x_as_complex = x_reshaped.float().reshape(*x_reshaped.shape[:-1], -1, 2) + x_real, x_imag = x_as_complex.unbind(-1) + + freqs_cos = freqs_cos.to(x.device).squeeze(1).unsqueeze(0).unsqueeze(2) + freqs_sin = freqs_sin.to(x.device).squeeze(1).unsqueeze(0).unsqueeze(2) + + x_out_real = x_real * freqs_cos - x_imag * freqs_sin + x_out_imag = x_real * freqs_sin + x_imag * freqs_cos + + x_out = torch.stack([x_out_real, x_out_imag], dim=-1).flatten(start_dim=-2) + + return rearrange(x_out, "b s h d -> b s (h d)").to(x.dtype) + + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def rope_apply_polar_op(x, freqs, num_heads): + x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) + x_out = torch.view_as_complex(x.to(torch.float64).reshape( + x.shape[0], x.shape[1], x.shape[2], -1, 2)) + freqs = freqs.to(x_out.device) + x_out = torch.view_as_real(x_out * freqs).flatten(2) + return x_out.to(x.dtype) + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + def forward(self, x): + dtype = x.dtype + return self.norm(x.float()).to(dtype) * self.weight + + +def _RMSNorm(normalized_shape, eps): + return RMSNorm(normalized_shape, eps=eps) + + +def RotaryPositionEmbedding(num_heads, head_dim): + if ENABLE_TENSORRT: + return RotaryPositionEmbeddingNoPolarOp(num_heads, head_dim) + else: + return RotaryPositionEmbeddingWithPolarOp(num_heads, head_dim) + + +def rope_apply(x, freqs, num_heads): + if ENABLE_TENSORRT: + return rope_apply_no_polar_op(x, freqs, num_heads) + else: + return rope_apply_polar_op(x, freqs, num_heads) + + +class RotaryPositionEmbeddingNoPolarOp(nn.Module): + def __init__(self, num_heads: int, head_dim: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + + self.freqs = self.precompute_freqs_cis_3d(head_dim) + + def precompute_freqs_cis_3d(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 3d rope precompute + f_freqs_cis = self.precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) + h_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + w_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + return {"f": f_freqs_cis, "h": h_freqs_cis, "w": w_freqs_cis} + + def precompute_freqs_cis(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 1d rope precompute + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + freqs = torch.outer(torch.arange(end, device=freqs.device), freqs) + freqs_real = torch.cos(freqs) + freqs_imag = torch.sin(freqs) + return (freqs_real, freqs_imag) + + def forward(self, f: int, h: int, w: int, a: int) -> torch.Tensor: + freqs_cos_3d = torch.cat( + [ + self.freqs["f"][0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs["h"][0][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs["w"][0][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(f * h * w, 1, -1) + freqs_sin_3d = torch.cat( + [ + self.freqs["f"][1][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs["h"][1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs["w"][1][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(f * h * w, 1, -1) + + + return torch.cat([freqs_cos_3d, freqs_sin_3d], dim=0) + + def post_initialize(self): + self.freqs = { + key: (value[0].to("cuda"), value[1].to("cuda")) for key, value in self.freqs.items() + } + + +class RotaryPositionEmbeddingWithPolarOp(nn.Module): + def __init__(self, num_heads: int, head_dim: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + + self.freqs = self.precompute_freqs_cis_3d(head_dim) + + def precompute_freqs_cis_3d(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 3d rope precompute + f_freqs_cis = self.precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) + h_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + w_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + return {"f": f_freqs_cis, "h": h_freqs_cis, "w": w_freqs_cis} + + def precompute_freqs_cis(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 1d rope precompute + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].double() / dim)) + freqs = torch.outer(torch.arange(end), freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 + return freqs_cis + + def forward(self, f: int, h: int, w: int, a: int) -> torch.Tensor: + freqs = torch.cat( + [ + self.freqs["f"][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs["h"][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs["w"][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(f * h * w, 1, -1) + return freqs + + def post_initialize(self): + self.freqs = {key: value.to(device="cuda") for key, value in self.freqs.items()} + + +class AttentionModule(nn.Module): + def __init__(self, num_heads): + super().__init__() + self.num_heads = num_heads + + def forward(self, q, k, v): + x = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads, compatibility_mode=FLASH_ATTN_COMPATIBILITY_MODE) + return x + + +class SelfAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = _RMSNorm(dim, eps=eps) + self.norm_k = _RMSNorm(dim, eps=eps) + + self.attn = AttentionModule(self.num_heads) + + def forward(self, x, freqs): + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(x)) + v = self.v(x) + q = rope_apply(q, freqs, self.num_heads) + k = rope_apply(k, freqs, self.num_heads) + x = self.attn(q, k, v) + return self.o(x) + + +class CrossAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, has_image_input: bool = False): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = _RMSNorm(dim, eps=eps) + self.norm_k = _RMSNorm(dim, eps=eps) + self.has_image_input = has_image_input + if has_image_input: + self.k_img = nn.Linear(dim, dim) + self.v_img = nn.Linear(dim, dim) + self.norm_k_img = _RMSNorm(dim, eps=eps) + + self.attn = AttentionModule(self.num_heads) + + def forward(self, x: torch.Tensor, y: torch.Tensor): + if self.has_image_input: + img = y[:, :257] + ctx = y[:, 257:] + else: + ctx = y + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(ctx)) + v = self.v(ctx) + x = self.attn(q, k, v) + if self.has_image_input: + k_img = self.norm_k_img(self.k_img(img)) + v_img = self.v_img(img) + y = flash_attention(q, k_img, v_img, num_heads=self.num_heads, compatibility_mode=FLASH_ATTN_COMPATIBILITY_MODE) + x = x + y + return self.o(x) + + +class GateModule(nn.Module): + def __init__(self,): + super().__init__() + + @torch.compile(disable=DISABLE_TORCH_COMPILE) + def forward(self, x, gate, residual): + return x + gate * residual + +class DiTBlock(nn.Module): + def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.ffn_dim = ffn_dim + + self.self_attn = SelfAttention(dim, num_heads, eps) + self.cross_attn = CrossAttention( + dim, num_heads, eps, has_image_input=has_image_input) + self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm3 = nn.LayerNorm(dim, eps=eps) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU( + approximate='tanh'), nn.Linear(ffn_dim, dim)) + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + self.gate = GateModule() + + def forward(self, x, context, t_mod, freqs): + # msa: multi-head self-attention mlp: multi-layer perceptron + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1) + input_x = modulate(self.norm1(x), shift_msa, scale_msa).to(dtype=t_mod.dtype) + x = self.gate(x, gate_msa, self.self_attn(input_x, freqs)) + x = x + self.cross_attn(self.norm3(x).to(dtype=t_mod.dtype), context.to(dtype=t_mod.dtype)) + input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) + x = self.gate(x, gate_mlp, self.ffn(input_x)) + return x + + +class MLP(torch.nn.Module): + def __init__(self, in_dim, out_dim, has_pos_emb=False): + super().__init__() + self.proj = torch.nn.Sequential( + nn.LayerNorm(in_dim), + nn.Linear(in_dim, in_dim), + nn.GELU(), + nn.Linear(in_dim, out_dim), + nn.LayerNorm(out_dim) + ) + self.has_pos_emb = has_pos_emb + if has_pos_emb: + self.emb_pos = torch.nn.Parameter(torch.zeros((1, 514, 1280))) + + def forward(self, x): + if self.has_pos_emb: + x = x + self.emb_pos.to(dtype=x.dtype, device=x.device) + return self.proj(x) + + +class Head(nn.Module): + def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float): + super().__init__() + self.dim = dim + self.patch_size = patch_size + self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.head = nn.Linear(dim, out_dim * math.prod(patch_size)) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, t_mod): + batch_size = x.shape[0] + shift, scale = (self.modulation.repeat(batch_size,1,1).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(1)).chunk(2, dim=1) + x = (self.head(self.norm(x) * (1 + scale) + shift)) + return x + + +class WanModel(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + dim: int, + in_dim: int, + ffn_dim: int, + out_dim: int, + freq_dim: int, + eps: float, + num_heads: int, + num_layers: int, + text_dim: int = 4096, + patch_size: Tuple[int, int, int] = [1, 2, 2], + has_image_input: bool = True, + has_image_pos_emb: bool = False, + has_ref_conv: bool = False, + add_control_adapter: bool = False, + in_dim_control_adapter: int = 24, + diffusion_model_pretrained_path: str = None, + ): + super().__init__() + self.dim = dim + self.freq_dim = freq_dim + self.has_image_input = has_image_input + self.patch_size = patch_size + self.diffusion_model_pretrained_path = diffusion_model_pretrained_path + + + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), + nn.GELU(approximate='tanh'), + nn.Linear(dim, dim) + ) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim) + ) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + self.blocks = nn.ModuleList([ + DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps) + for _ in range(num_layers) + ]) + self.head = Head(dim, out_dim, patch_size, eps) + head_dim = dim // num_heads + + self.rope = RotaryPositionEmbedding(num_heads=num_heads, head_dim=head_dim) + + if has_image_input: + self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb) # clip_feature_dim = 1280 + if has_ref_conv: + self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2)) + self.has_image_pos_emb = has_image_pos_emb + self.has_ref_conv = has_ref_conv + if add_control_adapter: + self.control_adapter = SimpleAdapter(in_dim_control_adapter, dim, kernel_size=patch_size[1:], stride=patch_size[1:]) + else: + self.control_adapter = None + + self.use_gradient_checkpointing = False + self.use_gradient_checkpointing_offload = False + + def patchify(self, x: torch.Tensor,control_camera_latents_input: torch.Tensor = None): + x = self.patch_embedding(x) + if self.control_adapter is not None and control_camera_latents_input is not None: + y_camera = self.control_adapter(control_camera_latents_input) + x = [u + v for u, v in zip(x, y_camera)] + x = x[0].unsqueeze(0) + grid_size = x.shape[2:] + x = rearrange(x, 'b c f h w -> b (f h w) c').contiguous() + return x, grid_size # x, grid_size: (f, h, w) + + def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor): + return rearrange( + x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)', + f=grid_size[0], h=grid_size[1], w=grid_size[2], + x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2] + ) + + def forward(self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + clip_feature: Optional[torch.Tensor] = None, + y: Optional[torch.Tensor] = None, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + **kwargs, + ): + t = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep)) + t_mod = self.time_projection(t).unflatten(1, (6, self.dim)) + context = self.text_embedding(context) + + if self.has_image_input: + x = torch.cat([x, y], dim=1) # (b, c_x + c_y, f, h, w) + clip_embdding = self.img_emb(clip_feature) + context = torch.cat([clip_embdding, context], dim=1) + # print("clip embedding shape", clip_embdding.shape) + # print("x before patchify", x.shape) + # clip embedding shape of B * 257 * dim (5120) + # context shape of B * 769 * dim (5120) - 257 + 512 + # x before patchify B * 36 * l_t * l_h * l_w + x, (f, h, w) = self.patchify(x) + # x after patchify B * 512 * 5120 + # f = 2, h = 16, w = 16 + # print("x and context shape", x.shape, context.shape, f,h,w) + + freqs = self.rope(f=f, h=h, w=w, a=x.shape[1]) + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + for block in self.blocks: + if self.training and self.use_gradient_checkpointing: + if self.use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = block(x, context, t_mod, freqs) + + # print("x before head", x.shape) + x = self.head(x, t) + # print("x before unpatchify", x.shape) + x = self.unpatchify(x, (f, h, w)) + # print("x after patchify", x.shape) + return x + def post_initialize(self): + self.rope.post_initialize() + @staticmethod + def state_dict_converter(): + return WanModelStateDictConverter() + + +class WanModelStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + rename_dict = { + "blocks.0.attn1.norm_k.weight": "blocks.0.self_attn.norm_k.weight", + "blocks.0.attn1.norm_q.weight": "blocks.0.self_attn.norm_q.weight", + "blocks.0.attn1.to_k.bias": "blocks.0.self_attn.k.bias", + "blocks.0.attn1.to_k.weight": "blocks.0.self_attn.k.weight", + "blocks.0.attn1.to_out.0.bias": "blocks.0.self_attn.o.bias", + "blocks.0.attn1.to_out.0.weight": "blocks.0.self_attn.o.weight", + "blocks.0.attn1.to_q.bias": "blocks.0.self_attn.q.bias", + "blocks.0.attn1.to_q.weight": "blocks.0.self_attn.q.weight", + "blocks.0.attn1.to_v.bias": "blocks.0.self_attn.v.bias", + "blocks.0.attn1.to_v.weight": "blocks.0.self_attn.v.weight", + "blocks.0.attn2.norm_k.weight": "blocks.0.cross_attn.norm_k.weight", + "blocks.0.attn2.norm_q.weight": "blocks.0.cross_attn.norm_q.weight", + "blocks.0.attn2.to_k.bias": "blocks.0.cross_attn.k.bias", + "blocks.0.attn2.to_k.weight": "blocks.0.cross_attn.k.weight", + "blocks.0.attn2.to_out.0.bias": "blocks.0.cross_attn.o.bias", + "blocks.0.attn2.to_out.0.weight": "blocks.0.cross_attn.o.weight", + "blocks.0.attn2.to_q.bias": "blocks.0.cross_attn.q.bias", + "blocks.0.attn2.to_q.weight": "blocks.0.cross_attn.q.weight", + "blocks.0.attn2.to_v.bias": "blocks.0.cross_attn.v.bias", + "blocks.0.attn2.to_v.weight": "blocks.0.cross_attn.v.weight", + "blocks.0.ffn.net.0.proj.bias": "blocks.0.ffn.0.bias", + "blocks.0.ffn.net.0.proj.weight": "blocks.0.ffn.0.weight", + "blocks.0.ffn.net.2.bias": "blocks.0.ffn.2.bias", + "blocks.0.ffn.net.2.weight": "blocks.0.ffn.2.weight", + "blocks.0.norm2.bias": "blocks.0.norm3.bias", + "blocks.0.norm2.weight": "blocks.0.norm3.weight", + "blocks.0.scale_shift_table": "blocks.0.modulation", + "condition_embedder.text_embedder.linear_1.bias": "text_embedding.0.bias", + "condition_embedder.text_embedder.linear_1.weight": "text_embedding.0.weight", + "condition_embedder.text_embedder.linear_2.bias": "text_embedding.2.bias", + "condition_embedder.text_embedder.linear_2.weight": "text_embedding.2.weight", + "condition_embedder.time_embedder.linear_1.bias": "time_embedding.0.bias", + "condition_embedder.time_embedder.linear_1.weight": "time_embedding.0.weight", + "condition_embedder.time_embedder.linear_2.bias": "time_embedding.2.bias", + "condition_embedder.time_embedder.linear_2.weight": "time_embedding.2.weight", + "condition_embedder.time_proj.bias": "time_projection.1.bias", + "condition_embedder.time_proj.weight": "time_projection.1.weight", + "patch_embedding.bias": "patch_embedding.bias", + "patch_embedding.weight": "patch_embedding.weight", + "scale_shift_table": "head.modulation", + "proj_out.bias": "head.head.bias", + "proj_out.weight": "head.head.weight", + } + state_dict_ = {} + for name, param in state_dict.items(): + if name in rename_dict: + state_dict_[rename_dict[name]] = param + else: + name_ = ".".join(name.split(".")[:1] + ["0"] + name.split(".")[2:]) + if name_ in rename_dict: + name_ = rename_dict[name_] + name_ = ".".join(name_.split(".")[:1] + [name.split(".")[1]] + name_.split(".")[2:]) + state_dict_[name_] = param + if hash_state_dict_keys(state_dict) == "cb104773c6c2cb6df4f9529ad5c60d0b": + config = { + "model_type": "t2v", + "patch_size": (1, 2, 2), + "text_len": 512, + "in_dim": 16, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "window_size": (-1, -1), + "qk_norm": True, + "cross_attn_norm": True, + "eps": 1e-6, + } + else: + config = {} + return state_dict_, config + + def from_civitai(self, state_dict): + state_dict = {name: param for name, param in state_dict.items() if not name.startswith("vace")} + if hash_state_dict_keys(state_dict) == "9269f8db9040a9d860eaca435be61814": + config = { + "has_image_input": False, + "patch_size": [1, 2, 2], + "in_dim": 16, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "aafcfd9672c3a2456dc46e1cb6e52c70": + config = { + "has_image_input": False, + "patch_size": [1, 2, 2], + "in_dim": 16, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "6bfcfb3b342cb286ce886889d519a77e": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "6d6ccde6845b95ad9114ab993d917893": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "6bfcfb3b342cb286ce886889d519a77e": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "349723183fc063b2bfc10bb2835cf677": + # 1.3B PAI control + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "efa44cddf936c70abd0ea28b6cbe946c": + # 14B PAI control + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "3ef3b1f8e1dab83d5b71fd7b617f859f": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6, + "has_image_pos_emb": True + } + elif hash_state_dict_keys(state_dict) == "70ddad9d3a133785da5ea371aae09504": + # 1.3B PAI control v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6, + "has_ref_conv": True + } + elif hash_state_dict_keys(state_dict) == "26bde73488a92e64cc20b0a7485b9e5b": + # 14B PAI control v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6, + "has_ref_conv": True + } + elif hash_state_dict_keys(state_dict) == "ac6a5aa74f4a0aab6f64eb9a72f19901": + # 1.3B PAI control-camera v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 32, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6, + "has_ref_conv": False, + "add_control_adapter": True, + "in_dim_control_adapter": 24, + } + elif hash_state_dict_keys(state_dict) == "b61c605c2adbd23124d152ed28e049ae": + # 14B PAI control-camera v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 32, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6, + "has_ref_conv": False, + "add_control_adapter": True, + "in_dim_control_adapter": 24, + } + else: + config = {} + return state_dict, config diff --git a/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py b/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..c5e2528401c0d7a1c69daafb29eb4ad50c06f30b --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py @@ -0,0 +1,2245 @@ +from typing import Any, TypeAlias + +from groot.vla.model.dreamzero.modules.wan2_1_attention import AttentionModule +from groot.vla.model.n1_5.modules.action_encoder import ( + SinusoidalPositionalEncoding, + swish, +) +from groot.vla.model.dreamzero.modules.wan2_1_submodule import ( + WanRMSNorm, + rope_action_apply, + WanLayerNorm, + WAN_CROSSATTENTION_CLASSES, + rope_params, + MLPProj, + sinusoidal_embedding_1d +) +from torch.nn.attention.flex_attention import create_block_mask, create_mask +from torch.nn.attention.flex_attention import BlockMask +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +import torch.nn as nn +import torch.nn.functional as F +import torch +import math +import torch.distributed as dist +import os + +ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" + + +class CategorySpecificLinear(nn.Module): + def __init__(self, num_categories, input_dim, hidden_dim): + super().__init__() + self.num_categories = num_categories + # For each category, we have separate weights and biases. + self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim)) + self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim)) + + def forward(self, x, cat_ids): + selected_W = self.W[cat_ids] + selected_b = self.b[cat_ids] + return torch.bmm(x, selected_W) + selected_b.unsqueeze(1) + + +class CategorySpecificMLP(nn.Module): + def __init__(self, num_categories, input_dim, hidden_dim, output_dim): + super().__init__() + self.num_categories = num_categories + self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim) + self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim) + + def forward(self, x, cat_ids): + hidden = F.relu(self.layer1(x, cat_ids)) + return self.layer2(hidden, cat_ids) + + +class MultiEmbodimentActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size, num_embodiments): + super().__init__() + self.hidden_size = hidden_size + self.num_embodiments = num_embodiments + + # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} + self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size) # (d -> w) + self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size) # (2w -> w) + self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size) # (w -> w) + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions, timesteps, cat_ids): + """ + actions: shape (B, T, action_dim) + timesteps: shape (B,) -- a single scalar per batch item + cat_ids: shape (B,) + returns: shape (B, T, hidden_size) + """ + B, T, _ = actions.shape + + # Standard action MLP step for shape => (B, T, w) + a_emb = self.W1(actions, cat_ids) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.W2(x, cat_ids)) + + # 5) Finally W3 => (B, T, w) + x = self.W3(x, cat_ids) + return x + + +def causal_rope_action_apply(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block, action_state_index): + if ENABLE_TENSORRT: + return causal_rope_action_apply_no_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block, action_state_index) + else: + return causal_rope_action_apply_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block, action_state_index) + + +def causal_rope_action_apply_no_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + num_action_per_block: int, + num_state_per_block: int, + action_state_index: int, +): + B, seq_len, n, d = x.shape + + # (B, seq_len, n, d) -> (B, seq_len, n, d/2, 2) + x = x.reshape(B, seq_len, n, -1, 2) + x_real = x[..., 0] + x_imag = x[..., 1] + + # Split freqs into cos and sin components + freqs = freqs.unsqueeze(0).view(1, freqs.shape[0], 1, -1, 2) + freqs_cos = freqs[..., 0] # Shape: (1, seq_len', 1, d/2) + freqs_sin = freqs[..., 1] # Shape: (1, seq_len', 1, d/2) + + # Handle the Action/State Register Frequencies + if action_register_length is not None: + assert action_register_length == (num_action_per_block + num_state_per_block) + + freqs_action_slice = freqs_action[ + action_state_index * num_action_per_block:(action_state_index + 1) * num_action_per_block + ] + freqs_state_slice = freqs_state[ + action_state_index * num_state_per_block:(action_state_index + 1) * num_state_per_block + ] + + # Combine the action/state tokens for this frame + freqs_1d = torch.cat([freqs_action_slice, freqs_state_slice], dim=0).view( + action_register_length, 1, -1, 2 + ) + + # Split the new action/state frequencies + freqs_cos_1d = freqs_1d[..., 0] + freqs_sin_1d = freqs_1d[..., 1] + + # Append the action/state register sin/cos to the main sequence sin/cos + freqs_cos = torch.cat([freqs_cos[0], freqs_cos_1d], dim=0).unsqueeze(0) + freqs_sin = torch.cat([freqs_sin[0], freqs_sin_1d], dim=0).unsqueeze(0) + + x_real_rotated = x_real * freqs_cos - x_imag * freqs_sin + x_imag_rotated = x_real * freqs_sin + x_imag * freqs_cos + + x_rotated = torch.stack((x_real_rotated, x_imag_rotated), dim=-1) + + return x_rotated.flatten(3) + +def causal_rope_action_apply_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + num_action_per_block: int, + num_state_per_block: int, + action_state_index: int, +): + B, seq_len, n, _ = x.shape + + # precompute multipliers + x = torch.view_as_complex( + x.to(torch.float64).reshape(B, seq_len, n, -1, 2) + ) + + if action_register_length is not None: + assert action_register_length == (num_action_per_block + num_state_per_block) + freqs_action = freqs_action[ + action_state_index * num_action_per_block:(action_state_index + 1) * num_action_per_block + ] + freqs_state = freqs_state[ + action_state_index * num_state_per_block:(action_state_index + 1) * num_state_per_block + ] + freqs_1d = torch.cat([freqs_action, freqs_state], dim=0).view(action_register_length, 1, -1) + freqs = torch.cat([freqs, freqs_1d], dim=0) + + # apply rotary embedding + freqs = freqs.unsqueeze(0) + x = torch.view_as_real(x * freqs).flatten(3) + + return x + + +class CausalWanSelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + frame_seqlen, + local_attn_size=-1, + sink_size=0, + num_frame_per_block=1, + qk_norm=True, + eps=1e-6, + num_action_per_block=32, + num_state_per_block=1): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.num_frame_per_block = num_frame_per_block + self.qk_norm = qk_norm + self.eps = eps + self.max_attention_size = 21 * frame_seqlen if local_attn_size == -1 else local_attn_size * frame_seqlen + self.frame_seqlen = frame_seqlen + self.num_action_per_block = num_action_per_block + self.num_state_per_block = num_state_per_block + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.attn = AttentionModule(num_heads=self.num_heads, head_dim=self.head_dim) + self.causal_attn = AttentionModule(num_heads=self.num_heads, head_dim=self.head_dim, causal=True) + + def _visualize_attention_mask(self, total_len, first_image_len, image_blocks_len, + action_len, state_len, num_image_blocks, + num_action_blocks, num_state_blocks, + num_frame_per_block, frame_seqlen, + num_action_per_block, num_state_per_block): + """ + Create and print a visualization of the attention mask pattern. + Returns a binary mask [total_len, total_len] where 1 = can attend, 0 = cannot attend. + """ + # Token ranges + first_image_start = 0 + first_image_end = first_image_len + image_blocks_start = first_image_end + image_blocks_end = image_blocks_start + image_blocks_len + action_start = image_blocks_end + action_end = action_start + action_len + state_start = action_end + state_end = state_start + state_len + + # Create mask tensor + mask = torch.zeros(total_len, total_len, dtype=torch.bool) + + # First image: self-attention only + mask[first_image_start:first_image_end, first_image_start:first_image_end] = True + + # Image blocks + for block_idx in range(num_image_blocks): + block_start = image_blocks_start + block_idx * num_frame_per_block * frame_seqlen + block_end = image_blocks_start + (block_idx + 1) * num_frame_per_block * frame_seqlen + + # Attend to first image + mask[block_start:block_end, first_image_start:first_image_end] = True + + # Attend to previous and current image blocks + if self.local_attn_size != -1: + image_kv_start = max(image_blocks_start, block_end - self.local_attn_size * frame_seqlen) + else: + image_kv_start = image_blocks_start + mask[block_start:block_end, image_kv_start:block_end] = True + + # Attend to current action block + action_block_start = action_start + block_idx * num_action_per_block + action_block_end = action_start + (block_idx + 1) * num_action_per_block + mask[block_start:block_end, action_block_start:action_block_end] = True + + # Attend to current state block + state_block_start = state_start + block_idx * num_state_per_block + state_block_end = state_start + (block_idx + 1) * num_state_per_block + mask[block_start:block_end, state_block_start:state_block_end] = True + + # Action blocks + for block_idx in range(num_action_blocks): + action_block_start = action_start + block_idx * num_action_per_block + action_block_end = action_start + (block_idx + 1) * num_action_per_block + + # Attend to first image + mask[action_block_start:action_block_end, first_image_start:first_image_end] = True + + # Attend to previous and current image blocks + image_block_end = image_blocks_start + (block_idx + 1) * num_frame_per_block * frame_seqlen + if self.local_attn_size != -1: + image_kv_start = max(image_blocks_start, image_block_end - self.local_attn_size * frame_seqlen) + else: + image_kv_start = image_blocks_start + mask[action_block_start:action_block_end, image_kv_start:image_block_end] = True + + # Self-attention + mask[action_block_start:action_block_end, action_block_start:action_block_end] = True + + # Attend to current state block + state_block_start = state_start + block_idx * num_state_per_block + state_block_end = state_start + (block_idx + 1) * num_state_per_block + mask[action_block_start:action_block_end, state_block_start:state_block_end] = True + + # State blocks: self-attention only + for block_idx in range(num_state_blocks): + state_block_start = state_start + block_idx * num_state_per_block + state_block_end = state_start + (block_idx + 1) * num_state_per_block + mask[state_block_start:state_block_end, state_block_start:state_block_end] = True + + return mask + + def _blockwise_causal_flash_attn(self, q, k, v, frame_seqlen, num_frame_per_block=1, + action_horizon=None, state_horizon=None, + num_action_per_block=None, num_state_per_block=None, + visualize_mask=False): + """ + Implement blockwise causal attention using flash_attention. + Matches the pattern from _prepare_blockwise_causal_attn_mask: + + Structure: + - First image: conditioning only, cannot attend to anything + - Image blocks: can attend to first image + previous image blocks + current action block + current state block + - Action blocks: can attend to previous image blocks + current image block + current state block + first image + - State blocks: conditioning only, cannot attend to anything + + Args: + q, k, v: Query, key, value tensors [B, L, num_heads, head_dim] + frame_seqlen: Number of tokens per frame + num_frame_per_block: Number of frames per attention block + action_horizon: Total number of action tokens (if None, no action/state tokens) + state_horizon: Total number of state tokens (if None, no action/state tokens) + num_action_per_block: Number of action tokens per block + num_state_per_block: Number of state tokens per block + visualize_mask: If True, print the attention mask pattern + + Returns: + Attention output [B, L, num_heads, head_dim] + """ + b, total_len, n, d = q.shape + + # Check if we have action/state tokens + has_action_state = (action_horizon is not None and state_horizon is not None) + + if not has_action_state: + # OPTIMIZED: Simple blockwise causal attention (without action/state tokens) + num_frames = total_len // frame_seqlen + block_size = frame_seqlen * num_frame_per_block + num_blocks = (num_frames - 1) // num_frame_per_block + + # Handle edge case when sequence is too short (no blocks to process) + if num_blocks <= 0: + # Process entire sequence as a single block + return self.attn(q, k, v) + + # OPTIMIZATION: For global attention, process all blocks in one call with causal masking + if self.local_attn_size == -1: + # Single flash_attention call with causal=True for all blocks at once + # This is much faster than looping! + return self.causal_attn(q, k, v) + + # With local attention, still need loop but optimize it + # Pre-allocate output tensor + output = torch.empty_like(q) + + # Pre-compute block boundaries + block_starts = [frame_seqlen + i * block_size for i in range(num_blocks)] + block_ends = [min(start + block_size, total_len) for start in block_starts] + kv_starts = [max(0, end - self.local_attn_size * frame_seqlen) for end in block_ends] + + for block_idx in range(num_blocks): + block_start = block_starts[block_idx] + block_end = block_ends[block_idx] + kv_start = kv_starts[block_idx] + + output[:, block_start:block_end] = self.attn( + q[:, block_start:block_end], + k[:, kv_start:block_end], + v[:, kv_start:block_end] + ) + + return output + + assert action_horizon is not None and state_horizon is not None + assert num_action_per_block is not None and num_state_per_block is not None + + # Multi-modal structure: [first image] [image blocks] [action blocks] [state blocks] + # Calculate block structure + first_image_len = frame_seqlen + action_len = action_horizon + state_len = state_horizon + image_blocks_len = total_len - first_image_len - action_len - state_len + + num_image_blocks = image_blocks_len // (num_frame_per_block * frame_seqlen) + num_action_blocks = action_horizon // num_action_per_block + num_state_blocks = state_horizon // num_state_per_block + + assert num_image_blocks == num_action_blocks == num_state_blocks + + # Token ranges + first_image_start = 0 + first_image_end = first_image_len + image_blocks_start = first_image_end + image_blocks_end = image_blocks_start + image_blocks_len + action_start = image_blocks_end + action_end = action_start + action_len + state_start = action_end + state_end = state_start + state_len + + # Visualize attention mask if requested + if visualize_mask: + mask = self._visualize_attention_mask( + total_len, first_image_len, image_blocks_len, + action_len, state_len, num_image_blocks, + num_action_blocks, num_state_blocks, + num_frame_per_block, frame_seqlen, + num_action_per_block, num_state_per_block + ) + + print("\n" + "="*80) + print("ATTENTION MASK VISUALIZATION") + print("="*80) + print(f"Total length: {total_len}") + print(f"First image: [{first_image_start}:{first_image_end}] (len={first_image_len})") + print(f"Image blocks: [{image_blocks_start}:{image_blocks_end}] (len={image_blocks_len}, num_blocks={num_image_blocks})") + print(f"Action tokens: [{action_start}:{action_end}] (len={action_len}, num_blocks={num_action_blocks})") + print(f"State tokens: [{state_start}:{state_end}] (len={state_len}, num_blocks={num_state_blocks})") + print(f"Local attention size: {self.local_attn_size}") + print("-"*80) + + # Print a downsampled version of the mask if it's too large + if total_len <= 100: + # Print full mask for small sequences + print("Attention mask (1=can attend, 0=cannot attend):") + print("Rows=Query tokens, Cols=Key tokens") + for i in range(total_len): + row = "".join(["1" if mask[i, j] else "." for j in range(total_len)]) + print(f"{i:4d}: {row}") + else: + # Print downsampled version for large sequences + downsample = max(1, total_len // 100) + print(f"Attention mask (downsampled by {downsample}x):") + print("Rows=Query tokens, Cols=Key tokens (1=can attend, .=cannot attend)") + for i in range(0, total_len, downsample): + row = "".join(["1" if mask[i, j] else "." for j in range(0, total_len, downsample)]) + print(f"{i:4d}: {row}") + + # Save mask as image + try: + import cv2 + import numpy as np + mask_np = mask.cpu().float().numpy() + # Resize for visualization if needed + if total_len > 1000: + mask_np = cv2.resize(mask_np, (1000, 1000), interpolation=cv2.INTER_NEAREST) + mask_img = (mask_np * 255).astype(np.uint8) + cv2.imwrite("attention_mask_blockwise_flash.png", mask_img) + print(f"\nMask saved to: attention_mask_blockwise_flash.png") + except Exception as e: + print(f"Could not save mask image: {e}") + + print("="*80 + "\n") + + # OPTIMIZED: Pre-allocate output tensor and pre-compute all indices + output = torch.empty_like(q) + + # Process first image (conditioning, can only self-attend) + output[:, first_image_start:first_image_end] = self.attn( + q[:, first_image_start:first_image_end], + k[:, first_image_start:first_image_end], + v[:, first_image_start:first_image_end] + ) + + # Pre-compute all block indices for image blocks + image_block_starts = [image_blocks_start + i * num_frame_per_block * frame_seqlen for i in range(num_image_blocks)] + image_block_ends = [image_blocks_start + (i + 1) * num_frame_per_block * frame_seqlen for i in range(num_image_blocks)] + if self.local_attn_size != -1: + image_kv_starts = [max(image_blocks_start, end - self.local_attn_size * frame_seqlen) for end in image_block_ends] + else: + image_kv_starts = [image_blocks_start] * num_image_blocks + + # Pre-compute action and state block indices + action_block_starts = [action_start + i * num_action_per_block for i in range(num_action_blocks)] + action_block_ends = [action_start + (i + 1) * num_action_per_block for i in range(num_action_blocks)] + state_block_starts = [state_start + i * num_state_per_block for i in range(num_state_blocks)] + state_block_ends = [state_start + (i + 1) * num_state_per_block for i in range(num_state_blocks)] + + # Process each image block + for block_idx in range(num_image_blocks): + block_start = image_block_starts[block_idx] + block_end = image_block_ends[block_idx] + image_kv_start = image_kv_starts[block_idx] + action_block_start = action_block_starts[block_idx] + action_block_end = action_block_ends[block_idx] + state_block_start = state_block_starts[block_idx] + state_block_end = state_block_ends[block_idx] + + # Build context: first image + relevant image blocks + current action + current state + k_context = torch.cat([ + k[:, first_image_start:first_image_end], # First image + k[:, image_kv_start:block_end], # Image blocks + k[:, action_block_start:action_block_end], # Current action block + k[:, state_block_start:state_block_end] # Current state block + ], dim=1) + v_context = torch.cat([ + v[:, first_image_start:first_image_end], + v[:, image_kv_start:block_end], + v[:, action_block_start:action_block_end], + v[:, state_block_start:state_block_end] + ], dim=1) + + output[:, block_start:block_end] = self.attn( + q[:, block_start:block_end], k_context, v_context + ) + + # Process each action block + for block_idx in range(num_action_blocks): + action_block_start = action_block_starts[block_idx] + action_block_end = action_block_ends[block_idx] + image_block_end = image_block_ends[block_idx] + state_block_start = state_block_starts[block_idx] + state_block_end = state_block_ends[block_idx] + + # Determine image context range + if self.local_attn_size != -1: + image_kv_start = max(image_blocks_start, image_block_end - self.local_attn_size * frame_seqlen) + else: + image_kv_start = image_blocks_start + + # Build context + k_context = torch.cat([ + k[:, first_image_start:first_image_end], # First image + k[:, image_kv_start:image_block_end], # Image blocks + k[:, action_block_start:action_block_end], # Current action block + k[:, state_block_start:state_block_end] # Current state block + ], dim=1) + v_context = torch.cat([ + v[:, first_image_start:first_image_end], + v[:, image_kv_start:image_block_end], + v[:, action_block_start:action_block_end], + v[:, state_block_start:state_block_end] + ], dim=1) + + output[:, action_block_start:action_block_end] = self.attn( + q[:, action_block_start:action_block_end], k_context, v_context + ) + + # Process state blocks (conditioning, can only self-attend) + for block_idx in range(num_state_blocks): + state_block_start = state_block_starts[block_idx] + state_block_end = state_block_ends[block_idx] + + output[:, state_block_start:state_block_end] = self.attn( + q[:, state_block_start:state_block_end], + k[:, state_block_start:state_block_end], + v[:, state_block_start:state_block_end] + ) + + return output + + def _process_clean_image_only(self, clean_image_q, clean_image_k, clean_image_v, clean_frames): + """Process clean image blocks with causal attention pattern - OPTIMIZED + + First frame: conditioning, cannot attend to anything (self-attention only) + Block i: attends to first frame + previous blocks (0 to i-1) + current block + + OPTIMIZATION: Instead of looping through blocks, we batch process them together + by using a single flash_attention call with properly structured KV cache. + """ + block_size = self.frame_seqlen * self.num_frame_per_block + num_blocks = (clean_frames - 1) // self.num_frame_per_block + + if num_blocks == 0: + # Only first frame - single attention call + return self.attn( + clean_image_q[:, :self.frame_seqlen], + clean_image_k[:, :self.frame_seqlen], + clean_image_v[:, :self.frame_seqlen] + ) + + # Pre-allocate output tensor (avoids list append + cat overhead) + b, total_len, n, d = clean_image_q.shape + output = torch.empty_like(clean_image_q) + + # First frame: conditioning, self-attention only + output[:, :self.frame_seqlen] = self.attn( + clean_image_q[:, :self.frame_seqlen], + clean_image_k[:, :self.frame_seqlen], + clean_image_v[:, :self.frame_seqlen] + ) + + # OPTIMIZATION: Process all blocks together with causal masking + # For global attention (no local_attn_size), we can process all blocks in one call + if self.local_attn_size == -1: + # Single attention call for all blocks! + # Each position can attend to first_frame + everything up to itself + blocks_q = clean_image_q[:, self.frame_seqlen:] + blocks_k = clean_image_k # Can attend to everything including first frame + blocks_v = clean_image_v + + # Use causal masking: each block token can see first frame + all previous tokens + output[:, self.frame_seqlen:] = self.causal_attn( + blocks_q, blocks_k, blocks_v + ) + else: + # With local attention, we still need to loop but with optimizations + # Pre-compute all block boundaries to reduce overhead + block_starts = [self.frame_seqlen + i * block_size for i in range(num_blocks)] + block_ends = [min(start + block_size, total_len) for start in block_starts] + + for block_idx in range(num_blocks): + block_start = block_starts[block_idx] + block_end = block_ends[block_idx] + + q_block = clean_image_q[:, block_start:block_end] + + # Context: first frame + recent blocks within local_attn_size + image_kv_start = max(self.frame_seqlen, block_end - self.local_attn_size * self.frame_seqlen) + k_context = torch.cat([ + clean_image_k[:, :self.frame_seqlen], # First frame + clean_image_k[:, image_kv_start:block_end] # Recent blocks + current + ], dim=1) + v_context = torch.cat([ + clean_image_v[:, :self.frame_seqlen], + clean_image_v[:, image_kv_start:block_end] + ], dim=1) + + output[:, block_start:block_end] = self.attn(q_block, k_context, v_context) + + return output + + def _process_state_blocks(self, state_q, state_k, state_v, state_horizon): + """Process state blocks: self-attention only - OPTIMIZED + + OPTIMIZATION: State blocks only do self-attention within each block. + Instead of looping, we can process all blocks in a single call with block-diagonal masking, + or even simpler: just one attention call since they're independent. + """ + num_blocks = state_horizon // self.num_state_per_block + + if num_blocks == 1: + # Single block - one attention call + return self.attn(state_q, state_k, state_v) + + # OPTIMIZATION: Since each state block only attends to itself (no cross-block attention), + # we can process all blocks in a single batched call. Flash attention will handle this + # efficiently. The blocks are independent, so this is safe. + # Alternative: reshape and process as separate batch items + + # Pre-allocate output + output = torch.empty_like(state_q) + + # Process all blocks (keeping loop for now due to block-diagonal pattern) + # This could be further optimized with custom masking + for block_idx in range(num_blocks): + state_block_start = block_idx * self.num_state_per_block + state_block_end = state_block_start + self.num_state_per_block + + output[:, state_block_start:state_block_end] = self.attn( + state_q[:, state_block_start:state_block_end], + state_k[:, state_block_start:state_block_end], + state_v[:, state_block_start:state_block_end] + ) + + return output + + def _process_noisy_image_blocks(self, noisy_image_q, noisy_image_k, noisy_image_v, + clean_image_k, clean_image_v, + noisy_action_k, noisy_action_v, noisy_state_k, noisy_state_v, + half_frames, action_horizon, state_horizon): + """Process noisy image blocks with teacher forcing pattern - OPTIMIZED + + First frame: conditioning, cannot attend to anything (self-attention only) + Block i: attends to action[i] + state[i] + first_clean_frame + clean_blocks[0:i] + current_noisy_block + + OPTIMIZATION: Pre-allocate output, pre-compute indices, reduce memory allocations + """ + block_size = self.frame_seqlen * self.num_frame_per_block + num_blocks = (half_frames - 1) // self.num_frame_per_block + + # Pre-allocate output tensor + output = torch.empty_like(noisy_image_q) + + # First noisy frame: conditioning, self-attention only + output[:, :self.frame_seqlen] = self.attn( + noisy_image_q[:, :self.frame_seqlen], + noisy_image_k[:, :self.frame_seqlen], + noisy_image_v[:, :self.frame_seqlen] + ) + + if num_blocks == 0: + return output + + # Pre-compute all block indices to reduce loop overhead + noisy_block_starts = [self.frame_seqlen + i * block_size for i in range(num_blocks)] + noisy_block_ends = [min(start + block_size, noisy_image_q.shape[1]) for start in noisy_block_starts] + clean_context_ends = [self.frame_seqlen + i * block_size for i in range(num_blocks)] + action_block_starts = [i * self.num_action_per_block for i in range(num_blocks)] + action_block_ends = [start + self.num_action_per_block for start in action_block_starts] + state_block_starts = [i * self.num_state_per_block for i in range(num_blocks)] + state_block_ends = [start + self.num_state_per_block for start in state_block_starts] + + # Process noisy image blocks + for block_idx in range(num_blocks): + noisy_start = noisy_block_starts[block_idx] + noisy_end = noisy_block_ends[block_idx] + clean_end = clean_context_ends[block_idx] + action_start = action_block_starts[block_idx] + action_end = action_block_ends[block_idx] + state_start = state_block_starts[block_idx] + state_end = state_block_ends[block_idx] + + q_block = noisy_image_q[:, noisy_start:noisy_end] + + # Build context: first_clean_frame + clean_blocks[0:i] + current_noisy_block + action[i] + state[i] + k_context = torch.cat([ + clean_image_k[:, :clean_end], + noisy_image_k[:, noisy_start:noisy_end], + noisy_action_k[:, action_start:action_end], + noisy_state_k[:, state_start:state_end] + ], dim=1) + v_context = torch.cat([ + clean_image_v[:, :clean_end], + noisy_image_v[:, noisy_start:noisy_end], + noisy_action_v[:, action_start:action_end], + noisy_state_v[:, state_start:state_end] + ], dim=1) + + output[:, noisy_start:noisy_end] = self.attn(q_block, k_context, v_context) + + return output + + def _process_noisy_action_blocks(self, noisy_action_q, noisy_action_k, noisy_action_v, + clean_image_k, clean_image_v, + noisy_image_k, noisy_image_v, + noisy_state_k, noisy_state_v, + half_frames, action_horizon, state_horizon): + """Process noisy action blocks with teacher forcing pattern - OPTIMIZED + + First action (for first frame): cannot attend to anything (self-attention only) + Action block i: attends to first_clean_frame + clean_blocks[0:i] + noisy_image[i] + action[i] + state[i] + + OPTIMIZATION: Pre-allocate output, pre-compute indices, reduce memory allocations + """ + num_blocks = (half_frames - 1) // self.num_frame_per_block + + if num_blocks == 0: + return torch.empty_like(noisy_action_q) + + # Pre-allocate output tensor + output = torch.empty_like(noisy_action_q) + + # Pre-compute all block indices + action_block_starts = [i * self.num_action_per_block for i in range(num_blocks)] + action_block_ends = [start + self.num_action_per_block for start in action_block_starts] + clean_context_ends = [self.frame_seqlen + i * self.frame_seqlen * self.num_frame_per_block for i in range(num_blocks)] + noisy_image_block_starts = [self.frame_seqlen + i * self.frame_seqlen * self.num_frame_per_block for i in range(num_blocks)] + noisy_image_block_ends = [start + self.frame_seqlen * self.num_frame_per_block for start in noisy_image_block_starts] + state_block_starts = [i * self.num_state_per_block for i in range(num_blocks)] + state_block_ends = [start + self.num_state_per_block for start in state_block_starts] + + # Process noisy action blocks + for block_idx in range(num_blocks): + action_start = action_block_starts[block_idx] + action_end = action_block_ends[block_idx] + clean_end = clean_context_ends[block_idx] + noisy_img_start = noisy_image_block_starts[block_idx] + noisy_img_end = noisy_image_block_ends[block_idx] + state_start = state_block_starts[block_idx] + state_end = state_block_ends[block_idx] + + q_block = noisy_action_q[:, action_start:action_end] + + # Build context: first_clean_frame + clean_blocks[0:i] + noisy_image[i] + action[i] + state[i] + k_context = torch.cat([ + clean_image_k[:, :clean_end], + noisy_image_k[:, noisy_img_start:noisy_img_end], + noisy_action_k[:, action_start:action_end], + noisy_state_k[:, state_start:state_end] + ], dim=1) + v_context = torch.cat([ + clean_image_v[:, :clean_end], + noisy_image_v[:, noisy_img_start:noisy_img_end], + noisy_action_v[:, action_start:action_end], + noisy_state_v[:, state_start:state_end] + ], dim=1) + + output[:, action_start:action_end] = self.attn(q_block, k_context, v_context) + + return output + + def forward( + self, + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + kv_cache: torch.Tensor | None = None, + current_start_frame: int = 0, + is_tf: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + updated_kv_cache: torch.Tensor | None = None + + if kv_cache is None: + if is_tf: + # Teacher forcing training. + if action_register_length is not None: + q_context = q[:, :(s-action_register_length)//2] + k_context = k[:, :(s-action_register_length)//2] + q_noisy = q[:, (s-action_register_length)//2:] + k_noisy = k[:, (s-action_register_length)//2:] + else: + q_context = q[:, :s//2] + k_context = k[:, :s//2] + q_noisy = q[:, s//2:] + k_noisy = k[:, s//2:] + roped_query = [] + roped_key = [] + + # rope should be same for clean and noisy parts + rq_context = rope_action_apply( + x=q_context, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=None, + ).type_as(v) + rk_context = rope_action_apply( + x=k_context, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=None, + ).type_as(v) + + rq_noisy = rope_action_apply( + x=q_noisy, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + rk_noisy = rope_action_apply( + x=k_noisy, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + + roped_query.append(rq_context) + roped_key.append(rk_context) + roped_query.append(rq_noisy) + roped_key.append(rk_noisy) + + roped_query = torch.cat(roped_query, dim=1) + roped_key = torch.cat(roped_key, dim=1) + # Calculate sequence dimensions + half_seq_len = (s - (action_register_length if action_register_length is not None else 0)) // 2 + + if action_register_length is not None: + # Teacher forcing structure: + # Clean half: [image tokens only] + # Noisy half: [image tokens][action tokens][state tokens] + # Causality only applies to image blocks! + + # Clean half contains ONLY image tokens + clean_image_seq_len = half_seq_len + clean_frames = clean_image_seq_len // self.frame_seqlen + + # Noisy half contains image + action + state tokens + noisy_image_seq_len = half_seq_len + noisy_frames = noisy_image_seq_len // self.frame_seqlen + num_image_blocks = (noisy_frames - 1) // self.num_frame_per_block + action_horizon = num_image_blocks * self.num_action_per_block + state_horizon = num_image_blocks * self.num_state_per_block + + # Block layout must match actual register length. For 5B use 320x176 so latent frame_seqlen=55. + if roped_query.shape[1] != half_seq_len + noisy_image_seq_len + action_horizon + state_horizon: + raise ValueError( + "Sequence length does not match block layout. " + "For 5B use 320x176 (e.g. data=dreamzero/droid_relative_wan22 or image_resolution_width=320, image_resolution_height=176). " + f"Got noisy_frames={noisy_frames}, num_image_blocks={num_image_blocks}, " + f"action_register_length={action_register_length}. " + "Ensure (noisy_frames - 1) // num_frame_per_block >= 1 and register length equals " + "num_blocks * (num_action_per_block + num_state_per_block)." + ) + + # Split clean and noisy parts + # Clean: [image tokens only] + clean_image_q = roped_query[:, :clean_image_seq_len] + clean_image_k = roped_key[:, :clean_image_seq_len] + clean_image_v = v[:, :clean_image_seq_len] + + # Noisy: [image tokens][action tokens][state tokens] + noisy_image_q = roped_query[:, half_seq_len:half_seq_len + noisy_image_seq_len] + noisy_action_q = roped_query[:, half_seq_len + noisy_image_seq_len:half_seq_len + noisy_image_seq_len + action_horizon] + noisy_state_q = roped_query[:, half_seq_len + noisy_image_seq_len + action_horizon:] + + noisy_image_k = roped_key[:, half_seq_len:half_seq_len + noisy_image_seq_len] + noisy_action_k = roped_key[:, half_seq_len + noisy_image_seq_len:half_seq_len + noisy_image_seq_len + action_horizon] + noisy_state_k = roped_key[:, half_seq_len + noisy_image_seq_len + action_horizon:] + + noisy_image_v = v[:, half_seq_len:half_seq_len + noisy_image_seq_len] + noisy_action_v = v[:, half_seq_len + noisy_image_seq_len:half_seq_len + noisy_image_seq_len + action_horizon] + noisy_state_v = v[:, half_seq_len + noisy_image_seq_len + action_horizon:] + + # ========== Process CLEAN (context) image tokens ========== + # Clean images: simple blockwise causal attention (no action/state) + clean_image_outputs = self._process_clean_image_only( + clean_image_q, clean_image_k, clean_image_v, clean_frames) + + # ========== Process NOISY tokens ========== + # Noisy image blocks: attend to previous clean image blocks + current noisy image + current noisy action + current noisy state + noisy_image_outputs = self._process_noisy_image_blocks( + noisy_image_q, noisy_image_k, noisy_image_v, + clean_image_k, clean_image_v, + noisy_action_k, noisy_action_v, noisy_state_k, noisy_state_v, + noisy_frames, action_horizon, state_horizon) + + # Noisy action blocks: attend to previous clean image blocks (including first) + current noisy image + current noisy action + same state + noisy_action_outputs = self._process_noisy_action_blocks( + noisy_action_q, noisy_action_k, noisy_action_v, + clean_image_k, clean_image_v, + noisy_image_k, noisy_image_v, + noisy_state_k, noisy_state_v, + noisy_frames, action_horizon, state_horizon) + + # Noisy state blocks: self-attention only + noisy_state_outputs = self._process_state_blocks( + noisy_state_q, noisy_state_k, noisy_state_v, state_horizon) + + # Concatenate all outputs in order: clean_img, noisy_img, noisy_act, noisy_state + x = torch.cat([ + clean_image_outputs, + noisy_image_outputs, noisy_action_outputs, noisy_state_outputs + ], dim=1) + else: + # No action/state tokens, fall back to simple image-only teacher forcing + half_frames = half_seq_len // self.frame_seqlen + clean_q = roped_query[:, :half_seq_len] + clean_k = roped_key[:, :half_seq_len] + clean_v = v[:, :half_seq_len] + noisy_q = roped_query[:, half_seq_len:] + noisy_k = roped_key[:, half_seq_len:] + noisy_v = v[:, half_seq_len:] + + # Process clean frames with blockwise causal attention + x_clean = self._blockwise_causal_flash_attn( + clean_q, clean_k, clean_v, self.frame_seqlen, self.num_frame_per_block, + action_horizon=None, state_horizon=None, + num_action_per_block=None, num_state_per_block=None, + visualize_mask=False) + + # Process noisy frames: attend to all clean frames + themselves + full_k = torch.cat([clean_k, noisy_k], dim=1) + full_v = torch.cat([clean_v, noisy_v], dim=1) + x_noisy = self.attn(noisy_q, full_k, full_v) + + x = torch.cat([x_clean, x_noisy], dim=1) + + else: + roped_query = rope_action_apply( + x=q, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + roped_key = rope_action_apply( + x=k, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + + # Calculate dynamic action and state horizons + if action_register_length is not None: + chunk_size = action_register_length // (self.num_action_per_block + self.num_state_per_block) + action_horizon = chunk_size * self.num_action_per_block + state_horizon = chunk_size * self.num_state_per_block + else: + action_horizon = None + state_horizon = None + + # Use blockwise causal flash attention without massive padding + visualize = False + x = self._blockwise_causal_flash_attn( + roped_query, roped_key, v, self.frame_seqlen, self.num_frame_per_block, + action_horizon=action_horizon, + state_horizon=state_horizon, + num_action_per_block=self.num_action_per_block if action_register_length else None, + num_state_per_block=self.num_state_per_block if action_register_length else None, + visualize_mask=visualize) + + else: + action_state_index = (current_start_frame - 1) // self.num_frame_per_block + + roped_query = causal_rope_action_apply( + x=q, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + action_state_index=action_state_index, + ).type_as(v) + roped_key = causal_rope_action_apply( + x=k, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + action_state_index=action_state_index, + ).type_as(v) + + # split roped_query and roped_action_query (the last action_register_length tokens) + roped_action_query: torch.Tensor | None = None + roped_action_key: torch.Tensor | None = None + action_v: torch.Tensor | None = None + + if action_register_length is not None: + roped_action_query = roped_query[:, -action_register_length:] + roped_query = roped_query[:, :-action_register_length] + roped_action_key = roped_key[:, -action_register_length:] + roped_key = roped_key[:, :-action_register_length] + action_v = v[:, -action_register_length:] + v = v[:, :-action_register_length] + assert roped_action_query is not None + assert roped_action_key is not None + assert action_v is not None + + num_new_tokens = roped_query.shape[1] + assert roped_key.shape[1] == num_new_tokens + assert v.shape[1] == num_new_tokens + + # If we are using local attention and the current KV cache size is larger + # than the local attention size, we need to truncate the KV cache + + updated_kv_cache = kv_cache + updated_k = updated_kv_cache[0] + updated_v = updated_kv_cache[1] + # Assign new keys/values directly up to current_end + new_k = torch.cat([updated_k, roped_key], dim=1) + new_v = torch.cat([updated_v, v], dim=1) + + # We may need to truncate the KV cache if it's size is larger than the max attention size. + new_k = new_k[:, -self.max_attention_size:] + new_v = new_v[:, -self.max_attention_size:] + + if action_register_length is not None: + x = self.attn( + torch.cat([roped_query, roped_action_query], dim=1), + torch.cat([new_k, roped_action_key], dim=1), + torch.cat([new_v, action_v], dim=1), + ) + else: + x = self.attn( + roped_query, + new_k, + new_v, + ) + updated_kv_cache = torch.stack([new_k, new_v], dim=0) + + + # output + x = x.flatten(2) + x = self.o(x) + return x, updated_kv_cache + + +class CausalWanAttentionBlock(nn.Module): + + def __init__(self, + cross_attn_type, + dim, + ffn_dim, + num_heads, + frame_seqlen, + local_attn_size=-1, + sink_size=0, + num_frame_per_block=1, + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + num_action_per_block=32, + num_state_per_block=1): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.local_attn_size = local_attn_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = CausalWanSelfAttention( + dim=dim, + num_heads=num_heads, + frame_seqlen=frame_seqlen, + local_attn_size=local_attn_size, + sink_size=sink_size, + num_frame_per_block=num_frame_per_block, + qk_norm=qk_norm, + eps=eps, + num_action_per_block=num_action_per_block, + num_state_per_block=num_state_per_block, + ) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, + num_heads, + (-1, -1), + qk_norm, + eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x: torch.Tensor, + e: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + context: torch.Tensor, + kv_cache: torch.Tensor | None = None, + crossattn_cache: torch.Tensor | None = None, + current_start_frame: int = 0, + is_tf: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, F, 6, C] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2) + + # Align modulation sequence length to x so mul/add broadcast (e.g. when F != L under compile) + L = x.shape[1] + aligned = [] + for part in e: + L_e = part.shape[1] + if L_e == L: + aligned.append(part) + elif L_e >= L: + aligned.append(part[:, :L]) + else: + repeat = (L + L_e - 1) // L_e + aligned.append(part.repeat_interleave(repeat, dim=1)[:, :L]) + e = tuple(aligned) + + # self-attention + y, updated_kv_cache = self.self_attn( + x=(self.norm1(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2)), + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + kv_cache=kv_cache, + is_tf=is_tf, + current_start_frame=current_start_frame, + ) + x = x + (y * e[2].squeeze(2)) + + # cross-attention & ffn function + def cross_attn_ffn(x, context, e): + x = x + self.cross_attn(self.norm3(x), context) + y = self.ffn( + (self.norm2(x) * (1 + e[4].squeeze(2)) + e[3].squeeze(2)) + ) + x = x + (y * e[5].squeeze(2)) + return x + + x = cross_attn_ffn(x, context, e) + return x, updated_kv_cache + + +class CausalHead(nn.Module): + + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, F, 1, C] + """ + e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2) + # Align modulation sequence length to x (e.g. when F != L1 under compile) + L = x.shape[1] + aligned = [] + for part in e: + L_e = part.shape[1] + if L_e == L: + aligned.append(part) + elif L_e >= L: + aligned.append(part[:, :L]) + else: + repeat = (L + L_e - 1) // L_e + aligned.append(part.repeat_interleave(repeat, dim=1)[:, :L]) + e = tuple(aligned) + x = (self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2))) + return x + + +class CausalWanModel(ModelMixin, ConfigMixin): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + ignore_for_config = [ + 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim' + ] + _no_split_modules = ['WanAttentionBlock'] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__(self, + model_type='t2v', + patch_size=(1, 2, 2), + frame_seqlen=220, + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + max_chunk_size=-1, + sink_size=0, + qk_norm=True, + cross_attn_norm=True, + eps=1e-6, + num_frame_per_block=1, + action_dim=32, + num_registers=8, + max_state_dim=64, + max_num_embodiments=32, + hidden_size=1024, + diffusion_model_pretrained_path=None, + num_action_per_block=32, + num_state_per_block=1, + concat_first_frame_latent=True): + r""" + Initialize the diffusion model backbone. + + Args: + concat_first_frame_latent (`bool`, *optional*, defaults to True): + If True, concat [x; y] before patch_embedding (14B I2V style). If False, latent only (5B pretrained style; first-frame via CLIP). + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + local_attn_size (`int`, *optional*, defaults to -1): + Window size for temporal local attention (-1 indicates global attention) + sink_size (`int`, *optional*, defaults to 0): + Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + assert model_type in ['t2v', 'i2v', 'ti2v'] + self.model_type = model_type + + self.patch_size = patch_size + self.frame_seqlen = frame_seqlen + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.local_attn_size = max_chunk_size * num_frame_per_block + 1 if max_chunk_size != -1 else -1 + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + self.num_frame_per_block = num_frame_per_block + self.diffusion_model_pretrained_path = diffusion_model_pretrained_path + self.action_dim = action_dim + self.num_registers = num_registers + self.max_state_dim = max_state_dim + self.max_num_embodiments = max_num_embodiments + self.hidden_size = hidden_size + self.num_action_per_block = num_action_per_block + self.num_state_per_block = num_state_per_block + self.concat_first_frame_latent = concat_first_frame_latent + + max_num_embodiments = 1 + + self.state_encoder = CategorySpecificMLP( + num_categories=max_num_embodiments, + input_dim=max_state_dim, + hidden_dim=self.hidden_size, + output_dim=self.dim, + ) + self.action_encoder = MultiEmbodimentActionEncoder( + action_dim=action_dim, + hidden_size=self.dim, + num_embodiments=max_num_embodiments, + ) + self.action_decoder = CategorySpecificMLP( + num_categories=max_num_embodiments, + input_dim=dim, + hidden_dim=self.hidden_size, + output_dim=action_dim, + ) + + # embeddings + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), + nn.Linear(dim, dim)) + + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn' + self.blocks = nn.ModuleList([ + CausalWanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, frame_seqlen, + self.local_attn_size, sink_size, num_frame_per_block, qk_norm, cross_attn_norm, eps, + num_action_per_block, num_state_per_block) + for _ in range(num_layers) + ]) + + # head + self.head = CausalHead(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + + self.freqs_action = rope_params(1024*10, d) + self.freqs_state = rope_params(1024, d) + self.freqs = [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + ] + if model_type in ('i2v', 'ti2v'): + self.img_emb = MLPProj(1280, dim) + + # initialize weights + self.init_weights() + + self.gradient_checkpointing = True + self.independent_first_frame = False if self.num_frame_per_block == 1 else True + + + def _set_gradient_checkpointing(self, module, value=False): + self.gradient_checkpointing = value + + @staticmethod + def _prepare_blockwise_causal_attn_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1, action_horizon=1, state_horizon=1, num_action_per_block=30, num_state_per_block=1 + ) -> BlockMask: + """ + We will divide the token sequence into the following format: + [first image (conditioning)] [image blocks] [action blocks] [state blocks] + + Structure: + - First image: conditioning only, cannot attend to anything + - Image blocks: can attend to first image + previous image block + current action block + current state block + - Action blocks: can attend to previous image block + current image block + current state block + - State blocks: conditioning only, cannot attend to anything + + Block alignment: + - num_image_blocks = (num_frames - 1) // num_frame_per_block + - num_action_blocks = action_horizon // num_action_per_block + - num_state_blocks = state_horizon // num_state_per_block + - num_image_blocks = num_action_blocks + 1 = num_state_blocks + 1 + """ + # Calculate block structure + num_image_blocks = (num_frames - 1) // num_frame_per_block + num_action_blocks = action_horizon // num_action_per_block + num_state_blocks = state_horizon // num_state_per_block + + # Verify the relationship: num_image_blocks = num_action_blocks + 1 = num_state_blocks + 1 + assert num_image_blocks == num_action_blocks, \ + f"image_blocks mismatch: {num_image_blocks} != {num_action_blocks}" + assert num_image_blocks == num_state_blocks, \ + f"image_blocks mismatch: {num_image_blocks} != {num_state_blocks}" + + # Token ranges + first_image_len = frame_seqlen # First image (conditioning) + image_blocks_len = num_image_blocks * num_frame_per_block * frame_seqlen + action_len = action_horizon + state_len = state_horizon + total_length = first_image_len + image_blocks_len + action_len + state_len + + # print("total_length", total_length, first_image_len, image_blocks_len, action_len, state_len) + # Padding to multiple of 128 + # padded_length = math.ceil(total_length / 128) * 128 - total_length + padded_length = math.ceil((local_attn_size * frame_seqlen + (local_attn_size - 1) + 32 * (local_attn_size - 1))/128) * 128 - total_length + total_padded_length = total_length + padded_length + # print("total_padded_length", total_padded_length, total_length, padded_length) + + # Define token ranges for each modality + first_image_start = 0 + first_image_end = first_image_len + image_blocks_start = first_image_end + image_blocks_end = image_blocks_start + image_blocks_len + action_start = image_blocks_end + action_end = action_start + action_len + state_start = action_end + state_end = state_start + state_len + + # Precompute block indices for each token + block_indices = torch.zeros(total_padded_length, device=device, dtype=torch.long) + + # First image gets special block index -1 (conditioning, cannot attend to anything) + block_indices[first_image_start:first_image_end] = -1 + + # Assign block indices for image blocks (0 to num_image_blocks-1) + for block_idx in range(num_image_blocks): + start_idx = image_blocks_start + block_idx * num_frame_per_block * frame_seqlen + end_idx = image_blocks_start + (block_idx + 1) * num_frame_per_block * frame_seqlen + block_indices[start_idx:end_idx] = block_idx + + # Assign block indices for action tokens (0 to num_action_blocks-1) + for block_idx in range(num_action_blocks): + start_idx = action_start + block_idx * num_action_per_block + end_idx = action_start + (block_idx + 1) * num_action_per_block + block_indices[start_idx:end_idx] = block_idx + + # Assign block indices for state tokens (0 to num_state_blocks-1) + for block_idx in range(num_state_blocks): + start_idx = state_start + block_idx * num_state_per_block + end_idx = state_start + (block_idx + 1) * num_state_per_block + block_indices[start_idx:end_idx] = block_idx + + # Padding tokens get block index of last block + 1 (won't attend to anything) + block_indices[total_length:] = num_image_blocks + + def attention_mask(b, h, q_idx, kv_idx): + # Self-attention + self_attn = (q_idx == kv_idx) + + # Determine which modality q and kv belong to + q_is_first_image = (q_idx >= first_image_start) & (q_idx < first_image_end) + q_is_image_block = (q_idx >= image_blocks_start) & (q_idx < image_blocks_end) + q_is_action = (q_idx >= action_start) & (q_idx < action_end) + q_is_state = (q_idx >= state_start) & (q_idx < state_end) + + kv_is_first_image = (kv_idx >= first_image_start) & (kv_idx < first_image_end) + kv_is_image_block = (kv_idx >= image_blocks_start) & (kv_idx < image_blocks_end) + kv_is_action = (kv_idx >= action_start) & (kv_idx < action_end) + kv_is_state = (kv_idx >= state_start) & (kv_idx < state_end) + + q_block = block_indices[q_idx] + kv_block = block_indices[kv_idx] + + # First image query (conditioning) - cannot attend to anything + first_image_mask = q_is_first_image & False + + # Image block query + image_to_first = q_is_image_block & kv_is_first_image # Image block to first image: always allowed + image_to_image = q_is_image_block & kv_is_image_block & (kv_block <= q_block) # Image block to image block: can attend to current and previous image blocks + image_to_action = q_is_image_block & kv_is_action & (kv_block == q_block) # Image block to action: can attend to current action block + image_to_state = q_is_image_block & kv_is_state & (kv_block == q_block) # Image block to state: can attend to current state block + + image_block_mask = image_to_first | image_to_image | image_to_action | image_to_state + + # Action query + action_to_image = q_is_action & kv_is_image_block & (kv_block <= q_block) # Action to image block: can attend to current and all previous image blocks + action_to_action = q_is_action & kv_is_action & (kv_block == q_block) # Action to action: only same block + action_to_state = q_is_action & kv_is_state & (kv_block == q_block) # Action to state: only same block + action_to_first = q_is_action & kv_is_first_image # Action to first image: always allowed + + action_mask = action_to_image | action_to_action | action_to_state | action_to_first + + # State query (conditioning) - cannot attend to anything + state_mask = q_is_state & False + + # Combine all masks + return self_attn | first_image_mask | image_block_mask | action_mask | state_mask + + block_mask = create_block_mask( + attention_mask, B=None, H=None, + Q_LEN=total_padded_length, + KV_LEN=total_padded_length, + _compile=False, device=device + ) + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f"Created blockwise causal attention mask:") + print(f" first_image_tokens={first_image_len} (conditioning)") + print(f" num_image_blocks={num_image_blocks} (blocks of {num_frame_per_block * frame_seqlen})") + print(f" num_action_blocks={num_action_blocks} (blocks of {num_action_per_block})") + print(f" num_state_blocks={num_state_blocks} (blocks of {num_state_per_block})") + print(f" total_length={total_length}, padded_length={padded_length}") + print(block_mask) + + # Debug: materialize a small slice of the mask into 0/1 strings + try: + dense_mask = create_mask( + attention_mask, + B=None, + H=None, + Q_LEN=total_padded_length, + KV_LEN=total_padded_length, + device=device, + )[0, 0] # [Q, K] + preview_q = min(979, dense_mask.shape[0]) + preview_k = min(979, dense_mask.shape[1]) + print("Block mask (preview):") + for qi in range(preview_q): + row = dense_mask[qi, :preview_k].to(torch.int8).tolist() + print(" ".join(str(int(v)) for v in row)) + except Exception as err: + print("[warn] Failed to materialize block mask preview:", err) + + return block_mask + + @staticmethod + def _prepare_teacher_forcing_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1 + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [1 latent frame] ... [1 latent frame] + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen * 2 + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(self.local_attn_size * frame_seqlen/128) * 128 - total_length + # padded_length = math.ceil(total_length / 128) * 128 - total_length + + clean_ends = num_frames * frame_seqlen + # for clean context frames, we can construct their flex attention mask based on a [start, end] interval + context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + # for noisy frames, we need two intervals to construct the flex attention mask [context_start, context_end] [noisy_start, noisy_end] + noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + attention_block_size = frame_seqlen * num_frame_per_block + frame_indices = torch.arange( + start=0, + end=num_frames * frame_seqlen, + step=attention_block_size, + device=device, dtype=torch.long + ) + + # attention for clean context frames + for start in frame_indices: + context_ends[start:start + attention_block_size] = start + attention_block_size + + noisy_image_start_list = torch.arange( + num_frames * frame_seqlen, total_length, + step=attention_block_size, + device=device, dtype=torch.long + ) + noisy_image_end_list = noisy_image_start_list + attention_block_size + + # attention for noisy frames + for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)): + # attend to noisy tokens within the same block + noise_noise_starts[start:end] = start + noise_noise_ends[start:end] = end + # attend to context tokens in previous blocks + # noise_context_starts[start:end] = 0 + noise_context_ends[start:end] = block_index * attention_block_size + + def attention_mask(b, h, q_idx, kv_idx): + # first design the mask for clean frames + clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx]) + # then design the mask for noisy frames + # noisy frames will attend to all clean preceeding clean frames + itself + C1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx]) + C2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx]) + noise_mask = (q_idx >= clean_ends) & (C1 | C2) + + eye_mask = q_idx == kv_idx + return eye_mask | clean_mask | noise_mask + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + + if DEBUG: + print(block_mask) + import imageio + import numpy as np + from torch.nn.attention.flex_attention import create_mask + + mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + + padded_length, KV_LEN=total_length + padded_length, device=device) + import cv2 + mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) + imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) + + return block_mask + + @staticmethod + def _prepare_blockwise_causal_attn_mask_i2v( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=4, local_attn_size=-1 + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [N latent frame] ... [N latent frame] + The first frame is separated out to support I2V generation + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(local_attn_size * frame_seqlen/128) * 128 - total_length + # padded_length = math.ceil(total_length / 128) * 128 - total_length + + ends = torch.zeros(total_length + padded_length, + device=device, dtype=torch.long) + + # special handling for the first frame + ends[:frame_seqlen] = frame_seqlen + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + frame_indices = torch.arange( + start=frame_seqlen, + end=total_length, + step=frame_seqlen * num_frame_per_block, + device=device + ) + + for idx, tmp in enumerate(frame_indices): + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ + frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) + else: + return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | \ + (q_idx == kv_idx) + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print( + f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + print(block_mask) + + return block_mask + + def _forward_blocks( + self, + x: torch.Tensor, + seq_len: int, + freqs: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + clip_feature: torch.Tensor | None, + embodiment_id: torch.Tensor | None, + action: torch.Tensor | None, + timestep_action: torch.Tensor | None, + state: torch.Tensor | None, + kv_cache: list[torch.Tensor], + current_start_frame: int, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: + r""" + Forward pass through the diffusion model blocks. + """ + x = x.flatten(start_dim=2).transpose(1, 2) + + B = x.shape[0] + F = timestep.shape[1] + + if action is not None: + embodiment_id = torch.tensor([0], device=x.device).repeat(x.shape[0]) + action_features = self.action_encoder(action, timestep_action, embodiment_id) + state_features = self.state_encoder(state, embodiment_id) + action_register = torch.cat([action_features, state_features], dim=1) + action_length = action_features.shape[1] + action_register_length = action_register.shape[1] + x = torch.cat([x, action_register], dim=1) + else: + action_features = None + state_features = None + action_length = 0 + action_register_length = None + + # time embeddings: expand to exactly seq_len so e matches x (5B: frame_seqlen=50, 1 frame -> 50 tokens) + if F <= seq_len: + repeat = (seq_len + F - 1) // F + timestep = timestep.repeat_interleave(repeat, dim=1)[:, :seq_len] + else: + indices = torch.linspace(0, F - 1, seq_len, device=timestep.device, dtype=torch.long) + timestep = timestep[:, indices] + + if action is not None: + assert timestep_action is not None + assert state_features is not None + stride = timestep_action.shape[1] // state_features.shape[1] + timestep_state = timestep_action[:, ::stride] + timestep = torch.cat([timestep, timestep_action, timestep_state], dim=1) + + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).type_as(x)) + e = e.unflatten(dim=0, sizes=(B, -1)) + e0 = self.time_projection(e) + e0 = e0.unflatten(dim=2, sizes=(6, self.dim)) + + # context + context = self.text_embedding(context) + + if clip_feature is not None: + clip_embedding = self.img_emb(clip_feature) + context = torch.cat([clip_embedding, context], dim=1) + + updated_kv_caches: list[torch.Tensor] = [] + for block_index, block in enumerate(self.blocks): + x, updated_kv_cache = block( + x=x, + e=e0, + freqs=freqs, + freqs_action=self.freqs_action, + freqs_state=self.freqs_state, + context=context, + action_register_length=action_register_length, + kv_cache=kv_cache[block_index], + current_start_frame=current_start_frame, + ) + updated_kv_caches.append(updated_kv_cache) + + if action is not None: + action_noise_pred = x[:, seq_len: seq_len + action_length] + action_noise_pred = self.action_decoder(action_noise_pred, embodiment_id) + else: + action_noise_pred = None + + # Build a tensor that contains only video tokens per sample with length = max(video_lens) + x_video = x[:, :seq_len] + e_video = e[:, :seq_len] + + # Unpatchify video-only tokens + x_video = self.head(x_video, e_video.unsqueeze(2)) + + return x_video, action_noise_pred, updated_kv_caches + + + def _forward_inference_trt( + self, + x, + timestep, + context, + kv_cache_packed: torch.Tensor, + y, + clip_feature, + action, + timestep_action, + state, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + + + frame_seqlen = 880 + seq_len = 2*frame_seqlen + kv_cache_seq_len = kv_cache_packed.shape[3] + current_start_frame = kv_cache_seq_len // frame_seqlen + + kv_cache_list = [] + for block_index in range(len(self.blocks)): + kv_cache_list.append(kv_cache_packed[block_index]) + + x_video, action_noise_pred, _ = self._forward_inference( + x=x, + timestep=timestep, + context=context, + seq_len=int(seq_len), + kv_cache=kv_cache_list, + crossattn_cache=None, + y=y, + clip_feature=clip_feature, + action=action, + timestep_action=timestep_action, + state=state, + current_start_frame = current_start_frame, + ) + + return x_video, action_noise_pred + + def _forward_inference_trt_droid( + self, + x, + timestep, + context, + kv_cache_packed: torch.Tensor, + y, + clip_feature, + action, + timestep_action, + state, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + + + frame_seqlen = 880 + seq_len = 2*frame_seqlen + kv_cache_seq_len = kv_cache_packed.shape[3] + current_start_frame = kv_cache_seq_len // frame_seqlen + + kv_cache_list = [] + for block_index in range(len(self.blocks)): + kv_cache_list.append(kv_cache_packed[block_index]) + + x_video, action_noise_pred, _ = self._forward_inference( + x=x, + timestep=timestep, + context=context, + seq_len=int(seq_len), + kv_cache=kv_cache_list, + crossattn_cache=None, + y=y, + clip_feature=clip_feature, + action=action, + timestep_action=timestep_action, + state=state, + current_start_frame = current_start_frame, + ) + + return x_video, action_noise_pred + + + def _forward_inference( + self, + x, + timestep, + context, + seq_len, + kv_cache: list[torch.Tensor], + crossattn_cache: list[torch.Tensor], + current_start_frame: int, + y=None, + clip_feature=None, + action=None, + timestep_action=None, + state=None, + embodiment_id=None, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: + r""" + Run the diffusion model with kv caching. + See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. + This function will be run for num_frame times. + Process the latent frames one by one (1560 tokens each) + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + timestep (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + action (Tensor, *optional*): + Action tensor of shape [B, H, D] + state (Tensor, *optional*): + State tensor of shape [B, H, D] + embodiment_id (Tensor, *optional*): + Embodiment ID tensor of shape [B] + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + clip_feature (Tensor, *optional*): + CLIP image features for image-to-video mode + timestep_action (Tensor, *optional*): + Action timestep tensor of shape [B] + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_feature is not None and y is not None + assert context.shape[1] == self.text_len + + # Concat [x; y] only when pretrained that way (14B). 5B uses latent only, first-frame via CLIP. + if y is not None and self.concat_first_frame_latent: + x = torch.cat([x, y.to(dtype=x.dtype)], dim=1) + + # embeddings + x = self.patch_embedding(x) + grid_size = torch.tensor(x.shape[2:], dtype=torch.long) + + freqs = self._create_freqs( + grid_size=grid_size, + start_frame=current_start_frame, + ) + + x_video, action_noise_pred, updated_kv_caches = self._forward_blocks( + x=x, + seq_len=seq_len, + freqs=freqs, + timestep=timestep, + context=context, + clip_feature=clip_feature, + embodiment_id=embodiment_id, + action=action, + timestep_action=timestep_action, + state=state, + kv_cache=kv_cache, + current_start_frame=current_start_frame, + ) + + # Copy the updated KV caches back to the original KV cache. + x_video = x_video.clone() + if action_noise_pred is not None: + action_noise_pred = action_noise_pred.clone() + #for block_index, updated_kv_cache in enumerate(updated_kv_caches): + # kv_cache[block_index] = updated_kv_cache.clone() + + video_noise_pred = self.unpatchify(x_video, grid_size) + + return video_noise_pred, action_noise_pred, updated_kv_caches + + def _forward_train( + self, + x, + timestep, + timestep_action, + context, + seq_len, + clean_x=None, + aug_t=None, + y=None, + clip_feature=None, + action=None, + state=None, + embodiment_id=None, + ): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + clip_fea (Tensor, *optional*): + CLIP image features for image-to-video mode + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_feature is not None and y is not None + + # Concat [x; y] only when pretrained that way (14B). 5B uses latent only, first-frame via CLIP. + if y is not None and self.concat_first_frame_latent: + x = torch.cat([x, y.to(dtype=x.dtype)], dim=1) + + # embeddings + x = self.patch_embedding(x) + + grid_size = torch.tensor(x.shape[2:], dtype=torch.long) + freqs = self._create_freqs( + grid_size=grid_size, + start_frame=0, + ) + + x = x.flatten(start_dim=2).transpose(1, 2) + assert x.shape[1] == seq_len + + B = x.shape[0] + F = timestep.shape[1] + + # time embeddings + if action is not None: + embodiment_id = torch.tensor([0]).repeat(x.shape[0]).to(device=embodiment_id.device) + action_features = self.action_encoder(action, timestep_action, embodiment_id) + action_length = action_features.shape[1] + state_features = self.state_encoder(state, embodiment_id) + action_register = torch.cat([action_features, state_features], dim=1) + action_register_length = action_register.shape[1] + x = torch.cat([x, action_register], dim=1) + else: + action_features = None + action_length = None + state_features = None + action_register = None + action_register_length = None + + # time embeddings + timestep = timestep.unsqueeze(-1).expand(B, F, seq_len // F).reshape(B, -1) + timestep_original = timestep.clone() + + if action is not None: + assert timestep_action is not None + assert state_features is not None + stride = timestep_action.shape[1] // state_features.shape[1] + timestep_state = timestep_action[:, ::stride] + timestep = torch.cat([timestep, timestep_action, timestep_state], dim=1) + + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).type_as(x)) + e = e.unflatten(dim=0, sizes=(B, -1)) + e0 = self.time_projection(e) + e0 = e0.unflatten(dim=2, sizes=(6, self.dim)) + + # context + assert context.shape[1] == self.text_len + context = self.text_embedding(context) + + if clip_feature is not None: + clip_embedding = self.img_emb(clip_feature) + context = torch.cat([clip_embedding, context], dim=1) + + if clean_x is not None: + if y is not None and self.concat_first_frame_latent: + clean_x = torch.cat([clean_x, y.to(dtype=clean_x.dtype)], dim=1) + clean_x = self.patch_embedding(clean_x) + clean_x = clean_x.flatten(start_dim=2).transpose(1, 2) + assert clean_x.shape[1] == seq_len + + x = torch.cat([clean_x, x], dim=1) + + if aug_t is None: + aug_t = torch.zeros_like(timestep_original) + assert aug_t is not None + + e_clean = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, aug_t.flatten()).type_as(x)) + e_clean = e_clean.unflatten(dim=0, sizes=timestep_original.shape) + e0_clean = self.time_projection(e_clean) + e0_clean = e0_clean.unflatten(dim=2, sizes=(6, self.dim)) + e0 = torch.cat([e0_clean, e0], dim=1) + + # arguments + kwargs = dict( + e=e0, + freqs=freqs, + freqs_action=self.freqs_action, + freqs_state=self.freqs_state, + action_register_length=action_register_length, + context=context, + is_tf=clean_x is not None, + ) + + def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + outputs, updated_kv_cache = module(*inputs, **kwargs) + assert updated_kv_cache is None + return outputs + return custom_forward + + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, **kwargs, + use_reentrant=False, + ) + else: + x = block(x, **kwargs) + + if clean_x is not None: + x = x[:, clean_x.shape[1]:] + + if action is not None: + action_noise_pred = x[:, seq_len: seq_len + action_length] + action_noise_pred = self.action_decoder(action_noise_pred, embodiment_id) + else: + action_noise_pred = None + + # Build a tensor that contains only video tokens per sample with length = max(video_lens) + x_video = x[:, :seq_len] + e_video = e[:, :seq_len] + + # Unpatchify video-only tokens + x_video = self.head(x_video, e_video.unsqueeze(2)) + video_noise_pred = self.unpatchify(x_video, grid_size) + + return video_noise_pred, action_noise_pred + + def forward( + self, + *args, + **kwargs + ): + if kwargs.get('kv_cache', None) is not None: + return self._forward_inference(*args, **kwargs) + else: + return self._forward_train(*args, **kwargs) + + def unpatchify(self, x, grid_size): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (Tensor): + Patchified features, with shape [B, L, C_out * prod(patch_size)]. + grid_size (Tensor): + Spatial-temporal grid dimensions before patching, with shape [3] + (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + Tensor: + Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] + """ + B = x.shape[0] + c = self.out_dim + grid_size = grid_size.tolist() + assert x.shape[1] == math.prod(grid_size) + x = x.view(B, *grid_size, *self.patch_size, c) + x = torch.einsum('bfhwpqrc->bcfphqwr', x) + x = x.reshape(B, c, *[i * j for i, j in zip(grid_size, self.patch_size)]) + return x + + def _create_freqs( + self, + grid_size: torch.Tensor, + start_frame: int, + ): + device = self.patch_embedding.weight.device + if any(freq.device != device for freq in self.freqs): + self.freqs = [freq.to(device) for freq in self.freqs] + if self.freqs_action.device != device: + self.freqs_action = self.freqs_action.to(device) + if self.freqs_state.device != device: + self.freqs_state = self.freqs_state.to(device) + + f, h, w = grid_size.tolist() + freqs = torch.cat( + [ + self.freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1 + ).reshape(f * h * w, 1, -1) + + return freqs + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/groot/vla/model/dreamzero/modules/wan_video_image_encoder.py b/groot/vla/model/dreamzero/modules/wan_video_image_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..ec0523226f3019b78d2d9d4dc6b17476e010384e --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_image_encoder.py @@ -0,0 +1,908 @@ +""" +Concise re-implementation of +``https://github.com/openai/CLIP'' and +``https://github.com/mlfoundations/open_clip''. +""" +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.transforms as T +from groot.vla.model.dreamzero.modules.wan_video_dit import flash_attention + + +class SelfAttention(nn.Module): + + def __init__(self, dim, num_heads, dropout=0.1, eps=1e-5): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, mask): + """ + x: [B, L, C]. + """ + b, s, c, n, d = *x.size(), self.num_heads, self.head_dim + + # compute query, key, value + q = self.q(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + k = self.k(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + v = self.v(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + + # compute attention + p = self.dropout.p if self.training else 0.0 + x = F.scaled_dot_product_attention(q, k, v, mask, p) + x = x.permute(0, 2, 1, 3).reshape(b, s, c) + + # output + x = self.o(x) + x = self.dropout(x) + return x + + +class AttentionBlock(nn.Module): + + def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.post_norm = post_norm + self.eps = eps + + # layers + self.attn = SelfAttention(dim, num_heads, dropout, eps) + self.norm1 = nn.LayerNorm(dim, eps=eps) + self.ffn = nn.Sequential( + nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim), + nn.Dropout(dropout)) + self.norm2 = nn.LayerNorm(dim, eps=eps) + + def forward(self, x, mask): + if self.post_norm: + x = self.norm1(x + self.attn(x, mask)) + x = self.norm2(x + self.ffn(x)) + else: + x = x + self.attn(self.norm1(x), mask) + x = x + self.ffn(self.norm2(x)) + return x + + +class XLMRoberta(nn.Module): + """ + XLMRobertaModel with no pooler and no LM head. + """ + + def __init__(self, + vocab_size=250002, + max_seq_len=514, + type_size=1, + pad_id=1, + dim=1024, + num_heads=16, + num_layers=24, + post_norm=True, + dropout=0.1, + eps=1e-5): + super().__init__() + self.vocab_size = vocab_size + self.max_seq_len = max_seq_len + self.type_size = type_size + self.pad_id = pad_id + self.dim = dim + self.num_heads = num_heads + self.num_layers = num_layers + self.post_norm = post_norm + self.eps = eps + + # embeddings + self.token_embedding = nn.Embedding(vocab_size, dim, padding_idx=pad_id) + self.type_embedding = nn.Embedding(type_size, dim) + self.pos_embedding = nn.Embedding(max_seq_len, dim, padding_idx=pad_id) + self.dropout = nn.Dropout(dropout) + + # blocks + self.blocks = nn.ModuleList([ + AttentionBlock(dim, num_heads, post_norm, dropout, eps) + for _ in range(num_layers) + ]) + + # norm layer + self.norm = nn.LayerNorm(dim, eps=eps) + + def forward(self, ids): + """ + ids: [B, L] of torch.LongTensor. + """ + b, s = ids.shape + mask = ids.ne(self.pad_id).long() + + # embeddings + x = self.token_embedding(ids) + \ + self.type_embedding(torch.zeros_like(ids)) + \ + self.pos_embedding(self.pad_id + torch.cumsum(mask, dim=1) * mask) + if self.post_norm: + x = self.norm(x) + x = self.dropout(x) + + # blocks + mask = torch.where( + mask.view(b, 1, 1, s).gt(0), 0.0, + torch.finfo(x.dtype).min) + for block in self.blocks: + x = block(x, mask) + + # output + if not self.post_norm: + x = self.norm(x) + return x + + +def xlm_roberta_large(pretrained=False, + return_tokenizer=False, + device='cpu', + **kwargs): + """ + XLMRobertaLarge adapted from Huggingface. + """ + # params + cfg = dict( + vocab_size=250002, + max_seq_len=514, + type_size=1, + pad_id=1, + dim=1024, + num_heads=16, + num_layers=24, + post_norm=True, + dropout=0.1, + eps=1e-5) + cfg.update(**kwargs) + + # init model + if pretrained: + from sora import DOWNLOAD_TO_CACHE + + # init a meta model + with torch.device('meta'): + model = XLMRoberta(**cfg) + + # load checkpoint + model.load_state_dict( + torch.load( + DOWNLOAD_TO_CACHE('models/xlm_roberta/xlm_roberta_large.pth'), + map_location=device), + assign=True) + else: + # init a model on device + with torch.device(device): + model = XLMRoberta(**cfg) + + # init tokenizer + if return_tokenizer: + from sora.data import HuggingfaceTokenizer + tokenizer = HuggingfaceTokenizer( + name='xlm-roberta-large', + seq_len=model.text_len, + clean='whitespace') + return model, tokenizer + else: + return model + + + +def pos_interpolate(pos, seq_len): + if pos.size(1) == seq_len: + return pos + else: + src_grid = int(math.sqrt(pos.size(1))) + tar_grid = int(math.sqrt(seq_len)) + n = pos.size(1) - src_grid * src_grid + return torch.cat([ + pos[:, :n], + F.interpolate( + pos[:, n:].float().reshape(1, src_grid, src_grid, -1).permute( + 0, 3, 1, 2), + size=(tar_grid, tar_grid), + mode='bicubic', + align_corners=False).flatten(2).transpose(1, 2) + ], + dim=1) + + +class QuickGELU(nn.Module): + + def forward(self, x): + return x * torch.sigmoid(1.702 * x) + + +class LayerNorm(nn.LayerNorm): + + def forward(self, x): + return super().forward(x).type_as(x) + + +class SelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + causal=False, + attn_dropout=0.0, + proj_dropout=0.0): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.causal = causal + self.attn_dropout = attn_dropout + self.proj_dropout = proj_dropout + + # layers + self.to_qkv = nn.Linear(dim, dim * 3) + self.proj = nn.Linear(dim, dim) + + def forward(self, x): + """ + x: [B, L, C]. + """ + # compute query, key, value + q, k, v = self.to_qkv(x).chunk(3, dim=-1) + + # compute attention + x = flash_attention(q, k, v, num_heads=self.num_heads, compatibility_mode=True) + + # output + x = self.proj(x) + x = F.dropout(x, self.proj_dropout, self.training) + return x + + +class SwiGLU(nn.Module): + + def __init__(self, dim, mid_dim): + super().__init__() + self.dim = dim + self.mid_dim = mid_dim + + # layers + self.fc1 = nn.Linear(dim, mid_dim) + self.fc2 = nn.Linear(dim, mid_dim) + self.fc3 = nn.Linear(mid_dim, dim) + + def forward(self, x): + x = F.silu(self.fc1(x)) * self.fc2(x) + x = self.fc3(x) + return x + + +class AttentionBlock(nn.Module): + + def __init__(self, + dim, + mlp_ratio, + num_heads, + post_norm=False, + causal=False, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + norm_eps=1e-5): + assert activation in ['quick_gelu', 'gelu', 'swi_glu'] + super().__init__() + self.dim = dim + self.mlp_ratio = mlp_ratio + self.num_heads = num_heads + self.post_norm = post_norm + self.causal = causal + self.norm_eps = norm_eps + + # layers + self.norm1 = LayerNorm(dim, eps=norm_eps) + self.attn = SelfAttention(dim, num_heads, causal, attn_dropout, + proj_dropout) + self.norm2 = LayerNorm(dim, eps=norm_eps) + if activation == 'swi_glu': + self.mlp = SwiGLU(dim, int(dim * mlp_ratio)) + else: + self.mlp = nn.Sequential( + nn.Linear(dim, int(dim * mlp_ratio)), + QuickGELU() if activation == 'quick_gelu' else nn.GELU(), + nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout)) + + def forward(self, x): + if self.post_norm: + x = x + self.norm1(self.attn(x)) + x = x + self.norm2(self.mlp(x)) + else: + x = x + self.attn(self.norm1(x)) + x = x + self.mlp(self.norm2(x)) + return x + + +class AttentionPool(nn.Module): + + def __init__(self, + dim, + mlp_ratio, + num_heads, + activation='gelu', + proj_dropout=0.0, + norm_eps=1e-5): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.mlp_ratio = mlp_ratio + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.proj_dropout = proj_dropout + self.norm_eps = norm_eps + + # layers + gain = 1.0 / math.sqrt(dim) + self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim)) + self.to_q = nn.Linear(dim, dim) + self.to_kv = nn.Linear(dim, dim * 2) + self.proj = nn.Linear(dim, dim) + self.norm = LayerNorm(dim, eps=norm_eps) + self.mlp = nn.Sequential( + nn.Linear(dim, int(dim * mlp_ratio)), + QuickGELU() if activation == 'quick_gelu' else nn.GELU(), + nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout)) + + def forward(self, x): + """ + x: [B, L, C]. + """ + b, s, c, n, d = *x.size(), self.num_heads, self.head_dim + + # compute query, key, value + q = self.to_q(self.cls_embedding).view(1, 1, n*d).expand(b, -1, -1) + k, v = self.to_kv(x).chunk(2, dim=-1) + + # compute attention + x = flash_attention(q, k, v, num_heads=self.num_heads, compatibility_mode=True) + x = x.reshape(b, 1, c) + + # output + x = self.proj(x) + x = F.dropout(x, self.proj_dropout, self.training) + + # mlp + x = x + self.mlp(self.norm(x)) + return x[:, 0] + + +class VisionTransformer(nn.Module): + + def __init__(self, + image_size=224, + patch_size=16, + dim=768, + mlp_ratio=4, + out_dim=512, + num_heads=12, + num_layers=12, + pool_type='token', + pre_norm=True, + post_norm=False, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + if image_size % patch_size != 0: + print( + '[WARNING] image_size is not divisible by patch_size', + flush=True) + assert pool_type in ('token', 'token_fc', 'attn_pool') + out_dim = out_dim or dim + super().__init__() + self.image_size = image_size + self.patch_size = patch_size + self.num_patches = (image_size // patch_size)**2 + self.dim = dim + self.mlp_ratio = mlp_ratio + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.pool_type = pool_type + self.post_norm = post_norm + self.norm_eps = norm_eps + + # embeddings + gain = 1.0 / math.sqrt(dim) + self.patch_embedding = nn.Conv2d( + 3, + dim, + kernel_size=patch_size, + stride=patch_size, + bias=not pre_norm) + if pool_type in ('token', 'token_fc'): + self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim)) + self.pos_embedding = nn.Parameter(gain * torch.randn( + 1, self.num_patches + + (1 if pool_type in ('token', 'token_fc') else 0), dim)) + self.dropout = nn.Dropout(embedding_dropout) + + # transformer + self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None + self.transformer = nn.Sequential(*[ + AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False, + activation, attn_dropout, proj_dropout, norm_eps) + for _ in range(num_layers) + ]) + self.post_norm = LayerNorm(dim, eps=norm_eps) + + # head + if pool_type == 'token': + self.head = nn.Parameter(gain * torch.randn(dim, out_dim)) + elif pool_type == 'token_fc': + self.head = nn.Linear(dim, out_dim) + elif pool_type == 'attn_pool': + self.head = AttentionPool(dim, mlp_ratio, num_heads, activation, + proj_dropout, norm_eps) + + def forward(self, x, interpolation=False, use_31_block=False): + b = x.size(0) + + # embeddings + x = self.patch_embedding(x).flatten(2).permute(0, 2, 1) + # print("x in forward: ", x[0,0,100:105], x.shape) + # print("patch_embedding: ", self.patch_embedding.module.weight[0:10, 0, 0, 3], self.patch_embedding.module.weight.shape) + if self.pool_type in ('token', 'token_fc'): + x = torch.cat([self.cls_embedding.expand(b, -1, -1).to(dtype=x.dtype, device=x.device), x], dim=1) + if interpolation: + e = pos_interpolate(self.pos_embedding, x.size(1)) + else: + e = self.pos_embedding + e = e.to(dtype=x.dtype, device=x.device) + x = self.dropout(x + e) + if self.pre_norm is not None: + x = self.pre_norm(x) + + # transformer + if use_31_block: + # print("x before transformer: ", x[0,0,100:105], x.shape) + x = self.transformer[:-1](x) + return x + else: + # print("x before transformer: ", x[0,0,100:105], x.shape) + x = self.transformer(x) + return x + + +class CLIP(nn.Module): + + def __init__(self, + embed_dim=512, + image_size=224, + patch_size=16, + vision_dim=768, + vision_mlp_ratio=4, + vision_heads=12, + vision_layers=12, + vision_pool='token', + vision_pre_norm=True, + vision_post_norm=False, + vocab_size=49408, + text_len=77, + text_dim=512, + text_mlp_ratio=4, + text_heads=8, + text_layers=12, + text_causal=True, + text_pool='argmax', + text_head_bias=False, + logit_bias=None, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + super().__init__() + self.embed_dim = embed_dim + self.image_size = image_size + self.patch_size = patch_size + self.vision_dim = vision_dim + self.vision_mlp_ratio = vision_mlp_ratio + self.vision_heads = vision_heads + self.vision_layers = vision_layers + self.vision_pool = vision_pool + self.vision_pre_norm = vision_pre_norm + self.vision_post_norm = vision_post_norm + self.vocab_size = vocab_size + self.text_len = text_len + self.text_dim = text_dim + self.text_mlp_ratio = text_mlp_ratio + self.text_heads = text_heads + self.text_layers = text_layers + self.text_causal = text_causal + self.text_pool = text_pool + self.text_head_bias = text_head_bias + self.norm_eps = norm_eps + + # models + self.visual = VisionTransformer( + image_size=image_size, + patch_size=patch_size, + dim=vision_dim, + mlp_ratio=vision_mlp_ratio, + out_dim=embed_dim, + num_heads=vision_heads, + num_layers=vision_layers, + pool_type=vision_pool, + pre_norm=vision_pre_norm, + post_norm=vision_post_norm, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.textual = TextTransformer( + vocab_size=vocab_size, + text_len=text_len, + dim=text_dim, + mlp_ratio=text_mlp_ratio, + out_dim=embed_dim, + num_heads=text_heads, + num_layers=text_layers, + causal=text_causal, + pool_type=text_pool, + head_bias=text_head_bias, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([])) + if logit_bias is not None: + self.logit_bias = nn.Parameter(logit_bias * torch.ones([])) + + # initialize weights + self.init_weights() + + def forward(self, imgs, txt_ids): + """ + imgs: [B, 3, H, W] of torch.float32. + - mean: [0.48145466, 0.4578275, 0.40821073] + - std: [0.26862954, 0.26130258, 0.27577711] + txt_ids: [B, L] of torch.long. Encoded by data.CLIPTokenizer. + """ + xi = self.visual(imgs) + xt = self.textual(txt_ids) + return xi, xt + + def init_weights(self): + # embeddings + nn.init.normal_(self.textual.token_embedding.weight, std=0.02) + nn.init.normal_(self.visual.patch_embedding.weight, std=0.1) + + # attentions + for modality in ['visual', 'textual']: + dim = self.vision_dim if modality == 'visual' else self.text_dim + transformer = getattr(self, modality).transformer + proj_gain = (1.0 / math.sqrt(dim)) * ( + 1.0 / math.sqrt(2 * len(transformer))) + attn_gain = 1.0 / math.sqrt(dim) + mlp_gain = 1.0 / math.sqrt(2.0 * dim) + for block in transformer: + nn.init.normal_(block.attn.to_qkv.weight, std=attn_gain) + nn.init.normal_(block.attn.proj.weight, std=proj_gain) + nn.init.normal_(block.mlp[0].weight, std=mlp_gain) + nn.init.normal_(block.mlp[2].weight, std=proj_gain) + + def param_groups(self): + groups = [{ + 'params': [ + p for n, p in self.named_parameters() + if 'norm' in n or n.endswith('bias') + ], + 'weight_decay': 0.0 + }, { + 'params': [ + p for n, p in self.named_parameters() + if not ('norm' in n or n.endswith('bias')) + ] + }] + return groups + + +class XLMRobertaWithHead(XLMRoberta): + + def __init__(self, **kwargs): + self.out_dim = kwargs.pop('out_dim') + super().__init__(**kwargs) + + # head + mid_dim = (self.dim + self.out_dim) // 2 + self.head = nn.Sequential( + nn.Linear(self.dim, mid_dim, bias=False), nn.GELU(), + nn.Linear(mid_dim, self.out_dim, bias=False)) + + def forward(self, ids): + # xlm-roberta + x = super().forward(ids) + + # average pooling + mask = ids.ne(self.pad_id).unsqueeze(-1).to(x) + x = (x * mask).sum(dim=1) / mask.sum(dim=1) + + # head + x = self.head(x) + return x + + +class XLMRobertaCLIP(nn.Module): + + def __init__(self, + embed_dim=1024, + image_size=224, + patch_size=14, + vision_dim=1280, + vision_mlp_ratio=4, + vision_heads=16, + vision_layers=32, + vision_pool='token', + vision_pre_norm=True, + vision_post_norm=False, + activation='gelu', + vocab_size=250002, + max_text_len=514, + type_size=1, + pad_id=1, + text_dim=1024, + text_heads=16, + text_layers=24, + text_post_norm=True, + text_dropout=0.1, + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + super().__init__() + self.embed_dim = embed_dim + self.image_size = image_size + self.patch_size = patch_size + self.vision_dim = vision_dim + self.vision_mlp_ratio = vision_mlp_ratio + self.vision_heads = vision_heads + self.vision_layers = vision_layers + self.vision_pre_norm = vision_pre_norm + self.vision_post_norm = vision_post_norm + self.activation = activation + self.vocab_size = vocab_size + self.max_text_len = max_text_len + self.type_size = type_size + self.pad_id = pad_id + self.text_dim = text_dim + self.text_heads = text_heads + self.text_layers = text_layers + self.text_post_norm = text_post_norm + self.norm_eps = norm_eps + + # models + self.visual = VisionTransformer( + image_size=image_size, + patch_size=patch_size, + dim=vision_dim, + mlp_ratio=vision_mlp_ratio, + out_dim=embed_dim, + num_heads=vision_heads, + num_layers=vision_layers, + pool_type=vision_pool, + pre_norm=vision_pre_norm, + post_norm=vision_post_norm, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.textual = None + self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([])) + + def forward(self, imgs, txt_ids): + """ + imgs: [B, 3, H, W] of torch.float32. + - mean: [0.48145466, 0.4578275, 0.40821073] + - std: [0.26862954, 0.26130258, 0.27577711] + txt_ids: [B, L] of torch.long. + Encoded by data.CLIPTokenizer. + """ + xi = self.visual(imgs) + xt = self.textual(txt_ids) + return xi, xt + + def param_groups(self): + groups = [{ + 'params': [ + p for n, p in self.named_parameters() + if 'norm' in n or n.endswith('bias') + ], + 'weight_decay': 0.0 + }, { + 'params': [ + p for n, p in self.named_parameters() + if not ('norm' in n or n.endswith('bias')) + ] + }] + return groups + + +def _clip(pretrained=False, + pretrained_name=None, + model_cls=CLIP, + return_transforms=False, + return_tokenizer=False, + tokenizer_padding='eos', + dtype=torch.float32, + device='cpu', + **kwargs): + # init model + if pretrained and pretrained_name: + from sora import BUCKET, DOWNLOAD_TO_CACHE + + # init a meta model + with torch.device('meta'): + model = model_cls(**kwargs) + + # checkpoint path + checkpoint = f'models/clip/{pretrained_name}' + if dtype in (torch.float16, torch.bfloat16): + suffix = '-' + { + torch.float16: 'fp16', + torch.bfloat16: 'bf16' + }[dtype] + if object_exists(BUCKET, f'{checkpoint}{suffix}.pth'): + checkpoint = f'{checkpoint}{suffix}' + checkpoint += '.pth' + + # load + model.load_state_dict( + torch.load(DOWNLOAD_TO_CACHE(checkpoint), map_location=device), + assign=True, + strict=False) + else: + # init a model on device + with torch.device(device): + model = model_cls(**kwargs) + + # set device + output = (model,) + + # init transforms + if return_transforms: + # mean and std + if 'siglip' in pretrained_name.lower(): + mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5] + else: + mean = [0.48145466, 0.4578275, 0.40821073] + std = [0.26862954, 0.26130258, 0.27577711] + + # transforms + transforms = T.Compose([ + T.Resize((model.image_size, model.image_size), + interpolation=T.InterpolationMode.BICUBIC), + T.ToTensor(), + T.Normalize(mean=mean, std=std) + ]) + output += (transforms,) + + # init tokenizer + if return_tokenizer: + from sora import data + if 'siglip' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name=f'timm/{pretrained_name}', + seq_len=model.text_len, + clean='canonicalize') + elif 'xlm' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name='xlm-roberta-large', + seq_len=model.max_text_len - 2, + clean='whitespace') + elif 'mba' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name='facebook/xlm-roberta-xl', + seq_len=model.max_text_len - 2, + clean='whitespace') + else: + tokenizer = data.CLIPTokenizer( + seq_len=model.text_len, padding=tokenizer_padding) + output += (tokenizer,) + return output[0] if len(output) == 1 else output + + +def clip_xlm_roberta_vit_h_14( + pretrained=False, + pretrained_name='open-clip-xlm-roberta-large-vit-huge-14', + **kwargs): + cfg = dict( + embed_dim=1024, + image_size=224, + patch_size=14, + vision_dim=1280, + vision_mlp_ratio=4, + vision_heads=16, + vision_layers=32, + vision_pool='token', + activation='gelu', + vocab_size=250002, + max_text_len=514, + type_size=1, + pad_id=1, + text_dim=1024, + text_heads=16, + text_layers=24, + text_post_norm=True, + text_dropout=0.1, + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0) + cfg.update(**kwargs) + return _clip(pretrained, pretrained_name, XLMRobertaCLIP, **cfg) + + +class WanImageEncoder(torch.nn.Module): + + def __init__(self, image_encoder_pretrained_path: str=None): + super().__init__() + # init model + self.model, self.transforms = clip_xlm_roberta_vit_h_14( + pretrained=False, + return_transforms=True, + return_tokenizer=False, + dtype=torch.float32, + device="cpu") + self.image_encoder_pretrained_path = image_encoder_pretrained_path + + def encode_image(self, videos): + # preprocess + size = (self.model.image_size,) * 2 + videos = torch.cat([ + F.interpolate( + u, + size=size, + mode='bicubic', + align_corners=False) for u in videos + ]) + videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5)) + + # forward + dtype = next(iter(self.model.visual.parameters())).dtype + videos = videos.to(dtype) + out = self.model.visual(videos, use_31_block=True) + # The outputs of torch compile always need to be cloned before being used. + out = out.clone() + return out + + @staticmethod + def state_dict_converter(): + return WanImageEncoderStateDictConverter() + + +class WanImageEncoderStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict + + def from_civitai(self, state_dict): + state_dict_ = {} + for name, param in state_dict.items(): + if name.startswith("textual."): + continue + name = "model." + name + state_dict_[name] = param + return state_dict_ diff --git a/groot/vla/model/dreamzero/modules/wan_video_text_encoder.py b/groot/vla/model/dreamzero/modules/wan_video_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..abeb018333f6f61f1af12aa39ab064105274f8d1 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_text_encoder.py @@ -0,0 +1,280 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def fp16_clamp(x): + if x.dtype == torch.float16 and torch.isinf(x).any(): + clamp = torch.finfo(x.dtype).max - 1000 + x = torch.clamp(x, min=-clamp, max=clamp) + return x + + +class GELU(nn.Module): + + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh( + math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + + +class T5LayerNorm(nn.Module): + + def __init__(self, dim, eps=1e-6): + super(T5LayerNorm, self).__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.type_as(self.weight) + return self.weight * x + + +class T5Attention(nn.Module): + + def __init__(self, dim, dim_attn, num_heads, dropout=0.1): + assert dim_attn % num_heads == 0 + super(T5Attention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.num_heads = num_heads + self.head_dim = dim_attn // num_heads + + # layers + self.q = nn.Linear(dim, dim_attn, bias=False) + self.k = nn.Linear(dim, dim_attn, bias=False) + self.v = nn.Linear(dim, dim_attn, bias=False) + self.o = nn.Linear(dim_attn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, context=None, mask=None, pos_bias=None): + """ + x: [B, L1, C]. + context: [B, L2, C] or None. + mask: [B, L2] or [B, L1, L2] or None. + """ + # check inputs + context = x if context is None else context + b, n, c = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.q(x).view(b, -1, n, c) + k = self.k(context).view(b, -1, n, c) + v = self.v(context).view(b, -1, n, c) + + # attention bias + attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) + if pos_bias is not None: + attn_bias += pos_bias + if mask is not None: + assert mask.ndim in [2, 3] + mask = mask.view(b, 1, 1, + -1) if mask.ndim == 2 else mask.unsqueeze(1) + attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min) + + # compute attention (T5 does not use scaling) + attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias + attn = F.softmax(attn.float(), dim=-1).type_as(attn) + x = torch.einsum('bnij,bjnc->binc', attn, v) + + # output + x = x.reshape(b, -1, n * c) + x = self.o(x) + x = self.dropout(x) + return x + + +class T5FeedForward(nn.Module): + + def __init__(self, dim, dim_ffn, dropout=0.1): + super(T5FeedForward, self).__init__() + self.dim = dim + self.dim_ffn = dim_ffn + + # layers + self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU()) + self.fc1 = nn.Linear(dim, dim_ffn, bias=False) + self.fc2 = nn.Linear(dim_ffn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + x = self.fc1(x) * self.gate(x) + x = self.dropout(x) + x = self.fc2(x) + x = self.dropout(x) + return x + + +class T5SelfAttention(nn.Module): + + def __init__(self, + dim, + dim_attn, + dim_ffn, + num_heads, + num_buckets, + shared_pos=True, + dropout=0.1): + super(T5SelfAttention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.norm1 = T5LayerNorm(dim) + self.attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm2 = T5LayerNorm(dim) + self.ffn = T5FeedForward(dim, dim_ffn, dropout) + self.pos_embedding = None if shared_pos else T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) + + def forward(self, x, mask=None, pos_bias=None): + e = pos_bias if self.shared_pos else self.pos_embedding( + x.size(1), x.size(1)) + x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) + x = fp16_clamp(x + self.ffn(self.norm2(x))) + # print("x after attn: ", x[0, 0:10], x.shape) + return x + + +class T5RelativeEmbedding(nn.Module): + + def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): + super(T5RelativeEmbedding, self).__init__() + self.num_buckets = num_buckets + self.num_heads = num_heads + self.bidirectional = bidirectional + self.max_dist = max_dist + + # layers + self.embedding = nn.Embedding(num_buckets, num_heads) + + def forward(self, lq, lk): + device = self.embedding.weight.device + # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \ + # torch.arange(lq).unsqueeze(1).to(device) + rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \ + torch.arange(lq, device=device).unsqueeze(1) + rel_pos = self._relative_position_bucket(rel_pos) + rel_pos_embeds = self.embedding(rel_pos) + rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze( + 0) # [1, N, Lq, Lk] + return rel_pos_embeds.contiguous() + + def _relative_position_bucket(self, rel_pos): + # preprocess + if self.bidirectional: + num_buckets = self.num_buckets // 2 + rel_buckets = (rel_pos > 0).long() * num_buckets + rel_pos = torch.abs(rel_pos) + else: + num_buckets = self.num_buckets + rel_buckets = 0 + rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) + + # embeddings for small and large positions + max_exact = num_buckets // 2 + rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) / + math.log(self.max_dist / max_exact) * + (num_buckets - max_exact)).long() + rel_pos_large = torch.min( + rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) + rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) + return rel_buckets + +def init_weights(m): + if isinstance(m, T5LayerNorm): + nn.init.ones_(m.weight) + elif isinstance(m, T5FeedForward): + nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5) + nn.init.normal_(m.fc1.weight, std=m.dim**-0.5) + nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5) + elif isinstance(m, T5Attention): + nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5) + nn.init.normal_(m.k.weight, std=m.dim**-0.5) + nn.init.normal_(m.v.weight, std=m.dim**-0.5) + nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5) + elif isinstance(m, T5RelativeEmbedding): + nn.init.normal_( + m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5) + + +class WanTextEncoder(torch.nn.Module): + + def __init__(self, + vocab: int | nn.Embedding = 256384, + dim=4096, + dim_attn=4096, + dim_ffn=10240, + num_heads=64, + num_layers=24, + num_buckets=32, + shared_pos=False, + dropout=0.1, + text_encoder_pretrained_path: str=None): + super(WanTextEncoder, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_layers = num_layers + self.num_buckets = num_buckets + self.shared_pos = shared_pos + self.text_encoder_pretrained_path = text_encoder_pretrained_path + + # layers + if isinstance(vocab, int): + self.token_embedding = nn.Embedding(vocab, dim) + else: + self.token_embedding = vocab + if shared_pos: + self.pos_embedding = T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) + else: + self.pos_embedding = None + self.dropout = nn.Dropout(dropout) + self.blocks = nn.ModuleList([ + T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, + shared_pos, dropout) for _ in range(num_layers) + ]) + self.norm = T5LayerNorm(dim) + + # initialize weights + self.apply(init_weights) + + def forward(self, ids, mask=None): + x = self.token_embedding(ids) + x = self.dropout(x) + if self.shared_pos: + assert self.pos_embedding is not None + e = self.pos_embedding(x.size(1), x.size(1)) + else: + e = None + for block in self.blocks: + x = block(x, mask, pos_bias=e) + x = self.norm(x) + x = self.dropout(x) + return x + + @staticmethod + def state_dict_converter(): + return WanTextEncoderStateDictConverter() + + +class WanTextEncoderStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict + + def from_civitai(self, state_dict): + return state_dict \ No newline at end of file diff --git a/groot/vla/model/dreamzero/modules/wan_video_vae.py b/groot/vla/model/dreamzero/modules/wan_video_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..4e2ff44f0db1e043b3ba9914962040f5ca0d36fa --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_vae.py @@ -0,0 +1,1373 @@ +from einops import rearrange, repeat + +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm import tqdm + +CACHE_T = 2 + + +def check_is_instance(model, module_class): + if isinstance(model, module_class): + return True + if hasattr(model, "module") and isinstance(model.module, module_class): + return True + return False + + +def block_causal_mask(x, block_size): + # params + b, n, s, _ = x.shape + assert s % block_size == 0 + num_blocks = s // block_size + + # build mask + mask = torch.zeros(b, n, s, s, dtype=torch.bool, device=x.device) + for i in range(num_blocks): + mask[:, :, + i * block_size:(i + 1) * block_size, :(i + 1) * block_size] = 1 + return mask + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], + self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + return super().forward(x) + + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else + -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + + def __init__(self, dim, mode): + assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', + 'downsample3d') + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == 'upsample2d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + elif mode == 'upsample3d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + self.time_conv = CausalConv3d(dim, + dim * 2, (3, 1, 1), + padding=(1, 0, 0)) + + elif mode == 'downsample2d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == 'downsample3d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, + dim, (3, 1, 1), + stride=(2, 1, 1), + padding=(0, 0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == 'upsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = 'Rep' + feat_idx[0] += 1 + else: + + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] != 'Rep': + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] == 'Rep': + cache_x = torch.cat([ + torch.zeros_like(cache_x), + cache_x, + ], dim=2) + if feat_cache[idx] == 'Rep': + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.resample(x) + x = rearrange(x, '(b t) c h w -> b c t h w', t=t) + + if self.mode == 'downsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + one_matrix = torch.eye(c1, c2) + init_matrix = one_matrix + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b c f (h q) (w r) -> b (c r q) f h w", + q=patch_size, + r=patch_size) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + return x + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b (c r q) f h w -> b c f (h q) (w r)", + q=patch_size, + r=patch_size) + return x + + +class Resample38(Resample): + + def __init__(self, dim, mode): + assert mode in ( + "none", + "upsample2d", + "upsample3d", + "downsample2d", + "downsample3d", + ) + super(Resample, self).__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + elif mode == "downsample3d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + self.time_conv = CausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) + ) + else: + self.resample = nn.Identity() + +class ResidualBlock(nn.Module): + + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # layers + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1)) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ + if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + # zero out the last layer params + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.norm(x) + # compute query, key, value + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( + 0, 1, 3, 2).contiguous().chunk(3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention( + q, + k, + v, + #attn_mask=block_causal_mask(q, block_size=h * w) + ) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + + # output + x = self.proj(x) + x = rearrange(x, '(b t) c h w-> b c t h w', t=t) + return x + identity + + +class AvgDown3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False + ): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample38(out_dim, mode=mode)) + + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + + return x + self.avg_shortcut(x_copy) + + +class Up_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False + ): + super().__init__() + # Shortcut path with upsample + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2 if up_flag else 1, + ) + else: + self.avg_shortcut = None + + # Main path with residual blocks and upsample + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final upsample block + if up_flag: + mode = "upsample3d" if temperal_upsample else "upsample2d" + upsamples.append(Resample38(out_dim, mode=mode)) + + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + x_main = x.clone() + for module in self.upsamples: + x_main = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + x_shortcut = self.avg_shortcut(x, first_chunk) + return x_main + x_shortcut + else: + return x_main + + +class Encoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = 'downsample3d' if temperal_downsample[ + i] else 'downsample2d' + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout)) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Encoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = ( + temperal_downsample[i] if i < len(temperal_downsample) else False + ) + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(dim_mult) - 1, + ) + ) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + + # # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + + def forward(self, x, feat_cache=None, feat_idx=[0]): + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + + return x + + +class Decoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2**(len(dim_mult) - 2) + + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i == 1 or i == 2 or i == 3: + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # upsample block + if i != len(dim_mult) - 1: + mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' + upsamples.append(Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 3, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + + +class Decoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False + upsamples.append( + Up_ResidualBlock(in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temperal_upsample=t_up_flag, + up_flag=i != len(dim_mult) - 1)) + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1)) + + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class VideoVAE_(nn.Module): + + def __init__(self, + dim=96, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + self._enc_conv_num = count_conv3d(self.encoder) + self._dec_conv_num = count_conv3d(self.decoder) + + def encode(self, x, scale): + feat_map = [None] * self._enc_conv_num + + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + + for i in range(1, iter_): + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], dim=2) + mu, _ = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + def decode(self, z, scale): + feat_map = [None] * self._dec_conv_num + + # z: [b,c,t,h,w] + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + + out = self.decoder( + x[:, :, 0:1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + + for i in range(1, iter_): + out_ = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], dim=2) + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + +class WanVideoVAE(nn.Module): + + def __init__(self, z_dim=16, vae_pretrained_path: str | None = None): + super().__init__() + + mean = [ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + ] + std = [ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + ] + self.mean = torch.tensor(mean, device='cuda') + self.std = torch.tensor(std, device='cuda') + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE_(z_dim=z_dim).eval().requires_grad_(False) + self.upsampling_factor = 8 + self.z_dim = z_dim + self.vae_pretrained_path = vae_pretrained_path + + + def build_1d_mask(self, length, left_bound, right_bound, border_width, device): + x = torch.ones((length,), device=device) + border = (torch.arange(border_width, device=device) + 1) + if not left_bound: + x[:border_width] = border / border_width + if not right_bound: + x[-border_width:] = torch.flip(border / border_width, dims=(0,)) + return x + + + def build_mask(self, data, is_bound, border_width): + _, _, _, H, W = data.shape + h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0], device=data.device) + w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1], device=data.device) + + h = repeat(h, "H -> H W", H=H, W=W) + w = repeat(w, "W -> H W", H=H, W=W) + + mask = torch.stack([h, w]).min(dim=0).values + mask = rearrange(mask, "H W -> 1 1 1 H W") + return mask + + + def tiled_decode(self, hidden_states, tile_size, tile_stride): + _, _, T, H, W = hidden_states.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + out_T = T * 4 - 3 + weight = torch.zeros( + (1, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=hidden_states.device + ) + values = torch.zeros( + (1, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"): + hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_] + hidden_states_batch = self.model.decode(hidden_states_batch, self.scale) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) * self.upsampling_factor, (size_w - stride_w) * self.upsampling_factor) + ).to(dtype=hidden_states.dtype) + + target_h = h * self.upsampling_factor + target_w = w * self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + values = values.clamp_(-1, 1) + return values + + def tiled_encode(self, video, tile_size, tile_stride): + _, _, T, H, W = video.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + out_T = (T + 3) // 4 + weight = torch.zeros( + (1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), + dtype=video.dtype, + device=video.device, + ) + values = torch.zeros( + (1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), + dtype=video.dtype, + device=video.device, + ) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"): + hidden_states_batch = video[:, :, :, h:h_, w:w_] + hidden_states_batch = self.model.encode(hidden_states_batch, self.scale) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) // self.upsampling_factor, (size_w - stride_w) // self.upsampling_factor) + ).to(dtype=video.dtype) + + target_h = h // self.upsampling_factor + target_w = w // self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + return values + + def single_encode(self, video): + x = self.model.encode(video, self.scale) + # The outputs of torch compile always need to be cloned before being used. + x = x.clone() + return x + + def single_decode(self, hidden_state): + video = self.model.decode(hidden_state, self.scale) + return video.clamp_(-1, 1) + + def encode(self, videos, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + hidden_states = [] + for video in videos: + video = video.unsqueeze(0) + if tiled: + tile_size = (tile_size[0] * self.upsampling_factor, tile_size[1] * self.upsampling_factor) + tile_stride = (tile_stride[0] * self.upsampling_factor, tile_stride[1] * self.upsampling_factor) + hidden_state = self.tiled_encode(video, tile_size, tile_stride) + else: + hidden_state = self.single_encode(video) + hidden_state = hidden_state.squeeze(0) + hidden_states.append(hidden_state) + hidden_states = torch.stack(hidden_states) + return hidden_states + + def decode(self, hidden_states, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + if tiled: + video = self.tiled_decode(hidden_states, tile_size, tile_stride) + else: + video = self.single_decode(hidden_states) + return video + + + @staticmethod + def state_dict_converter(): + return WanVideoVAEStateDictConverter() + + +class WanVideoVAEStateDictConverter: + + def __init__(self): + pass + + def from_civitai(self, state_dict): + state_dict_ = {} + if 'model_state' in state_dict: + state_dict = state_dict['model_state'] + for name in state_dict: + state_dict_['model.' + name] = state_dict[name] + return state_dict_ + + +class VideoVAE38_(VideoVAE_): + + def __init__(self, + dim=160, + z_dim=48, + dec_dim=256, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super(VideoVAE_, self).__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d_38(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d_38(dec_dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + self._enc_conv_num = count_conv3d(self.encoder) + self._dec_conv_num = count_conv3d(self.decoder) + + def encode(self, x, scale): + feat_map = [None] * self._enc_conv_num + + x = patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + + for i in range(1, iter_): + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], dim=2) + + mu, _ = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + def decode(self, z, scale): + feat_map = [None] * self._dec_conv_num + + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + + out = self.decoder( + x[:, :, 0:1, :, :], + feat_cache=feat_map, + feat_idx=[0], + first_chunk=True, + ) + + for i in range(1, iter_): + out_ = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], 2) + + out = unpatchify(out, patch_size=2) + return out + + +class WanVideoVAE38(WanVideoVAE): + + def __init__(self, z_dim=48, dim=160, vae_pretrained_path: str | None = None): + super(WanVideoVAE, self).__init__() + + mean = [ + -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, + -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, + -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, + -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, + -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, + 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667 + ] + std = [ + 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, + 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, + 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, + 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, + 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, + 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744 + ] + self.mean = torch.tensor(mean, device='cuda') + self.std = torch.tensor(std, device='cuda') + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE38_(z_dim=z_dim, dim=dim).eval().requires_grad_(False) + self.upsampling_factor = 16 + self.z_dim = z_dim + self.vae_pretrained_path = vae_pretrained_path \ No newline at end of file diff --git a/groot/vla/model/dreamzero/transform/__init__.py b/groot/vla/model/dreamzero/transform/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/transform/common.py b/groot/vla/model/dreamzero/transform/common.py new file mode 100644 index 0000000000000000000000000000000000000000..913add82986c026fb17f0c0ca6224102dd2a8cc1 --- /dev/null +++ b/groot/vla/model/dreamzero/transform/common.py @@ -0,0 +1,15 @@ +""" +Maybe there's a better place for the functions. +""" + +import re + + +def formalize_language(language: str) -> str: + """ + 1. Force lowercase + 2. Remove all punctuations + """ + language = language.lower() + language = re.sub(r"[^\w\s]", "", language) + return language diff --git a/groot/vla/model/dreamzero/transform/dreamzero_cotrain.py b/groot/vla/model/dreamzero/transform/dreamzero_cotrain.py new file mode 100644 index 0000000000000000000000000000000000000000..8b3b747d3d02127d1d23dde1417b7e7ffb305fbd --- /dev/null +++ b/groot/vla/model/dreamzero/transform/dreamzero_cotrain.py @@ -0,0 +1,634 @@ +import os +import random +from typing import Any, Dict, List, Optional + +from einops import rearrange +import numpy as np +from pydantic import Field, PrivateAttr +import torch +from transformers import AutoProcessor, ProcessorMixin, AutoTokenizer +from transformers.data.data_collator import DataCollatorMixin +from transformers.feature_extraction_utils import BatchFeature +import tree +import re +import ftfy +import html +import regex as re +import ast + +from groot.vla.data.schema import ( + EmbodimentTag, + DatasetMetadata, +) +from groot.vla.data.transform.base import InvertibleModalityTransform +from groot.vla.model.dreamzero.transform.common import formalize_language + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + +def whitespace_clean(text): + text = re.sub(r'\s+', ' ', text) + text = text.strip() + return text + + +class HuggingfaceTokenizer: + + def __init__(self, name, seq_len=None, clean=None, **kwargs): + assert clean in (None, 'whitespace') + self.name = name + self.seq_len = seq_len + self.clean = clean + + # When loading from a local checkpoint path (e.g. from training runs), pass + # local_files_only=True to avoid HFValidationError from validate_repo_id. + load_kwargs = dict(kwargs) + if os.path.isdir(name): + load_kwargs.setdefault("local_files_only", True) + # init tokenizer + self.tokenizer = AutoTokenizer.from_pretrained(name, **load_kwargs) + self.vocab_size = self.tokenizer.vocab_size + + def __call__(self, sequence, **kwargs): + return_mask = kwargs.pop('return_mask', False) + + # arguments + _kwargs = {'return_tensors': 'pt'} + if self.seq_len is not None: + _kwargs.update({ + 'padding': 'max_length', + 'truncation': True, + 'max_length': self.seq_len + }) + _kwargs.update(**kwargs) + + + # tokenization + if isinstance(sequence, str): + sequence = [sequence] + if self.clean: + sequence = [self._clean(u) for u in sequence] + ids = self.tokenizer(sequence, **_kwargs) + + # output + if return_mask: + return ids.input_ids, ids.attention_mask + else: + return ids.input_ids + + def _clean(self, text): + if self.clean == 'whitespace': + text = whitespace_clean(basic_clean(text)) + # elif self.clean == 'lower': + # text = whitespace_clean(basic_clean(text)).lower() + # elif self.clean == 'canonicalize': + # text = canonicalize(basic_clean(text)) + return text + + +def collate(features: List[dict], tokenizer: AutoTokenizer, num_views=3, embodiment_tag_mapping=None) -> dict: + batch = {} + keys = features[0].keys() + + for key in keys: + if key == "text": + output_values = [] + for elem in features: + item = elem[key] + try: + parsed_item = ast.literal_eval(item) + # Handle different return types from ast.literal_eval + if isinstance(parsed_item, (list, tuple)): + processed_item = str(parsed_item[0]) + else: + # If it's already a scalar (string, float, int, etc.), convert to string + processed_item = str(parsed_item) + + if num_views > 1 and elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.AGIBOT.value]: + processed_item = "A multi-view video shows that a robot " + processed_item.lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right hand, the bottom-left view shows the camera view from the left hand, and the bottom-right view is a black screen (inactive view). The robot " + processed_item.lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.OXE_DROID.value]: + processed_item = ( + "A multi-view video shows that a robot " + + processed_item.lower() + + " The video is split into three views: The top view shows the camera view from the robot's wrist, the bottom-left view shows the camera view from the left exterior camera, and the bottom-right view shows the camera view from the right exterior camera. During training, one of the two bottom exterior views may be a black screen (dropped view). The robot " + + processed_item.lower() + ) + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.GR1_UNIFIED.value]: + processed_item = "A single view video shows that a human " + processed_item.lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.MECKA_HANDS.value]: + processed_item = "A single view video shows that a human " + processed_item.lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.XDOF.value]: + processed_item = "A multi-view video shows that a robot " + processed_item.lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right hand, the bottom-left view shows the camera view from the left hand, and the bottom-right view is a black screen (inactive view). The robot " + processed_item.lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.YAM.value]: + processed_item = "A multi-view video shows that a robot " + processed_item.lower() + " The video is split into four views: The top-left view shows the top camera, the top-right view shows the right camera, the bottom-left view shows the left camera, and the bottom-right view is a black screen. The robot " + processed_item.lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.LIBERO_SIM.value]: + processed_item = "A single view video shows that a robot " + processed_item.lower() + else: + raise ValueError(f"Embodiment ID {elem['embodiment_id']} not supported.") + output_values.append(processed_item) + except (ValueError, SyntaxError, TypeError): + # If parsing fails or item is already a string, use it directly + if num_views > 1 and elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.AGIBOT.value]: + item = "A multi-view video shows that a robot " + str(item).lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right hand, the bottom-left view shows the camera view from the left hand, and the bottom-right view is a black screen (inactive view). The robot " + str(item).lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.OXE_DROID.value]: + item = ( + "A multi-view video shows that a robot " + + str(item).lower() + + " The video is split into three views: The top view shows the camera view from the robot's wrist, the bottom-left view shows the camera view from the left exterior camera, and the bottom-right view shows the camera view from the right exterior camera. During training, one of the two bottom exterior views may be a black screen (dropped view). The robot " + + str(item).lower() + ) + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.GR1_UNIFIED.value]: + item = "A single view video shows that a human " + str(item).lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.MECKA_HANDS.value]: + item = "A single view video shows that a human " + str(item).lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.XDOF.value]: + item = "A multi-view video shows that a robot " + str(item).lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right hand, the bottom-left view shows the camera view from the left hand, and the bottom-right view is a black screen (inactive view). The robot " + str(item).lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.YAM.value]: + item = "A multi-view video shows that a robot " + str(item).lower() + " The video is split into four views: The top-left view shows the top camera, the top-right view shows the right camera, the bottom-left view shows the left camera, and the bottom-right view is a black screen. The robot " + str(item).lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.LIBERO_SIM.value]: + item = "A single view video shows that a robot " + str(item).lower() + else: + raise ValueError(f"Embodiment ID {elem['embodiment_id']} not supported.") + output_values.append(item) + # print("output_values", output_values) + ids, mask = tokenizer(output_values, return_mask=True, add_special_tokens=True) + batch[key] = ids + batch['text_attention_mask'] = mask + elif key == "text_negative": + values = [elem[key] for elem in features] + ids, mask = tokenizer(values, return_mask=True, add_special_tokens=True) + batch[key] = ids + batch['text_attention_mask_negative'] = mask + else: + values = [elem[key] for elem in features] + batch[key] = torch.from_numpy(np.stack(values)) + return batch + + + +class DefaultDataCollator(DataCollatorMixin): + def __init__(self, tokenizer_path: str="google/umt5-xxl", max_length: int=512, num_views: int=1, embodiment_tag_mapping=None): + super().__init__() + self.tokenizer = HuggingfaceTokenizer(name=tokenizer_path, seq_len=max_length, clean='whitespace') + self.num_views = num_views + self.embodiment_tag_mapping = embodiment_tag_mapping + + def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]: + return collate(features, self.tokenizer, self.num_views, self.embodiment_tag_mapping) + + +class DreamTransform(InvertibleModalityTransform): + + # -- We inherit from ModalityTransform, so we keep apply_to as well -- + apply_to: list[str] = Field( + default_factory=list, description="Not used in this transform, kept for compatibility." + ) + training: bool = Field( + default=True, description="Whether to apply the transform in training mode." + ) + + formalize_language: bool = Field(default=False, description="Formalize language if True.") + + embodiment_tag_mapping: dict[str, int] = Field( + default_factory=dict, + description="The projector index of each embodiment tag.", + ) + + language_dropout_prob: float = Field( + default=0.0, + description="Dropout probability for language.", + ) + always_use_default_instruction: bool = Field( + default=False, + description="Whether to always use the default instruction. For studying how much the language helps.", + ) + + # Private attributes to keep track of shapes/dimensions across apply/unapply + _language_key: Optional[str] = PrivateAttr(default=None) + _language_keys: Optional[list[str]] = PrivateAttr(default=None) + + # XEmbDiT arguments + default_instruction: str + max_state_dim: int + max_action_dim: int + max_length: int = 512 + embodiment_tag: EmbodimentTag | None = None + state_horizon: int + action_horizon: int + num_views: int = 3 + + # Add tokenizer attribute + tokenizer_path: str = Field( + default="google/umt5-xxl", + description="Path to the tokenizer." + ) + _tokenizer: Optional[HuggingfaceTokenizer] = PrivateAttr(default=None) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + # Initialize the tokenizer + self._tokenizer = HuggingfaceTokenizer( + name=self.tokenizer_path, + seq_len=self.max_length, + clean='whitespace' + ) + + @property + def tokenizer(self): + return self._tokenizer + + def set_metadata( + self, dataset_metadata: DatasetMetadata + ): + self.embodiment_tag = dataset_metadata.embodiment_tag + + def get_embodiment_tag(self) -> int: + """Get the embodiment tag from the data.""" + assert ( + self.embodiment_tag is not None + ), "Embodiment tag not set. Please call set_metadata first." + return self.embodiment_tag_mapping[self.embodiment_tag.value] + + def check_keys_and_batch_size(self, data): + grouped_keys = {} + for key in data.keys(): + try: + modality, _ = key.split(".") + if "annotation" in key: + modality = "language" + except: # noqa: E722 + ### Handle language annotation special case + if "annotation" in key: + modality = "language" + else: + modality = "others" # will contain the video, state, and action + if modality not in grouped_keys: + grouped_keys[modality] = [] + grouped_keys[modality].append(key) + # Use video key to determine batch size. + video_ndim = data["video"].ndim + if video_ndim == 5: # Interpret as [T, V, H, W, C] + is_batched = False + batch_size = 1 + elif video_ndim == 6: # Interpret as [B, T, V, H, W, C] + is_batched = True + batch_size = data["video"].shape[0] + else: + raise ValueError(f"Unsupported video number of dimensions: {video_ndim}") + + # Handle language + if "language" in grouped_keys: + language_keys = grouped_keys["language"] + self._language_keys = language_keys # Store all keys for random selection + if len(language_keys) == 1: + self._language_key = language_keys[0] + else: + self._language_key = None # Will be selected randomly in _prepare_language + return is_batched, batch_size + + def _apply_vlm_processing(self, batch: dict) -> BatchFeature: + """ + Args: + batch: + video: [V, T, C, H, W] + Returns: required input with the format `BatchFeature` + """ + images = batch["images"] # [V, T, C, H, W] + + np_images = rearrange(images, "v t c h w -> (t v) h w c") + if "language" in batch: + lang = batch["language"] + if isinstance(lang, list) or isinstance(lang, np.ndarray): + lang = lang[0] + + inputs = {} + inputs["images"] = np_images + inputs["text"] = lang + + return inputs + + def _prepare_video(self, data: dict): + """Process, stack, and pad images from data['video'].""" + images = rearrange( + data["video"], + "t v h w c -> v t c h w", + ) + if images.shape[0] > 1: + v, t, c, h, w = images.shape + + # For DROID embodiment: 2x2 grid where the wrist view spans the full top row, + # and the two exterior views occupy the bottom row. + # + # View indices (expected): + # - View 0: left exterior + # - View 1: right exterior + # - View 2: wrist + # + # Layout: + # [wrist, wrist] (wrist duplicated to have 2x width) + # [left_ext | right_ext] + # + # Training-time augmentation: + # - Randomly drop (black out) either left_ext or right_ext. + if self.embodiment_tag == EmbodimentTag.OXE_DROID and v >= 3: + left_exterior = images[0] # (t, c, h, w) + right_exterior = images[1] # (t, c, h, w) + wrist_image = images[2] # (t, c, h, w) + + concat_images = np.zeros((1, t, c, 2 * h, 2 * w), dtype=images.dtype) + + # Top row: a SINGLE wrist view, resized to be 2x wider (same height). + # We use nearest-neighbor upscaling by repeating pixels along width. + wrist_wide = np.repeat(wrist_image, 2, axis=-1) # (t, c, h, 2w) + concat_images[0, :, :, :h, :] = wrist_wide + + # # Bottom row: left/right exteriors. + # drop_exterior_idx = None + # if self.training: + # # Always drop exactly one exterior view during training. + # drop_exterior_idx = random.choice([0, 1]) # 0=left, 1=right + + # if drop_exterior_idx != 0: + concat_images[0, :, :, h:, :w] = left_exterior + # if drop_exterior_idx != 1: + concat_images[0, :, :, h:, w:] = right_exterior + + return concat_images + + # For other embodiments: use 2x2 grid layout + # Layout: [head, right] + # [left, black] + + # Create output tensor with doubled height and width + concat_images = np.zeros((1, t, c, 2*h, 2*w), dtype=images.dtype) + + # Place images in the 2x2 grid + # Left upper: head image (view 0) + if v > 0: + concat_images[0, :, :, :h, :w] = images[0] + + # Left bottom: left image (view 1) + if v > 1: + concat_images[0, :, :, h:, :w] = images[1] + + # Right top: right image (view 2) + if v > 2: + concat_images[0, :, :, :h, w:] = images[2] + + # Right bottom: black pixels (already zeros from initialization) + + return concat_images + + return images + + def _prepare_language(self, data: dict): + """Tokenize data['language'] (or default_instruction if missing).""" + # Determine which language key to use + selected_key = self._language_key + + # For DROID embodiment during training, randomly select from available language keys + if (self._language_keys is not None and + len(self._language_keys) > 1 and + self.training and + self.embodiment_tag == EmbodimentTag.OXE_DROID): + selected_key = random.choice(self._language_keys) + elif self._language_keys is not None and len(self._language_keys) > 0 and selected_key is None: + selected_key = self._language_keys[0] + + if selected_key is not None: + raw_language = data[selected_key] + if isinstance(raw_language, np.ndarray): + raw_language = raw_language.item() if raw_language.size == 1 else raw_language[0] + if isinstance(raw_language, list): + raw_language = raw_language[0] + + # Language dropout + # WARNING: this is not compatible with LAPA and DREAM + if self.training and self.language_dropout_prob > 1e-9: + if random.random() < self.language_dropout_prob: + raw_language = self.default_instruction + else: + raw_language = self.default_instruction + + if "" in raw_language: + raw_language = raw_language.replace("", "") + is_lapa_instance = True + else: + is_lapa_instance = False + + if "" in raw_language: + raw_language = raw_language.replace("", "") + is_dream_instance = True + else: + is_dream_instance = False + + if "" in raw_language: + raw_language = raw_language.replace("", "") + is_cotrain_instance = True + else: + is_cotrain_instance = False + + if self.always_use_default_instruction: + raw_language = self.default_instruction + + # print("raw_language", raw_language) + + # Formalize language + if self.formalize_language: + formalized_language = formalize_language(raw_language) + return formalized_language, is_lapa_instance, is_dream_instance, is_cotrain_instance + else: + return raw_language, is_lapa_instance, is_dream_instance, is_cotrain_instance + + def _prepare_state(self, data: dict): + """ + Gathers final state from data['state'], then pads to max_state_dim. + Return (state, state_mask, n_state_tokens). + """ + + if "state" not in data: + state = np.zeros((self.state_horizon, self.max_state_dim)) + state_mask = np.zeros((self.state_horizon, self.max_state_dim), dtype=bool) + n_state_tokens = self.state_horizon + return state, state_mask, n_state_tokens + + state = data["state"] + assert state.shape[0] % self.state_horizon == 0, f"{state.shape=}, {self.state_horizon=}" + + n_state_dims = state.shape[-1] + + # Instead of asserting, just take the first max_state_dim dimensions if needed + if n_state_dims > self.max_state_dim: + state = state[:, : self.max_state_dim] + n_state_dims = self.max_state_dim + else: + # Pad up to max_state_dim if smaller + state = np.pad(state, ((0, 0), (0, self.max_state_dim - n_state_dims)), "constant") + + # Create mask for real state dims + state_mask = np.zeros_like(state).astype(bool) + state_mask[:, :n_state_dims] = True + + # We only have 1 "proprio" token to represent the entire state + n_state_tokens = state.shape[0] + return state, state_mask, n_state_tokens + + def _prepare_action(self, data: dict): + """ + Pad to max_action_dim, return masks. + """ + if "action" not in data: + actions = np.zeros((self.action_horizon, self.max_action_dim)) + actions_mask = np.zeros((self.action_horizon, self.max_action_dim), dtype=bool) + n_action_tokens = self.action_horizon + return actions, actions_mask, n_action_tokens + + actions = data["action"] + assert actions.shape[0] % self.action_horizon == 0, f"{actions.shape=}, {self.action_horizon=}" + + n_action_tokens = actions.shape[0] # T + n_action_dims = actions.shape[1] + + assert ( + n_action_dims <= self.max_action_dim + ), f"Action dim {n_action_dims} exceeds max allowed {self.max_action_dim}." + + # Pad the channel dimension + actions = np.pad(actions, ((0, 0), (0, self.max_action_dim - n_action_dims)), "constant") + + # Create mask: [T, max_action_dim] + actions_mask = np.zeros((n_action_tokens, self.max_action_dim), dtype=bool) + actions_mask[:, :n_action_dims] = True + + return actions, actions_mask, n_action_tokens + + def apply_single(self, data: dict) -> dict: + transformed_data = {} + + # 1) Prepare video and language with vlm processing. + images = self._prepare_video(data) + images = images.astype(np.uint8) + language, is_lapa_instance, is_dream_instance, is_cotrain_instance = self._prepare_language(data) + batch_data = {"images": images, "language": language} + vlm_outputs = self._apply_vlm_processing(batch_data) + + # 2) Prepare state + state, state_mask, _ = self._prepare_state(data) + transformed_data["state"] = state + transformed_data["state_mask"] = state_mask + + if self.training: + # 3) Prepare actions + is_detection_instance = self.embodiment_tag == EmbodimentTag.GR1_UNIFIED_SEGMENTATION + if is_detection_instance: + transformed_data["segmentation_target"] = data["action"][0, -3:-1] + transformed_data["segmentation_target_mask"] = data["action"][0, -1:] + transformed_data["has_real_action"] = np.zeros((), dtype=bool) + else: + transformed_data["segmentation_target"] = np.zeros((2,)) + transformed_data["segmentation_target_mask"] = np.zeros((1,)) + transformed_data["has_real_action"] = np.ones((), dtype=bool) + actions, actions_mask, _ = self._prepare_action(data) + transformed_data["action"] = actions + transformed_data["action_mask"] = actions_mask + + # default for lapa instance + transformed_data["lapa_action"] = np.zeros_like(transformed_data["action"]) + transformed_data["lapa_action_mask"] = np.zeros_like(transformed_data["action_mask"]) + # else: + transformed_data["text_negative"] = "Vibrant colors, overexposed, static, blurry details, text, subtitles, style, artwork, painting, image, still, grayscale, dull, worst quality, low quality, JPEG artifacts, ugly, mutilated, extra fingers, bad hands, bad face, deformed, disfigured, mutated limbs, fused fingers, stagnant image, cluttered background, three legs, many people in the background, walking backwards." + + for k, v in vlm_outputs.items(): + assert k not in transformed_data, f"Key {k} already exists in transformed_data." + transformed_data[k] = v + + transformed_data["embodiment_id"] = self.get_embodiment_tag() + + if self.embodiment_tag == EmbodimentTag.MECKA_HANDS: + is_cotrain_instance = True + else: + is_cotrain_instance = False + + transformed_data["has_lapa_action"] = np.zeros((), dtype=bool) + # print("dreamzero_fixed", is_cotrain_instance) + if is_cotrain_instance: + transformed_data["is_cotrain_instance"] = np.ones((), dtype=bool) + else: + transformed_data["is_cotrain_instance"] = np.zeros((), dtype=bool) + + if is_dream_instance: + assert "dream_actions" in data + transformed_data["embodiment_id"] = self.embodiment_tag_mapping["dream"] + transformed_data["state"] = np.zeros_like(transformed_data["state"]) + actions_shape = transformed_data["action"].shape + + # Treat the "dream" IDM action as a real action so that flow matching loss will be applied. + transformed_data["has_real_action"] = np.ones((), dtype=bool) + transformed_data["has_lapa_action"] = np.zeros((), dtype=bool) + + dream_actions = data["dream_actions"] + assert ( + dream_actions.size == actions_shape[0] * actions_shape[1] + ), f"dream_actions size {dream_actions.size} does not match action shape {actions_shape}" + transformed_data["action"] = dream_actions.reshape(actions_shape) + + if is_lapa_instance: + assert "lapa_action" in data + transformed_data["has_real_action"] = np.ones((), dtype=bool) + transformed_data["has_lapa_action"] = np.zeros((), dtype=bool) + transformed_data["embodiment_id"] = self.embodiment_tag_mapping["lapa"] + transformed_data["state"] = np.zeros_like(transformed_data["state"]) + actions_shape = transformed_data["action"].shape + lapa_actions = data["lapa_action"] + # Ensure total elements match before reshaping + assert ( + lapa_actions.size == actions_shape[0] * actions_shape[1] + ), f"Cannot reshape lapa_actions of size {lapa_actions.size} to {actions_shape}" + # Reshape the lapa_actions to match the expected shape + reshaped_lapa_actions = lapa_actions.reshape(actions_shape) + # lapa_action should be between -1 and 1 + assert np.all(reshaped_lapa_actions >= -1) and np.all( + reshaped_lapa_actions <= 1 + ), "LAPA action values should be between -1 and 1" + transformed_data["action"] = reshaped_lapa_actions + transformed_data["action_mask"] = np.ones(actions_shape, dtype=bool) + + if self.training: + action_and_mask_keys = ["action", "action_mask", "lapa_action", "lapa_action_mask"] + assert all( + transformed_data[key].shape == transformed_data["action"].shape + for key in action_and_mask_keys + ), f"Shape mismatch: {[(key, transformed_data[key].shape) for key in action_and_mask_keys]}" + + return transformed_data + + def apply_batch(self, data: dict, batch_size: int) -> dict: + # Split on batch dimension. + # delete lapa_action and lapa_action_mask + data.pop("lapa_action", None) + # data.pop("lapa_action_mask", None) + data.pop("dream_actions", None) + data_split = [tree.map_structure(lambda x: x[i], data) for i in range(batch_size)] + # Process each element. + data_split_processed = [self.apply_single(elem) for elem in data_split] + return collate(data_split_processed, self.tokenizer, self.num_views, self.embodiment_tag_mapping) + + def apply(self, data: dict) -> dict: + if not self.training and data["video"].ndim == 5: + data["video"] = data["video"][None, ...] + is_batched, batch_size = self.check_keys_and_batch_size(data) + if is_batched: + return self.apply_batch(data, batch_size) + else: + return self.apply_single(data) + + def unapply(self, data: dict) -> dict: + # Leave as is so that ConcatTransform can split the values + return data + + def __call__(self, data: dict) -> dict: + return self.apply(data) + diff --git a/groot/vla/model/n1_5/__init__.py b/groot/vla/model/n1_5/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/n1_5/action_head/__init__.py b/groot/vla/model/n1_5/action_head/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/n1_5/action_head/base_action_head.py b/groot/vla/model/n1_5/action_head/base_action_head.py new file mode 100644 index 0000000000000000000000000000000000000000..e8398622d4da362796b0fdb8a807a1592b55270e --- /dev/null +++ b/groot/vla/model/n1_5/action_head/base_action_head.py @@ -0,0 +1,32 @@ +from abc import ABC, abstractmethod + +from torch import nn +from transformers.feature_extraction_utils import BatchFeature + + +class ActionHead(ABC, nn.Module): + def __init__(self): + super(ActionHead, self).__init__() + + @abstractmethod + def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: + pass + + def get_action( + self, + backbone_output: BatchFeature, + action_input: BatchFeature, + num_action_samples: int = 1, + inference_batch_size: int = 32, + ) -> BatchFeature: + # Used for predicting actions during inference + # By default, the action head does the same thing as a normal forward pass + return self.forward(backbone_output, action_input) + + def prepare_input(self, batch: dict) -> BatchFeature: + pass + + def set_override_kwargs(self, **kwargs): + for key, value in kwargs.items(): + setattr(self.config, key, value) + setattr(self, key, value) diff --git a/groot/vla/model/n1_5/modules/__init__.py b/groot/vla/model/n1_5/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/n1_5/modules/action_encoder.py b/groot/vla/model/n1_5/modules/action_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..55a39303c2c69913be01c7a86fc5cea7af6921bb --- /dev/null +++ b/groot/vla/model/n1_5/modules/action_encoder.py @@ -0,0 +1,87 @@ +### Implemented based on Pi0 Action Encoding +import torch +import torch.nn as nn + + +def swish(x): + return x * torch.sigmoid(x) + + +class SinusoidalPositionalEncoding(nn.Module): + """ + Produces a sinusoidal encoding of shape (B, T, w) + given timesteps of shape (B, T). + """ + + def __init__(self, embedding_dim): + super().__init__() + self.embedding_dim = embedding_dim + + def forward(self, timesteps): + # timesteps: shape (B, T) + # We'll compute sin/cos frequencies across dim T + timesteps = timesteps.float() # ensure float + + B, T = timesteps.shape + device = timesteps.device + + half_dim = self.embedding_dim // 2 + # typical log space frequencies for sinusoidal encoding + exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * ( + torch.log(torch.tensor(10000.0)) / half_dim + ) + # Expand timesteps to (B, T, 1) then multiply + freqs = timesteps.unsqueeze(-1) * exponent.exp() # (B, T, half_dim) + + sin = torch.sin(freqs) + cos = torch.cos(freqs) + enc = torch.cat([sin, cos], dim=-1) # (B, T, w) + + return enc + + +class ActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size): + super().__init__() + self.hidden_size = hidden_size + + # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} + self.W1 = nn.Linear(action_dim, hidden_size) # (d -> w) + self.W2 = nn.Linear(2 * hidden_size, hidden_size) # (2w -> w) + self.W3 = nn.Linear(hidden_size, hidden_size) # (w -> w) + + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions, timesteps): + """ + actions: shape (B, T, action_dim) + timesteps: shape (B,) -- a single scalar per batch item + returns: shape (B, T, hidden_size) + """ + B, T, _ = actions.shape + + # 1) Expand each batch's single scalar time 'tau' across all T steps + # so that shape => (B, T) + # e.g. if timesteps is (B,), replicate across T + if timesteps.dim() == 1 and timesteps.shape[0] == B: + # shape (B,) => (B,T) + timesteps = timesteps.unsqueeze(1).expand(-1, T) + else: + raise ValueError( + "Expected `timesteps` to have shape (B,) so we can replicate across T." + ) + + # 2) Standard action MLP step for shape => (B, T, w) + a_emb = self.W1(actions) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.W2(x)) + + # 5) Finally W3 => (B, T, w) + x = self.W3(x) + + return x diff --git a/groot/vla/model/n1_5/sim_policy.py b/groot/vla/model/n1_5/sim_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..f1d01b510928deaeb69ad454c55388df97e6ef61 --- /dev/null +++ b/groot/vla/model/n1_5/sim_policy.py @@ -0,0 +1,1197 @@ +import importlib +import json +from pathlib import Path +from typing import Any, Callable + +import gymnasium as gym +from hydra.utils import instantiate +import numpy as np +from omegaconf import OmegaConf +from tianshou.data import Batch +from tianshou.policy import BasePolicy as BaseTianshouPolicy +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +import tree +import time + +from groot.vla.data.schema import DatasetMetadata, EmbodimentTag +from groot.vla.data.transform import ComposedModalityTransform + + +class ModelManager: + """ + Manages model loading/offloading to handle memory efficiently when using multiple models. + Modified to keep VLM always loaded and only manage text_encoder/vae components. + """ + def __init__(self): + self.active_components = None # Track which components are active for action_head + self.models = {} + self.vlm_policy = None # Keep VLM always loaded + self.action_head_policy = None # Action head policy for component management + + def register_model(self, name: str, policy_instance): + """Register a policy instance with the manager.""" + self.models[name] = policy_instance + + if name == "vlm": + self.vlm_policy = policy_instance + # Load VLM to GPU and keep it there + self.load_vlm_model() + elif name == "action_head": + self.action_head_policy = policy_instance + # Ensure vram management is enabled for action head + self.enable_action_head_vram_management() + # Initially offload action head components to save memory + self.offload_action_head_components() + + def load_vlm_model(self): + """Load VLM model to GPU and keep it there.""" + if self.vlm_policy is None: + return + + policy = self.vlm_policy + print(f"Loading VLM model to GPU (keeping it loaded)...") + + # Move VLM model to GPU + policy.trained_model.to(device=policy.device) + + # Apply bf16 if needed + if policy.eval_bf16: + policy.trained_model = policy.trained_model.to(dtype=torch.bfloat16) + + torch.cuda.empty_cache() # Clear cache after loading + + def activate_model(self, name: str): + """Activate a model - for VLM this is a no-op, for action_head this manages components.""" + if name not in self.models: + raise ValueError(f"Model {name} not registered") + + if name == "vlm": + # VLM is always loaded, just ensure it's ready + if self.vlm_policy: + self.vlm_policy.ensure_model_on_gpu() + return + elif name == "action_head": + # For action head, activate the text_encoder and vae components + self.activate_action_head_components(['text_encoder', 'vae']) + self.active_components = ['text_encoder', 'vae'] + + def activate_action_head_components(self, component_names): + """Activate specific components of the action head model.""" + if self.action_head_policy is None: + return + + policy = self.action_head_policy + print(f"Loading action head components {component_names} to GPU...") + + # Check if the action head has the load_models_to_device method + if (hasattr(policy.trained_model, "action_head") and + hasattr(policy.trained_model.action_head, "load_models_to_device")): + # Use the selective loading method + policy.trained_model.action_head.load_models_to_device(component_names) + else: + # Fallback to loading the entire action head + if hasattr(policy.trained_model, "action_head") and hasattr(policy.trained_model.action_head, "image_encoder"): + policy.trained_model.action_head.enable_vram_management() + else: + policy.trained_model.to(device=policy.device) + + # Apply bf16 if needed + if policy.eval_bf16: + policy.trained_model = policy.trained_model.to(dtype=torch.bfloat16) + + torch.cuda.empty_cache() # Clear cache after loading + + def offload_action_head_components(self): + """Offload action head components to CPU.""" + if self.action_head_policy is None: + return + + policy = self.action_head_policy + print(f"Offloading action head components to CPU...") + + # Check if the action head has the load_models_to_device method + if (hasattr(policy.trained_model, "action_head") and + hasattr(policy.trained_model.action_head, "load_models_to_device")): + # Use the selective offloading method - pass empty list to offload all + policy.trained_model.action_head.load_models_to_device([]) + else: + # Fallback to offloading the entire action head + if hasattr(policy.trained_model, "action_head") and hasattr(policy.trained_model.action_head, "image_encoder"): + if hasattr(policy.trained_model.action_head, 'disable_vram_management'): + policy.trained_model.action_head.disable_vram_management() + policy.trained_model.to(device='cpu') + else: + policy.trained_model.to(device='cpu') + + torch.cuda.empty_cache() # Clear cache after offloading + self.active_components = None + + def load_model(self, name: str): + """Load a model to GPU. (Deprecated - use activate_model instead)""" + print(f"Warning: load_model is deprecated. Use activate_model instead.") + self.activate_model(name) + + def offload_model(self, name: str): + """Offload a model to CPU. (Deprecated - components are managed automatically)""" + if name == "action_head": + self.offload_action_head_components() + # VLM is never offloaded in this new approach + + def get_status(self): + """Get current status of model manager for debugging.""" + status = { + "vlm_loaded": self.vlm_policy is not None, + "action_head_available": self.action_head_policy is not None, + "active_components": self.active_components + } + return status + + def enable_action_head_vram_management(self): + """Enable vram management for the action head model.""" + if self.action_head_policy is None: + return + + policy = self.action_head_policy + + # Check if the action head has vram management capability + if (hasattr(policy.trained_model, "action_head") and + hasattr(policy.trained_model.action_head, "enable_vram_management")): + print("Enabling vram management for action head...") + policy.trained_model.action_head.enable_vram_management() + else: + print("Warning: Action head does not support vram management") + + +class BaseGrootSimPolicy(BaseTianshouPolicy): + def __init__(self, embodiment_tag: EmbodimentTag, model_path: str, device: int | str): + super().__init__() + self.embodiment_tag = embodiment_tag + self.model_path = model_path + self.device = device + + def forward(self, batch, state=None, **kwargs): + raise NotImplementedError + + @property + def video_delta_indices(self) -> np.ndarray: + return np.array([0]) + + @property + def state_delta_indices(self) -> np.ndarray: + return np.array([0]) + + @property + def raw_data_image_transform(self) -> Callable: + return lambda x: x + + def on_env_init(self, env: gym.Env): + pass + + def learn(self, batch: Batch, **kwargs) -> dict: + """Dummy learn method for BasePolicy.learn as this is an inference-only wrapper. + + Args: + batch: Input batch of experiences + **kwargs: Additional arguments + + Returns: + Empty dict as no learning occurs + """ + return {} + + +def _update_tokenizer_path_in_config(cfg, new_path: str) -> None: + """Update tokenizer_path in the transforms subtree only (avoids trainer.model etc.).""" + from omegaconf import DictConfig, ListConfig + if isinstance(cfg, DictConfig): + if "tokenizer_path" in cfg: + cfg.tokenizer_path = new_path + # Only recurse via "transforms" to avoid triggering resolution of trainer, etc. + if "transforms" not in cfg: + return + sub = cfg.transforms + if isinstance(sub, DictConfig): + for v in sub.values(): + _update_tokenizer_path_in_config(v, new_path) + elif isinstance(sub, ListConfig): + for v in sub: + _update_tokenizer_path_in_config(v, new_path) + elif isinstance(cfg, ListConfig): + for v in cfg: + _update_tokenizer_path_in_config(v, new_path) + + +class GrootSimPolicy(BaseGrootSimPolicy): + def __init__( + self, + embodiment_tag: EmbodimentTag, + model_path: str, + device: int | str, + model_config_overrides: list[str] | None = [], + tokenizer_path_override: str | None = None, + skip_assert_delta_indices: bool = False, + skip_img_transform: bool = False, + lazy_load: bool = False, + device_mesh: DeviceMesh | None = None, + ): + """ + Initialize the GrootSimPolicy. + + Args: + env_name (str): The name of the environment. + model_path (str): Path to the model checkpoint. + device (int | str): Device to run the model on. + lazy_load (bool): If True, don't load model to GPU immediately. + device_mesh (DeviceMesh | None): Device mesh to parallelize the model across. + """ + super().__init__(embodiment_tag=embodiment_tag, model_path=model_path, device=device) + model_dir = Path(model_path) + self.rank = dist.get_rank() + + exp_cfg_dir = model_dir / "experiment_cfg" + train_cfg_path = exp_cfg_dir / "conf.yaml" + train_cfg = OmegaConf.load(train_cfg_path) + if tokenizer_path_override is not None: + _update_tokenizer_path_in_config(train_cfg, tokenizer_path_override) + self.train_cfg = train_cfg + self.lazy_load = lazy_load + + # Store model loading parameters for lazy loading + self.model_config_overrides = model_config_overrides + self.model_dir = model_dir + + # 1. Load the model + if ( + train_cfg.model._target_.endswith(".from_pretrained") + or train_cfg.model._target_.endswith(".from_pretrained_for_tuning") + or train_cfg.model._target_.endswith(".from_pretrained_with_wrapped_action_head") + ): + # Compatibility with the finetuned model + model_target = train_cfg.model._target_.rsplit(".", 1)[0] + else: + model_target = train_cfg.model._target_ + + self.model_target = model_target + + if model_config_overrides is not None and len(model_config_overrides) != 0: + print(f"Applying model config overrides: {model_config_overrides}") + # Only apply new logic if there are config overrides + + # Import the model class from the target module path + module_path, class_name = model_target.rsplit(".", 1) + if "lora" in class_name: + module_path, class_name = module_path.rsplit(".", 1) + module = importlib.import_module(module_path) + model_class = getattr(module, class_name) + model_config_class = model_class.config_class + + # Load the model config from the model directory + model_config = json.load(open(model_dir / "config.json", "r")) + model_config = OmegaConf.create(model_config) + model_config.merge_with_dotlist(list(model_config_overrides)) + model_config = OmegaConf.to_container( + model_config, resolve=True + ) # need dict for later model instantiation + model_config = model_config_class.from_dict(model_config) + + # Instantiate the model + if hasattr(train_cfg, "save_lora_only") and train_cfg.save_lora_only is True: + print(f"Loading LoRA weights from pretrained") + model = model_class.load_lora(model_path) + else: + print(f"Loading model from pretrained directly") + model = model_class.from_pretrained(model_path, config=model_config) + else: + print(f"No model config overrides provided") + # Otherwise, just call from_pretrained directly + cls_module, cls_name = model_target.rsplit(".", 1) + if 'lora' in cls_name: + cls_module, cls_name = cls_module.rsplit(".", 1) + if hasattr(train_cfg, "save_lora_only") and train_cfg.save_lora_only is True: + print(f"Loading LoRA weights from pretrained") + cls = getattr(importlib.import_module(cls_module), cls_name) + model = cls.load_lora(model_path) + else: + print(f"Loading model from pretrained directly") + cls = getattr(importlib.import_module(cls_module), cls_name) + from_pretrained = getattr(cls, "from_pretrained") + model = from_pretrained(model_path) + + model.eval() + model.requires_grad_(False) + if model.action_head.train_architecture == "lora": + print(f"Merging LoRA weights into main model weights") + # Merge the LoRA weights into the main model weights, and delete the LoRA + # weights to save memory. Note that the WanModel is in model.action_head.model. + model.action_head.model = model.action_head.model.merge_and_unload() + else: + print(f"Skipping merging LoRA weights into main model weights") + + self.eval_bf16 = self.train_cfg.get("eval_bf16", False) + if self.eval_bf16 and not lazy_load: + model = model.to(dtype=torch.bfloat16) + + # Store model initially on CPU if lazy loading + if lazy_load: + model.to(device='cpu') + else: + model.to(device=device) + + # Post initialize, move RoPE freqs to cuda. + model.post_initialize() + + # Parallelize the model across devices. + try: + model.parallelize(device_mesh=device_mesh) + except Exception as e: + print("Skipping parallelization") + + torch.cuda.empty_cache() + + self.trained_model = model + + # 2. Load the action, video, and state transforms + # 2.1. Load the metadata for normalization stats + # We have an assumption: one policy is only for rolling out one type of env, i.e., one embodiment_tag + # metadata_versions = train_cfg.metadata_versions + # metadata = get_metadata(self.embodiment_tag, metadata_versions[self.embodiment_tag.value]) + metadata_path = exp_cfg_dir / "metadata.json" + with open(metadata_path, "r") as f: + metadatas = json.load(f) + if "gr1_unified_offline_rl" in metadatas and self.embodiment_tag.value == "gr1_unified": + self.embodiment_tag = EmbodimentTag.GR1_UNIFIED_OFFLINE_RL + metadata = DatasetMetadata.model_validate(metadatas[self.embodiment_tag.value]) + + # If the model's action head has target_video_height/width (e.g. DreamZero Wan 5B), use that + # as the expected video resolution so the transform matches the model. metadata.json can + # otherwise contain a different resolution (e.g. 180x320) from dataset config. + if hasattr(self.trained_model, "action_head") and hasattr( + self.trained_model.action_head, "config" + ): + cfg = self.trained_model.action_head.config + target_h = getattr(cfg, "target_video_height", None) + target_w = getattr(cfg, "target_video_width", None) + if target_h is not None and target_w is not None and metadata.modalities.video: + for key in metadata.modalities.video.keys(): + metadata.modalities.video[key].resolution = (int(target_w), int(target_h)) + + # 2.2. Get the eval transforms + assert ( + self.embodiment_tag.value in train_cfg.transforms + ), f"{self.embodiment_tag.value=}, {train_cfg.transforms.keys()=}" + eval_transform_cfg = train_cfg.transforms[self.embodiment_tag.value] + if skip_img_transform: + for t in eval_transform_cfg.transforms: + if t._target_ == "groot.vla.data.transform.VideoCrop": + scale = t.scale + for key in metadata.modalities.video.keys(): + metadata.modalities.video[key].resolution = ( + int(metadata.modalities.video[key].resolution[0] * scale), + int(metadata.modalities.video[key].resolution[1] * scale), + ) + elif t._target_ == "groot.vla.data.transform.VideoResize": + height, width = t.height, t.width + for key in metadata.modalities.video.keys(): + metadata.modalities.video[key].resolution = (height, width) + + # assume it's always ComposedModalityTransform + assert ( + eval_transform_cfg._target_ == "groot.vla.data.transform.ComposedModalityTransform" + ) + skipped_transforms = [ + "groot.vla.data.transform.VideoCrop", + "groot.vla.data.transform.VideoResize", + "groot.vla.data.transform.VideoColorJitter", + ] + eval_transform_cfg.transforms = [ + t for t in eval_transform_cfg.transforms if t._target_ not in skipped_transforms + ] + + eval_transform = instantiate(train_cfg.transforms[self.embodiment_tag.value]) + assert isinstance(eval_transform, ComposedModalityTransform), f"{eval_transform=}" + eval_transform.set_metadata(metadata) + + # Set per-horizon statistics for PerHorizonActionTransform if using relative_action_per_horizon + relative_action_per_horizon = self.train_cfg.get('relative_action_per_horizon', False) + print(f"DEBUG: relative_action_per_horizon = {relative_action_per_horizon}") + if relative_action_per_horizon: + # Extract per-horizon statistics from metadata + # The metadata has format: {embodiment: {statistics: {action: {key: {stat: [[h0], [h1], ...]}}}} + action_stats = metadata.statistics.action + print(f"DEBUG: action_stats keys = {list(action_stats.keys())}") + per_horizon_stats = {} + for action_key in action_stats: + stats_dict = action_stats[action_key].model_dump() + print(f"DEBUG: action_key={action_key}, stats_dict keys={list(stats_dict.keys())}") + # Check if stats are per-horizon (2D lists) by examining q01/q99 + if 'q01' in stats_dict: + q01_val = stats_dict['q01'] + print(f"DEBUG: q01 type={type(q01_val)}, value sample={q01_val[:2] if hasattr(q01_val, '__getitem__') else q01_val}") + # Handle both list and numpy array + is_2d = False + if isinstance(q01_val, (list, np.ndarray)) and len(q01_val) > 0: + first_elem = q01_val[0] + print(f"DEBUG: q01[0] type={type(first_elem)}") + if isinstance(first_elem, (list, np.ndarray)): + is_2d = True + + if is_2d: + # This is per-horizon stats (2D array) - convert to list if numpy + if isinstance(q01_val, np.ndarray): + for k in stats_dict: + if isinstance(stats_dict[k], np.ndarray): + stats_dict[k] = stats_dict[k].tolist() + per_horizon_stats[action_key] = stats_dict + print(f"DEBUG: Added {action_key} to per_horizon_stats") + + if per_horizon_stats: + print(f"Setting per-horizon statistics for keys: {list(per_horizon_stats.keys())}") + eval_transform.set_per_horizon_statistics(per_horizon_stats) + else: + print(f"WARNING: No per-horizon statistics found despite relative_action_per_horizon=True") + + eval_transform.eval() + self.eval_transform = eval_transform + + # 3. Load horizons needed + if self.embodiment_tag.value in train_cfg.modality_configs: + self.modality_configs = instantiate( + train_cfg.modality_configs[self.embodiment_tag.value] + ) + else: + self.modality_configs = instantiate(train_cfg.modality_configs) + + self._video_delta_indices = np.array(self.modality_configs.video.eval_delta_indices) + # self._video_delta_indices = np.array([0]) + # self.assert_delta_indices(self._video_delta_indices) + self._video_horizon = len(self._video_delta_indices) + + # We might not use `state`, which will be a vision-only policy + if "state" in self.modality_configs: + self._state_delta_indices = np.array(self.modality_configs.state.eval_delta_indices) + if not skip_assert_delta_indices: + self.assert_delta_indices(self._state_delta_indices) + self._state_horizon = len(self._state_delta_indices) + else: + self._state_horizon = None + self._state_delta_indices = None + self._raw_data_image_transform = None + + def offload_to_cpu(self): + """Offload the model to CPU to free GPU memory.""" + if hasattr(self.trained_model, "action_head") and hasattr(self.trained_model.action_head, "image_encoder"): + # For models with vram management, disable it and move to CPU + if hasattr(self.trained_model.action_head, 'disable_vram_management'): + self.trained_model.action_head.disable_vram_management() + + self.trained_model.to(device='cpu') + torch.cuda.empty_cache() + print(f"Model offloaded to CPU") + + def load_to_gpu(self): + """Load the model to GPU for inference.""" + print(f"Loading model to GPU...") + + # Move model to GPU + if hasattr(self.trained_model, "action_head") and hasattr(self.trained_model.action_head, "image_encoder"): + self.trained_model.action_head.enable_vram_management() + else: + self.trained_model.to(device=self.device) + + # Apply bf16 if needed + if self.eval_bf16: + self.trained_model = self.trained_model.to(dtype=torch.bfloat16) + + torch.cuda.empty_cache() + + def ensure_model_on_gpu(self): + """Ensure the model is loaded on GPU before inference.""" + # Check if model is on CPU + model_device = next(self.trained_model.parameters()).device + if model_device.type == 'cpu': + self.load_to_gpu() + + def assert_delta_indices(self, delta_indices: np.ndarray): + # All delta indices should be non-positive because there's no way to get the future observations + assert np.all(delta_indices <= 0), f"{delta_indices=}" + # The last delta index should be 0 because it doesn't make sense to not use the latest observation + assert delta_indices[-1] == 0, f"{delta_indices=}" + if len(delta_indices) > 1: + # The step is consistent + assert np.all( + np.diff(delta_indices) == delta_indices[1] - delta_indices[0] + ), f"{delta_indices=}" + # And the step is positive + assert (delta_indices[1] - delta_indices[0]) > 0, f"{delta_indices=}" + + def apply(self, batch: Batch, **kwargs) -> Batch: + """Normalize inputs""" + obs = batch.obs + + normalized_input = self.eval_transform(obs) + batch.normalized_obs = normalized_input + return batch + + def unapply(self, batch: Batch, obs: dict = None, **kwargs): + """Unnormalize actions and convert relative actions to absolute if needed""" + unnormalized_action = self.eval_transform.unapply( + dict(action=batch.normalized_action.cpu()) + ) + + # Check if relative_action is enabled and convert relative to absolute + relative_action = self.train_cfg.get('relative_action', False) + relative_action_per_horizon = self.train_cfg.get('relative_action_per_horizon', False) + relative_action_keys = self.train_cfg.get('relative_action_keys', []) + print("relative_action_per_horizon", relative_action_per_horizon) + if (relative_action or relative_action_per_horizon) and relative_action_keys and obs is not None: + for key in relative_action_keys: + action_key = f"action.{key}" + state_key = f"state.{key}" + + if action_key not in unnormalized_action: + continue + + # Try to find the state data - check multiple possible key formats + last_state = None + + + if last_state is None and state_key in obs: + # Format 1: Direct key like "state.joint_position" + last_state = obs[state_key] + elif last_state is None: + # Format 2: Search for keys containing both "state" and the key name + for obs_key in obs.keys(): + if 'state' in obs_key and key in obs_key: + last_state = obs[obs_key] + break + + # Format 3: If key is "joint_position" and obs has "state" key directly + # This handles cases where the observation uses modality-level keys + if last_state is None and 'state' in obs: + state_data = obs['state'] + # Check if the state data shape matches the action shape + action_dim = unnormalized_action[action_key].shape[-1] + if torch.is_tensor(state_data): + state_dim = state_data.shape[-1] + elif isinstance(state_data, np.ndarray): + state_dim = state_data.shape[-1] + else: + state_dim = None + + if state_dim == action_dim: + last_state = state_data + + if last_state is None: + continue + + if torch.is_tensor(last_state): + last_state = last_state.cpu().numpy() + + # Shape is (B, T, D) or (T, D), we want the last timestep + # After indexing: (B, D) or (D,) + if len(last_state.shape) >= 2: + last_state = last_state[..., -1, :] # Get the last timestep + + # Action shape is (horizon, D) or (B, horizon, D) + # Expand dims to broadcast: (D,) -> (1, D) or (B, D) -> (B, 1, D) + if len(unnormalized_action[action_key].shape) > len(last_state.shape): + last_state = np.expand_dims(last_state, axis=-2) # Add horizon dimension + + # Add state to relative action to get absolute action + print("last_state", last_state.shape, "unnormalized_action[action_key]", unnormalized_action[action_key].shape) + unnormalized_action[action_key] = unnormalized_action[action_key] + last_state + + batch.act = unnormalized_action + return batch + + def forward(self, batch, state=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.get_action(normalized_input) + normalized_action = model_pred["action_pred"].float() + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch + + def joint_forward(self, batch, video=None, state=None, **kwargs): + # 0. Ensure model is on GPU + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.joint_video_action(normalized_input) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch, video_pred + + def lazy_joint_forward_causal(self, batch, video=None, latent_video=None, state=None, video_only=False, **kwargs): + + transform_start_time = time.perf_counter() + + # Save original observation before any modification (for relative action conversion) + original_obs_for_relative = {k: v.copy() if isinstance(v, np.ndarray) else v.clone() if torch.is_tensor(v) else v + for k, v in batch.obs.items()} + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + # Also unsqueeze the saved original obs + original_obs_for_relative = unsqueeze_dict_values(original_obs_for_relative) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + transform_time = time.perf_counter() - transform_start_time + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + model_start_time = time.perf_counter() + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.lazy_joint_video_action_causal(normalized_input, latent_video=latent_video) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + model_time = time.perf_counter() - model_start_time + + untransform_start_time = time.perf_counter() + + # 4. Unnormalize actions (pass obs for relative action conversion) + if not video_only: + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs_for_relative) + else: + batch = Batch(normalized_action=normalized_action) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + + untransform_time = time.perf_counter() - untransform_start_time + total_time = transform_time + model_time + untransform_time + + if self.rank == 0: + print(f"Inference Time: Total {total_time:.3f} seconds, " + f"Transform: {transform_time:.3f} seconds, " + f"Model: {model_time:.3f} seconds, " + f"Untransform: {untransform_time:.3f} seconds") + + return batch, video_pred + + def lazy_joint_forward_causal_gt_cond(self, batch, video=None, latent_video=None, state=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.lazy_joint_video_action_causal_gt_cond(normalized_input, latent_video=latent_video) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch, video_pred + + def lazy_joint_forward(self, batch, video=None, state=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.lazy_joint_video_action(normalized_input) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch, video_pred + + def lazy_joint_forward_efficient(self, batch, video=None, state=None, prompt_embs=None, prompt_emb_nega=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.lazy_joint_video_action_efficient(normalized_input, prompt_embs=prompt_embs, prompt_emb_nega=prompt_emb_nega) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch, video_pred + + def language_forward(self, batch, state=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.get_language(normalized_input) + # normalized_action = model_pred["action_pred"].float() + output = model_pred["output"] + # 4. Unnormalize actions + # batch = self.unapply(Batch(normalized_action=normalized_action)) + + # # 5. Remove batch dimension if we added it + # if not is_batched: + # batch.act = squeeze_dict_values(batch.act) + return output + # return batch + + + def gt_video_action_pred(self, batch, video=None, state=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.gt_video_action_pred(normalized_input) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch, video_pred + + def _check_state_is_batched(self, obs: dict[str, Any]) -> bool: + for k, v in obs.items(): + if "state" in k and len(v.shape) < 3: # (B, Time, Dim) + return False + return True + + @property + def raw_data_image_transform(self) -> Callable: + """ + Get the raw data image transform for the policy + """ + return lambda x: x + + @property + def state_delta_indices(self) -> np.ndarray: + return self._state_delta_indices + + @property + def video_delta_indices(self) -> np.ndarray: + return self._video_delta_indices + + +def unsqueeze_dict_values(data: dict[str, Any]) -> dict[str, Any]: + """ + Unsqueeze the values of a dictionary. + This converts the data to be batched of size 1. + """ + unsqueezed_data = {} + for k, v in data.items(): + if isinstance(v, np.ndarray): + unsqueezed_data[k] = np.expand_dims(v, axis=0) + elif isinstance(v, list): + unsqueezed_data[k] = np.array(v) + elif isinstance(v, torch.Tensor): + unsqueezed_data[k] = v.unsqueeze(0) + elif isinstance(v, str): + unsqueezed_data[k] = np.array([v]) + else: + unsqueezed_data[k] = v + return unsqueezed_data + + +def squeeze_dict_values(data: dict[str, Any]) -> dict[str, Any]: + """ + Squeeze the values of a dictionary. This removes the batch dimension. + """ + squeezed_data = {} + for k, v in data.items(): + if isinstance(v, np.ndarray): + squeezed_data[k] = np.squeeze(v) + elif isinstance(v, torch.Tensor): + squeezed_data[k] = v.squeeze() + else: + squeezed_data[k] = v + return squeezed_data + + +def tree_get_leading_dim(data: dict[str, Any]) -> int: + """ + Get the batch size (leading dimension) of all tensors in the structure. + + Args: + data: Nested structure containing tensors + strict: If True, raises error if batch sizes don't match. + If False, returns None if inconsistent. + + Returns: + The common batch size, or None if inconsistent (when strict=False) + + Raises: + ValueError: If batch sizes are inconsistent (when strict=True) + ValueError: If no tensors found + """ + batch_sizes = [] + tensor_paths = [] + + def collect_batch_sizes(path, x): + if hasattr(x, "shape") and len(x.shape) > 0: + batch_sizes.append(x.shape[0]) + tensor_paths.append(".".join(map(str, path))) + + # Collect all batch sizes with their paths + tree.map_structure_with_path(collect_batch_sizes, data) + + if not batch_sizes: + raise ValueError("No tensors found in the data structure") + + # Check if all batch sizes are the same + first_batch_size = batch_sizes[0] + + if not all(bs == first_batch_size for bs in batch_sizes): + inconsistent_info = [f"{path}: {bs}" for path, bs in zip(tensor_paths, batch_sizes)] + error_msg = "Inconsistent batch sizes found:\n" + "\n".join(inconsistent_info) + + raise ValueError(error_msg) + + return first_batch_size + + +class GrootSimRLPolicy(BaseGrootSimPolicy): + """ + A class for Offline RL policies, + """ + + def __init__( + self, + embodiment_tag: EmbodimentTag, + model_path: str, + device: int | str, + q_fn_path: str | None = None, + n_action_samples: int = 1, + n_denoising_steps: int = 4, + gaussian_std: float = 0.0, + inference_batch_size: int = 1, + model_config_overrides: list[str] | None = [], + ): + """ + Implicit Policy Extraction (IPE) GrootSimPolicy. + """ + super().__init__(embodiment_tag=embodiment_tag, model_path=model_path, device=device) + + self.base_policy = GrootSimPolicy( + embodiment_tag=embodiment_tag, + model_path=str(model_path), + device=device, + ) + self.base_policy.trained_model.action_head.num_inference_timesteps = n_denoising_steps + + if q_fn_path is None: + if n_action_samples > 1: + raise ValueError("If n_action_samples > 1, q_fn_path should be provided.") + self.qfn = None + else: + assert n_action_samples > 1, "If q_fn_path is provided, n_action_samples should be > 1." + self.qfn = GrootSimPolicy( + embodiment_tag=EmbodimentTag.GR1_UNIFIED_OFFLINE_RL, + model_path=q_fn_path, + device=device, + skip_assert_delta_indices=True, + ) + + self.n_action_samples = n_action_samples + self.gaussian_std = gaussian_std + self.inference_batch_size = inference_batch_size + + def batch_obs(self, obs: dict[str, Any]) -> tuple[dict[str, Any], bool]: + """ + Check if input is batched and add batch dimension if needed + """ + is_batched = self.base_policy._check_state_is_batched(obs) + if not is_batched: + obs = unsqueeze_dict_values(obs) + + return obs, is_batched + + def add_action_noise(self, action: torch.Tensor, std: float) -> torch.Tensor: + if std <= 0.0: + return action + noise = torch.randn_like(action) * std + noisy_action = action + noise + return noisy_action + + def sample_actions_from_base_policy( + self, + batch: Batch, + n_action_samples: int, # number of actions to sample PER item in the batch + inference_batch_size: int = 1, + ) -> torch.Tensor: + # a) Apply transforms/normalization + batch = self.base_policy.apply(batch) + normalized_input = batch.normalized_obs + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.base_policy.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + with torch.inference_mode(): + model_pred = self.base_policy.trained_model.get_action( + normalized_input, + num_action_samples=n_action_samples, + inference_batch_size=inference_batch_size, + validate=False, + ) + normalized_action = model_pred["action_pred"] + normalized_action = self.add_action_noise(normalized_action, self.gaussian_std) + return normalized_action # [num_envs, n_action_samples, horizon, action_dim] + + def get_action_values( + self, batch: Batch, actions: torch.Tensor, inference_batch_size: int = 1 + ) -> torch.Tensor: + assert self.qfn is not None, "Q-function must be provided to evaluate actions." + + # a) Apply transforms/normalization + batch = self.qfn.apply(batch) + normalized_input = batch.normalized_obs + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.qfn.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + with torch.inference_mode(): + # [num_envs, ...] + backbone_inputs, action_inputs = self.qfn.trained_model.prepare_input(normalized_input) + backbone_outputs = self.qfn.trained_model.backbone(backbone_inputs) + q_values = self.qfn.trained_model.action_head.get_q_pred( + backbone_outputs, + action_inputs, + actions, + inference_batch_size=inference_batch_size, + ) + + return q_values + + def forward(self, batch, state=None, **kwargs): + orig_obs = batch.obs.copy() + orig_obs, is_batched = self.batch_obs(orig_obs) + + # Get number of environments + num_envs = 1 if not is_batched else tree_get_leading_dim(orig_obs) + + # 1/ Sample normalized actions from the base policy + actions = self.sample_actions_from_base_policy( + Batch(obs=orig_obs), + n_action_samples=self.n_action_samples, + inference_batch_size=self.inference_batch_size, + ) # [num_envs, n_action_samples, horizon, action_dim] + + # 2/ Evaluate the sampled actions with the Q-function + if self.qfn is not None: + q_values = self.get_action_values( + Batch(obs=orig_obs), + actions=actions, + inference_batch_size=self.inference_batch_size, + ) + + best_indices = torch.argmax(q_values, dim=1) # [num_envs] + normalized_action = actions[ + torch.arange(num_envs), best_indices + ].float() # [num_envs, horizon, action_dim] + else: + normalized_action = actions.float() + + batch = self.base_policy.unapply(Batch(normalized_action=normalized_action), obs=orig_obs) + + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch + + @property + def video_delta_indices(self) -> np.ndarray: + return self.base_policy.video_delta_indices + + @property + def state_delta_indices(self) -> np.ndarray: + return self.base_policy.state_delta_indices + + @property + def raw_data_image_transform(self) -> Callable: + return lambda x: x diff --git a/groot/vla/utils/__init__.py b/groot/vla/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/utils/action_args_override_utils.py b/groot/vla/utils/action_args_override_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..67ecdbb307bf02ebef0282a1d5065f4d3bd8b24f --- /dev/null +++ b/groot/vla/utils/action_args_override_utils.py @@ -0,0 +1,68 @@ +""" +Utility functions for automatically updating action horizon and related configurations. +""" + +from omegaconf import DictConfig, open_dict + + +def update_action_horizon_configs(cfg: DictConfig, action_horizon: int) -> DictConfig: + """ + Automatically update action_horizon and corresponding delta_indices in all modality configs. + + Args: + cfg: The hydra configuration + action_horizon: The desired action horizon (e.g., 30) + + Returns: + Updated configuration with action_horizon and delta_indices set appropriately + """ + # Generate delta_indices for the given action_horizon [0, 1, 2, ..., action_horizon-1] + delta_indices = list(range(action_horizon)) + + # Update the global action_horizon + with open_dict(cfg): + cfg.action_horizon = action_horizon + if hasattr(cfg.model, "vla_override_kwargs"): + cfg.model.vla_override_kwargs.action_horizon = action_horizon + if hasattr(cfg.model, "action_head_override_kwargs"): + cfg.model.action_head_override_kwargs.action_horizon = action_horizon + + # Update delta_indices for all action modalities in modality_configs + if hasattr(cfg, "modality_configs"): + for embodiment_name, modality_config in cfg.modality_configs.items(): + if hasattr(modality_config, "action"): + # Update the action delta_indices + modality_config.action.delta_indices = delta_indices + print(f"Updated {embodiment_name}.action.delta_indices to {delta_indices}") + + return cfg + + +def update_action_dim_configs(cfg: DictConfig, new_action_dim: int) -> DictConfig: + """ + Update the action dimension in all modality configs. + """ + with open_dict(cfg): + cfg.max_action_dim = new_action_dim + return cfg + + +def apply_action_overrides(cfg: DictConfig) -> DictConfig: + """ + Apply action horizon overrides if action_horizon is specified in config. + This function should be called after the config is loaded but before model instantiation. + """ + if hasattr(cfg.model, "action_head_override_kwargs"): + action_horizon = cfg.action_horizon + print(f"Applying action_horizon={action_horizon} overrides...") + cfg = update_action_horizon_configs(cfg, action_horizon) + + if hasattr(cfg.model, "expand_action_head_kwargs"): + expand_action_head_kwargs = cfg.model.expand_action_head_kwargs + if "expand_action_dim" in expand_action_head_kwargs: + old_action_dim = expand_action_head_kwargs.expand_action_dim.old_action_dim + new_action_dim = expand_action_head_kwargs.expand_action_dim.new_action_dim + print(f"Applying expand_action_dim={old_action_dim}->{new_action_dim} overrides...") + cfg = update_action_dim_configs(cfg, new_action_dim) + + return cfg diff --git a/groot/vla/utils/timer.py b/groot/vla/utils/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..d4e17daa2caee54c85fd7869195136cd2b098a87 --- /dev/null +++ b/groot/vla/utils/timer.py @@ -0,0 +1,42 @@ +import time + + +class ContextTimer: + + def __init__(self, trainer): + self.last_key = None + self.trainer = trainer + self.start_times = {} + self.key_stack = [] + + def with_label(self, key): + self.last_key = key + return self + + def __enter__(self): + self.key_stack.append(self.last_key) # Push key to stack + self.start_times[self.last_key] = time.time() # Start timing for this key + return self + + def __exit__(self, exc_type, exc_value, traceback): + key = self.key_stack.pop() # Pop key from stack + diff = time.time() - self.start_times[key] + self.trainer.log({f"{key}_time": diff}) + # print(f"{key}: {diff:.2f} seconds") + + +if __name__ == "__main__": + + class MockTrainer: + def log(self, data): + print("Logging:", data) + + trainer = MockTrainer() + my_timer = ContextTimer(trainer) + + with my_timer.with_label("outer"): + time.sleep(1) + with my_timer.with_label("inner"): + time.sleep(2) + with my_timer.with_label("inner"): # Another inner block + time.sleep(1) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..6dff6220af5cc86bbf0061af39b0a3b72b344d7e --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,41 @@ +[build-system] +requires = ["setuptools>=67", "wheel", "pip"] +build-backend = "setuptools.build_meta" + +[project] +name = "dreamzero" +description = "DreamZero: World Action Models Are Zero-Shot Policies" +version = "1.0.0" +readme = "README.md" +requires-python = "~=3.11,<3.13" +dependencies = [ + "torch>=2.3,<2.7", + "torchvision>=0.18", + "transformers==4.51.3", + "diffusers>=0.30,<0.32", + "accelerate>=0.27", + "deepspeed>=0.12,<0.16", + "hydra-core>=1.3", + "omegaconf>=2.3", + "decord", + "opencv-python>=4.8", + "av", + "einops>=0.8", + "safetensors", + "sentencepiece", + "protobuf", + "zarr", + "h5py", + "pyarrow", + "pandas", + "scipy", + "numpy<2", + "tqdm", + "bitsandbytes>=0.41", + "wandb", + "tensorboard", +] + +[tool.setuptools.packages.find] +where = ["."] +include = ["groot*"] diff --git a/python3.11.80s-22193,autodl-container-lezfbu0bdx-b16a.btr b/python3.11.80s-22193,autodl-container-lezfbu0bdx-b16a.btr new file mode 100644 index 0000000000000000000000000000000000000000..d7f15b0ee30af5b9438529f0943ccf09196c2442 --- /dev/null +++ b/python3.11.80s-22193,autodl-container-lezfbu0bdx-b16a.btr @@ -0,0 +1,60 @@ + +python3.11:22193 terminated with signal 11 at PC=7f84cebaa0c0 SP=7f7e3effc450. Backtrace: +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine11ready_queueESt10shared_ptrINS0_10ReadyQueueEEN3c106DeviceE+0x100)[0x7f84cebaa0c0] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine17evaluate_functionERSt10shared_ptrINS0_9GraphTaskEEPNS0_4NodeERNS0_11InputBufferERKS2_INS0_10ReadyQueueEE+0x135d)[0x7f84cebb4bdd] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine11thread_mainERKSt10shared_ptrINS0_9GraphTaskEE+0x5c9)[0x7f84cebb5669] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine11thread_initEiRKSt10shared_ptrINS0_10ReadyQueueEEb+0x337)[0x7f84cebaa977] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x8117e2)[0x7f84dec0c7e2] +/usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0xdc253)[0x7f85686b0253] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x94ac3)[0x7f85acf65ac3] +/usr/lib/x86_64-linux-gnu/libc.so.6(clone+0x44)[0x7f85acff6a04] +S7_+0x755)[0x7f84cebb30b5] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x811b46)[0x7f84dec0cb46] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x811f56)[0x7f84dec0cf56] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x51ad67] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0x1c5)[0x5412a5] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c8b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyFunction_Vectorcall+0x173)[0x531823] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c8b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c8b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c8b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c852] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c8b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x5bc60c] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyObject_MakeTpCall+0x22c)[0x4e75dc] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Vectorcall+0x35)[0x50d745] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x8f2)[0x4fb152] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c852] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x62e1b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyEval_EvalCode+0x97)[0x4f3a67] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x647c37] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x645350] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x650d15] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyRun_SimpleFileObject+0x194)[0x650a64] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyRun_AnyFileObject+0x43)[0x650833] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(Py_RunMain+0x327)[0x64f787] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(Py_BytesMain+0x2d)[0x61ee0d] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x29d90)[0x7f85acefad90] +/usr/lib/x86_64-linux-gnu/libc.so.6(__libc_start_main+0x80)[0x7f85acefae40] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_start+0x25)[0x61ec95] diff --git a/python3.11.80s-22194,autodl-container-lezfbu0bdx-b16a.btr b/python3.11.80s-22194,autodl-container-lezfbu0bdx-b16a.btr new file mode 100644 index 0000000000000000000000000000000000000000..88a9afac6ac6bea44280e54c6dc42cea5227e284 --- /dev/null +++ b/python3.11.80s-22194,autodl-container-lezfbu0bdx-b16a.btr @@ -0,0 +1,79 @@ + +python3.11:22194 terminated with signal 6 at PC=7f47027f79fc SP=7f42e8a59600. Backtrace: +/usr/lib/x86_64-linux-gnu/libc.so.6(pthread_kill+0x12c)[0x7f47027f79fc] +/usr/lib/x86_64-linux-gnu/libc.so.6(raise+0x16)[0x7f47027a3476] +/usr/lib/x86_64-linux-gnu/libc.so.6(abort+0xd3)[0x7f47027897f3] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x89676)[0x7f47027ea676] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0xa0cfc)[0x7f4702801cfc] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0xa17cc)[0x7f47028027cc] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0xa1969)[0x7f4702802969] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0xa3bdb)[0x7f4702804bdb] +/usr/lib/x86_64-linux-gnu/libc.so.6(malloc+0x99)[0x7f4702806139] +/usr/lib/x86_64-linux-gnu/libc.so.6(__backtrace_symbols+0x192)[0x7f4702894d92] +/usr/lib/x86_64-linux-gnu/libucs.so.0(ucs_debug_backtrace_create+0x78)[0x7f428ddc60b8] +/usr/lib/x86_64-linux-gnu/libucs.so.0(ucs_debug_print_backtrace+0x39)[0x7f428ddc6459] +/usr/lib/x86_64-linux-gnu/libucs.so.0(ucs_handle_error+0x2e4)[0x7f428ddc6fc4] +/usr/lib/x86_64-linux-gnu/libucs.so.0(+0x24fec)[0x7f428ddcafec] +/usr/lib/x86_64-linux-gnu/libucs.so.0(+0x251aa)[0x7f428ddcb1aa] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x42520)[0x7f47027a3520] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x453fc)[0x7f47027a63fc] +/usr/lib/x86_64-linux-gnu/libc.so.6(on_exit+0x0)[0x7f47027a6610] +/usr/lib/x86_64-linux-gnu/libinfinipath.so.4(+0x42a7)[0x7f430a33c2a7] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x42520)[0x7f47027a3520] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x91117)[0x7f47027f2117] +/usr/lib/x86_64-linux-gnu/libc.so.6(pthread_cond_wait+0x211)[0x7f47027f4a41] +/usr/lib/x86_64-linux-gnu/libstdc++.so.6(_ZNSt18condition_variable4waitERSt11unique_lockISt5mutexE+0x17)[0x7f46bde80747] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(+0x5e48abb)[0x7f46243b0abb] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine11thread_mainERKSt10shared_ptrINS0_9GraphTaskEE+0x6f)[0x7f46243b510f] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine23execute_with_graph_taskERKSt10shared_ptrINS0_9GraphTaskEES2_INS0_4NodeEEONS0_11InputBufferE+0x423)[0x7f46243b02b3] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x811bb5)[0x7f463440cbb5] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine7executeERKSt6vectorINS0_4EdgeESaIS3_EERKS2_IN2at6TensorESaIS9_EEbbbS7_+0x755)[0x7f46243b30b5] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x811b46)[0x7f463440cb46] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x811f56)[0x7f463440cf56] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x51ad67] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0x1c5)[0x5412a5] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c8b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyFunction_Vectorcall+0x173)[0x531823] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c8b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c8b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c8b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c852] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c8b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0xb5)[0x541195] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x5bc60c] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyObject_MakeTpCall+0x22c)[0x4e75dc] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Vectorcall+0x35)[0x50d745] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x8f2)[0x4fb152] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55e097] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c852] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x62e1b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyEval_EvalCode+0x97)[0x4f3a67] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x647c37] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x645350] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x650d15] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyRun_SimpleFileObject+0x194)[0x650a64] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyRun_AnyFileObject+0x43)[0x650833] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(Py_RunMain+0x327)[0x64f787] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(Py_BytesMain+0x2d)[0x61ee0d] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x29d90)[0x7f470278ad90] +/usr/lib/x86_64-linux-gnu/libc.so.6(__libc_start_main+0x80)[0x7f470278ae40] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_start+0x25)[0x61ec95] diff --git a/python3.11.80s-2617,autodl-container-lezfbu0bdx-b16a.btr b/python3.11.80s-2617,autodl-container-lezfbu0bdx-b16a.btr new file mode 100644 index 0000000000000000000000000000000000000000..4c39378bbd74f691befdcba9524dc56364ec837b --- /dev/null +++ b/python3.11.80s-2617,autodl-container-lezfbu0bdx-b16a.btr @@ -0,0 +1,63 @@ + +python3.11:2617 terminated with signal 11 at PC=7fc01b0573fc SP=7fff96b24270. Backtrace: +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x453fc)[0x7fc01b0573fc] +/usr/lib/x86_64-linux-gnu/libc.so.6(on_exit+0x0)[0x7fc01b057610] +/usr/lib/x86_64-linux-gnu/libinfinipath.so.4(+0x42a7)[0x7fbc00c4a2a7] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x42520)[0x7fc01b054520] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x91117)[0x7fc01b0a3117] +/usr/lib/x86_64-linux-gnu/libc.so.6(pthread_cond_wait+0x211)[0x7fc01b0a5a41] +/usr/lib/x86_64-linux-gnu/libstdc++.so.6(_ZNSt18condition_variable4waitERSt11unique_lockISt5mutexE+0x17)[0x7fbfd6680747] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(+0x5e48abb)[0x7fbf3cbb0abb] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine11thread_mainERKSt10shared_ptrINS0_9GraphTaskEE+0x6f)[0x7fbf3cbb510f] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine23execute_with_graph_taskERKSt10shared_ptrINS0_9GraphTaskEES2_INS0_4NodeEEONS0_11InputBufferE+0x423)[0x7fbf3cbb02b3] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x811bb5)[0x7fbf4cc0cbb5] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine7executeERKSt6vectorINS0_4EdgeESaIS3_EERKS2_IN2at6TensorESaIS9_EEbbbS7_+0x755)[0x7fbf3cbb30b5] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x811b46)[0x7fbf4cc0cb46] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x811f56)[0x7fbf4cc0cf56] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x51ad67] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Call+0x1c5)[0x5412a5] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so(+0x379da1a)[0x7fbf04be2a1a] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so(+0x38aee78)[0x7fbf04cf3e78] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN2at18TensorIteratorBase26allocate_or_resize_outputsEv+0x234)[0x7fbf3878aa64] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN2at18TensorIteratorBase5buildERNS_20TensorIteratorConfigE+0x2fb)[0x7fbf3878e75b] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN2at18TensorIteratorBase50build_borrowing_except_last_argument_comparison_opERKNS_10TensorBaseES3_S3_+0xea)[0x7fbf3879007a] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN2at4meta20structured_eq_Scalar4metaERKNS_6TensorERKN3c106ScalarE+0x3e)[0x7fbf38b1df2e] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so(+0x38f1093)[0x7fbf04d36093] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so(+0x38f114f)[0x7fbf04d3614f] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN2at4_ops9eq_Scalar4callERKNS_6TensorERKN3c106ScalarE+0x186)[0x7fbf3954ba16] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN2at6native5isinfERKNS_6TensorE+0x270)[0x7fbf38f20690] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(+0x31bd51c)[0x7fbf39f2551c] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN2at4_ops5isinf10redispatchEN3c1014DispatchKeySetERKNS_6TensorE+0xde)[0x7fbf3983b8ae] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(+0x54055a8)[0x7fbf3c16d5a8] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN2at4_ops5isinf4callERKNS_6TensorE+0x16d)[0x7fbf398d8c3d] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x46f363)[0x7fbf4c86a363] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x5321ce] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyObject_Vectorcall+0x35)[0x50d745] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x8f2)[0x4fb152] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyFunction_Vectorcall+0x173)[0x531823] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x82cf57)[0x7fbf4cc27f57] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x82e8d4)[0x7fbf4cc298d4] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd14AccumulateGrad5applyEOSt6vectorIN2at6TensorESaIS4_EE+0xdb)[0x7fbf3cbbd35b] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(+0x5e52ecb)[0x7fbf3cbbaecb] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine17evaluate_functionERSt10shared_ptrINS0_9GraphTaskEEPNS0_4NodeERNS0_11InputBufferERKS2_INS0_10ReadyQueueEE+0x1117)[0x7fbf3cbb4997] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine11thread_mainERKSt10shared_ptrINS0_9GraphTaskEE+0x5c9)[0x7fbf3cbb5669] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch8autograd6Engine11thread_initEiRKSt10shared_ptrINS0_10ReadyQueueEEb+0x337)[0x7fbf3cbaa977] +/root/autodl-tmp/venvs/dreamzero/lib/python3.11/site-packages/torch/lib/libtorch_python.so(+0x8117e2)[0x7fbf4cc0c7e2] +/usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0xdc253)[0x7fbfd66b0253] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x94ac3)[0x7fc01b0a6ac3] +/usr/lib/x86_64-linux-gnu/libc.so.6(clone+0x44)[0x7fc01b137a04] +[0x0] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x55c852] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyEval_EvalFrameDefault+0x4781)[0x4fefe1] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x62e1b4] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(PyEval_EvalCode+0x97)[0x4f3a67] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x647c37] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x645350] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11[0x650d15] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyRun_SimpleFileObject+0x194)[0x650a64] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_PyRun_AnyFileObject+0x43)[0x650833] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(Py_RunMain+0x327)[0x64f787] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(Py_BytesMain+0x2d)[0x61ee0d] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x29d90)[0x7fc01b03bd90] +/usr/lib/x86_64-linux-gnu/libc.so.6(__libc_start_main+0x80)[0x7fc01b03be40] +/root/autodl-tmp/venvs/dreamzero/bin/python3.11(_start+0x25)[0x61ec95] diff --git a/python3.11.80s-2618,autodl-container-lezfbu0bdx-b16a.btr b/python3.11.80s-2618,autodl-container-lezfbu0bdx-b16a.btr new file mode 100644 index 0000000000000000000000000000000000000000..2c22297d997fb4a0588dff67b092d4c95df28347 --- /dev/null +++ b/python3.11.80s-2618,autodl-container-lezfbu0bdx-b16a.btr @@ -0,0 +1,12 @@ + +python3.11:2618 terminated with signal 11 at PC=7f9ff63113fc SP=7f9ec2691ef0. Backtrace: +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x453fc)[0x7f9ff63113fc] +/usr/lib/x86_64-linux-gnu/libc.so.6(on_exit+0x0)[0x7f9ff6311610] +/usr/lib/x86_64-linux-gnu/libinfinipath.so.4(+0x42a7)[0x7f9c0c6cc2a7] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x42520)[0x7f9ff630e520] +/usr/lib/x86_64-linux-gnu/libc.so.6(__poll+0x4f)[0x7f9ff63e4bcf] +/usr/lib/x86_64-linux-gnu/libcuda.so.1(+0x2f8527)[0x7f9fb927d527] +/usr/lib/x86_64-linux-gnu/libcuda.so.1(+0x3d54e7)[0x7f9fb935a4e7] +/usr/lib/x86_64-linux-gnu/libcuda.so.1(+0x2e3f33)[0x7f9fb9268f33] +/usr/lib/x86_64-linux-gnu/libc.so.6(+0x94ac3)[0x7f9ff6360ac3] +/usr/lib/x86_64-linux-gnu/libc.so.6(clone+0x44)[0x7f9ff63f1a04] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..c10a2462cd9b5a7f01e8d39e88b5a3b9a1f53deb --- /dev/null +++ b/requirements.txt @@ -0,0 +1,46 @@ +# DreamZero 依赖清单 +# 安装: pip install -r requirements.txt +# 或完整安装: pip install -e .[train] + +# 核心深度学习框架 +torch>=2.3,<2.7 +torchvision>=0.18 +transformers==4.51.3 +diffusers>=0.30,<0.32 +accelerate>=0.27 + +# 分布式训练 +deepspeed>=0.12,<0.16 + +# 配置 +hydra-core>=1.3 +omegaconf>=2.3 + +# 数据/视频处理 +decord +opencv-python>=4.8 +av +einops>=0.8 + +# 序列化 +safetensors +sentencepiece +protobuf + +# 数据格式 +zarr +h5py +pyarrow +pandas +scipy + +# 数值 +numpy<2 +tqdm + +# 显存优化 +bitsandbytes>=0.41 + +# 实验追踪 +wandb +tensorboard diff --git a/scripts/cluster/slurm_train.sh b/scripts/cluster/slurm_train.sh new file mode 100644 index 0000000000000000000000000000000000000000..d84b1ab116aab742accf58866aa7947547fe3895 --- /dev/null +++ b/scripts/cluster/slurm_train.sh @@ -0,0 +1,95 @@ +#!/bin/bash +# DreamZero SLURM 多节点训练启动脚本 +# +# 用法: +# sbatch scripts/cluster/slurm_train.sh libero # 默认 4 节点 × 8 GPU +# sbatch --nodes=8 scripts/cluster/slurm_train.sh manifeel # 8 节点 +# +# 或者覆盖参数: +# sbatch scripts/cluster/slurm_train.sh robotwin \ +# --max_steps=100000 --training_args.learning_rate=5e-6 + +#SBATCH --job-name=dreamzero-sft +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=1 +#SBATCH --gpus-per-node=8 +#SBATCH --cpus-per-task=64 +#SBATCH --mem-per-cpu=8G +#SBATCH --time=72:00:00 +#SBATCH --output=logs/%x-%j.out +#SBATCH --error=logs/%x-%j.err +#SBATCH --partition=gpu + +set -euo pipefail + +# ============ 参数解析 ============ +BENCHMARK=${1:-libero} +shift || true + +# ============ 环境变量 ============ +export NCCL_DEBUG=WARN +export NCCL_IB_DISABLE=0 +export NCCL_IB_TIMEOUT=22 +export NCCL_IB_RETRY_CNT=4 +export NCCL_SOCKET_IFNAME=^docker0,lo +export HF_ENDPOINT=https://hf-mirror.com + +NODES=${SLURM_NNODES:-4} +GPUS_PER_NODE=${SLURM_GPUS_PER_NODE:-8} +TOTAL_GPUS=$((NODES * GPUS_PER_NODE)) + +# 自动获取 MASTER_ADDR +if [ -n "${SLURM_NODELIST:-}" ]; then + MASTER_ADDR=$(scontrol show hostname "$SLURM_NODELIST" | head -n1) +else + MASTER_ADDR="localhost" +fi +MASTER_PORT=${MASTER_PORT:-29500} + +echo "============================================" +echo " DreamZero SLURM Training" +echo "============================================" +echo " Nodes: $NODES" +echo " GPUs/node: $GPUS_PER_NODE" +echo " Total GPUs: $TOTAL_GPUS" +echo " Master: $MASTER_ADDR:$MASTER_PORT" +echo " Benchmark: $BENCHMARK" +echo "============================================" + +# ============ 路径 ============ +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/output/${BENCHMARK}_full_${NODES}nodes}" + +# ============ 运行训练 ============ +torchrun \ + --nnodes="$NODES" \ + --nproc_per_node="$GPUS_PER_NODE" \ + --rdzv_id="dz_sft_${BENCHMARK}" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + "$REPO_ROOT/groot/vla/experiment/experiment.py" \ + report_to=wandb \ + data="dreamzero/${BENCHMARK}" \ + wandb_project=dreamzero-sft \ + train_architecture=full \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ + model/dreamzero/transform=dreamzero_cotrain \ + per_device_train_batch_size=1 \ + global_batch_size=$((TOTAL_GPUS * 4)) \ + save_steps=2000 \ + save_total_limit=4 \ + training_args.learning_rate=1e-5 \ + training_args.warmup_ratio=0.05 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero3_multinode.json" \ + training_args.bf16=true \ + training_args.tf32=true \ + training_args.eval_bf16=true \ + output_dir="$OUTPUT_DIR" \ + model.gradient_checkpointing=true \ + dataloader_num_workers=8 \ + dataloader_pin_memory=true \ + optim=adamw_bnb_8bit \ + bf16=true tf32=true eval_bf16=true \ + "$@" diff --git a/scripts/cluster/torchrun_multinode.sh b/scripts/cluster/torchrun_multinode.sh new file mode 100644 index 0000000000000000000000000000000000000000..daa4de783de901b795abc392f67a7e42ac5712ec --- /dev/null +++ b/scripts/cluster/torchrun_multinode.sh @@ -0,0 +1,79 @@ +#!/bin/bash +# DreamZero 多节点 torchrun 启动脚本（不依赖 SLURM） +# +# 用法: +# bash scripts/cluster/torchrun_multinode.sh [额外参数] +# +# 示例: +# bash scripts/cluster/torchrun_multinode.sh "node1,node2,node3,node4" libero +# bash scripts/cluster/torchrun_multinode.sh "10.0.0.1,10.0.0.2" manifeel --max_steps=100000 +# +# 前置条件: +# - 所有节点可互相 SSH 免密访问 +# - 代码和数据在所有节点上路径一致 +# - 所有节点已安装所需依赖 + +set -euo pipefail + +if [ $# -lt 2 ]; then + echo "用法: $0 [额外参数]" + echo "示例: $0 node1,node2,node3,node4 libero" + exit 1 +fi + +NODE_LIST="$1" +BENCHMARK="$2" +shift 2 + +# ============ 解析节点列表 ============ +IFS=',' read -ra NODES <<< "$NODE_LIST" +NNODES=${#NODES[@]} +FIRST_NODE="${NODES[0]}" +MASTER_PORT=${MASTER_PORT:-29500} + +echo "============================================" +echo " DreamZero Multi-Node Training" +echo "============================================" +echo " Nodes: $NNODES ($NODE_LIST)" +echo " Master: $FIRST_NODE:$MASTER_PORT" +echo " Benchmark: $BENCHMARK" +echo "============================================" + +# ============ 路径 ============ +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/output/${BENCHMARK}_full_${NNODES}nodes}" + +# ============ 环境变量 ============ +export NCCL_DEBUG=WARN +export NCCL_IB_DISABLE=0 +export NCCL_IB_TIMEOUT=22 +export NCCL_SOCKET_IFNAME=^docker0,lo +export HF_ENDPOINT=https://hf-mirror.com + +# ============ 运行 ============ +torchrun \ + --nnodes="$NNODES" \ + --nproc_per_node=8 \ + --rdzv_id="dz_${BENCHMARK}" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="$FIRST_NODE:$MASTER_PORT" \ + "$REPO_ROOT/groot/vla/experiment/experiment.py" \ + report_to=wandb \ + data="dreamzero/${BENCHMARK}" \ + wandb_project=dreamzero-sft \ + train_architecture=full \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ + model/dreamzero/transform=dreamzero_cotrain \ + per_device_train_batch_size=1 \ + save_steps=2000 \ + save_total_limit=4 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero3_multinode.json" \ + training_args.bf16=true \ + training_args.tf32=true \ + output_dir="$OUTPUT_DIR" \ + dataloader_num_workers=8 \ + optim=adamw_bnb_8bit \ + "$@" diff --git a/scripts/compare_loss.py b/scripts/compare_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..7b79ee1f8a9b900d34bd80ec3e8da428cc8b91d4 --- /dev/null +++ b/scripts/compare_loss.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Compare loss curves between LoRA and full fine-tuning runs. + +Usage: + python scripts/compare_loss.py \ + --lora-log ./checkpoints/dreamzero_droid_lora/loss_log.jsonl \ + --full-log ./checkpoints/dreamzero_droid_full_finetune/loss_log.jsonl \ + [--plot loss_comparison.png] +""" + +import argparse +import json + + +def load_loss_log(path): + entries = [] + with open(path) as f: + for line in f: + line = line.strip() + if line: + entries.append(json.loads(line)) + return entries + + +def print_comparison_table(lora_entries, full_entries): + # Index by step + lora_by_step = {e["step"]: e for e in lora_entries} + full_by_step = {e["step"]: e for e in full_entries} + all_steps = sorted(set(lora_by_step.keys()) | set(full_by_step.keys())) + + header = f"{'Step':>6} {'LoRA Loss':>10} {'Full Loss':>10} {'LoRA Dyn':>10} {'Full Dyn':>10} {'LoRA Act':>10} {'Full Act':>10}" + print(header) + print("-" * len(header)) + + for step in all_steps: + lora = lora_by_step.get(step, {}) + full = full_by_step.get(step, {}) + + def fmt(d, key): + v = d.get(key) + return f"{v:10.4f}" if v is not None else f"{'—':>10}" + + print( + f"{step:>6} " + f"{fmt(lora, 'loss')} {fmt(full, 'loss')} " + f"{fmt(lora, 'dynamics_loss_avg')} {fmt(full, 'dynamics_loss_avg')} " + f"{fmt(lora, 'action_loss_avg')} {fmt(full, 'action_loss_avg')}" + ) + + +def plot_comparison(lora_entries, full_entries, output_path): + try: + import matplotlib.pyplot as plt + except ImportError: + print("matplotlib not installed, skipping plot generation.") + print("Install with: pip install matplotlib") + return + + metrics = [ + ("loss", "Total Loss"), + ("dynamics_loss_avg", "Dynamics Loss"), + ("action_loss_avg", "Action Loss"), + ] + + fig, axes = plt.subplots(1, len(metrics), figsize=(5 * len(metrics), 4)) + if len(metrics) == 1: + axes = [axes] + + for ax, (key, title) in zip(axes, metrics): + lora_steps = [e["step"] for e in lora_entries if key in e] + lora_vals = [e[key] for e in lora_entries if key in e] + full_steps = [e["step"] for e in full_entries if key in e] + full_vals = [e[key] for e in full_entries if key in e] + + if lora_steps: + ax.plot(lora_steps, lora_vals, label="LoRA", marker="o", markersize=3) + if full_steps: + ax.plot(full_steps, full_vals, label="Full FT", marker="s", markersize=3) + + ax.set_title(title) + ax.set_xlabel("Step") + ax.set_ylabel("Loss") + ax.legend() + ax.grid(True, alpha=0.3) + + fig.tight_layout() + fig.savefig(output_path, dpi=150) + print(f"Plot saved to {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="Compare LoRA vs full fine-tuning loss curves") + parser.add_argument("--lora-log", required=True, help="Path to LoRA run loss_log.jsonl") + parser.add_argument("--full-log", required=True, help="Path to full FT run loss_log.jsonl") + parser.add_argument("--plot", default=None, help="Output path for comparison plot (e.g., loss_comparison.png)") + args = parser.parse_args() + + lora_entries = load_loss_log(args.lora_log) + full_entries = load_loss_log(args.full_log) + + print(f"LoRA: {len(lora_entries)} log entries") + print(f"Full: {len(full_entries)} log entries") + print() + + print_comparison_table(lora_entries, full_entries) + + if args.plot: + plot_comparison(lora_entries, full_entries, args.plot) + + +if __name__ == "__main__": + main() diff --git a/scripts/data/convert_agibot.py b/scripts/data/convert_agibot.py new file mode 100644 index 0000000000000000000000000000000000000000..8d888cc17e402ff7045060f200438bc3536fd7a0 --- /dev/null +++ b/scripts/data/convert_agibot.py @@ -0,0 +1,1271 @@ +import argparse +from copy import deepcopy +from functools import partial +import gc +import json +import logging +from math import ceil +import os +from pathlib import Path +import shutil +from typing import Callable + +import einops +import h5py +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset +from lerobot.common.datasets.utils import ( + STATS_PATH, + check_timestamps_sync, + get_episode_data_index, + serialize_dict, + write_json, +) +import numpy as np +import torch +from tqdm import tqdm +from tqdm.contrib.concurrent import process_map + + +def generate_modality_json(output_dir: str) -> None: + """Generate modality.json file defining field mappings for the dataset.""" + modality_config = { + "state": { + "left_arm_joint_position": { + "original_key": "observation.state", + "start": 0, + "end": 7, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "right_arm_joint_position": { + "original_key": "observation.state", + "start": 7, + "end": 14, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "left_effector_position": { + "original_key": "observation.state", + "start": 14, + "end": 15, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "right_effector_position": { + "original_key": "observation.state", + "start": 15, + "end": 16, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "head_position": { + "original_key": "observation.state", + "start": 16, + "end": 18, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "waist_pitch": { + "original_key": "observation.state", + "start": 18, + "end": 19, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "waist_lift": { + "original_key": "observation.state", + "start": 19, + "end": 20, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + }, + "action": { + "left_arm_joint_position": { + "original_key": "action", + "start": 0, + "end": 7, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "right_arm_joint_position": { + "original_key": "action", + "start": 7, + "end": 14, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "left_effector_position": { + "original_key": "action", + "start": 14, + "end": 15, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "right_effector_position": { + "original_key": "action", + "start": 15, + "end": 16, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "head_position": { + "original_key": "action", + "start": 16, + "end": 18, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "waist_pitch": { + "original_key": "action", + "start": 18, + "end": 19, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "waist_lift": { + "original_key": "action", + "start": 19, + "end": 20, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "robot_velocity": { + "original_key": "action", + "start": 20, + "end": 22, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + }, + "video": { + "top_head": { + "original_key": "observation.images.top_head" + }, + "hand_left": { + "original_key": "observation.images.hand_left" + }, + "hand_right": { + "original_key": "observation.images.hand_right" + }, + }, + "annotation": { + "language.action_text": { + "original_key": "task_index" + }, + "agibot.sub_task": { + "original_key": "annotation.agibot.sub_task" + }, + "frame_type": { + "original_key": "annotation.frame_type" + }, + }, + } + + modality_path = os.path.join(output_dir, "modality.json") + with open(modality_path, "w") as f: + json.dump(modality_config, f, indent=4) + print(f"Generated modality.json at {modality_path}") + + +HEAD_COLOR = "head_color.mp4" +HAND_LEFT_COLOR = "hand_left_color.mp4" +HAND_RIGHT_COLOR = "hand_right_color.mp4" + +FEATURES = { + "observation.images.top_head": { + "dtype": "video", + "shape": [480, 640, 3], + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": 30.0, + "video.codec": "av1", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + }, + "observation.images.hand_left": { + "dtype": "video", + "shape": [480, 640, 3], + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": 30.0, + "video.codec": "av1", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + }, + "observation.images.hand_right": { + "dtype": "video", + "shape": [480, 640, 3], + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": 30.0, + "video.codec": "av1", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + }, + "observation.state": { + "dtype": "float32", + "shape": [20], + }, + "action": { + "dtype": "float32", + "shape": [22], + }, + "annotation.language.action_text": { + "dtype": "int64", + "shape": [1], + "names": None, + }, + "annotation.agibot.tasks": { + "dtype": "int64", + "shape": [1], + "names": None, + }, + "episode_index": { + "dtype": "int64", + "shape": [1], + "names": None, + }, + "frame_index": { + "dtype": "int64", + "shape": [1], + "names": None, + }, + "index": { + "dtype": "int64", + "shape": [1], + "names": None, + }, + "task_index": { + "dtype": "int64", + "shape": [1], + "names": None, + }, +} + + +def get_stats_einops_patterns(dataset, num_workers=0): + """These einops patterns will be used to aggregate batches and compute statistics. + + Note: We assume the images are in channel first format + """ + + dataloader = torch.utils.data.DataLoader( + dataset, + num_workers=num_workers, + batch_size=2, + shuffle=False, + ) + batch = next(iter(dataloader)) + + stats_patterns = {} + + for key in dataset.features: + # sanity check that tensors are not float64 + assert batch[key].dtype != torch.float64 + + # if isinstance(feats_type, (VideoFrame, Image)): + if key in dataset.meta.camera_keys: + # sanity check that images are channel first + _, c, h, w = batch[key].shape + assert c < h and c < w, f"expect channel first images, but instead {batch[key].shape}" + assert ( + batch[key].dtype == torch.float32 + ), f"expect torch.float32, but instead {batch[key].dtype=}" + # assert batch[key].max() <= 1, f"expect pixels lower than 1, but instead {batch[key].max()=}" + # assert batch[key].min() >= 0, f"expect pixels greater than 1, but instead {batch[key].min()=}" + stats_patterns[key] = "b c h w -> c 1 1" + elif batch[key].ndim == 2: + stats_patterns[key] = "b c -> c " + elif batch[key].ndim == 1: + stats_patterns[key] = "b -> 1" + else: + raise ValueError(f"{key}, {batch[key].shape}") + + return stats_patterns + + +def compute_stats(dataset, batch_size=8, num_workers=4, max_num_samples=None): + """Compute mean/std and min/max statistics of all data keys in a LeRobotDataset.""" + if max_num_samples is None: + max_num_samples = len(dataset) + + # for more info on why we need to set the same number of workers, see `load_from_videos` + stats_patterns = get_stats_einops_patterns(dataset, num_workers) + + # mean and std will be computed incrementally while max and min will track the running value. + mean, std, max, min = {}, {}, {}, {} + for key in stats_patterns: + mean[key] = torch.tensor(0.0).float() + std[key] = torch.tensor(0.0).float() + max[key] = torch.tensor(-float("inf")).float() + min[key] = torch.tensor(float("inf")).float() + + def create_seeded_dataloader(dataset, batch_size, seed): + generator = torch.Generator() + generator.manual_seed(seed) + dataloader = torch.utils.data.DataLoader( + dataset, + num_workers=num_workers, + batch_size=batch_size, + shuffle=True, + drop_last=False, + generator=generator, + ) + return dataloader + + # Note: Due to be refactored soon. The point of storing `first_batch` is to make sure we don't get + # surprises when rerunning the sampler. + first_batch = None + running_item_count = 0 # for online mean computation + dataloader = create_seeded_dataloader(dataset, batch_size, seed=1337) + for i, batch in enumerate( + tqdm( + dataloader, + total=ceil(max_num_samples / batch_size), + desc="Compute mean, min, max", + ) + ): + this_batch_size = len(batch["index"]) + running_item_count += this_batch_size + if first_batch is None: + first_batch = deepcopy(batch) + for key, pattern in stats_patterns.items(): + batch[key] = batch[key].float() + # Numerically stable update step for mean computation. + batch_mean = einops.reduce(batch[key], pattern, "mean") + # Hint: to update the mean we need x̄ₙ = (Nₙ₋₁x̄ₙ₋₁ + Bₙxₙ) / Nₙ, where the subscript represents + # the update step, N is the running item count, B is this batch size, x̄ is the running mean, + # and x is the current batch mean. Some rearrangement is then required to avoid risking + # numerical overflow. Another hint: Nₙ₋₁ = Nₙ - Bₙ. Rearrangement yields + # x̄ₙ = x̄ₙ₋₁ + Bₙ * (xₙ - x̄ₙ₋₁) / Nₙ + mean[key] = mean[key] + this_batch_size * (batch_mean - mean[key]) / running_item_count + max[key] = torch.maximum(max[key], einops.reduce(batch[key], pattern, "max")) + min[key] = torch.minimum(min[key], einops.reduce(batch[key], pattern, "min")) + + if i == ceil(max_num_samples / batch_size) - 1: + break + + first_batch_ = None + running_item_count = 0 # for online std computation + dataloader = create_seeded_dataloader(dataset, batch_size, seed=1337) + for i, batch in enumerate( + tqdm(dataloader, total=ceil(max_num_samples / batch_size), desc="Compute std") + ): + this_batch_size = len(batch["index"]) + running_item_count += this_batch_size + # Sanity check to make sure the batches are still in the same order as before. + if first_batch_ is None: + first_batch_ = deepcopy(batch) + for key in stats_patterns: + assert torch.equal(first_batch_[key], first_batch[key]) + for key, pattern in stats_patterns.items(): + batch[key] = batch[key].float() + # Numerically stable update step for mean computation (where the mean is over squared + # residuals).See notes in the mean computation loop above. + batch_std = einops.reduce((batch[key] - mean[key]) ** 2, pattern, "mean") + std[key] = std[key] + this_batch_size * (batch_std - std[key]) / running_item_count + + if i == ceil(max_num_samples / batch_size) - 1: + break + + for key in stats_patterns: + std[key] = torch.sqrt(std[key]) + + stats = {} + for key in stats_patterns: + stats[key] = { + "mean": mean[key], + "std": std[key], + "max": max[key], + "min": min[key], + } + return stats + + +class AgiBotDataset(LeRobotDataset): + def __init__( + self, + repo_id: str, + root: str | Path | None = None, + episodes: list[int] | None = None, + image_transforms: Callable | None = None, + delta_timestamps: dict[list[float]] | None = None, + tolerance_s: float = 1e-4, + download_videos: bool = True, + local_files_only: bool = False, + video_backend: str | None = None, + ): + super().__init__( + repo_id=repo_id, + root=root, + episodes=episodes, + image_transforms=image_transforms, + delta_timestamps=delta_timestamps, + tolerance_s=tolerance_s, + download_videos=download_videos, + local_files_only=local_files_only, + video_backend=video_backend, + ) + + def save_episode( + self, task: str, episode_data: dict | None = None, videos: dict | None = None + ) -> None: + """ + We rewrite this method to copy mp4 videos to the target position + """ + if not episode_data: + episode_buffer = self.episode_buffer + + episode_length = episode_buffer.pop("size") + episode_index = episode_buffer["episode_index"] + if episode_index != self.meta.total_episodes: + # TODO(aliberts): Add option to use existing episode_index + raise NotImplementedError( + "You might have manually provided the episode_buffer with an episode_index that doesn't " + "match the total number of episodes in the dataset. This is not supported for now." + ) + + if episode_length == 0: + raise ValueError( + "You must add one or several frames with `add_frame` before calling `add_episode`." + ) + + # Use our custom task indexing instead of LeRobot's built-in mechanism + task_index = getattr(self, "_custom_task_to_index", {}).get(task, 0) + + # Remove the 'task' key if it exists (it's passed as a parameter, not needed in buffer) + episode_buffer.pop("task", None) + + if not set(episode_buffer.keys()) == set(self.features): + raise ValueError() + + for key, ft in self.features.items(): + if key == "index": + episode_buffer[key] = np.arange( + self.meta.total_frames, self.meta.total_frames + episode_length + ) + elif key == "episode_index": + episode_buffer[key] = np.full((episode_length,), episode_index) + elif key == "task_index": + episode_buffer[key] = np.full((episode_length,), task_index) + elif ft["dtype"] in ["image", "video"]: + continue + elif ft["dtype"] == "string": + pass + elif len(ft["shape"]) == 1 and ft["shape"][0] == 1: + episode_buffer[key] = np.array(episode_buffer[key], dtype=ft["dtype"]) + elif len(ft["shape"]) == 1 and ft["shape"][0] > 1: + episode_buffer[key] = np.stack(episode_buffer[key]) + else: + raise ValueError(key) + + self._wait_image_writer() + self._save_episode_table(episode_buffer, episode_index) + + # Copy videos first before calling meta.save_episode which might try to read them + for key in self.meta.video_keys: + video_path = self.root / self.meta.get_video_file_path(episode_index, key) + episode_buffer[key] = video_path + video_path.parent.mkdir(parents=True, exist_ok=True) + # Copy video files to target location + shutil.copyfile(str(videos[key]), str(video_path)) + + try: + # FIX: Call meta.save_episode with correct parameters + # Note: We pass an empty task list to prevent duplicate entries in tasks.jsonl + # since we create our own tasks.jsonl file with proper indexing + self.meta.save_episode(episode_index, episode_length, [], {}) + except AttributeError as e: + if "'NoneType' object has no attribute 'items'" in str(e): + # Handle the episode stats issue - skip episode stats for now + print( + f"Warning: Episode stats computation failed, proceeding without stats " + f"for episode {episode_index}" + ) + # Just skip the problematic save_episode call - the core data is already saved + pass + else: + raise + if not episode_data: # Reset the buffer + self.episode_buffer = self.create_episode_buffer() + self.consolidated = False + + def consolidate(self, run_compute_stats: bool = True, keep_image_files: bool = False) -> None: + self.hf_dataset = self.load_hf_dataset() + self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes) + check_timestamps_sync(self.hf_dataset, self.episode_data_index, self.fps, self.tolerance_s) + if len(self.meta.video_keys) > 0: + self.meta.write_video_info() + + if not keep_image_files: + img_dir = self.root / "images" + if img_dir.is_dir(): + shutil.rmtree(self.root / "images") + video_files = list(self.root.rglob("*.mp4")) + assert len(video_files) == self.num_episodes * len(self.meta.video_keys) + + parquet_files = list(self.root.rglob("*.parquet")) + assert len(parquet_files) == self.num_episodes + + if run_compute_stats: + self.stop_image_writer() + self.meta.stats = compute_stats(self) + serialized_stats = serialize_dict(self.meta.stats) + write_json(serialized_stats, self.root / STATS_PATH) + self.consolidated = True + else: + logging.warning( + "Skipping computation of the dataset statistics, dataset is not fully consolidated." + ) + + def add_frame(self, frame: dict) -> None: + """ + This function only adds the frame to the episode_buffer. Apart from images — which are written in a + temporary directory — nothing is written to disk. To save those frames, the 'save_episode()' method + then needs to be called. + """ + # TODO(aliberts, rcadene): Add sanity check for the input, check it's numpy or torch, + # check the dtype and shape matches, etc. + + if self.episode_buffer is None: + self.episode_buffer = self.create_episode_buffer() + + frame_index = self.episode_buffer["size"] + timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps + self.episode_buffer["frame_index"].append(frame_index) + self.episode_buffer["timestamp"].append(timestamp) + + for key in frame: + if key not in self.features: + raise ValueError(key) + item = frame[key].numpy() if isinstance(frame[key], torch.Tensor) else frame[key] + self.episode_buffer[key].append(item) + + self.episode_buffer["size"] += 1 + + +def detect_dataset_format(src_path: str) -> str: + """Detect whether the dataset follows old or new format structure""" + src_path = Path(src_path) + + # Check for old format indicators + if (src_path / "task_info").exists() and (src_path / "proprio_stats").exists(): + return "old" + + # Check for new format indicators + # Look for pattern: job_id/robot_id/episode_id with aligned_joints.h5 + # The test_data/2810125 directory contains job_id directories + subdirs = [d for d in src_path.iterdir() if d.is_dir()] + if subdirs: + # Check if we have nested structure with aligned_joints.h5 + for job_dir in subdirs: # These are job_id directories like 3335477 + if not job_dir.is_dir(): + continue + for ( + robot_dir + ) in job_dir.iterdir(): # These are robot_id directories like A2D0015AB00061 + if not robot_dir.is_dir(): + continue + for ( + episode_dir + ) in robot_dir.iterdir(): # These are episode_id directories like 12052353 + if episode_dir.is_dir() and (episode_dir / "aligned_joints.h5").exists(): + return "new" + + return "unknown" + + +def load_local_dataset_old_format(episode_id: int, src_path: str, task_id: int) -> list | None: + """Load local dataset from old format and return a dict with observations and actions""" + + # --- Load task info for this specific call --- + task_json_path = Path(src_path) / f"task_info/task_{task_id}.json" + task_info_list = None + if task_json_path.exists(): + try: + with open(task_json_path, "r") as f: + task_info_list = json.load(f) + except json.JSONDecodeError: + print(f"Warning: Failed to decode JSON {task_json_path} for episode {episode_id}") + task_info_list = [] # Treat as empty if decode fails + else: + print(f"Warning: Task info JSON not found at {task_json_path} for episode {episode_id}") + task_info_list = [] # Treat as empty if not found + + # --- Find action_config for this episode_id --- + episode_action_config = None + if isinstance(task_info_list, list): # Check if loading succeeded and it's a list + for item in task_info_list: + # Ensure episode_id exists and compare as int + if "episode_id" in item and int(item["episode_id"]) == episode_id: + episode_action_config = item.get("label_info", {}).get("action_config") + break + + default_action_text = "N/A" + + ob_dir = Path(src_path) / f"observations/{task_id}/{episode_id}" + proprio_dir = Path(src_path) / f"proprio_stats/{task_id}/{episode_id}" + + with h5py.File(proprio_dir / "proprio_stats.h5") as f: + state_joint = np.array(f["state/joint/position"]) + state_effector = np.clip((np.array(f["state/effector/position"]) - 35.0) / (120.0 - 35.0), 0.0, 1.0) + state_head = np.array(f["state/head/position"]) + state_waist = np.array(f["state/waist/position"]) + action_joint = np.array(f["action/joint/position"]) + action_effector = np.clip((np.array(f["action/effector/position"]) - 35.0) / (120.0 - 35.0), 0.0, 1.0) + action_head = np.array(f["action/head/position"]) + action_waist = np.array(f["action/waist/position"]) + action_velocity = np.array(f["action/robot/velocity"]) + + # State (20 DOF): joint(14) + effector(2) + head(2) + waist(2) + states_value = np.hstack( + [ + state_joint, + state_effector, + state_head, + state_waist, + ] + ).astype(np.float32) + assert ( + action_joint.shape[0] == action_effector.shape[0] + ), f"shape of action_joint:{action_joint.shape};shape of action_effector:{action_effector.shape}" + # Action (22 DOF): joint(14) + effector(2) + head(2) + waist(2) + velocity(2) + action_value = np.hstack( + [ + action_joint, + action_effector, + action_head, + action_waist, + action_velocity, + ] + ).astype(np.float32) + + num_frames = len(states_value) + + # --- Create frame -> action_text mapping --- + frame_action_texts = [default_action_text] * num_frames + if episode_action_config: # Only proceed if config was found + for action in episode_action_config: + start = action.get("start_frame") + end = action.get("end_frame") + text = action.get("action_text", default_action_text) + + if start is None or end is None: + continue + + clamped_start = max(0, start) + clamped_end = min(num_frames, end) + for i in range(clamped_start, clamped_end): + frame_action_texts[i] = text + + frames = [ + { + "observation.state": states_value[i], + "action": action_value[i], + "annotation.language.action_text": [frame_action_texts[i]], # Add action_text here + } + for i in range(num_frames) + ] + + v_path = ob_dir / "videos" + videos = { + "observation.images.top_head": v_path / HEAD_COLOR, + "observation.images.hand_left": v_path / HAND_LEFT_COLOR, + "observation.images.hand_right": v_path / HAND_RIGHT_COLOR, + } + return frames, videos + + +def load_local_dataset_new_format(episode_path: str) -> list | None: + """Load local dataset from new format and return a dict with observations and actions""" + + episode_dir = Path(episode_path) + + # Load data info JSON + data_info_path = episode_dir / "data_info.json" + episode_action_config = None + default_action_text = "N/A" + + if data_info_path.exists(): + try: + with open(data_info_path, "r") as f: + data_info = json.load(f) + episode_action_config = data_info.get("label_info", {}).get("action_config") + except json.JSONDecodeError: + print(f"Warning: Failed to decode JSON {data_info_path}") + + # Load aligned joints data + joints_path = episode_dir / "aligned_joints.h5" + if not joints_path.exists(): + print(f"Warning: aligned_joints.h5 not found at {joints_path}") + return None + + with h5py.File(joints_path) as f: + # Extract state data - using same mapping as old format for compatibility + state_joint = np.array(f["state/joint/position"]) + state_head = np.array(f["state/head/position"]) + state_waist = np.array(f["state/waist/position"]) + + # For new format, use separate left/right effector position data + # Normalize from raw [35, 120] to [0, 1] + state_left_effector = np.array(f["state/left_effector/position"]) # Shape: (N, 1) + state_right_effector = np.array(f["state/right_effector/position"]) # Shape: (N, 1) + state_effector = np.clip( + np.column_stack([state_left_effector.flatten(), state_right_effector.flatten()]) - 35.0, + 0.0, 85.0 + ) / 85.0 # Shape: (N, 2), range [0, 1] + + # Extract action data + action_joint = np.array(f["action/joint/position"]) + action_head = np.array(f["action/head/position"]) + action_waist = np.array(f["action/waist/position"]) + + # For new format, use separate left/right effector position data for actions + # Normalize from raw [35, 120] to [0, 1] + action_left_effector = np.array(f["action/left_effector/position"]) # Shape: (N, 1) + action_right_effector = np.array(f["action/right_effector/position"]) # Shape: (N, 1) + action_effector = np.clip( + np.column_stack([action_left_effector.flatten(), action_right_effector.flatten()]) - 35.0, + 0.0, 85.0 + ) / 85.0 # Shape: (N, 2), range [0, 1] + + # Get robot velocity (N,) or (N, 2) + action_velocity_raw = np.array(f["action/robot/velocity"]) + if action_velocity_raw.ndim == 1: + # Scalar velocity: pad second component with zeros (linear_x only) + action_velocity = np.column_stack( + [action_velocity_raw, np.zeros_like(action_velocity_raw)] + ) + else: + action_velocity = action_velocity_raw[:, :2] + + # State (20 DOF): joint(14) + effector(2) + head(2) + waist(2) + states_value = np.hstack( + [ + state_joint, + state_effector, + state_head, + state_waist, + ] + ).astype(np.float32) + # Action (22 DOF): joint(14) + effector(2) + head(2) + waist(2) + velocity(2) + action_value = np.hstack( + [ + action_joint, + action_effector, + action_head, + action_waist, + action_velocity, + ] + ).astype(np.float32) + + num_frames = len(states_value) + + # --- Create frame -> action_text mapping --- + frame_action_texts = [default_action_text] * num_frames + if episode_action_config: # Only proceed if config was found + for action in episode_action_config: + start = action.get("start_frame") + end = action.get("end_frame") + # Use English action text if available, otherwise use Chinese + text = action.get("english_action_text") or action.get( + "action_text", default_action_text + ) + + if start is None or end is None: + continue + + clamped_start = max(0, start) + clamped_end = min(num_frames, end) + for i in range(clamped_start, clamped_end): + frame_action_texts[i] = text + + frames = [ + { + "observation.state": states_value[i], + "action": action_value[i], + "annotation.language.action_text": [frame_action_texts[i]], + } + for i in range(num_frames) + ] + + # Videos are at episode level in new format + videos = { + "observation.images.top_head": episode_dir / HEAD_COLOR, + "observation.images.hand_left": episode_dir / HAND_LEFT_COLOR, + "observation.images.hand_right": episode_dir / HAND_RIGHT_COLOR, + } + return frames, videos + + +def load_local_dataset( + episode_id: int, + src_path: str, + task_id: int = None, + episode_path: str = None, + format_type: str = "old", +) -> list | None: + """Load local dataset and return a dict with observations and actions + + Args: + episode_id: Episode ID (used for old format) + src_path: Source path (used for old format) + task_id: Task ID (used for old format) + episode_path: Full path to episode directory (used for new format) + format_type: "old" or "new" format + """ + if format_type == "old": + return load_local_dataset_old_format(episode_id, src_path, task_id) + elif format_type == "new": + return load_local_dataset_new_format(episode_path) + else: + raise ValueError(f"Unknown format type: {format_type}") + + +def get_task_instruction_old_format(task_json_path: str) -> str: + """Get task language instruction from old format""" + with open(task_json_path, "r") as f: + task_info = json.load(f) + task_name = task_info[0]["task_name"] + task_init_scene = task_info[0]["init_scene_text"] + task_instruction = f"{task_name}.{task_init_scene}" + print(f"Get Task Instruction <{task_instruction}>") + return task_instruction + + +def get_task_instruction_new_format(episode_paths: list) -> str: + """Get task language instruction from new format - use first episode's data_info.json""" + if not episode_paths: + return "Unknown Task" + + first_episode_path = Path(episode_paths[0]) + data_info_path = first_episode_path / "data_info.json" + + if data_info_path.exists(): + try: + with open(data_info_path, "r") as f: + data_info = json.load(f) + # Use English task name if available, otherwise use Chinese + task_name = data_info.get("english_task_name") or data_info.get( + "task_name", "Unknown Task" + ) + task_instruction = task_name + print( + f"Get Task Instruction <{task_instruction}> " + f"(english_task_name: {data_info.get('english_task_name')}, " + f"task_name: {data_info.get('task_name')})" + ) + return task_instruction + except json.JSONDecodeError: + print(f"Warning: Failed to decode JSON {data_info_path}") + + return "Unknown Task" + + +def get_task_instruction( + task_json_path: str = None, episode_paths: list = None, format_type: str = "old" +) -> str: + """Get task language instruction""" + if format_type == "old": + return get_task_instruction_old_format(task_json_path) + elif format_type == "new": + return get_task_instruction_new_format(episode_paths) + else: + raise ValueError(f"Unknown format type: {format_type}") + + +def load_new_format_episode(episode_path): + """Helper function for multiprocessing - load new format episode""" + return load_local_dataset( + episode_id=0, src_path="", episode_path=episode_path, format_type="new" + ) + + +def create_tasks_jsonl(tgt_path: str, repo_id: str, task_name: str, all_action_texts: set) -> dict: + """Create tasks.jsonl file with unique task names and action texts.""" + meta_path = os.path.join(tgt_path, repo_id, "meta") + os.makedirs(meta_path, exist_ok=True) + + tasks_jsonl_path = os.path.join(meta_path, "tasks.jsonl") + + # Create a list of unique tasks combining task name and action texts + # Remove task_name from action texts if it exists to avoid duplicates + unique_action_texts = all_action_texts - {task_name} + + tasks = [task_name] # First entry is the main task + tasks.extend(sorted(unique_action_texts)) # Then all unique action texts (excluding task name) + + # Check if file already exists and has correct content + should_write = True + if os.path.exists(tasks_jsonl_path): + try: + with open(tasks_jsonl_path, "r") as f: + existing_content = f.read().strip() + + # Generate expected content + expected_lines = [] + for i, task in enumerate(tasks): + task_entry = {"task_index": i, "task": task} + expected_lines.append(json.dumps(task_entry)) + expected_content = "\n".join(expected_lines) + + # If content matches, don't rewrite + if existing_content == expected_content: + should_write = False + except Exception as e: + print(f"Warning: Failed to read tasks.jsonl: {e}") + # If there's any issue reading, we'll rewrite + pass + + if should_write: + # Write tasks.jsonl (overwrite to ensure clean content) + with open(tasks_jsonl_path, "w") as f: + for i, task in enumerate(tasks): + task_entry = {"task_index": i, "task": task} + f.write(json.dumps(task_entry) + "\n") + + print(f"Created tasks.jsonl with {len(tasks)} entries at {tasks_jsonl_path}") + else: + print(f"tasks.jsonl already exists with correct content at {tasks_jsonl_path}") + + # Create mapping for lookups + task_to_index = {task: i for i, task in enumerate(tasks)} + + return task_to_index + + +def main( + src_path: str, + tgt_path: str, + task_id: int = None, + repo_id: str = None, + task_info_json: str = None, + debug: bool = False, + chunk_size: int = 10, +): + # Detect dataset format + format_type = detect_dataset_format(src_path) + print(f"Detected dataset format: {format_type}") + + if format_type == "unknown": + raise ValueError(f"Unable to detect dataset format for path: {src_path}") + + # Collect all unique action texts first + all_action_texts = set() + + # Initialize dataset + if not repo_id: + if format_type == "old": + repo_id = f"agibotworld/task_{task_id}" + else: # new format + # Use the top-level directory name as task_id + task_id = Path(src_path).name + repo_id = f"agibotworld/task_{task_id}" + + dataset = AgiBotDataset.create( + repo_id=repo_id, + root=f"{tgt_path}/{repo_id}", + fps=30, + robot_type="a2d", + features=FEATURES, + ) + + if format_type == "old": + # Old format processing + task_name = get_task_instruction(task_json_path=task_info_json, format_type="old") + + all_subdir = sorted( + [f.as_posix() for f in Path(src_path).glob(f"observations/{task_id}/*") if f.is_dir()] + ) + + if debug: + all_subdir = all_subdir[:2] + + # Get all episode id + all_subdir_eids = [int(Path(path).name) for path in all_subdir] + all_subdir_episode_desc = [task_name] * len(all_subdir_eids) + + # First pass: collect all unique action texts + print("Collecting unique action texts...") + for episode_id in tqdm(all_subdir_eids, desc="Scanning for action texts"): + frames_data, _ = load_local_dataset( + episode_id, src_path=src_path, task_id=task_id, format_type="old" + ) + if frames_data: + for frame in frames_data: + action_text = frame["annotation.language.action_text"][0] + all_action_texts.add(action_text) + + # Create tasks.jsonl with all unique texts + task_to_index = create_tasks_jsonl(tgt_path, repo_id, task_name, all_action_texts) + + # Pass the task mapping to the dataset for custom indexing + dataset._custom_task_to_index = task_to_index + + # Process in chunks to reduce memory usage + for chunk_start in tqdm( + range(0, len(all_subdir_eids), chunk_size), desc="Processing chunks" + ): + chunk_end = min(chunk_start + chunk_size, len(all_subdir_eids)) + chunk_eids = all_subdir_eids[chunk_start:chunk_end] + chunk_descs = all_subdir_episode_desc[chunk_start:chunk_end] + + # Process only this chunk + if debug: + raw_datasets_chunk = [ + load_local_dataset( + subdir, src_path=src_path, task_id=task_id, format_type="old" + ) + for subdir in tqdm(chunk_eids, desc="Loading chunk data") + ] + else: + raw_datasets_chunk = process_map( + partial( + load_local_dataset, src_path=src_path, task_id=task_id, format_type="old" + ), + chunk_eids, + max_workers=os.cpu_count() // 2, + desc=f"Loading chunk {chunk_start//chunk_size + 1}/" + f"{(len(all_subdir_eids) + chunk_size - 1)//chunk_size}", + ) + + # Filter out None results + valid_datasets = [ + (ds, desc) for ds, desc in zip(raw_datasets_chunk, chunk_descs) if ds is not None + ] + + # Process each dataset in the chunk + for raw_dataset, episode_desc in tqdm( + valid_datasets, desc="Processing episodes in chunk" + ): + for raw_dataset_sub in tqdm(raw_dataset[0], desc="Processing frames", leave=False): + # Convert string annotation to int index + action_text = raw_dataset_sub["annotation.language.action_text"][0] + raw_dataset_sub["annotation.language.action_text"] = [ + task_to_index[action_text] + ] + raw_dataset_sub["annotation.agibot.tasks"] = [task_to_index[episode_desc]] + dataset.add_frame(raw_dataset_sub) + dataset.save_episode(task=episode_desc, videos=raw_dataset[1]) + + # Clear memory after each chunk + raw_datasets_chunk = None + valid_datasets = None + gc.collect() + + else: # new format + # Find all episode directories + all_episode_paths = [] + src_path = Path(src_path) + + # Walk through job_id/robot_id/episode_id structure (test_data/2810125 contains job_id directories) + for job_dir in src_path.iterdir(): # These are job_id directories like 3335477 + if not job_dir.is_dir(): + continue + for ( + robot_dir + ) in job_dir.iterdir(): # These are robot_id directories like A2D0015AB00061 + if not robot_dir.is_dir(): + continue + for ( + episode_dir + ) in robot_dir.iterdir(): # These are episode_id directories like 12052353 + if episode_dir.is_dir() and (episode_dir / "aligned_joints.h5").exists(): + all_episode_paths.append(str(episode_dir)) + + all_episode_paths = sorted(all_episode_paths) + + if debug: + all_episode_paths = all_episode_paths[:2] + + # Get task name from first episode + task_name = get_task_instruction(episode_paths=all_episode_paths, format_type="new") + all_episode_descs = [task_name] * len(all_episode_paths) + + # First pass: collect all unique action texts + print("Collecting unique action texts...") + for episode_path in tqdm(all_episode_paths, desc="Scanning for action texts"): + frames_data, _ = load_local_dataset( + episode_id=0, src_path="", episode_path=episode_path, format_type="new" + ) + if frames_data: + for frame in frames_data: + action_text = frame["annotation.language.action_text"][0] + all_action_texts.add(action_text) + + # Create tasks.jsonl with all unique texts + task_to_index = create_tasks_jsonl(tgt_path, repo_id, task_name, all_action_texts) + + # Pass the task mapping to the dataset for custom indexing + dataset._custom_task_to_index = task_to_index + + # Process in chunks to reduce memory usage + for chunk_start in tqdm( + range(0, len(all_episode_paths), chunk_size), desc="Processing chunks" + ): + chunk_end = min(chunk_start + chunk_size, len(all_episode_paths)) + chunk_paths = all_episode_paths[chunk_start:chunk_end] + chunk_descs = all_episode_descs[chunk_start:chunk_end] + + # Process only this chunk + if debug: + raw_datasets_chunk = [ + load_local_dataset( + episode_id=0, src_path="", episode_path=episode_path, format_type="new" + ) + for episode_path in tqdm(chunk_paths, desc="Loading chunk data") + ] + else: + raw_datasets_chunk = process_map( + load_new_format_episode, + chunk_paths, + max_workers=os.cpu_count() // 2, + desc=f"Loading chunk {chunk_start//chunk_size + 1}/" + f"{(len(all_episode_paths) + chunk_size - 1)//chunk_size}", + ) + + # Filter out None results + valid_datasets = [ + (ds, desc) for ds, desc in zip(raw_datasets_chunk, chunk_descs) if ds is not None + ] + + # Process each dataset in the chunk + for raw_dataset, episode_desc in tqdm( + valid_datasets, desc="Processing episodes in chunk" + ): + for raw_dataset_sub in tqdm(raw_dataset[0], desc="Processing frames", leave=False): + # Convert string annotation to int index + action_text = raw_dataset_sub["annotation.language.action_text"][0] + raw_dataset_sub["annotation.language.action_text"] = [ + task_to_index[action_text] + ] + raw_dataset_sub["annotation.agibot.tasks"] = [task_to_index[episode_desc]] + dataset.add_frame(raw_dataset_sub) + dataset.save_episode(task=episode_desc, videos=raw_dataset[1]) + + # Clear memory after each chunk + raw_datasets_chunk = None + valid_datasets = None + gc.collect() + + # Only consolidate at the end + try: + dataset.consolidate(run_compute_stats=False) + except Exception as e: + print(f"Warning: Consolidation failed with error: {e}") + print("Dataset conversion completed but may not be fully consolidated.") + + # Generate modality.json file in the dataset meta directory + meta_path = os.path.join(tgt_path, repo_id, "meta") + os.makedirs(meta_path, exist_ok=True) + generate_modality_json(meta_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Convert AgiBot dataset to LeRobot format. Supports both old and new format datasets." + ) + parser.add_argument( + "--src_path", type=str, required=True, help="Path to source dataset directory" + ) + parser.add_argument( + "--task_id", + type=str, + required=False, + help="Task ID (required for old format, optional for new format)", + ) + parser.add_argument( + "--tgt_path", type=str, required=True, help="Path to target output directory" + ) + parser.add_argument( + "--repo_id", + type=str, + required=False, + help="Repository ID for the dataset (auto-generated if not provided)", + ) + parser.add_argument( + "--debug", action="store_true", help="Process only first 2 episodes for debugging" + ) + parser.add_argument( + "--chunk_size", type=int, default=10, help="Number of episodes to process at once" + ) + args = parser.parse_args() + + # Detect format first + format_type = detect_dataset_format(args.src_path) + print(f"Auto-detected format: {format_type}") + + # Validate arguments based on format + if format_type == "old": + if not args.task_id: + parser.error("--task_id is required for old format datasets") + + task_id = int(args.task_id) + json_file = f"{args.src_path}/task_info/task_{args.task_id}.json" + + if not Path(json_file).exists(): + parser.error(f"Cannot find task info file: {json_file}") + + main( + src_path=args.src_path, + tgt_path=args.tgt_path, + task_id=task_id, + repo_id=args.repo_id, + task_info_json=json_file, + debug=args.debug, + chunk_size=args.chunk_size, + ) + + elif format_type == "new": + main( + src_path=args.src_path, + tgt_path=args.tgt_path, + task_id=args.task_id, + repo_id=args.repo_id, + task_info_json=None, + debug=args.debug, + chunk_size=args.chunk_size, + ) + + else: + parser.error( + f"Unknown dataset format. Please check the directory structure at: {args.src_path}" + ) + diff --git a/scripts/data/convert_droid.py b/scripts/data/convert_droid.py new file mode 100644 index 0000000000000000000000000000000000000000..2a548e196f6912b341c3906c02b9dcd89584f5af --- /dev/null +++ b/scripts/data/convert_droid.py @@ -0,0 +1,679 @@ +""" +Convert DROID 1.0.1 (RLDS/TFDS format) to LeRobot format with idle frame filtering. + +This script takes the raw DROID dataset in RLDS format, applies idle frame filtering +using a pre-computed JSON file of non-idle frame ranges, filters out failed episodes +and episodes without language annotations, and outputs the dataset in LeRobot v2.0 format. + +The idle filtering is based on Physical Intelligence's approach (see openpi): + https://github.com/Physical-Intelligence/openpi/blob/main/examples/droid/README_train.md + +The pre-computed idle filter ranges can be downloaded from: + gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json + +Usage: + python scripts/data/convert_droid.py \\ + --keep-ranges-path \\ + [--fps 15] [--first-n N] [-n 16] [--filter-failed] + +Example: + # Download DROID 1.0.1 raw dataset + gsutil -m cp -r gs://gresearch/robotics/droid/1.0.1 ./data/droid/1.0.1 + + # Download idle filter ranges from openpi + gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json ./data/keep_ranges.json + + # Run conversion + python scripts/data/convert_droid.py ./data/droid/1.0.1 ./data/droid_lerobot \\ + --keep-ranges-path ./data/keep_ranges.json --filter-failed + +Original dataset structure (RLDS): + - 3 camera views: exterior_image_1_left, exterior_image_2_left, wrist_image_left + - State: cartesian_position (6), gripper_position (1), joint_position (7) + - Action: cartesian_position (6), cartesian_velocity (6), gripper_position (1), + gripper_velocity (1), joint_position (7), joint_velocity (7) + - Language instructions (up to 3 per episode) + +Credits: + - Original conversion script by Loic Magne (NVIDIA) + - Idle filtering by Scott Reed (NVIDIA), based on Physical Intelligence's approach +""" + +from concurrent.futures import ProcessPoolExecutor, as_completed +import json +import multiprocessing as mp +import os +from pathlib import Path + +import av +import numpy as np +import polars as pl +import tensorflow as tf +import tensorflow_datasets as tfds +import torch +import tqdm + +# Limit thread counts to avoid oversubscription in multiprocessing +os.environ["MKL_NUM_THREADS"] = "1" +os.environ["NUMEXPR_NUM_THREADS"] = "1" +os.environ["OMP_NUM_THREADS"] = "1" +os.environ["MPI_NUM_THREADS"] = "1" +os.environ["TF_NUM_INTRAOP_THREADS"] = "1" +os.environ["TF_NUM_INTEROP_THREADS"] = "1" +os.environ["OPENBLAS_NUM_THREADS"] = "1" +os.environ["VECLIB_MAXIMUM_THREADS"] = "1" + +tf.get_logger().setLevel("WARN") +tf.config.threading.set_inter_op_parallelism_threads(1) +tf.config.threading.set_intra_op_parallelism_threads(1) +tf.config.set_soft_device_placement(True) + + +def tf_to_torch(data): + return torch.from_numpy(data.numpy()) + + +def tf_img_convert(img): + if img.dtype == tf.string: + img = tf.io.decode_image(img, expand_animations=False, dtype=tf.uint8) + elif img.dtype != tf.uint8: + raise ValueError(f"Unsupported image dtype: found with dtype {img.dtype}") + return img.numpy() + + +def _broadcast_metadata_rlds(i: tf.Tensor, traj: dict) -> dict: + steps = traj.pop("steps") + traj_len = tf.shape(tf.nest.flatten(steps)[0])[0] + metadata = tf.nest.map_structure(lambda x: tf.repeat(x, traj_len), traj) + traj = {**steps, "traj_metadata": metadata} + traj["_len"] = tf.repeat(traj_len, traj_len) + traj["_traj_index"] = tf.repeat(i, traj_len) + traj["_frame_index"] = tf.range(traj_len) + return traj + + +def concat_state_or_action(modality_dict, keys, compute_concat_info=False): + arrays = [] + if compute_concat_info: + concat_info = {} + start_index = 0 + for key in keys: + array = tf_to_torch(modality_dict[key]) + arrays.append(array) + if compute_concat_info: + D = array.shape[1] + data_dtype = array.numpy().dtype + if np.issubdtype(data_dtype, bool): + data_dtype = "int64" + data_range = [0, 1] + else: + data_dtype = data_dtype.name + data_range = None + concat_info[key] = { + "start": start_index, + "end": start_index + D, + } + if data_dtype != "float64": + concat_info[key]["dtype"] = data_dtype + if data_range is not None: + concat_info[key]["range"] = data_range + start_index += D + concatenated = torch.cat(arrays, dim=1) + ret_dict = {} + if compute_concat_info: + ret_dict["concat_info"] = concat_info + ret_dict["concatenated"] = concatenated + return ret_dict + + +def encode_video(frames: np.ndarray, output_path: Path, fps: int) -> None: + """Encode a sequence of frames to a video file using PyAV.""" + options = { + "threads": "1", + "thread_type": "slice", + "preset": "ultrafast", + "tune": "zerolatency", + "crf": "23", + } + + container = av.open(str(output_path), mode="w") + stream = container.add_stream("h264", rate=fps, options=options) + stream.width = frames.shape[2] + stream.height = frames.shape[1] + stream.pix_fmt = "yuv420p" + + video_frame = av.VideoFrame(width=stream.width, height=stream.height, format="rgb24") + frame_array = video_frame.to_ndarray(format="rgb24") + + for frame in frames: + frame_array[:] = frame + packet = stream.encode(video_frame) + container.mux(packet) + + packet = stream.encode(None) + container.mux(packet) + container.close() + + +def process_tfrecord( + ith_shard, + raw_dir, + output_path, + fps, + all_tasks, + state_keys, + action_keys, + lang_keys, + image_keys, + start_episode_idx, + kept_registry, + keep_ranges_path, +): + config = tfds.ReadConfig( + try_autocache=False, + num_parallel_calls_for_decode=1, + num_parallel_calls_for_interleave_files=1, + interleave_cycle_length=1, + shuffle_reshuffle_each_iteration=False, + ) + + ds_builder = tfds.builder_from_directory(str(raw_dir)) + dataset = ds_builder.as_dataset( + split=f"train[{ith_shard}shard]", + decoders={"steps": tfds.decode.SkipDecoding()}, + read_config=config, + ) + + dataset = dataset.enumerate().map(_broadcast_metadata_rlds) + all_keep_ranges = json.load(open(keep_ranges_path, "r")) + + episodes_data = [] + for local_idx, episode in enumerate(dataset): + + # Add keep frame info to episode. + file_path = ( + episode["traj_metadata"]["episode_metadata"]["file_path"][0].numpy().decode("utf-8") + ) + recording_folderpath = ( + episode["traj_metadata"]["episode_metadata"]["recording_folderpath"][0] + .numpy() + .decode("utf-8") + ) + idle_key = f"{recording_folderpath}--{file_path}" + keep_ranges = all_keep_ranges[idle_key] + + global_episode_idx = start_episode_idx + local_idx + + # check if the episode has been filtered + if global_episode_idx not in kept_registry: + continue + + episode_idx = kept_registry[global_episode_idx] + episode_data = process_sample( + episode_idx, + episode, + output_path, + fps, + all_tasks, + state_keys, + action_keys, + lang_keys, + image_keys, + keep_ranges, + ) + episodes_data.append(episode_data) + return episodes_data + + +def process_sample( + ep_idx, + episode, + output_path, + fps, + all_tasks, + state_keys, + action_keys, + lang_keys, + image_keys, + keep_ranges, +): + chunk_idx = ep_idx // 1000 + + # Create chunk directory + (output_path / f"data/chunk-{chunk_idx:03d}").mkdir(parents=True, exist_ok=True) + for img_key in image_keys: + (output_path / f"videos/chunk-{chunk_idx:03d}/observation.images.{img_key}").mkdir( + parents=True, exist_ok=True + ) + + # Use concat_state_or_action for state and action + state_dict = concat_state_or_action(episode["observation"], state_keys) + action_dict = concat_state_or_action(episode["action_dict"], action_keys) + + # Count number of non-idle frames. + num_frames = len(episode["observation"][state_keys[0]]) + actual_num_frames = 0 + for start_ix, end_ix in keep_ranges: + actual_num_frames += end_ix - start_ix + + # Build episode data dictionary + episode_dict = { + "observation.state": state_dict["concatenated"].numpy(), + "action": action_dict["concatenated"].numpy(), + "next.reward": tf_to_torch(episode["reward"]).numpy(), + "next.done": tf_to_torch(episode["is_last"]).numpy(), + "is_terminal": tf_to_torch(episode["is_terminal"]).numpy(), + "is_first": tf_to_torch(episode["is_first"]).numpy(), + "discount": tf_to_torch(episode["discount"]).numpy(), + "timestamp": np.arange(actual_num_frames) / fps, + "episode_index": np.full(actual_num_frames, ep_idx), + "frame_index": np.arange(actual_num_frames), + } + + # Initialize all annotation columns with default value + for lang_key in lang_keys: + episode_dict[f"annotation.language.{lang_key}"] = np.full( + num_frames, all_tasks["not provided"], dtype=np.int64 + ) + + # Add language instruction indices to parquet + episode_tasks = [] + for lang_key in lang_keys: + if lang_key in episode: + task = episode[lang_key][0].numpy().decode("utf-8") + if task and len(task) > 1: + episode_tasks.append(task) + task_idx = all_tasks[task] + episode_dict[f"annotation.language.{lang_key}"] = np.full( + num_frames, task_idx, dtype=np.int64 + ) + + # Set task_index to match the first language instruction annotation + episode_dict["task_index"] = episode_dict[f"annotation.language.{lang_keys[0]}"].copy() + + # Filter idle frames from episode_dict. + for key in episode_dict: + if key in ["timestamp", "episode_index", "frame_index"]: + continue + tensor_parts = [] + for start_ix, end_ix in keep_ranges: + tensor_parts.append(episode_dict[key][start_ix:end_ix]) + episode_dict[key] = np.concatenate(tensor_parts, axis=0) + + # Filter idle frames from observation images. + for img_key in image_keys: + video_parts = [] + all_frames = np.stack( + [tf_img_convert(episode["observation"][img_key][i]) for i in range(num_frames)] + ) + for start_ix, end_ix in keep_ranges: + video_parts.append(all_frames[start_ix:end_ix]) + new_video = np.concatenate(video_parts, axis=0) + assert new_video.shape[0] == actual_num_frames + episode["observation"][img_key] = new_video + + # Save to parquet using polars + df = pl.DataFrame(episode_dict) + parquet_path = output_path / f"data/chunk-{chunk_idx:03d}/episode_{ep_idx:06d}.parquet" + df.write_parquet(parquet_path) + + # Process videos for each image key + for img_key in image_keys: + frames = episode["observation"][img_key] + video_path = ( + output_path + / f"videos/chunk-{chunk_idx:03d}/observation.images.{img_key}/episode_{ep_idx:06d}.mp4" + ) + encode_video(frames, video_path, fps) + + episode_data = { + "episode_index": ep_idx, + "tasks": episode_tasks, + "length": actual_num_frames, + "success": bool(np.any(tf_to_torch(episode["reward"]).numpy() != 0)), + } + return episode_data + + +def convert_droid_dataset( + raw_dir: str, + output_dir: str, + keep_ranges_path: str, + fps: int = 15, + first_n: int | None = None, + max_workers: int = 16, + filter_failed: bool = False, +): + """ + Convert DROID 1.0.1 RLDS dataset to LeRobot format with idle filtering. + + Args: + raw_dir: Path to raw DROID RLDS dataset (e.g., ./data/droid/1.0.1) + output_dir: Path to output directory for LeRobot dataset + keep_ranges_path: Path to JSON file containing idle filter ranges. + Download from: gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json + fps: Frames per second for output videos + first_n: Only process the first N tfrecord shards (for debugging) + max_workers: Max workers for multiprocessing + filter_failed: Whether to filter out failed episodes (all zero rewards) + """ + output_path = Path(output_dir) + + # Validate keep_ranges_path exists + if not os.path.exists(keep_ranges_path): + raise FileNotFoundError( + f"Keep ranges file not found: {keep_ranges_path}\n" + "Download it with: gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json " + ) + + # Load dataset + config = tfds.ReadConfig( + try_autocache=False, + num_parallel_calls_for_decode=1, + num_parallel_calls_for_interleave_files=1, + interleave_cycle_length=1, + shuffle_reshuffle_each_iteration=False, + ) + ds_builder = tfds.builder_from_directory(str(raw_dir)) + split_str = f"train[:{first_n}shard]" if first_n is not None else "train" + dataset = ds_builder.as_dataset( + split=split_str, + decoders={"steps": tfds.decode.SkipDecoding()}, + read_config=config, + ) + dataset_info = ds_builder.info + dataset = dataset.enumerate().map(_broadcast_metadata_rlds) + + # Extract keys + image_keys = [] + state_keys = [ + "cartesian_position", + "gripper_position", + "joint_position", + ] + action_keys = [ + "cartesian_position", + "cartesian_velocity", + "gripper_position", + "gripper_velocity", + "joint_position", + "joint_velocity", + ] + lang_keys = [ + "language_instruction", + "language_instruction_2", + "language_instruction_3", + ] + + observation_info = dataset_info.features["steps"]["observation"] + for key in observation_info: + if len(observation_info[key].shape) == 3: + if observation_info[key].dtype == tf.uint8: + image_keys.append(key) + else: + assert key in state_keys, f"{key=}, {state_keys=}" + + print(f"Found image keys: {image_keys}") + print(f"Using state keys: {state_keys}") + print(f"Using action keys: {action_keys}") + + (output_path / "meta").mkdir(parents=True, exist_ok=True) + + # Get concat info for modality.json from first episode + first_episode = next(iter(dataset)) + state_info = concat_state_or_action( + first_episode["observation"], state_keys, compute_concat_info=True + ) + action_info = concat_state_or_action( + first_episode["action_dict"], action_keys, compute_concat_info=True + ) + + # Generate modality.json + modality_config = { + "state": state_info["concat_info"], + "action": action_info["concat_info"], + "video": {k: {"original_key": f"observation.images.{k}"} for k in image_keys}, + "annotation": {f"language.{lang_key}": {} for lang_key in lang_keys}, + } + + with open(output_path / "meta/modality.json", "w") as f: + json.dump(modality_config, f, indent=4) + + # Get file instructions from TFDS + ds_builder = tfds.builder_from_directory(str(raw_dir)) + file_instructions = ds_builder.info.splits["train"].file_instructions + if first_n is not None: + file_instructions = file_instructions[:first_n] + + # First pass: collect unique tasks and determine which episodes to keep + all_tasks = {} # task string -> task index + task_counter = 0 + print(f"First pass: collecting unique tasks from {len(dataset)} episodes") + # Add a default "not provided" task + all_tasks["not provided"] = task_counter + task_counter += 1 + + # kept_registry maps global episode index -> filtered episode index + kept_registry = {} + kept_count = 0 + all_keep_ranges = json.load(open(keep_ranges_path, "r")) + + for i, episode in enumerate(tqdm.tqdm(dataset)): + # filter out failed episodes + filtered = False + if filter_failed: + if not np.any(tf_to_torch(episode["reward"]).numpy() != 0): + filtered = True + + # Check language annotations + has_lang = False + for lang_key in lang_keys: + if lang_key in episode: + task = episode[lang_key][0].numpy().decode("utf-8") + if task and (len(task) > 1) and task not in all_tasks: + has_lang = True + all_tasks[task] = task_counter + task_counter += 1 + if not has_lang: + # Do not include episodes missing language annotations + filtered = True + + # Filter out episodes that are only idle + file_path = ( + episode["traj_metadata"]["episode_metadata"]["file_path"][0].numpy().decode("utf-8") + ) + recording_folderpath = ( + episode["traj_metadata"]["episode_metadata"]["recording_folderpath"][0] + .numpy() + .decode("utf-8") + ) + idle_key = f"{recording_folderpath}--{file_path}" + keep_ranges = all_keep_ranges[idle_key] + if len(keep_ranges) == 0: + filtered = True + + if not filtered: + kept_registry[i] = kept_count + kept_count += 1 + + print(f"Kept {len(kept_registry)}/{len(dataset)} episodes") + + # Write tasks.jsonl + with open(output_path / "meta/tasks.jsonl", "w") as f: + for task, task_idx in all_tasks.items(): + f.write(json.dumps({"task_index": task_idx, "task": task}) + "\n") + + if max_workers > 1: + # Calculate process args with cumulative indices + cumsum = 0 + process_args = [] + for i, instruction in enumerate(file_instructions): + args = ( + i, + raw_dir, + output_path, + fps, + all_tasks, + state_keys, + action_keys, + lang_keys, + image_keys, + cumsum, + kept_registry, + keep_ranges_path, + ) + process_args.append(args) + cumsum += instruction.examples_in_shard + + ctx = mp.get_context("spawn") + with ProcessPoolExecutor(mp_context=ctx, max_workers=max_workers) as executor: + futures = [executor.submit(process_tfrecord, *args) for args in process_args] + episodes_data = [] + for future in tqdm.tqdm(as_completed(futures), total=len(futures)): + episodes_data.extend(future.result()) + else: + episodes_data = [] + cumsum = 0 + for i, instruction in enumerate(file_instructions): + episodes_data.extend( + process_tfrecord( + i, + raw_dir, + output_path, + fps, + all_tasks, + state_keys, + action_keys, + lang_keys, + image_keys, + cumsum, + kept_registry, + keep_ranges_path, + ) + ) + cumsum += instruction.examples_in_shard + + # Order episodes by episode index + episodes_data = sorted(episodes_data, key=lambda x: x["episode_index"]) + + # Generate episodes.jsonl + with open(output_path / "meta/episodes.jsonl", "w") as f: + for episode in episodes_data: + f.write(json.dumps(episode) + "\n") + + # Generate info.json + ds_length = len(episodes_data) + num_chunks = (ds_length // 1000) + (1 if ds_length % 1000 else 0) + info = { + "codebase_version": "v2.0", + "robot_type": "droid", + "total_episodes": ds_length, + "total_frames": sum(ep["length"] for ep in episodes_data), + "total_tasks": len(all_tasks), + "total_videos": len(image_keys), + "total_chunks": num_chunks, + "chunks_size": 1000, + "fps": fps, + "splits": {"train": "0:100"}, + "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + "video_path": "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4", + "features": { + # Video features + **{ + f"observation.images.{k}": { + "dtype": "video", + "shape": list(tf_img_convert(first_episode["observation"][k][0]).shape), + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": fps, + "video.codec": "h264", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + } + for k in image_keys + }, + # State feature + "observation.state": { + "dtype": "float64", + "shape": [state_info["concatenated"].shape[1]], + "names": state_keys, + }, + # Action feature + "action": { + "dtype": "float64", + "shape": [action_info["concatenated"].shape[1]], + "names": action_keys, + }, + # Single value features + "timestamp": {"dtype": "float64", "shape": [1]}, + "task_index": {"dtype": "int64", "shape": [1]}, + "episode_index": {"dtype": "int64", "shape": [1]}, + "index": {"dtype": "int64", "shape": [1]}, + "next.reward": {"dtype": "float64", "shape": [1]}, + "next.done": {"dtype": "bool", "shape": [1]}, + "is_terminal": {"dtype": "bool", "shape": [1]}, + "is_first": {"dtype": "bool", "shape": [1]}, + "discount": {"dtype": "float64", "shape": [1]}, + # Language annotation features + **{f"annotation.language.{k}": {"dtype": "int64", "shape": [1]} for k in lang_keys}, + }, + } + + with open(output_path / "meta/info.json", "w") as f: + json.dump(info, f, indent=4) + + # Sanity check: chunk directories should contain exactly 1000 episodes (except last) + for i in range(num_chunks): + chunk_path = output_path / f"data/chunk-{i:03d}" + episodes = list(chunk_path.glob("episode_*.parquet")) + assert ( + len(episodes) == 1000 if i != num_chunks - 1 else len(episodes) <= 1000 + ), f"chunk-{i:03d} contains {len(episodes)} episodes" + + for img_key in image_keys: + img_path = output_path / f"videos/chunk-{i:03d}/observation.images.{img_key}" + episodes = list(img_path.glob("episode_*.mp4")) + assert ( + len(episodes) == 1000 if i != num_chunks - 1 else len(episodes) <= 1000 + ), f"{img_path} contains {len(episodes)} episodes" + + print("Sanity check passed.") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description="Convert DROID 1.0.1 (RLDS) to LeRobot format with idle filtering." + ) + parser.add_argument("raw_dir", help="Path to raw DROID RLDS dataset (e.g., ./data/droid/1.0.1)") + parser.add_argument("output_dir", help="Path to output directory for LeRobot dataset") + parser.add_argument( + "--keep-ranges-path", + required=True, + help="Path to idle filter JSON file. Download with: " + "gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json ", + ) + parser.add_argument("--fps", type=int, default=15, help="Frames per second for videos") + parser.add_argument( + "--first-n", type=int, help="Only convert first N tfrecord shards (for debugging)" + ) + parser.add_argument("-n", type=int, default=16, help="Max workers for multiprocessing") + parser.add_argument( + "--filter-failed", + action="store_true", + help="Whether to filter out failed episodes (i.e., episodes with all zero rewards)", + ) + args = parser.parse_args() + + convert_droid_dataset( + args.raw_dir, + args.output_dir, + args.keep_ranges_path, + args.fps, + args.first_n, + args.n, + args.filter_failed, + ) \ No newline at end of file diff --git a/scripts/data/convert_lerobot_to_gear.py b/scripts/data/convert_lerobot_to_gear.py new file mode 100644 index 0000000000000000000000000000000000000000..6576b61249aae9524e9b01dba77047c50228f21b --- /dev/null +++ b/scripts/data/convert_lerobot_to_gear.py @@ -0,0 +1,634 @@ +""" +Convert a standard LeRobot v2 dataset to the GEAR/DreamZero training format. + +This script takes a dataset collected with LeRobot v2 and generates/augments the +metadata files required by DreamZero's training pipeline: + + - meta/modality.json (state/action/video/annotation key mapping) + - meta/embodiment.json (embodiment tag for the training pipeline) + - meta/stats.json (dataset-level statistics: mean, std, min, max, q01, q99) + - meta/relative_stats_dreamzero.json (relative action statistics) + - meta/tasks.jsonl (task descriptions) + - meta/episodes.jsonl (episode-level metadata) + +The script does NOT modify parquet files or videos -- it only creates metadata. + +Usage: + # Auto-detect state/action structure, default embodiment tag 'xdof': + python scripts/data/convert_lerobot_to_gear.py --dataset-path ./Dataset/my_robot_data + + # Explicit modality mapping via JSON: + python scripts/data/convert_lerobot_to_gear.py \\ + --dataset-path ./Dataset/my_robot_data \\ + --embodiment-tag xdof \\ + --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 \\ + --task-key annotation.task + + # Copy to a new output directory instead of modifying in-place: + python scripts/data/convert_lerobot_to_gear.py \\ + --dataset-path ./Dataset/my_robot_data \\ + --output-path ./Dataset/my_robot_data_gear +""" + +from __future__ import annotations + +import argparse +import json +import logging +import shutil +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +from tqdm import tqdm + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + +VALID_EMBODIMENT_TAGS = [ + "real_gr1_arms_only", "real_gr1_arms_only_annotated", + "real_gr1_arms_waist", "real_gr1_arms_waist_annotated", + "dexmg_gr1_arms_only_inspire", "dexmg_gr1_arms_only_fourier", + "dexmg_gr1_arms_waist_fourier", + "robocasa_single_arm", "onex_eve_gripper", + "robocasa_gr1_arms_only_inspire_hands", "robocasa_gr1_arms_only_fourier_hands", + "robocasa_gr1_fixed_lower_body_inspire_hands", "robocasa_gr1_fixed_lower_body_fourier_hands", + "robocasa_panda_omron", + "robocasa_bimanual_panda_parallel_gripper", "robocasa_bimanual_panda_inspire_hand", + "oxe_droid", "oxe_fractal", "oxe_language_table", "oxe_bridge", + "real_panda_single_arm", "hot3d_hands_only", + "gr1_unified", "robocasa_gr1_arms_waist_fourier_hands", + "agibot", "lapa", "oxe_mutex", "oxe_roboset", "oxe_plex", + "dream", "yam", "xdof", + "gr1_unified_segmentation", "language_table_sim", "gr1_isaac", + "sim_behavior_r1_pro", "mecka_hands", "real_r1_pro_sharpa", "libero_sim", +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def load_info(dataset_path: Path) -> dict: + info_path = dataset_path / "meta" / "info.json" + if not info_path.exists(): + log.error("meta/info.json not found at %s", info_path) + sys.exit(1) + with open(info_path) as f: + return json.load(f) + + +def get_parquet_paths(dataset_path: Path, info: dict) -> list[Path]: + pattern = info.get("data_path", "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet") + total_episodes = info["total_episodes"] + chunks_size = info.get("chunks_size", 1000) + paths = [] + for ep_idx in range(total_episodes): + chunk_idx = ep_idx // chunks_size + p = dataset_path / pattern.format(episode_chunk=chunk_idx, episode_index=ep_idx) + if p.exists(): + paths.append(p) + return sorted(paths) + + +def detect_features(info: dict) -> dict: + """Return categorised feature names from info.json.""" + features = info.get("features", {}) + state_keys = [k for k in features if k.startswith("observation.state")] + action_keys = [k for k in features if k == "action" or k.startswith("action.")] + video_keys = [k for k in features if features[k].get("dtype") == "video"] + annotation_keys = [k for k in features if k.startswith("annotation")] + return { + "state": state_keys, + "action": action_keys, + "video": video_keys, + "annotation": annotation_keys, + "features": features, + } + + +def parse_key_mapping(raw: str | None) -> dict[str, list[int]] | None: + """Parse a JSON string like '{"joint_pos": [0, 6], "gripper": [6, 7]}'.""" + if raw is None: + return None + try: + mapping = json.loads(raw) + except json.JSONDecodeError as e: + log.error("Invalid JSON for key mapping: %s", e) + sys.exit(1) + for name, bounds in mapping.items(): + if not isinstance(bounds, list) or len(bounds) != 2: + log.error("Each entry must be [start, end]. Got %s for '%s'", bounds, name) + sys.exit(1) + return mapping + + +# --------------------------------------------------------------------------- +# Modality JSON +# --------------------------------------------------------------------------- + +def build_modality_json( + info: dict, + detected: dict, + state_mapping: dict[str, list[int]] | None, + action_mapping: dict[str, list[int]] | None, + task_key: str | None, +) -> dict: + """Build the modality.json structure expected by GEAR/DreamZero.""" + features = detected["features"] + modality: dict = {"state": {}, "action": {}, "video": {}, "annotation": {}} + + # --- State --- + state_col = detected["state"][0] if detected["state"] else None + if state_col and state_mapping: + for name, (start, end) in state_mapping.items(): + dtype = features[state_col].get("dtype", "float64") + modality["state"][name] = { + "original_key": state_col, + "start": start, + "end": end, + "rotation_type": None, + "absolute": True, + "dtype": dtype, + "range": None, + } + elif state_col: + shape = features[state_col].get("shape", [1]) + dim = shape[0] if isinstance(shape, list) else shape + dtype = features[state_col].get("dtype", "float64") + modality["state"]["state"] = { + "original_key": state_col, + "start": 0, + "end": dim, + "rotation_type": None, + "absolute": True, + "dtype": dtype, + "range": None, + } + + # --- Action --- + action_col = detected["action"][0] if detected["action"] else None + if action_col and action_mapping: + for name, (start, end) in action_mapping.items(): + dtype = features[action_col].get("dtype", "float64") + modality["action"][name] = { + "original_key": action_col, + "start": start, + "end": end, + "rotation_type": None, + "absolute": True, + "dtype": dtype, + "range": None, + } + elif action_col: + shape = features[action_col].get("shape", [1]) + dim = shape[0] if isinstance(shape, list) else shape + dtype = features[action_col].get("dtype", "float64") + modality["action"]["action"] = { + "original_key": action_col, + "start": 0, + "end": dim, + "rotation_type": None, + "absolute": True, + "dtype": dtype, + "range": None, + } + + # --- Video --- + for vk in detected["video"]: + short_name = vk.replace("observation.images.", "") + modality["video"][short_name] = {"original_key": vk} + + # --- Annotation --- + if task_key: + short = task_key.replace("annotation.", "") + modality["annotation"][short] = {"original_key": task_key} + else: + for ak in detected["annotation"]: + short = ak.replace("annotation.", "") + modality["annotation"][short] = {"original_key": ak} + + return modality + + +# --------------------------------------------------------------------------- +# Stats computation +# --------------------------------------------------------------------------- + +def compute_stats(parquet_paths: list[Path], columns: list[str]) -> dict: + """Compute mean/std/min/max/q01/q99 for numeric columns across all episodes.""" + all_data: dict[str, list] = {col: [] for col in columns} + for pp in tqdm(parquet_paths, desc="Computing stats"): + df = pd.read_parquet(pp) + for col in columns: + if col not in df.columns: + continue + arr = np.stack(df[col].values) + if arr.ndim == 1: + arr = arr.reshape(-1, 1) + all_data[col].append(arr) + + stats = {} + for col in columns: + if not all_data[col]: + continue + data = np.concatenate(all_data[col], axis=0).astype(np.float64) + stats[col] = { + "mean": np.mean(data, axis=0).tolist(), + "std": np.std(data, axis=0).tolist(), + "min": np.min(data, axis=0).tolist(), + "max": np.max(data, axis=0).tolist(), + "q01": np.quantile(data, 0.01, axis=0).tolist(), + "q99": np.quantile(data, 0.99, axis=0).tolist(), + } + return stats + + +def compute_relative_stats( + parquet_paths: list[Path], + modality: dict, + relative_action_keys: list[str], + action_horizon: int = 24, +) -> dict: + """Compute relative-action statistics: (action - reference_state) for each key. + + This replicates the logic in groot/vla/data/dataset/lerobot.py + _calculate_relative_stats_for_key. + """ + stats: dict = {} + for rel_key in relative_action_keys: + if rel_key not in modality["action"]: + log.warning("Relative action key '%s' not found in action modality, skipping", rel_key) + continue + if rel_key not in modality["state"]: + log.warning( + "Relative action key '%s' has no matching state key -- " + "relative stats require a corresponding state key with the same name. Skipping.", + rel_key, + ) + continue + + action_meta = modality["action"][rel_key] + state_meta = modality["state"][rel_key] + + all_relative = [] + for pp in tqdm(parquet_paths, desc=f"Relative stats [{rel_key}]"): + df = pd.read_parquet(pp) + action_col = action_meta["original_key"] + state_col = state_meta["original_key"] + if action_col not in df.columns or state_col not in df.columns: + continue + + action_data = np.stack(df[action_col].values).astype(np.float64) + state_data = np.stack(df[state_col].values).astype(np.float64) + if action_data.ndim == 1: + action_data = action_data.reshape(-1, 1) + if state_data.ndim == 1: + state_data = state_data.reshape(-1, 1) + + a_start, a_end = action_meta["start"], action_meta["end"] + s_start, s_end = state_meta["start"], state_meta["end"] + + action_slice = action_data[:, a_start:a_end] + state_slice = state_data[:, s_start:s_end] + + traj_len = len(df) + usable = traj_len - action_horizon + for i in range(max(usable, 0)): + ref_state = state_slice[i] + chunk_end = min(i + action_horizon, traj_len) + actions = action_slice[i:chunk_end] + relative = actions - ref_state + all_relative.extend(relative) + + if not all_relative: + log.warning("No relative actions computed for '%s'", rel_key) + continue + + data = np.array(all_relative) + stats[rel_key] = { + "max": np.max(data, axis=0).tolist(), + "min": np.min(data, axis=0).tolist(), + "mean": np.mean(data, axis=0).tolist(), + "std": np.std(data, axis=0).tolist(), + "q01": np.quantile(data, 0.01, axis=0).tolist(), + "q99": np.quantile(data, 0.99, axis=0).tolist(), + } + + return stats + + +# --------------------------------------------------------------------------- +# Tasks & episodes +# --------------------------------------------------------------------------- + +def build_tasks(parquet_paths: list[Path], task_key: str | None) -> list[dict]: + """Build tasks.jsonl entries from the dataset.""" + if task_key is None: + return [{"task_index": 0, "task": ""}] + + task_set: dict[str, int] = {} + for pp in tqdm(parquet_paths, desc="Extracting tasks"): + df = pd.read_parquet(pp) + if task_key not in df.columns: + continue + for val in df[task_key].unique(): + text = str(val) if not isinstance(val, str) else val + if text not in task_set: + task_set[text] = len(task_set) + + if not task_set: + return [{"task_index": 0, "task": ""}] + + return [{"task_index": idx, "task": text} for text, idx in sorted(task_set.items(), key=lambda x: x[1])] + + +def build_episodes(parquet_paths: list[Path], info: dict, task_key: str | None, tasks: list[dict]) -> list[dict]: + """Build episodes.jsonl entries.""" + task_text_to_idx = {t["task"]: t["task_index"] for t in tasks} + episodes = [] + for ep_idx, pp in enumerate(tqdm(parquet_paths, desc="Building episodes")): + df = pd.read_parquet(pp) + length = len(df) + + ep_tasks: list[str] = [] + if task_key and task_key in df.columns: + unique_tasks = df[task_key].unique() + for t in unique_tasks: + text = str(t) if not isinstance(t, str) else t + if text and text in task_text_to_idx: + ep_tasks.append(text) + if not ep_tasks: + ep_tasks = [""] + + episodes.append({ + "episode_index": ep_idx, + "tasks": ep_tasks, + "length": length, + }) + + return episodes + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + +def validate_dataset(dataset_path: Path, info: dict, modality: dict) -> list[str]: + """Run basic validation and return a list of warnings.""" + warnings = [] + + # Check required directories + for subdir in ["data", "videos", "meta"]: + if not (dataset_path / subdir).exists(): + warnings.append(f"Missing directory: {subdir}/") + + # Check at least one video key exists + if not modality["video"]: + warnings.append("No video features detected -- DreamZero requires at least one camera view") + + # Check state/action exist + if not modality["state"]: + warnings.append("No state modality keys defined") + if not modality["action"]: + warnings.append("No action modality keys defined") + + # Check total_episodes > 0 + if info.get("total_episodes", 0) == 0: + warnings.append("total_episodes is 0 in info.json") + + # Check FPS + if info.get("fps") is None: + warnings.append("fps not set in info.json") + + return warnings + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser( + description="Convert a LeRobot v2 dataset to GEAR/DreamZero training format.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument("--dataset-path", type=str, required=True, help="Path to the LeRobot v2 dataset") + parser.add_argument("--output-path", type=str, default=None, help="Output path (default: modify in-place)") + parser.add_argument( + "--embodiment-tag", type=str, default="xdof", + help=f"Embodiment tag (default: xdof). Valid: {', '.join(sorted(set(VALID_EMBODIMENT_TAGS)))}" + ) + parser.add_argument( + "--state-keys", type=str, default=None, + help='JSON mapping of state sub-keys to [start, end] index ranges, ' + 'e.g. \'{"joint_pos": [0, 6], "gripper_pos": [6, 7]}\'' + ) + parser.add_argument( + "--action-keys", type=str, default=None, + help='JSON mapping of action sub-keys to [start, end] index ranges' + ) + parser.add_argument( + "--relative-action-keys", type=str, nargs="*", default=None, + help="Action sub-key names to compute relative stats for (e.g. joint_pos gripper_pos). " + "Each key must also exist in --state-keys. If omitted, skips relative stats." + ) + parser.add_argument("--task-key", type=str, default=None, help="Column name for language annotations (auto-detected if not set)") + parser.add_argument("--fps", type=float, default=None, help="Override FPS (default: use dataset FPS from info.json)") + parser.add_argument("--action-horizon", type=int, default=24, help="Action horizon for relative stats (default: 24)") + parser.add_argument("--force", action="store_true", help="Overwrite existing GEAR metadata files") + + args = parser.parse_args() + + dataset_path = Path(args.dataset_path).resolve() + if not dataset_path.exists(): + log.error("Dataset path does not exist: %s", dataset_path) + sys.exit(1) + + # Validate embodiment tag + if args.embodiment_tag not in VALID_EMBODIMENT_TAGS: + log.error( + "Invalid embodiment tag '%s'. Valid tags:\n %s", + args.embodiment_tag, + "\n ".join(sorted(set(VALID_EMBODIMENT_TAGS))), + ) + sys.exit(1) + + # Output path handling + if args.output_path: + output_path = Path(args.output_path).resolve() + if output_path != dataset_path: + log.info("Copying dataset to %s", output_path) + if output_path.exists(): + if not args.force: + log.error("Output path already exists. Use --force to overwrite.") + sys.exit(1) + shutil.rmtree(output_path) + shutil.copytree(dataset_path, output_path) + dataset_path = output_path + else: + output_path = dataset_path + + meta_dir = output_path / "meta" + meta_dir.mkdir(parents=True, exist_ok=True) + + # 1. Load info.json + info = load_info(dataset_path) + detected = detect_features(info) + + log.info("Dataset: %s", dataset_path.name) + log.info(" Episodes: %d", info.get("total_episodes", 0)) + log.info(" FPS: %s", info.get("fps", "not set")) + log.info(" State columns: %s", detected["state"]) + log.info(" Action columns: %s", detected["action"]) + log.info(" Video features: %d camera(s)", len(detected["video"])) + log.info(" Annotation columns: %s", detected["annotation"]) + + if args.fps is not None: + info["fps"] = args.fps + with open(output_path / "meta" / "info.json", "w") as f: + json.dump(info, f, indent=4) + log.info(" Overriding FPS to %s", args.fps) + + # Parse user-provided key mappings + state_mapping = parse_key_mapping(args.state_keys) + action_mapping = parse_key_mapping(args.action_keys) + + # Auto-detect task key if not provided + task_key = args.task_key + if task_key is None and detected["annotation"]: + for candidate in ["annotation.task", "annotation.language.language_instruction"]: + if candidate in detected["annotation"]: + task_key = candidate + break + if task_key is None: + task_key = detected["annotation"][0] + log.info(" Auto-detected task key: %s", task_key) + + # 2. Build modality.json + modality = build_modality_json(info, detected, state_mapping, action_mapping, task_key) + + modality_path = meta_dir / "modality.json" + if modality_path.exists() and not args.force: + log.info(" modality.json already exists, skipping (use --force to overwrite)") + else: + with open(modality_path, "w") as f: + json.dump(modality, f, indent=4) + log.info(" Wrote modality.json (%d state keys, %d action keys, %d video keys)", + len(modality["state"]), len(modality["action"]), len(modality["video"])) + + # 3. Write embodiment.json + embodiment = {"robot_type": args.embodiment_tag, "embodiment_tag": args.embodiment_tag} + embodiment_path = meta_dir / "embodiment.json" + if embodiment_path.exists() and not args.force: + log.info(" embodiment.json already exists, skipping") + else: + with open(embodiment_path, "w") as f: + json.dump(embodiment, f, indent=4) + log.info(" Wrote embodiment.json (tag=%s)", args.embodiment_tag) + + # 4. Get parquet file paths + parquet_paths = get_parquet_paths(output_path, info) + if not parquet_paths: + log.error("No parquet files found. Check dataset structure.") + sys.exit(1) + log.info(" Found %d parquet files", len(parquet_paths)) + + # 5. Compute stats.json + stats_path = meta_dir / "stats.json" + numeric_cols = detected["state"] + detected["action"] + if "timestamp" in info.get("features", {}): + numeric_cols.append("timestamp") + + if stats_path.exists() and not args.force: + log.info(" stats.json already exists, skipping") + else: + log.info(" Computing dataset statistics...") + stats = compute_stats(parquet_paths, numeric_cols) + with open(stats_path, "w") as f: + json.dump(stats, f, indent=4) + log.info(" Wrote stats.json (%d features)", len(stats)) + + # 6. Compute relative_stats_dreamzero.json + rel_stats_path = meta_dir / "relative_stats_dreamzero.json" + if args.relative_action_keys: + if rel_stats_path.exists() and not args.force: + log.info(" relative_stats_dreamzero.json already exists, skipping") + else: + log.info(" Computing relative action statistics for keys: %s", args.relative_action_keys) + rel_stats = compute_relative_stats( + parquet_paths, modality, args.relative_action_keys, + action_horizon=args.action_horizon, + ) + if rel_stats: + with open(rel_stats_path, "w") as f: + json.dump(rel_stats, f, indent=4) + log.info(" Wrote relative_stats_dreamzero.json (%d keys)", len(rel_stats)) + else: + log.warning(" No relative stats computed (check key names match between state and action)") + else: + log.info(" Skipping relative stats (no --relative-action-keys provided)") + + # 7. Build tasks.jsonl + tasks_path = meta_dir / "tasks.jsonl" + if tasks_path.exists() and not args.force: + log.info(" tasks.jsonl already exists, skipping") + else: + tasks = build_tasks(parquet_paths, task_key) + with open(tasks_path, "w") as f: + for t in tasks: + f.write(json.dumps(t) + "\n") + log.info(" Wrote tasks.jsonl (%d tasks)", len(tasks)) + + # 8. Build episodes.jsonl + episodes_path = meta_dir / "episodes.jsonl" + if episodes_path.exists() and not args.force: + log.info(" episodes.jsonl already exists, skipping") + else: + tasks = [] + if tasks_path.exists(): + with open(tasks_path) as f: + for line in f: + tasks.append(json.loads(line.strip())) + if not tasks: + tasks = [{"task_index": 0, "task": ""}] + episodes = build_episodes(parquet_paths, info, task_key, tasks) + with open(episodes_path, "w") as f: + for ep in episodes: + f.write(json.dumps(ep) + "\n") + log.info(" Wrote episodes.jsonl (%d episodes)", len(episodes)) + + # 9. Validation + warnings = validate_dataset(output_path, info, modality) + if warnings: + log.warning("Validation warnings:") + for w in warnings: + log.warning(" - %s", w) + else: + log.info("Validation passed -- no warnings") + + # Summary + print("\n" + "=" * 60) + print("Conversion complete!") + print(f" Output: {output_path}") + print(f" Embodiment tag: {args.embodiment_tag}") + print(f" State keys: {list(modality['state'].keys())}") + print(f" Action keys: {list(modality['action'].keys())}") + print(f" Video keys: {list(modality['video'].keys())}") + print(f" Task key: {task_key or '(none)'}") + if args.relative_action_keys: + print(f" Relative action keys: {args.relative_action_keys}") + print("=" * 60) + print("\nNext steps:") + print(" 1. Create a YAML data config in groot/vla/configs/data/dreamzero/") + print(" 2. Add modality configs to base_48_wan_fine_aug_relative.yaml") + print(" 3. Create a training script in scripts/train/") + print(" See docs/CUSTOM_EMBODIMENT_TRAINING.md for the full guide.") + + +if __name__ == "__main__": + main() diff --git a/scripts/data/convert_libero_to_hf.py b/scripts/data/convert_libero_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..92967e6317bbe552740ba9faa6d0f83ea732d117 --- /dev/null +++ b/scripts/data/convert_libero_to_hf.py @@ -0,0 +1,281 @@ +#!/usr/bin/env python3 +"""将 LIBERO（LeRobot parquet 格式）转换为 HuggingFace datasets 格式。 + +用法: + python scripts/data/convert_libero_to_hf.py \\ + --input /path/to/libero_data \\ + --output /path/to/hf_libero \\ + --num-workers 8 + +输出结构: + hf_libero/ + ├── data/ + │ ├── train-00000.parquet # {episode_idx, frame_idx, state, action, action_mask, text, video_path} + │ └── ... + ├── videos/ + │ └── view_0/ + │ ├── episode_000000.mp4 + │ └── ... + ├── dataset_info.json + └── README.md +""" + +import os, json, argparse, logging, multiprocessing as mp +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import cv2 +from tqdm import tqdm + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger(__name__) + +# DreamZero 标准参数 +NUM_FRAMES = 12 +ACTION_HORIZON = 12 +MAX_STATE_DIM = 44 +MAX_ACTION_DIM = 32 +VIDEO_HEIGHT = 160 +VIDEO_WIDTH = 320 +FPS = 30 + + +def pad_state(state: np.ndarray, max_dim: int = MAX_STATE_DIM) -> np.ndarray: + """Pad state vector to max_dim with zeros.""" + d = state.shape[-1] + if d >= max_dim: + return state[..., :max_dim].astype(np.float32) + padded = np.zeros((max_dim,), dtype=np.float32) + padded[:d] = state.astype(np.float32) + return padded + + +def pad_action(action: np.ndarray, max_dim: int = MAX_ACTION_DIM) -> np.ndarray: + """Pad action vector to max_dim.""" + d = action.shape[-1] + if d >= max_dim: + return action[..., :max_dim].astype(np.float32) + padded = np.zeros((*action.shape[:-1], max_dim), dtype=np.float32) + padded[..., :d] = action.astype(np.float32) + return padded + + +def process_episode(args): + """处理单个 episode: 读取视频帧、state、写入 HF 格式。""" + ep_idx, parquet_path, video_dir, output_dir = args + + try: + df = pd.read_parquet(parquet_path) + except Exception as e: + logger.warning(f"无法读取 {parquet_path}: {e}") + return None + + output_video_dir = output_dir / "videos" / "view_0" + output_video_dir.mkdir(parents=True, exist_ok=True) + + # 检测 episode 边界 + episode_indices = [] + if "episode_index" in df.columns: + episode_indices = df["episode_index"].unique() + else: + episode_indices = [0] + + records = [] + for local_ep_idx, ep_val in enumerate(episode_indices): + ep_mask = df["episode_index"] == ep_val if "episode_index" in df.columns else slice(None) + ep_df = df[ep_mask].reset_index(drop=True) + n_frames = len(ep_df) + if n_frames < NUM_FRAMES + ACTION_HORIZON: + continue + + global_ep_idx = ep_idx * 1000 + local_ep_idx + + # 提取视频帧并编码为 mp4 + frames = [] + for i in range(n_frames): + row = ep_df.iloc[i] + # LeRobot 格式: 图像存为 {'bytes': b'...'} dict + raw_img = row.get("image", row.get("observation.image", None)) + if isinstance(raw_img, dict) and "bytes" in raw_img: + img_bytes = raw_img["bytes"] + elif isinstance(raw_img, bytes): + img_bytes = raw_img + else: + continue + + img_array = np.frombuffer(img_bytes, dtype=np.uint8) + img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) + if img is None: + continue + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + img = cv2.resize(img, (VIDEO_WIDTH, VIDEO_HEIGHT)) + frames.append(img) + + if len(frames) < NUM_FRAMES + ACTION_HORIZON: + continue + + # 写入 mp4 + video_filename = f"episode_{global_ep_idx:06d}.mp4" + video_path = output_video_dir / video_filename + out_writer = cv2.VideoWriter( + str(video_path), + cv2.VideoWriter_fourcc(*"mp4v"), + FPS, + (VIDEO_WIDTH, VIDEO_HEIGHT), + ) + for f in frames: + out_writer.write(cv2.cvtColor(f, cv2.COLOR_RGB2BGR)) + out_writer.release() + + # 提取 state/action + for t in range(n_frames - ACTION_HORIZON): + state = pad_state(ep_df.iloc[t].get("state", ep_df.iloc[t].get("observation.state", np.zeros(14))).astype(np.float32)) + action_chunk = [] + for a in range(ACTION_HORIZON): + act = ep_df.iloc[t + a].get("action", ep_df.iloc[t + a].get("action.joint_position", np.zeros(14))).astype(np.float32) + action_chunk.append(act) + action_chunk = np.stack(action_chunk) # [horizon, D] + action_chunk = pad_action(action_chunk) + + records.append({ + "episode_index": global_ep_idx, + "frame_index": t, + "state": state.tolist(), + "action": action_chunk.tolist(), + "action_mask": [True] * ACTION_HORIZON, + "text": ep_df.iloc[t].get( + "task_description", + ep_df.iloc[t].get("annotation.human.action.task_description", + "Perform the task")), + "video_path": f"videos/view_0/{video_filename}", + }) + + return records + + +def main(): + parser = argparse.ArgumentParser(description="Convert LIBERO to HF datasets format") + parser.add_argument("--input", "-i", required=True, help="LIBERO 数据目录 (parquet)") + parser.add_argument("--output", "-o", required=True, help="HF 数据集输出目录") + parser.add_argument("--num-workers", "-w", type=int, default=4, help="并行工作进程数") + args = parser.parse_args() + + input_dir = Path(args.input) + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + # 扫描 parquet 文件 + parquet_files = sorted(input_dir.glob("file-*.parquet")) + if not parquet_files: + # 尝试其他命名模式 + parquet_files = sorted(input_dir.glob("*.parquet")) + logger.info(f"找到 {len(parquet_files)} 个 parquet 文件") + + # 并行处理 + video_dir = input_dir / "videos" + tasks = [(i, str(pf), str(video_dir), output_dir) for i, pf in enumerate(parquet_files)] + + all_records = [] + with mp.Pool(args.num_workers) as pool: + for result in tqdm( + pool.imap_unordered(process_episode, tasks), + total=len(tasks), + desc="Converting episodes", + ): + if result: + all_records.extend(result) + + logger.info(f"生成 {len(all_records)} 条训练样本") + + if not all_records: + logger.error("未生成任何样本！") + return + + # 写入 parquet + df = pd.DataFrame(all_records) + data_dir = output_dir / "data" + data_dir.mkdir(parents=True, exist_ok=True) + + table = pa.Table.from_pandas(df) + pq.write_table(table, data_dir / "train-00000.parquet") + + # 写入 dataset_info.json + dataset_info = { + "description": "LIBERO benchmark dataset for DreamZero", + "features": { + "episode_index": {"dtype": "int64", "_type": "Value"}, + "frame_index": {"dtype": "int64", "_type": "Value"}, + "state": {"dtype": "float32", "shape": [MAX_STATE_DIM], "_type": "Sequence"}, + "action": {"dtype": "float32", "shape": [ACTION_HORIZON, MAX_ACTION_DIM], "_type": "Sequence"}, + "action_mask": {"dtype": "bool", "shape": [ACTION_HORIZON], "_type": "Sequence"}, + "text": {"dtype": "string", "_type": "Value"}, + "video_path": {"dtype": "string", "_type": "Value"}, + }, + "splits": {"train": {"num_examples": len(all_records)}}, + "homepage": "https://huggingface.co/datasets/dreamzero/libero", + } + with open(output_dir / "dataset_info.json", "w") as f: + json.dump(dataset_info, f, indent=2) + + # 写入 README.md (dataset card) + readme_template = """--- +license: cc-by-4.0 +task_categories: +- robotics +tags: +- robot-vla +- flow-matching +- libero +--- + +# DreamZero - LIBERO + +## Description +LIBERO benchmark dataset converted to HuggingFace datasets format for DreamZero training. + +## Schema +| Column | Type | Shape | Description | +|--------|------|-------|-------------| +| video_path | string | - | Path to video file | +| state | float32 | [{STATE_DIM}] | Robot state (padded) | +| action | float32 | [{ACT_HORIZON}, {ACT_DIM}] | Action chunks (padded) | +| action_mask | bool | [{ACT_HORIZON}] | Valid action mask | +| text | string | - | Task instruction | +| episode_index | int64 | - | Episode ID | + +## Statistics +- Total samples: {TOTAL_SAMPLES} +- Views: 1 +- Video resolution: {W}x{H} +- Frames per sample: {FRAMES} + +## Citation +``` +@inproceedings{libero2023, + title={LIBERO: Benchmarking Knowledge Transfer in Lifelong Robot Learning}, + ... +} +``` +""" + readme = readme_template.format( + STATE_DIM=MAX_STATE_DIM, + ACT_HORIZON=ACTION_HORIZON, + ACT_DIM=MAX_ACTION_DIM, + TOTAL_SAMPLES=len(all_records), + W=VIDEO_WIDTH, + H=VIDEO_HEIGHT, + FRAMES=NUM_FRAMES, + ) + with open(output_dir / "README.md", "w") as f: + f.write(readme) + + logger.info(f"转换完成！输出目录: {output_dir}") + logger.info(f" 数据: {data_dir}/train-00000.parquet") + logger.info(f" 视频: {output_dir}/videos/view_0/") + + +if __name__ == "__main__": + main() diff --git a/scripts/data/convert_libero_v3_to_gear.py b/scripts/data/convert_libero_v3_to_gear.py new file mode 100644 index 0000000000000000000000000000000000000000..391b7310397997d3748d7ca5e3c52f5c8139c8d2 --- /dev/null +++ b/scripts/data/convert_libero_v3_to_gear.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +""" +Convert LeRobot v3 LIBERO (chunked parquet + combined MP4) to GEAR format (per-episode). + +Usage: + python3 convert_libero_v3_to_gear.py \ + --input /root/autodl-tmp/data/libero \ + --output /root/autodl-tmp/data/libero_gear \ + --workers 4 +""" +import argparse, json, os, subprocess, sys, time, logging +from pathlib import Path +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +from concurrent.futures import ProcessPoolExecutor, as_completed + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +log = logging.getLogger("convert_libero_v3") + +SUITES = ["libero_spatial", "libero_object", "libero_goal", "libero_10", "libero_90"] +VIDEO_KEYS = ["observation.images.image", "observation.images.wrist_image"] + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--input", required=True) + p.add_argument("--output", required=True) + p.add_argument("--workers", type=int, default=4) + return p.parse_args() + +def process_suite(suite_name, input_dir, output_dir): + """Convert one suite (e.g. libero_spatial) to GEAR format.""" + suite_in = input_dir / suite_name + suite_out = output_dir / suite_name + + # Read info + with open(suite_in / "meta" / "info.json") as f: + info = json.load(f) + + # Read tasks + tasks_df = pd.read_parquet(suite_in / "meta" / "tasks.parquet") + tasks = {row["task_index"]: row["task_index"] for _, row in tasks_df.iterrows()} + # tasks.parquet in LeRobot v3 uses index as task name + tasks_parquet = suite_in / "meta" / "tasks.parquet" + tdf = pd.read_parquet(tasks_parquet) + task_names = {} + for idx, row in tdf.iterrows(): + task_names[int(row["task_index"])] = str(idx) + + # Find parquet files + data_dir = suite_in / "data" / "chunk-000" + parquet_files = sorted(data_dir.glob("file-*.parquet")) + if not parquet_files: + log.warning(f"No parquet files in {data_dir}, skipping {suite_name}") + return 0 + + # Find video files + video_paths = {} + for vk in VIDEO_KEYS: + vp = suite_in / "videos" / vk / "chunk-000" / "file-000.mp4" + if vp.exists(): + video_paths[vk] = vp + else: + log.warning(f"Video not found: {vp}") + + # Read episodes metadata + ep_meta_dir = suite_in / "meta" / "episodes" / "chunk-000" + ep_meta_files = sorted(ep_meta_dir.glob("*.parquet")) + ep_meta = pd.read_parquet(ep_meta_files[0]) if ep_meta_files else None + + # Process each parquet file - build episode index + log.info(f"Processing {suite_name}: {len(parquet_files)} parquet files, {len(video_paths)} video streams") + + ep_data_dir = suite_out / "data" / "chunk-000" + ep_data_dir.mkdir(parents=True, exist_ok=True) + + total_episodes = 0 + total_frames = 0 + ep_entries = [] + + for pf in parquet_files: + table = pq.read_table(pf) + df = table.to_pandas() + + # Group by episode_index + for ep_idx, ep_df in df.groupby("episode_index"): + ep_df = ep_df.sort_values("frame_index").reset_index(drop=True) + n_frames = len(ep_df) + + # Get task index + task_idx = int(ep_df["task_index"].iloc[0]) + task_name = task_names.get(task_idx, f"task_{task_idx}") + + # Save per-episode parquet + ep_parquet = ep_data_dir / f"episode_{total_episodes:06d}.parquet" + ep_parquet_table = pa.Table.from_pandas(ep_df[["observation.state", "action", "timestamp", "frame_index", "index"]]) + pq.write_table(ep_parquet_table, ep_parquet) + + # Get video timestamp from episodes metadata + from_ts = 0.0 + to_ts = float(n_frames) * 0.05 # 20fps = 0.05s per frame + if ep_meta is not None: + match = ep_meta[ep_meta["episode_index"] == ep_idx] + if len(match) > 0: + for vk in VIDEO_KEYS: + ft_col = f"videos/{vk}/from_timestamp" + tt_col = f"videos/{vk}/to_timestamp" + if ft_col in match.columns: + from_ts = float(match[ft_col].iloc[0]) + to_ts = float(match[tt_col].iloc[0]) + break + + # Extract video segments + for vk, vp in video_paths.items(): + vk_short = vk.split(".")[-1] # "image" or "wrist_image" + out_vid_dir = suite_out / "videos" / vk / "chunk-000" + out_vid_dir.mkdir(parents=True, exist_ok=True) + out_mp4 = out_vid_dir / f"episode_{total_episodes:06d}.mp4" + + # Use ffmpeg to extract segment (copy codec, no re-encode) + cmd = [ + "ffmpeg", "-y", "-ss", str(from_ts), "-to", str(to_ts), + "-i", str(vp), "-c", "copy", "-avoid_negative_ts", "make_zero", + str(out_mp4) + ] + # Try keyframe-accurate extraction first + result = subprocess.run(cmd, capture_output=True, text=True, timeout=120) + if result.returncode != 0: + log.warning(f"ffmpeg failed for {suite_name} ep {total_episodes}: {result.stderr[:200]}") + # Fallback: copy whole video (will be trimmed in dataloader) + # Actually just re-extract with re-encode + cmd2 = [ + "ffmpeg", "-y", "-ss", str(from_ts), "-to", str(to_ts), + "-i", str(vp), "-c:v", "libx264", "-preset", "fast", + "-crf", "23", "-pix_fmt", "yuv420p", + str(out_mp4) + ] + subprocess.run(cmd2, capture_output=True, text=True, timeout=300) + + ep_entries.append({ + "episode_index": total_episodes, + "task_index": task_idx, + "task": task_name, + "length": n_frames, + "from_timestamp": from_ts, + "to_timestamp": to_ts, + }) + + total_episodes += 1 + total_frames += n_frames + + if total_episodes % 200 == 0: + log.info(f" {suite_name}: {total_episodes} episodes done") + + # Write metadata files for GEAR + meta_dir = suite_out / "meta" + meta_dir.mkdir(parents=True, exist_ok=True) + + # episodes.jsonl + with open(meta_dir / "episodes.jsonl", "w") as f: + for ep in ep_entries: + f.write(json.dumps(ep) + "\n") + + # tasks.jsonl + seen_tasks = {} + for ep in ep_entries: + ti = ep["task_index"] + if ti not in seen_tasks: + seen_tasks[ti] = {"task_index": ti, "task": ep["task"]} + with open(meta_dir / "tasks.jsonl", "w") as f: + for t in sorted(seen_tasks.values(), key=lambda x: x["task_index"]): + f.write(json.dumps(t) + "\n") + + # info.json + info_out = { + "total_episodes": total_episodes, + "total_frames": total_frames, + "total_tasks": len(seen_tasks), + "fps": 20, + "chunks_size": 2000, + "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + "video_path": "videos/{video_key}/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.mp4", + "features": info.get("features", {}), + } + with open(meta_dir / "info.json", "w") as f: + json.dump(info_out, f, indent=2) + + # modality.json + modality = { + "state": {"state": {"original_key": "observation.state"}}, + "action": {"action": {"original_key": "action"}}, + "video": {}, + "annotation": {"task": {"original_key": "annotation.task"}}, + } + for vk in video_paths: + short = vk.split(".")[-1] + modality["video"][short] = {"original_key": vk} + with open(meta_dir / "modality.json", "w") as f: + json.dump(modality, f, indent=2) + + # embodiment.json + embodiment = {"embodiment_tag": "libero_sim"} + with open(meta_dir / "embodiment.json", "w") as f: + json.dump(embodiment, f, indent=2) + + log.info(f"{suite_name}: {total_episodes} episodes, {total_frames} frames, {len(seen_tasks)} tasks") + return total_episodes + +def main(): + args = parse_args() + input_dir = Path(args.input) + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + total = 0 + for suite in SUITES: + t0 = time.time() + n = process_suite(suite, input_dir, output_dir) + dt = time.time() - t0 + log.info(f"{suite}: {n} episodes in {dt:.0f}s ({n/dt:.1f} ep/s)") + total += n + + log.info(f"Total: {total} episodes across {len(SUITES)} suites") + log.info(f"Output: {output_dir}") + +if __name__ == "__main__": + main() diff --git a/scripts/data/convert_manifeel_to_hf.py b/scripts/data/convert_manifeel_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..34e6923e5bcbdca9cb0d4a436f1dd0a66b7762ce --- /dev/null +++ b/scripts/data/convert_manifeel_to_hf.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +"""将 ManiFeel (zarr 格式) 转换为 HuggingFace datasets 格式。 + +用法: + python scripts/data/convert_manifeel_to_hf.py \\ + --input /path/to/manifeel_zarr \\ + --output /path/to/hf_manifeel \\ + --num-workers 8 + +ManiFeel zarr 结构: + / + ├── front/ # zarr 数组 [T, H, W, C] + ├── wrist/ + ├── side/ + ├── state/ # [T, 7] + ├── action/ # [T, 6] + └── meta/ + └── episode_ends # episode 边界索引 +""" + +import os, json, argparse, logging +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import zarr +import cv2 +from tqdm import tqdm + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger(__name__) + +# DreamZero 标准参数 +ACTION_HORIZON = 12 +MAX_STATE_DIM = 44 +MAX_ACTION_DIM = 32 +VIDEO_HEIGHT = 160 +VIDEO_WIDTH = 320 +FPS = 30 + + +def pad_state(state: np.ndarray) -> np.ndarray: + 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: + """action_chunk: [horizon, D]""" + 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 + + +def process_task(args): + """处理一个 ManiFeel 任务目录，返回训练样本列表。""" + task_dir_str, output_dir = args + task_dir = Path(task_dir_str) + task_name = task_dir.name + + try: + # 读取 zarr + front = zarr.open(str(task_dir / "front"), mode="r") + wrist = zarr.open(str(task_dir / "wrist"), mode="r") + side = zarr.open(str(task_dir / "side"), mode="r") + state_arr = zarr.open(str(task_dir / "state"), mode="r") + action_arr = zarr.open(str(task_dir / "action"), mode="r") + episode_ends = zarr.open(str(task_dir / "meta" / "episode_ends"), mode="r") + except Exception as e: + logger.warning(f"无法读取 {task_dir}: {e}") + return [] + + # zarr 读取为 numpy + front_np = np.array(front) # [T, H, W, C] + wrist_np = np.array(wrist) + side_np = np.array(side) + state_np = np.array(state_arr) # [T, D_state] + action_np = np.array(action_arr) # [T, D_action] + ends = np.array(episode_ends) # episode 边界索引 + + records = [] + video_out_dir = output_dir / "videos" + # 3 views + for view_idx, view_name in enumerate(["front", "wrist", "side"]): + (video_out_dir / f"view_{view_idx}").mkdir(parents=True, exist_ok=True) + + prev_end = 0 + for ep_idx, end in enumerate(ends): + start = prev_end + prev_end = end + ep_len = end - start + if ep_len < ACTION_HORIZON + 1: + continue + + # 为每个视角编码 mp4 + video_paths = [] + for view_idx, view_data in enumerate([front_np, wrist_np, side_np]): + ep_frames = view_data[start:end] # [ep_len, H, W, C] + video_filename = f"episode_{task_name}_{ep_idx:04d}_view{view_idx}.mp4" + video_path = video_out_dir / f"view_{view_idx}" / video_filename + + # resize + encode + h, w = ep_frames.shape[1], ep_frames.shape[2] + writer = cv2.VideoWriter( + str(video_path), + cv2.VideoWriter_fourcc(*"mp4v"), + FPS, (VIDEO_WIDTH, VIDEO_HEIGHT), + ) + for f_idx in range(ep_len): + frame = ep_frames[f_idx] + if h != VIDEO_HEIGHT or w != VIDEO_WIDTH: + frame = cv2.resize(frame, (VIDEO_WIDTH, VIDEO_HEIGHT)) + if frame.shape[-1] == 3: + frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) + writer.write(frame) + writer.release() + video_paths.append(f"videos/view_{view_idx}/{video_filename}") + + # 为每个时间窗口生成样本 + ep_state = state_np[start:end] + ep_action = action_np[start:end] + + for t in range(ep_len - ACTION_HORIZON): + state = pad_state(ep_state[t]) + action_chunk = ep_action[t:t + ACTION_HORIZON] # [horizon, D] + action_chunk = pad_action(action_chunk) + + # ManiFeel 所有视角拼接到一条样本，以 dict 形式存储 + records.append({ + "task": task_name, + "episode_index": ep_idx, + "frame_index": t, + "state": state.tolist(), + "action": action_chunk.tolist(), + "action_mask": [True] * ACTION_HORIZON, + "text": task_name.replace("_", " ").replace("-", " "), + "video_path_0": video_paths[0], + "video_path_1": video_paths[1], + "video_path_2": video_paths[2], + }) + + logger.info(f" {task_name}: {len(records)} samples") + return records + + +def main(): + parser = argparse.ArgumentParser(description="Convert ManiFeel zarr to HF datasets format") + parser.add_argument("--input", "-i", required=True, help="ManiFeel 数据根目录（包含任务子目录）") + parser.add_argument("--output", "-o", required=True, help="HF 数据集输出目录") + parser.add_argument("--num-workers", "-w", type=int, default=4, help="并行任务数") + args = parser.parse_args() + + input_dir = Path(args.input) + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + # 扫描任务目录 + task_dirs = sorted([d for d in input_dir.iterdir() if d.is_dir()]) + logger.info(f"找到 {len(task_dirs)} 个任务目录") + + # 处理每个任务 + all_records = [] + for task_dir in tqdm(task_dirs, desc="Processing tasks"): + records = process_task((str(task_dir), output_dir)) + all_records.extend(records) + + logger.info(f"总共生成 {len(all_records)} 条训练样本") + + if not all_records: + logger.error("未生成任何样本！") + return + + # 写入 parquet + df = pd.DataFrame(all_records) + data_dir = output_dir / "data" + data_dir.mkdir(parents=True, exist_ok=True) + + table = pa.Table.from_pandas(df) + pq.write_table(table, data_dir / "train-00000.parquet") + + # dataset_info.json + dataset_info = { + "description": "ManiFeel benchmark dataset for DreamZero", + "features": { + "task": {"dtype": "string", "_type": "Value"}, + "episode_index": {"dtype": "int64", "_type": "Value"}, + "frame_index": {"dtype": "int64", "_type": "Value"}, + "state": {"dtype": "float32", "shape": [MAX_STATE_DIM], "_type": "Sequence"}, + "action": {"dtype": "float32", "shape": [ACTION_HORIZON, MAX_ACTION_DIM], "_type": "Sequence"}, + "action_mask": {"dtype": "bool", "shape": [ACTION_HORIZON], "_type": "Sequence"}, + "text": {"dtype": "string", "_type": "Value"}, + "video_path_0": {"dtype": "string", "_type": "Value"}, + "video_path_1": {"dtype": "string", "_type": "Value"}, + "video_path_2": {"dtype": "string", "_type": "Value"}, + }, + "num_views": 3, + "splits": {"train": {"num_examples": len(all_records)}}, + } + with open(output_dir / "dataset_info.json", "w") as f: + json.dump(dataset_info, f, indent=2) + + # README + readme = f"""--- +license: cc-by-4.0 +--- + +# DreamZero - ManiFeel + +## Description +ManiFeel benchmark dataset (3 views: front, wrist, side) converted for DreamZero. + +## Schema +| Column | Type | Description | +|--------|------|-------------| +| video_path_0/1/2 | string | Front/wrist/side video files | +| state | float32[{MAX_STATE_DIM}] | Robot state (padded) | +| action | float32[{ACTION_HORIZON},{MAX_ACTION_DIM}] | Action chunks (padded) | +| text | string | Task description | +| episode_index | int64 | Episode ID | + +## Statistics +- Total samples: {len(all_records)} +- Views: 3 (front, wrist, side) +- Video resolution: {VIDEO_WIDTH}x{VIDEO_HEIGHT} +""" + with open(output_dir / "README.md", "w") as f: + f.write(readme) + + logger.info(f"转换完成！输出: {output_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/data/convert_robotwin_to_hf.py b/scripts/data/convert_robotwin_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..7a1a17880e8c070755d3dbc54d04d78c52a37dd1 --- /dev/null +++ b/scripts/data/convert_robotwin_to_hf.py @@ -0,0 +1,258 @@ +#!/usr/bin/env python3 +"""将 RoboTwin (qpos.pt + video.mp4 + metas/*.txt) 转换为 HuggingFace datasets 格式。 + +用法: + python scripts/data/convert_robotwin_to_hf.py \\ + --input /path/to/robotwin_data \\ + --output /path/to/hf_robotwin \\ + --num-workers 8 + +RoboTwin 目录结构: + / + ├── qpos/ + │ ├── episode0.pt # [T_state, 14] float32 + │ └── episode1.pt + ├── videos/ + │ ├── episode0.mp4 # [T_video, H, W, C] + │ └── episode1.mp4 + └── metas/ + ├── task_0.txt # 任务描述 + └── task_1.txt +""" + +import os, json, argparse, logging +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import torch +import cv2 +from tqdm import tqdm + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger(__name__) + +# DreamZero 标准参数 +ACTION_HORIZON = 12 +MAX_STATE_DIM = 44 +MAX_ACTION_DIM = 32 +VIDEO_HEIGHT = 160 +VIDEO_WIDTH = 320 +FPS = 30 + + +def pad_state(state: np.ndarray) -> np.ndarray: + 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: + 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 + + +def process_episode(args): + """处理一个 RoboTwin episode。""" + task_dir_str, ep_name, output_dir, ep_idx = args + task_dir = Path(task_dir_str) + task_name = task_dir.name + video_out_dir = output_dir / "videos" / "view_0" + video_out_dir.mkdir(parents=True, exist_ok=True) + + try: + # 读取 qpos + qpos_path = task_dir / "qpos" / f"{ep_name}.pt" + qpos_data = torch.load(str(qpos_path)).numpy() # [T_state, 14] + + # 读取视频 + video_path = task_dir / "videos" / f"{ep_name}.mp4" + cap = cv2.VideoCapture(str(video_path)) + n_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + cap.release() + + # 读取任务描述 + meta_dir = task_dir / "metas" + task_desc = "Manipulate the object on the table" + if meta_dir.exists(): + meta_files = sorted(meta_dir.glob("*.txt")) + ep_num = int(ep_name.replace("episode", "")) + if ep_num < len(meta_files): + with open(meta_files[ep_num]) as f: + task_desc = f.read().strip() + except Exception as e: + logger.warning(f" 无法读取 {task_name}/{ep_name}: {e}") + return [] + + if n_video_frames < ACTION_HORIZON + 1 or len(qpos_data) < ACTION_HORIZON + 1: + return [] + + # 视频帧数 vs state 帧数对齐 + video_to_qpos_ratio = len(qpos_data) / max(n_video_frames, 1) + + records = [] + + # 读取所有视频帧到内存 + cap = cv2.VideoCapture(str(video_path)) + all_frames = [] + while True: + ret, frame = cap.read() + if not ret: + break + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame = cv2.resize(frame, (VIDEO_WIDTH, VIDEO_HEIGHT)) + all_frames.append(frame) + cap.release() + + if len(all_frames) < ACTION_HORIZON + 1: + return [] + + # 将整个 episode 编码为一个 mp4 文件 + video_filename = f"episode_{task_name}_{ep_idx:06d}.mp4" + video_path_out = video_out_dir / video_filename + writer = cv2.VideoWriter( + str(video_path_out), + cv2.VideoWriter_fourcc(*"mp4v"), + FPS, (VIDEO_WIDTH, VIDEO_HEIGHT), + ) + for frame in all_frames: + writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) + writer.release() + + # 为每个时间窗口生成样本 + for t in range(len(all_frames) - ACTION_HORIZON): + # 对齐 qpos 索引 + qpos_t = min(int(t * video_to_qpos_ratio), len(qpos_data) - 1) + + state = pad_state(qpos_data[min(qpos_t, len(qpos_data) - 1)]) + + # 从 qpos 差值计算 action (velocity) + action_chunk = np.zeros((ACTION_HORIZON, 14), dtype=np.float32) + for a in range(ACTION_HORIZON): + fi = min(qpos_t + a + 1, len(qpos_data) - 1) + si = min(qpos_t + a, len(qpos_data) - 1) + action_chunk[a] = qpos_data[fi] - qpos_data[si] + + action_chunk = pad_action(action_chunk) + + records.append({ + "task": task_name, + "episode_index": ep_idx, + "frame_index": t, + "state": state.tolist(), + "action": action_chunk.tolist(), + "action_mask": [True] * ACTION_HORIZON, + "text": task_desc, + "video_path": f"videos/view_0/{video_filename}", + }) + + logger.info(f" {task_name}/{ep_name}: {len(records)} samples") + return records + + +def main(): + parser = argparse.ArgumentParser(description="Convert RoboTwin to HF datasets format") + parser.add_argument("--input", "-i", required=True, help="RoboTwin 数据根目录（包含任务子目录）") + parser.add_argument("--output", "-o", required=True, help="HF 数据集输出目录") + parser.add_argument("--num-workers", "-w", type=int, default=4, help="并行进程数") + args = parser.parse_args() + + input_dir = Path(args.input) + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + # 扫描任务目录 + task_dirs = sorted([d for d in input_dir.iterdir() if d.is_dir()]) + logger.info(f"找到 {len(task_dirs)} 个任务目录") + + all_records = [] + global_ep_idx = 0 + + for task_dir in tqdm(task_dirs, desc="Processing tasks"): + video_files = sorted((task_dir / "videos").glob("episode*.mp4")) + + for vf in video_files: + ep_name = vf.stem + qf = task_dir / "qpos" / f"{ep_name}.pt" + if not qf.exists(): + continue + + records = process_episode((str(task_dir), ep_name, output_dir, global_ep_idx)) + all_records.extend(records) + global_ep_idx += 1 + + logger.info(f"总共生成 {len(all_records)} 条训练样本") + + if not all_records: + logger.error("未生成任何样本！") + return + + # 写入 parquet + df = pd.DataFrame(all_records) + data_dir = output_dir / "data" + data_dir.mkdir(parents=True, exist_ok=True) + + table = pa.Table.from_pandas(df) + # 分片（每 10 万条一个文件） + num_shards = max(1, len(all_records) // 100_000) + if num_shards > 1: + shard_size = len(all_records) // num_shards + for i in range(num_shards): + start = i * shard_size + end = start + shard_size if i < num_shards - 1 else len(all_records) + shard_table = table.slice(start, end - start) + pq.write_table(shard_table, data_dir / f"train-{i:05d}.parquet") + else: + pq.write_table(table, data_dir / "train-00000.parquet") + + dataset_info = { + "description": "RoboTwin benchmark dataset for DreamZero", + "features": { + "task": {"dtype": "string", "_type": "Value"}, + "episode_index": {"dtype": "int64", "_type": "Value"}, + "frame_index": {"dtype": "int64", "_type": "Value"}, + "state": {"dtype": "float32", "shape": [MAX_STATE_DIM], "_type": "Sequence"}, + "action": {"dtype": "float32", "shape": [ACTION_HORIZON, MAX_ACTION_DIM], "_type": "Sequence"}, + "action_mask": {"dtype": "bool", "shape": [ACTION_HORIZON], "_type": "Sequence"}, + "text": {"dtype": "string", "_type": "Value"}, + "video_path": {"dtype": "string", "_type": "Value"}, + }, + "splits": {"train": {"num_examples": len(all_records)}}, + } + with open(output_dir / "dataset_info.json", "w") as f: + json.dump(dataset_info, f, indent=2) + + readme = f"""--- +license: cc-by-4.0 +--- + +# DreamZero - RoboTwin + +## Description +RoboTwin benchmark dataset converted for DreamZero training. + +## Schema +| Column | Type | Description | +|--------|------|-------------| +| video_path | string | Video file path | +| state | float32[{MAX_STATE_DIM}] | Robot state (padded) | +| action | float32[{ACTION_HORIZON},{MAX_ACTION_DIM}] | Action chunks (padded, velocity) | +| text | string | Task description | + +## Statistics +- Total samples: {len(all_records)} +- Views: 1 +""" + with open(output_dir / "README.md", "w") as f: + f.write(readme) + + logger.info(f"转换完成！输出: {output_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/data/download_checkpoints.sh b/scripts/data/download_checkpoints.sh new file mode 100644 index 0000000000000000000000000000000000000000..32839df369bfbfc7ce2843e14069adb0982b2983 --- /dev/null +++ b/scripts/data/download_checkpoints.sh @@ -0,0 +1,122 @@ +#!/bin/bash +# DreamZero 模型权重一键下载脚本 +# +# 用法: +# bash scripts/data/download_checkpoints.sh +# bash scripts/data/download_checkpoints.sh --mirror # 使用 HF 镜像（中国用户） +# bash scripts/data/download_checkpoints.sh --offline /path/to/cache # 从本地缓存加载 +# +# 下载内容: +# - Wan2.2-TI2V-5B: Wan2.2 视频 DiT 骨干 (~32G) +# - umt5-xxl: T5 tokenizer (~49G) +# - CLIP encoder: 图像编码器 (~2.4G) + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +CHECKPOINT_DIR="${CHECKPOINT_DIR:-$REPO_ROOT/checkpoints}" + +MIRROR=false +OFFLINE_CACHE="" + +# Parse arguments +while [[ $# -gt 0 ]]; do + case $1 in + --mirror) + MIRROR=true + shift + ;; + --offline) + OFFLINE_CACHE="$2" + shift 2 + ;; + *) + echo "Unknown option: $1" + echo "Usage: $0 [--mirror] [--offline /path/to/cache]" + exit 1 + ;; + esac +done + +# ============ 网络配置 ============ +if [ "$MIRROR" = true ]; then + export HF_ENDPOINT=https://hf-mirror.com + echo ">>> 使用 HF 镜像: hf-mirror.com" +fi + +# ============ 离线模式 ============ +if [ -n "$OFFLINE_CACHE" ]; then + if [ ! -d "$OFFLINE_CACHE" ]; then + echo "ERROR: 离线缓存目录不存在: $OFFLINE_CACHE" + exit 1 + fi + echo ">>> 离线模式，从 $OFFLINE_CACHE 复制..." + cp -r "$OFFLINE_CACHE"/* "$CHECKPOINT_DIR/" + echo ">>> 复制完成！" + exit 0 +fi + +# ============ 下载 Wan2.2-TI2V-5B ============ +WAN22_DIR="$CHECKPOINT_DIR/Wan2.2-TI2V-5B" +if [ -d "$WAN22_DIR" ] && [ -n "$(ls -A "$WAN22_DIR" 2>/dev/null)" ]; then + echo "[SKIP] Wan2.2-TI2V-5B 已存在: $WAN22_DIR" +else + echo "[DOWNLOAD] Wan2.2-TI2V-5B → $WAN22_DIR" + mkdir -p "$WAN22_DIR" + huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_DIR" + echo "[OK] Wan2.2-TI2V-5B 下载完成" +fi + +# ============ 下载 umt5-xxl ============ +TOKENIZER_DIR="$CHECKPOINT_DIR/umt5-xxl" +if [ -d "$TOKENIZER_DIR" ] && [ -n "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "[SKIP] umt5-xxl 已存在: $TOKENIZER_DIR" +else + echo "[DOWNLOAD] umt5-xxl → $TOKENIZER_DIR" + mkdir -p "$TOKENIZER_DIR" + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" + echo "[OK] umt5-xxl 下载完成" +fi + +# ============ 下载 CLIP encoder ============ +CLIP_DIR="$CHECKPOINT_DIR/clip-encoder" +CLIP_FILE="$CLIP_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" +if [ -f "$CLIP_FILE" ]; then + echo "[SKIP] CLIP encoder 已存在: $CLIP_FILE" +else + echo "[DOWNLOAD] CLIP encoder → $CLIP_DIR" + mkdir -p "$CLIP_DIR" + + # 尝试从 dreamzero 独立仓库下载 + if huggingface-cli download poet70/clip-encoder \ + --local-dir "$CLIP_DIR" 2>/dev/null; then + echo "[OK] CLIP encoder 下载完成" + else + echo "[WARN] 独立 CLIP 仓库暂不可用，尝试从 Wan2.1 提取..." + # 后备：下载 Wan2.1-I2V-14B-480P（包含 CLIP 权重） + WAN21_DIR="$CHECKPOINT_DIR/Wan2.1-I2V-14B-480P" + if [ ! -d "$WAN21_DIR" ]; then + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN21_DIR" + fi + # 复制 CLIP 权重 + if [ -f "$WAN21_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then + cp "$WAN21_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" "$CLIP_DIR/" + echo "[OK] CLIP encoder 从 Wan2.1 提取完成" + else + echo "[ERROR] 无法获取 CLIP encoder" + echo "请手动从 DreamZero-DROID 检查点提取:" + echo " cp /path/to/DreamZero-DROID/models_clip_*.pth $CLIP_DIR/" + exit 1 + fi + fi +fi + +echo "" +echo "============================================" +echo " 所有权重下载完成！" +echo "============================================" +echo "Wan2.2: $WAN22_DIR" +echo "umt5-xxl: $TOKENIZER_DIR" +echo "CLIP: $CLIP_DIR" +echo "============================================" diff --git a/scripts/data/download_droid_hf.py b/scripts/data/download_droid_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..0862e39d15acd3414afe183ef566951b9bb568f9 --- /dev/null +++ b/scripts/data/download_droid_hf.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +""" +Download DreamZero DROID dataset from Hugging Face with rate-limit handling. + +Use this script when `huggingface-cli download` hits 429 (Too Many Requests). +It uses a single worker and retries with backoff to stay within the 3000 req/5min limit. + +Usage: + python scripts/data/download_droid_hf.py [--local-dir ./data/droid_lerobot] [--max-workers 1] + + # Or with env (same as CLI default): + DROID_DATA_ROOT=./data/droid_lerobot python scripts/data/download_droid_hf.py +""" + +from __future__ import annotations + +import argparse +import os +import sys +import time + +try: + from huggingface_hub import snapshot_download +except ImportError: + print("Install huggingface_hub: pip install huggingface_hub", file=sys.stderr) + sys.exit(1) + +REPO_ID = "GEAR-Dreams/DreamZero-DROID-Data" +REPO_TYPE = "dataset" + + +def main() -> None: + p = argparse.ArgumentParser(description="Download DreamZero DROID dataset with rate-limit handling.") + p.add_argument( + "--local-dir", + default=os.environ.get("DROID_DATA_ROOT", "./data/droid_lerobot"), + help="Local directory to download into (default: DROID_DATA_ROOT or ./data/droid_lerobot)", + ) + p.add_argument( + "--max-workers", + type=int, + default=1, + help="Concurrent download threads (default 1 to reduce API requests and avoid 429)", + ) + p.add_argument( + "--retry-wait", + type=int, + default=320, + help="Seconds to wait on 429 before retry (default 320 ≈ 5min)", + ) + args = p.parse_args() + + local_dir = os.path.abspath(args.local_dir) + os.makedirs(local_dir, exist_ok=True) + + attempt = 0 + while True: + attempt += 1 + try: + print(f"Download attempt {attempt} (max_workers={args.max_workers})...") + snapshot_download( + repo_id=REPO_ID, + repo_type=REPO_TYPE, + local_dir=local_dir, + local_dir_use_symlinks=False, + max_workers=args.max_workers, + resume_download=True, + ) + print(f"Done. Dataset at: {local_dir}") + return + except Exception as e: + err_str = str(e).lower() + if "429" in err_str or "too many requests" in err_str or "rate limit" in err_str: + print(f"Rate limited (429). Waiting {args.retry_wait}s before retry...", file=sys.stderr) + time.sleep(args.retry_wait) + continue + raise + + +if __name__ == "__main__": + main() diff --git a/scripts/data/prepare_libero_gear.py b/scripts/data/prepare_libero_gear.py new file mode 100644 index 0000000000000000000000000000000000000000..f1cd020fb8d14740d3b65ea0673fc5adfe1c6f39 --- /dev/null +++ b/scripts/data/prepare_libero_gear.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +""" +Prepare LIBERO dataset for DreamZero GEAR pipeline. + +Converts chunked LeRobot v2 format (multiple episodes per parquet, images as PNG bytes) +into individual-episode format (one parquet + one mp4 per episode) expected by +convert_lerobot_to_gear.py. + +Usage: + python3 prepare_libero_gear.py \ + --input-dir /root/autodl-tmp/data/libero \ + --output-dir /root/autodl-tmp/data/libero_gear +""" + +import argparse +import json +import os +import io +import sys +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +from PIL import Image +import cv2 + + +def parse_args(): + parser = argparse.ArgumentParser(description="Prepare LIBERO data for DreamZero GEAR pipeline") + parser.add_argument("--input-dir", required=True, help="Path to original LIBERO dataset") + parser.add_argument("--output-dir", required=True, help="Path for output GEAR-ready dataset") + parser.add_argument("--num-workers", type=int, default=4, help="Number of parallel workers") + parser.add_argument("--skip-video", action="store_true", help="Skip video encoding (test only)") + return parser.parse_args() + + +def load_info(info_path: Path) -> dict: + with open(info_path) as f: + return json.load(f) + + +def save_info(info: dict, output_path: Path, num_episodes: int, total_frames: int): + """Update info.json for individual-episode format.""" + info["total_episodes"] = num_episodes + info["total_frames"] = total_frames + info["chunks_size"] = 2000 # All episodes in chunk-000 + info["data_path"] = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet" + info["video_path"] = "videos/{video_key}/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.mp4" + + # Remove meta/episodes path reference since we'll embed tasks directly + info.pop("splits", None) + + with open(output_path / "meta" / "info.json", "w") as f: + json.dump(info, f, indent=2) + + +def decode_png_to_rgb(png_bytes: bytes) -> np.ndarray: + """Decode PNG bytes to RGB numpy array (H, W, 3) uint8.""" + img = Image.open(io.BytesIO(png_bytes)) + return np.array(img.convert("RGB")) + + +def extract_episode_metadata(input_dir: Path) -> tuple[pd.DataFrame, dict]: + """Read tasks from the episodes metadata.""" + meta_dir = input_dir / "meta" / "episodes" + if meta_dir.exists(): + ep_files = sorted(meta_dir.rglob("*.parquet")) + if ep_files: + df = pd.read_parquet(ep_files[0]) + tasks = {} + for _, row in df.iterrows(): + ep_idx = row["episode_index"] + tasks[ep_idx] = row["tasks"] + return df, tasks + + # Fallback: scan parquet files for task_index + return None, {} + + +def process_parquet_file( + parquet_path: Path, + output_data_dir: Path, + output_video_dir: Path, + fps: float, + skip_video: bool = False, +) -> tuple[int, int]: + """ + Process a single chunked parquet file. + Returns (num_episodes_processed, num_frames_processed). + """ + # Read the parquet file + df = pd.read_parquet(parquet_path) + + # Group by episode_index + episodes_processed = 0 + frames_processed = 0 + + for ep_idx, group in df.groupby("episode_index"): + ep_idx = int(ep_idx) + group = group.reset_index(drop=True) + n_frames = len(group) + + # Output parquet path + ep_parquet_path = output_data_dir / f"episode_{ep_idx:06d}.parquet" + + # Drop the image columns for the parquet (they're in the video now) + # But keep them for now — the official stats computation only uses numeric columns + # We need to keep image columns as they might be needed by the dataset loader + # Actually, for the GEAR format, images should ONLY be in videos. + # Remove image columns to avoid confusion. + parquet_cols = [c for c in group.columns + if not c.startswith("observation.images.")] + df_out = group[parquet_cols].copy() + + # Write parquet + table = pa.Table.from_pandas(df_out) + pq.write_table(table, ep_parquet_path) + + if not skip_video: + # Decode and write video for observation.images.image (first camera) + frames = [] + for _, row in group.iterrows(): + img_bytes = row["observation.images.image"]["bytes"] + frame = decode_png_to_rgb(img_bytes) + frames.append(frame) + + # Write mp4 video + ep_video_path = output_video_dir / f"episode_{ep_idx:06d}.mp4" + height, width = frames[0].shape[:2] + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + out = cv2.VideoWriter( + str(ep_video_path), fourcc, fps, (width, height) + ) + for frame in frames: + # cv2 uses BGR order + out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) + out.release() + + episodes_processed += 1 + frames_processed += n_frames + + return episodes_processed, frames_processed + + +def main(): + args = parse_args() + input_dir = Path(args.input_dir) + output_dir = Path(args.output_dir) + + # Create output structure + output_data_dir = output_dir / "data" / "chunk-000" + output_video_dir = output_dir / "videos" / "observation.images.image" / "chunk-000" + output_meta_dir = output_dir / "meta" + + output_data_dir.mkdir(parents=True, exist_ok=True) + output_video_dir.mkdir(parents=True, exist_ok=True) + output_meta_dir.mkdir(parents=True, exist_ok=True) + + # Load original info.json + info = load_info(input_dir / "meta" / "info.json") + fps = info.get("fps", 10.0) + + # Find all data parquet files + data_dir = input_dir / "data" / "chunk-000" + parquet_files = sorted(data_dir.glob("file-*.parquet")) + print(f"Found {len(parquet_files)} parquet files") + + # Process each file + total_episodes = 0 + total_frames = 0 + start_time = time.time() + + for i, pf in enumerate(parquet_files): + n_eps, n_frames = process_parquet_file( + pf, output_data_dir, output_video_dir, fps, + skip_video=args.skip_video, + ) + total_episodes += n_eps + total_frames += n_frames + + elapsed = time.time() - start_time + rate = (i + 1) / elapsed if elapsed > 0 else 0 + eta = (len(parquet_files) - i - 1) / rate if rate > 0 else 0 + print( + f" [{i+1}/{len(parquet_files)}] {pf.name}: " + f"{n_eps} eps, {n_frames} frames " + f"({rate:.1f} files/min, ETA {eta/60:.0f}min)" + ) + + # Write info.json + save_info(info, output_dir, total_episodes, total_frames) + + # Copy tasks metadata if available + meta_ep_dir = input_dir / "meta" / "episodes" / "chunk-000" + if meta_ep_dir.exists(): + ep_files = sorted(meta_ep_dir.glob("*.parquet")) + if ep_files: + ep_meta_df = pd.read_parquet(ep_files[0]) + # Extract task_index → task mapping + tasks = {} + for _, row in ep_meta_df.iterrows(): + ep_idx = int(row["episode_index"]) + task_text = row["tasks"] + if isinstance(task_text, np.ndarray): + task_text = task_text.item() if task_text.size > 0 else "" + elif isinstance(task_text, bytes): + task_text = task_text.decode("utf-8") + tasks[ep_idx] = str(task_text) + + # Write tasks.jsonl for GEAR format + unique_tasks = sorted(set(tasks.values())) + with open(output_meta_dir / "tasks.jsonl", "w") as f: + for ti, task in enumerate(unique_tasks): + f.write(json.dumps({"task_index": ti, "task": task}) + "\n") + + # Write episodes.jsonl + with open(output_meta_dir / "episodes.jsonl", "w") as f: + for _, row in ep_meta_df.iterrows(): + ep_idx = int(row["episode_index"]) + length = int(row["length"]) + task_text = tasks.get(ep_idx, "") + task_index = unique_tasks.index(task_text) if task_text in unique_tasks else -1 + f.write(json.dumps({ + "episode_index": ep_idx, + "length": length, + "task_index": task_index, + }) + "\n") + + print(f"Wrote {len(unique_tasks)} tasks and {total_episodes} episode entries") + + print(f"\nDone! {total_episodes} episodes, {total_frames} frames") + print(f"Output: {output_dir}") + print(f"Time: {(time.time() - start_time)/60:.1f} minutes") + print(f"\nNext step: run convert_lerobot_to_gear.py on the output dir") + + +if __name__ == "__main__": + main() diff --git a/scripts/data/prepare_robotwin.py b/scripts/data/prepare_robotwin.py new file mode 100644 index 0000000000000000000000000000000000000000..df8c95777a1b26fb6d07d07a7e96a0e91bbec91b --- /dev/null +++ b/scripts/data/prepare_robotwin.py @@ -0,0 +1,317 @@ +#!/usr/bin/env python3 +"""RoboTwin 数据准备脚本：将原始 ZIP 格式转换为 RobotWinDataset 期望的 qpos+videos+metas 格式。 + +用法: + python scripts/data/prepare_robotwin.py \\ + --input /root/autol-tmp/data/robotwin3/dataset \\ + --output /root/autol-tmp/data/robotwin_gear + +RoboTwin 原始结构 (ZIP): + / + ├── franka_clean_50.zip + │ └── franka_clean_50/ + │ ├── scene_info.json + │ ├── instructions/episodeN.json # {"seen": [...], "unseen": [...]} + │ ├── video/episodeN.mp4 + │ └── _traj_data/episodeN.pkl # {arm_key: [seg1, seg2, ...]} + ├── aloha-agilex_clean_50.zip + └── ... + +输出结构: + / + ├── _/ + │ ├── qpos/ + │ │ ├── episode0.pt # [T_qpos, 14] float32 + │ │ └── episode1.pt + │ ├── videos/ + │ │ ├── episode0.mp4 + │ │ └── episode1.mp4 + │ └── metas/ + │ ├── task_0.txt + │ └── task_1.txt + └── ... +""" + +import os, sys, json, argparse, logging, pickle, shutil +from pathlib import Path + +import numpy as np +import torch +import zipfile + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger(__name__) + +# 14-dim layout: 前 7 = right arm, 后 7 = left arm +# 单臂机器人：使用的手臂放前 N 位，其余补 0 +ROBOT_CONFIG = { + "franka": {"arm": "right", "dims": 7}, + "aloha-agilex": {"arm": "left", "dims": 6}, + "arx-x5": {"arm": "left", "dims": 6}, + "ur5": {"arm": "left", "dims": 6}, + "piper": {"arm": "left", "dims": 6}, +} + +OUTPUT_STATE_DIM = 14 + + +def get_robot_name(zip_path: Path) -> str | None: + """从 ZIP 文件名提取机器人名称（去掉 _clean_50.zip 后缀）。""" + name = zip_path.stem # e.g. franka_clean_50 + # Remove _clean_50 or _50 suffix + for suffix in ["_clean_50", "_50"]: + if name.endswith(suffix): + name = name[:-len(suffix)] + break + return name + + +def get_arm_key(robot: str) -> str: + """根据机器人类型返回 pkl 中的 arm key。""" + cfg = ROBOT_CONFIG.get(robot) + if cfg is None: + logger.warning(f"Unknown robot {robot}, trying right_joint_path") + return "right_joint_path" + return f"{cfg['arm']}_joint_path" + + +def merge_trajectory_segments(segments: list) -> np.ndarray: + """合并多段轨迹为一个连续数组。 + + Args: + segments: list of dict, each with "position" key [T_i, D] + + Returns: + concatenated position array [sum(T_i), D] + """ + arrays = [seg["position"] for seg in segments if seg.get("position") is not None] + if not arrays: + return np.zeros((0, arrays[0].shape[1])) if arrays else np.zeros((0, 1)) + return np.concatenate(arrays, axis=0) + + +def pad_to_14dim(arr: np.ndarray, joint_dim: int) -> np.ndarray: + """Pad joint positions to 14-dim。 + + 布局：前 7 = right arm, 后 7 = left arm。 + - franka (7, right): 放在前 7 维 + - aloha (6, left): 放在后 6 维（前补 0） + """ + if arr.ndim == 1: + padded = np.zeros(OUTPUT_STATE_DIM, dtype=np.float32) + if joint_dim == 7: + padded[:joint_dim] = arr.astype(np.float32) + else: # 6-dim → 后 6 维 + padded[OUTPUT_STATE_DIM - joint_dim:] = arr.astype(np.float32) + else: + padded = np.zeros((arr.shape[0], OUTPUT_STATE_DIM), dtype=np.float32) + if joint_dim == 7: + padded[:, :joint_dim] = arr.astype(np.float32) + else: # 6-dim → 后 6 维 + padded[:, OUTPUT_STATE_DIM - joint_dim:] = arr.astype(np.float32) + return padded + + +def get_task_description(zf: zipfile.ZipFile, robot_prefix: str, ep_idx: int) -> str: + """从 instruction JSON 读取任务描述。 + + Returns: + 第一条 "seen" 指令，或 fallback 文本 + """ + try: + inst_path = f"{robot_prefix}/instructions/episode{ep_idx}.json" + with zf.open(inst_path) as f: + inst = json.load(f) + seen = inst.get("seen", []) + if seen: + return seen[0] + unseen = inst.get("unseen", []) + if unseen: + return unseen[0] + except Exception as e: + logger.debug(f" Cannot read instruction: {e}") + return "Manipulate the object on the table" + + +def process_robot_zip( + zip_path: Path, + output_dir: Path, + delete_after: bool = False, +) -> int: + """处理一个 RoboTwin ZIP 文件。 + + Returns: + 成功处理的 episode 数 + """ + robot = get_robot_name(zip_path) + if robot is None: + logger.warning(f"Cannot determine robot from {zip_path.name}, skipping") + return 0 + + # 输出目录：{task}_{robot} + task_name = zip_path.parent.name + out_subdir = output_dir / f"{task_name}_{robot}" + qpos_dir = out_subdir / "qpos" + video_dir = out_subdir / "videos" + meta_dir = out_subdir / "metas" + qpos_dir.mkdir(parents=True, exist_ok=True) + video_dir.mkdir(parents=True, exist_ok=True) + meta_dir.mkdir(parents=True, exist_ok=True) + + cfg = ROBOT_CONFIG.get(robot) + if cfg is None: + logger.warning(f"Unknown robot {robot}, skipping {zip_path}") + return 0 + + joint_dim = cfg["dims"] + arm_key = get_arm_key(robot) + robot_prefix = f"{robot}_clean_50" # ZIP 内部目录名 + + try: + zf = zipfile.ZipFile(str(zip_path)) + except Exception as e: + logger.error(f"Cannot open {zip_path}: {e}") + return 0 + + # 从 scene_info.json 获取 episode 列表 + try: + with zf.open(f"{robot_prefix}/scene_info.json") as f: + scene_info = json.load(f) + except Exception as e: + logger.warning(f"No scene_info in {zip_path}: {e}") + zf.close() + return 0 + + # 列出所有 episode + episode_keys = sorted([k for k in scene_info if k.startswith("episode_")], + key=lambda x: int(x.split("_")[1])) + if not episode_keys: + logger.warning(f"No episodes in scene_info of {zip_path}") + zf.close() + return 0 + + success_count = 0 + for ep_key in episode_keys: + ep_idx = int(ep_key.split("_")[1]) + ep_name = f"episode{ep_idx}" + + # 检查是否有轨迹数据 + traj_path = f"{robot_prefix}/_traj_data/{ep_name}.pkl" + video_path_in = f"{robot_prefix}/video/{ep_name}.mp4" + + if traj_path not in zf.namelist(): + logger.debug(f" No traj data for {ep_name} in {zip_path.name}") + continue + if video_path_in not in zf.namelist(): + logger.debug(f" No video for {ep_name} in {zip_path.name}") + continue + + try: + # --- 轨迹处理 --- + with zf.open(traj_path) as f: + traj_data = pickle.load(f) + + segments = traj_data.get(arm_key, []) + if not segments: + logger.debug(f" No {arm_key} segments for {ep_name}") + continue + + # 合并多段轨迹 + pos = merge_trajectory_segments(segments) # [T, joint_dim] + if pos.shape[0] < 24: # 至少需要 num_frames + action_horizon = 24 + logger.debug(f" Too few frames ({pos.shape[0]}) for {ep_name}") + continue + + # Pad 到 14-dim + pos_14 = pad_to_14dim(pos, joint_dim) # [T, 14] + + # 保存为 .pt + qpos_path = qpos_dir / f"{ep_name}.pt" + torch.save(torch.from_numpy(pos_14), qpos_path) + + # --- 视频提取 --- + video_out_path = video_dir / f"{ep_name}.mp4" + with zf.open(video_path_in) as src, open(video_out_path, "wb") as dst: + shutil.copyfileobj(src, dst) + + # --- 任务描述 --- + task_desc = get_task_description(zf, robot_prefix, ep_idx) + meta_path = meta_dir / f"task_{ep_idx}.txt" + with open(meta_path, "w") as f: + f.write(task_desc) + + success_count += 1 + + except Exception as e: + logger.warning(f" Error processing {ep_name} in {zip_path.name}: {e}") + continue + + zf.close() + + # 删除 ZIP 释放空间 + if delete_after and success_count > 0: + try: + zip_path.unlink() + logger.info(f" Deleted {zip_path.name}") + except Exception as e: + logger.warning(f" Cannot delete {zip_path.name}: {e}") + + if success_count > 0: + logger.info(f"{zip_path.name}: {success_count}/{len(episode_keys)} episodes extracted -> {out_subdir}") + + return success_count + + +def main(): + parser = argparse.ArgumentParser(description="Prepare RoboTwin data for DreamZero training") + parser.add_argument("--input", "-i", required=True, + help="RoboTwin 数据目录 (含任务子目录)") + parser.add_argument("--output", "-o", required=True, + help="输出目录") + parser.add_argument("--delete-zip", action="store_true", + help="处理完成后删除 ZIP 文件（节省空间）") + parser.add_argument("--max-tasks", type=int, default=None, + help="最多处理前 N 个任务（用于测试）") + parser.add_argument("--num-workers", type=int, default=4, + help="并行处理数（暂未实现）") + args = parser.parse_args() + + input_dir = Path(args.input) + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + # 扫描所有任务目录 + task_dirs = sorted([d for d in input_dir.iterdir() if d.is_dir()]) + logger.info(f"Found {len(task_dirs)} task directories") + + if args.max_tasks: + task_dirs = task_dirs[:args.max_tasks] + logger.info(f"Limited to {args.max_tasks} tasks") + + total_episodes = 0 + total_zips = 0 + + for task_dir in task_dirs: + # 找到所有 ZIP 文件 + zip_files = sorted(task_dir.glob("*_clean_50.zip")) + if not zip_files: + logger.warning(f"No ZIP files in {task_dir}") + continue + + for zip_path in zip_files: + try: + ep_count = process_robot_zip( + zip_path, output_dir, delete_after=args.delete_zip + ) + total_episodes += ep_count + total_zips += 1 + except Exception as e: + logger.error(f"Fatal error processing {zip_path}: {e}") + continue + + logger.info(f"Done! Processed {total_zips} ZIPs, {total_episodes} episodes") + logger.info(f"Output: {output_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/inference/build_trt_engine.py b/scripts/inference/build_trt_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..227906c7f50f669323c970de1b0427d5b2809d83 --- /dev/null +++ b/scripts/inference/build_trt_engine.py @@ -0,0 +1,243 @@ +"""Build a TensorRT engine from a DreamZero checkpoint. + +Must be launched via build_trt_engine.sh (or with ENABLE_TENSORRT=true already +set) so that flash-attention compatibility mode is active before any groot model +modules are imported. + +Launched via torchrun so that RANK / WORLD_SIZE / MASTER_* env vars exist for +GrootSimPolicy's distributed initialisation. + +Calibration: + For quantized precisions (nvfp4, fp8), ModelOpt calibrates quantization + parameters by observing activation statistics during forward passes. Using + real dataset trajectories produces a significantly more accurate engine than + random dummy inputs. Pass --dataset-path to enable real calibration. +""" + +import os +import sys +import argparse +import logging +from types import SimpleNamespace + +# Verify ENABLE_TENSORRT was exported before any groot imports occur. +if os.getenv("ENABLE_TENSORRT", "").lower() != "true": + print( + "ERROR: ENABLE_TENSORRT must be 'true' before importing this script.\n" + "Use build_trt_engine.sh instead of calling this script directly.", + file=sys.stderr, + ) + sys.exit(1) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", +) +logger = logging.getLogger(__name__) + +import numpy as np +import torch +import torch.distributed as dist +from tianshou.data import Batch +from torch.distributed.device_mesh import init_device_mesh + +from groot.vla.data.schema import EmbodimentTag +from groot.vla.model.n1_5.sim_policy import GrootSimPolicy +from groot.control.tensorrt_utils import ( + wan_trt_quantize_and_load_engine, + create_wan_test_inputs, +) + +# DreamZero-DROID uses the ar_14B_droid model type in tensorrt_utils. +_MODEL_TYPE = "ar_14B_droid" + + +def _init_single_gpu_mesh(): + """Initialise a single-GPU device mesh (launched via torchrun --nproc_per_node=1).""" + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(rank) + mesh = init_device_mesh( + device_type="cuda", + mesh_shape=(world_size,), + mesh_dim_names=("ip",), + ) + return mesh + + +def _make_dummy_forward_loop(): + """Fallback calibration using random dummy inputs. + + Acceptable for fp16 (no quantization), but may reduce accuracy for + nvfp4/fp8 since the activation distribution differs from real data. + Prefer _make_dataset_forward_loop when a dataset is available. + """ + def forward_loop(model): + trt_forward = getattr(model, "_forward_inference_trt_droid", model.forward) + test_inputs = create_wan_test_inputs(None, device="cuda", model_type=_MODEL_TYPE) + for _ in range(16): + with torch.no_grad(): + trt_forward(*test_inputs) + + return forward_loop + + +def _make_dataset_forward_loop(policy, dataset_path: str, num_calibration_trajs: int = 2): + """Real-data calibration loop — mirrors the internal droid_video_pred.sh approach. + + Loads ``num_calibration_trajs`` trajectories from the LeRobot dataset and + runs ``policy.lazy_joint_forward_causal`` at each action-horizon step, + exercising the DiT model with realistic activation distributions. + """ + from groot.vla.data.dataset.lerobot import LeRobotSingleDataset + + def forward_loop(model): + logger.info( + "Calibration: loading dataset from %s (%d trajs)", dataset_path, num_calibration_trajs + ) + dataset = LeRobotSingleDataset( + dataset_path=dataset_path, + modality_configs=policy.modality_configs, + embodiment_tag=policy.embodiment_tag, + video_backend="torchvision_av", + video_backend_kwargs=None, + transforms=None, # policy.lazy_joint_forward_causal applies transforms + use_global_metadata=False, + ) + + action_horizon = policy.trained_model.action_head.action_horizon + num_frame_per_block = policy.trained_model.action_head.num_frame_per_block + torch._dynamo.config.recompile_limit = 500 + + for traj_id in range(min(num_calibration_trajs, len(dataset.trajectory_lengths))): + logger.info("Calibration trajectory %d / %d", traj_id + 1, num_calibration_trajs) + traj_len = int(dataset.trajectory_lengths[traj_id]) + latent_video = None + + # Step through the trajectory at action-horizon intervals (same cadence as + # real inference) for up to 5 chunks — enough to cover the KV-cache build-up + # and the cached inference path that the TRT engine will handle. + max_steps = min(traj_len, 5 * action_horizon) + for step in range(0, max_steps, action_horizon): + # Clamp delta indices to valid range for this trajectory. + indices = { + k: np.clip(v + step, 0, traj_len - 1) + for k, v in dataset.delta_indices.items() + } + data_point = dataset.get_step_data(traj_id, indices) + batch = Batch(obs=data_point) + + dist.barrier() + with torch.no_grad(): + result_batch, video_pred = policy.lazy_joint_forward_causal( + batch, latent_video=latent_video + ) + dist.barrier() + + # Feed the last generated frame back as context for the next step, + # matching autoregressive inference behaviour. + if video_pred is not None: + latent_video = video_pred[:, :, -num_frame_per_block:] + + # Reset AR state between trajectories. + policy.trained_model.action_head.current_start_frame = 0 + policy.trained_model.action_head.kv_cache1 = None + policy.trained_model.action_head.kv_cache_neg = None + policy.trained_model.action_head.crossattn_cache = None + policy.trained_model.action_head.crossattn_cache_neg = None + + return forward_loop + + +def main(): + parser = argparse.ArgumentParser( + description="Build TensorRT engine for the DreamZero DiT model." + ) + parser.add_argument("--model-path", required=True, help="Path to checkpoint directory.") + parser.add_argument( + "--tensorrt", + required=True, + choices=["nvfp4", "fp8", "fp16"], + help="TensorRT quantization / precision format.", + ) + parser.add_argument( + "--dataset-path", + default=None, + help=( + "Path to a LeRobot-format DROID dataset for real calibration. " + "Strongly recommended for nvfp4/fp8 — random dummy inputs are used as " + "fallback but may reduce quantization accuracy." + ), + ) + parser.add_argument( + "--num-calibration-trajs", + type=int, + default=2, + help="Number of dataset trajectories used for calibration (default: 2).", + ) + args = parser.parse_args() + + if args.tensorrt in ("nvfp4", "fp8") and args.dataset_path is None: + logger.warning( + "No --dataset-path provided for %s quantization. " + "Falling back to random dummy inputs — this may reduce engine accuracy. " + "Re-run with --dataset-path for best results.", + args.tensorrt, + ) + + engine_dir = os.path.join(args.model_path, "tensorrt", "wan") + engine_path = os.path.join(engine_dir, f"WanModel_{args.tensorrt}.trt") + onnx_path = os.path.join(engine_dir, f"CausalWanModel.onnx") + os.makedirs(engine_dir, exist_ok=True) + + if os.path.exists(engine_path): + logger.info("TRT engine already exists: %s", engine_path) + logger.info("Delete it first if you want to rebuild.") + return + + logger.info("Loading DreamZero policy from : %s", args.model_path) + logger.info("Target engine path : %s", engine_path) + logger.info("Quantization precision : %s", args.tensorrt) + + device_mesh = _init_single_gpu_mesh() + + policy = GrootSimPolicy( + embodiment_tag=EmbodimentTag("oxe_droid"), + model_path=args.model_path, + device="cuda" if torch.cuda.is_available() else "cpu", + device_mesh=device_mesh, + ) + + # Build calibration forward loop — prefer real data for quantized precisions. + if args.dataset_path is not None: + forward_loop = _make_dataset_forward_loop( + policy, args.dataset_path, args.num_calibration_trajs + ) + logger.info( + "Calibration: using %d real trajectories from %s", + args.num_calibration_trajs, + args.dataset_path, + ) + else: + forward_loop = _make_dummy_forward_loop() + logger.info("Calibration: using random dummy inputs (no --dataset-path given).") + + # cfg mimics the Hydra config used by the internal eval script. + cfg = SimpleNamespace(inference_mode="trt_build", quantize_dtype=args.tensorrt) + + logger.info("Building TensorRT engine (ONNX export + trtexec, may take 10-30 min) ...") + wan_trt_quantize_and_load_engine( + policy=policy, + cfg=cfg, + onnx_path=onnx_path, + engine_path=engine_path, + model_type=_MODEL_TYPE, + forward_loop=forward_loop, + ) + + logger.info("TRT engine saved to: %s", engine_path) + + +if __name__ == "__main__": + main() diff --git a/scripts/inference/build_trt_engine.sh b/scripts/inference/build_trt_engine.sh new file mode 100644 index 0000000000000000000000000000000000000000..5c943cc6bf670a18930f853d1d1a31dbf20ada4e --- /dev/null +++ b/scripts/inference/build_trt_engine.sh @@ -0,0 +1,142 @@ +#!/usr/bin/env bash +# Build a TensorRT engine from a DreamZero checkpoint. +# +# Usage (recommended — with real calibration data): +# bash scripts/inference/build_trt_engine.sh \ +# --model-path ./checkpoints/DreamZero-DROID \ +# --tensorrt nvfp4 \ +# --dataset-path ./data/droid_lerobot \ +# --cuda-device 0 +# +# Usage (without dataset — acceptable for fp16, not recommended for nvfp4/fp8): +# bash scripts/inference/build_trt_engine.sh \ +# --model-path ./checkpoints/DreamZero-DROID \ +# --tensorrt nvfp4 \ +# --cuda-device 0 +# +# The engine is saved to: +# {model_path}/tensorrt/wan/WanModel_{precision}.trt +# +# Supported precisions: nvfp4 (recommended), fp8, fp16 +# +# For quantized precisions (nvfp4, fp8), ModelOpt calibrates quantization +# parameters using real forward passes. Providing --dataset-path is strongly +# recommended — random dummy inputs are used as fallback but reduce accuracy. +# +# ENABLE_TENSORRT=true must be set before any groot modules are imported +# (it controls flash-attention compatibility mode for ONNX/TRT export). +# This script sets it and launches the Python build script via torchrun so +# that RANK / WORLD_SIZE env vars are available for GrootSimPolicy init. + +# export HF_HUB_CACHE=/mnt/aws-lfs-02/shared/ckpts +set -euo pipefail + +MODEL_PATH="" +TENSORRT_PRECISION="" +CUDA_DEVICE="0" +DATASET_PATH="" +NUM_CALIBRATION_TRAJS="2" + +while [[ $# -gt 0 ]]; do + case $1 in + --model-path) + MODEL_PATH="$2" + shift 2 + ;; + --tensorrt) + TENSORRT_PRECISION="$2" + shift 2 + ;; + --cuda-device) + CUDA_DEVICE="$2" + shift 2 + ;; + --dataset-path) + DATASET_PATH="$2" + shift 2 + ;; + --num-calibration-trajs) + NUM_CALIBRATION_TRAJS="$2" + shift 2 + ;; + -h|--help) + echo "Usage: $0 --model-path --tensorrt [OPTIONS]" + echo "" + echo "Options:" + echo " --model-path PATH Path to DreamZero checkpoint directory" + echo " --tensorrt PRECISION TRT precision: nvfp4 (recommended), fp8, fp16" + echo " --dataset-path PATH LeRobot dataset for real calibration (recommended for nvfp4/fp8)" + echo " --num-calibration-trajs N Number of calibration trajectories (default: 2)" + echo " --cuda-device ID CUDA device index (default: 0)" + exit 0 + ;; + *) + echo "Unknown argument: $1" >&2 + exit 1 + ;; + esac +done + +if [[ -z "$MODEL_PATH" ]]; then + echo "Error: --model-path is required" >&2 + exit 1 +fi + +if [[ -z "$TENSORRT_PRECISION" ]]; then + echo "Error: --tensorrt is required (e.g. nvfp4, fp8, fp16)" >&2 + exit 1 +fi + +if [[ ! -d "$MODEL_PATH" ]]; then + echo "Error: checkpoint directory not found: $MODEL_PATH" >&2 + exit 1 +fi + +ENGINE_PATH="${MODEL_PATH}/tensorrt/wan/WanModel_${TENSORRT_PRECISION}.trt" + +echo "==========================================" +echo "DreamZero TensorRT Engine Builder" +echo " Checkpoint : $MODEL_PATH" +echo " Precision : $TENSORRT_PRECISION" +echo " CUDA device : $CUDA_DEVICE" +echo " Dataset (calibrate): ${DATASET_PATH:-}" +echo " Calibration trajs : $NUM_CALIBRATION_TRAJS" +echo " Output : $ENGINE_PATH" +echo "==========================================" + +# ENABLE_TENSORRT must be set before Python imports any groot model modules +# (it activates flash-attention compatibility mode required for ONNX/TRT export). +export ENABLE_TENSORRT=true +export CUDA_VISIBLE_DEVICES="$CUDA_DEVICE" +export ATTENTION_BACKEND="TE" +export HYDRA_FULL_ERROR=1 + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" + +# Build the Python argument list. +PYTHON_ARGS=( + --model-path "$MODEL_PATH" + --tensorrt "$TENSORRT_PRECISION" + --num-calibration-trajs "$NUM_CALIBRATION_TRAJS" +) +if [[ -n "$DATASET_PATH" ]]; then + PYTHON_ARGS+=(--dataset-path "$DATASET_PATH") +fi + +# torchrun sets RANK / WORLD_SIZE / MASTER_ADDR / MASTER_PORT which are +# required by GrootSimPolicy's distributed init. +torchrun \ + --standalone \ + --nproc_per_node=1 \ + "${REPO_ROOT}/scripts/inference/build_trt_engine_droid.py" \ + "${PYTHON_ARGS[@]}" + +echo "==========================================" +echo "Engine built successfully: $ENGINE_PATH" +echo "" +echo "Run inference with:" +echo " CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.run --standalone --nproc_per_node=2 \\" +echo " socket_test_optimized_AR.py --port 5000 --enable-dit-cache \\" +echo " --model-path ${MODEL_PATH} --tensorrt ${TENSORRT_PRECISION}" +echo "==========================================" diff --git a/scripts/open_loop_yam.py b/scripts/open_loop_yam.py new file mode 100644 index 0000000000000000000000000000000000000000..b5a5e989ea56ab77e77cb050a7189c78bb766769 --- /dev/null +++ b/scripts/open_loop_yam.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 +"""Offline open-loop evaluation for DreamZero on YAM data. + +Loads a model checkpoint directly (no server needed), reads YAM dataset +(parquet + MP4), runs inference, and compares predicted vs ground-truth actions. + +Usage: + python scripts/open_loop_yam.py \ + --model_path /path/to/checkpoint \ + --dataset_path Dataset/YAM_play_data \ + --device cuda:0 \ + --num_samples 200 +""" + +import torch._dynamo +torch._dynamo.config.disable = True + +import argparse +import glob +import os +import time + +import cv2 +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pyarrow.parquet as pq +import torch +import torch.distributed as dist +from tianshou.data import Batch + +from groot.vla.data.schema import EmbodimentTag +from groot.vla.model.n1_5.sim_policy import GrootSimPolicy + + +# --------------------------------------------------------------------------- +# YAM layout (from Dataset/YAM_play_data/meta/modality.json) +# --------------------------------------------------------------------------- + +VIDEO_CAMERAS = { + "video.top_camera-images-rgb": "observation.images.top_camera-images-rgb", + "video.left_camera-images-rgb": "observation.images.left_camera-images-rgb", + "video.right_camera-images-rgb": "observation.images.right_camera-images-rgb", +} + +STATE_SLICES = { + "state.left_joint_pos": (34, 40), + "state.left_gripper_pos": (32, 33), + "state.right_joint_pos": (40, 46), + "state.right_gripper_pos": (33, 34), +} + +ACTION_SLICES = { + "action.left_joint_pos": (34, 40), + "action.left_gripper_pos": (32, 33), + "action.right_joint_pos": (40, 46), + "action.right_gripper_pos": (33, 34), +} + +ACTION_KEY_ORDER = [ + "action.left_joint_pos", + "action.left_gripper_pos", + "action.right_joint_pos", + "action.right_gripper_pos", +] + + +# --------------------------------------------------------------------------- +# Dataset reader (LeRobot chunked format) +# --------------------------------------------------------------------------- + +class YAMDataset: + """Reads LeRobot-style chunked parquet + MP4.""" + + def __init__(self, dataset_path: str): + self.root = dataset_path + + data_dir = os.path.join(dataset_path, "data") + parquet_files = sorted(glob.glob(os.path.join(data_dir, "**", "episode_*.parquet"), recursive=True)) + if not parquet_files: + raise FileNotFoundError(f"No episode_*.parquet found under {data_dir}") + + self.episodes = [] + self.cum_lengths = [0] + for pf in parquet_files: + table = pq.read_table(pf) + self.episodes.append(table) + self.cum_lengths.append(self.cum_lengths[-1] + table.num_rows) + self.total_rows = self.cum_lengths[-1] + + videos_root = os.path.join(dataset_path, "videos") + self.video_dirs = {} + for server_key, folder_name in VIDEO_CAMERAS.items(): + candidates = sorted(glob.glob(os.path.join(videos_root, "**", folder_name), recursive=True)) + if candidates: + self.video_dirs[server_key] = candidates[0] + + print(f"YAMDataset: {len(self.episodes)} episodes, " + f"{self.total_rows} rows, {len(self.video_dirs)} cameras") + + def __len__(self): + return self.total_rows + + def _locate(self, idx): + for ep in range(len(self.episodes)): + if idx < self.cum_lengths[ep + 1]: + return ep, idx - self.cum_lengths[ep] + raise IndexError(f"Index {idx} out of range ({self.total_rows})") + + def get_state(self, idx) -> np.ndarray: + ep, row = self._locate(idx) + return np.array(self.episodes[ep].column("observation.state")[row].as_py(), dtype=np.float64) + + def get_action(self, idx) -> np.ndarray: + ep, row = self._locate(idx) + return np.array(self.episodes[ep].column("action")[row].as_py(), dtype=np.float64) + + def get_task(self, idx) -> str: + ep, row = self._locate(idx) + try: + return str(self.episodes[ep].column("annotation.task")[row].as_py()) + except Exception: + return "" + + def get_frame(self, idx, server_key) -> np.ndarray: + """Read one video frame → (H, W, 3) uint8 RGB.""" + ep, row = self._locate(idx) + mp4 = os.path.join(self.video_dirs[server_key], f"episode_{ep:06d}.mp4") + cap = cv2.VideoCapture(mp4) + cap.set(cv2.CAP_PROP_POS_FRAMES, row) + ret, frame = cap.read() + cap.release() + if not ret: + raise RuntimeError(f"Failed to read frame {row} from {mp4}") + return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + +# --------------------------------------------------------------------------- +# Observation builder +# --------------------------------------------------------------------------- + +def build_obs(dataset: YAMDataset, idx: int, prompt: str) -> dict: + """Build an obs dict matching what GrootSimPolicy.forward() expects.""" + obs = {} + + for server_key in dataset.video_dirs: + frame = dataset.get_frame(idx, server_key) + obs[server_key] = frame[np.newaxis, ...].astype(np.uint8) # (1, H, W, C) + + state = dataset.get_state(idx) + for key, (start, end) in STATE_SLICES.items(): + obs[key] = state[start:end].reshape(1, -1).astype(np.float64) # (1, D) + + obs["annotation.task"] = prompt + + return obs + + +def get_gt_action_dict(dataset: YAMDataset, idx: int) -> dict: + """Split the flat GT action vector into per-key arrays.""" + action_flat = dataset.get_action(idx) + gt = {} + for key in ACTION_KEY_ORDER: + s, e = ACTION_SLICES[key] + gt[key] = action_flat[s:e] + return gt + + +# --------------------------------------------------------------------------- +# Plotting +# --------------------------------------------------------------------------- + +def save_plots(all_preds, all_gts, key_names, output_dir): + """Plot pred vs gt for each action dimension across all keys.""" + pred_flat = np.concatenate([all_preds[k] for k in key_names], axis=-1) + gt_flat = np.concatenate([all_gts[k] for k in key_names], axis=-1) + D = pred_flat.shape[1] + mse_dim = np.mean((pred_flat - gt_flat) ** 2, axis=0) + + for d in range(D): + plt.figure(figsize=(10, 4)) + plt.plot(gt_flat[:, d], label="gt", alpha=0.8) + plt.plot(pred_flat[:, d], label="pred", alpha=0.8) + plt.title(f"Action dim {d} (MSE={mse_dim[d]:.6f})") + plt.xlabel("sample index"); plt.ylabel("value") + plt.legend(); plt.grid(True, alpha=0.3); plt.tight_layout() + plt.savefig(os.path.join(output_dir, f"action_dim_{d}.png"), dpi=150) + plt.close() + + ncols = 4 + nrows = (D + ncols - 1) // ncols + fig, axes = plt.subplots(nrows, ncols, figsize=(5 * ncols, 3.5 * nrows), squeeze=False) + overall_mse = float(np.mean(mse_dim)) + fig.suptitle(f"All action dims (overall MSE={overall_mse:.6f})", fontsize=14) + for d in range(D): + ax = axes[d // ncols][d % ncols] + ax.plot(gt_flat[:, d], label="gt", alpha=0.7, lw=0.8) + ax.plot(pred_flat[:, d], label="pred", alpha=0.7, lw=0.8) + ax.set_title(f"dim {d} (MSE={mse_dim[d]:.4f})", fontsize=9) + ax.tick_params(labelsize=7); ax.grid(True, alpha=0.2) + if d == 0: ax.legend(fontsize=7) + for d in range(D, nrows * ncols): + axes[d // ncols][d % ncols].set_visible(False) + fig.tight_layout(rect=[0, 0, 1, 0.96]) + fig.savefig(os.path.join(output_dir, "all_action_dims.png"), dpi=200) + plt.close(fig) + + # Per-key summary plot + fig2, axes2 = plt.subplots(1, len(key_names), figsize=(5 * len(key_names), 4), squeeze=False) + for i, k in enumerate(key_names): + ax = axes2[0][i] + p, g = all_preds[k], all_gts[k] + for d in range(p.shape[1]): + ax.plot(g[:, d], '--', alpha=0.5, lw=0.8) + ax.plot(p[:, d], alpha=0.7, lw=0.8) + key_mse = float(np.mean((p - g) ** 2)) + ax.set_title(f"{k}\nMSE={key_mse:.6f}", fontsize=9) + ax.grid(True, alpha=0.2); ax.tick_params(labelsize=7) + fig2.suptitle("Per-key pred (solid) vs gt (dashed)", fontsize=12) + fig2.tight_layout(rect=[0, 0, 1, 0.94]) + fig2.savefig(os.path.join(output_dir, "per_key_summary.png"), dpi=200) + plt.close(fig2) + + return mse_dim, overall_mse + + +# --------------------------------------------------------------------------- +# Main evaluation +# --------------------------------------------------------------------------- + +def evaluate(args): + # Single-process distributed init (GrootSimPolicy uses dist.get_rank()) + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29500") + dist.init_process_group(backend="gloo", world_size=1, rank=0) + + print(f"Loading model from {args.model_path} ...") + policy = GrootSimPolicy( + embodiment_tag=EmbodimentTag.YAM, + model_path=args.model_path, + device=args.device, + ) + print("Model loaded.") + + dataset = YAMDataset(args.dataset_path) + os.makedirs(args.output_dir, exist_ok=True) + + num = min(args.num_samples, len(dataset)) + preds_per_key = {k: [] for k in ACTION_KEY_ORDER} + gts_per_key = {k: [] for k in ACTION_KEY_ORDER} + times = [] + + print(f"\nEvaluating {num} samples (start={args.start_idx}) ...") + print("-" * 60) + + for i in range(num): + idx = args.start_idx + i + + prompt = args.prompt + if args.use_dataset_prompt: + task = dataset.get_task(idx) + if task: + prompt = task + + obs = build_obs(dataset, idx, prompt) + + t0 = time.perf_counter() + with torch.inference_mode(): + result, _ = policy.lazy_joint_forward_causal(Batch(obs=obs)) + elapsed = time.perf_counter() - t0 + times.append(elapsed) + + gt = get_gt_action_dict(dataset, idx) + + for k in ACTION_KEY_ORDER: + if k in result.act: + pred_val = result.act[k] + if isinstance(pred_val, torch.Tensor): + pred_val = pred_val.cpu().numpy() + # First timestep: squeeze may collapse (1,24,1) -> (24,) + pred_val = np.atleast_1d(pred_val[0]).flatten() + preds_per_key[k].append(pred_val) + gts_per_key[k].append(gt[k]) + + if i % args.log_every == 0: + if i == 0: + print(f" Action keys in output: {list(result.act.keys())}") + for k in ACTION_KEY_ORDER: + if k in result.act: + v = result.act[k] + shape = v.shape if hasattr(v, 'shape') else "?" + print(f" {k}: pred_shape={shape}, gt_shape={gt[k].shape}") + print(f" [{i:>5d}/{num}] idx={idx} infer={elapsed:.3f}s prompt={prompt!r:.60}") + + # Stack results + valid_keys = [k for k in ACTION_KEY_ORDER if len(preds_per_key[k]) > 0] + if not valid_keys: + print("No predictions!"); return + + stacked_preds = {k: np.stack(preds_per_key[k]) for k in valid_keys} + stacked_gts = {k: np.stack(gts_per_key[k]) for k in valid_keys} + + pred_all = np.concatenate([stacked_preds[k] for k in valid_keys], axis=-1) + gt_all = np.concatenate([stacked_gts[k] for k in valid_keys], axis=-1) + overall_mse = float(np.mean((pred_all - gt_all) ** 2)) + + print(f"\n{'='*60}") + print(f"Overall MSE: {overall_mse:.6f} | Avg inference time: {np.mean(times):.4f}s") + for k in valid_keys: + k_mse = float(np.mean((stacked_preds[k] - stacked_gts[k]) ** 2)) + print(f" {k}: MSE={k_mse:.6f}") + print(f"{'='*60}") + + mse_dim, _ = save_plots(stacked_preds, stacked_gts, valid_keys, args.output_dir) + + with open(os.path.join(args.output_dir, "mse.txt"), "w") as f: + f.write(f"overall_mse,{overall_mse}\n") + for k in valid_keys: + k_mse = float(np.mean((stacked_preds[k] - stacked_gts[k]) ** 2)) + f.write(f"{k},{k_mse}\n") + for d, v in enumerate(mse_dim): + f.write(f"dim_{d},{v}\n") + + print(f"Results saved to {os.path.abspath(args.output_dir)}/") + + +def main(): + p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + p.add_argument("--model_path", required=True, + help="Path to model checkpoint dir (contains config.json, model.safetensors, experiment_cfg/)") + p.add_argument("--dataset_path", required=True, + help="Root of YAM dataset (contains data/, videos/, meta/)") + p.add_argument("--device", default="cuda:0") + p.add_argument("--prompt", default="pick up the object") + p.add_argument("--use_dataset_prompt", action="store_true", + help="Read task annotation from parquet instead of --prompt") + p.add_argument("--num_samples", type=int, default=300) + p.add_argument("--start_idx", type=int, default=0) + p.add_argument("--output_dir", default="results_yam") + p.add_argument("--log_every", type=int, default=10) + main_args = p.parse_args() + evaluate(main_args) + + +if __name__ == "__main__": + main() diff --git a/scripts/setup_env.sh b/scripts/setup_env.sh new file mode 100644 index 0000000000000000000000000000000000000000..82f3fcf51b9d1fb7858f11706ecdc4b44e5579c6 --- /dev/null +++ b/scripts/setup_env.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# DreamZero 环境快速设置脚本 +# +# 用法: +# source scripts/setup_env.sh # 中国网络（镜像站） +# source scripts/setup_env.sh --offline # 无网络环境 + +set -euo pipefail + +MODE="${1:-mirror}" + +case "$MODE" in + mirror) + echo ">>> 设置 HF 镜像 (hf-mirror.com)" + export HF_ENDPOINT=https://hf-mirror.com + export PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple + export PIP_TRUSTED_HOST=pypi.tuna.tsinghua.edu.cn + ;; + offline) + echo ">>> 离线模式" + export HF_DATASETS_OFFLINE=1 + export TRANSFORMERS_OFFLINE=1 + export HF_HUB_DISABLE_TELEMETRY=1 + ;; + *) + echo "Usage: source scripts/setup_env.sh [mirror|offline]" + return 1 + ;; +esac + +echo ">>> 环境配置完成 (MODE=$MODE)" diff --git a/scripts/train/agibot_training.sh b/scripts/train/agibot_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..e8702b45e2c739ff3cce16079f1f59a547109453 --- /dev/null +++ b/scripts/train/agibot_training.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# DreamZero AGIbot Training Script +# +# Usage: +# # Set your dataset path and output directory, then run: +# bash scripts/train/agibot_training.sh +# +# Prerequisites: +# - AGIbot dataset in LeRobot format at AGIBOT_DATA_ROOT (state 32, action 22, 3 views: top_head, hand_left, hand_right) +# See docs/DATASET_TO_GEAR_AND_TRAIN.md for conversion instructions +# - Wan2.1-I2V-14B-480P weights (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl +# - DreamZero-AgiBot pretrained checkpoint (for loading LoRA weights before fine-tuning) +# git clone https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot ./checkpoints/DreamZero-AgiBot + +export HYDRA_FULL_ERROR=1 + +# ============ CHANGE THESE VARIABLES ============ +# Dataset path (AGIbot in LeRobot format: state 32, action 22, videos top_head, hand_left, hand_right) +AGIBOT_DATA_ROOT=${AGIBOT_DATA_ROOT:-"./data/agibot_lerobot"} + +# Output directory for training checkpoints +OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero_agibot_lora_5k"} + +# Number of GPUs to use (default: all visible GPUs, so 4-GPU machines use 4 without setting NUM_GPUS) +if [ -z "${NUM_GPUS}" ]; then + NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +fi +NUM_GPUS=${NUM_GPUS:-8} + +# Model weight paths (download from HuggingFace if not already present) +WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.1-I2V-14B-480P not found at $WAN_CKPT_DIR. Downloading from HuggingFace..." + 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 + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +# ================================================ + +# Validate dataset exists +if [ ! -d "$AGIBOT_DATA_ROOT" ]; then + echo "ERROR: AGIbot dataset not found at $AGIBOT_DATA_ROOT" + echo "Set AGIBOT_DATA_ROOT to your LeRobot-format AGIbot dataset (e.g. Dataset/3222_raw_assemble)" + exit 1 +fi + +torchrun --nproc_per_node $NUM_GPUS --standalone groot/vla/experiment/experiment.py \ + report_to=wandb \ + data=dreamzero/agibot_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=2500 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=1 \ + max_steps=5000 \ + 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 \ + agibot_data_root=$AGIBOT_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 diff --git a/scripts/train/droid_training_full_finetune.sh b/scripts/train/droid_training_full_finetune.sh new file mode 100644 index 0000000000000000000000000000000000000000..2e7c2e65e08bc9ae9c7da71bd39a63184a79cdb8 --- /dev/null +++ b/scripts/train/droid_training_full_finetune.sh @@ -0,0 +1,93 @@ +#!/bin/bash +# DreamZero DROID Full Fine-Tuning Script (8x H100, ZeRO-2 + CPU Offload) +# +# Usage: +# bash scripts/train/droid_training_full_finetune.sh +# +# Prerequisites: +# - DROID dataset in LeRobot format at DROID_DATA_ROOT +# Download: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +# Or convert from scratch: see scripts/data/convert_droid.py +# - Wan2.1-I2V-14B-480P weights (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# ============ USER CONFIGURATION ============ +# Dataset path (DROID in LeRobot format) +DROID_DATA_ROOT=${DROID_DATA_ROOT:-"./data/droid_lerobot"} + +# Output directory for training checkpoints +OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero_droid_full_finetune"} + +# Number of GPUs to use (8x H100 for ZeRO-2 + CPU offload full fine-tuning) +NUM_GPUS=${NUM_GPUS:-8} + +# Model weight paths (download from HuggingFace if not already present) +WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.1-I2V-14B-480P not found at $WAN_CKPT_DIR. Downloading from HuggingFace..." + 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 + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +# ================================================ + +# Validate dataset exists +if [ ! -d "$DROID_DATA_ROOT" ]; then + echo "ERROR: DROID dataset not found at $DROID_DATA_ROOT" + echo "Download with: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir $DROID_DATA_ROOT" + exit 1 +fi + +torchrun --nproc_per_node $NUM_GPUS --standalone groot/vla/experiment/experiment.py \ + report_to=none \ + data=dreamzero/droid_relative \ + wandb_project=dreamzero \ + train_architecture=full \ + 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_offload.json" \ + save_steps=1000 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=1 \ + max_steps=100 \ + 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=false \ + max_chunk_size=4 \ + frame_seqlen=880 \ + save_strategy=no \ + droid_data_root=$DROID_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 diff --git a/scripts/train/droid_training_full_finetune_wan21.sh b/scripts/train/droid_training_full_finetune_wan21.sh new file mode 100644 index 0000000000000000000000000000000000000000..a43a8d8d4a323b738f6b33b5c47ed2c0563a3984 --- /dev/null +++ b/scripts/train/droid_training_full_finetune_wan21.sh @@ -0,0 +1,139 @@ +#!/bin/bash +# DreamZero DROID Full Fine-Tuning Script with Wan2.1-I2V-14B-480P backbone +# +# Usage: +# bash scripts/train/droid_training_full_finetune_wan21.sh +# +# Prerequisites: +# - DROID dataset in LeRobot format at DROID_DATA_ROOT +# Download: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +# - Wan2.1-I2V-14B-480P weights (auto-downloaded or pre-downloaded from HuggingFace) +# huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded) +# huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# Repo root (same logic as droid_training_wan22.sh) +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SCRIPT_REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +if [ -n "$DREAMZERO_ROOT" ] && [ -d "$DREAMZERO_ROOT/groot" ]; then + : +elif [ -d "/root/yejink/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/yejink/dreamzero +elif [ -d "/root/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/dreamzero +elif [ -d "$SCRIPT_REPO_ROOT/groot" ]; then + DREAMZERO_ROOT="$SCRIPT_REPO_ROOT" +else + DREAMZERO_ROOT="${DREAMZERO_ROOT:-/root/yejink/dreamzero}" +fi +if [ ! -d "$DREAMZERO_ROOT/groot" ]; then + echo "ERROR: No groot/ under $DREAMZERO_ROOT. Set DREAMZERO_ROOT to the dreamzero repo root that contains groot/." + exit 1 +fi + +# ============ USER CONFIGURATION ============ +DROID_DATA_ROOT=${DROID_DATA_ROOT:-"$DREAMZERO_ROOT/data/droid_lerobot"} +if [ "$DROID_DATA_ROOT" = "./data/droid_lerobot" ]; then + DROID_DATA_ROOT="$DREAMZERO_ROOT/data/droid_lerobot" +fi +# Fixed output dir so Beaker restarts resume from latest checkpoint in this folder. +# For concurrent runs use: OUTPUT_DIR=.../dreamzero_droid_wan21_full_finetune_$(date +%Y%m%d_%H%M%S) +# +# Resume from 100k to 200k: use this same OUTPUT_DIR and run this script with max_steps=200000. +# If the 100k run completed (saved final model), the code will see config.json and skip training. +# To force resume: remove config.json from OUTPUT_DIR so the latest checkpoint-* is used, then run. +OUTPUT_DIR=${OUTPUT_DIR:-"$DREAMZERO_ROOT/checkpoints/dreamzero_droid_wan21_full_finetune"} + +NUM_GPUS=${NUM_GPUS:-4} +PER_DEVICE_BS=${PER_DEVICE_BS:-1} +# Global batch: default = NUM_GPUS * PER_DEVICE_BS. Override for larger effective batch, e.g. GLOBAL_BATCH_SIZE=128 +# (uses gradient accumulation; memory stays same. With 4 GPUs, 128 → grad_accum=32.) +GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-$((NUM_GPUS * PER_DEVICE_BS))} + +# Wan2.1-I2V-14B-480P checkpoint (includes CLIP image encoder, T5 text encoder, and VAE) +WAN21_CKPT_DIR=${WAN21_CKPT_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"$DREAMZERO_ROOT/checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN21_CKPT_DIR" ] || [ -z "$(ls -A "$WAN21_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.1-I2V-14B-480P not found at $WAN21_CKPT_DIR. Downloading from HuggingFace..." + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN21_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +# ================================================ + +if [ ! -d "$DROID_DATA_ROOT" ]; then + echo "ERROR: DROID dataset not found at $DROID_DATA_ROOT" + echo "Download with: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir $DROID_DATA_ROOT" + exit 1 +fi + +EXPERIMENT_PY="$DREAMZERO_ROOT/groot/vla/experiment/experiment.py" +if [ ! -f "$EXPERIMENT_PY" ]; then + echo "ERROR: Not found: $EXPERIMENT_PY" + exit 1 +fi +PYTHON_311="/usr/bin/python3.11" +if [ -x "$PYTHON_311" ]; then + if [ -n "${FIX_NUMPY_IN_SCRIPT:-}" ]; then + "$PYTHON_311" -m pip install "numpy==1.26.4" --force-reinstall -q 2>/dev/null || true + fi + RUN_CMD=( "$PYTHON_311" -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using image Python 3.11: $PYTHON_311" +else + RUN_CMD=( python3 -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using: $(command -v python3)" +fi +cd "$DREAMZERO_ROOT" + +# Full fine-tune: train_architecture=full, save_lora_only=false, ZeRO-2 (+ optional CPU offload) +DEEPSPEED_CFG=${DEEPSPEED_CFG:-zero2_offload} +"${RUN_CMD[@]}" \ + report_to=wandb \ + data=dreamzero/droid_relative \ + wandb_project=dreamzero \ + train_architecture=full \ + 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/${DEEPSPEED_CFG}.json" \ + save_steps=500 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=$PER_DEVICE_BS \ + global_batch_size=$GLOBAL_BATCH_SIZE \ + max_steps=200000 \ + weight_decay=1e-5 \ + save_total_limit=10 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=true \ + dataloader_num_workers=4 \ + image_resolution_width=320 \ + image_resolution_height=176 \ + save_lora_only=false \ + max_chunk_size=4 \ + save_strategy=steps \ + droid_data_root=$DROID_DATA_ROOT \ + dit_version=$WAN21_CKPT_DIR \ + text_encoder_pretrained_path=$WAN21_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \ + image_encoder_pretrained_path=$WAN21_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ + vae_pretrained_path=$WAN21_CKPT_DIR/Wan2.1_VAE.pth \ + tokenizer_path=$TOKENIZER_DIR diff --git a/scripts/train/droid_training_full_finetune_wan22.sh b/scripts/train/droid_training_full_finetune_wan22.sh new file mode 100644 index 0000000000000000000000000000000000000000..fa10a25df3fa09e7c8c3222d17ad6f4e420a8ca1 --- /dev/null +++ b/scripts/train/droid_training_full_finetune_wan22.sh @@ -0,0 +1,147 @@ +#!/bin/bash +# DreamZero DROID Full Fine-Tuning Script with Wan2.2-TI2V-5B backbone +# +# Usage: +# bash scripts/train/droid_training_full_finetune_wan22.sh +# +# Prerequisites: +# - DROID dataset in LeRobot format at DROID_DATA_ROOT +# Download: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +# - Wan2.2-TI2V-5B weights (auto-downloaded or pre-downloaded from HuggingFace) +# huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir ./checkpoints/Wan2.2-TI2V-5B +# - Image encoder (CLIP) from Wan2.1 - Wan2.2-TI2V-5B does not include it +# huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded) +# huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# Repo root (same logic as droid_training_wan22.sh) +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SCRIPT_REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +if [ -n "$DREAMZERO_ROOT" ] && [ -d "$DREAMZERO_ROOT/groot" ]; then + : +elif [ -d "/root/yejink/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/yejink/dreamzero +elif [ -d "/root/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/dreamzero +elif [ -d "$SCRIPT_REPO_ROOT/groot" ]; then + DREAMZERO_ROOT="$SCRIPT_REPO_ROOT" +else + DREAMZERO_ROOT="${DREAMZERO_ROOT:-/root/yejink/dreamzero}" +fi +if [ ! -d "$DREAMZERO_ROOT/groot" ]; then + echo "ERROR: No groot/ under $DREAMZERO_ROOT. Set DREAMZERO_ROOT to the dreamzero repo root that contains groot/." + exit 1 +fi + +# ============ USER CONFIGURATION ============ +DROID_DATA_ROOT=${DROID_DATA_ROOT:-"$DREAMZERO_ROOT/data/droid_lerobot"} +if [ "$DROID_DATA_ROOT" = "./data/droid_lerobot" ]; then + DROID_DATA_ROOT="$DREAMZERO_ROOT/data/droid_lerobot" +fi +# Fixed output dir so Beaker restarts resume from latest checkpoint in this folder. +# For concurrent runs use: OUTPUT_DIR=.../dreamzero_droid_wan22_full_finetune_$(date +%Y%m%d_%H%M%S) +# +# Resume from 100k to 200k: use this same OUTPUT_DIR and run this script with max_steps=200000. +# If the 100k run completed (saved final model), the code will see config.json and skip training. +# To force resume: remove config.json from OUTPUT_DIR so the latest checkpoint-* is used, then run. +OUTPUT_DIR=${OUTPUT_DIR:-"$DREAMZERO_ROOT/checkpoints/dreamzero_droid_wan22_full_finetune"} + +NUM_GPUS=${NUM_GPUS:-4} +PER_DEVICE_BS=${PER_DEVICE_BS:-1} +# Global batch: default = NUM_GPUS * PER_DEVICE_BS. Override for larger effective batch, e.g. GLOBAL_BATCH_SIZE=128 +# (uses gradient accumulation; memory stays same. With 4 GPUs, 128 → grad_accum=32.) +GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-$((NUM_GPUS * PER_DEVICE_BS))} + +# Wan2.2-TI2V-5B checkpoint +WAN22_CKPT_DIR=${WAN22_CKPT_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.2-TI2V-5B"} +IMAGE_ENCODER_DIR=${IMAGE_ENCODER_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"$DREAMZERO_ROOT/checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN22_CKPT_DIR" ] || [ -z "$(ls -A "$WAN22_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.2-TI2V-5B not found at $WAN22_CKPT_DIR. Downloading from HuggingFace..." + huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi + +if [ ! -f "$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then + echo "Image encoder not found. Downloading Wan2.1-I2V-14B-480P (for CLIP only)..." + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$IMAGE_ENCODER_DIR" +fi +# ================================================ + +if [ ! -d "$DROID_DATA_ROOT" ]; then + echo "ERROR: DROID dataset not found at $DROID_DATA_ROOT" + echo "Download with: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir $DROID_DATA_ROOT" + exit 1 +fi + +EXPERIMENT_PY="$DREAMZERO_ROOT/groot/vla/experiment/experiment.py" +if [ ! -f "$EXPERIMENT_PY" ]; then + echo "ERROR: Not found: $EXPERIMENT_PY" + exit 1 +fi +PYTHON_311="/usr/bin/python3.11" +if [ -x "$PYTHON_311" ]; then + if [ -n "${FIX_NUMPY_IN_SCRIPT:-}" ]; then + "$PYTHON_311" -m pip install "numpy==1.26.4" --force-reinstall -q 2>/dev/null || true + fi + RUN_CMD=( "$PYTHON_311" -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using image Python 3.11: $PYTHON_311" +else + RUN_CMD=( python3 -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using: $(command -v python3)" +fi +cd "$DREAMZERO_ROOT" + +# Full fine-tune: train_architecture=full, save_lora_only=false, ZeRO-2 (+ optional CPU offload) +DEEPSPEED_CFG=${DEEPSPEED_CFG:-zero2_offload} +"${RUN_CMD[@]}" \ + report_to=wandb \ + data=dreamzero/droid_relative_wan22 \ + wandb_project=dreamzero \ + train_architecture=full \ + num_frames=33 \ + action_horizon=24 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ + 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/${DEEPSPEED_CFG}.json" \ + save_steps=1000 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=$PER_DEVICE_BS \ + global_batch_size=$GLOBAL_BATCH_SIZE \ + max_steps=200000 \ + weight_decay=1e-5 \ + save_total_limit=10 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=true \ + dataloader_num_workers=4 \ + image_resolution_width=320 \ + image_resolution_height=160 \ + save_lora_only=false \ + max_chunk_size=4 \ + save_strategy=steps \ + droid_data_root=$DROID_DATA_ROOT \ + 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 \ + tokenizer_path=$TOKENIZER_DIR diff --git a/scripts/train/droid_training_lora.sh b/scripts/train/droid_training_lora.sh new file mode 100644 index 0000000000000000000000000000000000000000..978693bc352d899aa099f4c60c422fd36f3c1d39 --- /dev/null +++ b/scripts/train/droid_training_lora.sh @@ -0,0 +1,94 @@ +#!/bin/bash +# DreamZero DROID Training Script +# +# Usage: +# # Set your dataset path and output directory, then run: +# bash scripts/train/droid_training.sh +# +# Prerequisites: +# - DROID dataset in LeRobot format at DROID_DATA_ROOT +# Download: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +# Or convert from scratch: see scripts/data/convert_droid.py +# - Wan2.1-I2V-14B-480P weights (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# ============ USER CONFIGURATION ============ +# Dataset path (DROID in LeRobot format) +DROID_DATA_ROOT=${DROID_DATA_ROOT:-"./data/droid_lerobot"} + +# Output directory for training checkpoints +OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero_droid_lora"} + +# Number of GPUs to use +NUM_GPUS=${NUM_GPUS:-8} + +# Model weight paths (download from HuggingFace if not already present) +WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.1-I2V-14B-480P not found at $WAN_CKPT_DIR. Downloading from HuggingFace..." + 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 + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +# ================================================ + +# Validate dataset exists +if [ ! -d "$DROID_DATA_ROOT" ]; then + echo "ERROR: DROID dataset not found at $DROID_DATA_ROOT" + echo "Download with: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir $DROID_DATA_ROOT" + exit 1 +fi + +torchrun --nproc_per_node $NUM_GPUS --standalone groot/vla/experiment/experiment.py \ + report_to=none \ + data=dreamzero/droid_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-4 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \ + save_steps=1000 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=1 \ + max_steps=100 \ + 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=no \ + droid_data_root=$DROID_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 \ No newline at end of file diff --git a/scripts/train/droid_training_wan22.sh b/scripts/train/droid_training_wan22.sh new file mode 100644 index 0000000000000000000000000000000000000000..da56a65b95ac96f66e0a874445dd9c61a6c67bde --- /dev/null +++ b/scripts/train/droid_training_wan22.sh @@ -0,0 +1,141 @@ +#!/bin/bash +# DreamZero DROID Training Script with Wan2.2-TI2V-5B backbone +# +# Usage: +# bash scripts/train/droid_training_wan22.sh +# +# Prerequisites: +# - DROID dataset in LeRobot format at DROID_DATA_ROOT +# Download: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +# - Wan2.2-TI2V-5B weights (download from HuggingFace) +# huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir ./checkpoints/Wan2.2-TI2V-5B +# - Image encoder (CLIP) from Wan2.1 - Wan2.2-TI2V-5B does not include it +# Option A: huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# Option B: Set IMAGE_ENCODER_DIR to a path containing models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded) +# huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# Repo root: must be a directory that contains groot/ (so experiment.py can be found). +# Beaker/weka uses /root/yejink/dreamzero; image has /root/dreamzero; else script location. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SCRIPT_REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +if [ -n "$DREAMZERO_ROOT" ] && [ -d "$DREAMZERO_ROOT/groot" ]; then + : # keep existing and valid +elif [ -d "/root/yejink/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/yejink/dreamzero +elif [ -d "/root/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/dreamzero +elif [ -d "$SCRIPT_REPO_ROOT/groot" ]; then + DREAMZERO_ROOT="$SCRIPT_REPO_ROOT" +else + DREAMZERO_ROOT="${DREAMZERO_ROOT:-/root/yejink/dreamzero}" +fi +if [ ! -d "$DREAMZERO_ROOT/groot" ]; then + echo "ERROR: No groot/ under $DREAMZERO_ROOT. Set DREAMZERO_ROOT to the dreamzero repo root that contains groot/." + exit 1 +fi + +# ============ USER CONFIGURATION ============ +NUM_GPUS=${NUM_GPUS:-8} +DROID_DATA_ROOT=${DROID_DATA_ROOT:-"$DREAMZERO_ROOT/data/droid_lerobot"} +# If env set the old relative default, resolve to repo root (e.g. Beaker image env) +if [ "$DROID_DATA_ROOT" = "./data/droid_lerobot" ]; then + DROID_DATA_ROOT="$DREAMZERO_ROOT/data/droid_lerobot" +fi +OUTPUT_DIR=${OUTPUT_DIR:-"$DREAMZERO_ROOT/checkpoints/dreamzero_droid_wan22_lora"} + +# Wan2.2-TI2V-5B checkpoint (contains: diffusion weights, T5, VAE) +WAN22_CKPT_DIR=${WAN22_CKPT_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.2-TI2V-5B"} + +# Image encoder: Wan2.2-TI2V-5B does NOT include CLIP - use Wan2.1's or standalone +IMAGE_ENCODER_DIR=${IMAGE_ENCODER_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.1-I2V-14B-480P"} + +TOKENIZER_DIR=${TOKENIZER_DIR:-"$DREAMZERO_ROOT/checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN22_CKPT_DIR" ] || [ -z "$(ls -A "$WAN22_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.2-TI2V-5B not found at $WAN22_CKPT_DIR. Downloading from HuggingFace..." + huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi + +# Image encoder: download Wan2.1 if not present (only need CLIP from it) +if [ ! -f "$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then + echo "Image encoder not found. Downloading Wan2.1-I2V-14B-480P (for CLIP only)..." + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$IMAGE_ENCODER_DIR" +fi +# ================================================ + +# Validate dataset exists +if [ ! -d "$DROID_DATA_ROOT" ]; then + echo "ERROR: DROID dataset not found at $DROID_DATA_ROOT" + echo "Download with: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir $DROID_DATA_ROOT" + exit 1 +fi + +# Use image Python 3.11 when available (Dockerfile installs dreamzero with python3.11 -m pip). +# Absolute path so worker processes open the correct file even if their cwd differs. +EXPERIMENT_PY="$DREAMZERO_ROOT/groot/vla/experiment/experiment.py" +if [ ! -f "$EXPERIMENT_PY" ]; then + echo "ERROR: Not found: $EXPERIMENT_PY" + exit 1 +fi +PYTHON_311="/usr/bin/python3.11" +if [ -x "$PYTHON_311" ]; then + if [ -n "${FIX_NUMPY_IN_SCRIPT:-}" ]; then + "$PYTHON_311" -m pip install "numpy==1.26.4" --force-reinstall -q 2>/dev/null || true + fi + RUN_CMD=( "$PYTHON_311" -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using image Python 3.11: $PYTHON_311" +else + RUN_CMD=( python3 -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using: $(command -v python3)" +fi +cd "$DREAMZERO_ROOT" + +"${RUN_CMD[@]}" \ + report_to=wandb \ + data=dreamzero/droid_relative_wan22 \ + 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_wan22 \ + 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=1000 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=1 \ + max_steps=100 \ + 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 \ + save_lora_only=true \ + max_chunk_size=4 \ + save_strategy=no \ + droid_data_root=$DROID_DATA_ROOT \ + 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 \ + tokenizer_path=$TOKENIZER_DIR diff --git a/scripts/train/libero_training.sh b/scripts/train/libero_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..5aa4b255b0431401c0dac3b8b812530497a84207 --- /dev/null +++ b/scripts/train/libero_training.sh @@ -0,0 +1,89 @@ +#!/bin/bash +# DreamZero LIBERO Full Fine-Tuning Script (Wan2.2-TI2V-5B, GEAR pipeline) +# Prerequisites: LIBERO dataset in GEAR format, Wan2.2-TI2V-5B weights + +export HYDRA_FULL_ERROR=1 + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" + +# Activate DreamZero virtual environment +VENV_DIR="/root/autodl-tmp/venvs/dreamzero" +if [ -f "$VENV_DIR/bin/activate" ]; then + source "$VENV_DIR/bin/activate" + echo "Activated venv: $VENV_DIR" +fi + +# ============ CONFIGURATION ============ +LIBERO_DATA_ROOT=${LIBERO_DATA_ROOT:-"/root/autodl-tmp/data/libero_gear"} +OUTPUT_DIR=${OUTPUT_DIR:-"$REPO_ROOT/checkpoints/dreamzero_libero_full"} + +if [ -z "${NUM_GPUS:-}" ]; then + NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +fi +PER_DEVICE_BS=${PER_DEVICE_BS:-1} + +CKPT_DIR="${CHECKPOINT_DIR:-/root/autodl-tmp/checkpoints}" +WAN22_DIR="$CKPT_DIR/Wan2.2-TI2V-5B" +TOKENIZER_DIR="$CKPT_DIR/umt5-xxl" +CLIP_DIR="$CKPT_DIR/clip-encoder" +# ======================================== + +# Auto-download weights if missing (with academic proxy) +source /etc/network_turbo 2>/dev/null +if [ ! -d "$WAN22_DIR" ] || [ -z "$(ls -A "$WAN22_DIR" 2>/dev/null)" ]; then + echo "Downloading Wan2.2-TI2V-5B..." + huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_DIR" +fi +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "Downloading umt5-xxl..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +if [ ! -d "$LIBERO_DATA_ROOT" ]; then + echo "ERROR: LIBERO dataset not found at $LIBERO_DATA_ROOT"; exit 1 +fi + +cd "$REPO_ROOT" + +# ZeRO config +if [ "$NUM_GPUS" -le 2 ]; then + DEEPSPEED_CFG=${DEEPSPEED_CFG:-"groot/vla/configs/deepspeed/zero2_offload.json"} +else + DEEPSPEED_CFG=${DEEPSPEED_CFG:-"groot/vla/configs/deepspeed/zero2.json"} +fi + +torchrun --standalone --nproc_per_node "$NUM_GPUS" \ + groot/vla/experiment/experiment.py \ + report_to="${REPORT_TO:-none}" \ + data=dreamzero/libero_sim_relative \ + wandb_project="${WANDB_PROJECT:-dreamzero-libero-sft}" \ + train_architecture=full \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frames=33 action_horizon=12 num_views=1 \ + num_frame_per_block=2 num_action_per_block=24 \ + num_state_per_block=1 max_chunk_size=4 frame_seqlen=50 \ + image_resolution_width=320 image_resolution_height=160 \ + seed=42 \ + training_args.learning_rate="${LR:-1e-5}" \ + training_args.deepspeed="$DEEPSPEED_CFG" \ + save_steps="${SAVE_STEPS:-2000}" \ + training_args.warmup_ratio=0.05 \ + output_dir="$OUTPUT_DIR" \ + per_device_train_batch_size="$PER_DEVICE_BS" \ + max_steps="${MAX_STEPS:-200000}" \ + weight_decay=1e-5 save_total_limit=2 \ + upload_checkpoints=false bf16=true tf32=true eval_bf16=true \ + dataloader_pin_memory=true dataloader_num_workers=4 \ + save_lora_only=false save_strategy=steps \ + libero_gear_root="$LIBERO_DATA_ROOT" \ + dit_version="$WAN22_DIR" \ + text_encoder_pretrained_path="$WAN22_DIR/models_t5_umt5-xxl-enc-bf16.pth" \ + image_encoder_pretrained_path="$CLIP_DIR/models_clip-open-clip-xlm-roberta-large-vit-huge-14.pth" \ + vae_pretrained_path="$WAN22_DIR/Wan2.2_VAE.pth" \ + tokenizer_path="$TOKENIZER_DIR" + +echo "============================================" +echo " LIBERO training finished! Output: $OUTPUT_DIR" +echo "============================================" diff --git a/scripts/train/manifeel_training.sh b/scripts/train/manifeel_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..da2122cc54dd4f557296030c00f4285667df9d3f --- /dev/null +++ b/scripts/train/manifeel_training.sh @@ -0,0 +1,92 @@ +#!/bin/bash +# DreamZero ManiFeel Full Fine-Tuning Script (Wan2.2-TI2V-5B, custom ManiFeelDataset) +# +# Usage: +# bash scripts/train/manifeel_training.sh +# +# Prerequisites: +# - ManiFeel data extracted at MANIFEEL_DATA_DIR +# (zarr format, each task dir has .zgroup, data/{front,side,wrist,state,action}, meta/episode_ends) +# - Wan2.2-TI2V-5B weights at CHECKPOINT_DIR +# - umt5-xxl tokenizer at CHECKPOINT_DIR +# - DreamZero-DROID checkpoint (for CLIP encoder) at CHECKPOINT_DIR + +export HYDRA_FULL_ERROR=1 + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" + +# ============ CONFIGURATION ============ +MANIFEEL_DATA_DIR=${MANIFEEL_DATA_DIR:-"/root/autodl-tmp/tmp/manifeel_extracted"} +OUTPUT_DIR=${OUTPUT_DIR:-"$REPO_ROOT/checkpoints/dreamzero_manifeel_full"} + +if [ -z "${NUM_GPUS:-}" ]; then + NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +fi +PER_DEVICE_BS=${PER_DEVICE_BS:-1} + +CKPT_DIR="${CHECKPOINT_DIR:-/root/autodl-tmp/checkpoints}" +WAN22_DIR="$CKPT_DIR/Wan2.2-TI2V-5B" +TOKENIZER_DIR="$CKPT_DIR/umt5-xxl" +CLIP_DIR="$CKPT_DIR/DreamZero-DROID" +# ======================================== + +# Auto-download weights if missing (with academic proxy) +source /etc/network_turbo 2>/dev/null +if [ ! -d "$WAN22_DIR" ] || [ -z "$(ls -A "$WAN22_DIR" 2>/dev/null)" ]; then + echo "Downloading Wan2.2-TI2V-5B..." + huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_DIR" +fi +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "Downloading umt5-xxl..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +if [ ! -d "$MANIFEEL_DATA_DIR" ]; then + echo "ERROR: ManiFeel data not found at $MANIFEEL_DATA_DIR"; exit 1 +fi + +cd "$REPO_ROOT" + +# ZeRO config: must override with zero2.json (no offload) when using 1-2 GPUs +if [ "$NUM_GPUS" -le 2 ]; then + DEEPSPEED_CFG=${DEEPSPEED_CFG:-"groot/vla/configs/deepspeed/zero2_offload.json"} +else + DEEPSPEED_CFG=${DEEPSPEED_CFG:-"groot/vla/configs/deepspeed/zero2.json"} +fi + +torchrun --standalone --nproc_per_node "$NUM_GPUS" \ + groot/vla/experiment/experiment.py \ + report_to="${REPORT_TO:-none}" \ + data=dreamzero/manifeel \ + wandb_project="${WANDB_PROJECT:-dreamzero-manifeel-sft}" \ + train_architecture=full \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frames=12 action_horizon=12 state_horizon=1 num_views=3 \ + num_frame_per_block=2 num_action_per_block=12 \ + num_state_per_block=1 max_chunk_size=4 frame_seqlen=50 \ + image_resolution_width=320 image_resolution_height=160 \ + max_state_dim=44 max_action_dim=32 \ + seed=42 \ + training_args.learning_rate="${LR:-1e-5}" \ + training_args.deepspeed="$DEEPSPEED_CFG" \ + save_steps="${SAVE_STEPS:-2000}" \ + training_args.warmup_ratio=0.05 \ + output_dir="$OUTPUT_DIR" \ + per_device_train_batch_size="$PER_DEVICE_BS" \ + max_steps="${MAX_STEPS:-200000}" \ + weight_decay=1e-5 save_total_limit=2 \ + upload_checkpoints=false bf16=true tf32=true eval_bf16=true \ + dataloader_pin_memory=true dataloader_num_workers=4 \ + save_lora_only=false save_strategy=steps \ + manifeel_dataset_dir="$MANIFEEL_DATA_DIR" \ + dit_version="$WAN22_DIR" \ + text_encoder_pretrained_path="$WAN22_DIR/models_t5_umt5-xxl-enc-bf16.pth" \ + image_encoder_pretrained_path="$CLIP_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \ + vae_pretrained_path="$WAN22_DIR/Wan2.2_VAE.pth" \ + tokenizer_path="$TOKENIZER_DIR" + +echo "============================================" +echo " ManiFeel training finished! Output: $OUTPUT_DIR" +echo "============================================" diff --git a/scripts/train/quick_start.sh b/scripts/train/quick_start.sh new file mode 100644 index 0000000000000000000000000000000000000000..0a60405fb5dc8fbc526ce31f2344473f5c924679 --- /dev/null +++ b/scripts/train/quick_start.sh @@ -0,0 +1,144 @@ +#!/bin/bash +# DreamZero 单节点快速启动脚本 (5 步验证) +# +# 用法: +# bash scripts/train/quick_start.sh libero # 单卡验证 LIBERO +# bash scripts/train/quick_start.sh manifeel 4 # 4 卡 ManiFeel +# bash scripts/train/quick_start.sh robotwin 2 # 2 卡 RoboTwin +# +# 前置条件: 权重文件已下载到 checkpoints/ 目录 + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" + +BENCHMARK="${1:-libero}" +NUM_GPUS="${2:-1}" +OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/output/quickstart_${BENCHMARK}}" +PER_DEVICE_BS="${PER_DEVICE_BS:-1}" +MAX_STEPS="${MAX_STEPS:-5}" + +echo "============================================" +echo " DreamZero Quick Start" +echo "============================================" +echo " Benchmark: $BENCHMARK" +echo " GPUs: $NUM_GPUS" +echo " Max steps: $MAX_STEPS" +echo " Output: $OUTPUT_DIR" +echo "============================================" + +# ============ 验证权重文件 ============ +CKPT_DIR="${CHECKPOINT_DIR:-$REPO_ROOT/checkpoints}" +WAN22_DIR="$CKPT_DIR/Wan2.2-TI2V-5B" +TOKENIZER_DIR="$CKPT_DIR/umt5-xxl" +CLIP_DIR="$CKPT_DIR/clip-encoder" + +if [ ! -d "$WAN22_DIR" ]; then + echo "ERROR: 未找到 Wan2.2 权重: $WAN22_DIR" + echo "请先运行: bash scripts/data/download_checkpoints.sh" + exit 1 +fi +if [ ! -d "$TOKENIZER_DIR" ]; then + echo "ERROR: 未找到 tokenizer: $TOKENIZER_DIR" + echo "请先运行: bash scripts/data/download_checkpoints.sh" + exit 1 +fi +if [ ! -f "$CLIP_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then + echo "ERROR: 未找到 CLIP encoder: $CLIP_DIR" + echo "请先运行: bash scripts/data/download_checkpoints.sh" + exit 1 +fi + +# ============ 数据集选择 ============ +case "$BENCHMARK" in + libero) + DATA_CFG="dreamzero/libero" + NUM_FRAMES=12 + ACTION_HORIZON=12 + NUM_VIEWS=1 + MAX_STATE_DIM=44 + MAX_ACTION_DIM=32 + NUM_FRAME_PER_BLOCK=2 + NUM_ACTION_PER_BLOCK=12 + LR=1e-5 + ;; + manifeel) + DATA_CFG="dreamzero/manifeel" + NUM_FRAMES=12 + ACTION_HORIZON=12 + NUM_VIEWS=3 + MAX_STATE_DIM=44 + MAX_ACTION_DIM=32 + NUM_FRAME_PER_BLOCK=2 + NUM_ACTION_PER_BLOCK=12 + LR=1e-5 + ;; + robotwin) + DATA_CFG="dreamzero/robotwin" + NUM_FRAMES=12 + ACTION_HORIZON=12 + NUM_VIEWS=1 + MAX_STATE_DIM=44 + MAX_ACTION_DIM=32 + NUM_FRAME_PER_BLOCK=2 + NUM_ACTION_PER_BLOCK=12 + LR=1e-5 + ;; + *) + echo "ERROR: 未知 benchmark: $BENCHMARK (可选: libero, manifeel, robotwin)" + exit 1 + ;; +esac + +# ============ 自动选择 DeepSpeed 配置 ============ +if [ "$NUM_GPUS" -le 2 ]; then + DEEPSPEED_CFG="groot/vla/configs/deepspeed/zero2.json" +elif [ "$NUM_GPUS" -le 8 ]; then + DEEPSPEED_CFG="groot/vla/configs/deepspeed/zero2_offload.json" +else + DEEPSPEED_CFG="groot/vla/configs/deepspeed/zero3_multinode.json" +fi + +cd "$REPO_ROOT" + +torchrun --standalone --nproc_per_node "$NUM_GPUS" \ + groot/vla/experiment/experiment.py \ + report_to=none \ + data="$DATA_CFG" \ + train_architecture=full \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frames="$NUM_FRAMES" \ + action_horizon="$ACTION_HORIZON" \ + num_views="$NUM_VIEWS" \ + max_state_dim="$MAX_STATE_DIM" \ + max_action_dim="$MAX_ACTION_DIM" \ + num_frame_per_block="$NUM_FRAME_PER_BLOCK" \ + num_action_per_block="$NUM_ACTION_PER_BLOCK" \ + num_state_per_block=1 \ + max_chunk_size=4 \ + frame_seqlen=50 \ + image_resolution_width=320 \ + image_resolution_height=160 \ + per_device_train_batch_size="$PER_DEVICE_BS" \ + max_steps="$MAX_STEPS" \ + save_strategy=no \ + optim=adamw_bnb_8bit \ + training_args.learning_rate="$LR" \ + training_args.deepspeed="$DEEPSPEED_CFG" \ + training_args.bf16=true \ + training_args.tf32=true \ + training_args.eval_bf16=true \ + output_dir="$OUTPUT_DIR" \ + dit_version="$WAN22_DIR" \ + text_encoder_pretrained_path="$WAN22_DIR/models_t5_umt5-xxl-enc-bf16.pth" \ + image_encoder_pretrained_path="$CLIP_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \ + vae_pretrained_path="$WAN22_DIR/Wan2.2_VAE.pth" \ + tokenizer_path="$TOKENIZER_DIR" + +echo "" +echo "============================================" +echo " Quick start 完成! (benchmark: $BENCHMARK)" +echo "============================================" diff --git a/scripts/train/robotwin_training.sh b/scripts/train/robotwin_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..e60c17487498505350c8cda0d370a5dc464d4c78 --- /dev/null +++ b/scripts/train/robotwin_training.sh @@ -0,0 +1,87 @@ +#!/bin/bash +# DreamZero RoboTwin Full Fine-Tuning Script (Wan2.2-TI2V-5B, custom RobotWinDataset) +export HYDRA_FULL_ERROR=1 + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" + +# Activate DreamZero virtual environment +VENV_DIR="/root/autodl-tmp/venvs/dreamzero" +if [ -f "$VENV_DIR/bin/activate" ]; then + source "$VENV_DIR/bin/activate" +fi + +# ============ CONFIGURATION ============ +ROBOTWIN_DATA_DIR=${ROBOTWIN_DATA_DIR:-"/root/autodl-tmp/data/robotwin_gear"} +OUTPUT_DIR=${OUTPUT_DIR:-"$REPO_ROOT/checkpoints/dreamzero_robotwin_full"} + +if [ -z "${NUM_GPUS:-}" ]; then + NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +fi +PER_DEVICE_BS=${PER_DEVICE_BS:-1} + +CKPT_DIR="${CHECKPOINT_DIR:-/root/autodl-tmp/checkpoints}" +WAN22_DIR="$CKPT_DIR/Wan2.2-TI2V-5B" +TOKENIZER_DIR="$CKPT_DIR/umt5-xxl" +CLIP_DIR="$CKPT_DIR/DreamZero-DROID" +# ======================================== + +# Auto-download weights if missing +source /etc/network_turbo 2>/dev/null +if [ ! -d "$WAN22_DIR" ] || [ -z "$(ls -A "$WAN22_DIR" 2>/dev/null)" ]; then + echo "Downloading Wan2.2-TI2V-5B..." + huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_DIR" +fi +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "Downloading umt5-xxl..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +if [ ! -d "$ROBOTWIN_DATA_DIR" ]; then + echo "ERROR: RoboTwin data not found at $ROBOTWIN_DATA_DIR"; exit 1 +fi + +cd "$REPO_ROOT" + +# ZeRO config +if [ "$NUM_GPUS" -le 2 ]; then + DEEPSPEED_CFG=${DEEPSPEED_CFG:-"groot/vla/configs/deepspeed/zero2_offload.json"} +else + DEEPSPEED_CFG=${DEEPSPEED_CFG:-"groot/vla/configs/deepspeed/zero2.json"} +fi + +torchrun --standalone --nproc_per_node "$NUM_GPUS" \ + groot/vla/experiment/experiment.py \ + report_to="${REPORT_TO:-none}" \ + data=dreamzero/robotwin \ + wandb_project="${WANDB_PROJECT:-dreamzero-robotwin-sft}" \ + train_architecture=full \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ + model/dreamzero/transform=dreamzero_cotrain \ + 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_chunk_size=4 frame_seqlen=50 \ + image_resolution_width=320 image_resolution_height=160 \ + max_state_dim=44 max_action_dim=32 \ + seed=42 \ + training_args.learning_rate="${LR:-1e-5}" \ + training_args.deepspeed="$DEEPSPEED_CFG" \ + save_steps="${SAVE_STEPS:-2000}" \ + training_args.warmup_ratio=0.05 \ + output_dir="$OUTPUT_DIR" \ + per_device_train_batch_size="$PER_DEVICE_BS" \ + max_steps="${MAX_STEPS:-200000}" \ + weight_decay=1e-5 save_total_limit=2 \ + upload_checkpoints=false bf16=true tf32=true eval_bf16=true \ + dataloader_pin_memory=true dataloader_num_workers=4 \ + save_lora_only=false save_strategy=steps \ + robotwin_dataset_dir="$ROBOTWIN_DATA_DIR" \ + dit_version="$WAN22_DIR" \ + text_encoder_pretrained_path="$WAN22_DIR/models_t5_umt5-xxl-enc-bf16.pth" \ + image_encoder_pretrained_path="$CLIP_DIR/models_clip-open-clip-xlm-roberta-large-vit-huge-14.pth" \ + vae_pretrained_path="$WAN22_DIR/Wan2.2_VAE.pth" \ + tokenizer_path="$TOKENIZER_DIR" + +echo "============================================" +echo " RoboTwin training finished! Output: $OUTPUT_DIR" +echo "============================================" diff --git a/scripts/train/yam_training.sh b/scripts/train/yam_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..49f83079a585b4b1bf71ac8dea3660a9840bae85 --- /dev/null +++ b/scripts/train/yam_training.sh @@ -0,0 +1,103 @@ +#!/bin/bash +# DreamZero YAM Training Script +# +# Usage: +# # Set your dataset path and output directory, then run: +# bash scripts/train/yam_training.sh +# +# Prerequisites: +# - YAM dataset in LeRobot format at YAM_DATA_ROOT (state 14, action 14, 3 views: top, left, right) +# meta/embodiment.json must have "embodiment_tag": "yam" +# modality: state (left_joint_pos, left_gripper_pos, right_joint_pos, right_gripper_pos), +# action (same keys), video (top, left, right), annotation.human.task_description +# - Wan2.1-I2V-14B-480P weights (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl +# - DreamZero-AgiBot pretrained checkpoint (for loading LoRA weights before fine-tuning) +# git clone https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot ./checkpoints/DreamZero-AgiBot + +export HYDRA_FULL_ERROR=1 + +# ============ CHANGE THESE VARIABLES ============ +# Dataset path (YAM in LeRobot format: state 14, action 14, videos top, left, right) +YAM_DATA_ROOT=${YAM_DATA_ROOT:-"./data/yam_lerobot"} + +# Output directory for training checkpoints +OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero_yam_lora_dz_pretrained_100k_folding"} + +# Number of GPUs to use (default: all visible GPUs, so 4-GPU machines use 4 without setting NUM_GPUS) +if [ -z "${NUM_GPUS}" ]; then + NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +fi +NUM_GPUS=${NUM_GPUS:-8} + +# Model weight paths (download from HuggingFace if not already present) +WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.1-I2V-14B-480P not found at $WAN_CKPT_DIR. Downloading from HuggingFace..." + 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 + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +# ================================================ + +# Validate dataset exists +if [ ! -d "$YAM_DATA_ROOT" ]; then + echo "ERROR: YAM dataset not found at $YAM_DATA_ROOT" + echo "Set YAM_DATA_ROOT to your LeRobot-format YAM dataset (meta/embodiment.json with embodiment_tag: yam)" + exit 1 +fi + +torchrun --nproc_per_node $NUM_GPUS --standalone groot/vla/experiment/experiment.py \ + report_to=wandb \ + data=dreamzero/yam_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 \ + yam_data_root=$YAM_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 diff --git a/socket_test_optimized_AR.py b/socket_test_optimized_AR.py new file mode 100644 index 0000000000000000000000000000000000000000..5b93e05e9a8f7f873b81541ed8eddf4cb47f83ad --- /dev/null +++ b/socket_test_optimized_AR.py @@ -0,0 +1,836 @@ +import dataclasses +import logging +import socket +import asyncio +import os +import http +import logging +import time +import traceback +import torch +import tyro +from einops import rearrange +import datetime + +from groot.vla.model.n1_5.sim_policy import GrootSimPolicy +from groot.vla.data.schema import EmbodimentTag +import imageio +import numpy as np + +from openpi_client import base_policy as _base_policy +from openpi_client import msgpack_numpy +import websockets.asyncio.server as _server +import websockets.frames +from tianshou.data import Batch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh + +# Use roboarena policy server interface +from eval_utils.policy_server import WebsocketPolicyServer as RoboarenaServer +from eval_utils.policy_server import PolicyServerConfig + +logger = logging.getLogger(__name__) + +@dataclasses.dataclass +class Args: + port: int = 8000 + timeout_seconds: int = 50000 # 10 hours default, configurable + model_path: str = "./checkpoints/dreamzero" + enable_dit_cache: bool = False + index: int = 0 + max_chunk_size: int | None = None # If None, use config value. Otherwise override max_chunk_size for inference. + + +class ARDroidRoboarenaPolicy: + """Wrapper policy that implements roboarena.policy.BasePolicy interface for AR_droid. + + Handles: + - Observation format conversion (roboarena -> AR_droid format) + - Frame accumulation across calls (roboarena sends single frames, AR_droid expects multi-frame video) + - Action format conversion (AR_droid dict -> roboarena array format) + - Distributed inference coordination + """ + + # Number of frames to accumulate after the first call + FRAMES_PER_CHUNK = 4 + + def __init__( + self, + groot_policy: GrootSimPolicy, + signal_group: dist.ProcessGroup, + output_dir: str | None = None, + ) -> None: + self._policy = groot_policy + self._signal_group = signal_group + self._output_dir = output_dir + + # Frame buffers for accumulation (per camera view) + self._frame_buffers: dict[str, list[np.ndarray]] = { + "video.exterior_image_1_left": [], + "video.exterior_image_2_left": [], + "video.wrist_image_left": [], + } + self._call_count = 0 + self._is_first_call = True + + # Session tracking - reset state when new session starts + self._current_session_id: str | None = None + + # Video across time for saving (similar to original server) + self.video_across_time = [] + self._msg_index = 0 + + # Create output directory if specified + if self._output_dir: + os.makedirs(self._output_dir, exist_ok=True) + + def _convert_observation(self, obs: dict) -> dict: + """Convert roboarena observation format to AR_droid format. + + Roboarena format: + - observation/exterior_image_0_left: (H, W, 3) single frame + - observation/exterior_image_1_left: (H, W, 3) single frame + - observation/wrist_image_left: (H, W, 3) single frame + - observation/joint_position: (7,) + - observation/gripper_position: (1,) + - prompt: str + + AR_droid format: + - video.exterior_image_1_left: (T, H, W, 3) multi-frame + - video.exterior_image_2_left: (T, H, W, 3) multi-frame + - video.wrist_image_left: (T, H, W, 3) multi-frame + - state.joint_position: (1, 7) + - state.gripper_position: (1, 1) + - annotation.language.action_text: str + """ + converted = {} + + # Map image keys (roboarena uses 0-indexed, AR_droid uses 1-indexed) + image_key_mapping = { + "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", + } + + # Accumulate frames for each camera view + for roboarena_key, droid_key in image_key_mapping.items(): + if roboarena_key in obs: + data = obs[roboarena_key] + if isinstance(data, np.ndarray): + if data.ndim == 4: + # Multiple frames (T, H, W, 3) + self._frame_buffers[droid_key].extend(list(data)) + else: + # Single frame (H, W, 3) + self._frame_buffers[droid_key].append(data) + + # Determine how many frames to use + if self._is_first_call: + # First call: use only 1 frame + num_frames = 1 + else: + # Subsequent calls: use exactly FRAMES_PER_CHUNK frames + num_frames = self.FRAMES_PER_CHUNK + + # Build video tensors from accumulated frames + for droid_key, buffer in self._frame_buffers.items(): + if len(buffer) > 0: + if len(buffer) >= num_frames: + # Take the last num_frames frames + frames_to_use = buffer[-num_frames:] + else: + # Pad by repeating the first frame to reach num_frames + frames_to_use = buffer.copy() + while len(frames_to_use) < num_frames: + # Prepend the first frame to pad + frames_to_use.insert(0, buffer[0]) + # Stack to (T, H, W, C) + video = np.stack(frames_to_use, axis=0) + converted[droid_key] = video + + # Convert state observations + if "observation/joint_position" in obs: + joint_pos = obs["observation/joint_position"] + # Reshape to (1, 7) if needed + if joint_pos.ndim == 1: + joint_pos = joint_pos.reshape(1, -1) + converted["state.joint_position"] = joint_pos.astype(np.float64) + else: + converted["state.joint_position"] = np.zeros((1, 7), dtype=np.float64) + + if "observation/gripper_position" in obs: + gripper_pos = obs["observation/gripper_position"] + # Reshape to (1, 1) if needed + if gripper_pos.ndim == 1: + gripper_pos = gripper_pos.reshape(1, -1) + converted["state.gripper_position"] = gripper_pos.astype(np.float64) + else: + converted["state.gripper_position"] = np.zeros((1, 1), dtype=np.float64) + + # Convert prompt + if "prompt" in obs: + converted["annotation.language.action_text"] = obs["prompt"] + else: + converted["annotation.language.action_text"] = "" + + return converted + + def _convert_action(self, action_dict: dict) -> np.ndarray: + """Convert AR_droid action dict to roboarena action array. + + AR_droid format: + - action.joint_position: (N, 7) + - action.gripper_position: (N,) or (N, 1) + + Roboarena format: + - action: (N, 8) - 7 joint positions + 1 gripper + """ + joint_action = None + gripper_action = None + + # Extract actions from dict + for key, value in action_dict.items(): + if "joint_position" in key: + joint_action = value + elif "gripper_position" in key or "gripper" in key: + gripper_action = value + + if joint_action is None: + # Fallback: return zeros + return np.zeros((1, 8), dtype=np.float32) + + # Convert to numpy if tensor + if isinstance(joint_action, torch.Tensor): + joint_action = joint_action.cpu().numpy() + + # Ensure 2D shape (N, 7) + if joint_action.ndim == 1: + joint_action = joint_action.reshape(1, -1) + + N = joint_action.shape[0] + + # Handle gripper action + if gripper_action is not None: + if isinstance(gripper_action, torch.Tensor): + gripper_action = gripper_action.cpu().numpy() + # Reshape to (N, 1) if needed + if gripper_action.ndim == 1: + gripper_action = gripper_action.reshape(-1, 1) + elif gripper_action.ndim == 0: + gripper_action = gripper_action.reshape(1, 1) + else: + gripper_action = np.zeros((N, 1), dtype=np.float32) + + # Concatenate: (N, 7) + (N, 1) -> (N, 8) + action = np.concatenate([joint_action, gripper_action], axis=-1).astype(np.float32) + + return action + + def _broadcast_batch_to_workers(self, obs: dict) -> None: + """Broadcast batch data from rank 0 to all other ranks.""" + import pickle + + # Serialize the obs + serialized = pickle.dumps(obs) + data_size = len(serialized) + + # Broadcast size first + size_tensor = torch.tensor([data_size], dtype=torch.int64, device='cuda') + dist.broadcast(size_tensor, src=0) + + # Broadcast data + data_tensor = torch.frombuffer(serialized, dtype=torch.uint8).cuda() + dist.broadcast(data_tensor, src=0) + + def infer(self, obs: dict) -> np.ndarray: + """Infer actions from observations. + + Args: + obs: Observation dict in roboarena format + + Returns: + action: (N, 8) action array + """ + # Check for session change - reset state if new session + session_id = obs.get("session_id", None) + if session_id is not None and session_id != self._current_session_id: + if self._current_session_id is not None: + logger.info(f"Session changed from '{self._current_session_id}' to '{session_id}', resetting state") + # Reset state for new session + self._reset_state() + else: + logger.info(f"New session started: '{session_id}'") + self._current_session_id = session_id + + self._msg_index += 1 + self._call_count += 1 + + # Convert observation format + converted_obs = self._convert_observation(obs) + + # Signal workers to continue (0 = continue) + signal_tensor = torch.zeros(1, dtype=torch.int32, device='cpu') + dist.broadcast(signal_tensor, src=0, group=self._signal_group) + + # Broadcast obs to workers + self._broadcast_batch_to_workers(converted_obs) + + # Create batch for policy + batch = Batch(obs=converted_obs) + + # Distributed forward pass + dist.barrier() + with torch.no_grad(): + result_batch, video_pred = self._policy.lazy_joint_forward_causal(batch) + dist.barrier() + + # Store video predictions for potential saving + self.video_across_time.append(video_pred) + + # Extract and convert action + action_chunk_dict = result_batch.act + + # Convert Batch to dict + action_dict = {} + 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) + + # Update first call flag + if self._is_first_call: + self._is_first_call = False + + return action + + def _reset_state(self, save_video: bool = True) -> None: + """Internal method to reset policy state. + + Args: + save_video: Whether to save accumulated video before reset. + """ + # Optionally save accumulated video before reset + if save_video and len(self.video_across_time) > 0 and self._output_dir: + try: + frame_list = [] + video_across_time_cat = torch.cat(self.video_across_time, dim=2) + frames = self._policy.trained_model.action_head.vae.decode( + video_across_time_cat, + tiled=self._policy.trained_model.action_head.tiled, + tile_size=(self._policy.trained_model.action_head.tile_size_height, self._policy.trained_model.action_head.tile_size_width), + tile_stride=(self._policy.trained_model.action_head.tile_stride_height, self._policy.trained_model.action_head.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C") + frames = frames[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + for frame in frames: + frame_list.append(frame) + + if len(frame_list) > 0: + sample_frame = frame_list[0] + if len(sample_frame.shape) == 3 and sample_frame.shape[2] in [1, 3, 4]: + save_dir = self._output_dir + os.makedirs(save_dir, exist_ok=True) + all_mp4_files = [f for f in os.listdir(save_dir) if f.endswith(".mp4")] + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + num_frames = len(frame_list) + n = (num_frames - 1) // 8 + output_path = os.path.join(save_dir, f'{len(all_mp4_files):06}_{timestamp}_n{n}.mp4') + imageio.mimsave(output_path, frame_list, fps=5, codec='libx264') + logger.info(f"Saved video on reset to: {output_path}") + except Exception as e: + logger.warning(f"Failed to save video on reset: {e}") + + # Clear frame buffers + for key in self._frame_buffers: + self._frame_buffers[key] = [] + + self._call_count = 0 + self._is_first_call = True + self.video_across_time = [] + + def reset(self, reset_info: dict) -> None: + """Reset the policy state for a new episode. + + Clears frame buffers and resets call count. + """ + self._reset_state(save_video=True) + + +class WebsocketPolicyServer: + """Serves a policy using the websocket protocol. See websocket_client_policy.py for a client implementation. + Currently only implements the `load` and `infer` methods. + """ + + def __init__( + self, + policy: _base_policy.BasePolicy, + host: str = "0.0.0.0", + port: int | None = None, + metadata: dict | None = None, + output_dir: str | None = None, + signal_group: dist.ProcessGroup | None = None, + ) -> None: + self._policy = policy + self._host = host + self._port = port + self._metadata = metadata or {} + self._output_dir = output_dir + logging.getLogger("websockets.server").setLevel(logging.INFO) + self.video_across_time = [] + self._msg_index = 0 + self._signal_group = signal_group + # Create output directory if specified + if self._output_dir: + os.makedirs(self._output_dir, exist_ok=True) + os.makedirs(os.path.join(self._output_dir, "inputs"), exist_ok=True) + + def _save_input_obs(self, obs: dict) -> None: + """Save incoming observation images per message. + + Expected format: THWC (Time, Height, Width, Channel) with 4 frames. + Saves each frame as a separate PNG image: HWC format (uint8). + + Directory structure: + output_dir/inputs/{msg_index:06d}_{timestamp}/{obs_key}/f{frame_idx:02d}.png + """ + if not self._output_dir: + return + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + base_dir = os.path.join(self._output_dir, "inputs", f"{self._msg_index:06d}_{timestamp}") + try: + os.makedirs(base_dir, exist_ok=True) + except Exception: + return + + for key in ("video.exterior_image_1_left", "video.exterior_image_2_left", "video.wrist_image_left"): + if key not in obs: + continue + value = obs[key] + try: + # Convert to numpy if tensor + if isinstance(value, torch.Tensor): + arr = value.detach().cpu().numpy() + else: + arr = np.asarray(value) + + # Expected format: THWC (Time, Height, Width, Channel) + if arr.ndim != 4: + logger.warning(f"obs key '{key}' has shape {arr.shape}, expected 4D (T,H,W,C)") + continue + + # arr is (T, H, W, C) + T, H, W, C = arr.shape + + # Normalize to uint8 + if arr.dtype == np.uint8: + frames_u8 = arr + else: + f = arr.astype(np.float32) + # Common conventions: [-1,1] or [0,1] + min_val = float(np.nanmin(f)) + max_val = float(np.nanmax(f)) + if min_val >= -1.1 and max_val <= 1.1: + # Assume [-1,1] range + frames_u8 = ((f + 1.0) * 127.5).clip(0, 255).astype(np.uint8) + else: + # Min-max scaling + denom = (max_val - min_val) if (max_val - min_val) > 1e-6 else 1.0 + frames_u8 = ((f - min_val) / denom * 255.0).clip(0, 255).astype(np.uint8) + + # Save each frame: frames_u8[i] is (H, W, C) + key_dir = os.path.join(base_dir, key.replace("/", "_")) + os.makedirs(key_dir, exist_ok=True) + for frame_idx in range(T): + frame = frames_u8[frame_idx] # (H, W, C) + # Handle grayscale (H, W) -> (H, W, 1) + if frame.ndim == 2: + frame = np.expand_dims(frame, axis=-1) + imageio.imwrite(os.path.join(key_dir, f"f{frame_idx:02d}.png"), frame) + + except Exception as e: + logger.warning(f"Failed to save obs key '{key}': {e}") + continue + + + + def serve_forever(self, rank: int = 0) -> None: + asyncio.run(self.run(rank)) + + async def run(self, rank: int = 0): + if rank == 0: + async with _server.serve( + self._handler, + self._host, + self._port, + compression=None, + max_size=None, + process_request=_health_check, + ping_interval=None, + ) as server: + await server.serve_forever() + else: + # Non-rank-0 processes run a worker loop + await self._worker_loop() + + async def _worker_loop(self): + """Worker loop for non-rank-0 processes to participate in distributed inference.""" + logger.info(f"Worker loop started for rank {dist.get_rank()}") + signal_tensor = torch.zeros(1, dtype=torch.int32, device='cpu') + while True: + try: + # Wait for obs broadcast from rank 0 + # Create a dummy obs dict structure - will be filled by broadcast + # obs = {} + + dist.broadcast(signal_tensor, src=0, group=self._signal_group) + + signal = signal_tensor.item() + if signal == 1: + logger.info(f"Rank {dist.get_rank()} received shutdown signal") + break + + # --- ADD THIS ELIF BLOCK --- + elif signal == 2: + logger.info(f"Rank {dist.get_rank()} received idle signal. Waiting for next client.") + # Loop back to the top and wait for the next signal + continue + + # Receive the batch data via broadcast/gather mechanism + # This is a simplified version - the actual obs structure needs to be broadcasted + batch = self._receive_batch_from_rank0() + # Participate in distributed forward pass + dist.barrier() + with torch.no_grad(): + result_batch, video_pred = self._policy.lazy_joint_forward_causal(batch) + dist.barrier() + + except Exception as e: + logger.error(f"Worker loop error on rank {dist.get_rank()}: {e}") + traceback.print_exc() + break + + def _receive_batch_from_rank0(self): + """Receive batch data from rank 0 using torch.distributed primitives.""" + import pickle + + # Receive the size of the pickled data first + size_tensor = torch.zeros(1, dtype=torch.int64, device='cuda') + dist.broadcast(size_tensor, src=0) + data_size = size_tensor.item() + + # Receive the actual data + data_tensor = torch.zeros(data_size, dtype=torch.uint8, device='cuda') + dist.broadcast(data_tensor, src=0) + + # Deserialize + obs = pickle.loads(data_tensor.cpu().numpy().tobytes()) + return Batch(obs=obs) + + def _broadcast_batch_to_workers(self, obs): + """Broadcast batch data from rank 0 to all other ranks.""" + import pickle + + # Serialize the obs + serialized = pickle.dumps(obs) + data_size = len(serialized) + + # Broadcast size first + size_tensor = torch.tensor([data_size], dtype=torch.int64, device='cuda') + dist.broadcast(size_tensor, src=0) + + # Broadcast data + data_tensor = torch.frombuffer(serialized, dtype=torch.uint8).cuda() + dist.broadcast(data_tensor, src=0) + + async def _handler(self, websocket: _server.ServerConnection): + logger.info(f"Connection from {websocket.remote_address} opened") + packer = msgpack_numpy.Packer() + + await websocket.send(packer.pack(self._metadata)) + + prev_total_time = None + signal_tensor = torch.zeros(1, dtype=torch.int32, device='cpu') + + try: + while True: + try: + start_time = time.perf_counter() + data = await websocket.recv() + recv_done = time.perf_counter() + obs = msgpack_numpy.unpackb(data) + print(f"Wait Time: {recv_done - start_time:.2f} seconds") + self._msg_index += 1 + + infer_start_time = time.perf_counter() + + # Signal other ranks to continue (0 = continue) + signal_tensor.zero_() + dist.broadcast(signal_tensor, src=0, group=self._signal_group) # <-- USE GLOO GROUP + + # Broadcast the obs to all ranks for distributed inference + self._broadcast_batch_to_workers(obs) + batch = Batch(obs=obs) + + # All ranks need to participate in the forward pass + dist.barrier() + forward_start_time = time.perf_counter() + with torch.no_grad(): + result_batch, video_pred = self._policy.lazy_joint_forward_causal(batch) + dist.barrier() + print(f"Forward Time: {time.perf_counter() - forward_start_time:.2f} seconds") + + action_chunk_dict = result_batch.act + video_chunk = video_pred + + print(f"Inference Time: {time.perf_counter() - infer_start_time:.2f} seconds") + + self.video_across_time.append(video_chunk) + + if len(self.video_across_time) > 10: + frame_list = [] + video_across_time_cat = torch.cat(self.video_across_time, dim=2) + frames = self._policy.trained_model.action_head.vae.decode( + video_across_time_cat, + tiled=self._policy.trained_model.action_head.tiled, + tile_size=(self._policy.trained_model.action_head.tile_size_height, self._policy.trained_model.action_head.tile_size_width), + tile_stride=(self._policy.trained_model.action_head.tile_stride_height, self._policy.trained_model.action_head.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C") + frames = frames[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + # Add each frame individually to the list + for frame in frames: + frame_list.append(frame) + + sample_frame = frame_list[0] + if len(sample_frame.shape) == 3 and sample_frame.shape[2] in [1, 3, 4]: + # Save all frames as a single MP4 file + save_dir = self._output_dir if self._output_dir else "." + os.makedirs(save_dir, exist_ok=True) + all_mp4_files = [f for f in os.listdir(save_dir) if f.endswith(".mp4")] + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + num_frames = len(frame_list) + n = (num_frames - 1) // 8 # num_frames = 8n+1, so n = (num_frames-1)/8 + output_path = os.path.join(save_dir, f'{len(all_mp4_files):06}_{timestamp}_n{n}.mp4') + imageio.mimsave(output_path, frame_list, fps=5, codec='libx264') + print(f"Saved video to: {output_path}") + else: + print(f"Warning: Invalid frame shape {sample_frame.shape}. Expected (H, W, C) with C in [1, 3, 4]. Skipping video save.") + + self.video_across_time = [] + elif self._policy.trained_model.action_head.current_start_frame == 1 + self._policy.trained_model.action_head.num_frame_per_block and len(self.video_across_time) > 1: + print("current_start_frame == 1 + num_frame_per_block and len(self.video_across_time) > 1") + frame_list = [] + video_across_time_cat = torch.cat(self.video_across_time[:-1], dim=2) + frames = self._policy.trained_model.action_head.vae.decode( + video_across_time_cat, + tiled=self._policy.trained_model.action_head.tiled, + tile_size=(self._policy.trained_model.action_head.tile_size_height, self._policy.trained_model.action_head.tile_size_width), + tile_stride=(self._policy.trained_model.action_head.tile_stride_height, self._policy.trained_model.action_head.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C") + frames = frames[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + # Add each frame individually to the list + for frame in frames: + frame_list.append(frame) + sample_frame = frame_list[0] + if len(sample_frame.shape) == 3 and sample_frame.shape[2] in [1, 3, 4]: + # Save all frames as a single MP4 file + save_dir = self._output_dir if self._output_dir else "." + os.makedirs(save_dir, exist_ok=True) + all_mp4_files = [f for f in os.listdir(save_dir) if f.endswith(".mp4")] + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + num_frames = len(frame_list) + n = (num_frames - 1) // 8 # num_frames = 8n+1, so n = (num_frames-1)/8 + output_path = os.path.join(save_dir, f'{len(all_mp4_files):06}_{timestamp}_n{n}.mp4') + imageio.mimsave(output_path, frame_list, fps=5, codec='libx264') + print(f"Saved video to: {output_path}") + self.video_across_time = [video_chunk] + + + def batch_to_dict(batch): + out = {} + for k in dir(batch): + if not k.startswith("action."): + continue + out[k] = getattr(batch, k) + return out + action_chunk_dict = batch_to_dict(action_chunk_dict) + await websocket.send(packer.pack(action_chunk_dict)) + + except websockets.ConnectionClosed: + logger.info(f"Connection from {websocket.remote_address} closed") + if len(self.video_across_time) > 0: + frame_list = [] + video_across_time_cat = torch.cat(self.video_across_time, dim=2) + frames = self._policy.trained_model.action_head.vae.decode( + video_across_time_cat, + tiled=self._policy.trained_model.action_head.tiled, + tile_size=(self._policy.trained_model.action_head.tile_size_height, self._policy.trained_model.action_head.tile_size_width), + tile_stride=(self._policy.trained_model.action_head.tile_stride_height, self._policy.trained_model.action_head.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C") + frames = frames[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + # Add each frame individually to the list + for frame in frames: + frame_list.append(frame) + + sample_frame = frame_list[0] + if len(sample_frame.shape) == 3 and sample_frame.shape[2] in [1, 3, 4]: + # Save all frames as a single MP4 file + save_dir = self._output_dir if self._output_dir else "." + os.makedirs(save_dir, exist_ok=True) + all_mp4_files = [f for f in os.listdir(save_dir) if f.endswith(".mp4")] + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + num_frames = len(frame_list) + n = (num_frames - 1) // 8 # num_frames = 8n+1, so n = (num_frames-1)/8 + output_path = os.path.join(save_dir, f'{len(all_mp4_files):06}_{timestamp}_n{n}.mp4') + imageio.mimsave(output_path, frame_list, fps=5, codec='libx264') + print(f"Saved video to: {output_path}") + else: + print(f"Warning: Invalid frame shape {sample_frame.shape}. Expected (H, W, C) with C in [1, 3, 4]. Skipping video save.") + + self.video_across_time = [] + 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 + finally: + logger.info(f"Rank 0: Client session ended. Sending idle signal (2) to workers.") + signal_tensor.fill_(2) # Set tensor value to 2 + dist.broadcast(signal_tensor, src=0, group=self._signal_group) + # When connection closes, signal other ranks to continue waiting for next connection + # (or implement proper shutdown if needed) + + +def init_mesh() -> DeviceMesh: + # env vars set by torchrun + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + print(f"Rank {rank}/{world_size} (PID: {os.getpid()}) setting device to {rank}") + + torch.cuda.set_device(rank) + device = torch.device(f"cuda:{rank}") + + mesh = init_device_mesh( + device_type="cuda", + mesh_shape=(world_size, ), + mesh_dim_names=("ip", ), + ) + print(f"Rank {rank}/{world_size} (PID: {os.getpid()}) using device {device}") + + return mesh + +def _health_check(connection: _server.ServerConnection, request: _server.Request) -> _server.Response | None: + if request.path == "/healthz": + return connection.respond(http.HTTPStatus.OK, "OK\n") + # Continue with the normal request handling. + return None + + +def main(args: Args) -> None: + # Set environment variable for DIT cache. + os.environ["ENABLE_DIT_CACHE"] = "true" if args.enable_dit_cache else "false" + + # Use TE cuDNN backend for attention. + os.environ["ATTENTION_BACKEND"] = "TE" + + # Increase the recompile limit to 100 for inference due + # to autoregressive nature of the model (several possible shapes). + torch._dynamo.config.recompile_limit = 800 + + embodiment_tag = "oxe_droid" + model_path = args.model_path + policy_metadata = { + "embodiment": embodiment_tag, + "model_name": "dreamzero", + "model_path": model_path, + } + + device_mesh = init_mesh() + rank = dist.get_rank() + + timeout_delta = datetime.timedelta(seconds=args.timeout_seconds) + signal_group = dist.new_group(backend="gloo", timeout=timeout_delta) + logger.info(f"Rank {rank} initialized signal_group (gloo)") + + policy = GrootSimPolicy( + embodiment_tag=EmbodimentTag(embodiment_tag), + model_path=model_path, + device="cuda" if torch.cuda.is_available() else "cpu", + device_mesh=device_mesh, + ) + + # Create server for all ranks - rank 0 handles websocket, others run worker loop + hostname = socket.gethostname() + local_ip = socket.gethostbyname(hostname) + + if rank == 0: + logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip) + # Create output directory for videos + # Extract parent directory and checkpoint name from model_path + parent_dir = os.path.dirname(model_path) + date_suffix = datetime.datetime.now().strftime("%Y%m%d") + checkpoint_name = os.path.basename(model_path) + output_dir = os.path.join(parent_dir, f"real_world_eval_gen_{date_suffix}_{args.index}", checkpoint_name) + os.makedirs(output_dir, exist_ok=True) + logging.info("Videos will be saved to: %s", output_dir) + else: + output_dir = None + logging.info(f"Rank {rank} starting as worker for distributed inference...") + + # Create wrapper policy that converts between roboarena and AR_droid formats + wrapper_policy = ARDroidRoboarenaPolicy( + groot_policy=policy, + signal_group=signal_group, + output_dir=output_dir, + ) + + # Configure server for AR_droid (2 external cameras, wrist camera, joint position actions) + server_config = PolicyServerConfig( + image_resolution=(180, 320), # AR_droid expects 180x320 images + needs_wrist_camera=True, + n_external_cameras=2, + needs_stereo_camera=False, + needs_session_id=True, # Track session to reset state for new clients + action_space="joint_position", + ) + + if rank == 0: + logging.info("Using roboarena policy server interface") + logging.info(f"Server config: {server_config}") + roboarena_server = RoboarenaServer( + policy=wrapper_policy, + server_config=server_config, + host="0.0.0.0", + port=args.port, + ) + roboarena_server.serve_forever() + else: + # Non-rank-0 processes need to run worker loop for distributed inference + # We'll use the existing WebsocketPolicyServer's worker loop mechanism + server = WebsocketPolicyServer( + policy=policy, + host="0.0.0.0", + port=args.port, + metadata=policy_metadata, + output_dir=output_dir, + signal_group=signal_group, + ) + asyncio.run(server._worker_loop()) + + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, force=True) + args = tyro.cli(Args) + main(args) \ No newline at end of file diff --git a/test_client_AR.py b/test_client_AR.py new file mode 100644 index 0000000000000000000000000000000000000000..e52370c21483ff5c1ce92db112206d9a83bff7bd --- /dev/null +++ b/test_client_AR.py @@ -0,0 +1,331 @@ +#!/usr/bin/env python3 +"""Test client for AR_droid policy server using roboarena interface. + +Sends real video frames from debug_image/ directory instead of zero dummy images. + +Frame schedule (matching debug_inference.py): + - Step 0 (initial): send frame [0] (1 frame, H W 3) + - Step 1: send frames [0, 7, 15, 23] (4 frames, 4 H W 3) + - Step 2: send frames [24, 31, 39, 47] (4 frames) + - Step 3: send frames [48, 55, 63, 71] (4 frames) + - ... + +Expected server configuration: + - image_resolution: (180, 320) + - n_external_cameras: 2 + - needs_wrist_camera: True + - action_space: "joint_position" + +Usage: + # Start server with roboarena interface: + torchrun --nproc_per_node=8 socket_test_optimized_AR.py --port 8000 + + # Run this test: + python test_client_AR.py --host --port 8000 + + # Use zero images instead of real video (old behavior): + python test_client_AR.py --host --port 8000 --use-zero-images +""" + +import argparse +import logging +import os +import time +import uuid + +import cv2 +import numpy as np + +import eval_utils.policy_server as policy_server +from eval_utils.policy_client import WebsocketClientPolicy + +VIDEO_DIR = os.path.join(os.path.dirname(__file__), "debug_image") + +# roboarena key -> video filename +CAMERA_FILES = { + "observation/exterior_image_0_left": "exterior_image_1_left.mp4", + "observation/exterior_image_1_left": "exterior_image_2_left.mp4", + "observation/wrist_image_left": "wrist_image_left.mp4", +} + +# Frame schedule constants (matching debug_inference.py) +RELATIVE_OFFSETS = [-23, -16, -8, 0] +ACTION_HORIZON = 24 + + +def load_all_frames(video_path: str) -> np.ndarray: + """Load all frames from a video file. Returns (N, H, W, 3) uint8 array (RGB).""" + cap = cv2.VideoCapture(video_path) + frames = [] + while True: + ret, frame = cap.read() + if not ret: + break + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + cap.release() + if not frames: + raise RuntimeError(f"No frames loaded from {video_path}") + return np.stack(frames, axis=0) + + +def load_camera_frames() -> dict[str, np.ndarray]: + """Load all video frames for each camera from the debug_image/ directory. + + Returns: + Dict mapping roboarena camera keys to (N, H, W, 3) uint8 arrays. + """ + camera_frames: dict[str, np.ndarray] = {} + for cam_key, fname in CAMERA_FILES.items(): + path = os.path.join(VIDEO_DIR, fname) + camera_frames[cam_key] = load_all_frames(path) + logging.info(f"Loaded {cam_key}: {camera_frames[cam_key].shape}") + return camera_frames + + +def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]: + """Build the frame index schedule for multi-frame chunks. + + Returns a list of frame-index lists. Each inner list has 4 indices. + """ + chunks: list[list[int]] = [] + current_frame = 23 # first anchor frame + for _ in range(num_chunks): + indices = [max(current_frame + off, 0) for off in RELATIVE_OFFSETS] + if indices[-1] >= total_frames: + logging.info( + f"Frame {indices[-1]} >= {total_frames}, stopping at {len(chunks)} chunks" + ) + break + chunks.append(indices) + current_frame += ACTION_HORIZON + return chunks + + +def _make_obs_from_video( + camera_frames: dict[str, np.ndarray], + frame_indices: list[int], + prompt: str, + session_id: str, +) -> dict: + """Build an observation dict from real video frames. + + For 1 frame: each image key is (H, W, 3). + For 4 frames: each image key is (4, H, W, 3). + """ + obs: dict = {} + for cam_key, all_frames in camera_frames.items(): + selected = all_frames[frame_indices] # (T, H, W, 3) + if len(frame_indices) == 1: + selected = selected[0] # (H, W, 3) + obs[cam_key] = selected + + obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) + obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) + obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) + obs["prompt"] = prompt + obs["session_id"] = session_id + return obs + + +def _make_zero_observation( + server_config: policy_server.PolicyServerConfig, + prompt: str = "pick up the object", + session_id: str | None = None, +) -> dict: + """Create a dummy observation matching AR_droid expectations. + + AR_droid expects: + - 2 external cameras (exterior_image_0_left, exterior_image_1_left) + - 1 wrist camera (wrist_image_left) + - Image resolution: 180x320 (H x W) + - joint_position: 7 DoF + - gripper_position: 1 DoF + """ + obs = {} + + # Determine image resolution + if server_config.image_resolution is not None: + h, w = server_config.image_resolution + else: + # Default for AR_droid + h, w = 180, 320 + + # External cameras (0-indexed in roboarena) + for i in range(server_config.n_external_cameras): + obs[f"observation/exterior_image_{i}_left"] = np.zeros((h, w, 3), dtype=np.uint8) + if server_config.needs_stereo_camera: + obs[f"observation/exterior_image_{i}_right"] = np.zeros((h, w, 3), dtype=np.uint8) + + # Wrist camera + if server_config.needs_wrist_camera: + obs["observation/wrist_image_left"] = np.zeros((h, w, 3), dtype=np.uint8) + if server_config.needs_stereo_camera: + obs["observation/wrist_image_right"] = np.zeros((h, w, 3), dtype=np.uint8) + + # Session ID - should be passed in to ensure consistency within a session + if server_config.needs_session_id: + import uuid + # Generate unique session ID if not provided + obs["session_id"] = session_id if session_id else str(uuid.uuid4()) + + # State observations (AR_droid: 7 DoF arm + 1 gripper) + obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) + obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) + obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) + + # Language prompt + obs["prompt"] = prompt + + return obs + + +def test_ar_droid_policy_server( + host: str = "localhost", + port: int = 8000, + num_chunks: int = 15, + prompt: str = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan", + use_zero_images: bool = False, +): + """Test the AR_droid policy server with roboarena interface. + + When use_zero_images is False (default), loads real video frames from + debug_image/ and follows the frame schedule from debug_inference.py. + """ + logging.info(f"Connecting to AR_droid server at {host}:{port}...") + + client = WebsocketClientPolicy(host=host, port=port) + + # Validate server metadata + metadata = client.get_server_metadata() + logging.info(f"Server metadata: {metadata}") + assert isinstance(metadata, dict), "Metadata should be a dict" + + try: + server_config = policy_server.PolicyServerConfig(**metadata) + except Exception as e: + logging.error(f"Error parsing metadata: {e}") + raise e + + # Validate expected AR_droid configuration + logging.info(f"Server config: {server_config}") + assert server_config.n_external_cameras == 2, f"Expected 2 external cameras, got {server_config.n_external_cameras}" + assert server_config.needs_wrist_camera, "Expected wrist camera to be enabled" + assert server_config.action_space == "joint_position", f"Expected joint_position action space, got {server_config.action_space}" + + logging.info("Server configuration validated for AR_droid") + + # Generate unique session ID for this test run + import uuid + session_id = str(uuid.uuid4()) + logging.info(f"Session ID: {session_id}") + + # ── Zero-image fallback mode ────────────────────────────────────── + if use_zero_images: + logging.info("Using ZERO dummy images (legacy mode)") + for i in range(num_chunks): + obs = _make_zero_observation(server_config, prompt=prompt, session_id=session_id) + logging.info(f"Inference {i + 1}/{num_chunks}: prompt='{prompt}'") + t0 = time.time() + actions = client.infer(obs) + dt = time.time() - t0 + _log_action(actions, dt) + + logging.info("Sending reset...") + client.reset({}) + logging.info("Done (zero-image mode).") + return + + # ── Real video frame mode ───────────────────────────────────────── + logging.info("Loading real video frames from debug_image/ directory") + camera_frames = load_camera_frames() + + total_frames = min(v.shape[0] for v in camera_frames.values()) + logging.info(f"Total frames available: {total_frames}") + + # Build frame schedule + chunks = build_frame_schedule(total_frames, num_chunks) + + logging.info("Frame schedule:") + logging.info(" Initial: [0]") + for i, indices in enumerate(chunks): + logging.info(f" Chunk {i}: {indices}") + + # Step 0: initial single frame + logging.info("=== Initial: frame [0] ===") + obs = _make_obs_from_video(camera_frames, [0], prompt, session_id) + t0 = time.time() + actions = client.infer(obs) + dt = time.time() - t0 + _log_action(actions, dt) + + # Subsequent chunks: send 4 frames at a time + for chunk_idx, frame_indices in enumerate(chunks): + logging.info(f"=== Chunk {chunk_idx}: frames {frame_indices} ===") + obs = _make_obs_from_video(camera_frames, frame_indices, prompt, session_id) + t0 = time.time() + actions = client.infer(obs) + dt = time.time() - t0 + _log_action(actions, dt) + + # Reset triggers video save on the server + logging.info("Sending reset to save video...") + client.reset({}) + + logging.info("Done.") + + +def _log_action(actions: np.ndarray, dt: float) -> None: + """Pretty-print action shape, range, and timing.""" + assert isinstance(actions, np.ndarray), f"Expected numpy array, got {type(actions)}" + assert actions.ndim == 2, f"Expected 2D array, got shape {actions.shape}" + assert actions.shape[-1] == 8, ( + f"Expected 8 action dims (7 joints + 1 gripper), got {actions.shape[-1]}" + ) + logging.info( + f" Action shape: {actions.shape}, " + f"range: [{actions.min():.4f}, {actions.max():.4f}], " + f"time: {dt:.2f}s" + ) + + +def main(): + parser = argparse.ArgumentParser( + description="Test AR_droid policy server with real video frames from debug_image/" + ) + parser.add_argument("--host", default="localhost", help="Server hostname") + parser.add_argument("--port", type=int, default=8000, help="Server port") + parser.add_argument( + "--num-chunks", + type=int, + default=15, + help="Number of 4-frame chunks to send after the initial frame (default: 15)", + ) + parser.add_argument( + "--prompt", + default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan", + help="Language prompt for the policy", + ) + parser.add_argument( + "--use-zero-images", + action="store_true", + help="Use zero dummy images instead of real video frames (legacy mode)", + ) + + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + ) + + test_ar_droid_policy_server( + host=args.host, + port=args.port, + num_chunks=args.num_chunks, + prompt=args.prompt, + use_zero_images=args.use_zero_images, + ) + + +if __name__ == "__main__": + main()