From d0d8193625864e9588b1d4a95564e8f37d9294db Mon Sep 17 00:00:00 2001 From: Simon Alibert Date: Thu, 17 Oct 2024 23:33:51 +0200 Subject: [PATCH] Add unitreeh and umi --- .../v2/batch_convert_dataset_v1_to_v2.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 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 73939bd1..00a47473 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 @@ -90,6 +90,7 @@ # 'lerobot/xarm_push_medium_replay', # 'lerobot/xarm_push_medium_replay_image', +import traceback from pathlib import Path from lerobot import available_datasets @@ -139,10 +140,23 @@ def batch_convert(): 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 "unitreeh1" in name: + if "fold_clothes" in name: + single_task = "Fold the sweatshirt." + elif "rearrange_objects" in name or "rearrange_objects" in name: + single_task = "Put the object into the bin." + elif "two_robot_greeting" in name: + single_task = "Greet the other robot with a high five." + elif "warehouse" in name: + single_task = ( + "Grab the spray paint on the shelf and place it in the bin on top of the robot dog." + ) 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." + elif name == "umi_cup_in_the_wild": + single_task = "Put the cup on the plate." else: tasks_col = "language_instruction" @@ -157,8 +171,8 @@ def batch_convert(): status = f"{repo_id}: success." with open(logfile, "a") as file: file.write(status + "\n") - except Exception as e: - status = f"{repo_id}: {e}" + except Exception: + status = f"{repo_id}: failed\n {traceback.format_exc()}" with open(logfile, "a") as file: file.write(status + "\n") continue