add another option

r-pad · Jun 12, 2024 · 8772cc4 · 8772cc4
1 parent 66f8547
commit 8772cc4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/rpad/rlbench_utils/placement_dataset.py b/src/rpad/rlbench_utils/placement_dataset.py
@@ -222,6 +222,7 @@ def filter_out_names(rgb, point_cloud, mask, handlemapping, names=BACKGROUND_NAM
 class ActionMode(str, Enum):
     GRIPPER_AND_OBJECT = "gripper_and_object"
     OBJECT = "object"
+    GRIPPER = "gripper"
 
 
 class AnchorMode(str, Enum):
@@ -308,6 +309,8 @@ def get_action_points(
         action_handles = action_handles + gripper_handles
     elif action_mode == ActionMode.OBJECT:
         pass
+    elif action_mode == ActionMode.GRIPPER:
+        action_handles = gripper_handles
     else:
         raise ValueError("Action mode must be one of the ActionMode enum values.")
 

diff --git a/src/rpad/rlbench_utils/task_info.py b/src/rpad/rlbench_utils/task_info.py
@@ -496,7 +496,6 @@
 SAFE_OBJ_NAMES = ["safe_body", "safe_door"]
 MONEY_OBJ_NAMES = ["dollar_stack0", "dollar_back_visual0"]
 SAFE_POSE_NAME = "safe_body"
-# MONEY_POSE_NAME = "dollar_stack0"
 MONEY_POSE_NAME = "dollar_front_visual0"