feature(wrh): add robotwin converting data and training

[ruiheng123]

🚀 Feature(wrh): Adding training code for robotwin 2.0 implementation for adapter.

---

This PR commits my implementation of finetuning VLA-adapter on RoboTwin 2.0 benchmark (onlty training code).

Todo list:

• RoboTwin 2.0 generated data (HDF5) → RLDS format conversion
• Register new robotwin_dataset so adapter dataloader can load it
• Deploy & benchmark results (in progress)

Process:

RoboTwin 2.0 documentation

Setup RoboTwin 2.0 environment and install requirements by following here

I first generate the dataset. (reference: here. RoboTwin provides 50 bimanual tasks in total; for brevity, only one representative task is detailed below—the remaining 49 tasks are processed in exactly the same way.

NOTE: I give beat_block_hammer as an example here. More tasks like adjust_bottle, click_bell, etc. have similar process!

1. Executing generate dataset shell script

bash collect_data.sh ${task_name} ${task_config} ${gpu_id}
# bash collect_data.sh beat_block_hammer demo_clean 0 
# task config optional: demo_clean or demo_randomized

and you can find generated data in data folder (in the robotwin working directory) .

~/vla/RoboTwin/data$ tree -L 3
.
├── beat_block_hammer
│   ├── demo_clean
│   │   ├── data
│   │   ├──     ***.hdf5
│   │   ├── instructions
│   │   ├──     ***.json
│   │   ├── scene_info.json
│   │   ├── seed.txt
│   │   ├── _traj_data
│   │   └── video
│   │   ├──     ***.mp4
│   └──
└── process_stuck.py

I collect 1000 episodes per task, num_episodes can be modified in task_config/demo_clean.yml key episode_num for demo_clean and similar to demo_randomized.

2. convert hdf5 to rlds format

Since VLA-adapter is built upon OpenVLA-OFT, which supports rlds tensorflow dataset format, we need to convert generated hdf5 format data to rlds format.

Back to adapter working directory, I give an example of converting beat_block_hammer data to rlds format by adding this folder:

~/vla/VLA-Adapter/beat_block_hammer_rt$ tree 
.
├── beat_block_hammer_rt.py
└── __init__.py

Remember to replace the dataset paths!

(line 174)  instruction_path = os.path.join(<your/path/to/robotwin/task/demo/instructions>, instruction_filename)
(line 215)  base_path = "<your/path/to/robotwin/task/demo/data>"

you can change the name. Here I add postfix rt means robotwin dataset, which need to be recognized as aloha platform in adapter training scripts.

To keep naming consistent, use CamelCase for the class name (BeatBlockHammerRt) and convert it to lowercase with underscores for the file name (beat_block_hammer_rt.py).

class BeatBlockHammerRt(tfds.core.GeneratorBasedBuilder):

and now execute

~/vla/VLA-Adapter$  tfds build beat_block_hammer_rt --overwrite

This command will find beat_block_hammer_rt/beat_block_hammer_rt.py file to build the rlds dataset. After this you can find converted data in ~/tensoflow_datasets directory:

~/tensorflow_datasets$ tree -L 2
.
├── beat_block_hammer_rt
│   └── 1.0.0
│   └────  beat_block_hammer_rt-train.tfrecord-00000-of-00064
.... (more files)

which is the converted rlds dataset.

REMARK: Please make sure all keys are aligned!

3. Register in oxe dataset list

We add 3 settings in: (refering aloha bimanual platform setup)

# prismatic/vla/datasets/rlds/oxe/configs.py
    "beat_block_hammer_rt": {
        "image_obs_keys": {"primary": "image", "secondary": None, "left_wrist": "left_wrist_image", "right_wrist": "right_wrist_image"},
        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
        "state_obs_keys": ["state"],
        "state_encoding": StateEncoding.JOINT_BIMANUAL,
        "action_encoding": ActionEncoding.JOINT_POS_BIMANUAL,
    }


# prismatic/vla/datasets/rlds/oxe/mixtures.py
"beat_block_hammer_rt": [
        ("beat_block_hammer_rt", 1.0),
    ],


# prismatic/vla/datasets/rlds/oxe/transforms.py
"beat_block_hammer_rt": aloha_dataset_transform

I added some modification to enable beat block hammer dataset can be recognized as aloha platform. Then you can input

# run.sh
data_name=beat_block_hammer_rt

--dataset_name $data_name

--data_root_dir ~/tensorflow_datasets

and the dataset can be built by our RLDSDataset successfully!

Moreover, I also added a jupyter notebook for minimal checking, making sure if the converted rlds dataset can be loaded by Adapter dataloader. This can be found in vla-scripts/loading_data.ipynb, including loading + visualizing 3 views of data.

Deploying (evaluating) : to be continued! This only gives a typical case of collecting--converting--loading process for VLA-adapter training on RoboTwin 2.0 tasks!