changed naming convention oxe to openx

huggingface · Aug 12, 2024 · a845254 · a845254
1 parent 87b5609
commit a845254
Show file tree

Hide file tree

Showing 9 changed files with 172 additions and 130 deletions.
diff --git a/lerobot/__init__.py b/lerobot/__init__.py
@@ -129,6 +129,47 @@
     "lerobot/unitreeh1_rearrange_objects",
     "lerobot/unitreeh1_two_robot_greeting",
     "lerobot/unitreeh1_warehouse",
+    "lerobot/nyu_rot_dataset",
+    "lerobot/utokyo_saytap",
+    "lerobot/imperialcollege_sawyer_wrist_cam",
+    "lerobot/utokyo_xarm_bimanual",
+    "lerobot/tokyo_u_lsmo",
+    "lerobot/utokyo_pr2_opening_fridge",
+    "lerobot/cmu_franka_exploration_dataset",
+    "lerobot/cmu_stretch",
+    "lerobot/asu_table_top",
+    "lerobot/utokyo_pr2_tabletop_manipulation",
+    "lerobot/utokyo_xarm_pick_and_place",
+    "lerobot/ucsd_kitchen_dataset",
+    "lerobot/austin_buds_dataset",
+    "lerobot/dlr_sara_grid_clamp",
+    "lerobot/conq_hose_manipulation",
+    "lerobot/columbia_cairlab_pusht_real",
+    "lerobot/dlr_sara_pour",
+    "lerobot/dlr_edan_shared_control",
+    "lerobot/ucsd_pick_and_place_dataset",
+    "lerobot/berkeley_cable_routing",
+    "lerobot/nyu_franka_play_dataset",
+    "lerobot/austin_sirius_dataset",
+    "lerobot/cmu_play_fusion",
+    "lerobot/berkeley_gnm_sac_son",
+    "lerobot/nyu_door_opening_surprising_effectiveness",
+    "lerobot/berkeley_fanuc_manipulation",
+    "lerobot/jaco_play",
+    "lerobot/viola",
+    "lerobot/kaist_nonprehensile",
+    "lerobot/berkeley_mvp",
+    "lerobot/uiuc_d3field",
+    "lerobot/berkeley_gnm_recon",
+    "lerobot/austin_sailor_dataset",
+    "lerobot/utaustin_mutex",
+    "lerobot/roboturk",
+    "lerobot/stanford_hydra_dataset",
+    "lerobot/berkeley_autolab_ur5",
+    "lerobot/stanford_robocook",
+    "lerobot/toto",
+    "lerobot/fmb",
+
 ]
 
 available_datasets = list(

diff --git a/lerobot/common/datasets/push_dataset_to_hub/_download_raw.py b/lerobot/common/datasets/push_dataset_to_hub/_download_raw.py
@@ -60,69 +60,78 @@
     "lerobot-raw/aloha_static_vinh_cup_left_raw": "aloha_hdf5",
     "lerobot-raw/aloha_static_vinh_cup_raw": "aloha_hdf5",
     "lerobot-raw/aloha_static_ziploc_slide_raw": "aloha_hdf5",
+    "lerobot-raw/umi_cup_in_the_wild_raw": "umi_zarr",
     "lerobot-raw/pusht_raw": "pusht_zarr",
-    "lerobot-raw/fractal20220817_data": "oxe_rlds.fractal20220817_data",
-    "lerobot-raw/kuka": "oxe_rlds.kuka",
-    "lerobot-raw/bridge_oxe": "oxe_rlds.bridge_oxe",
-    "lerobot-raw/taco_play": "oxe_rlds.taco_play",
-    "lerobot-raw/jaco_play": "oxe_rlds.jaco_play",
-    "lerobot-raw/berkeley_cable_routing": "oxe_rlds.berkeley_cable_routing",
-    "lerobot-raw/roboturk": "oxe_rlds.roboturk",
-    "lerobot-raw/nyu_door_opening_surprising_effectiveness": "oxe_rlds.nyu_door_opening_surprising_effectiveness",
-    "lerobot-raw/viola": "oxe_rlds.viola",
-    "lerobot-raw/berkeley_autolab_ur5": "oxe_rlds.berkeley_autolab_ur5",
-    "lerobot-raw/toto": "oxe_rlds.toto",
-    "lerobot-raw/language_table": "oxe_rlds.language_table",
-    "lerobot-raw/columbia_cairlab_pusht_real": "oxe_rlds.columbia_cairlab_pusht_real",
-    "lerobot-raw/stanford_kuka_multimodal_dataset_converted_externally_to_rlds": "oxe_rlds.stanford_kuka_multimodal_dataset_converted_externally_to_rlds",
-    "lerobot-raw/nyu_rot_dataset_converted_externally_to_rlds": "oxe_rlds.nyu_rot_dataset_converted_externally_to_rlds",
-    "lerobot-raw/io_ai_tech": "oxe_rlds.io_ai_tech",
-    "lerobot-raw/stanford_hydra_dataset_converted_externally_to_rlds": "oxe_rlds.stanford_hydra_dataset_converted_externally_to_rlds",
-    "lerobot-raw/austin_buds_dataset_converted_externally_to_rlds": "oxe_rlds.austin_buds_dataset_converted_externally_to_rlds",
-    "lerobot-raw/nyu_franka_play_dataset_converted_externally_to_rlds": "oxe_rlds.nyu_franka_play_dataset_converted_externally_to_rlds",
-    "lerobot-raw/maniskill_dataset_converted_externally_to_rlds": "oxe_rlds.maniskill_dataset_converted_externally_to_rlds",
-    "lerobot-raw/furniture_bench_dataset_converted_externally_to_rlds": "oxe_rlds.furniture_bench_dataset_converted_externally_to_rlds",
-    "lerobot-raw/cmu_franka_exploration_dataset_converted_externally_to_rlds": "oxe_rlds.cmu_franka_exploration_dataset_converted_externally_to_rlds",
-    "lerobot-raw/ucsd_kitchen_dataset_converted_externally_to_rlds": "oxe_rlds.ucsd_kitchen_dataset_converted_externally_to_rlds",
-    "lerobot-raw/ucsd_pick_and_place_dataset_converted_externally_to_rlds": "oxe_rlds.ucsd_pick_and_place_dataset_converted_externally_to_rlds",
-    "lerobot-raw/spoc": "oxe_rlds.spoc",
-    "lerobot-raw/austin_sailor_dataset_converted_externally_to_rlds": "oxe_rlds.austin_sailor_dataset_converted_externally_to_rlds",
-    "lerobot-raw/austin_sirius_dataset_converted_externally_to_rlds": "oxe_rlds.austin_sirius_dataset_converted_externally_to_rlds",
-    "lerobot-raw/bc_z": "oxe_rlds.bc_z",
-    "lerobot-raw/utokyo_pr2_opening_fridge_converted_externally_to_rlds": "oxe_rlds.utokyo_pr2_opening_fridge_converted_externally_to_rlds",
-    "lerobot-raw/utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds": "oxe_rlds.utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds",
-    "lerobot-raw/utokyo_xarm_pick_and_place_converted_externally_to_rlds": "oxe_rlds.utokyo_xarm_pick_and_place_converted_externally_to_rlds",
-    "lerobot-raw/utokyo_xarm_bimanual_converted_externally_to_rlds": "oxe_rlds.utokyo_xarm_bimanual_converted_externally_to_rlds",
-    "lerobot-raw/robo_net": "oxe_rlds.robo_net",
-    "lerobot-raw/robo_set": "oxe_rlds.robo_set",
-    "lerobot-raw/berkeley_mvp_converted_externally_to_rlds": "oxe_rlds.berkeley_mvp_converted_externally_to_rlds",
-    "lerobot-raw/berkeley_rpt_converted_externally_to_rlds": "oxe_rlds.berkeley_rpt_converted_externally_to_rlds",
-    "lerobot-raw/kaist_nonprehensile_converted_externally_to_rlds": "oxe_rlds.kaist_nonprehensile_converted_externally_to_rlds",
-    "lerobot-raw/stanford_mask_vit_converted_externally_to_rlds": "oxe_rlds.stanford_mask_vit_converted_externally_to_rlds",
-    "lerobot-raw/tokyo_u_lsmo_converted_externally_to_rlds": "oxe_rlds.tokyo_u_lsmo_converted_externally_to_rlds",
-    "lerobot-raw/dlr_sara_pour_converted_externally_to_rlds": "oxe_rlds.dlr_sara_pour_converted_externally_to_rlds",
-    "lerobot-raw/dlr_sara_grid_clamp_converted_externally_to_rlds": "oxe_rlds.dlr_sara_grid_clamp_converted_externally_to_rlds",
-    "lerobot-raw/dlr_edan_shared_control_converted_externally_to_rlds": "oxe_rlds.dlr_edan_shared_control_converted_externally_to_rlds",
-    "lerobot-raw/asu_table_top_converted_externally_to_rlds": "oxe_rlds.asu_table_top_converted_externally_to_rlds",
-    "lerobot-raw/stanford_robocook_converted_externally_to_rlds": "oxe_rlds.stanford_robocook_converted_externally_to_rlds",
-    "lerobot-raw/imperialcollege_sawyer_wrist_cam": "oxe_rlds.imperialcollege_sawyer_wrist_cam",
-    "lerobot-raw/iamlab_cmu_pickup_insert_converted_externally_to_rlds": "oxe_rlds.iamlab_cmu_pickup_insert_converted_externally_to_rlds",
-    "lerobot-raw/uiuc_d3field": "oxe_rlds.uiuc_d3field",
-    "lerobot-raw/utaustin_mutex": "oxe_rlds.utaustin_mutex",
-    "lerobot-raw/berkeley_fanuc_manipulation": "oxe_rlds.berkeley_fanuc_manipulation",
-    "lerobot-raw/cmu_playing_with_food": "oxe_rlds.cmu_playing_with_food",
-    "lerobot-raw/cmu_play_fusion": "oxe_rlds.cmu_play_fusion",
-    "lerobot-raw/cmu_stretch": "oxe_rlds.cmu_stretch",
-    "lerobot-raw/berkeley_gnm_recon": "oxe_rlds.berkeley_gnm_recon",
-    "lerobot-raw/berkeley_gnm_cory_hall": "oxe_rlds.berkeley_gnm_cory_hall",
-    "lerobot-raw/berkeley_gnm_sac_son": "oxe_rlds.berkeley_gnm_sac_son",
-    "lerobot-raw/droid": "oxe_rlds.droid",
-    "lerobot-raw/droid100": "oxe_rlds.droid100",
-    "lerobot-raw/fmb": "oxe_rlds.fmb",
-    "lerobot-raw/dobbe": "oxe_rlds.dobbe",
-    "lerobot-raw/usc_cloth_sim_converted_externally_to_rlds": "oxe_rlds.usc_cloth_sim_converted_externally_to_rlds",
-    "lerobot-raw/plex_robosuite": "oxe_rlds.plex_robosuite",
-    "lerobot-raw/conq_hose_manipulation": "oxe_rlds.conq_hose_manipulation",
+    "lerobot-raw/unitreeh1_fold_clothes_raw": "aloha_hdf5",
+    "lerobot-raw/unitreeh1_rearrange_objects_raw": "aloha_hdf5",
+    "lerobot-raw/unitreeh1_two_robot_greeting_raw": "aloha_hdf5",
+    "lerobot-raw/unitreeh1_warehouse_raw": "aloha_hdf5",
+    "lerobot-raw/xarm_lift_medium_raw": "xarm_pkl",
+    "lerobot-raw/xarm_lift_medium_replay_raw": "xarm_pkl",
+    "lerobot-raw/xarm_push_medium_raw": "xarm_pkl",
+    "lerobot-raw/xarm_push_medium_replay_raw": "xarm_pkl",
+    "lerobot-raw/fractal20220817_data": "openx_rlds.fractal20220817_data",
+    "lerobot-raw/kuka": "openx_rlds.kuka",
+    "lerobot-raw/bridge_openx": "openx_rlds.bridge_openx",
+    "lerobot-raw/taco_play": "openx_rlds.taco_play",
+    "lerobot-raw/jaco_play": "openx_rlds.jaco_play",
+    "lerobot-raw/berkeley_cable_routing": "openx_rlds.berkeley_cable_routing",
+    "lerobot-raw/roboturk": "openx_rlds.roboturk",
+    "lerobot-raw/nyu_door_opening_surprising_effectiveness": "openx_rlds.nyu_door_opening_surprising_effectiveness",
+    "lerobot-raw/viola": "openx_rlds.viola",
+    "lerobot-raw/berkeley_autolab_ur5": "openx_rlds.berkeley_autolab_ur5",
+    "lerobot-raw/toto": "openx_rlds.toto",
+    "lerobot-raw/language_table": "openx_rlds.language_table",
+    "lerobot-raw/columbia_cairlab_pusht_real": "openx_rlds.columbia_cairlab_pusht_real",
+    "lerobot-raw/stanford_kuka_multimodal_dataset_converted_externally_to_rlds": "openx_rlds.stanford_kuka_multimodal_dataset_converted_externally_to_rlds",
+    "lerobot-raw/nyu_rot_dataset_converted_externally_to_rlds": "openx_rlds.nyu_rot_dataset_converted_externally_to_rlds",
+    "lerobot-raw/io_ai_tech": "openx_rlds.io_ai_tech",
+    "lerobot-raw/stanford_hydra_dataset_converted_externally_to_rlds": "openx_rlds.stanford_hydra_dataset_converted_externally_to_rlds",
+    "lerobot-raw/austin_buds_dataset_converted_externally_to_rlds": "openx_rlds.austin_buds_dataset_converted_externally_to_rlds",
+    "lerobot-raw/nyu_franka_play_dataset_converted_externally_to_rlds": "openx_rlds.nyu_franka_play_dataset_converted_externally_to_rlds",
+    "lerobot-raw/maniskill_dataset_converted_externally_to_rlds": "openx_rlds.maniskill_dataset_converted_externally_to_rlds",
+    "lerobot-raw/furniture_bench_dataset_converted_externally_to_rlds": "openx_rlds.furniture_bench_dataset_converted_externally_to_rlds",
+    "lerobot-raw/cmu_franka_exploration_dataset_converted_externally_to_rlds": "openx_rlds.cmu_franka_exploration_dataset_converted_externally_to_rlds",
+    "lerobot-raw/ucsd_kitchen_dataset_converted_externally_to_rlds": "openx_rlds.ucsd_kitchen_dataset_converted_externally_to_rlds",
+    "lerobot-raw/ucsd_pick_and_place_dataset_converted_externally_to_rlds": "openx_rlds.ucsd_pick_and_place_dataset_converted_externally_to_rlds",
+    "lerobot-raw/spoc": "openx_rlds.spoc",
+    "lerobot-raw/austin_sailor_dataset_converted_externally_to_rlds": "openx_rlds.austin_sailor_dataset_converted_externally_to_rlds",
+    "lerobot-raw/austin_sirius_dataset_converted_externally_to_rlds": "openx_rlds.austin_sirius_dataset_converted_externally_to_rlds",
+    "lerobot-raw/bc_z": "openx_rlds.bc_z",
+    "lerobot-raw/utokyo_pr2_opening_fridge_converted_externally_to_rlds": "openx_rlds.utokyo_pr2_opening_fridge_converted_externally_to_rlds",
+    "lerobot-raw/utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds": "openx_rlds.utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds",
+    "lerobot-raw/utokyo_xarm_pick_and_place_converted_externally_to_rlds": "openx_rlds.utokyo_xarm_pick_and_place_converted_externally_to_rlds",
+    "lerobot-raw/utokyo_xarm_bimanual_converted_externally_to_rlds": "openx_rlds.utokyo_xarm_bimanual_converted_externally_to_rlds",
+    "lerobot-raw/robo_net": "openx_rlds.robo_net",
+    "lerobot-raw/robo_set": "openx_rlds.robo_set",
+    "lerobot-raw/berkeley_mvp_converted_externally_to_rlds": "openx_rlds.berkeley_mvp_converted_externally_to_rlds",
+    "lerobot-raw/berkeley_rpt_converted_externally_to_rlds": "openx_rlds.berkeley_rpt_converted_externally_to_rlds",
+    "lerobot-raw/kaist_nonprehensile_converted_externally_to_rlds": "openx_rlds.kaist_nonprehensile_converted_externally_to_rlds",
+    "lerobot-raw/stanford_mask_vit_converted_externally_to_rlds": "openx_rlds.stanford_mask_vit_converted_externally_to_rlds",
+    "lerobot-raw/tokyo_u_lsmo_converted_externally_to_rlds": "openx_rlds.tokyo_u_lsmo_converted_externally_to_rlds",
+    "lerobot-raw/dlr_sara_pour_converted_externally_to_rlds": "openx_rlds.dlr_sara_pour_converted_externally_to_rlds",
+    "lerobot-raw/dlr_sara_grid_clamp_converted_externally_to_rlds": "openx_rlds.dlr_sara_grid_clamp_converted_externally_to_rlds",
+    "lerobot-raw/dlr_edan_shared_control_converted_externally_to_rlds": "openx_rlds.dlr_edan_shared_control_converted_externally_to_rlds",
+    "lerobot-raw/asu_table_top_converted_externally_to_rlds": "openx_rlds.asu_table_top_converted_externally_to_rlds",
+    "lerobot-raw/stanford_robocook_converted_externally_to_rlds": "openx_rlds.stanford_robocook_converted_externally_to_rlds",
+    "lerobot-raw/imperialcollege_sawyer_wrist_cam": "openx_rlds.imperialcollege_sawyer_wrist_cam",
+    "lerobot-raw/iamlab_cmu_pickup_insert_converted_externally_to_rlds": "openx_rlds.iamlab_cmu_pickup_insert_converted_externally_to_rlds",
+    "lerobot-raw/uiuc_d3field": "openx_rlds.uiuc_d3field",
+    "lerobot-raw/utaustin_mutex": "openx_rlds.utaustin_mutex",
+    "lerobot-raw/berkeley_fanuc_manipulation": "openx_rlds.berkeley_fanuc_manipulation",
+    "lerobot-raw/cmu_playing_with_food": "openx_rlds.cmu_playing_with_food",
+    "lerobot-raw/cmu_play_fusion": "openx_rlds.cmu_play_fusion",
+    "lerobot-raw/cmu_stretch": "openx_rlds.cmu_stretch",
+    "lerobot-raw/berkeley_gnm_recon": "openx_rlds.berkeley_gnm_recon",
+    "lerobot-raw/berkeley_gnm_cory_hall": "openx_rlds.berkeley_gnm_cory_hall",
+    "lerobot-raw/berkeley_gnm_sac_son": "openx_rlds.berkeley_gnm_sac_son",
+    "lerobot-raw/droid": "openx_rlds.droid",
+    "lerobot-raw/droid100": "openx_rlds.droid100",
+    "lerobot-raw/fmb": "openx_rlds.fmb",
+    "lerobot-raw/dobbe": "openx_rlds.dobbe",
+    "lerobot-raw/usc_cloth_sim_converted_externally_to_rlds": "openx_rlds.usc_cloth_sim_converted_externally_to_rlds",
+    "lerobot-raw/plex_robosuite": "openx_rlds.plex_robosuite",
+    "lerobot-raw/conq_hose_manipulation": "openx_rlds.conq_hose_manipulation",
 }
 
 

