add hanger

r-pad · Nov 21, 2023 · 49a5345 · 49a5345
1 parent c65dfb8
commit 49a5345
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 3 deletions.
diff --git a/notebooks/explore_dset2.ipynb b/notebooks/explore_dset2.ipynb
@@ -27,12 +27,13 @@
    "outputs": [],
    "source": [
     "dset = RLBenchPlacementDataset(\n",
-    "    dataset_root=\"/home/beisner/datasets/rlbench/\",\n",
-    "    task_name=\"stack_wine\",\n",
+    "    dataset_root=\"/home/beisner/datasets/rlbench_larger/\",\n",
+    "    # task_name=\"stack_wine\",\n",
     "    # task_name=\"insert_onto_square_peg\",\n",
     "    # task_name=\"insert_usb_in_computer\",\n",
     "    # task_name=\"phone_on_base\",\n",
     "    # task_name=\"put_toilet_roll_on_stand\",\n",
+    "    task_name=\"place_hanger_on_rack\",\n",
     "    n_demos=10,\n",
     "    phase=\"place\",\n",
     ")"
@@ -249,7 +250,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "load_handle_mapping(dset.dataset_root, \"put_toilet_roll_on_stand\", 0)"
+    "load_handle_mapping(\"/home/beisner/datasets/rlbench_larger/\", \"place_hanger_on_rack\", 0)"
    ]
   },
   {

diff --git a/src/rpad/rlbench_utils/placement_dataset.py b/src/rpad/rlbench_utils/placement_dataset.py
@@ -19,6 +19,7 @@
     InsertOntoSquarePeg,
     InsertUsbInComputer,
     PhoneOnBase,
+    PlaceHangerOnRack,
     PutToiletRollOnStand,
     StackWine,
 )
@@ -133,6 +134,33 @@ class StackWinePhase(str, Enum):
             },
         },
     },
+    "place_hanger_on_rack": {
+        "task_class": PlaceHangerOnRack,
+        "phase": {
+            "grasp": {
+                "action_obj_names": [
+                    "Panda_leftfinger_visual",
+                    "Panda_rightfinger_visual",
+                    "Panda_gripper_visual",
+                ],
+                "anchor_obj_names": ["clothes_hanger_visual"],
+                "action_pose_name": "gripper",
+                "anchor_pose_name": "clothes_hanger_visual",
+            },
+            "place": {
+                "action_obj_names": ["clothes_hanger_visual"],
+                "anchor_obj_names": [
+                    "clothes_rack",
+                    "clothes_rack_sub2",
+                    "clothes_rack_sub3",
+                    "hanger_holder",
+                    "hanger_holder0",
+                ],
+                "action_pose_name": "clothes_hanger_visual",
+                "anchor_pose_name": "clothes_rack",
+            },
+        },
+    },
 }