remove gripper from initial observation, and add wrist cam

r-pad · May 21, 2024 · 81dff31 · 81dff31
1 parent 9a4185f
commit 81dff31
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 2 deletions.
diff --git a/configs/dataset/rlbench/_default.yaml b/configs/dataset/rlbench/_default.yaml
@@ -20,6 +20,7 @@ train_dset:
     num_points: 512
     action_mode: "gripper_and_object"
     anchor_mode: "background_robot_removed"
+    include_wrist_cam: True
 
 val_dset:
   demo_dset:
@@ -37,3 +38,4 @@ val_dset:
     num_points: ${...train_dset.demo_dset.num_points}
     action_mode: ${...train_dset.demo_dset.action_mode}
     anchor_mode: ${...train_dset.demo_dset.anchor_mode}
+    include_wrist_cam: ${...train_dset.demo_dset.include_wrist_cam}
diff --git a/pyproject.toml b/pyproject.toml
@@ -60,7 +60,7 @@ rlbench = [
   # These should be installed manually...
   # "pyrep",
   # "rlbench @ git+https://github.com/stepjam/RLBench.git",
-  "rpad-rlbench-utils @ git+https://github.com/r-pad/rlbench_utils.git@46bba6ea62fbacc14a359186ee4c8adf1075a17a",
+  "rpad-rlbench-utils @ git+https://github.com/r-pad/rlbench_utils.git@634f5b617eaad3d12d966778f412bd7a5bdc06a9",
 ]
 
 

diff --git a/scripts/eval_rlbench.py b/scripts/eval_rlbench.py
@@ -551,8 +551,9 @@ def obs_to_input(
     action_mode: ActionMode,
     anchor_mode: AnchorMode,
     n_pts: Optional[int] = 1024,
+    include_wrist_cam=False,
 ):
-    rgb, pc, mask = obs_to_rgb_point_cloud(obs)
+    rgb, pc, mask = obs_to_rgb_point_cloud(obs, include_wrist_cam)
 
     ########################################
     # Separate the action and anchor points.

diff --git a/taxpose/datasets/rlbench.py b/taxpose/datasets/rlbench.py
@@ -261,6 +261,9 @@ class RLBenchPointCloudDatasetConfig:
     anchor_mode: AnchorMode = AnchorMode.SINGLE_OBJECT
     action_mode: ActionMode = ActionMode.OBJECT
 
+    # Whether to include the wrist camera.
+    include_wrist_cam: bool = False
+
 
 class RLBenchPointCloudDataset(Dataset[PlacementPointCloudData]):
     def __init__(self, cfg: RLBenchPointCloudDatasetConfig):
@@ -273,6 +276,7 @@ def __init__(self, cfg: RLBenchPointCloudDatasetConfig):
             use_first_as_init_keyframe=cfg.use_first_as_init_keyframe,
             anchor_mode=cfg.anchor_mode,
             action_mode=cfg.action_mode,
+            include_wrist_cam=cfg.include_wrist_cam,
         )
 
         self.cfg = cfg