diff --git a/...sets/push_dataset_to_hub/oxe/configs.yaml → ...ts/push_dataset_to_hub/openx/configs.yaml b/...sets/push_dataset_to_hub/oxe/configs.yaml → ...ts/push_dataset_to_hub/openx/configs.yaml
@@ -13,7 +13,7 @@ ActionEncoding:
   EEF_VEL: 5              # EEF XYZ vel (3) + angular vel (3) + Gripper Open/Close (1)
   WHEELED: 6              # For wheeled robots
 
-OXE_DATASET_CONFIGS:
+OPENX_DATASET_CONFIGS:
   fractal20220817_data:
     image_obs_keys:
       - image
@@ -38,7 +38,7 @@ OXE_DATASET_CONFIGS:
     action_encoding: EEF_POS_7
     fps: 10
 
-  bridge_oxe:
+  bridge_openx:
     image_obs_keys:
       - image
     depth_obs_keys:
@@ -264,13 +264,13 @@ OXE_DATASET_CONFIGS:
   furniture_bench_dataset_converted_externally_to_rlds:
     image_obs_keys:
       - image
-      - wrist
+      - wrist_image
     depth_obs_keys:
       - null
     state_obs_keys:
       - state
     state_encoding: POS_QUAT
-    action_encoding: EEF_VEL
+    actionencoding: EEF_VEL
     fps: 10
 
   cmu_franka_exploration_dataset_converted_externally_to_rlds:

