From 110264000f74b7b44e0040873100a178f373af7f Mon Sep 17 00:00:00 2001 From: Simon Alibert Date: Tue, 15 Oct 2024 19:03:11 +0200 Subject: [PATCH] Add fixes for batch convert --- .../v2/batch_convert_dataset_v1_to_v2.py | 291 +++++++++--------- .../datasets/v2/convert_dataset_v1_to_v2.py | 17 +- 2 files changed, 169 insertions(+), 139 deletions(-) diff --git a/lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py b/lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py index a5c5440f..dcb949c2 100644 --- a/lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py +++ b/lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py @@ -1,144 +1,161 @@ -from pprint import pprint +# 'lerobot/aloha_mobile_cabinet', +# 'lerobot/aloha_mobile_chair', +# 'lerobot/aloha_mobile_elevator', +# 'lerobot/aloha_mobile_shrimp', +# 'lerobot/aloha_mobile_wash_pan', +# 'lerobot/aloha_mobile_wipe_wine', +# 'lerobot/aloha_sim_insertion_human', +# 'lerobot/aloha_sim_insertion_human_image', +# 'lerobot/aloha_sim_insertion_scripted', +# 'lerobot/aloha_sim_insertion_scripted_image', +# 'lerobot/aloha_sim_transfer_cube_human', +# 'lerobot/aloha_sim_transfer_cube_human_image', +# 'lerobot/aloha_sim_transfer_cube_scripted', +# 'lerobot/aloha_sim_transfer_cube_scripted_image', +# 'lerobot/aloha_static_battery', +# 'lerobot/aloha_static_candy', +# 'lerobot/aloha_static_coffee', +# 'lerobot/aloha_static_coffee_new', +# 'lerobot/aloha_static_cups_open', +# 'lerobot/aloha_static_fork_pick_up', +# 'lerobot/aloha_static_pingpong_test', +# 'lerobot/aloha_static_pro_pencil', +# 'lerobot/aloha_static_screw_driver', +# 'lerobot/aloha_static_tape', +# 'lerobot/aloha_static_thread_velcro', +# 'lerobot/aloha_static_towel', +# 'lerobot/aloha_static_vinh_cup', +# 'lerobot/aloha_static_vinh_cup_left', +# 'lerobot/aloha_static_ziploc_slide', +# 'lerobot/asu_table_top', +# 'lerobot/austin_buds_dataset', +# 'lerobot/austin_sailor_dataset', +# 'lerobot/austin_sirius_dataset', +# 'lerobot/berkeley_autolab_ur5', +# 'lerobot/berkeley_cable_routing', +# 'lerobot/berkeley_fanuc_manipulation', +# 'lerobot/berkeley_gnm_cory_hall', +# 'lerobot/berkeley_gnm_recon', +# 'lerobot/berkeley_gnm_sac_son', +# 'lerobot/berkeley_mvp', +# 'lerobot/berkeley_rpt', +# 'lerobot/cmu_franka_exploration_dataset', +# 'lerobot/cmu_play_fusion', +# 'lerobot/cmu_stretch', +# 'lerobot/columbia_cairlab_pusht_real', +# 'lerobot/conq_hose_manipulation', +# 'lerobot/dlr_edan_shared_control', +# 'lerobot/dlr_sara_grid_clamp', +# 'lerobot/dlr_sara_pour', +# 'lerobot/droid_100', +# 'lerobot/fmb', +# 'lerobot/iamlab_cmu_pickup_insert', +# 'lerobot/imperialcollege_sawyer_wrist_cam', +# 'lerobot/jaco_play', +# 'lerobot/kaist_nonprehensile', +# 'lerobot/nyu_door_opening_surprising_effectiveness', +# 'lerobot/nyu_franka_play_dataset', +# 'lerobot/nyu_rot_dataset', +# 'lerobot/pusht', +# 'lerobot/pusht_image', +# 'lerobot/roboturk', +# 'lerobot/stanford_hydra_dataset', +# 'lerobot/stanford_kuka_multimodal_dataset', +# 'lerobot/stanford_robocook', +# 'lerobot/taco_play', +# 'lerobot/tokyo_u_lsmo', +# 'lerobot/toto', +# 'lerobot/ucsd_kitchen_dataset', +# 'lerobot/ucsd_pick_and_place_dataset', +# 'lerobot/uiuc_d3field', +# 'lerobot/umi_cup_in_the_wild', +# 'lerobot/unitreeh1_fold_clothes', +# 'lerobot/unitreeh1_rearrange_objects', +# 'lerobot/unitreeh1_two_robot_greeting', +# 'lerobot/unitreeh1_warehouse', +# 'lerobot/usc_cloth_sim', +# 'lerobot/utaustin_mutex', +# 'lerobot/utokyo_pr2_opening_fridge', +# 'lerobot/utokyo_pr2_tabletop_manipulation', +# 'lerobot/utokyo_saytap', +# 'lerobot/utokyo_xarm_bimanual', +# 'lerobot/utokyo_xarm_pick_and_place', +# 'lerobot/viola', +# 'lerobot/xarm_lift_medium', +# 'lerobot/xarm_lift_medium_image', +# 'lerobot/xarm_lift_medium_replay', +# 'lerobot/xarm_lift_medium_replay_image', +# 'lerobot/xarm_push_medium', +# 'lerobot/xarm_push_medium_image', +# 'lerobot/xarm_push_medium_replay', +# 'lerobot/xarm_push_medium_replay_image', + +from pathlib import Path from lerobot import available_datasets +from lerobot.common.datasets.v2.convert_dataset_v1_to_v2 import convert_dataset, parse_robot_config -# from lerobot.common.datasets.v2.convert_dataset_v1_to_v2 import convert_dataset +# import tensorflow_datasets as tfds +# builder = tfds.builder("columbia_cairlab_pusht_real") +# builder.info.features -pprint(available_datasets) - -for repo_id in available_datasets: - name = repo_id.split("/")[1] - if "aloha" in name: - if "insertion" in name: - single_task = "Insert the peg into the socket." - elif "transfer" in name: - single_task = "Pick up the cube with the right arm and transfer it to the left arm." - elif "battery" in name: - single_task = "Place the battery into the slot of the remote controller." - elif "candy" in name: - single_task = "Pick up the candy and unwrap it." - elif "coffee_new" in name: - single_task = "Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray, then push the 'Hot Water' and 'Travel Mug' buttons." - elif "coffee" in name: - single_task = "Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray." - elif "cups_open" in name: - single_task = "Pick up the plastic cup and open its lid." - elif "fork_pick_up" in name: - single_task = "Pick up the fork and place it on the plate." - elif "pingpong_test" in name: - single_task = "Transfer one of the two balls in the right glass into the left glass, then transfer it back to the right glass." - elif "pro_pencil" in name: - single_task = "Pick up the pencil with the right arm, hand it over to the left arm then place it back onto the table." - elif "screw_driver" in name: - single_task = "Pick up the screwdriver with the right arm, hand it over to the left arm then place it into the cup." - elif "tape" in name: - single_task = ( - "Cut a small piece of tape from the tape dispenser then place it on the cardboard box's edge." - ) - elif "towel" in name: - single_task = "Pick up a piece of paper towel and place it on the spilled liquid." - elif "vinh_cup_left" in name: - single_task = "Pick up the platic cup with the right arm, then pop its lid open with the left arm" - elif "thread_velcro" in name: - single_task = "Pick up the velcro cable tie with the left arm, then insert the end of the velcro tie into the other end's loop with the right arm." - elif "shrimp" in name: - single_task = "Sauté the raw shrimp on both sides, then serve it in the bowl." - elif "wash_pan" in name: - single_task = "" +LOCAL_DIR = Path("data/") +ALOHA_SINGLE_TASKS_REAL = { + "aloha_mobile_cabinet": "Open the top cabinet, store the pot inside it then close the cabinet.", + "aloha_mobile_chair": "Push the chairs in front of the desk to place them against it.", + "aloha_mobile_elevator": "Take the elevator to the 1st floor.", + # Alternative version, not sure what's best. + # 'aloha_mobile_elevator': "Navigate to the elevator and call it. When it arrives, get inside et push the 1st floor button.", + "aloha_mobile_shrimp": "Sauté the raw shrimp on both sides, then serve it in the bowl.", + "aloha_mobile_wash_pan": "Pick up the pan, rinse it in the sink and then place it in the drying rack.", + "aloha_mobile_wipe_wine": "Pick up the wet cloth on the faucet and use it to clean the spilled wine on the table and underneath the glass.", + "aloha_static_battery": "Place the battery into the slot of the remote controller.", + "aloha_static_candy": "Pick up the candy and unwrap it.", + "aloha_static_coffee": "Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray, then push the 'Hot Water' and 'Travel Mug' buttons.", + "aloha_static_coffee_new": "Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray.", + "aloha_static_cups_open": "Pick up the plastic cup and open its lid.", + "aloha_static_fork_pick_up": "Pick up the fork and place it on the plate.", + "aloha_static_pingpong_test": "Transfer one of the two balls in the right glass into the left glass, then transfer it back to the right glass.", + "aloha_static_pro_pencil": "Pick up the pencil with the right arm, hand it over to the left arm then place it back onto the table.", + "aloha_static_screw_driver": "Pick up the screwdriver with the right arm, hand it over to the left arm then place it into the cup.", + "aloha_static_tape": "Cut a small piece of tape from the tape dispenser then place it on the cardboard box's edge.", + "aloha_static_thread_velcro": "Pick up the velcro cable tie with the left arm, then insert the end of the velcro tie into the other end's loop with the right arm.", + "aloha_static_towel": "Pick up a piece of paper towel and place it on the spilled liquid.", + "aloha_static_vinh_cup": "Pick up the platic cup with the right arm, then pop its lid open with the left arm.", + "aloha_static_vinh_cup_left": "Pick up the platic cup with the left arm, then pop its lid open with the right arm.", + "aloha_static_ziploc_slide": "Slide open the ziploc bag.", +} +ALOHA_CONFIG = Path("lerobot/configs/robot/aloha.yaml") -# datasets = [ -# 'lerobot/aloha_mobile_cabinet', -# 'lerobot/aloha_mobile_chair', -# 'lerobot/aloha_mobile_elevator', -# 'lerobot/aloha_mobile_shrimp', -# 'lerobot/aloha_mobile_wash_pan', -# 'lerobot/aloha_mobile_wipe_wine', -# 'lerobot/aloha_sim_insertion_human', -# 'lerobot/aloha_sim_insertion_human_image', -# 'lerobot/aloha_sim_insertion_scripted', -# 'lerobot/aloha_sim_insertion_scripted_image', -# 'lerobot/aloha_sim_transfer_cube_human', -# 'lerobot/aloha_sim_transfer_cube_human_image', -# 'lerobot/aloha_sim_transfer_cube_scripted', -# 'lerobot/aloha_sim_transfer_cube_scripted_image', -# 'lerobot/aloha_static_battery', -# 'lerobot/aloha_static_candy', -# 'lerobot/aloha_static_coffee', -# 'lerobot/aloha_static_coffee_new', -# 'lerobot/aloha_static_cups_open', -# 'lerobot/aloha_static_fork_pick_up', -# 'lerobot/aloha_static_pingpong_test', -# 'lerobot/aloha_static_pro_pencil', -# 'lerobot/aloha_static_screw_driver', -# 'lerobot/aloha_static_tape', -# 'lerobot/aloha_static_thread_velcro', -# 'lerobot/aloha_static_towel', -# 'lerobot/aloha_static_vinh_cup', -# 'lerobot/aloha_static_vinh_cup_left', -# 'lerobot/aloha_static_ziploc_slide', -# 'lerobot/asu_table_top', -# 'lerobot/austin_buds_dataset', -# 'lerobot/austin_sailor_dataset', -# 'lerobot/austin_sirius_dataset', -# 'lerobot/berkeley_autolab_ur5', -# 'lerobot/berkeley_cable_routing', -# 'lerobot/berkeley_fanuc_manipulation', -# 'lerobot/berkeley_gnm_cory_hall', -# 'lerobot/berkeley_gnm_recon', -# 'lerobot/berkeley_gnm_sac_son', -# 'lerobot/berkeley_mvp', -# 'lerobot/berkeley_rpt', -# 'lerobot/cmu_franka_exploration_dataset', -# 'lerobot/cmu_play_fusion', -# 'lerobot/cmu_stretch', -# 'lerobot/columbia_cairlab_pusht_real', -# 'lerobot/conq_hose_manipulation', -# 'lerobot/dlr_edan_shared_control', -# 'lerobot/dlr_sara_grid_clamp', -# 'lerobot/dlr_sara_pour', -# 'lerobot/droid_100', -# 'lerobot/fmb', -# 'lerobot/iamlab_cmu_pickup_insert', -# 'lerobot/imperialcollege_sawyer_wrist_cam', -# 'lerobot/jaco_play', -# 'lerobot/kaist_nonprehensile', -# 'lerobot/nyu_door_opening_surprising_effectiveness', -# 'lerobot/nyu_franka_play_dataset', -# 'lerobot/nyu_rot_dataset', -# 'lerobot/pusht', -# 'lerobot/pusht_image', -# 'lerobot/roboturk', -# 'lerobot/stanford_hydra_dataset', -# 'lerobot/stanford_kuka_multimodal_dataset', -# 'lerobot/stanford_robocook', -# 'lerobot/taco_play', -# 'lerobot/tokyo_u_lsmo', -# 'lerobot/toto', -# 'lerobot/ucsd_kitchen_dataset', -# 'lerobot/ucsd_pick_and_place_dataset', -# 'lerobot/uiuc_d3field', -# 'lerobot/umi_cup_in_the_wild', -# 'lerobot/unitreeh1_fold_clothes', -# 'lerobot/unitreeh1_rearrange_objects', -# 'lerobot/unitreeh1_two_robot_greeting', -# 'lerobot/unitreeh1_warehouse', -# 'lerobot/usc_cloth_sim', -# 'lerobot/utaustin_mutex', -# 'lerobot/utokyo_pr2_opening_fridge', -# 'lerobot/utokyo_pr2_tabletop_manipulation', -# 'lerobot/utokyo_saytap', -# 'lerobot/utokyo_xarm_bimanual', -# 'lerobot/utokyo_xarm_pick_and_place', -# 'lerobot/viola', -# 'lerobot/xarm_lift_medium', -# 'lerobot/xarm_lift_medium_image', -# 'lerobot/xarm_lift_medium_replay', -# 'lerobot/xarm_lift_medium_replay_image', -# 'lerobot/xarm_push_medium', -# 'lerobot/xarm_push_medium_image', -# 'lerobot/xarm_push_medium_replay', -# 'lerobot/xarm_push_medium_replay_image', -# ] +def batch_convert(): + for num, repo_id in enumerate(available_datasets): + print(f"Converting {repo_id} ({num}/{len(available_datasets)})") + name = repo_id.split("/")[1] + single_task, tasks_col, robot_config = None, None, None + if "aloha" in name: + robot_config = parse_robot_config(ALOHA_CONFIG) + if "sim_insertion" in name: + single_task = "Insert the peg into the socket." + elif "sim_transfer" in name: + single_task = "Pick up the cube with the right arm and transfer it to the left arm." + else: + single_task = ALOHA_SINGLE_TASKS_REAL[name] + elif name != "columbia_cairlab_pusht_real" and "pusht" in name: + single_task = "Push the T-shaped block onto the T-shaped target." + elif "xarm_lift" in name or "xarm_push" in name: + single_task = "Pick up the cube and lift it." + else: + tasks_col = "language_instruction" -# convert_dataset(repo_id=repo_id) + convert_dataset( + repo_id=repo_id, + local_dir=LOCAL_DIR, + single_task=single_task, + tasks_col=tasks_col, + robot_config=robot_config, + ) + + +if __name__ == "__main__": + batch_convert() diff --git a/lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py b/lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py index cecab0df..ffd9f05e 100644 --- a/lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py +++ b/lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py @@ -88,6 +88,7 @@ import contextlib import json import math import subprocess +import warnings from pathlib import Path import datasets @@ -414,6 +415,14 @@ def convert_dataset( dataset = datasets.load_dataset("parquet", data_dir=v1x_dir / "data", split="train") keys = get_keys(dataset) + if single_task and "language_instruction" in dataset.column_names: + warnings.warn( + "'single_task' provided but 'language_instruction' tasks_col found. Using 'language_instruction'.", + stacklevel=1, + ) + single_task = None + tasks_col = "language_instruction" + # Episodes episode_indices = sorted(dataset.unique("episode_index")) total_episodes = len(episode_indices) @@ -462,6 +471,10 @@ def convert_dataset( if robot_config is not None: robot_type = robot_config["robot_type"] names = robot_config["names"] + if "observation.effort" in keys["sequence"]: + names["observation.effort"] = names["observation.state"] + if "observation.velocity" in keys["sequence"]: + names["observation.velocity"] = names["observation.state"] repo_tags = [robot_type] else: robot_type = "unknown" @@ -500,10 +513,10 @@ def convert_dataset( convert_stats_to_json(v1x_dir / "meta_data", v20_dir / "meta") #### TODO: delete - repo_id = f"aliberts/{repo_id.split('/')[1]}" + # repo_id = f"aliberts/{repo_id.split('/')[1]}" # if hub_api.repo_exists(repo_id=repo_id, repo_type="dataset"): # hub_api.delete_repo(repo_id=repo_id, repo_type="dataset") - hub_api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True) + # hub_api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True) #### with contextlib.suppress(EntryNotFoundError):