Upload 10 files

.gitignore

@@ -0,0 +1,184 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+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/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# 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
+
+# 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/#use-with-ide
+.pdm.toml
+
+# 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
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# 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/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Some encoder paths
+facebook/
+openai/
+google/
+
+# Log
+logs/
+
+# Output
+outs/
+
+# Checkpoints
+checkpoints/
+
+# VSC
+.vscode/
+
+# Wandb
+wandb/
+
+# Distributed leaning
+hostfile.txt
+.deepspeed_env

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 TSAIL group
+
+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.

README.md

@@ -1,3 +1,366 @@
----
-license: apache-2.0
----
+# RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation
+
+### 📝[Paper](https://arxiv.org/pdf/2410.07864) | 🌍[Project Page](https://rdt-robotics.github.io/rdt-robotics/) | 🤗[Model](https://huggingface.co/robotics-diffusion-transformer/rdt-1b) | 🛢️[Data](https://huggingface.co/datasets/robotics-diffusion-transformer/rdt-ft-data) | 🏞️[Poster](./assets/iclr2025_poster.png)
+
+![](./assets/head.png)
+
+RDT-1B is a **1B**-parameter (*largest* to date) imitation learning **Diffusion Transformer** pre-trained on **1M+** (*largest* to date) multi-robot episodes. Given language instruction and RGB images of up to three views, RDT can predict the next $64$ robot actions. RDT is inherently compatible with **almost all kinds of modern mobile manipulators**, from single-arm to dual-arm, joint to EEF, position to velocity, and even with wheeled locomotion.
+
+We have fine-tuned RDT on **6K+** (one of the *largest*) self-collected bimanual episodes and deployed it on the ALOHA **dual-arm** robot. It has achieved state-of-the-art performance in terms of dexterity, zero-shot generalizability, and few-shot learning. You can find Demo videos on our [project page](https://rdt-robotics.github.io/rdt-robotics/).
+
+This repo is an official PyTorch implementation of RDT, containing:
+
+- 🛠️Model [implementation](models/rdt_runner.py) of RDT
+- 🤗1M-step [checkpoint](https://huggingface.co/robotics-diffusion-transformer/rdt-1b) of RDT-1B pre-trained on multi-robot data
+- 🤗500K-step [checkpoint](https://huggingface.co/robotics-diffusion-transformer/rdt-170m) of RDT-170M (RDT(small) in [ablation](https://arxiv.org/pdf/2410.07864))
+- 📈Training and sampling [scripts](train/train.py) (with DeepSpeed)
+- 🤖An [example](scripts/agilex_inference.py) of real-robot deployment
+- 🕹️Simulation benchmark from [Maniskill](https://github.com/haosulab/ManiSkill) environment
+
+The following guides include the [installation](#installation), [fine-tuning](#fine-tuning-on-your-own-dataset), and [deployment](#deployment-on-real-robots). Please refer to [pre-training](docs/pretrain.md) for a detailed list of pre-training datasets and a pre-training guide.
+
+## 📰 News
+- [2025/04/04] [Poster](./assets/iclr2025_poster.png) is uploaded.
+- [2024/12/17] 🔥 [Scripts](#simulation-benchmark) for evaluating RDT in Maniskill Simulation Benchmark is released!
+- [2024/10/23] 🔥 **RDT-170M** (Smaller) model is released, a more VRAM-friendly solution 🚀💻.
+
+## Installation
+
+1. Clone this repo and install prerequisites:
+
+    ```bash
+    # Clone this repo
+    git clone git@github.com:thu-ml/RoboticsDiffusionTransformer.git
+    cd RoboticsDiffusionTransformer
+    
+    # Create a Conda environment
+    conda create -n rdt python=3.10.0
+    conda activate rdt
+    
+    # Install pytorch
+    # Look up https://pytorch.org/get-started/previous-versions/ with your cuda version for a correct command
+    pip install torch==2.1.0 torchvision==0.16.0  --index-url https://download.pytorch.org/whl/cu121
+    
+    # Install packaging
+    pip install packaging==24.0
+    
+    # Install flash-attn
+    pip install flash-attn --no-build-isolation
+    
+    # Install other prequisites
+    pip install -r requirements.txt
+    ```
+
+2. Download off-the-shelf multi-modal encoders:
+
+   You can download the encoders from the following links:
+
+   - `t5-v1_1-xxl`: [link](https://huggingface.co/google/t5-v1_1-xxl/tree/main)🤗
+   - `siglip`: [link](https://huggingface.co/google/siglip-so400m-patch14-384)🤗
+
+   And link the encoders to the repo directory:
+
+   ```bash
+   # Under the root directory of this repo
+   mkdir -p google
+   
+   # Link the downloaded encoders to this repo
+   ln -s /path/to/t5-v1_1-xxl google/t5-v1_1-xxl
+   ln -s /path/to/siglip-so400m-patch14-384 google/siglip-so400m-patch14-384
+   ```
+3. Fill the missing argument in [this file](configs/base.yaml#L22):
+   
+   Note that this buffer will only be used during pre-training. See [this doc](docs/pretrain.md) for more details.
+   ```
+   # ...
+   
+   dataset:
+   # ...
+   # ADD YOUR buf_path: the path to the buffer (at least 400GB)
+      buf_path: /path/to/buffer
+   # ...
+   ```
+
+## Fine-Tuning on Your Own Dataset
+
+If your fine-tuning dataset is in the [Open X-Embodiment](https://robotics-transformer-x.github.io/) or the collection of our pre-training datasets (see [this doc](docs/pretrain.md#download-and-prepare-datasets)), you can also fine-tune RDT through the pre-trained pipeline. You need to remove other redundant datasets in the parameters. We refer to [this guide](docs/pretrain.md) (pre-training).
+
+1. Prepare your dataset:
+
+   You need to download your dataset to the disk and give it a name `my_cool_dataset`.
+
+   Then, you can link your dataset to the repo directory:
+
+   ```bash
+   # Under the root directory of this repo
+   cd data
+   mkdir -p datasets
+   
+   # Link the downloaded dataset to this repo
+   ln -s /path/to/my_cool_dataset datasets/my_cool_dataset
+   ```
+
+2. Implement the dataset loader:
+
+   You need to:
+
+   1. Register the configuration of `my_cool_dataset`:
+
+      Append the control frequency of `my_cool_dataset` in [this file](configs/dataset_control_freq.json). Write the name of `my_cool_dataset` in [this file](configs/finetune_datasets.json) and [this file](configs/finetune_sample_weights.json), where the value of the sampling weight doesn't matter since you only have one dataset. In these two files, we leave a placeholder of `agilex`; you can simply replace it with `my_cool_dataset`.
+
+   2. Re-Implement the class of `HDF5VLADataset`:
+
+      You can find this class in [this file](data/hdf5_vla_dataset.py). In this file, we provide an example of loading the fine-tuning dataset used in our paper (see [this link](https://huggingface.co/datasets/robotics-diffusion-transformer/rdt-ft-data)).
+
+      To adapt it to your dataset, you need to: (a) modify the `HDF5_DIR` (directory to `my_cool_dataset`) and `DATASET_NAME` (should be `"my_cool_dataset"`) in L21 and L22; (b) Implement the two functions of `parse_hdf5_file()` and `parse_hdf5_file_state_only()`. Please take a look at the original file for detailed comments and examples.
+
+      Note 1: Despite its name, you don't necessarily need to use HDF5 to store your data. Just make sure that the class is correctly implemented.
+
+      Note 2: During implementation, you may need to fill your robot action into the unified action vector (L180-194). Please refer to [this file](configs/state_vec.py) for an explanation of each element in the unified vector. We have reserved enough slots for each physical quantity. For example, we have reserved ten slots for joint angles. If your robot arm has six degrees of freedom, you only need to fill in the first six. 
+
+      **IMPORTANT 1:** If your robot is single-arm, please fill its action into the *right-arm* portion of the unified action vector, aligning with our pre-training datasets.
+
+      **IMPORTANT 2:** We use [6D representation](https://arxiv.org/pdf/1812.07035) for EEF rotation. If your action space contains EEF rotation (angle or quaternion), please refer to [this file](docs/test_6drot.py) for conversion. We note that this mapping is not reversible. Different Euler angles may be equivalent and correspond to the same 6D representation.
+
+      **IMPORTANT 3:** No physical quantities (except the gripper width) are normalized during pre-training. This can preserve each physical quantity's meaning, thereby promoting generalization across robots. Therefore, we encourage you not to normalize any physical quantities but to choose appropriate units for them. Generally, we use the International System of Units, which ensures that most values fall within [-1,1]. As an exception, we perform min-max normalization on the gripper width to [0,1].
+
+      **IMPORTANT 4:** If you use RTX 4090 (or lower), the GPU memory may be too low to load the `t5-v1_1-xxl` encoder. Instead, we recommend you precompute the language embeddings (see [this file](scripts/encode_lang_batch.py) for an example script) and load them during training. In this way, you need to specify the path to the embeddings in the `HDF5VLADataset` (see L148) rather than the natural language.
+
+   3. Compute the dataset statistics information for `my_cool_dataset`:
+
+      ```bash
+      # Under the root directory of this repo
+      # Use -h to see the full usage
+      python -m data.compute_dataset_stat_hdf5
+      ```
+
+3. Start fine-tuning:
+
+   Configurations relevant to model architecture and data processing are in [this file](configs/base.yaml). Normally, you do not need to modify these configurations; otherwise, it will cause errors in loading the pre-training checkpoint. Configurations relevant to training are passed through *Command Line Arguments*. Use `python main.py -h ` to see the descriptions. We provide an example of a fine-tuning script in [this file](finetune.sh) (`finetune.sh`). You may need to modify some of the parameters in this file, such as `CUTLASS_PATH` and `WANDB_PROJECT`.
+
+   Use this to start fine-tuning:
+
+   ```bash
+   source finetune.sh
+   ```
+
+   with `finetune.sh` detailed as below:
+
+   ```bash
+      deepspeed --hostfile=hostfile.txt main.py \
+         --deepspeed="./configs/zero2.json" \   # If you want to use DeepSpeed, which is strongly recommended
+         --pretrained_model_name_or_path=<MODEL ID | DIRECTORY OF MODEL WEIGHTS | PATH TO MODEL CHECKPOINT> \
+         --pretrained_text_encoder_name_or_path=<MODEL ID | PATH TO MODEL DIRECTORY > \   # e.g., google/t5-v1_1-xxl
+         --pretrained_vision_encoder_name_or_path=<MODEL ID | PATH TO MODEL DIRECTORY> \  # e.g., google/siglip-so400m-patch14-384
+         --output_dir=<DIRECTORY to SAVE CHECKPOINTS> \ # e.g., checkpoints/rdt-1b-agilex
+         --train_batch_size=32 \
+         --sample_batch_size=64 \   # batch size for diffusion sampling in validation 
+         --max_train_steps=200000 \
+         --checkpointing_period=1000 \
+         --sample_period=500 \   # sample period for validation
+         --checkpoints_total_limit=40 \
+         --lr_scheduler="constant" \
+         --learning_rate=1e-4 \
+         --mixed_precision="bf16" \ # If you want to use mixed precision, bf16 is recommended
+         --dataloader_num_workers=8 \
+         --image_aug \  # If you want to use image augmentation
+         --dataset_type="finetune" \
+         --state_noise_snr=40 \  # If you want to add noise to the state
+         --load_from_hdf5 \   # If you use HDF5 to store your data
+         --report_to=wandb
+   ```
+
+   **IMPORTANT**: If you have already chosen to precompute the language embeddings, please specify `--precomp_lang_embed` in the `finetune.sh`.
+
+   Note 1: `pretrained_model_name_or_path` can one of:
+
+      - a string, the *model id* of a pre-trained model hosted inside a model repo on HuggingFace. Please fill with `"robotics-diffusion-transformer/rdt-1b"`, which is the officially-released [RDT-1B model](https://huggingface.co/robotics-diffusion-transformer/rdt-1b)🤗 at HuggingFace. (recommended)
+      - a string, the path to a *directory* containing the manually downloaded model weights from HuggingFace, e.g., `"/path/to/rdt-1b"`. You should first manually download the `rdt-1b` directory from this [link](https://huggingface.co/robotics-diffusion-transformer/rdt-1b)🤗.
+      - a string, the path to a *directory* containing model weights saved using [`~RDTRunner.save_pretrained`] method. This can be either:
+        -  `"checkpoints/rdt-pretrain-1b/checkpoint-<STEP NUMBER>"`: This is the path to the checkpoint saved in the `<STEP NUMBE>` iteration during pre-training. Refer to [this file](docs/pretrain.md) for a tutorial on how to start your own pre-training.
+        - `"checkpoints/rdt-pretrain-1b"`: If the pre-training completes normally without any exception, you can specify this path to load the last checkpoint.
+      - a string, the path to model checkpoint (`*.pt`) saved by DeepSpeed, e.g., `"checkpoints/rdt-pretrain-1b/checkpoint-<STEP NUMBER>/pytorch_model/mp_rank_00_model_states.pt"` (verified)
+      - `None` if you want to randomly initialize the model using configuration at `config_path`.
+
+   Note 2: You can monitor the training process by observing `loss` (through a long window moving average) and `overall_avg_sample_mse` in [Wandb](https://wandb.ai/site) or [TensorBoard](https://www.tensorflow.org/tensorboard). We empirically found that the lower the `overall_avg_sample_mse`, the better the model performs. Usually, fine-tuning is over when this value converges.
+
+   Note 3: If the training oscillates, you can increase the batch size by adding more GPUs or setting a larger `--gradient_accumulation_steps`.
+
+   Note 4: Please specify `--load_from_hdf5` in your script when finetuning with an HDF5 dataset.
+
+## Deployment on Real-Robots
+
+We have encapsulated the inference of the model into a class named `RoboticDiffusionTransformerModel` (see [this file](scripts/agilex_model.py#L38)). You can call this class's `step()` method for inference. However, you may need to re-implement some parts according to your specific robot. You should at least modify the `_format_joint_to_state()` (L164) and `_unformat_action_to_joint()` (L196) to convert between robot raw actions and unified action vectors that RDT accepts. You may also specify the control frequency of your robot (L49).
+
+**IMPORTANT**: When you feed the images into `step()`, remember the order MUST be `[ext_{t-1}, right_wrist_{t-1}, left_wrist_{t-1}, ext_{t}, right_wrist_{t}, left_wrist_{t}]`.
+
+We provide an example hardware code in [this file](scripts/agilex_inference.py) for deployment on Mobile ALOHA, and the corresponding running script in [this file](inference.sh) (`inference.sh`), which is detailed below;
+
+   ```bash
+      python -m scripts.agilex_inference \
+         --use_actions_interpolation \
+         --pretrained_model_name_or_path=<PATH TO MODEL CHECKPOINT> \  # your finetuned checkpoint: e.g., checkpoints/rdt-finetune-1b/checkpoint-<STEP NUMBER>, checkpoints/rdt-finetune-1b/checkpoint-<STEP NUMBER>/pytorch_model/mp_rank_00_model_states.pt, the same before
+         --lang_embeddings_path=<PATH TO YOUR INSTURCTION EMBEDDINGS> \ # e.g. outs/lang_embeddings/your_instr.pt"
+         --ctrl_freq=25    # your control frequency
+   ```
+
+**IMPORTANT**: If you on-board GPU memory is not enough to encode the language, please refer to [this file](scripts/encode_lang.py) for precomputation and specify the language embedding path in `inference.sh`. Detail instructions are provided below:
+
+   1. Set Required Parameters in `scripts/encode_lang.py`
+
+      ```python
+      # ...
+
+      GPU = 0
+      MODEL_PATH = "google/t5-v1_1-xxl"
+      CONFIG_PATH = "configs/base.yaml"
+      SAVE_DIR = "outs/"   # output directory
+
+      # Modify this to your task name and instruction
+      TASK_NAME = "handover_pan"
+      INSTRUCTION = "Pick up the black marker on the right and put it into the packaging box on the left."
+
+      # Note: if your GPU VRAM is less than 24GB, 
+      # it is recommended to enable offloading by specifying an offload directory. 
+      OFFLOAD_DIR = None  # Specify your offload directory here, ensuring the directory exists.
+
+      # ...
+      ```
+
+   2. Run the script
+      ```
+      python -m scripts.encode_lang
+      ```
+
+Note: If you want to deploy on the Mobile ALOHA robot, don't forget to install the hardware prerequisites (see [this repo](https://github.com/MarkFzp/mobile-aloha)).
+
+## Simulation Benchmark
+
+We comprehensively evaluate RDT against baseline methods using the ManiSkill simulation benchmark. Specifically, we focus on five benchmark tasks: `PegInsertionSide`, `PickCube`, `StackCube`, `PlugCharger`, and `PushCube`. Here's a brief overview of the evaluation setup:
+
+**Evaluation Setup:**
+
+1. **Install ManiSkill:**  
+   Within the [RDT environment](#installation), install ManiSkill as follows:
+   ```bash
+   conda activate rdt
+   pip install --upgrade mani_skill
+   ```
+
+2. **Configure Vulkan:**  
+   Follow the [ManiSkill documentation](https://maniskill.readthedocs.io/en/latest/user_guide/getting_started/installation.html#vulkan) to properly set up Vulkan。
+
+3. **Obtain Model Weights:**  
+   Download the fine-tuned model weights from [this Hugging Face repository](https://huggingface.co/robotics-diffusion-transformer/maniskill-model/tree/main/rdt). Download the precomputed language embeddings from [here](https://huggingface.co/robotics-diffusion-transformer/maniskill-model/tree/main/lang_embeds) to the root directory of this repo.
+   
+4. **Run Evaluation Scripts:**  
+   After completing the setup steps, execute the provided evaluation scripts to assess RDT on the selected tasks.
+
+```
+conda activate rdt 
+python -m eval_sim.eval_rdt_maniskill \
+--pretrained_path PATH_TO_PRETRAINED_MODEL
+```
+
+### Implementation Details
+
+#### Data
+
+Utilizing the [official ManiSkill repository](https://github.com/haosulab/ManiSkill), we generated 5,000 trajectories through motion planning. The initial action mode of these trajectories is absolute joint position control and we subsequently converted them into delta end-effector pose control to align with the pre-training action space of OpenVLA and Octo. We strictly adhered to the official codebases of OpenVLA and Octo, modifying only the dataset-loading scripts. Consequently, we finetuned OpenVLA and Octo using the delta end-effector pose data. For RDT and Diffusion-Policy we leverage joint position control data for training which is aligned with our pre-training stage as well.
+
+####  Training
+- OpenVLA is fine-tuned from the officially released pre-trained checkpoint with LoRA-rank 32 until converge.
+- Octo is fine-tuned from the officially released pre-trained checkpoint for 1M iterations until converge. 
+- Diffusion-Policy is trained from scratch for 1000 epochs. We select the checkpoint of 700 epoch which has the lowest validation sample loss of 1e-3.
+- RDT is fine-tuned from our released pre-trained checkpoint for 300ks iterations.
+
+#### Results
+
+Each method is evaluated over 250 trials (10 random seeds with 25 trials per seed). The quantitative results, including success rate mean and std value across 10 random seeds are presented below:
+
+
+||PegInsertionSide|PickCube|StackCube|PlugCharger|PushCube|Mean|
+|---|---|---|---|---|---|---|
+|RDT|**13.2±0.29%**|**77.2±0.48%**|74.0±0.30%|**1.2±0.07%**|**100±0.00%**|**53.6±0.52%**|
+|OpenVLA|0.0±0.00%|8±0.00%|8±0.00%|0.0±0.00%|8±0.00%|4.8±0.00%|
+|Octo|0.0±0.00%|0.0±0.00%|0.0±0.00%|0.0±0.00%|0.0±0.00%|0.0±0.00%|
+|Diffusion-Policy|0.0±0.00%|40.0±0.00%|**80.0±0.00%**|0.0%±0.00%|88.0±0.00%|30.2±0.00%|
+
+#### Finetune RDT with Maniskill Data
+
+To fine-tune RDT with Maniskill data, first download the Maniskill data from [here](https://huggingface.co/robotics-diffusion-transformer/maniskill-model) and extract it to `data/datasets/rdt-ft-data`. Then copy the code in `data/hdf5_vla_dataset.py` to `data/hdf5_maniskill_dataset.py` and run the following script:
+
+```
+bash finetune_maniskill.sh
+```
+
+#### Reproducing Baseline Results
+
+Download and extract the fine-tuned model weights from [here](https://huggingface.co/robotics-diffusion-transformer/maniskill-model) to `eval_sim/`.
+
+- OpenVLA: Clone [OpenVLA repo](https://github.com/openvla/openvla) in `./eval_sim/` and install its environment & ManiSkill. Then run the following script:
+```
+python -m eval_sim.eval_openvla --pretrained_path PATH_TO_PRETRAINED_MODEL
+```
+- Octo: Clone [Octo repo](https://github.com/octo-models/octo.git) in `./eval_sim/` and install its environment & ManiSkill. The run the following script:
+```
+python -m eval_sim.eval_octo --pretrained_path PATH_TO_PRETRAINED_MODEL
+```
+- Diffusion-Policy: Clone our simplified [Diffusion-Policy repo](https://github.com/LBG21/RDT-Eval-Diffusion-Policy) in `./eval_sim/` and run:
+```
+python -m eval_sim.eval_dp --pretrained_path PATH_TO_PRETRAINED_MODEL
+```
+
+### RDT on [RoboTwin Benchmark](https://robotwin-platform.github.io/) 
+
+RoboTwin is a benchmark simulator based on Sapien, comprising 50 common dual-arm tasks. All policies are trained on 50 trajectories collected in a clean environment for each of the 50 tasks, and then deployed and tested for success rates in environments corresponding to either easy (clean) or hard (randomized) configurations for the same 50 tasks. You can refer to ​​[https://robotwin-platform.github.io/doc/usage/RDT.html#1-environment-setup](https://robotwin-platform.github.io/doc/usage/RDT.html#1-environment-setup)​​ to set up the environment required for RoboTwin and RDT. According to the official test results of RoboTwin ([RoboTwin 2.0 Leaderboard](https://robotwin-platform.github.io/leaderboard)), RDT ranks second only to Pi0 (excluding the DP3 policy, which utilizes ground truth point clouds).
+<img width="2048" height="1224" alt="image" src="https://github.com/user-attachments/assets/e721ffab-3dde-42f0-b36e-1593aa964a99" />
+
+
+## FAQ
+
+### 1. How can I fine-tune RDTs with limited VRAM?
+
+- **Use a Smaller Model**: Opt for the [RDT-170M model](https://huggingface.co/robotics-diffusion-transformer/rdt-170m), which requires less VRAM.
+  
+- **Select a Memory-Efficient ZeRO Stage**: Choose a more memory-efficient ZeRO stage based on your needs:
+  - **ZeRO-3 with Offload** > **ZeRO-3** > **ZeRO-2 with Offload** > **ZeRO-2** > **ZeRO-1**
+  - By default, we use [ZeRO-2](https://github.com/thu-ml/RoboticsDiffusionTransformer/blob/c68398ed526733faca4eec52cc1a7d15a9f8fea7/finetune.sh#L29) for a balance between speed and memory efficiency. Find more details on ZeRO stages [here](https://huggingface.co/docs/transformers/main/deepspeed#select-a-zero-stage) and [here](https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training).
+
+- **Enable 8-bit Adam Optimization**: Activate 8-bit Adam by setting [`use_8bit_adam=True`](https://github.com/thu-ml/RoboticsDiffusionTransformer/blob/c68398ed526733faca4eec52cc1a7d15a9f8fea7/main.py#L195) for reduced memory usage during training.
+
+- **Apply 4-bit or 8-bit Quantization**: Quantizing model weights can significantly reduce VRAM requirements.
+
+- **Use [XFormers](https://github.com/facebookresearch/xformers)**: This library provides optimized transformers with efficient memory usage.
+
+- **Enable Gradient Checkpointing**: Implement `gradient_checkpointing` manually to save memory during backpropagation. See [here](https://deepspeed.readthedocs.io/en/latest/activation-checkpointing.html) for instructions. Once you have successfully implemented this feature, we welcome you to submit a PR👏.
+- **Gradient Accumulation**: Set a larger `--gradient_accumulation_steps=<num_steps>`. This will accumulate the gradients of `<num_steps>` batches for backpropagation. Equivalently, this will increase the batch size by `<num_steps>` times, at the cost of `<num_steps>` times the running time.
+
+### 2. How many steps are recommended for fine-tuning RDT?
+
+Regardless of the batch size you select, it is recommended to train for at least 150K steps to achieve optimal results.
+
+### 3. What to do if t5-xxL is too large to store in GPU memory?
+
+1. Do not load T5-XXL in your GPU memory when training. Pre-compute language embeddings in advance.
+2. Set `OFFLOAD_DIR` to enable CPU offloading in `scripts/encode_lang_batch.py` and `scripts/encode_lang.py`.
+3. Use smaller versions of t5 like t5-base instead of t5-xxL.
+
+## Citation
+
+If you find our work helpful, please cite us:
+
+```bibtex
+@article{liu2024rdt,
+  title={RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation},
+  author={Liu, Songming and Wu, Lingxuan and Li, Bangguo and Tan, Hengkai and Chen, Huayu and Wang, Zhengyi and Xu, Ke and Su, Hang and Zhu, Jun},
+  journal={arXiv preprint arXiv:2410.07864},
+  year={2024}
+}
+```
+
+Thank you!
+
+## License
+
+All the code, model weights, and data are licensed under [MIT license](./LICENSE).

finetune.sh

@@ -0,0 +1,54 @@
+export NCCL_IB_HCA=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_7:1,mlx5_8:1,mlx5_9:1
+export NCCL_IB_DISABLE=0
+export NCCL_SOCKET_IFNAME=bond0
+export NCCL_DEBUG=INFO
+export NCCL_NVLS_ENABLE=0
+
+export TEXT_ENCODER_NAME="google/t5-v1_1-xxl"
+export VISION_ENCODER_NAME="google/siglip-so400m-patch14-384"
+export OUTPUT_DIR="./checkpoints/rdt-finetune-1b"
+export CFLAGS="-I/usr/include"
+export LDFLAGS="-L/usr/lib/x86_64-linux-gnu"
+export CUTLASS_PATH="/path/to/cutlass"
+
+export WANDB_PROJECT="robotics_diffusion_transformer"
+
+if [ ! -d "$OUTPUT_DIR" ]; then
+    mkdir "$OUTPUT_DIR"
+    echo "Folder '$OUTPUT_DIR' created"
+else
+    echo "Folder '$OUTPUT_DIR' already exists"
+fi
+
+# For run in a single node/machine
+# accelerate launch main.py \
+#     --deepspeed="./configs/zero2.json" \
+#     ...
+
+deepspeed --hostfile=hostfile.txt main.py \
+    --deepspeed="./configs/zero2.json" \
+    --pretrained_model_name_or_path="robotics-diffusion-transformer/rdt-1b" \
+    --pretrained_text_encoder_name_or_path=$TEXT_ENCODER_NAME \
+    --pretrained_vision_encoder_name_or_path=$VISION_ENCODER_NAME \
+    --output_dir=$OUTPUT_DIR \
+    --train_batch_size=32 \
+    --sample_batch_size=64 \
+    --max_train_steps=200000 \
+    --checkpointing_period=1000 \
+    --sample_period=500 \
+    --checkpoints_total_limit=40 \
+    --lr_scheduler="constant" \
+    --learning_rate=1e-4 \
+    --mixed_precision="bf16" \
+    --dataloader_num_workers=8 \
+    --image_aug \
+    --dataset_type="finetune" \
+    --state_noise_snr=40 \
+    --load_from_hdf5 \
+    --report_to=wandb
+
+    # Use this to resume training from some previous checkpoint
+    # --resume_from_checkpoint="checkpoint-36000" \
+    # Use this to load from saved lanuage instruction embeddings,
+    # instead of calculating it during training
+    # --precomp_lang_embed \

finetune_maniskill.sh

@@ -0,0 +1,48 @@
+export NCCL_IB_HCA=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_7:1,mlx5_8:1,mlx5_9:1
+export NCCL_IB_DISABLE=0
+export NCCL_SOCKET_IFNAME=bond0
+export NCCL_DEBUG=INFO
+export NCCL_NVLS_ENABLE=0
+
+export TEXT_ENCODER_NAME="google/t5-v1_1-xxl"
+export VISION_ENCODER_NAME="google/siglip-so400m-patch14-384"
+export OUTPUT_DIR="./checkpoints/rdt-finetune-1b-sim"
+export CFLAGS="-I/usr/include"
+export LDFLAGS="-L/usr/lib/x86_64-linux-gnu"
+export CUTLASS_PATH="/data/lingxuan/cutlass"
+
+export WANDB_PROJECT="robotic_diffusion_transformer"
+
+if [ ! -d "$OUTPUT_DIR" ]; then
+    mkdir "$OUTPUT_DIR"
+    echo "Folder '$OUTPUT_DIR' created"
+else
+    echo "Folder '$OUTPUT_DIR' already exists"
+fi
+# For run in a single node/machine
+# accelerate launch main.py \
+#     --deepspeed="./configs/zero2.json" \
+#     ...
+
+accelerate launch main.py \
+    --deepspeed="./configs/zero2.json" \
+    --pretrained_model_name_or_path="robotics-diffusion-transformer/rdt-1b" \
+    --pretrained_text_encoder_name_or_path=$TEXT_ENCODER_NAME \
+    --pretrained_vision_encoder_name_or_path=$VISION_ENCODER_NAME \
+    --output_dir=$OUTPUT_DIR \
+    --train_batch_size=24 \
+    --sample_batch_size=32 \
+    --max_train_steps=400000 \
+    --checkpointing_period=10000 \
+    --sample_period=500 \
+    --checkpoints_total_limit=40 \
+    --lr_scheduler="constant" \
+    --learning_rate=1e-4 \
+    --mixed_precision="bf16" \
+    --dataloader_num_workers=8 \
+    --image_aug \
+    --dataset_type="finetune" \
+    --state_noise_snr=40 \
+    --load_from_hdf5 \
+    --report_to=wandb
+

inference.sh

@@ -0,0 +1,5 @@
+python -m scripts.agilex_inference \
+    --use_actions_interpolation \
+    --pretrained_model_name_or_path="checkpoints/your_finetuned_ckpt.pt" \  # your finetuned checkpoint: e.g., checkpoints/rdt-finetune-1b/checkpoint-<STEP NUMBER>, checkpoints/rdt-finetune-1b/checkpoint-<STEP NUMBER>/pytorch_model/mp_rank_00_model_states.pt,
+    --lang_embeddings_path="outs/lang_embeddings/your_instr.pt" \
+    --ctrl_freq=25    # your control frequency

main.py

@@ -0,0 +1,300 @@
+import argparse
+import os
+from train.train import train
+
+from accelerate.logging import get_logger
+
+
+def parse_args(input_args=None):
+    parser = argparse.ArgumentParser(description="Main script for training RDT.")
+    parser.add_argument(
+        "--config_path",
+        type=str,
+        default="configs/base.yaml",
+        help="Path to the configuration file. Default is `configs/base.yaml`.",
+    )
+    parser.add_argument(
+        "--deepspeed",
+        type=str,
+        default=None,
+        help="Enable DeepSpeed and pass the path to its config file or an already initialized DeepSpeed config dictionary",
+    )
+    parser.add_argument(
+        "--pretrained_text_encoder_name_or_path",
+        type=str,
+        default=None,
+        help="Pretrained text encoder name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--pretrained_vision_encoder_name_or_path",
+        type=str,
+        default=None,
+        help="Pretrained vision encoder name or path if not the same as model_name",
+    )
+    
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="checkpoints",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+
+    parser.add_argument(
+        "--load_from_hdf5",
+        action="store_true",
+        default=False,
+        help=(
+            "Whether to load the dataset directly from HDF5 files. "
+            "If False, the dataset will be loaded using producer-consumer pattern, "
+            "where the producer reads TFRecords and saves them to buffer, and the consumer reads from buffer."
+        )
+    )
+    parser.add_argument(
+        "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument(
+        "--sample_batch_size", type=int, default=8, help="Batch size (per device) for the sampling dataloader."
+    )
+    parser.add_argument(
+        "--num_sample_batches", type=int, default=2, help="Number of batches to sample from the dataset."
+    )
+    parser.add_argument("--num_train_epochs", type=int, default=1)
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--checkpointing_period",
+        type=int,
+        default=500,
+        help=(
+            "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "
+            "In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference."
+            "Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components."
+            "See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step"
+            "instructions."
+        ),
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=(
+            "Max number of checkpoints to store. Passed as `total_limit` to the `Accelerator` `ProjectConfiguration`."
+            " See Accelerator::save_state https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state"
+            " for more details"
+        ),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+            ' `--checkpointing_period`, or `"latest"` to automatically select the last available checkpoint.'
+        ),
+    )
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        help=(
+            "Path or name of a pretrained checkpoint to load the model from.\n"
+            "   This can be either:\n"
+            "   - a string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co, e.g., `robotics-diffusion-transformer/rdt-1b`,\n"
+            "   - a path to a *directory* containing model weights saved using [`~RDTRunner.save_pretrained`] method, e.g., `./my_model_directory/`.\n"
+            "   - a path to model checkpoint (*.pt), .e.g, `my_model_directory/checkpoint-10000/pytorch_model/mp_rank_00_model_states.pt`"
+            "   - `None` if you are randomly initializing model using configuration at `config_path`."
+        )
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-6,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--cond_mask_prob",
+        type=float,
+        default=0.1,
+        help=(
+            "The probability to randomly mask the conditions (except states) during training. "
+            "If set to 0, the conditions are not masked."
+        ),
+    )
+    parser.add_argument(
+        "--cam_ext_mask_prob",
+        type=float,
+        default=-1.0,
+        help=(
+            "The probability to randomly mask the external camera image during training. "
+            "If set to < 0, the external camera image is masked with the probability of `cond_mask_prob`."
+        ),
+    )
+    parser.add_argument(
+        "--state_noise_snr",
+        type=float,
+        default=None,
+        help=(
+            "The signal-to-noise ratio (SNR, unit: dB) for adding noise to the states. "
+            "Default is None, which means no noise is added."
+        ),
+    )
+    parser.add_argument(
+        "--image_aug",
+        action="store_true",
+        default=False,
+        help="Whether or not to apply image augmentation (ColorJitter, blur, noise, etc) to the input images.",
+    )
+    parser.add_argument(
+        "--precomp_lang_embed",
+        action="store_true",
+        default=False,
+        help="Whether or not to use precomputed language embeddings.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--lr_num_cycles",
+        type=int,
+        default=1,
+        help="Number of hard resets of the lr in cosine_with_restarts scheduler.",
+    )
+    parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.")
+    parser.add_argument(
+        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    parser.add_argument("--alpha", type=float, default=0.9, help="The moving average coefficient for each dataset's loss.")
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--hub_model_id",
+        type=str,
+        default=None,
+        help="The name of the repository to keep in sync with the local `output_dir`.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument(
+        "--sample_period",
+        type=int,
+        default=-1,
+        help=(
+            "Run sampling every X steps. During the sampling phase, the model will sample a trajectory"
+            " and report the error between the sampled trajectory and groud-truth trajectory"
+            " in the training batch."
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+
+    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
+    parser.add_argument(
+        "--set_grads_to_none",
+        action="store_true",
+        help=(
+            "Save more memory by using setting grads to None instead of zero. Be aware, that this changes certain"
+            " behaviors, so disable this argument if it causes any problems. More info:"
+            " https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html"
+        ),
+    )
+
+    parser.add_argument('--dataset_type', 
+        type=str, 
+        default="pretrain",
+        required=False,
+        help="Whether to load the pretrain dataset or finetune dataset."
+    )
+
+    if input_args is not None:
+        args = parser.parse_args(input_args)
+    else:
+        args = parser.parse_args()
+
+    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
+    if env_local_rank != -1 and env_local_rank != args.local_rank:
+        args.local_rank = env_local_rank
+
+    return args
+
+
+if __name__ == "__main__":
+    logger = get_logger(__name__)
+    args = parse_args()
+    train(args, logger)

pretrain.sh

@@ -0,0 +1,47 @@
+export NCCL_IB_HCA=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_7:1,mlx5_8:1,mlx5_9:1
+export NCCL_IB_DISABLE=0
+export NCCL_SOCKET_IFNAME=bond0
+export NCCL_DEBUG=INFO
+export NCCL_NVLS_ENABLE=0
+
+export TEXT_ENCODER_NAME="google/t5-v1_1-xxl"
+export VISION_ENCODER_NAME="google/siglip-so400m-patch14-384"
+export OUTPUT_DIR="./checkpoints/rdt-pretrain-1b"
+export CFLAGS="-I/usr/include"
+export LDFLAGS="-L/usr/lib/x86_64-linux-gnu"
+export CUTLASS_PATH="/path/to/cutlass"
+
+export WANDB_PROJECT="robotics_diffusion_transformer"
+
+if [ ! -d "$OUTPUT_DIR" ]; then
+    mkdir "$OUTPUT_DIR"
+    echo "Folder '$OUTPUT_DIR' created"
+else
+    echo "Folder '$OUTPUT_DIR' already exists"
+fi
+
+# For run in a single node/machine
+# accelerate launch main.py \
+#     --deepspeed="./configs/zero2.json" \
+#     ...
+
+deepspeed --hostfile=hostfile.txt main.py \
+    --deepspeed="./configs/zero2.json" \
+    --pretrained_text_encoder_name_or_path=$TEXT_ENCODER_NAME \
+    --pretrained_vision_encoder_name_or_path=$VISION_ENCODER_NAME \
+    --output_dir=$OUTPUT_DIR \
+    --train_batch_size=32 \
+    --sample_batch_size=64 \
+    --max_train_steps=1000000 \
+    --checkpointing_period=1000 \
+    --sample_period=500 \
+    --checkpoints_total_limit=40 \
+    --lr_scheduler="constant" \
+    --learning_rate=1e-4 \
+    --mixed_precision="bf16" \
+    --dataloader_num_workers=8 \
+    --dataset_type="pretrain" \
+    --report_to=wandb
+
+    # Use this to resume training from some previous checkpoint
+    # --resume_from_checkpoint="checkpoint-1000" \

requirements.txt

@@ -0,0 +1,11 @@
+packaging==24.0
+wandb==0.17.0
+deepspeed==0.14.2   
+accelerate==0.30.1
+diffusers==0.27.2
+timm==1.0.3
+transformers==4.41.0
+sentencepiece==0.2.0
+h5py==3.11.0
+opencv-python==4.9.0.80
+imgaug==0.4.0

requirements_data.txt

@@ -0,0 +1,9 @@
+tfds-nightly==4.9.4.dev202402070044
+gsutil==5.27
+tensorflow==2.15.0.post1
+pillow==10.2.0
+pyyaml==6.0.1
+opencv-python==4.9.0.80
+tensorflow-graphics==2021.12.3
+imageio==2.34.0
+imageio-ffmpeg==0.4.9