diff --git a/...ets/push_dataset_to_hub/oxe/data_utils.py → ...s/push_dataset_to_hub/openx/data_utils.py b/...ets/push_dataset_to_hub/oxe/data_utils.py → ...s/push_dataset_to_hub/openx/data_utils.py
@@ -1,3 +1,18 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the Licens    e.
+# 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.
 """
 NOTE(YL): Adapted from:
     Octo: https://github.com/octo-models/octo/blob/main/octo/data/utils/data_utils.py

diff --git a/...ts/push_dataset_to_hub/oxe/droid_utils.py → .../push_dataset_to_hub/openx/droid_utils.py b/...ts/push_dataset_to_hub/oxe/droid_utils.py → .../push_dataset_to_hub/openx/droid_utils.py
diff --git a/...ets/push_dataset_to_hub/oxe/transforms.py → ...s/push_dataset_to_hub/openx/transforms.py b/...ets/push_dataset_to_hub/oxe/transforms.py → ...s/push_dataset_to_hub/openx/transforms.py
@@ -38,7 +38,7 @@
 
 import tensorflow as tf
 
-from lerobot.common.datasets.push_dataset_to_hub.oxe.data_utils import (
+from lerobot.common.datasets.push_dataset_to_hub.openx.data_utils import (
     binarize_gripper_actions,
     invert_gripper_actions,
     rel2abs_gripper_actions,
@@ -47,12 +47,12 @@
 
 
 def droid_baseact_transform_fn():
-    from lerobot.common.datasets.push_dataset_to_hub.oxe.droid_utils import droid_baseact_transform
+    from lerobot.common.datasets.push_dataset_to_hub.openx.droid_utils import droid_baseact_transform
 
     return droid_baseact_transform
 
 
-def bridge_oxe_dataset_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]:
+def bridge_openx_dataset_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]:
     """
     Applies to version of Bridge V2 in Open X-Embodiment mixture.
 
@@ -783,8 +783,8 @@ def identity_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]:
 
 
 # === Registry ===
-OXE_STANDARDIZATION_TRANSFORMS = {
-    "bridge_oxe": bridge_oxe_dataset_transform,
+OPENX_STANDARDIZATION_TRANSFORMS = {
+    "bridge_openx": bridge_openx_dataset_transform,
     "bridge_orig": bridge_orig_dataset_transform,
     "bridge_dataset": bridge_orig_dataset_transform,
     "ppgm": ppgm_dataset_transform,